diff --git a/Cargo.lock b/Cargo.lock index 25fb98611d0..13849ea4104 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9515,6 +9515,7 @@ dependencies = [ name = "vortex-btrblocks" version = "0.1.0" dependencies = [ + "arrow-array 58.4.0", "codspeed-divan-compat", "insta", "itertools 0.14.0", @@ -9522,8 +9523,11 @@ dependencies = [ "rand 0.10.2", "rstest", "test-with", + "tpchgen", + "tpchgen-arrow", "vortex-alp", "vortex-array", + "vortex-arrow", "vortex-buffer", "vortex-compressor", "vortex-datetime-parts", @@ -9531,6 +9535,7 @@ dependencies = [ "vortex-error", "vortex-fastlanes", "vortex-fsst", + "vortex-mask", "vortex-onpair", "vortex-pco", "vortex-runend", @@ -10303,6 +10308,7 @@ version = "0.1.0" dependencies = [ "arbitrary", "codspeed-divan-compat", + "insta", "itertools 0.14.0", "mimalloc", "num-traits", diff --git a/encodings/runend/Cargo.toml b/encodings/runend/Cargo.toml index 7ec149b9dd9..5e607b78cc6 100644 --- a/encodings/runend/Cargo.toml +++ b/encodings/runend/Cargo.toml @@ -29,6 +29,7 @@ workspace = true [dev-dependencies] divan = { workspace = true } +insta = { workspace = true } itertools = { workspace = true } mimalloc = { workspace = true } rand = { workspace = true } diff --git a/encodings/runend/src/lib.rs b/encodings/runend/src/lib.rs index ed3ff946d31..b991609c19c 100644 --- a/encodings/runend/src/lib.rs +++ b/encodings/runend/src/lib.rs @@ -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 { diff --git a/encodings/runend/src/trace_tests.rs b/encodings/runend/src/trace_tests.rs new file mode 100644 index 00000000000..96f2afff50a --- /dev/null +++ b/encodings/runend/src/trace_tests.rs @@ -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::(); + // 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 { + 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::(&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::(&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::(&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(()) +} diff --git a/vortex-array/src/arrays/scalar_fn/vtable/mod.rs b/vortex-array/src/arrays/scalar_fn/vtable/mod.rs index 96dedb15713..fa98bdb9683 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/mod.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/mod.rs @@ -155,7 +155,6 @@ impl VTable for ScalarFn { } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - ctx.log(format_args!("scalar_fn({}): executing", array.scalar_fn())); let args = VecExecutionArgs::new(array.children(), array.len()); array .scalar_fn() diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 54af1a8d5f2..4e87ff5e7dd 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -46,6 +46,7 @@ use crate::optimizer::kernels::ParentExecutionKernels; use crate::optimizer::kernels::execute_parent_key; use crate::stats::ArrayStats; use crate::stats::StatsSet; +use crate::trace_op; /// Returns the maximum number of iterations to attempt when executing an array before giving up and returning /// an error, can be by the `VORTEX_MAX_ITERATIONS` env variables, otherwise defaults to 2^22. @@ -163,6 +164,7 @@ impl ArrayRef { /// partially consumes `current_array`: some slots already live in the builder, so a /// parent rewrite would observe inconsistent state and could discard accumulated builder /// data. + #[allow(clippy::cognitive_complexity)] pub fn execute_until(self, ctx: &mut ExecutionCtx) -> VortexResult { let mut current_array = self; let mut current_builder: Option> = None; @@ -171,23 +173,40 @@ impl ArrayRef { let kernels = execute_parent_kernels.as_ref(); let max_iterations = max_iterations(); - for _ in 0..max_iterations { + trace_op!(record_execute_until_start::(¤t_array)); + + for _iteration in 0..max_iterations { + trace_op!(record_execute_until_iteration( + _iteration, + ¤t_array, + stack + .last() + .map(|frame| (&frame.parent_array, frame.slot_idx)), + current_builder.is_some(), + )); + let is_done = stack .last() .map_or(M::matches as DonePredicate, |frame| frame.done); - if is_done(¤t_array) || AnyCanonical::matches(¤t_array) { + let done_target = is_done(¤t_array); + let done_canonical = AnyCanonical::matches(¤t_array); + trace_op!(record_execute_until_done_check(done_target, done_canonical)); + + if done_target || done_canonical { match stack.pop() { None => { debug_assert!( current_builder.is_none(), "root activation should not retain a builder" ); - ctx.log(format_args!("-> {}", current_array)); + trace_op!(record_execute_until_return(¤t_array)); return Ok(current_array); } Some(frame) => { + let _slot_idx = frame.slot_idx; (current_array, current_builder) = pop_frame(frame, current_array)?; + trace_op!(record_execute_until_pop_frame(_slot_idx, ¤t_array)); continue; } } @@ -207,6 +226,7 @@ impl ArrayRef { && let Some(frame) = stack.last() && let Some(result) = { execute_parent_for_child( + "stack_execute_parent", &frame.parent_array, ¤t_array, frame.slot_idx, @@ -215,42 +235,48 @@ impl ArrayRef { )? } { - ctx.log(format_args!( - "execute_parent (stack) rewrote {} -> {}", - current_array, result - )); let frame = stack.pop().vortex_expect("just peeked"); - current_array = result.optimize_ctx(ctx.session())?; + let optimized = result.optimize_ctx(ctx.session())?; + trace_op!(record_execute_optimized(&result, &optimized)); + current_array = optimized; current_builder = frame.parent_builder; continue; } + if current_builder.is_none() && stack.last().is_some() { + trace_op!(record_execute_parent_none( + "stack_execute_parent", + ¤t_array, + )); + } // Step 2b: execute_parent against current_array's own children. if current_builder.is_none() && let Some(rewritten) = try_execute_parent(¤t_array, kernels, ctx)? { - ctx.log(format_args!( - "execute_parent rewrote {} -> {}", - current_array, rewritten - )); - current_array = rewritten.optimize_ctx(ctx.session())?; + let optimized = rewritten.optimize_ctx(ctx.session())?; + trace_op!(record_execute_optimized(&rewritten, &optimized)); + current_array = optimized; continue; } + if current_builder.is_none() { + trace_op!(record_execute_parent_none( + "child_execute_parent", + ¤t_array, + )); + } - // execute step let expected_len = current_array.len(); let expected_dtype = current_array.dtype().clone(); let stats = current_array.statistics().to_array_stats(); let encoding_id = current_array.encoding_id(); + trace_op!(record_execute_encoding(¤t_array)); let result = current_array.execute_encoding_unchecked(ctx)?; let (array, step) = result.into_parts(); match step { ExecutionStep::ExecuteSlot(i, done) => { let (parent, child) = unsafe { array.take_slot_unchecked(i) }?; - ctx.log(format_args!( - "ExecuteSlot({i}): pushing {}, focusing on {}", - parent, child - )); + + trace_op!(record_execute_slot(i, &parent, &child)); stack.push(StackFrame { parent_array: parent, parent_builder: current_builder.take(), @@ -264,6 +290,7 @@ impl ArrayRef { } ExecutionStep::AppendChild(i) => { if current_builder.is_none() { + trace_op!(record_builder_start(&array)); current_builder = Some(builder_with_capacity_in( ctx.allocator(), array.dtype(), @@ -271,10 +298,10 @@ impl ArrayRef { )); } let (parent, child) = unsafe { array.take_slot_unchecked(i) }?; - ctx.log(format_args!( - "AppendChild({i}): appending {} into builder", - child - )); + + trace_op!(record_append_child(i, &parent, &child)); + trace_op!(record_builder_append(&child)); + // TODO(joe)[7674]: replace with a builder kernel registry so we don't // need to go through the VTable append_to_builder indirection. child.append_to_builder( @@ -286,7 +313,8 @@ impl ArrayRef { current_array = parent; } ExecutionStep::Done => { - ctx.log(format_args!("Done: {}", array)); + let had_builder = current_builder.is_some(); + trace_op!(record_execute_done(&array)); (current_array, current_builder) = finalize_done( array, current_builder, @@ -295,6 +323,9 @@ impl ArrayRef { stats, encoding_id, )?; + if had_builder { + trace_op!(record_builder_finish(¤t_array)); + } } } } @@ -419,40 +450,49 @@ impl Drop for ExecutionCtx { /// `AppendChild` is returned. impl Executable for ArrayRef { fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + trace_op!(record_single_step_start(&array)); + if let Some(canonical) = array.as_opt::() { - ctx.log(format_args!("-> canonical {}", array)); - return Ok(Canonical::from(canonical).into_array()); + let output = Canonical::from(canonical).into_array(); + trace_op!(record_single_step_applied("canonical", &array, &output)); + return Ok(output); } + trace_op!(record_single_step_phase_none("canonical", &array)); if let Some(reduced) = array.reduce()? { - ctx.log(format_args!("reduce: rewrote {} -> {}", array, reduced)); reduced.statistics().inherit_from(array.statistics()); + trace_op!(record_single_step_applied("reduce", &array, &reduced)); return Ok(reduced); } + trace_op!(record_single_step_phase_none("reduce", &array)); for (slot_idx, slot) in array.slots().iter().enumerate() { let Some(child) = slot else { continue }; if let Some(reduced_parent) = child.reduce_parent(&array, slot_idx)? { - ctx.log(format_args!( - "reduce_parent: slot[{}]({}) rewrote {} -> {}", - slot_idx, - child.encoding_id(), - array, - reduced_parent - )); reduced_parent.statistics().inherit_from(array.statistics()); + trace_op!(record_single_step_applied( + "reduce_parent", + &array, + &reduced_parent, + )); return Ok(reduced_parent); } } + trace_op!(record_single_step_phase_none("reduce_parent", &array)); let execute_parent_kernels = Arc::clone(&ctx.execute_parent_kernels); let kernels = execute_parent_kernels.as_ref(); for (slot_idx, slot) in array.slots().iter().enumerate() { let Some(child) = slot else { continue }; - if let Some(executed_parent) = - execute_parent_for_child(&array, child, slot_idx, kernels, ctx)? - { + if let Some(executed_parent) = execute_parent_for_child( + "single_step_execute_parent", + &array, + child, + slot_idx, + kernels, + ctx, + )? { ctx.log(format_args!( "execute_parent: slot[{}]({}) rewrote {} -> {}", slot_idx, @@ -463,16 +503,22 @@ impl Executable for ArrayRef { executed_parent .statistics() .inherit_from(array.statistics()); + trace_op!(record_single_step_applied( + "execute_parent", + &array, + &executed_parent, + )); return Ok(executed_parent); } } + trace_op!(record_single_step_phase_none("execute_parent", &array)); + trace_op!(record_execute_encoding(&array)); - ctx.log(format_args!("executing {}", array)); let result = array.execute_encoding(ctx)?; let (array, step) = result.into_parts(); match step { ExecutionStep::Done => { - ctx.log(format_args!("-> {}", array)); + trace_op!(record_execute_done(&array)); Ok(array) } ExecutionStep::ExecuteSlot(i, _) => { @@ -484,9 +530,12 @@ impl Executable for ArrayRef { } ExecutionStep::AppendChild(_) => { // Single-step: build the entire parent via the builder path. + trace_op!(record_builder_start(&array)); let builder = builder_with_capacity_in(ctx.allocator(), array.dtype(), array.len()); let mut builder = execute_into_builder(array, builder, ctx)?; - Ok(builder.finish()) + let output = builder.finish(); + trace_op!(record_builder_finish(&output)); + Ok(output) } } } @@ -562,6 +611,7 @@ fn finalize_done( } fn execute_parent_for_child( + _phase: &'static str, parent: &ArrayRef, child: &ArrayRef, slot_idx: usize, @@ -570,7 +620,8 @@ fn execute_parent_for_child( ) -> VortexResult> { let key = execute_parent_key(parent.encoding_id(), child.encoding_id()); if let Some(plugins) = kernels.get(&key) { - for plugin in plugins.as_ref() { + #[allow(clippy::unused_enumerate_index)] + for (_plugin_idx, plugin) in plugins.as_ref().iter().enumerate() { if let Some(result) = plugin.execute_parent(child, parent, slot_idx, ctx)? { if cfg!(debug_assertions) { vortex_ensure!( @@ -582,8 +633,23 @@ fn execute_parent_for_child( "Executed parent canonical dtype mismatch" ); } + trace_op!(record_session_execute_parent_applied( + _phase, + parent, + child, + slot_idx, + _plugin_idx, + &result, + )); return Ok(Some(result)); } + trace_op!(record_session_execute_parent_declined( + _phase, + parent, + child, + slot_idx, + _plugin_idx, + )); } } @@ -599,7 +665,7 @@ fn try_execute_parent( for (slot_idx, slot) in array.slots().iter().enumerate() { let Some(child) = slot else { continue }; if let Some(executed_parent) = - execute_parent_for_child(array, child, slot_idx, kernels, ctx)? + execute_parent_for_child("child_execute_parent", array, child, slot_idx, kernels, ctx)? { ctx.log(format_args!( "execute_parent: slot[{}]({}) rewrote {} -> {}", @@ -709,8 +775,9 @@ impl ExecutionResult { /// /// The provided array is the (possibly modified) parent that still needs its slot executed. pub fn execute_slot(array: impl IntoArray, slot_idx: usize) -> Self { + let array = array.into_array(); Self { - array: array.into_array(), + array, step: ExecutionStep::ExecuteSlot(slot_idx, M::matches), } } @@ -719,8 +786,9 @@ impl ExecutionResult { /// activation's canonical builder, and leave the returned parent as the next /// `current_array`. pub fn append_child(array: impl IntoArray, slot_idx: usize) -> Self { + let array = array.into_array(); Self { - array: array.into_array(), + array, step: ExecutionStep::AppendChild(slot_idx), } } diff --git a/vortex-array/src/lib.rs b/vortex-array/src/lib.rs index ea61c18dfd8..9439dc7dd1b 100644 --- a/vortex-array/src/lib.rs +++ b/vortex-array/src/lib.rs @@ -105,6 +105,8 @@ use crate::stats::session::StatsSession; pub mod aggregate_fn; #[doc(hidden)] pub mod aliases; +mod trace_macros; +pub(crate) use trace_macros::trace_op; mod array; pub mod arrays; pub mod buffer; diff --git a/vortex-array/src/optimizer/mod.rs b/vortex-array/src/optimizer/mod.rs index 9b51a468ee8..41a6bb69b25 100644 --- a/vortex-array/src/optimizer/mod.rs +++ b/vortex-array/src/optimizer/mod.rs @@ -23,6 +23,7 @@ use vortex_session::VortexSession; use crate::ArrayRef; use crate::optimizer::kernels::ArrayKernelsExt; +use crate::trace_op; pub mod kernels; pub mod rules; @@ -73,22 +74,24 @@ fn try_optimize( let mut any_optimizations = false; let array_ref = session.map(|s| s.kernels()); + trace_op!(record_optimize_start(array, session.is_some())); + // Apply reduction rules to the current array until no more rules apply. - let mut loop_counter = 0; - 'outer: loop { - if loop_counter > 100 { - vortex_bail!("Exceeded maximum optimization iterations (possible infinite loop)"); - } - loop_counter += 1; + for _ in 0..=100 { + trace_op!(record_optimize_loop_start(¤t_array)); if let Some(new_array) = current_array.reduce()? { current_array = new_array; any_optimizations = true; + trace_op!(record_optimize_loop_end()); continue; } + trace_op!(record_optimize_reduce_none(¤t_array)); + // Apply parent reduction rules to each slot in the context of the current array. // Its important to take all slots here, as `current_array` can change inside the loop. + let mut parent_reduced = None; for (slot_idx, slot) in current_array.slots().iter().enumerate() { let Some(child) = slot else { continue }; @@ -97,34 +100,59 @@ fn try_optimize( && let Some(plugins) = array_ref.find_reduce_parent(current_array.encoding_id(), child.encoding_id()) { - for plugin in plugins.as_ref() { + #[allow(clippy::unused_enumerate_index)] + for (_plugin_idx, plugin) in plugins.as_ref().iter().enumerate() { if let Some(new_array) = plugin(child, ¤t_array, slot_idx)? { - current_array = new_array; - any_optimizations = true; - continue 'outer; + trace_op!(record_session_parent_reduce_applied( + ¤t_array, + child, + slot_idx, + _plugin_idx, + &new_array, + )); + parent_reduced = Some(new_array); + break; } + trace_op!(record_session_parent_reduce_declined( + ¤t_array, + child, + slot_idx, + _plugin_idx, + )); + } + if parent_reduced.is_some() { + break; } } if let Some(new_array) = child.reduce_parent(¤t_array, slot_idx)? { - // If the parent was replaced, then we attempt to reduce it again. - current_array = new_array; - any_optimizations = true; - - // Continue to the start of the outer loop - continue 'outer; + parent_reduced = Some(new_array); + break; } } + if let Some(new_array) = parent_reduced { + // If the parent was replaced, then we attempt to reduce it again. + current_array = new_array; + any_optimizations = true; + trace_op!(record_optimize_loop_end()); + continue; + } + + trace_op!(record_optimize_parent_reduce_none(¤t_array)); + trace_op!(record_optimize_loop_end()); + // No more optimizations can be applied - break; - } + trace_op!(record_optimize_done(¤t_array, any_optimizations)); - if any_optimizations { - Ok(Some(current_array)) - } else { - Ok(None) + if any_optimizations { + return Ok(Some(current_array)); + } else { + return Ok(None); + } } + + vortex_bail!("Exceeded maximum optimization iterations (possible infinite loop)"); } fn try_optimize_recursive( @@ -134,6 +162,8 @@ fn try_optimize_recursive( let mut current_array = array.clone(); let mut any_optimizations = false; + trace_op!(record_optimize_recursive_start(array)); + if let Some(new_array) = try_optimize(¤t_array, Some(session))? { current_array = new_array; any_optimizations = true; @@ -145,6 +175,11 @@ fn try_optimize_recursive( match slot { Some(child) => { if let Some(new_child) = try_optimize_recursive(child, session)? { + trace_op!(record_optimize_recursive_slot( + new_slots.len(), + child, + &new_child, + )); new_slots.push(Some(new_child)); any_slot_optimized = true; } else { diff --git a/vortex-array/src/optimizer/rules.rs b/vortex-array/src/optimizer/rules.rs index e505b21a199..8539e39aea2 100644 --- a/vortex-array/src/optimizer/rules.rs +++ b/vortex-array/src/optimizer/rules.rs @@ -28,6 +28,7 @@ use crate::ArrayRef; use crate::array::ArrayView; use crate::array::VTable; use crate::matcher::Matcher; +use crate::trace_op; /// A metadata-only rewrite rule that transforms an array based on its own structure (Layer 1). /// @@ -133,8 +134,10 @@ impl ReduceRuleSet { pub fn evaluate(&self, array: ArrayView<'_, V>) -> VortexResult> { for rule in self.rules.iter() { if let Some(reduced) = rule.reduce(array)? { + trace_op!(record_reduce_applied(array.array(), *rule, &reduced)); return Ok(Some(reduced)); } + trace_op!(record_reduce_declined(array.array(), *rule)); } Ok(None) } @@ -176,6 +179,12 @@ impl ParentRuleSet { ) -> VortexResult> { for rule in self.rules.iter() { if !rule.matches(parent) { + trace_op!(record_static_parent_reduce_no_match( + parent, + child.array(), + child_idx, + *rule, + )); continue; } if let Some(reduced) = rule.reduce_parent(child, parent, child_idx)? { @@ -198,8 +207,21 @@ impl ParentRuleSet { ); } + trace_op!(record_static_parent_reduce_applied( + parent, + child.array(), + child_idx, + *rule, + &reduced, + )); return Ok(Some(reduced)); } + trace_op!(record_static_parent_reduce_declined( + parent, + child.array(), + child_idx, + *rule, + )); } Ok(None) } diff --git a/vortex-array/src/test_harness.rs b/vortex-array/src/test_harness/mod.rs similarity index 97% rename from vortex-array/src/test_harness.rs rename to vortex-array/src/test_harness/mod.rs index 98be3a85ce4..d3a6b829e80 100644 --- a/vortex-array/src/test_harness.rs +++ b/vortex-array/src/test_harness/mod.rs @@ -12,6 +12,9 @@ use crate::ExecutionCtx; use crate::arrays::BoolArray; use crate::arrays::bool::BoolArrayExt; +#[cfg(not(codspeed))] +pub mod trace; + /// Check that a named metadata matches its previous versioning. /// /// Goldenfile takes care of checking for equality against a checked-in file. diff --git a/vortex-array/src/test_harness/trace/mod.rs b/vortex-array/src/test_harness/trace/mod.rs new file mode 100644 index 00000000000..709259f0bcb --- /dev/null +++ b/vortex-array/src/test_harness/trace/mod.rs @@ -0,0 +1,1257 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Snapshot-friendly tracing harness for the optimizer and executor. +//! +//! # What this records +//! +//! [`trace_op`] runs a closure with a thread-local recorder installed. While the recorder is +//! active, calls to the `trace_op!` macro inside the optimizer and executor push +//! structured events into the recorder. The recorder produces a [`TraceDisplay`] that renders +//! as a deterministic, hierarchical text trace suitable for `insta` snapshot assertions. +//! +//! Events cover: +//! +//! - **Optimization**: optimize/recursive-optimize entry, fixpoint loop iterations, applied +//! reduce rules, applied parent-reduce rules. +//! - **Execution**: `execute_until` iterations, single-step entries, parent kernel attempts and +//! matches, slot transitions, builder start/append/finish, and the eventual canonical output. +//! +//! Despite the name `trace_op`, the harness is *not* a generic logging facility: it is closely +//! coupled to the optimizer/executor state machines so that the resulting trace is stable enough +//! to commit as a snapshot. +//! +//! # When to use it +//! +//! Use [`trace_op`] to write a regression test that asserts on the sequence of optimizer +//! rewrites or executor steps an array goes through. Typical scenarios: +//! +//! - A reduce rule should fire exactly once on a specific input shape. +//! - A parent kernel should be tried in a specific order and the first match should win. +//! - The executor should walk into a slot, finish it, and pop back to the parent without +//! building a canonical intermediate. +//! - A chunked array should drive the builder path rather than the stack path. +//! +//! Two resolutions are available: +//! +//! - [`TraceResolution::ExecutedOnly`] (default) — only events that actually fired (rule +//! rewrites that matched, kernels that succeeded, execution steps that ran). Optimizer +//! passes that produced no change are elided. +//! - [`TraceResolution::Attempts`] — also records declined rule attempts, kernels that did +//! not match, and per-loop bookkeeping. Use this when ordering or fall-through matters. +//! +//! # Cost and scope +//! +//! - Capture is thread-local. Worker threads spawned inside `f` do not inherit the recorder. +//! - Nested captures return an error so that unrelated traces never merge. +//! - In release builds and CodSpeed benchmark builds, every `trace_op!` invocation is compiled +//! away by the macro's `cfg` gating; this module is then unused. See +//! `trace_op!` for the gating rules. +//! +//! # Example +//! +//! ```ignore +//! use vortex_array::test_harness::trace::trace_op; +//! +//! let traced = trace_op(|| filter_array.optimize())?; +//! assert!(traced.output.is::()); +//! insta::assert_snapshot!(traced.trace.to_string(), @r" +//! optimize root=vortex.filter(i32, len=4) session=false +//! reduce TrivialFilterRule: vortex.filter(i32, len=4) -> vortex.primitive(i32, len=4) +//! done output=vortex.primitive(i32, len=4) +//! "); +//! ``` + +use std::cell::Cell; +use std::cell::RefCell; +use std::fmt; +use std::fmt::Debug; +use std::fmt::Display; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::ArrayRef; + +/// Controls how much rule and kernel resolution detail is captured. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub enum TraceResolution { + /// Record only the operations that actually executed. + #[default] + ExecutedOnly, + /// Also record rule and kernel attempts that matched but declined, or did not match. + Attempts, +} + +/// Options for [`trace_op_with`]. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub struct TraceOptions { + /// The amount of rule and kernel resolution detail to include. + pub resolution: TraceResolution, +} + +/// The result of a traced operation. +#[derive(Clone, Debug)] +pub struct Traced { + /// The value returned by the traced closure. + pub output: T, + /// A stable, snapshot-friendly rendering of optimizer and execution activity. + pub trace: TraceDisplay, +} + +/// A stable, snapshot-friendly trace. +#[derive(Clone, Debug, Default)] +pub struct TraceDisplay { + options: TraceOptions, + events: Vec, +} + +impl Display for TraceDisplay { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let hidden_events = self.hidden_events(); + let mut optimize_depth = 0usize; + let mut wrote_event = false; + + for (idx, event) in self.events.iter().enumerate() { + if hidden_events[idx] { + continue; + } + + if event.closes_before(self.options.resolution) { + optimize_depth = optimize_depth.saturating_sub(1); + } + + if event.is_hidden(self.options.resolution) { + continue; + } + + if wrote_event { + writeln!(f)?; + } else { + wrote_event = true; + } + + write_indent( + f, + optimize_depth + event.relative_indent(self.options.resolution, optimize_depth > 0), + )?; + event.fmt_line(f, self.options.resolution)?; + + if event.opens_after(self.options.resolution) { + optimize_depth += 1; + } + if event.closes_after(self.options.resolution) { + optimize_depth = optimize_depth.saturating_sub(1); + } + } + Ok(()) + } +} + +impl TraceDisplay { + fn hidden_events(&self) -> Vec { + let mut hidden = vec![false; self.events.len()]; + if self.options.resolution != TraceResolution::ExecutedOnly { + return hidden; + } + + let mut optimize_stack = Vec::new(); + for (idx, event) in self.events.iter().enumerate() { + match event { + TraceEvent::OptimizeStart { .. } => optimize_stack.push(idx), + TraceEvent::OptimizeDone { changed, .. } => { + let Some(start) = optimize_stack.pop() else { + continue; + }; + if !changed { + hidden[start..=idx].fill(true); + } + } + _ => {} + } + } + hidden + } +} + +fn write_indent(f: &mut fmt::Formatter<'_>, depth: usize) -> fmt::Result { + for _ in 0..depth { + f.write_str(" ")?; + } + Ok(()) +} + +/// Run `f` while capturing a trace of the optimizer and executor work it performs. +/// +/// `f` typically invokes an operation that drives the executor or optimizer, such as +/// [`ArrayOptimizer::optimize`][crate::optimizer::ArrayOptimizer::optimize] or +/// [`ArrayRef::execute`][crate::ArrayRef::execute]. While `f` runs, +/// the optimizer and executor emit structured events via the `trace_op!` +/// macro into a thread-local recorder. When `f` returns, the recorder is finalized and +/// returned alongside the closure's output as a [`Traced`]. +/// +/// The default resolution ([`TraceResolution::ExecutedOnly`]) records the rule rewrites, +/// parent kernels, execution steps, and builder activity that actually executed. Optimizer +/// passes that produced no change are hidden from the rendered trace. Use [`trace_op_with`] +/// with [`TraceResolution::Attempts`] when a test needs to assert on declined rule attempts, +/// kernels that did not match, or other fall-through detail. +/// +/// # Examples +/// +/// ```ignore +/// let traced = trace_op(|| filter_array.optimize())?; +/// assert!(traced.output.is::()); +/// insta::assert_snapshot!(traced.trace.to_string(), @r" +/// optimize root=vortex.filter(i32, len=4) session=false +/// reduce TrivialFilterRule: vortex.filter(i32, len=4) -> vortex.primitive(i32, len=4) +/// done output=vortex.primitive(i32, len=4) +/// "); +/// ``` +/// +/// # Errors +/// +/// Returns whatever error `f` produces. Returns an error if a recorder is already active on +/// the current thread — nested traces are not supported. +pub fn trace_op(f: impl FnOnce() -> VortexResult) -> VortexResult> { + trace_op_with(TraceOptions::default(), f) +} + +/// Run `f` while capturing a trace using `options`. +/// +/// See [`trace_op`] for the common case. Use this when you need to override the default +/// [`TraceResolution`] to capture declined rules and unmatched kernels. +/// +/// Trace capture is thread-local and intentionally does not propagate to worker threads. Nested +/// trace captures return an error so tests do not accidentally merge unrelated traces. +pub fn trace_op_with( + options: TraceOptions, + f: impl FnOnce() -> VortexResult, +) -> VortexResult> { + let interest = TraceInterest::from(options.resolution); + ACTIVE_TRACE.with(|active| { + let mut active = active.borrow_mut(); + if active.is_some() { + return Err(vortex_err!("trace_op captures cannot be nested")); + } + *active = Some(TraceRecorder::new(options)); + Ok(()) + })?; + TRACE_INTEREST.with(|trace_interest| trace_interest.set(interest)); + ACTIVE_TRACE_COUNT.fetch_add(1, Ordering::Relaxed); + if interest == TraceInterest::Attempts { + ATTEMPTS_TRACE_COUNT.fetch_add(1, Ordering::Relaxed); + } + + let guard = ActiveTraceGuard { interest }; + let output = f(); + let recorder = ACTIVE_TRACE.with(|active| { + active + .borrow_mut() + .take() + .vortex_expect("trace recorder must be installed") + }); + drop(guard); + + output.map(|output| Traced { + output, + trace: recorder.finish(), + }) +} + +/// Returns true when the current thread has an active trace recorder. +#[inline(always)] +pub(crate) fn is_active() -> bool { + if ACTIVE_TRACE_COUNT.load(Ordering::Relaxed) == 0 { + return false; + } + TRACE_INTEREST.with(|interest| interest.get().is_active()) +} + +#[inline(always)] +fn attempts_enabled() -> bool { + if ATTEMPTS_TRACE_COUNT.load(Ordering::Relaxed) == 0 { + return false; + } + TRACE_INTEREST.with(|interest| interest.get() == TraceInterest::Attempts) +} + +#[derive(Clone, Copy, Debug)] +enum TraceSource { + Static, + Session(usize), +} + +impl Display for TraceSource { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + TraceSource::Static => f.write_str("static"), + TraceSource::Session(idx) => write!(f, "session[{idx}]"), + } + } +} + +#[derive(Clone, Copy, Debug)] +enum AttemptOutcome { + Declined, + NoMatch, +} + +impl Display for AttemptOutcome { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + AttemptOutcome::Declined => f.write_str("declined"), + AttemptOutcome::NoMatch => f.write_str("no-match"), + } + } +} + +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +enum TraceInterest { + #[default] + Off, + ExecutedOnly, + Attempts, +} + +impl TraceInterest { + #[inline] + fn is_active(self) -> bool { + self != Self::Off + } +} + +impl From for TraceInterest { + fn from(resolution: TraceResolution) -> Self { + match resolution { + TraceResolution::ExecutedOnly => Self::ExecutedOnly, + TraceResolution::Attempts => Self::Attempts, + } + } +} + +/// Snapshot-friendly wrapper around [`ArrayRef`] that renders the encoding, dtype, and length +/// using the canonical [`Display`] format (`vortex.primitive(i32, len=4)`). +/// +/// Carries a clone of the [`ArrayRef`] (a cheap [`Arc`] bump) instead of duplicating fields, +/// so trace events stay small and share the same rendering as every other `{array}` print in +/// the codebase. +#[derive(Clone, Debug)] +pub(crate) struct ArraySummary(ArrayRef); + +impl ArraySummary { + pub(crate) fn new(array: &ArrayRef) -> Self { + Self(array.clone()) + } +} + +impl Display for ArraySummary { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + Display::fmt(&self.0, f) + } +} + +pub(crate) fn record_optimize_start(root: &ArrayRef, session: bool) { + record(TraceEvent::OptimizeStart { + root: ArraySummary::new(root), + session, + }); +} + +pub(crate) fn record_optimize_loop_start(array: &ArrayRef) { + if !attempts_enabled() { + return; + } + record(TraceEvent::OptimizeLoopStart { + array: ArraySummary::new(array), + }); +} + +pub(crate) fn record_optimize_loop_end() { + if !attempts_enabled() { + return; + } + record(TraceEvent::OptimizeLoopEnd); +} + +pub(crate) fn record_optimize_reduce_none(array: &ArrayRef) { + if !attempts_enabled() { + return; + } + record(TraceEvent::PhaseNone { + indent: 0, + phase: "reduce", + subject: "array", + array: ArraySummary::new(array), + }); +} + +pub(crate) fn record_optimize_parent_reduce_none(array: &ArrayRef) { + if !attempts_enabled() { + return; + } + record(TraceEvent::PhaseNone { + indent: 0, + phase: "reduce_parent", + subject: "array", + array: ArraySummary::new(array), + }); +} + +pub(crate) fn record_optimize_done(output: &ArrayRef, changed: bool) { + record(TraceEvent::OptimizeDone { + output: ArraySummary::new(output), + changed, + }); +} + +pub(crate) fn record_optimize_recursive_start(root: &ArrayRef) { + record(TraceEvent::OptimizeRecursiveStart { + root: ArraySummary::new(root), + }); +} + +pub(crate) fn record_optimize_recursive_slot(slot_idx: usize, input: &ArrayRef, output: &ArrayRef) { + record(TraceEvent::OptimizeRecursiveSlot { + slot_idx, + input: ArraySummary::new(input), + output: ArraySummary::new(output), + }); +} + +pub(crate) fn record_reduce_applied(array: &ArrayRef, rule: &dyn Debug, output: &ArrayRef) { + record(TraceEvent::ReduceApplied { + array: ArraySummary::new(array), + rule: compact_label(rule), + output: ArraySummary::new(output), + }); +} + +pub(crate) fn record_reduce_declined(array: &ArrayRef, rule: &dyn Debug) { + if !attempts_enabled() { + return; + } + record(TraceEvent::ReduceAttempt { + array: ArraySummary::new(array), + rule: compact_label(rule), + outcome: AttemptOutcome::Declined, + }); +} + +pub(crate) fn record_session_parent_reduce_applied( + parent: &ArrayRef, + child: &ArrayRef, + slot_idx: usize, + plugin_idx: usize, + output: &ArrayRef, +) { + record_parent_reduce_applied( + parent, + child, + slot_idx, + TraceSource::Session(plugin_idx), + "reduce_parent_fn", + output, + ); +} + +pub(crate) fn record_session_parent_reduce_declined( + parent: &ArrayRef, + child: &ArrayRef, + slot_idx: usize, + plugin_idx: usize, +) { + record_parent_reduce_attempt( + parent, + child, + slot_idx, + TraceSource::Session(plugin_idx), + "reduce_parent_fn", + AttemptOutcome::Declined, + ); +} + +pub(crate) fn record_static_parent_reduce_no_match( + parent: &ArrayRef, + child: &ArrayRef, + slot_idx: usize, + rule: &dyn Debug, +) { + record_parent_reduce_attempt( + parent, + child, + slot_idx, + TraceSource::Static, + compact_label(rule), + AttemptOutcome::NoMatch, + ); +} + +pub(crate) fn record_static_parent_reduce_applied( + parent: &ArrayRef, + child: &ArrayRef, + slot_idx: usize, + rule: &dyn Debug, + output: &ArrayRef, +) { + record_parent_reduce_applied( + parent, + child, + slot_idx, + TraceSource::Static, + compact_label(rule), + output, + ); +} + +pub(crate) fn record_static_parent_reduce_declined( + parent: &ArrayRef, + child: &ArrayRef, + slot_idx: usize, + rule: &dyn Debug, +) { + record_parent_reduce_attempt( + parent, + child, + slot_idx, + TraceSource::Static, + compact_label(rule), + AttemptOutcome::Declined, + ); +} + +fn record_parent_reduce_attempt( + parent: &ArrayRef, + child: &ArrayRef, + slot_idx: usize, + source: TraceSource, + rule: impl Into, + outcome: AttemptOutcome, +) { + if !attempts_enabled() { + return; + } + record(TraceEvent::ParentReduceAttempt { + parent: ArraySummary::new(parent), + child: ArraySummary::new(child), + slot_idx, + source, + rule: rule.into(), + outcome, + }); +} + +fn record_parent_reduce_applied( + parent: &ArrayRef, + child: &ArrayRef, + slot_idx: usize, + source: TraceSource, + rule: impl Into, + output: &ArrayRef, +) { + record(TraceEvent::ParentReduceApplied { + parent: ArraySummary::new(parent), + child: ArraySummary::new(child), + slot_idx, + source, + rule: rule.into(), + output: ArraySummary::new(output), + }); +} + +pub(crate) fn record_execute_until_start(root: &ArrayRef) { + record(TraceEvent::ExecuteUntilStart { + target: short_type_name::(), + root: ArraySummary::new(root), + }); +} + +pub(crate) fn record_execute_until_iteration( + iteration: usize, + current: &ArrayRef, + stack_parent: Option<(&ArrayRef, usize)>, + builder_active: bool, +) { + record(TraceEvent::ExecuteUntilIteration { + iteration, + current: ArraySummary::new(current), + stack_parent: stack_parent.map(|(array, slot_idx)| (ArraySummary::new(array), slot_idx)), + builder_active, + }); +} + +pub(crate) fn record_execute_until_done_check(target: bool, canonical: bool) { + if !attempts_enabled() { + return; + } + record(TraceEvent::ExecuteUntilDoneCheck { target, canonical }); +} + +pub(crate) fn record_execute_until_return(output: &ArrayRef) { + record(TraceEvent::ExecuteUntilReturn { + output: ArraySummary::new(output), + }); +} + +pub(crate) fn record_execute_until_pop_frame(slot_idx: usize, output: &ArrayRef) { + record(TraceEvent::ExecuteUntilPopFrame { + slot_idx, + output: ArraySummary::new(output), + }); +} + +pub(crate) fn record_session_execute_parent_applied( + phase: &'static str, + parent: &ArrayRef, + child: &ArrayRef, + slot_idx: usize, + plugin_idx: usize, + output: &ArrayRef, +) { + record_execute_parent_applied( + phase, + parent, + child, + slot_idx, + TraceSource::Session(plugin_idx), + "execute_parent_fn", + output, + ); +} + +pub(crate) fn record_session_execute_parent_declined( + phase: &'static str, + parent: &ArrayRef, + child: &ArrayRef, + slot_idx: usize, + plugin_idx: usize, +) { + record_execute_parent_attempt( + phase, + parent, + child, + slot_idx, + TraceSource::Session(plugin_idx), + "execute_parent_fn", + AttemptOutcome::Declined, + ); +} + +fn record_execute_parent_attempt( + phase: &'static str, + parent: &ArrayRef, + child: &ArrayRef, + slot_idx: usize, + source: TraceSource, + kernel: impl Into, + outcome: AttemptOutcome, +) { + if !attempts_enabled() { + return; + } + record(TraceEvent::ExecuteParentAttempt { + phase, + parent: ArraySummary::new(parent), + child: ArraySummary::new(child), + slot_idx, + source, + kernel: kernel.into(), + outcome, + }); +} + +fn record_execute_parent_applied( + phase: &'static str, + parent: &ArrayRef, + child: &ArrayRef, + slot_idx: usize, + source: TraceSource, + kernel: impl Into, + output: &ArrayRef, +) { + record(TraceEvent::ExecuteParentApplied { + phase, + parent: ArraySummary::new(parent), + child: ArraySummary::new(child), + slot_idx, + source, + kernel: kernel.into(), + output: ArraySummary::new(output), + }); +} + +pub(crate) fn record_execute_parent_none(phase: &'static str, current: &ArrayRef) { + if !attempts_enabled() { + return; + } + record(TraceEvent::PhaseNone { + indent: 2, + phase, + subject: "current", + array: ArraySummary::new(current), + }); +} + +pub(crate) fn record_execute_optimized(input: &ArrayRef, output: &ArrayRef) { + let changed = !ArrayRef::ptr_eq(input, output); + if !changed && !attempts_enabled() { + return; + } + record(TraceEvent::ExecuteOptimized { + input: ArraySummary::new(input), + output: ArraySummary::new(output), + changed, + }); +} + +pub(crate) fn record_execute_encoding(array: &ArrayRef) { + if !attempts_enabled() { + return; + } + record(TraceEvent::ExecuteEncoding { + array: ArraySummary::new(array), + }); +} + +pub(crate) fn record_execute_slot(slot_idx: usize, parent: &ArrayRef, child: &ArrayRef) { + record(TraceEvent::SlotTransition { + step: "ExecuteSlot", + slot_idx, + parent: ArraySummary::new(parent), + child: ArraySummary::new(child), + }); +} + +pub(crate) fn record_builder_start(array: &ArrayRef) { + record(TraceEvent::BuilderEvent { + action: "start", + subject: "array", + array: ArraySummary::new(array), + }); +} + +pub(crate) fn record_append_child(slot_idx: usize, parent: &ArrayRef, child: &ArrayRef) { + record(TraceEvent::SlotTransition { + step: "AppendChild", + slot_idx, + parent: ArraySummary::new(parent), + child: ArraySummary::new(child), + }); +} + +pub(crate) fn record_builder_append(child: &ArrayRef) { + record(TraceEvent::BuilderEvent { + action: "append", + subject: "child", + array: ArraySummary::new(child), + }); +} + +pub(crate) fn record_execute_done(array: &ArrayRef) { + record(TraceEvent::ExecuteDone { + array: ArraySummary::new(array), + }); +} + +pub(crate) fn record_builder_finish(output: &ArrayRef) { + record(TraceEvent::BuilderEvent { + action: "finish", + subject: "output", + array: ArraySummary::new(output), + }); +} + +pub(crate) fn record_single_step_start(array: &ArrayRef) { + record(TraceEvent::SingleStepStart { + array: ArraySummary::new(array), + }); +} + +pub(crate) fn record_single_step_phase_none(phase: &'static str, array: &ArrayRef) { + if !attempts_enabled() { + return; + } + record(TraceEvent::PhaseNone { + indent: 1, + phase, + subject: "array", + array: ArraySummary::new(array), + }); +} + +pub(crate) fn record_single_step_applied(phase: &'static str, input: &ArrayRef, output: &ArrayRef) { + record(TraceEvent::SingleStepApplied { + phase, + input: ArraySummary::new(input), + output: ArraySummary::new(output), + }); +} + +fn record(event: TraceEvent) { + ACTIVE_TRACE.with(|active| { + if let Some(recorder) = active.borrow_mut().as_mut() { + recorder.events.push(event); + } + }); +} + +fn compact_label(value: &dyn Debug) -> String { + let label = format!("{value:?}"); + if let Some(label) = adapter_field(&label, "rule") { + return label.to_string(); + } + if let Some(label) = adapter_field(&label, "kernel") { + return label.to_string(); + } + label +} + +fn adapter_field<'a>(label: &'a str, field: &str) -> Option<&'a str> { + let marker = format!("{field}: "); + let start = label.find(&marker)? + marker.len(); + let rest = &label[start..]; + let end = rest.rfind(" }")?; + Some(&rest[..end]) +} + +fn short_type_name() -> String { + std::any::type_name::() + .rsplit("::") + .next() + .vortex_expect("type names are never empty") + .to_string() +} + +thread_local! { + static TRACE_INTEREST: Cell = const { Cell::new(TraceInterest::Off) }; + static ACTIVE_TRACE: RefCell> = const { RefCell::new(None) }; +} + +static ACTIVE_TRACE_COUNT: AtomicUsize = AtomicUsize::new(0); +static ATTEMPTS_TRACE_COUNT: AtomicUsize = AtomicUsize::new(0); + +struct ActiveTraceGuard { + interest: TraceInterest, +} + +impl Drop for ActiveTraceGuard { + fn drop(&mut self) { + if self.interest == TraceInterest::Attempts { + ATTEMPTS_TRACE_COUNT.fetch_sub(1, Ordering::Relaxed); + } + ACTIVE_TRACE_COUNT.fetch_sub(1, Ordering::Relaxed); + TRACE_INTEREST.with(|interest| interest.set(TraceInterest::Off)); + ACTIVE_TRACE.with(|active| { + active.borrow_mut().take(); + }); + } +} + +#[derive(Debug)] +struct TraceRecorder { + options: TraceOptions, + events: Vec, +} + +impl TraceRecorder { + fn new(options: TraceOptions) -> Self { + Self { + options, + events: Vec::new(), + } + } + + fn finish(self) -> TraceDisplay { + TraceDisplay { + options: self.options, + events: self.events, + } + } +} + +#[derive(Clone, Debug)] +enum TraceEvent { + OptimizeStart { + root: ArraySummary, + session: bool, + }, + OptimizeLoopStart { + array: ArraySummary, + }, + OptimizeLoopEnd, + OptimizeDone { + output: ArraySummary, + changed: bool, + }, + OptimizeRecursiveStart { + root: ArraySummary, + }, + OptimizeRecursiveSlot { + slot_idx: usize, + input: ArraySummary, + output: ArraySummary, + }, + ReduceAttempt { + array: ArraySummary, + rule: String, + outcome: AttemptOutcome, + }, + ReduceApplied { + array: ArraySummary, + rule: String, + output: ArraySummary, + }, + ParentReduceAttempt { + parent: ArraySummary, + child: ArraySummary, + slot_idx: usize, + source: TraceSource, + rule: String, + outcome: AttemptOutcome, + }, + ParentReduceApplied { + parent: ArraySummary, + child: ArraySummary, + slot_idx: usize, + source: TraceSource, + rule: String, + output: ArraySummary, + }, + ExecuteUntilStart { + target: String, + root: ArraySummary, + }, + ExecuteUntilIteration { + iteration: usize, + current: ArraySummary, + stack_parent: Option<(ArraySummary, usize)>, + builder_active: bool, + }, + ExecuteUntilDoneCheck { + target: bool, + canonical: bool, + }, + ExecuteUntilReturn { + output: ArraySummary, + }, + ExecuteUntilPopFrame { + slot_idx: usize, + output: ArraySummary, + }, + ExecuteParentAttempt { + phase: &'static str, + parent: ArraySummary, + child: ArraySummary, + slot_idx: usize, + source: TraceSource, + kernel: String, + outcome: AttemptOutcome, + }, + ExecuteParentApplied { + phase: &'static str, + parent: ArraySummary, + child: ArraySummary, + slot_idx: usize, + source: TraceSource, + kernel: String, + output: ArraySummary, + }, + PhaseNone { + indent: usize, + phase: &'static str, + subject: &'static str, + array: ArraySummary, + }, + ExecuteOptimized { + input: ArraySummary, + output: ArraySummary, + changed: bool, + }, + ExecuteEncoding { + array: ArraySummary, + }, + SlotTransition { + step: &'static str, + slot_idx: usize, + parent: ArraySummary, + child: ArraySummary, + }, + BuilderEvent { + action: &'static str, + subject: &'static str, + array: ArraySummary, + }, + ExecuteDone { + array: ArraySummary, + }, + SingleStepStart { + array: ArraySummary, + }, + SingleStepApplied { + phase: &'static str, + input: ArraySummary, + output: ArraySummary, + }, +} + +impl TraceEvent { + fn is_hidden(&self, resolution: TraceResolution) -> bool { + match resolution { + TraceResolution::Attempts => matches!(self, TraceEvent::OptimizeLoopEnd), + TraceResolution::ExecutedOnly => matches!( + self, + TraceEvent::OptimizeLoopStart { .. } + | TraceEvent::OptimizeLoopEnd + | TraceEvent::PhaseNone { .. } + | TraceEvent::ExecuteUntilDoneCheck { .. } + | TraceEvent::ExecuteEncoding { .. } + | TraceEvent::ExecuteOptimized { changed: false, .. } + | TraceEvent::ExecuteParentAttempt { .. } + | TraceEvent::ReduceAttempt { .. } + | TraceEvent::ParentReduceAttempt { .. } + ), + } + } + + fn opens_after(&self, resolution: TraceResolution) -> bool { + match resolution { + TraceResolution::Attempts => matches!( + self, + TraceEvent::OptimizeStart { .. } | TraceEvent::OptimizeLoopStart { .. } + ), + TraceResolution::ExecutedOnly => matches!(self, TraceEvent::OptimizeStart { .. }), + } + } + + fn closes_before(&self, resolution: TraceResolution) -> bool { + match resolution { + TraceResolution::Attempts => matches!(self, TraceEvent::OptimizeLoopEnd), + TraceResolution::ExecutedOnly => false, + } + } + + fn closes_after(&self, _resolution: TraceResolution) -> bool { + matches!(self, TraceEvent::OptimizeDone { .. }) + } + + fn relative_indent(&self, _resolution: TraceResolution, in_optimize_scope: bool) -> usize { + match self { + TraceEvent::OptimizeStart { .. } + | TraceEvent::OptimizeLoopStart { .. } + | TraceEvent::OptimizeDone { .. } => 0, + TraceEvent::ReduceAttempt { .. } + | TraceEvent::ReduceApplied { .. } + | TraceEvent::ParentReduceAttempt { .. } + | TraceEvent::ParentReduceApplied { .. } + if in_optimize_scope => + { + 0 + } + TraceEvent::PhaseNone { indent, .. } => *indent, + TraceEvent::ReduceAttempt { .. } + | TraceEvent::ReduceApplied { .. } + | TraceEvent::ParentReduceAttempt { .. } + | TraceEvent::ParentReduceApplied { .. } + | TraceEvent::ExecuteUntilDoneCheck { .. } + | TraceEvent::ExecuteUntilPopFrame { .. } + | TraceEvent::ExecuteParentAttempt { .. } + | TraceEvent::ExecuteParentApplied { .. } + | TraceEvent::ExecuteOptimized { .. } + | TraceEvent::ExecuteEncoding { .. } + | TraceEvent::SlotTransition { .. } + | TraceEvent::BuilderEvent { .. } + | TraceEvent::ExecuteDone { .. } => 2, + TraceEvent::OptimizeRecursiveSlot { .. } + | TraceEvent::ExecuteUntilIteration { .. } + | TraceEvent::ExecuteUntilReturn { .. } + | TraceEvent::SingleStepApplied { .. } => 1, + TraceEvent::OptimizeLoopEnd + | TraceEvent::OptimizeRecursiveStart { .. } + | TraceEvent::ExecuteUntilStart { .. } + | TraceEvent::SingleStepStart { .. } => 0, + } + } + + fn fmt_line(&self, f: &mut fmt::Formatter<'_>, resolution: TraceResolution) -> fmt::Result { + match self { + TraceEvent::OptimizeStart { root, session } => { + write!(f, "optimize root={root} session={session}") + } + TraceEvent::OptimizeLoopStart { array } => { + write!(f, "loop input={array}") + } + TraceEvent::OptimizeLoopEnd => Ok(()), + TraceEvent::OptimizeDone { output, changed } => match resolution { + TraceResolution::Attempts => write!(f, "done output={output} changed={changed}"), + TraceResolution::ExecutedOnly => write!(f, "done output={output}"), + }, + TraceEvent::OptimizeRecursiveStart { root } => { + write!(f, "optimize_recursive root={root}") + } + TraceEvent::OptimizeRecursiveSlot { + slot_idx, + input, + output, + } => write!(f, "recursive slot={slot_idx} input={input} output={output}"), + TraceEvent::ReduceAttempt { + array, + rule, + outcome, + } => write!( + f, + "reduce attempt array={array} source=static rule={rule} outcome={outcome}" + ), + TraceEvent::ReduceApplied { + array, + rule, + output, + } => match resolution { + TraceResolution::Attempts => write!( + f, + "reduce applied array={array} source=static rule={rule} output={output}" + ), + TraceResolution::ExecutedOnly => { + write!(f, "reduce {rule}: {array} -> {output}") + } + }, + TraceEvent::ParentReduceAttempt { + parent, + child, + slot_idx, + source, + rule, + outcome, + } => write!( + f, + "reduce_parent attempt slot={slot_idx} parent={parent} child={child} source={source} rule={rule} outcome={outcome}" + ), + TraceEvent::ParentReduceApplied { + parent, + child, + slot_idx, + source, + rule, + output, + } => match resolution { + TraceResolution::Attempts => write!( + f, + "reduce_parent applied slot={slot_idx} parent={parent} child={child} source={source} rule={rule} output={output}" + ), + TraceResolution::ExecutedOnly => write!( + f, + "reduce_parent {source}:{rule} slot={slot_idx} parent={parent} child={child} -> {output}" + ), + }, + TraceEvent::ExecuteUntilStart { target, root } => { + write!(f, "execute_until target={target} root={root}") + } + TraceEvent::ExecuteUntilIteration { + iteration, + current, + stack_parent, + builder_active, + } => { + write!(f, "iter {iteration} current={current}")?; + if let Some((parent, slot_idx)) = stack_parent { + write!(f, " stack_parent={parent} slot={slot_idx}")?; + } + write!(f, " builder_active={builder_active}") + } + TraceEvent::ExecuteUntilDoneCheck { target, canonical } => { + write!(f, "done_check target={target} canonical={canonical}") + } + TraceEvent::ExecuteUntilReturn { output } => { + write!(f, "return output={output}") + } + TraceEvent::ExecuteUntilPopFrame { slot_idx, output } => { + write!(f, "pop_frame slot={slot_idx} output={output}") + } + TraceEvent::ExecuteParentAttempt { + phase, + parent, + child, + slot_idx, + source, + kernel, + outcome, + } => write!( + f, + "{phase} attempt slot={slot_idx} parent={parent} child={child} source={source} kernel={kernel} outcome={outcome}" + ), + TraceEvent::ExecuteParentApplied { + phase, + parent, + child, + slot_idx, + source, + kernel, + output, + } => match resolution { + TraceResolution::Attempts => write!( + f, + "{phase} applied slot={slot_idx} parent={parent} child={child} source={source} kernel={kernel} output={output}" + ), + TraceResolution::ExecutedOnly => write!( + f, + "{phase} {source}:{kernel} slot={slot_idx} parent={parent} child={child} -> {output}" + ), + }, + TraceEvent::PhaseNone { + phase, + subject, + array, + .. + } => { + write!(f, "{phase} none {subject}={array}") + } + TraceEvent::ExecuteOptimized { + input, + output, + changed, + } => match resolution { + TraceResolution::Attempts => write!( + f, + "optimize_ctx input={input} output={output} changed={changed}" + ), + TraceResolution::ExecutedOnly => write!(f, "optimize_ctx {input} -> {output}"), + }, + TraceEvent::ExecuteEncoding { array } => { + write!(f, "execute encoding={array}") + } + TraceEvent::SlotTransition { + step, + slot_idx, + parent, + child, + } => write!(f, "{step} slot={slot_idx} parent={parent} child={child}"), + TraceEvent::ExecuteDone { array } => { + write!(f, "Done array={array}") + } + TraceEvent::BuilderEvent { + action, + subject, + array, + } => { + write!(f, "builder {action} {subject}={array}") + } + TraceEvent::SingleStepStart { array } => { + write!(f, "execute_step input={array}") + } + TraceEvent::SingleStepApplied { + phase, + input, + output, + } => match resolution { + TraceResolution::Attempts => { + write!(f, "{phase} applied input={input} output={output}") + } + TraceResolution::ExecutedOnly => write!(f, "{phase} {input} -> {output}"), + }, + } + } +} + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/test_harness/trace/tests.rs b/vortex-array/src/test_harness/trace/tests.rs new file mode 100644 index 00000000000..0cda75b7640 --- /dev/null +++ b/vortex-array/src/test_harness/trace/tests.rs @@ -0,0 +1,771 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt::Display; +use std::fmt::Formatter; +use std::hash::Hasher; + +use rstest::fixture; +use rstest::rstest; +use smallvec::smallvec; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_panic; +use vortex_mask::Mask; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::ArrayEq; +use crate::ArrayHash; +use crate::ArrayRef; +use crate::Canonical; +use crate::EqMode; +use crate::ExecutionCtx; +use crate::ExecutionResult; +use crate::IntoArray; +use crate::VTable; +use crate::array::Array; +use crate::array::ArrayId; +use crate::array::ArrayParts; +use crate::array::ArrayView; +use crate::array::vtable::NotSupported; +use crate::array::vtable::ValidityVTable; +use crate::array::vtable::with_empty_buffers; +use crate::arrays::BoolArray; +use crate::arrays::ChunkedArray; +use crate::arrays::ConstantArray; +use crate::arrays::DictArray; +use crate::arrays::Filter; +use crate::arrays::FilterArray; +use crate::arrays::Primitive; +use crate::arrays::PrimitiveArray; +use crate::arrays::StructArray; +use crate::arrays::VarBinViewArray; +use crate::arrays::filter::FilterArraySlotsExt; +use crate::arrays::scalar_fn::ScalarFnFactoryExt; +use crate::assert_arrays_eq; +use crate::buffer::BufferHandle; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::kernel::ExecuteParentKernel; +use crate::matcher::Matcher; +use crate::optimizer::ArrayOptimizer; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::scalar::Scalar; +use crate::scalar_fn::fns::binary::Binary; +use crate::scalar_fn::fns::like::Like; +use crate::scalar_fn::fns::like::LikeOptions; +use crate::scalar_fn::fns::operators::Operator; +use crate::serde::ArrayChildren; +use crate::session::ArraySession; +use crate::test_harness::trace::TraceOptions; +use crate::test_harness::trace::TraceResolution; +use crate::test_harness::trace::trace_op; +use crate::test_harness::trace::trace_op_with; +use crate::validity::Validity; + +#[fixture] +fn stack_parent_fixture() -> VortexResult { + stack_parent(stack_child()?) +} + +/// Build a session with the `StackChild` parent kernels registered. +/// +/// The declining kernel is registered first so strict trace snapshots can assert that both +/// session kernels are attempted in registration order. +#[fixture] +fn stack_parent_session() -> VortexSession { + let session = VortexSession::empty().with::(); + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(StackParent.id(), StackChild, StackDeclineKernel); + kernels.register_execute_parent_kernel(StackParent.id(), StackChild, StackParentKernel); + drop(kernels); + session +} + +fn stack_child() -> VortexResult { + Ok( + Array::try_from_parts(ArrayParts::new(StackChild, test_dtype(), 3, StackChildData))? + .into_array(), + ) +} + +fn stack_parent(child: ArrayRef) -> VortexResult { + Ok(Array::try_from_parts( + ArrayParts::new( + StackParent, + child.dtype().clone(), + child.len(), + StackParentData, + ) + .with_slots(smallvec![Some(child)]), + )? + .into_array()) +} + +fn test_dtype() -> DType { + DType::Primitive(PType::I32, Nullability::NonNullable) +} + +#[derive(Clone, Debug)] +struct StackParent; + +#[derive(Clone, Debug)] +struct StackParentData; + +impl Display for StackParentData { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.write_str("stack-parent") + } +} + +impl ArrayHash for StackParentData { + fn array_hash(&self, _state: &mut H, _eq_mode: EqMode) {} +} + +impl ArrayEq for StackParentData { + fn array_eq(&self, _other: &Self, _eq_mode: EqMode) -> bool { + true + } +} + +impl ValidityVTable for StackParent { + fn validity(_array: ArrayView<'_, StackParent>) -> VortexResult { + Ok(Validity::NonNullable) + } +} + +impl VTable for StackParent { + type TypedArrayData = StackParentData; + type OperationsVTable = NotSupported; + type ValidityVTable = Self; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.test.stack-parent"); + *ID + } + + fn validate( + &self, + _data: &Self::TypedArrayData, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + vortex_ensure!(dtype == &test_dtype(), "unexpected stack parent dtype"); + vortex_ensure!(len == 3, "unexpected stack parent length"); + vortex_ensure!(slots.len() == 1, "stack parent must have one child slot"); + let Some(child) = &slots[0] else { + vortex_bail!("stack parent child slot is missing"); + }; + vortex_ensure!(child.dtype() == dtype, "stack parent child dtype mismatch"); + vortex_ensure!(child.len() == len, "stack parent child length mismatch"); + Ok(()) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 0 + } + + fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { + vortex_panic!("StackParent buffer index {idx} out of bounds") + } + + fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option { + None + } + + fn serialize( + _array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + Ok(None) + } + + fn deserialize( + &self, + _dtype: &DType, + _len: usize, + _metadata: &[u8], + _buffers: &[BufferHandle], + _children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + vortex_bail!("StackParent cannot be deserialized") + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + with_empty_buffers(self, array, buffers) + } + + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + match idx { + 0 => "child".to_string(), + _ => vortex_panic!("StackParent slot index {idx} out of bounds"), + } + } + + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + let Some(child) = array.slots()[0].as_ref() else { + vortex_bail!("stack parent child slot is missing"); + }; + if !child.is::() { + return Ok(ExecutionResult::execute_slot::(array, 0)); + } + + Ok(ExecutionResult::done(child.clone())) + } +} + +#[derive(Clone, Debug)] +struct StackChild; + +#[derive(Clone, Debug)] +struct StackChildData; + +impl Display for StackChildData { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.write_str("stack-child") + } +} + +impl ArrayHash for StackChildData { + fn array_hash(&self, _state: &mut H, _eq_mode: EqMode) {} +} + +impl ArrayEq for StackChildData { + fn array_eq(&self, _other: &Self, _eq_mode: EqMode) -> bool { + true + } +} + +impl ValidityVTable for StackChild { + fn validity(_array: ArrayView<'_, StackChild>) -> VortexResult { + Ok(Validity::NonNullable) + } +} + +impl VTable for StackChild { + type TypedArrayData = StackChildData; + type OperationsVTable = NotSupported; + type ValidityVTable = Self; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.test.stack-child"); + *ID + } + + fn validate( + &self, + _data: &Self::TypedArrayData, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + vortex_ensure!(dtype == &test_dtype(), "unexpected stack child dtype"); + vortex_ensure!(len == 3, "unexpected stack child length"); + vortex_ensure!(slots.is_empty(), "stack child must not have slots"); + Ok(()) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 0 + } + + fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { + vortex_panic!("StackChild buffer index {idx} out of bounds") + } + + fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option { + None + } + + fn serialize( + _array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + Ok(None) + } + + fn deserialize( + &self, + _dtype: &DType, + _len: usize, + _metadata: &[u8], + _buffers: &[BufferHandle], + _children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + vortex_bail!("StackChild cannot be deserialized") + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + with_empty_buffers(self, array, buffers) + } + + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + vortex_panic!("StackChild slot index {idx} out of bounds") + } + + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + debug_assert!(array.slots().is_empty()); + Ok(ExecutionResult::done(PrimitiveArray::from_iter([ + 99i32, 99, 99, + ]))) + } +} + +#[derive(Debug)] +struct StackDeclineKernel; + +impl ExecuteParentKernel for StackDeclineKernel { + type Parent = StackParent; + + fn execute_parent( + &self, + _array: ArrayView<'_, StackChild>, + _parent: ::Match<'_>, + _child_idx: usize, + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(None) + } +} + +#[derive(Debug)] +struct StackParentKernel; + +impl ExecuteParentKernel for StackParentKernel { + type Parent = StackParent; + + fn execute_parent( + &self, + _array: ArrayView<'_, StackChild>, + parent: ::Match<'_>, + child_idx: usize, + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if parent + .slots() + .get(child_idx) + .is_some_and(|slot| slot.is_none()) + { + return Ok(Some(PrimitiveArray::from_iter([1i32, 2, 3]).into_array())); + } + + Ok(None) + } +} + +#[test] +fn trace_optimize_reduce_fixpoint() -> VortexResult<()> { + let values = PrimitiveArray::from_iter([0i32, 1, 2, 3]).into_array(); + let filter = FilterArray::try_new(values.clone(), Mask::new_true(values.len()))?.into_array(); + + let traced = trace_op(|| filter.optimize())?; + + assert!(traced.output.is::()); + assert_arrays_eq!(traced.output, values, &mut execution_ctx()); + insta::assert_snapshot!(traced.trace.to_string(), @r" +optimize root=vortex.filter(i32, len=4) session=false + reduce TrivialFilterRule: vortex.filter(i32, len=4) -> vortex.primitive(i32, len=4) + done output=vortex.primitive(i32, len=4) +"); + + Ok(()) +} + +#[test] +fn trace_optimize_parent_reduce_fixpoint_attempts() -> VortexResult<()> { + let values = PrimitiveArray::from_iter([0i32, 1, 2, 3, 4, 5]).into_array(); + let inner = FilterArray::try_new( + values, + Mask::from_iter([true, false, true, true, false, true]), + )? + .into_array(); + let outer = + FilterArray::try_new(inner, Mask::from_iter([false, true, true, false]))?.into_array(); + + let traced = trace_op_with( + TraceOptions { + resolution: TraceResolution::ExecutedOnly, + }, + || outer.optimize(), + )?; + + let optimized_filter = traced.output.as_::(); + assert!(optimized_filter.child().is::()); + assert_arrays_eq!( + traced.output, + PrimitiveArray::from_iter([2i32, 3]), + &mut execution_ctx() + ); + insta::assert_snapshot!(traced.trace.to_string(), @r" + optimize root=vortex.filter(i32, len=2) session=false + reduce_parent static:FilterReduceAdaptor(Filter) slot=0 parent=vortex.filter(i32, len=2) child=vortex.filter(i32, len=4) -> vortex.filter(i32, len=2) + done output=vortex.filter(i32, len=2) + "); + + let mut ctx = ExecutionCtx::new(VortexSession::empty().with::()); + let traced = trace_op_with( + TraceOptions { + resolution: TraceResolution::ExecutedOnly, + }, + || outer.execute::(&mut ctx), + )?; + + insta::assert_snapshot!(traced.trace.to_string(), @r" + execute_until target=AnyCanonical root=vortex.filter(i32, len=2) + iter 0 current=vortex.filter(i32, len=2) builder_active=false + ExecuteSlot slot=0 parent=vortex.filter(i32, len=2) child=vortex.filter(i32, len=4) + iter 1 current=vortex.filter(i32, len=4) stack_parent=vortex.filter(i32, len=2) slot=0 builder_active=false + Done array=vortex.primitive(i32, len=4) + iter 2 current=vortex.primitive(i32, len=4) stack_parent=vortex.filter(i32, len=2) slot=0 builder_active=false + pop_frame slot=0 output=vortex.filter(i32, len=2) + iter 3 current=vortex.filter(i32, len=2) builder_active=false + Done array=vortex.primitive(i32, len=2) + iter 4 current=vortex.primitive(i32, len=2) builder_active=false + return output=vortex.primitive(i32, len=2) + "); + + Ok(()) +} + +#[rstest] +fn trace_execution_stack_parent_kernel_attempts( + stack_parent_fixture: VortexResult, + stack_parent_session: VortexSession, +) -> VortexResult<()> { + let mut ctx = ExecutionCtx::new(stack_parent_session); + let parent = stack_parent_fixture?; + + let traced = trace_op_with( + TraceOptions { + resolution: TraceResolution::Attempts, + }, + || parent.execute::(&mut ctx), + )?; + + assert_arrays_eq!( + traced.output, + PrimitiveArray::from_iter([1i32, 2, 3]), + &mut ctx + ); + insta::assert_snapshot!(traced.trace.to_string(), @" + execute_until target=AnyCanonical root=vortex.test.stack-parent(i32, len=3) + iter 0 current=vortex.test.stack-parent(i32, len=3) builder_active=false + done_check target=false canonical=false + child_execute_parent attempt slot=0 parent=vortex.test.stack-parent(i32, len=3) child=vortex.test.stack-child(i32, len=3) source=session[0] kernel=execute_parent_fn outcome=declined + child_execute_parent attempt slot=0 parent=vortex.test.stack-parent(i32, len=3) child=vortex.test.stack-child(i32, len=3) source=session[1] kernel=execute_parent_fn outcome=declined + child_execute_parent none current=vortex.test.stack-parent(i32, len=3) + execute encoding=vortex.test.stack-parent(i32, len=3) + ExecuteSlot slot=0 parent=vortex.test.stack-parent(i32, len=3) child=vortex.test.stack-child(i32, len=3) + iter 1 current=vortex.test.stack-child(i32, len=3) stack_parent=vortex.test.stack-parent(i32, len=3) slot=0 builder_active=false + done_check target=false canonical=false + stack_execute_parent attempt slot=0 parent=vortex.test.stack-parent(i32, len=3) child=vortex.test.stack-child(i32, len=3) source=session[0] kernel=execute_parent_fn outcome=declined + stack_execute_parent applied slot=0 parent=vortex.test.stack-parent(i32, len=3) child=vortex.test.stack-child(i32, len=3) source=session[1] kernel=execute_parent_fn output=vortex.primitive(i32, len=3) + optimize root=vortex.primitive(i32, len=3) session=true + loop input=vortex.primitive(i32, len=3) + reduce none array=vortex.primitive(i32, len=3) + reduce_parent none array=vortex.primitive(i32, len=3) + done output=vortex.primitive(i32, len=3) changed=false + optimize_ctx input=vortex.primitive(i32, len=3) output=vortex.primitive(i32, len=3) changed=false + iter 2 current=vortex.primitive(i32, len=3) builder_active=false + done_check target=true canonical=true + return output=vortex.primitive(i32, len=3) + "); + + Ok(()) +} + +#[test] +fn trace_execution_chunked_append_child_flow() -> VortexResult<()> { + let chunks = vec![ + PrimitiveArray::from_iter([1i32, 2]).into_array(), + PrimitiveArray::from_iter([3i32]).into_array(), + PrimitiveArray::from_iter([4i32, 5]).into_array(), + ]; + let dtype = chunks[0].dtype().clone(); + let chunked = ChunkedArray::try_new(chunks, dtype)?.into_array(); + let mut ctx = ExecutionCtx::new(VortexSession::empty().with::()); + + let traced = trace_op(|| { + chunked + .execute::(&mut ctx) + .map(IntoArray::into_array) + })?; + + assert_arrays_eq!( + traced.output, + PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]), + &mut ctx + ); + insta::assert_snapshot!(traced.trace.to_string(), @" + execute_until target=AnyCanonical root=vortex.chunked(i32, len=5) + iter 0 current=vortex.chunked(i32, len=5) builder_active=false + builder start array=vortex.chunked(i32, len=5) + AppendChild slot=1 parent=vortex.chunked(i32, len=5) child=vortex.primitive(i32, len=2) + builder append child=vortex.primitive(i32, len=2) + iter 1 current=vortex.chunked(i32, len=5) builder_active=true + AppendChild slot=2 parent=vortex.chunked(i32, len=5) child=vortex.primitive(i32, len=1) + builder append child=vortex.primitive(i32, len=1) + iter 2 current=vortex.chunked(i32, len=5) builder_active=true + AppendChild slot=3 parent=vortex.chunked(i32, len=5) child=vortex.primitive(i32, len=2) + builder append child=vortex.primitive(i32, len=2) + iter 3 current=vortex.chunked(i32, len=5) builder_active=true + Done array=vortex.primitive(i32, len=0) + builder finish output=vortex.primitive(i32, len=5) + iter 4 current=vortex.primitive(i32, len=5) builder_active=false + return output=vortex.primitive(i32, len=5) + "); + + Ok(()) +} + +/// A dictionary of strings: codes `[0, 1, 2, 1, 0, 2]` over values +/// `["alpha", "beta", "gamma"]`. +fn dict_of_strings() -> VortexResult { + Ok(DictArray::try_new( + PrimitiveArray::from_iter([0u32, 1, 2, 1, 0, 2]).into_array(), + VarBinViewArray::from_iter_str(["alpha", "beta", "gamma"]).into_array(), + )? + .into_array()) +} + +fn execution_ctx() -> ExecutionCtx { + ExecutionCtx::new(VortexSession::empty().with::()) +} + +#[test] +fn trace_take_on_chunked() -> VortexResult<()> { + let chunked = ChunkedArray::try_new( + vec![ + PrimitiveArray::from_iter([10i32, 11]).into_array(), + PrimitiveArray::from_iter([12i32, 13, 14]).into_array(), + ], + DType::Primitive(PType::I32, Nullability::NonNullable), + )? + .into_array(); + + // A take is expressed as a `DictArray` whose codes are the take indices. + let indices = PrimitiveArray::from_iter([4u64, 0, 2, 2]).into_array(); + let take = DictArray::try_new(indices, chunked)?.into_array(); + + // No reduce rule rewrites a take over a chunked array: the work all happens at execution + // time, so the optimizer trace is empty. + let traced = trace_op(|| take.optimize())?; + insta::assert_snapshot!(traced.trace.to_string(), @""); + + let optimized = traced.output; + let traced = trace_op(|| { + optimized + .execute::(&mut execution_ctx()) + .map(IntoArray::into_array) + })?; + + assert_arrays_eq!( + traced.output, + PrimitiveArray::from_iter([14i32, 10, 12, 12]), + &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 + execute_until target=AnyCanonical root=vortex.chunked(i32, len=3) + iter 0 current=vortex.chunked(i32, len=3) builder_active=false + builder start array=vortex.chunked(i32, len=3) + AppendChild slot=1 parent=vortex.chunked(i32, len=3) child=vortex.filter(i32, len=1) + builder append child=vortex.filter(i32, len=1) + execute_until target=AnyCanonical root=vortex.filter(i32, len=1) + iter 0 current=vortex.filter(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) + iter 1 current=vortex.chunked(i32, len=3) builder_active=true + AppendChild slot=2 parent=vortex.chunked(i32, len=3) child=vortex.filter(i32, len=2) + builder append child=vortex.filter(i32, len=2) + execute_until target=AnyCanonical root=vortex.filter(i32, len=2) + iter 0 current=vortex.filter(i32, len=2) builder_active=false + Done array=vortex.primitive(i32, len=2) + iter 1 current=vortex.primitive(i32, len=2) builder_active=false + return output=vortex.primitive(i32, len=2) + iter 2 current=vortex.chunked(i32, len=3) builder_active=true + Done array=vortex.primitive(i32, len=0) + builder finish output=vortex.primitive(i32, len=3) + iter 3 current=vortex.primitive(i32, len=3) builder_active=false + return output=vortex.primitive(i32, len=3) + child_execute_parent session[0]:execute_parent_fn slot=1 parent=vortex.dict(i32, len=4) child=vortex.chunked(i32, len=5) -> 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_filter_on_struct_with_complex_children() -> VortexResult<()> { + let names = DictArray::try_new( + PrimitiveArray::from_iter([0u32, 1, 2, 1, 0]).into_array(), + VarBinViewArray::from_iter_str(["alpha", "beta", "gamma"]).into_array(), + )? + .into_array(); + let scores = ChunkedArray::try_new( + vec![ + PrimitiveArray::from_iter([1i64, 2]).into_array(), + PrimitiveArray::from_iter([3i64, 4, 5]).into_array(), + ], + DType::Primitive(PType::I64, Nullability::NonNullable), + )? + .into_array(); + let struct_ = StructArray::from_fields(&[("name", names), ("score", scores)])?.into_array(); + + let filtered = + FilterArray::try_new(struct_, Mask::from_iter([true, false, true, true, false]))? + .into_array(); + + let traced = trace_op(|| filtered.optimize())?; + insta::assert_snapshot!(traced.trace.to_string(), @" + optimize root=vortex.filter({name=utf8, score=i64}, len=3) session=false + optimize root=vortex.filter(utf8, len=3) session=false + reduce_parent static:FilterReduceAdaptor(Dict) slot=0 parent=vortex.filter(utf8, len=3) child=vortex.dict(utf8, len=5) -> vortex.dict(utf8, len=3) + done output=vortex.dict(utf8, len=3) + reduce FilterStructRule: vortex.filter({name=utf8, score=i64}, len=3) -> vortex.struct({name=utf8, score=i64}, len=3) + done output=vortex.struct({name=utf8, score=i64}, len=3) + "); + + let optimized = traced.output; + let traced = trace_op(|| { + optimized + .execute::(&mut execution_ctx()) + .map(IntoArray::into_array) + })?; + + let expected = StructArray::from_fields(&[ + ( + "name", + VarBinViewArray::from_iter_str(["alpha", "gamma", "beta"]).into_array(), + ), + ( + "score", + PrimitiveArray::from_iter([1i64, 3, 4]).into_array(), + ), + ])? + .into_array(); + assert_arrays_eq!(traced.output, expected, &mut execution_ctx()); + insta::assert_snapshot!(traced.trace.to_string(), @" + execute_until target=AnyCanonical root=vortex.struct({name=utf8, score=i64}, len=3) + iter 0 current=vortex.struct({name=utf8, score=i64}, len=3) builder_active=false + return output=vortex.struct({name=utf8, score=i64}, len=3) + "); + + Ok(()) +} + +#[test] +fn trace_compare_on_dict() -> VortexResult<()> { + let dict = DictArray::try_new( + PrimitiveArray::from_iter([0u32, 1, 2, 1, 0]).into_array(), + PrimitiveArray::from_iter([10i32, 20, 30]).into_array(), + )? + .into_array(); + let rhs = ConstantArray::new(Scalar::from(20i32), dict.len()).into_array(); + + let compared = Binary.try_new_array(dict.len(), Operator::Eq, [dict, rhs])?; + + let traced = trace_op(|| compared.optimize())?; + insta::assert_snapshot!(traced.trace.to_string(), @" + optimize root=vortex.binary(bool, len=5) session=false + reduce_parent static:DictionaryScalarFnValuesPushDownRule slot=0 parent=vortex.binary(bool, len=5) child=vortex.dict(i32, len=5) -> vortex.dict(bool, len=5) + done output=vortex.dict(bool, len=5) + "); + + let optimized = traced.output; + let traced = trace_op(|| { + optimized + .execute::(&mut execution_ctx()) + .map(IntoArray::into_array) + })?; + + assert_arrays_eq!( + traced.output, + BoolArray::from_iter([false, true, false, true, false]), + &mut execution_ctx() + ); + insta::assert_snapshot!(traced.trace.to_string(), @" + execute_until target=AnyCanonical root=vortex.dict(bool, len=5) + iter 0 current=vortex.dict(bool, len=5) builder_active=false + ExecuteSlot slot=1 parent=vortex.dict(bool, len=5) child=vortex.binary(bool, len=3) + iter 1 current=vortex.binary(bool, len=3) stack_parent=vortex.dict(bool, len=5) slot=1 builder_active=false + Done array=vortex.bool(bool, len=3) + iter 2 current=vortex.bool(bool, len=3) stack_parent=vortex.dict(bool, len=5) slot=1 builder_active=false + pop_frame slot=1 output=vortex.dict(bool, len=5) + iter 3 current=vortex.dict(bool, len=5) builder_active=false + child_execute_parent session[0]:execute_parent_fn slot=1 parent=vortex.dict(bool, len=5) child=vortex.bool(bool, len=3) -> vortex.bool(bool, len=5) + iter 4 current=vortex.bool(bool, len=5) builder_active=false + return output=vortex.bool(bool, len=5) + "); + + Ok(()) +} + +#[test] +fn trace_like_on_dict() -> VortexResult<()> { + let strings = dict_of_strings()?; + let pattern = ConstantArray::new(Scalar::from("b%"), strings.len()).into_array(); + + let like = Like.try_new_array( + strings.len(), + LikeOptions { + negated: false, + case_insensitive: false, + }, + [strings, pattern], + )?; + + let traced = trace_op(|| like.optimize())?; + insta::assert_snapshot!(traced.trace.to_string(), @" + optimize root=vortex.like(bool, len=6) session=false + reduce_parent static:LikeReduceAdaptor(Dict) slot=0 parent=vortex.like(bool, len=6) child=vortex.dict(utf8, len=6) -> vortex.dict(bool, len=6) + done output=vortex.dict(bool, len=6) + "); + + let optimized = traced.output; + let traced = trace_op(|| { + optimized + .execute::(&mut execution_ctx()) + .map(IntoArray::into_array) + })?; + + assert_arrays_eq!( + traced.output, + BoolArray::from_iter([false, true, false, true, false, false]), + &mut execution_ctx() + ); + insta::assert_snapshot!(traced.trace.to_string(), @" + execute_until target=AnyCanonical root=vortex.dict(bool, len=6) + iter 0 current=vortex.dict(bool, len=6) builder_active=false + ExecuteSlot slot=1 parent=vortex.dict(bool, len=6) child=vortex.like(bool, len=3) + iter 1 current=vortex.like(bool, len=3) stack_parent=vortex.dict(bool, len=6) slot=1 builder_active=false + Done array=vortex.bool(bool, len=3) + iter 2 current=vortex.bool(bool, len=3) stack_parent=vortex.dict(bool, len=6) slot=1 builder_active=false + pop_frame slot=1 output=vortex.dict(bool, len=6) + iter 3 current=vortex.dict(bool, len=6) builder_active=false + child_execute_parent session[0]:execute_parent_fn slot=1 parent=vortex.dict(bool, len=6) child=vortex.bool(bool, len=3) -> vortex.bool(bool, len=6) + iter 4 current=vortex.bool(bool, len=6) builder_active=false + return output=vortex.bool(bool, len=6) + "); + + Ok(()) +} diff --git a/vortex-array/src/trace_macros.rs b/vortex-array/src/trace_macros.rs new file mode 100644 index 00000000000..fd66f38b201 --- /dev/null +++ b/vortex-array/src/trace_macros.rs @@ -0,0 +1,38 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Dispatch macro used by the optimizer and executor to emit trace events. +//! +//! See [`test_harness::trace`][crate::test_harness::trace] for the harness that consumes the +//! events recorded here. + +/// Dispatch a trace event from inside the optimizer or executor. +/// +/// This macro is the only call-site interface for the trace harness. It is designed so that +/// release builds (and CodSpeed benchmark builds) compile every invocation away to nothing, +/// while debug-test builds with an active recorder route the call into +/// [`test_harness::trace`][crate::test_harness::trace]. +/// +/// # Gating +/// +/// Each invocation expands to a `cfg`-gated branch: +/// +/// - When `test` OR `_test-harness` is on, AND `codspeed` is off, the body is compiled in and +/// guarded by a runtime check for an active recorder. +/// - In all other configurations, the body is replaced with `()`. +/// +/// The macro is named `trace_op` because it records *operations* (optimizer rewrites, parent +/// kernels, execution steps, builder activity) rather than array values. The events it emits +/// describe what work the optimizer/executor did, not the contents of any array. +macro_rules! trace_op { + ($event:ident $(::<$($generic:ty),*>)?($($arg:expr),* $(,)?)) => {{ + #[cfg(all(any(test, feature = "_test-harness"), not(codspeed)))] + { + if $crate::test_harness::trace::is_active() { + $crate::test_harness::trace::$event $(::<$($generic),*>)?($($arg),*); + } + } + }}; +} + +pub(crate) use trace_op; diff --git a/vortex-btrblocks/Cargo.toml b/vortex-btrblocks/Cargo.toml index 255d2fbd075..4e22f042adf 100644 --- a/vortex-btrblocks/Cargo.toml +++ b/vortex-btrblocks/Cargo.toml @@ -39,11 +39,16 @@ vortex-zigzag = { workspace = true } vortex-zstd = { workspace = true, optional = true } [dev-dependencies] +arrow-array = { workspace = true } divan = { workspace = true } insta = { workspace = true } rstest = { workspace = true } test-with = { workspace = true } +tpchgen = { workspace = true } +tpchgen-arrow = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } +vortex-arrow = { workspace = true } +vortex-mask = { workspace = true } vortex-session = { workspace = true } [features] diff --git a/vortex-btrblocks/src/lib.rs b/vortex-btrblocks/src/lib.rs index 37915b427b6..1ca05c86b4e 100644 --- a/vortex-btrblocks/src/lib.rs +++ b/vortex-btrblocks/src/lib.rs @@ -70,6 +70,9 @@ mod builder; mod canonical_compressor; /// Compression scheme implementations. pub mod schemes; +#[cfg(test)] +#[cfg(not(codspeed))] +mod trace_tests; // Re-export framework types from vortex-compressor for backwards compatibility. // Btrblocks-specific exports. diff --git a/vortex-btrblocks/src/trace_tests.rs b/vortex-btrblocks/src/trace_tests.rs new file mode 100644 index 00000000000..ee91ffab298 --- /dev/null +++ b/vortex-btrblocks/src/trace_tests.rs @@ -0,0 +1,448 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Snapshot tests tracing how TPC-H style scan operations reduce and execute over +//! BtrBlocks-compressed lineitem columns. +//! +//! The lineitem table (SF 0.001, deterministic) compresses to the same encodings a real +//! file scan would see: bitpacked integers, `decimal_byte_parts` decimals with dictionary +//! or bitpacked byte parts, `ext(date)` over FoR + bitpacking, dictionary-of-FSST for +//! low-cardinality strings, and FSST for comments. The tests below pin down which reduce +//! rules and execute kernels fire for the scan operations TPC-H queries perform over those +//! encodings. + +use std::sync::LazyLock; + +use arrow_array::RecordBatch; +use tpchgen::distribution::Distributions; +use tpchgen::generators::LineItemGenerator; +use tpchgen::text::TextPool; +use tpchgen_arrow::LineItemArrow; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::DictArray; +use vortex_array::arrays::FilterArray; +use vortex_array::arrays::Patched; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::Struct; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::patched::use_experimental_patches; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::optimizer::ArrayOptimizer; +use vortex_array::scalar::DecimalValue; +use vortex_array::scalar::PValue; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::fns::binary::Binary; +use vortex_array::scalar_fn::fns::like::Like; +use vortex_array::scalar_fn::fns::like::LikeOptions; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::session::ArraySession; +use vortex_array::session::ArraySessionExt; +use vortex_array::test_harness::trace::Traced; +use vortex_array::test_harness::trace::trace_op; +use vortex_arrow::FromArrowArray; +use vortex_error::VortexResult; +use vortex_mask::Mask; +use vortex_session::VortexSession; + +use crate::BtrBlocksCompressor; + +/// A session with the default Vortex encodings registered. +/// +/// Reduce-parent and execute-parent kernels live in the session registry, so an encoding whose +/// `initialize` was never called still decodes correctly but silently contributes none of its +/// pushdown kernels — the traces below would degrade to plain canonicalization without failing +/// any value assertion. +/// +/// This mirrors `vortex_file::register_default_encodings`, copied rather than called so that +/// `vortex-btrblocks` does not depend on `vortex-file` (which depends on this crate). Keep the +/// two in step when encodings are added. `bytebool` and `tensor` are the only entries omitted: +/// this crate does not depend on them, so the compressor cannot emit them. +fn trace_session() -> VortexSession { + let session = VortexSession::empty().with::(); + + vortex_fsst::initialize(&session); + #[cfg(feature = "unstable_encodings")] + vortex_onpair::initialize(&session); + vortex_zigzag::initialize(&session); + + { + let arrays = session.arrays(); + #[cfg(feature = "pco")] + arrays.register(vortex_pco::Pco); + #[cfg(feature = "zstd")] + arrays.register(vortex_zstd::Zstd); + #[cfg(all(feature = "zstd", feature = "unstable_encodings"))] + arrays.register(vortex_zstd::ZstdBuffers); + if use_experimental_patches() { + arrays.register(Patched); + } + } + + vortex_alp::initialize(&session); + vortex_datetime_parts::initialize(&session); + vortex_decimal_byte_parts::initialize(&session); + vortex_fastlanes::initialize(&session); + vortex_runend::initialize(&session); + vortex_sequence::initialize(&session); + vortex_sparse::initialize(&session); + + session +} + +fn execution_ctx() -> ExecutionCtx { + ExecutionCtx::new(trace_session()) +} + +/// A 1MiB TPC-H text pool. The spec-default 300MiB pool takes several seconds to initialize in +/// debug builds; a smaller pool keeps the same text distribution and is deterministic. +static TEXT_POOL: LazyLock = + LazyLock::new(|| TextPool::new(1 << 20, Distributions::static_default())); + +/// The first 4096 rows of TPC-H lineitem at scale factor 0.001. Both tpchgen and the +/// compressor (fixed sampling seed) are deterministic, so the resulting encodings are stable. +fn lineitem() -> VortexResult { + let generator = LineItemGenerator::new_with_distributions_and_text_pool( + 0.001, + 1, + 1, + Distributions::static_default(), + &TEXT_POOL, + ); + let batch: RecordBatch = LineItemArrow::new(generator) + .with_batch_size(1 << 12) + .next() + .expect("at least one batch"); + ArrayRef::from_arrow(&batch, false) +} + +fn compressed_lineitem() -> VortexResult { + BtrBlocksCompressor::default().compress(&lineitem()?, &mut execution_ctx()) +} + +fn field(array: &ArrayRef, name: &str) -> VortexResult { + Ok(array.as_::().unmasked_field_by_name(name)?.clone()) +} + +/// Trace the optimize pass and then the execution of `array`, asserting that the canonical +/// result matches running the same operation over the uncompressed column. +fn optimize_then_execute( + array: ArrayRef, + expected: &ArrayRef, +) -> VortexResult<[Traced; 2]> { + let optimized = trace_op(|| array.optimize())?; + let optimized_output = optimized.output.clone(); + let executed = trace_op(|| { + optimized_output + .execute::(&mut execution_ctx()) + .map(IntoArray::into_array) + })?; + let expected = expected + .clone() + .execute::(&mut execution_ctx())? + .into_array(); + assert_arrays_eq!(executed.output, expected, &mut execution_ctx()); + Ok([ + Traced { + output: optimized.output, + trace: optimized.trace, + }, + executed, + ]) +} + +/// Q6-style predicate over the shipdate column: `l_shipdate >= DATE '1994-01-01'`. +/// +/// The column compresses to `ext(date) -> for -> bitpacked`. +fn shipdate_predicate(column: ArrayRef, len: usize) -> VortexResult { + let DType::Extension(ext) = column.dtype().clone() else { + panic!("expected extension dtype for l_shipdate") + }; + // 8766 days since the epoch = 1994-01-01. + let cutoff = Scalar::extension_ref( + ext, + Scalar::primitive_value(PValue::I32(8766), PType::I32, Nullability::NonNullable), + ); + Binary.try_new_array( + len, + Operator::Gte, + [column, ConstantArray::new(cutoff, len).into_array()], + ) +} + +#[test] +fn trace_scan_compare_on_compressed_shipdate() -> VortexResult<()> { + let compressed = compressed_lineitem()?; + let len = compressed.len(); + + let lazy = shipdate_predicate(field(&compressed, "l_shipdate")?, len)?; + let expected = shipdate_predicate(field(&lineitem()?, "l_shipdate")?, len)?; + let [optimized, executed] = optimize_then_execute(lazy, &expected)?; + + // No reduce rule rewrites a comparison over the extension array; the extension compare + // kernel handles it at execution time by comparing the underlying storage. + insta::assert_snapshot!(optimized.trace.to_string(), @""); + insta::assert_snapshot!(executed.trace.to_string(), @" + execute_until target=AnyCanonical root=vortex.binary(bool, len=4096) + iter 0 current=vortex.binary(bool, len=4096) builder_active=false + child_execute_parent session[0]:execute_parent_fn slot=0 parent=vortex.binary(bool, len=4096) child=vortex.ext(vortex.date[days](i32), len=4096) -> vortex.binary(bool, len=4096) + iter 1 current=vortex.binary(bool, len=4096) builder_active=false + execute_until target=AnyCanonical root=fastlanes.for(i32, len=4096) + iter 0 current=fastlanes.for(i32, len=4096) builder_active=false + execute_until target=AnyCanonical root=fastlanes.bitpacked(i32, len=4096) + iter 0 current=fastlanes.bitpacked(i32, len=4096) builder_active=false + Done array=vortex.primitive(i32, len=4096) + iter 1 current=vortex.primitive(i32, len=4096) builder_active=false + return output=vortex.primitive(i32, len=4096) + Done array=vortex.primitive(i32, len=4096) + iter 1 current=vortex.primitive(i32, len=4096) builder_active=false + return output=vortex.primitive(i32, len=4096) + Done array=vortex.bool(bool, len=4096) + iter 2 current=vortex.bool(bool, len=4096) builder_active=false + return output=vortex.bool(bool, len=4096) + "); + + Ok(()) +} + +/// Q6-style predicate over the quantity column: `l_quantity < 24`. +/// +/// The column compresses to `decimal_byte_parts -> dict -> bitpacked/sequence`. +fn quantity_predicate(column: ArrayRef, len: usize) -> VortexResult { + let cutoff = Scalar::decimal( + DecimalValue::I128(2400), + DecimalDType::new(15, 2), + Nullability::NonNullable, + ); + Binary.try_new_array( + len, + Operator::Lt, + [column, ConstantArray::new(cutoff, len).into_array()], + ) +} + +#[test] +fn trace_scan_compare_on_compressed_quantity() -> VortexResult<()> { + let compressed = compressed_lineitem()?; + let len = compressed.len(); + + let lazy = quantity_predicate(field(&compressed, "l_quantity")?, len)?; + let expected = quantity_predicate(field(&lineitem()?, "l_quantity")?, len)?; + let [optimized, executed] = optimize_then_execute(lazy, &expected)?; + + // No reduce rule rewrites a comparison over decimal_byte_parts; its compare kernel fires + // at execution time and pushes the comparison into the byte-parts dictionary, whose values + // are then compared and decoded. + insta::assert_snapshot!(optimized.trace.to_string(), @""); + insta::assert_snapshot!(executed.trace.to_string(), @" + execute_until target=AnyCanonical root=vortex.binary(bool, len=4096) + iter 0 current=vortex.binary(bool, len=4096) builder_active=false + optimize root=vortex.binary(bool, len=4096) session=false + reduce_parent static:DictionaryScalarFnValuesPushDownRule slot=0 parent=vortex.binary(bool, len=4096) child=vortex.dict(i16, len=4096) -> vortex.dict(bool, len=4096) + done output=vortex.dict(bool, len=4096) + child_execute_parent session[0]:execute_parent_fn slot=0 parent=vortex.binary(bool, len=4096) child=vortex.decimal_byte_parts(decimal(15,2), len=4096) -> vortex.dict(bool, len=4096) + iter 1 current=vortex.dict(bool, len=4096) builder_active=false + ExecuteSlot slot=0 parent=vortex.dict(bool, len=4096) child=fastlanes.bitpacked(u8, len=4096) + iter 2 current=fastlanes.bitpacked(u8, len=4096) stack_parent=vortex.dict(bool, len=4096) slot=0 builder_active=false + Done array=vortex.primitive(u8, len=4096) + iter 3 current=vortex.primitive(u8, len=4096) stack_parent=vortex.dict(bool, len=4096) slot=0 builder_active=false + pop_frame slot=0 output=vortex.dict(bool, len=4096) + iter 4 current=vortex.dict(bool, len=4096) builder_active=false + ExecuteSlot slot=1 parent=vortex.dict(bool, len=4096) child=vortex.binary(bool, len=50) + iter 5 current=vortex.binary(bool, len=50) stack_parent=vortex.dict(bool, len=4096) slot=1 builder_active=false + execute_until target=AnyCanonical root=vortex.sequence(i16, len=50) + iter 0 current=vortex.sequence(i16, len=50) builder_active=false + Done array=vortex.primitive(i16, len=50) + iter 1 current=vortex.primitive(i16, len=50) builder_active=false + return output=vortex.primitive(i16, len=50) + Done array=vortex.bool(bool, len=50) + iter 6 current=vortex.bool(bool, len=50) stack_parent=vortex.dict(bool, len=4096) slot=1 builder_active=false + pop_frame slot=1 output=vortex.dict(bool, len=4096) + iter 7 current=vortex.dict(bool, len=4096) builder_active=false + child_execute_parent session[0]:execute_parent_fn slot=1 parent=vortex.dict(bool, len=4096) child=vortex.bool(bool, len=50) -> vortex.bool(bool, len=4096) + iter 8 current=vortex.bool(bool, len=4096) builder_active=false + return output=vortex.bool(bool, len=4096) + "); + + Ok(()) +} + +/// Q12-style predicate over the shipmode column: `l_shipmode = 'AIR'`. +/// +/// The column compresses to `dict -> {bitpacked codes, fsst values}`. +fn shipmode_predicate(column: ArrayRef, len: usize) -> VortexResult { + Binary.try_new_array( + len, + Operator::Eq, + [ + column, + ConstantArray::new(Scalar::from("AIR"), len).into_array(), + ], + ) +} + +#[test] +fn trace_scan_compare_on_compressed_shipmode() -> VortexResult<()> { + let compressed = compressed_lineitem()?; + let len = compressed.len(); + + let lazy = shipmode_predicate(field(&compressed, "l_shipmode")?, len)?; + let expected = shipmode_predicate(field(&lineitem()?, "l_shipmode")?, len)?; + let [optimized, executed] = optimize_then_execute(lazy, &expected)?; + + insta::assert_snapshot!(optimized.trace.to_string(), @" + optimize root=vortex.binary(bool, len=4096) session=false + reduce_parent static:DictionaryScalarFnValuesPushDownRule slot=0 parent=vortex.binary(bool, len=4096) child=vortex.dict(utf8, len=4096) -> vortex.dict(bool, len=4096) + done output=vortex.dict(bool, len=4096) + "); + insta::assert_snapshot!(executed.trace.to_string(), @" + execute_until target=AnyCanonical root=vortex.dict(bool, len=4096) + iter 0 current=vortex.dict(bool, len=4096) builder_active=false + ExecuteSlot slot=0 parent=vortex.dict(bool, len=4096) child=fastlanes.bitpacked(u8, len=4096) + iter 1 current=fastlanes.bitpacked(u8, len=4096) stack_parent=vortex.dict(bool, len=4096) slot=0 builder_active=false + Done array=vortex.primitive(u8, len=4096) + iter 2 current=vortex.primitive(u8, len=4096) stack_parent=vortex.dict(bool, len=4096) slot=0 builder_active=false + pop_frame slot=0 output=vortex.dict(bool, len=4096) + iter 3 current=vortex.dict(bool, len=4096) builder_active=false + ExecuteSlot slot=1 parent=vortex.dict(bool, len=4096) child=vortex.binary(bool, len=7) + iter 4 current=vortex.binary(bool, len=7) stack_parent=vortex.dict(bool, len=4096) slot=1 builder_active=false + child_execute_parent session[0]:execute_parent_fn slot=0 parent=vortex.binary(bool, len=7) child=vortex.fsst(utf8, len=7) -> vortex.binary(bool, len=7) + iter 5 current=vortex.binary(bool, len=7) stack_parent=vortex.dict(bool, len=4096) slot=1 builder_active=false + child_execute_parent session[0]:execute_parent_fn slot=0 parent=vortex.binary(bool, len=7) child=vortex.varbin(binary, len=7) -> vortex.bool(bool, len=7) + iter 6 current=vortex.bool(bool, len=7) stack_parent=vortex.dict(bool, len=4096) slot=1 builder_active=false + pop_frame slot=1 output=vortex.dict(bool, len=4096) + iter 7 current=vortex.dict(bool, len=4096) builder_active=false + child_execute_parent session[0]:execute_parent_fn slot=1 parent=vortex.dict(bool, len=4096) child=vortex.bool(bool, len=7) -> vortex.bool(bool, len=4096) + iter 8 current=vortex.bool(bool, len=4096) builder_active=false + return output=vortex.bool(bool, len=4096) + "); + + Ok(()) +} + +/// Q13-style predicate over the comment column: `l_comment LIKE '%special%'`. +/// +/// The column compresses to `fsst -> bitpacked lengths/offsets`. +fn comment_predicate(column: ArrayRef, len: usize) -> VortexResult { + Like.try_new_array( + len, + LikeOptions { + negated: false, + case_insensitive: false, + }, + [ + column, + ConstantArray::new(Scalar::from("%special%"), len).into_array(), + ], + ) +} + +#[test] +fn trace_scan_like_on_compressed_comment() -> VortexResult<()> { + let compressed = compressed_lineitem()?; + let len = compressed.len(); + + let lazy = comment_predicate(field(&compressed, "l_comment")?, len)?; + let expected = comment_predicate(field(&lineitem()?, "l_comment")?, len)?; + let [optimized, executed] = optimize_then_execute(lazy, &expected)?; + + // No reduce rule rewrites a like over FSST; the FSST like kernel compiles the pattern and + // matches in compressed space at execution time. + insta::assert_snapshot!(optimized.trace.to_string(), @""); + insta::assert_snapshot!(executed.trace.to_string(), @" + execute_until target=AnyCanonical root=vortex.like(bool, len=4096) + iter 0 current=vortex.like(bool, len=4096) builder_active=false + child_execute_parent session[0]:execute_parent_fn slot=0 parent=vortex.like(bool, len=4096) child=vortex.fsst(utf8, len=4096) -> vortex.bool(bool, len=4096) + iter 1 current=vortex.bool(bool, len=4096) builder_active=false + return output=vortex.bool(bool, len=4096) + "); + + Ok(()) +} + +/// A narrowed projection of lineitem with one column per interesting compressed encoding: +/// decimal_byte_parts, ext-over-FoR dates, and dict-of-FSST strings. +fn project(table: &ArrayRef) -> VortexResult { + Ok(StructArray::from_fields(&[ + ("l_quantity", field(table, "l_quantity")?), + ("l_shipdate", field(table, "l_shipdate")?), + ("l_shipmode", field(table, "l_shipmode")?), + ])? + .into_array()) +} + +#[test] +fn trace_scan_filter_on_compressed_table() -> VortexResult<()> { + let compressed = project(&compressed_lineitem()?)?; + let len = compressed.len(); + let mask = Mask::from_iter((0..len).map(|i| i % 97 == 0)); + + let lazy = FilterArray::try_new(compressed, mask.clone())?.into_array(); + let expected = FilterArray::try_new(project(&lineitem()?)?, mask)?.into_array(); + let [optimized, executed] = optimize_then_execute(lazy, &expected)?; + + insta::assert_snapshot!(optimized.trace.to_string(), @" + optimize root=vortex.filter({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=43) session=false + optimize root=vortex.filter(decimal(15,2), len=43) session=false + optimize root=vortex.filter(i16, len=43) session=false + reduce_parent static:FilterReduceAdaptor(Dict) slot=0 parent=vortex.filter(i16, len=43) child=vortex.dict(i16, len=4096) -> vortex.dict(i16, len=43) + done output=vortex.dict(i16, len=43) + reduce_parent static:DecimalBytePartsFilterPushDownRule slot=0 parent=vortex.filter(decimal(15,2), len=43) child=vortex.decimal_byte_parts(decimal(15,2), len=4096) -> vortex.decimal_byte_parts(decimal(15,2), len=43) + done output=vortex.decimal_byte_parts(decimal(15,2), len=43) + optimize root=vortex.filter(vortex.date[days](i32), len=43) session=false + optimize root=vortex.filter(i32, len=43) session=false + reduce_parent static:FoRFilterPushDownRule slot=0 parent=vortex.filter(i32, len=43) child=fastlanes.for(i32, len=4096) -> fastlanes.for(i32, len=43) + done output=fastlanes.for(i32, len=43) + reduce_parent static:ExtensionFilterPushDownRule slot=0 parent=vortex.filter(vortex.date[days](i32), len=43) child=vortex.ext(vortex.date[days](i32), len=4096) -> vortex.ext(vortex.date[days](i32), len=43) + done output=vortex.ext(vortex.date[days](i32), len=43) + optimize root=vortex.filter(utf8, len=43) session=false + reduce_parent static:FilterReduceAdaptor(Dict) slot=0 parent=vortex.filter(utf8, len=43) child=vortex.dict(utf8, len=4096) -> vortex.dict(utf8, len=43) + done output=vortex.dict(utf8, len=43) + reduce FilterStructRule: vortex.filter({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=43) -> vortex.struct({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=43) + done output=vortex.struct({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=43) + "); + insta::assert_snapshot!(executed.trace.to_string(), @" + execute_until target=AnyCanonical root=vortex.struct({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=43) + iter 0 current=vortex.struct({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=43) builder_active=false + return output=vortex.struct({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=43) + "); + + Ok(()) +} + +#[test] +fn trace_scan_take_on_compressed_table() -> VortexResult<()> { + let compressed = project(&compressed_lineitem()?)?; + let len = compressed.len() as u64; + // A take is expressed as a `DictArray` whose codes are the take indices. + let indices = PrimitiveArray::from_iter((0..64u64).map(|i| (i * 941) % len)).into_array(); + + let lazy = DictArray::try_new(indices.clone(), compressed)?.into_array(); + let expected = DictArray::try_new(indices, project(&lineitem()?)?)?.into_array(); + let [optimized, executed] = optimize_then_execute(lazy, &expected)?; + + insta::assert_snapshot!(optimized.trace.to_string(), @" + optimize root=vortex.dict({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=64) session=false + optimize root=vortex.dict(vortex.date[days](i32), len=64) session=false + reduce_parent static:TakeReduceAdaptor(Extension) slot=1 parent=vortex.dict(vortex.date[days](i32), len=64) child=vortex.ext(vortex.date[days](i32), len=4096) -> vortex.ext(vortex.date[days](i32), len=64) + done output=vortex.ext(vortex.date[days](i32), len=64) + reduce_parent static:TakeReduceAdaptor(Struct) slot=1 parent=vortex.dict({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=64) child=vortex.struct({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=4096) -> vortex.struct({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=64) + done output=vortex.struct({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=64) + "); + insta::assert_snapshot!(executed.trace.to_string(), @" + execute_until target=AnyCanonical root=vortex.struct({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=64) + iter 0 current=vortex.struct({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=64) builder_active=false + return output=vortex.struct({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=64) + "); + + Ok(()) +}