Chore: add todo and clean a bit in pipeline step (#5334)

Signed-off-by: Connor Tsui <[email protected]>

Author: connortsui20

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/pipeline/driver/mod.rs

@@ -10,7 +10,7 @@ use std::hash::{BuildHasher, Hash, Hasher};
 
 use itertools::Itertools;
 use vortex_dtype::DType;
-use vortex_error::{VortexResult, vortex_bail};
+use vortex_error::{VortexResult, vortex_ensure};
 use vortex_mask::Mask;
 use vortex_utils::aliases::hash_map::{HashMap, RandomState};
 use vortex_vector::{Vector, VectorMut, VectorMutOps};
@@ -270,7 +270,7 @@ impl Pipeline {
 
     /// Perform a single step of the pipeline.
     fn step(&mut self, selection: &BitView, output: &mut VectorMut) -> VortexResult<()> {
-        // Loop over the kernels in toposorted execution order
+        // Loop over the kernels in toposorted execution order.
         for &node_idx in self.exec_order.iter() {
             let kernel = &mut self.kernels[node_idx];
 
@@ -283,38 +283,43 @@ impl Pipeline {
                     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")
                     }
 
-                    // 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());
 
-                    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);
                 }
             };
         }