Zero-copy pipeline input

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

CLAUDE.md

@@ -4,9 +4,13 @@
 
 * project is a monorepo Rust workspace, java bindings in `/java`, python bindings in `/vortex-python`
 * run `cargo build -p` to build a specific crate
-* use `cargo clippy --all-targets --all-features` to make sure a project is free of lint issues. Please do this every time you reach a stopping point or think you've finished work.
-* run `cargo +nightly fmt --all` to format Rust source files. Please do this every time you reach a stopping point or think you've finished work.
-* you can try running `cargo fix --lib --allow-dirty --allow-staged && cargo clippy --fix --lib --allow-dirty --allow-staged` to automatically many fix minor errors.
+* use `cargo clippy --all-targets --all-features` to make sure a project is free of lint issues. Please do this every
+  time you reach a stopping point or think you've finished work.
+* run `cargo +nightly fmt --all` to format Rust source files. Please do this every time you reach a stopping point or
+  think you've finished work.
+* you can try running
+  `cargo fix --lib --allow-dirty --allow-staged && cargo clippy --fix --lib --allow-dirty --allow-staged` to
+  automatically many fix minor errors.
 
 ## Architecture
 
@@ -31,8 +35,10 @@
 * Use `vortex_err!` to create a `VortexError` with a format string and `vortex_bail!` to do the same but immediately
   return it as a `VortexResult<T>` to the surrounding context.
 * When writing tests, strongly consider using `rstest` cases to parameterize repetitive test logic.
-* If you want to create a large number of tests to an existing file module called `foo.rs`, and if you think doing so would
-  be too many to inline in a `tests` submodule within `foo.rs`, then first promote `foo` to a directory module. You can do
+* If you want to create a large number of tests to an existing file module called `foo.rs`, and if you think doing so
+  would
+  be too many to inline in a `tests` submodule within `foo.rs`, then first promote `foo` to a directory module. You can
+  do
   this by running `mkdir foo && mv foo.rs foo/mod.rs`. Then, you can create a test file `foo/tests.rs` that you include
   in `foo/mod.rs` with the appropriate test config flag.
 * If you encounter clippy errors in tests that should only pertain to production code (e.g., prohibiting panic/unwrap,
@@ -45,3 +51,7 @@
 ## Other
 
 * When summarizing your work, please produce summaries in valid Markdown that can be easily copied/pasted to Github.
+
+## Commits
+
+* All commits must be signed of by the committers in the form `Signed-off-by: "COMMITTER" <COMMITTER_EMAIL>`.

encodings/alp/benches/alp_compress.rs

@@ -123,7 +123,7 @@ fn decompress_alp_pipeline<T: ALPFloat + NativePType>(bencher: Bencher, args: (u
                     validity
                         .as_array()
                         .map(|a| Validity::copy_from_array(a))
-                        .unwrap_or(validity.clone()),
+                        .unwrap_or_else(|| validity.clone()),
                 ),
                 None,
             )

encodings/alp/src/alp/operator.rs

@@ -185,7 +185,7 @@ mod test {
         let values = values.freeze();
 
         // We take a copy of the values to ensure we can into_mut
-        let array = PrimitiveArray::new(Buffer::copy_from(&values), validity.clone());
+        let array = PrimitiveArray::new(Buffer::copy_from(&values), validity);
         let array = alp_encode(&array, None).unwrap().into_array();
 
         let vector = array.execute().unwrap();

encodings/fastlanes/src/bitpacking/array/bitpack_pipeline.rs

@@ -217,6 +217,7 @@ mod tests {
         }
     }
 
+    #[ignore = "TODO(connor): need to filter in pipeline driver step"]
     #[test]
     fn test_bitpack_pipeline_dense_75_percent() {
         // Create exactly 1024 elements (0 to 1023).

vortex-array/src/arrow/convert.rs

@@ -29,7 +29,7 @@ use vortex_error::{VortexExpect as _, vortex_panic};
 use vortex_scalar::i256;
 
 use crate::arrays::{
-    BoolArray, DecimalArray, FixedSizeListArray, ListArray, ListViewArray, NullArray,
+    BoolArray, DecimalArray, DictArray, FixedSizeListArray, ListArray, ListViewArray, NullArray,
     PrimitiveArray, StructArray, TemporalArray, VarBinArray, VarBinViewArray,
 };
 use crate::arrow::FromArrowArray;
@@ -492,6 +492,36 @@ impl FromArrowArray<&dyn ArrowArray> for ArrayRef {
             DataType::Decimal256(..) => {
                 Self::from_arrow(array.as_primitive::<Decimal256Type>(), nullable)
             }
+            DataType::Dictionary(key_type, _) => match key_type.as_ref() {
+                DataType::Int8 => {
+                    DictArray::from_arrow(array.as_dictionary::<Int8Type>(), nullable).into_array()
+                }
+                DataType::Int16 => {
+                    DictArray::from_arrow(array.as_dictionary::<Int16Type>(), nullable).into_array()
+                }
+                DataType::Int32 => {
+                    DictArray::from_arrow(array.as_dictionary::<Int32Type>(), nullable).into_array()
+                }
+                DataType::Int64 => {
+                    DictArray::from_arrow(array.as_dictionary::<Int64Type>(), nullable).into_array()
+                }
+                DataType::UInt8 => {
+                    DictArray::from_arrow(array.as_dictionary::<UInt8Type>(), nullable).into_array()
+                }
+                DataType::UInt16 => {
+                    DictArray::from_arrow(array.as_dictionary::<UInt16Type>(), nullable)
+                        .into_array()
+                }
+                DataType::UInt32 => {
+                    DictArray::from_arrow(array.as_dictionary::<UInt32Type>(), nullable)
+                        .into_array()
+                }
+                DataType::UInt64 => {
+                    DictArray::from_arrow(array.as_dictionary::<UInt64Type>(), nullable)
+                        .into_array()
+                }
+                key_dt => vortex_panic!("Unsupported dictionary key type: {key_dt}"),
+            },
             dt => vortex_panic!("Array encoding not implemented for Arrow data type {dt}"),
         }
     }

vortex-array/src/pipeline/driver/mod.rs

@@ -10,16 +10,16 @@ use std::hash::{BuildHasher, Hash, Hasher};
 
 use itertools::Itertools;
 use vortex_dtype::DType;
-use vortex_error::{vortex_bail, VortexResult};
+use vortex_error::{VortexResult, vortex_ensure};
 use vortex_mask::Mask;
 use vortex_utils::aliases::hash_map::{HashMap, RandomState};
 use vortex_vector::{Vector, VectorMut, VectorMutOps};
 
 use crate::pipeline::bit_view::{BitView, BitViewExt};
-use crate::pipeline::driver::allocation::{allocate_vectors, OutputTarget};
+use crate::pipeline::driver::allocation::{OutputTarget, allocate_vectors};
 use crate::pipeline::driver::bind::bind_kernels;
 use crate::pipeline::driver::toposort::topological_sort;
-use crate::pipeline::{Kernel, KernelCtx, PipelineInputs, N};
+use crate::pipeline::{Kernel, KernelCtx, N, PipelineInputs};
 use crate::{Array, ArrayEq, ArrayHash, ArrayOperator, ArrayRef, ArrayVisitor, Precision};
 
 /// A pipeline driver takes a Vortex array and executes it into a canonical vector.
@@ -272,53 +272,59 @@ impl Pipeline {
 
     /// Perform a single step of the pipeline.
     fn step(&mut self, selection: &BitView, output: &mut VectorMut) -> VortexResult<()> {
-        // Loop over the kernels in execution order.
-        for node_id in &self.exec_order {
-            let kernel = &mut self.kernels[*node_id];
+        // Loop over the kernels in toposorted execution order.
+        for &node_id in self.exec_order.iter() {
+            let kernel = &mut self.kernels[node_id];
 
             // Depending on the output target, either write directly to the pipeline output, or
             // take the intermediate vector and write into that.
-            match &self.output_targets[*node_id] {
+            match &self.output_targets[node_id] {
                 OutputTarget::ExternalOutput => {
                     // We split off the next N elements of capacity from the external output vector.
                     let mut tail = output.split_off(output.len());
                     debug_assert!(tail.is_empty());
 
                     kernel.step(&self.ctx, selection, &mut tail)?;
-                    if tail.len() != N && tail.len() != selection.true_count() {
-                        vortex_bail!(
-                            "Kernel produced incorrect number of output elements, expected either {} or {}, got {}",
-                            N,
-                            selection.true_count(),
-                            tail.len()
-                        );
+
+                    let len = tail.len();
+                    vortex_ensure!(
+                        len == N || len == selection.true_count(),
+                        "Kernel produced incorrect number of output elements, \
+                            expected either {N} or {}, got {len}",
+                        selection.true_count(),
+                    );
+
+                    // Since we are writing to the final vector, there are no other kernels who we
+                    // can delegate filtering the selection mask out to, so check if we need to do
+                    // a final filter before we return.
+                    if selection.true_count() < N && len == N {
+                        // tail.filter(selection_mask)
+                        todo!("Filter via a bit mask")
                     }
 
                     // FIXME(ngates): if true_count < N && tail.len() == N, we need to
                     //  filter-in-place to ensure the pipeline's output vector is contiguous
 
-                    // Now we append the produced output back to the main output vector.
+                    // Now we join the produced output back to the main output vector.
                     output.unsplit(tail);
                 }
                 OutputTarget::IntermediateVector(vector_id) => {
                     let mut out_vector = self.ctx.take_output(vector_id);
                     out_vector.clear();
+                    debug_assert!(out_vector.is_empty());
 
                     kernel.step(&self.ctx, selection, &mut out_vector)?;
 
-                    match out_vector.len() {
-                        // Valid cases are all N elements, or only the selected elements.
-                        n if n == N || n == selection.true_count() => {
-                            // If the kernel added N elements, the output is in-place.
-                            self.ctx.replace_output(vector_id, out_vector);
-                        }
-                        _ => vortex_bail!(
-                            "Kernel produced incorrect number of output elements, expected either {} or {}, got {}",
-                            N,
-                            selection.true_count(),
-                            out_vector.len()
-                        ),
-                    }
+                    let len = out_vector.len();
+                    vortex_ensure!(
+                        len == N || len == selection.true_count(),
+                        "Kernel produced incorrect number of output elements, \
+                            expected either {N} or {}, got {len}",
+                        selection.true_count(),
+                    );
+
+                    // If the kernel added N elements, the output is in-place.
+                    self.ctx.replace_output(vector_id, out_vector);
                 }
             };
         }
@@ -336,7 +342,7 @@ impl Hash for ArrayKey<'_> {
 }
 impl PartialEq for ArrayKey<'_> {
     fn eq(&self, other: &Self) -> bool {
-        self.0.array_eq(&other.0, Precision::Ptr)
+        self.0.array_eq(other.0, Precision::Ptr)
     }
 }
 impl Eq for ArrayKey<'_> {}