Skip to content
6 changes: 6 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions encodings/runend/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ workspace = true

[dev-dependencies]
divan = { workspace = true }
insta = { workspace = true }
itertools = { workspace = true }
mimalloc = { workspace = true }
rand = { workspace = true }
Expand Down
3 changes: 3 additions & 0 deletions encodings/runend/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,9 @@ mod iter;
mod kernel;
pub mod ops;
mod rules;
#[cfg(test)]
#[cfg(not(codspeed))]
mod trace_tests;

#[doc(hidden)]
pub mod _benchmarking {
Expand Down
148 changes: 148 additions & 0 deletions encodings/runend/src/trace_tests.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,148 @@
// 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 {
let session = VortexSession::empty().with::<ArraySession>();
// Execute-parent kernels live in the session registry; without registering RunEnd's own
// kernels the traces below would show plain canonicalization instead of the pushdowns.
crate::initialize(&session);
ExecutionCtx::new(session)
}

/// 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]),
&mut execution_ctx()
);
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
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]),
&mut execution_ctx()
);
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 session[0]:execute_parent_fn 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 session[0]:execute_parent_fn 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]),
&mut execution_ctx()
);
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 session[0]:execute_parent_fn 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 session[0]:execute_parent_fn 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(())
}
1 change: 0 additions & 1 deletion vortex-array/src/arrays/scalar_fn/vtable/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -155,7 +155,6 @@ impl VTable for ScalarFn {
}

fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
ctx.log(format_args!("scalar_fn({}): executing", array.scalar_fn()));
let args = VecExecutionArgs::new(array.children(), array.len());
array
.scalar_fn()
Expand Down
Loading
Loading