diff --git a/Cargo.lock b/Cargo.lock index 43862759c38..57b348280ba 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10447,6 +10447,7 @@ dependencies = [ name = "vortex-pco" version = "0.1.0" dependencies = [ + "num-traits", "pco", "prost 0.14.4", "rstest", diff --git a/encodings/pco/Cargo.toml b/encodings/pco/Cargo.toml index 1683d4875e9..2201ea439e5 100644 --- a/encodings/pco/Cargo.toml +++ b/encodings/pco/Cargo.toml @@ -26,6 +26,7 @@ vortex-mask = { workspace = true } vortex-session = { workspace = true } [dev-dependencies] +num-traits = { workspace = true } rstest = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } vortex-arrow = { workspace = true } diff --git a/encodings/pco/src/array.rs b/encodings/pco/src/array.rs index 7dc3a371bda..e01470caa06 100644 --- a/encodings/pco/src/array.rs +++ b/encodings/pco/src/array.rs @@ -276,13 +276,14 @@ pub(crate) fn number_type_from_ptype(ptype: PType) -> NumberType { PType::F16 => NumberType::F16, PType::F32 => NumberType::F32, PType::F64 => NumberType::F64, + PType::I8 => NumberType::I8, PType::I16 => NumberType::I16, PType::I32 => NumberType::I32, PType::I64 => NumberType::I64, + PType::U8 => NumberType::U8, PType::U16 => NumberType::U16, PType::U32 => NumberType::U32, PType::U64 => NumberType::U64, - _ => unreachable!("PType not supported by Pco: {:?}", ptype), } } @@ -398,7 +399,6 @@ impl Display for PcoData { impl PcoData { /// Validate dtype, validity, slice, and Pco component invariants. pub fn validate(&self, dtype: &DType, len: usize, validity: &Validity) -> VortexResult<()> { - let _ = number_type_from_ptype(self.ptype); vortex_ensure!( dtype.as_ptype() == self.ptype, "expected ptype {}, got {}", @@ -543,7 +543,12 @@ impl PcoData { // perhaps one day we can make this more configurable let chunk_config = ChunkConfig::default() .with_compression_level(level) - .with_paging_spec(PagingSpec::EqualPagesUpTo(values_per_page)); + .with_paging_spec(PagingSpec::EqualPagesUpTo(values_per_page)) + // Pco refuses 8-bit types by default to stop callers compressing symbolic data such + // as UTF-8 text. Vortex only ever reaches here with a numeric primitive array, so the + // guard does not apply; whether Pco is the right scheme for a given 8-bit array is + // decided by the compressor's sampling estimate. + .with_enable_8_bit(true); let values = collect_valid(parray, ctx)?; let n_values = values.len(); diff --git a/encodings/pco/src/tests.rs b/encodings/pco/src/tests.rs index 8a0e1083b16..692e6e95204 100644 --- a/encodings/pco/src/tests.rs +++ b/encodings/pco/src/tests.rs @@ -4,6 +4,8 @@ use std::sync::LazyLock; +use num_traits::NumCast; +use rstest::rstest; use vortex_array::ArrayContext; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; @@ -12,8 +14,10 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::assert_nth_scalar; use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; +use vortex_array::match_each_native_ptype; use vortex_array::serde::SerializeOptions; use vortex_array::serde::SerializedArray; use vortex_array::session::ArraySessionExt; @@ -221,3 +225,80 @@ fn test_serde() -> VortexResult<()> { assert!(pco_arrow == decoded_arrow); Ok(()) } + +/// Round-trip `values` through Pco compression, checking both full decompression and a slice. +fn assert_pco_roundtrip(values: Vec) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = PrimitiveArray::from_iter(values.clone()); + let compressed = Pco::from_primitive(array.as_view(), 3, 0, &mut ctx)?; + + let unsliced_validity = compressed.unsliced_validity(); + let decompressed = compressed.decompress(&unsliced_validity, &mut ctx)?; + assert_arrays_eq!( + decompressed, + PrimitiveArray::from_iter(values.clone()), + &mut ctx + ); + + let slice = compressed.slice(1..values.len() - 1)?; + assert_arrays_eq!( + slice, + PrimitiveArray::from_iter(values[1..values.len() - 1].to_vec()), + &mut ctx + ); + Ok(()) +} + +/// Round-trip a nullable array of `values` with every third entry null. +fn assert_pco_nullable_roundtrip(values: Vec) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let options = values + .iter() + .enumerate() + .map(|(i, v)| (i % 3 != 0).then_some(*v)) + .collect::>(); + let array = PrimitiveArray::from_option_iter(options.clone()); + let compressed = Pco::from_primitive(array.as_view(), 3, 0, &mut ctx)?; + + assert_arrays_eq!( + compressed, + PrimitiveArray::from_option_iter(options), + &mut ctx + ); + Ok(()) +} + +#[test] +fn test_roundtrip_u8() -> VortexResult<()> { + let values: Vec = (0..=u8::MAX).collect(); + assert_pco_roundtrip(values.clone())?; + assert_pco_nullable_roundtrip(values) +} + +#[test] +fn test_roundtrip_i8() -> VortexResult<()> { + let values: Vec = (i8::MIN..=i8::MAX).collect(); + assert_pco_roundtrip(values.clone())?; + assert_pco_nullable_roundtrip(values) +} + +#[rstest] +#[case(PType::U8)] +#[case(PType::U16)] +#[case(PType::U32)] +#[case(PType::U64)] +#[case(PType::I8)] +#[case(PType::I16)] +#[case(PType::I32)] +#[case(PType::I64)] +#[case(PType::F16)] +#[case(PType::F32)] +#[case(PType::F64)] +fn test_roundtrip_each_ptype(#[case] ptype: PType) -> VortexResult<()> { + match_each_native_ptype!(ptype, |T| { + let values = (0..100_u8) + .map(|i| ::from(i).vortex_expect("0..100 fits in every native ptype")) + .collect::>(); + assert_pco_roundtrip(values) + }) +} diff --git a/vortex-btrblocks/src/schemes/integer/pco.rs b/vortex-btrblocks/src/schemes/integer/pco.rs index 675a112d44f..31d8b0774b8 100644 --- a/vortex-btrblocks/src/schemes/integer/pco.rs +++ b/vortex-btrblocks/src/schemes/integer/pco.rs @@ -11,7 +11,6 @@ use vortex_array::IntoArray; use vortex_array::VTable; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; -use vortex_compressor::scheme::EstimateVerdict; use vortex_error::VortexResult; use crate::ArrayAndStats; @@ -38,17 +37,10 @@ impl Scheme for PcoScheme { fn expected_compression_ratio( &self, - data: &ArrayAndStats, + _data: &ArrayAndStats, _compress_ctx: CompressorContext, _exec_ctx: &mut ExecutionCtx, ) -> CompressionEstimate { - use vortex_array::dtype::PType; - - // Pco does not support I8 or U8. - if matches!(data.array_as_primitive().ptype(), PType::I8 | PType::U8) { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); - } - CompressionEstimate::Deferred(DeferredEstimate::Sample) } diff --git a/vortex-btrblocks/tests/snapshots/golden__compact__binary_low_cardinality.snap b/vortex-btrblocks/tests/snapshots/golden__compact__binary_low_cardinality.snap index d68474ce61c..0f4831201d5 100644 --- a/vortex-btrblocks/tests/snapshots/golden__compact__binary_low_cardinality.snap +++ b/vortex-btrblocks/tests/snapshots/golden__compact__binary_low_cardinality.snap @@ -3,9 +3,9 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: binary, len=16384, nbytes=315856 -root: vortex.dict(binary, len=16384) nbytes=6199 +root: vortex.dict(binary, len=16384) nbytes=4840 metadata: all_values_referenced: true - codes: fastlanes.bitpacked(u8, len=16384) nbytes=6144 - metadata: bit_width: 3, offset: 0 + codes: vortex.pco(u8, len=16384) nbytes=4785 + metadata: ptype: u8, nrows: 16384, slice: 0..16384 values: vortex.zstd(binary, len=5) nbytes=55 metadata: nrows: 5, slice: 0..5 diff --git a/vortex-btrblocks/tests/snapshots/golden__compact__int_low_cardinality.snap b/vortex-btrblocks/tests/snapshots/golden__compact__int_low_cardinality.snap index 0cbe2e06814..5683409fc1c 100644 --- a/vortex-btrblocks/tests/snapshots/golden__compact__int_low_cardinality.snap +++ b/vortex-btrblocks/tests/snapshots/golden__compact__int_low_cardinality.snap @@ -3,9 +3,9 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: i64, len=16384, nbytes=131072 -root: vortex.dict(i64, len=16384) nbytes=6177 +root: vortex.dict(i64, len=16384) nbytes=5359 metadata: all_values_referenced: true - codes: fastlanes.bitpacked(u8, len=16384) nbytes=6144 - metadata: bit_width: 3, offset: 0 + codes: vortex.pco(u8, len=16384) nbytes=5326 + metadata: ptype: u8, nrows: 16384, slice: 0..16384 values: vortex.pco(i64, len=6) nbytes=33 metadata: ptype: i64, nrows: 6, slice: 0..6 diff --git a/vortex-btrblocks/tests/snapshots/golden__compact__string_low_cardinality.snap b/vortex-btrblocks/tests/snapshots/golden__compact__string_low_cardinality.snap index a65054ac178..a39c05ca709 100644 --- a/vortex-btrblocks/tests/snapshots/golden__compact__string_low_cardinality.snap +++ b/vortex-btrblocks/tests/snapshots/golden__compact__string_low_cardinality.snap @@ -3,9 +3,9 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: utf8, len=16384, nbytes=262144 -root: vortex.dict(utf8, len=16384) nbytes=8287 +root: vortex.dict(utf8, len=16384) nbytes=7488 metadata: all_values_referenced: true - codes: fastlanes.bitpacked(u8, len=16384) nbytes=8192 - metadata: bit_width: 4, offset: 0 + codes: vortex.pco(u8, len=16384) nbytes=7393 + metadata: ptype: u8, nrows: 16384, slice: 0..16384 values: vortex.zstd(utf8, len=12) nbytes=95 metadata: nrows: 12, slice: 0..12 diff --git a/vortex-btrblocks/tests/snapshots/golden__compact__struct_mixed.snap b/vortex-btrblocks/tests/snapshots/golden__compact__struct_mixed.snap index 83e34c583dc..083e4bcfd24 100644 --- a/vortex-btrblocks/tests/snapshots/golden__compact__struct_mixed.snap +++ b/vortex-btrblocks/tests/snapshots/golden__compact__struct_mixed.snap @@ -3,14 +3,14 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: {id=i64, category=utf8, value=f64}, len=16384, nbytes=524288 -root: vortex.struct({id=i64, category=utf8, value=f64}, len=16384) nbytes=55986 +root: vortex.struct({id=i64, category=utf8, value=f64}, len=16384) nbytes=55187 metadata: id: vortex.sequence(i64, len=16384) nbytes=0 metadata: base: 10000i64, multiplier: 7i64 - category: vortex.dict(utf8, len=16384) nbytes=8287 + category: vortex.dict(utf8, len=16384) nbytes=7488 metadata: all_values_referenced: true - codes: fastlanes.bitpacked(u8, len=16384) nbytes=8192 - metadata: bit_width: 4, offset: 0 + codes: vortex.pco(u8, len=16384) nbytes=7393 + metadata: ptype: u8, nrows: 16384, slice: 0..16384 values: vortex.zstd(utf8, len=12) nbytes=95 metadata: nrows: 12, slice: 0..12 value: vortex.alp(f64, len=16384) nbytes=47699 diff --git a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/mod.rs b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/mod.rs index 830b50450da..b2c4ecb9742 100644 --- a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/mod.rs +++ b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/mod.rs @@ -18,6 +18,7 @@ mod for_; mod fsst; mod patched; mod pco; +mod pco_8bit; mod rle; mod runend; mod sequence; @@ -46,6 +47,7 @@ pub fn fixtures() -> Vec> { // TODO(aduffy): add back once we stabilized Patched array // Box::new(patched::PatchedFixture), Box::new(pco::PcoFixture), + Box::new(pco_8bit::Pco8BitFixture), Box::new(rle::RleFixture), Box::new(runend::RunEndFixture), Box::new(sequence::SequenceFixture), diff --git a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/pco_8bit.rs b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/pco_8bit.rs new file mode 100644 index 00000000000..d448e374c98 --- /dev/null +++ b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/pco_8bit.rs @@ -0,0 +1,64 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex::array::ArrayId; +use vortex::array::ArrayRef; +use vortex::array::ArrayVTable; +use vortex::array::IntoArray; +use vortex::array::arrays::PrimitiveArray; +use vortex::array::arrays::StructArray; +use vortex::array::dtype::FieldNames; +use vortex::array::validity::Validity; +use vortex::encodings::pco::Pco; +use vortex::error::VortexResult; +use vortex_array::ExecutionCtx; + +use super::N; +use crate::fixtures::FlatLayoutFixture; + +/// 8-bit Pco patterns. +/// +/// Kept separate from [`PcoFixture`](super::pco::PcoFixture) because a published fixture's +/// schema is frozen: `check --mode superset` decodes each stored file and compares it against +/// a freshly built one of the same name, so adding fields to an existing fixture fails against +/// every version already in the store. A new file is simply skipped by older versions. +pub struct Pco8BitFixture; + +impl FlatLayoutFixture for Pco8BitFixture { + fn name(&self) -> &str { + "pco_8bit.vortex" + } + + fn description(&self) -> &str { + "8-bit integer patterns for Pco encoding" + } + + fn expected_encodings(&self) -> Vec { + vec![Pco.id()] + } + + fn build(&self, ctx: &mut ExecutionCtx) -> VortexResult { + let gradient_u8: PrimitiveArray = (0..N).map(|i| (i % 251) as u8).collect(); + let saturated_u8: PrimitiveArray = (0..N) + .map(|i| if i % 64 == 0 { u8::MAX } else { 0 }) + .collect(); + let negative_i8: PrimitiveArray = (0..N).map(|i| (-128 + (i % 256) as i32) as i8).collect(); + let nullable_i8 = PrimitiveArray::from_option_iter( + (0..N).map(|i| (i % 5 != 0).then_some((-64 + (i % 129) as i32) as i8)), + ); + + let arr = StructArray::try_new( + FieldNames::from(["gradient_u8", "saturated_u8", "negative_i8", "nullable_i8"]), + vec![ + Pco::from_primitive(gradient_u8.as_view(), 8, 0, ctx)?.into_array(), + Pco::from_primitive(saturated_u8.as_view(), 8, 0, ctx)?.into_array(), + Pco::from_primitive(negative_i8.as_view(), 8, 0, ctx)?.into_array(), + Pco::from_primitive(nullable_i8.as_view(), 8, 0, ctx)?.into_array(), + ], + N, + Validity::NonNullable, + )?; + + Ok(arr.into_array()) + } +}