From b2d9c4d6f7cb7ad1b6daa14b523fe78042a3cde7 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 24 Aug 2026 14:11:23 +0000 Subject: [PATCH 1/2] feat(pco): support 8-bit integers Pco gained u8/i8 support in the 1.0 release, but the Vortex integration still carried the 0.4-era assumption that 8-bit types were unavailable: `number_type_from_ptype` had no I8/U8 arms and would have hit its `unreachable!` arm, and the BtrBlocks integer scheme skipped Pco outright for those ptypes. Map both ptypes onto `NumberType`, which makes the match exhaustive and turns a future ptype addition into a compile error rather than a runtime panic, so the guard call in `PcoData::validate` is no longer needed. Pco additionally refuses 8-bit input unless `enable_8_bit` is set on the `ChunkConfig`, a guard against compressing symbolic data such as UTF-8 text. Vortex only reaches the encoder with numeric primitive arrays, so enable it and leave scheme choice to the compressor's sampling estimate. Dropping the skip lets the compact compressor pick Pco for u8 dictionary codes, where it beats bitpacking on every affected golden entry: int_low_cardinality 6177 -> 5359 bytes string_low_cardinality 8287 -> 7488 bytes binary_low_cardinality 6199 -> 4840 bytes struct_mixed 55986 -> 55187 bytes Signed-off-by: Joe Isaacs --- Cargo.lock | 1 + encodings/pco/Cargo.toml | 1 + encodings/pco/src/array.rs | 11 ++- encodings/pco/src/tests.rs | 81 +++++++++++++++++++ vortex-btrblocks/src/schemes/integer/pco.rs | 10 +-- ...lden__compact__binary_low_cardinality.snap | 6 +- .../golden__compact__int_low_cardinality.snap | 6 +- ...lden__compact__string_low_cardinality.snap | 6 +- .../golden__compact__struct_mixed.snap | 8 +- .../arrays/synthetic/encodings/pco.rs | 8 ++ 10 files changed, 113 insertions(+), 25 deletions(-) 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/pco.rs b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/pco.rs index 356e873d4a9..d8fc20edf99 100644 --- a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/pco.rs +++ b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/pco.rs @@ -60,6 +60,10 @@ impl FlatLayoutFixture for PcoFixture { }) .collect(); let narrow_i16: PrimitiveArray = (0..N as i16).map(|i| (i % 17) - 8).collect(); + let gradient_u8: PrimitiveArray = (0..N).map(|i| (i % 251) as u8).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([ @@ -71,6 +75,8 @@ impl FlatLayoutFixture for PcoFixture { "constant_u16", "spike_outliers", "narrow_i16", + "gradient_u8", + "nullable_i8", ]), vec![ Pco::from_primitive(irregular_i64.as_view(), 8, 0, ctx)?.into_array(), @@ -81,6 +87,8 @@ impl FlatLayoutFixture for PcoFixture { Pco::from_primitive(constant_u16.as_view(), 8, 0, ctx)?.into_array(), Pco::from_primitive(spike_outliers.as_view(), 8, 0, ctx)?.into_array(), Pco::from_primitive(narrow_i16.as_view(), 8, 0, ctx)?.into_array(), + Pco::from_primitive(gradient_u8.as_view(), 8, 0, ctx)?.into_array(), + Pco::from_primitive(nullable_i8.as_view(), 8, 0, ctx)?.into_array(), ], N, Validity::NonNullable, From e95e4836bb6203762178712e52fb70176d9c47fb Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 25 Aug 2026 12:35:47 +0100 Subject: [PATCH 2/2] fix(compat): keep 8-bit fixtures in a new file Adding `gradient_u8` and `nullable_i8` to the existing `pco.vortex` fixture broke `compat.py check --mode last`, which CI runs on every PR. That check downloads the fixtures published by the last release, rebuilds each one from the current fixture definitions, and compares the two with `assert_arrays_eq!`. A published fixture's schema is therefore frozen: the stored v0.85.0 `pco.vortex` has eight fields and the rebuilt one had ten, so the comparison panicked and the job aborted. Move the 8-bit columns into their own `pco_8bit.vortex` fixture. In superset mode a file absent from an older version's manifest is skipped, so the new fixture is inert until it is itself published, and `pco.vortex` returns to its published schema. Signed-off-by: Joe Isaacs --- .../arrays/synthetic/encodings/mod.rs | 2 + .../arrays/synthetic/encodings/pco.rs | 8 --- .../arrays/synthetic/encodings/pco_8bit.rs | 64 +++++++++++++++++++ 3 files changed, 66 insertions(+), 8 deletions(-) create mode 100644 vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/pco_8bit.rs 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.rs b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/pco.rs index d8fc20edf99..356e873d4a9 100644 --- a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/pco.rs +++ b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/pco.rs @@ -60,10 +60,6 @@ impl FlatLayoutFixture for PcoFixture { }) .collect(); let narrow_i16: PrimitiveArray = (0..N as i16).map(|i| (i % 17) - 8).collect(); - let gradient_u8: PrimitiveArray = (0..N).map(|i| (i % 251) as u8).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([ @@ -75,8 +71,6 @@ impl FlatLayoutFixture for PcoFixture { "constant_u16", "spike_outliers", "narrow_i16", - "gradient_u8", - "nullable_i8", ]), vec![ Pco::from_primitive(irregular_i64.as_view(), 8, 0, ctx)?.into_array(), @@ -87,8 +81,6 @@ impl FlatLayoutFixture for PcoFixture { Pco::from_primitive(constant_u16.as_view(), 8, 0, ctx)?.into_array(), Pco::from_primitive(spike_outliers.as_view(), 8, 0, ctx)?.into_array(), Pco::from_primitive(narrow_i16.as_view(), 8, 0, ctx)?.into_array(), - Pco::from_primitive(gradient_u8.as_view(), 8, 0, ctx)?.into_array(), - Pco::from_primitive(nullable_i8.as_view(), 8, 0, ctx)?.into_array(), ], N, Validity::NonNullable, 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()) + } +}