feat(pco): support 8-bit integers

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 }

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();

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<T: NativePType>(values: Vec<T>) -> 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<T: NativePType>(values: Vec<T>) -> VortexResult<()> {
+    let mut ctx = SESSION.create_execution_ctx();
+    let options = values
+        .iter()
+        .enumerate()
+        .map(|(i, v)| (i % 3 != 0).then_some(*v))
+        .collect::<Vec<_>>();
+    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<u8> = (0..=u8::MAX).collect();
+    assert_pco_roundtrip(values.clone())?;
+    assert_pco_nullable_roundtrip(values)
+}
+
+#[test]
+fn test_roundtrip_i8() -> VortexResult<()> {
+    let values: Vec<i8> = (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| <T as NumCast>::from(i).vortex_expect("0..100 fits in every native ptype"))
+            .collect::<Vec<T>>();
+        assert_pco_roundtrip(values)
+    })
+}

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)
     }
 

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

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

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

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

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<Box<dyn FlatLayoutFixture>> {
         // 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),

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<ArrayId> {
+        vec![Pco.id()]
+    }
+
+    fn build(&self, ctx: &mut ExecutionCtx) -> VortexResult<ArrayRef> {
+        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())
+    }
+}