feat: pipeline-pcodec

Signed-off-by: Alexander Droste <[email protected]>

Author: 0ax1

Cargo.lock

@@ -22,11 +22,20 @@ pco = { workspace = true }
 prost = { workspace = true }
 vortex-array = { workspace = true }
 vortex-buffer = { workspace = true }
+vortex-compute = { workspace = true }
 vortex-dtype = { workspace = true }
 vortex-error = { workspace = true }
 vortex-mask = { workspace = true }
 vortex-scalar = { workspace = true }
+vortex-vector = { workspace = true }
 
 [dev-dependencies]
+divan = { workspace = true }
+mimalloc = { workspace = true }
+rand = { workspace = true }
 rstest = { workspace = true }
 vortex-array = { workspace = true, features = ["test-harness"] }
+
+[[bench]]
+name = "pco"
+harness = false

encodings/pco/Cargo.toml

@@ -0,0 +1,86 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+#![allow(clippy::unwrap_used)]
+
+use divan::Bencher;
+use mimalloc::MiMalloc;
+use rand::prelude::StdRng;
+use rand::{Rng, SeedableRng};
+use vortex_array::compute::{filter, warm_up_vtables};
+use vortex_array::{IntoArray, ToCanonical};
+use vortex_buffer::{BitBuffer, BufferMut};
+use vortex_mask::Mask;
+use vortex_pco::PcoArray;
+
+#[global_allocator]
+static GLOBAL: MiMalloc = MiMalloc;
+
+pub fn main() {
+    warm_up_vtables();
+    divan::main();
+}
+
+#[divan::bench(args = [
+    (10_000, 0.1),
+    (10_000, 0.5),
+    (10_000, 0.9),
+    (10_000, 1.0),
+    (50_000, 0.1),
+    (50_000, 0.5),
+    (50_000, 0.9),
+    (50_000, 1.0),
+    (100_000, 0.1),
+    (100_000, 0.5),
+    (100_000, 0.9),
+    (100_000, 1.0)]
+)]
+pub fn pco_pipeline(bencher: Bencher, (size, selectivity): (usize, f64)) {
+    let mut rng = StdRng::seed_from_u64(42);
+    let values = (0..size)
+        .map(|i| (i % 10000) as i32)
+        .collect::<BufferMut<i32>>()
+        .into_array()
+        .to_primitive();
+
+    let pco_array = PcoArray::from_primitive(&values, 3, 0).unwrap();
+    let mask = (0..size)
+        .map(|_| rng.random_bool(selectivity))
+        .collect::<BitBuffer>();
+
+    bencher
+        .with_inputs(|| Mask::from_buffer(mask.clone()))
+        .bench_local_values(|mask| pco_array.execute_with_selection(&mask).unwrap());
+}
+
+#[divan::bench(args = [
+    (10_000, 0.1),
+    (10_000, 0.5),
+    (10_000, 0.9),
+    (10_000, 1.0),
+    (50_000, 0.1),
+    (50_000, 0.5),
+    (50_000, 0.9),
+    (50_000, 1.0),
+    (100_000, 0.1),
+    (100_000, 0.5),
+    (100_000, 0.9),
+    (100_000, 1.0)]
+)]
+pub fn pco_canonical(bencher: Bencher, (size, selectivity): (usize, f64)) {
+    let mut rng = StdRng::seed_from_u64(42);
+    let values = (0..size)
+        .map(|i| (i % 10000) as i32)
+        .collect::<BufferMut<i32>>()
+        .into_array()
+        .to_primitive();
+
+    let pco_array = PcoArray::from_primitive(&values, 3, 0).unwrap();
+    let mask = (0..size)
+        .map(|_| rng.random_bool(selectivity))
+        .collect::<BitBuffer>();
+
+    bencher
+        .with_inputs(|| Mask::from_buffer(mask.clone()))
+        .bench_values(|mask| filter(pco_array.to_canonical().as_ref(), &mask).unwrap());
+}

encodings/pco/benches/pco.rs

@@ -13,12 +13,14 @@ use pco::{ChunkConfig, PagingSpec, match_number_enum};
 use prost::Message;
 use vortex_array::arrays::{PrimitiveArray, PrimitiveVTable};
 use vortex_array::compute::filter;
+use vortex_array::pipeline::PipelinedNode;
 use vortex_array::serde::ArrayChildren;
 use vortex_array::stats::{ArrayStats, StatsSetRef};
 use vortex_array::validity::Validity;
 use vortex_array::vtable::{
-    ArrayVTable, CanonicalVTable, EncodeVTable, NotSupported, OperationsVTable, VTable,
-    ValidityHelper, ValiditySliceHelper, ValidityVTableFromValiditySliceHelper, VisitorVTable,
+    ArrayVTable, CanonicalVTable, EncodeVTable, NotSupported, OperationsVTable, OperatorVTable,
+    VTable, ValidityHelper, ValiditySliceHelper, ValidityVTableFromValiditySliceHelper,
+    VisitorVTable,
 };
 use vortex_array::{
     ArrayBufferVisitor, ArrayChildVisitor, ArrayEq, ArrayHash, ArrayRef, Canonical, EncodingId,
@@ -67,7 +69,7 @@ impl VTable for PcoVTable {
     type VisitorVTable = Self;
     type ComputeVTable = NotSupported;
     type EncodeVTable = Self;
-    type OperatorVTable = NotSupported;
+    type OperatorVTable = Self;
 
     fn id(_encoding: &Self::Encoding) -> EncodingId {
         EncodingId::new_ref("vortex.pco")
@@ -129,7 +131,7 @@ impl VTable for PcoVTable {
     }
 }
 
-fn number_type_from_dtype(dtype: &DType) -> NumberType {
+pub(crate) fn number_type_from_dtype(dtype: &DType) -> NumberType {
     let ptype = dtype.as_ptype();
     match ptype {
         PType::F16 => NumberType::F16,
@@ -150,7 +152,7 @@ fn collect_valid(parray: &PrimitiveArray) -> VortexResult<PrimitiveArray> {
     Ok(filter(&parray.to_array(), &mask)?.to_primitive())
 }
 
-fn vortex_err_from_pco(err: PcoError) -> VortexError {
+pub(crate) fn vortex_err_from_pco(err: PcoError) -> VortexError {
     use pco::errors::ErrorKind::*;
     match err.kind {
         Io(io_kind) => VortexError::from(std::io::Error::new(io_kind, err.message)),
@@ -509,6 +511,12 @@ impl VisitorVTable<PcoVTable> for PcoVTable {
     }
 }
 
+impl OperatorVTable<PcoVTable> for PcoVTable {
+    fn pipeline_node(array: &PcoArray) -> Option<&dyn PipelinedNode> {
+        Some(array)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use vortex_array::arrays::PrimitiveArray;

encodings/pco/src/array.rs

@@ -3,6 +3,7 @@
 
 mod array;
 mod compute;
+mod pipeline;
 #[cfg(test)]
 mod test;