test: trace reductions and executions of scan ops over complex and compressed arrays

encodings/runend/Cargo.toml

@@ -32,6 +32,7 @@ workspace = true
 arrow-array = { workspace = true }
 arrow-schema = { workspace = true }
 divan = { workspace = true }
+insta = { workspace = true }
 itertools = { workspace = true }
 rand = { workspace = true }
 rstest = { workspace = true }

encodings/runend/src/lib.rs

@@ -18,6 +18,9 @@ mod iter;
 mod kernel;
 mod ops;
 mod rules;
+#[cfg(test)]
+#[cfg(not(codspeed))]
+mod trace_tests;
 
 #[doc(hidden)]
 pub mod _benchmarking {

encodings/runend/src/trace_tests.rs

@@ -0,0 +1,149 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Snapshot tests tracing how reductions and executions flow through run-end arrays.
+
+use vortex_array::ArrayRef;
+use vortex_array::Canonical;
+use vortex_array::ExecutionCtx;
+use vortex_array::IntoArray;
+use vortex_array::arrays::BoolArray;
+use vortex_array::arrays::ConstantArray;
+use vortex_array::arrays::DictArray;
+use vortex_array::arrays::FilterArray;
+use vortex_array::arrays::PrimitiveArray;
+use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt;
+use vortex_array::assert_arrays_eq;
+use vortex_array::optimizer::ArrayOptimizer;
+use vortex_array::scalar::Scalar;
+use vortex_array::scalar_fn::fns::binary::Binary;
+use vortex_array::scalar_fn::fns::operators::Operator;
+use vortex_array::session::ArraySession;
+use vortex_array::test_harness::trace::trace_op;
+use vortex_error::VortexResult;
+use vortex_mask::Mask;
+use vortex_session::VortexSession;
+
+use crate::RunEnd;
+
+fn execution_ctx() -> ExecutionCtx {
+    ExecutionCtx::new(VortexSession::empty().with::<ArraySession>())
+}
+
+/// Run-end encoding of `[1, 1, 1, 2, 2, 3, 3, 3, 3]`: ends `[3, 5, 9]`, values `[1, 2, 3]`.
+fn runend_array() -> VortexResult<ArrayRef> {
+    Ok(RunEnd::encode(
+        PrimitiveArray::from_iter([1i32, 1, 1, 2, 2, 3, 3, 3, 3]).into_array(),
+        &mut execution_ctx(),
+    )?
+    .into_array())
+}
+
+#[test]
+fn trace_compare_on_runend() -> VortexResult<()> {
+    let runend = runend_array()?;
+    let rhs = ConstantArray::new(Scalar::from(2i32), runend.len()).into_array();
+    let compared = Binary.try_new_array(runend.len(), Operator::Eq, [runend, rhs])?;
+
+    let traced = trace_op(|| compared.optimize())?;
+    insta::assert_snapshot!(traced.trace.to_string(), @"
+    optimize root=vortex.binary(bool, len=9) session=false
+      reduce_parent static:RunEndScalarFnRule slot=0 parent=vortex.binary(bool, len=9) child=vortex.runend(i32, len=9) -> vortex.runend(bool, len=9)
+      done output=vortex.runend(bool, len=9)
+    ");
+
+    let optimized = traced.output;
+    let traced = trace_op(|| {
+        optimized
+            .execute::<Canonical>(&mut execution_ctx())
+            .map(IntoArray::into_array)
+    })?;
+
+    assert_arrays_eq!(
+        traced.output,
+        BoolArray::from_iter([false, false, false, true, true, false, false, false, false])
+    );
+    insta::assert_snapshot!(traced.trace.to_string(), @"
+    execute_until target=AnyCanonical root=vortex.runend(bool, len=9)
+      iter 0 current=vortex.runend(bool, len=9) builder_active=false
+    execute_until target=AnyCanonical root=vortex.binary(bool, len=3)
+      iter 0 current=vortex.binary(bool, len=3) builder_active=false
+    optimize root=vortex.cast(i32?, len=3) session=false
+      reduce_parent static:CastReduceAdaptor(Primitive) slot=0 parent=vortex.cast(i32?, len=3) child=vortex.primitive(i32, len=3) -> vortex.primitive(i32?, len=3)
+      done output=vortex.primitive(i32?, len=3)
+    optimize root=vortex.slice(i32, len=1) session=false
+      reduce_parent static:SliceReduceAdaptor(Constant) slot=0 parent=vortex.slice(i32, len=1) child=vortex.constant(i32, len=3) -> vortex.constant(i32, len=1)
+      done output=vortex.constant(i32, len=1)
+    optimize root=vortex.cast(i32?, len=1) session=false
+      reduce_parent static:CastReduceAdaptor(Constant) slot=0 parent=vortex.cast(i32?, len=1) child=vortex.constant(i32, len=1) -> vortex.constant(i32?, len=1)
+      done output=vortex.constant(i32?, len=1)
+    execute_until target=AnyCanonical root=vortex.constant(i32?, len=1)
+      iter 0 current=vortex.constant(i32?, len=1) builder_active=false
+        Done array=vortex.primitive(i32?, len=1)
+      iter 1 current=vortex.primitive(i32?, len=1) builder_active=false
+      return output=vortex.primitive(i32?, len=1)
+        Done array=vortex.bool(bool, len=3)
+      iter 1 current=vortex.bool(bool, len=3) builder_active=false
+      return output=vortex.bool(bool, len=3)
+        Done array=vortex.bool(bool, len=9)
+      iter 1 current=vortex.bool(bool, len=9) builder_active=false
+      return output=vortex.bool(bool, len=9)
+    ");
+
+    Ok(())
+}
+
+#[test]
+fn trace_filter_on_runend() -> VortexResult<()> {
+    let runend = runend_array()?;
+    let filtered = FilterArray::try_new(
+        runend,
+        Mask::from_iter([true, false, false, true, true, false, false, true, false]),
+    )?
+    .into_array();
+
+    let traced = trace_op(|| {
+        filtered
+            .execute::<Canonical>(&mut execution_ctx())
+            .map(IntoArray::into_array)
+    })?;
+
+    assert_arrays_eq!(traced.output, PrimitiveArray::from_iter([1i32, 2, 2, 3]));
+    insta::assert_snapshot!(traced.trace.to_string(), @"
+    execute_until target=AnyCanonical root=vortex.filter(i32, len=4)
+      iter 0 current=vortex.filter(i32, len=4) builder_active=false
+        child_execute_parent static:kernel[2] slot=0 parent=vortex.filter(i32, len=4) child=vortex.runend(i32, len=9) -> vortex.dict(i32, len=4)
+      iter 1 current=vortex.dict(i32, len=4) builder_active=false
+        child_execute_parent static:kernel[3] slot=1 parent=vortex.dict(i32, len=4) child=vortex.primitive(i32, len=3) -> vortex.primitive(i32, len=4)
+      iter 2 current=vortex.primitive(i32, len=4) builder_active=false
+      return output=vortex.primitive(i32, len=4)
+    ");
+
+    Ok(())
+}
+
+#[test]
+fn trace_take_on_runend() -> VortexResult<()> {
+    let runend = runend_array()?;
+    // A take is expressed as a `DictArray` whose codes are the take indices.
+    let indices = PrimitiveArray::from_iter([8u64, 0, 4, 4]).into_array();
+    let take = DictArray::try_new(indices, runend)?.into_array();
+
+    let traced = trace_op(|| {
+        take.execute::<Canonical>(&mut execution_ctx())
+            .map(IntoArray::into_array)
+    })?;
+
+    assert_arrays_eq!(traced.output, PrimitiveArray::from_iter([3i32, 1, 2, 2]));
+    insta::assert_snapshot!(traced.trace.to_string(), @"
+    execute_until target=AnyCanonical root=vortex.dict(i32, len=4)
+      iter 0 current=vortex.dict(i32, len=4) builder_active=false
+        child_execute_parent static:kernel[3] slot=1 parent=vortex.dict(i32, len=4) child=vortex.runend(i32, len=9) -> vortex.dict(i32, len=4)
+      iter 1 current=vortex.dict(i32, len=4) builder_active=false
+        child_execute_parent static:kernel[3] slot=1 parent=vortex.dict(i32, len=4) child=vortex.primitive(i32, len=3) -> vortex.primitive(i32, len=4)
+      iter 2 current=vortex.primitive(i32, len=4) builder_active=false
+      return output=vortex.primitive(i32, len=4)
+    ");
+
+    Ok(())
+}