From 5fb0d436d3c38d5f5058d92f74536d8980463350 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Wed, 6 May 2026 14:54:51 +0100 Subject: [PATCH 01/15] Exectuion and optimization tracing harness for testing Signed-off-by: Adam Gutglick --- .../src/arrays/scalar_fn/vtable/mod.rs | 1 - vortex-array/src/executor.rs | 182 ++- vortex-array/src/lib.rs | 4 + vortex-array/src/optimizer/mod.rs | 81 +- vortex-array/src/optimizer/rules.rs | 30 + vortex-array/src/test_harness.rs | 3 + vortex-array/src/test_harness/trace.rs | 1278 +++++++++++++++++ vortex-array/src/test_harness/trace_arrays.rs | 350 +++++ vortex-array/src/trace_macros.rs | 47 + 9 files changed, 1909 insertions(+), 67 deletions(-) create mode 100644 vortex-array/src/test_harness/trace.rs create mode 100644 vortex-array/src/test_harness/trace_arrays.rs create mode 100644 vortex-array/src/trace_macros.rs 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..62b29d472c1 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -171,23 +171,56 @@ impl ArrayRef { let kernels = execute_parent_kernels.as_ref(); let max_iterations = max_iterations(); - for _ in 0..max_iterations { + crate::trace_array!(record_execute_until_start::(¤t_array)); + + for iteration in 0..max_iterations { + crate::trace_array_use!(iteration); + crate::trace_array!(record_execute_until_iteration( + iteration + 1, + ¤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); + crate::trace_array!(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)); + crate::trace_array!(record_execute_until_return(¤t_array)); return Ok(current_array); } Some(frame) => { + let trace_pop_frame = crate::trace_array_value!( + Some(( + frame.parent_array.clone(), + current_array.clone(), + frame.slot_idx + )), + None::<(ArrayRef, ArrayRef, usize)> + ); (current_array, current_builder) = pop_frame(frame, current_array)?; + if let Some((parent_before, child_before, slot_idx)) = trace_pop_frame { + crate::trace_array_use!(parent_before, child_before, slot_idx,); + crate::trace_array!(record_execute_until_pop_frame( + &parent_before, + slot_idx, + &child_before, + ¤t_array, + )); + } continue; } } @@ -207,6 +240,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 +249,49 @@ 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())?; + crate::trace_array!(record_execute_optimized(&result, &optimized)); + current_array = optimized; current_builder = frame.parent_builder; continue; } + if current_builder.is_none() && stack.last().is_some() { + crate::trace_array!(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())?; + crate::trace_array!(record_execute_optimized(&rewritten, &optimized)); + current_array = optimized; continue; } + if current_builder.is_none() { + crate::trace_array!(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(); + crate::trace_array!(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 - )); + + crate::trace_array!(record_execute_slot(i, &parent, &child)); stack.push(StackFrame { parent_array: parent, parent_builder: current_builder.take(), @@ -264,6 +305,7 @@ impl ArrayRef { } ExecutionStep::AppendChild(i) => { if current_builder.is_none() { + crate::trace_array!(record_builder_start(&array)); current_builder = Some(builder_with_capacity_in( ctx.allocator(), array.dtype(), @@ -271,10 +313,10 @@ impl ArrayRef { )); } let (parent, child) = unsafe { array.take_slot_unchecked(i) }?; - ctx.log(format_args!( - "AppendChild({i}): appending {} into builder", - child - )); + + crate::trace_array!(record_append_child(i, &parent, &child)); + crate::trace_array!(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 +328,8 @@ impl ArrayRef { current_array = parent; } ExecutionStep::Done => { - ctx.log(format_args!("Done: {}", array)); + let had_builder = current_builder.is_some(); + crate::trace_array!(record_execute_done(&array)); (current_array, current_builder) = finalize_done( array, current_builder, @@ -295,6 +338,9 @@ impl ArrayRef { stats, encoding_id, )?; + if had_builder { + crate::trace_array!(record_builder_finish(¤t_array)); + } } } } @@ -419,40 +465,49 @@ impl Drop for ExecutionCtx { /// `AppendChild` is returned. impl Executable for ArrayRef { fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + crate::trace_array!(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(); + crate::trace_array!(record_single_step_applied("canonical", &array, &output)); + return Ok(output); } + crate::trace_array!(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()); + crate::trace_array!(record_single_step_applied("reduce", &array, &reduced)); return Ok(reduced); } + crate::trace_array!(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()); + crate::trace_array!(record_single_step_applied( + "reduce_parent", + &array, + &reduced_parent, + )); return Ok(reduced_parent); } } + crate::trace_array!(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 +518,22 @@ impl Executable for ArrayRef { executed_parent .statistics() .inherit_from(array.statistics()); + crate::trace_array!(record_single_step_applied( + "execute_parent", + &array, + &executed_parent, + )); return Ok(executed_parent); } } + crate::trace_array!(record_single_step_phase_none("execute_parent", &array)); + crate::trace_array!(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)); + crate::trace_array!(record_execute_done(&array)); Ok(array) } ExecutionStep::ExecuteSlot(i, _) => { @@ -484,9 +545,12 @@ impl Executable for ArrayRef { } ExecutionStep::AppendChild(_) => { // Single-step: build the entire parent via the builder path. + crate::trace_array!(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(); + crate::trace_array!(record_builder_finish(&output)); + Ok(output) } } } @@ -562,15 +626,18 @@ fn finalize_done( } fn execute_parent_for_child( + phase: &'static str, parent: &ArrayRef, child: &ArrayRef, slot_idx: usize, kernels: &ParentExecutionKernels, ctx: &mut ExecutionCtx, ) -> VortexResult> { + crate::trace_array_use!(phase); let key = execute_parent_key(parent.encoding_id(), child.encoding_id()); if let Some(plugins) = kernels.get(&key) { - for plugin in plugins.as_ref() { + for (plugin_idx, plugin) in plugins.as_ref().iter().enumerate() { + crate::trace_array_use!(plugin_idx); if let Some(result) = plugin.execute_parent(child, parent, slot_idx, ctx)? { if cfg!(debug_assertions) { vortex_ensure!( @@ -582,8 +649,26 @@ fn execute_parent_for_child( "Executed parent canonical dtype mismatch" ); } + crate::trace_array!(record_execute_parent_applied( + phase, + parent, + child, + slot_idx, + crate::test_harness::trace::TraceSource::Session(plugin_idx), + "execute_parent_fn", + &result, + )); return Ok(Some(result)); } + crate::trace_array!(record_execute_parent_attempt( + phase, + parent, + child, + slot_idx, + crate::test_harness::trace::TraceSource::Session(plugin_idx), + "execute_parent_fn", + crate::test_harness::trace::AttemptOutcome::Declined, + )); } } @@ -598,9 +683,14 @@ fn try_execute_parent( ) -> VortexResult> { 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( + "child_execute_parent", + array, + child, + slot_idx, + kernels, + ctx, + )? { ctx.log(format_args!( "execute_parent: slot[{}]({}) rewrote {} -> {}", slot_idx, @@ -709,8 +799,10 @@ 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(); + crate::trace_array!(record_execute_step_request::(&array, slot_idx)); Self { - array: array.into_array(), + array, step: ExecutionStep::ExecuteSlot(slot_idx, M::matches), } } @@ -719,8 +811,10 @@ 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(); + crate::trace_array!(record_append_child_request(&array, slot_idx)); 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..a860884dffb 100644 --- a/vortex-array/src/lib.rs +++ b/vortex-array/src/lib.rs @@ -105,6 +105,10 @@ use crate::stats::session::StatsSession; pub mod aggregate_fn; #[doc(hidden)] pub mod aliases; +mod trace_macros; +pub(crate) use trace_macros::trace_array; +pub(crate) use trace_macros::trace_array_use; +pub(crate) use trace_macros::trace_array_value; 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..e09ece28a76 100644 --- a/vortex-array/src/optimizer/mod.rs +++ b/vortex-array/src/optimizer/mod.rs @@ -73,22 +73,24 @@ fn try_optimize( let mut any_optimizations = false; let array_ref = session.map(|s| s.kernels()); + crate::trace_array!(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 { + crate::trace_array!(record_optimize_loop_start(¤t_array)); if let Some(new_array) = current_array.reduce()? { current_array = new_array; any_optimizations = true; + crate::trace_array!(record_optimize_loop_end()); continue; } + crate::trace_array!(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 +99,62 @@ fn try_optimize( && let Some(plugins) = array_ref.find_reduce_parent(current_array.encoding_id(), child.encoding_id()) { - for plugin in plugins.as_ref() { + for (plugin_idx, plugin) in plugins.as_ref().iter().enumerate() { + crate::trace_array_use!(plugin_idx); if let Some(new_array) = plugin(child, ¤t_array, slot_idx)? { - current_array = new_array; - any_optimizations = true; - continue 'outer; + crate::trace_array!(record_parent_reduce_applied( + ¤t_array, + child, + slot_idx, + crate::test_harness::trace::TraceSource::Session(plugin_idx), + "reduce_parent_fn", + &new_array, + )); + parent_reduced = Some(new_array); + break; } + crate::trace_array!(record_parent_reduce_attempt( + ¤t_array, + child, + slot_idx, + crate::test_harness::trace::TraceSource::Session(plugin_idx), + "reduce_parent_fn", + crate::test_harness::trace::AttemptOutcome::Declined, + )); + } + 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; + crate::trace_array!(record_optimize_loop_end()); + continue; + } + + crate::trace_array!(record_optimize_parent_reduce_none(¤t_array)); + crate::trace_array!(record_optimize_loop_end()); + // No more optimizations can be applied - break; - } + crate::trace_array!(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 +164,8 @@ fn try_optimize_recursive( let mut current_array = array.clone(); let mut any_optimizations = false; + crate::trace_array!(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 +177,11 @@ fn try_optimize_recursive( match slot { Some(child) => { if let Some(new_child) = try_optimize_recursive(child, session)? { + crate::trace_array!(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..5ccd48e10c3 100644 --- a/vortex-array/src/optimizer/rules.rs +++ b/vortex-array/src/optimizer/rules.rs @@ -133,8 +133,14 @@ impl ReduceRuleSet { pub fn evaluate(&self, array: ArrayView<'_, V>) -> VortexResult> { for rule in self.rules.iter() { if let Some(reduced) = rule.reduce(array)? { + crate::trace_array!(record_reduce_applied(array.array(), *rule, &reduced)); return Ok(Some(reduced)); } + crate::trace_array!(record_reduce_attempt( + array.array(), + *rule, + crate::test_harness::trace::AttemptOutcome::Declined, + )); } Ok(None) } @@ -176,6 +182,14 @@ impl ParentRuleSet { ) -> VortexResult> { for rule in self.rules.iter() { if !rule.matches(parent) { + crate::trace_array!(record_parent_reduce_attempt( + parent, + child.array(), + child_idx, + crate::test_harness::trace::TraceSource::Static, + crate::test_harness::trace::compact_label(*rule), + crate::test_harness::trace::AttemptOutcome::NoMatch, + )); continue; } if let Some(reduced) = rule.reduce_parent(child, parent, child_idx)? { @@ -198,8 +212,24 @@ impl ParentRuleSet { ); } + crate::trace_array!(record_parent_reduce_applied( + parent, + child.array(), + child_idx, + crate::test_harness::trace::TraceSource::Static, + crate::test_harness::trace::compact_label(*rule), + &reduced, + )); return Ok(Some(reduced)); } + crate::trace_array!(record_parent_reduce_attempt( + parent, + child.array(), + child_idx, + crate::test_harness::trace::TraceSource::Static, + crate::test_harness::trace::compact_label(*rule), + crate::test_harness::trace::AttemptOutcome::Declined, + )); } Ok(None) } diff --git a/vortex-array/src/test_harness.rs b/vortex-array/src/test_harness.rs index 98be3a85ce4..04b9ad8d888 100644 --- a/vortex-array/src/test_harness.rs +++ b/vortex-array/src/test_harness.rs @@ -12,6 +12,9 @@ use crate::ExecutionCtx; use crate::arrays::BoolArray; use crate::arrays::bool::BoolArrayExt; +pub mod trace; +pub mod trace_arrays; + /// 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.rs b/vortex-array/src/test_harness/trace.rs new file mode 100644 index 00000000000..c8e60a6457f --- /dev/null +++ b/vortex-array/src/test_harness/trace.rs @@ -0,0 +1,1278 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::cell::Cell; +use std::cell::RefCell; +use std::fmt; +use std::fmt::Debug; +use std::fmt::Display; + +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_array_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 default trace output. +/// +/// The default resolution records the rule rewrites, parent kernels, execution steps, and builder +/// activity that actually executed. Use [`trace_array_with`] and [`TraceResolution::Attempts`] +/// when a test needs to assert on every declined rule or kernel attempt. +pub fn trace_array(f: impl FnOnce() -> VortexResult) -> VortexResult> { + trace_array_with(TraceOptions::default(), f) +} + +/// Run `f` while capturing trace output using `options`. +/// +/// 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_array_with( + options: TraceOptions, + f: impl FnOnce() -> VortexResult, +) -> VortexResult> { + ACTIVE_TRACE.with(|active| { + let mut active = active.borrow_mut(); + if active.is_some() { + return Err(vortex_err!("trace_array captures cannot be nested")); + } + *active = Some(TraceRecorder::new(options)); + Ok(()) + })?; + TRACE_INTEREST.with(|interest| interest.set(TraceInterest::from(options.resolution))); + + let guard = ActiveTraceGuard; + 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] +pub(crate) fn is_active() -> bool { + TRACE_INTEREST.with(|interest| interest.get().is_active()) +} + +#[inline] +pub(crate) fn if_active(enabled: impl FnOnce() -> R, disabled: impl FnOnce() -> R) -> R { + if is_active() { enabled() } else { disabled() } +} + +#[inline] +fn attempts_enabled() -> bool { + TRACE_INTEREST.with(|interest| interest.get() == TraceInterest::Attempts) +} + +#[derive(Clone, Copy, Debug)] +pub(crate) 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)] +pub(crate) 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, + } + } +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub(crate) struct ArraySummary { + encoding: String, + len: usize, + dtype: String, +} + +impl ArraySummary { + pub(crate) fn new(array: &ArrayRef) -> Self { + Self { + encoding: array.encoding_id().to_string(), + len: array.len(), + dtype: array.dtype().to_string(), + } + } +} + +impl Display for ArraySummary { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{} len={} dtype={}", self.encoding, self.len, self.dtype) + } +} + +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_attempt(array: &ArrayRef, rule: &dyn Debug, outcome: AttemptOutcome) { + if !attempts_enabled() { + return; + } + record_attempt(TraceEvent::ReduceAttempt { + array: ArraySummary::new(array), + rule: compact_label(rule), + outcome, + }); +} + +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_parent_reduce_attempt( + parent: &ArrayRef, + child: &ArrayRef, + slot_idx: usize, + source: TraceSource, + rule: impl Into, + outcome: AttemptOutcome, +) { + if !attempts_enabled() { + return; + } + record_attempt(TraceEvent::ParentReduceAttempt { + parent: ArraySummary::new(parent), + child: ArraySummary::new(child), + slot_idx, + source, + rule: rule.into(), + outcome, + }); +} + +pub(crate) 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( + parent: &ArrayRef, + slot_idx: usize, + child: &ArrayRef, + output: &ArrayRef, +) { + record(TraceEvent::ExecuteUntilPopFrame { + parent: ArraySummary::new(parent), + slot_idx, + child: ArraySummary::new(child), + output: ArraySummary::new(output), + }); +} + +pub(crate) 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_attempt(TraceEvent::ExecuteParentAttempt { + phase, + parent: ArraySummary::new(parent), + child: ArraySummary::new(child), + slot_idx, + source, + kernel: kernel.into(), + outcome, + }); +} + +pub(crate) 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_step_request(array: &ArrayRef, slot_idx: usize) { + if !attempts_enabled() { + return; + } + record(TraceEvent::ExecutionRequest { + step: "ExecuteSlot", + parent: ArraySummary::new(array), + slot_idx, + target: Some(short_type_name::()), + }); +} + +pub(crate) fn record_append_child_request(array: &ArrayRef, slot_idx: usize) { + if !attempts_enabled() { + return; + } + record(TraceEvent::ExecutionRequest { + step: "AppendChild", + parent: ArraySummary::new(array), + slot_idx, + target: None, + }); +} + +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 record_attempt(event: TraceEvent) { + ACTIVE_TRACE.with(|active| { + if let Some(recorder) = active.borrow_mut().as_mut() + && recorder.options.resolution == TraceResolution::Attempts + { + recorder.events.push(event); + } + }); +} + +pub(crate) 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) }; +} + +struct ActiveTraceGuard; + +impl Drop for ActiveTraceGuard { + fn drop(&mut self) { + 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 { + parent: ArraySummary, + slot_idx: usize, + child: ArraySummary, + 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, + }, + ExecutionRequest { + step: &'static str, + parent: ArraySummary, + slot_idx: usize, + target: Option, + }, + 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::ExecutionRequest { .. } + | 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::ExecutionRequest { .. } + | 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 { + parent, + slot_idx, + child, + output, + } => write!( + f, + "pop_frame slot={slot_idx} parent={parent} child={child} 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::ExecutionRequest { + step, + parent, + slot_idx, + target, + } => { + write!(f, "request {step} slot={slot_idx}")?; + if let Some(target) = target { + write!(f, " target={target}")?; + } + write!(f, " parent={parent}") + } + 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 { + use vortex_error::VortexResult; + use vortex_mask::Mask; + use vortex_session::VortexSession; + + use crate::Canonical; + use crate::ExecutionCtx; + use crate::IntoArray; + use crate::arrays::ChunkedArray; + use crate::arrays::Filter; + use crate::arrays::FilterArray; + use crate::arrays::Primitive; + use crate::arrays::PrimitiveArray; + use crate::arrays::filter::FilterArraySlotsExt; + use crate::assert_arrays_eq; + use crate::optimizer::ArrayOptimizer; + use crate::session::ArraySession; + use crate::test_harness::trace::TraceOptions; + use crate::test_harness::trace::TraceResolution; + use crate::test_harness::trace::trace_array; + use crate::test_harness::trace::trace_array_with; + use crate::test_harness::trace_arrays::stack_parent_fixture; + use crate::test_harness::trace_arrays::stack_parent_session; + + #[test] + fn trace_optimize_reduce_fixpoint() -> VortexResult<()> { + let mut ctx = ExecutionCtx::new(VortexSession::empty().with::()); + 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_array(|| filter.optimize())?; + + assert!(traced.output.is::()); + assert_arrays_eq!(traced.output, values, &mut ctx); + insta::assert_snapshot!(traced.trace.to_string(), @r" +optimize root=vortex.filter len=4 dtype=i32 session=false + reduce TrivialFilterRule: vortex.filter len=4 dtype=i32 -> vortex.primitive len=4 dtype=i32 + done output=vortex.primitive len=4 dtype=i32 +"); + + Ok(()) + } + + #[test] + fn trace_optimize_parent_reduce_fixpoint_attempts() -> VortexResult<()> { + let mut ctx = ExecutionCtx::new(VortexSession::empty().with::()); + 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_array_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 ctx); + insta::assert_snapshot!(traced.trace.to_string(), @" + optimize root=vortex.filter len=2 dtype=i32 session=false + reduce_parent static:FilterReduceAdaptor(Filter) slot=0 parent=vortex.filter len=2 dtype=i32 child=vortex.filter len=4 dtype=i32 -> vortex.filter len=2 dtype=i32 + done output=vortex.filter len=2 dtype=i32 + "); + + let traced = trace_array_with( + TraceOptions { + resolution: TraceResolution::ExecutedOnly, + }, + || outer.execute::(&mut ctx), + )?; + + insta::assert_snapshot!(traced.trace.to_string(), @r" +execute_until target=AnyCanonical root=vortex.filter len=2 dtype=i32 + iter 1 current=vortex.filter len=2 dtype=i32 builder_active=false + ExecuteSlot slot=0 parent=vortex.filter len=2 dtype=i32 child=vortex.filter len=4 dtype=i32 + iter 2 current=vortex.filter len=4 dtype=i32 stack_parent=vortex.filter len=2 dtype=i32 slot=0 builder_active=false + Done array=vortex.primitive len=4 dtype=i32 + iter 3 current=vortex.primitive len=4 dtype=i32 stack_parent=vortex.filter len=2 dtype=i32 slot=0 builder_active=false + pop_frame slot=0 parent=vortex.filter len=2 dtype=i32 child=vortex.primitive len=4 dtype=i32 output=vortex.filter len=2 dtype=i32 + iter 4 current=vortex.filter len=2 dtype=i32 builder_active=false + Done array=vortex.primitive len=2 dtype=i32 + iter 5 current=vortex.primitive len=2 dtype=i32 builder_active=false + return output=vortex.primitive len=2 dtype=i32 +"); + + Ok(()) + } + + #[test] + fn trace_execution_stack_parent_kernel_attempts() -> VortexResult<()> { + let mut ctx = ExecutionCtx::new(stack_parent_session()); + let parent = stack_parent_fixture()?; + + let traced = trace_array_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 len=3 dtype=i32 + iter 1 current=vortex.test.stack-parent len=3 dtype=i32 builder_active=false + done_check target=false canonical=false + child_execute_parent attempt slot=0 parent=vortex.test.stack-parent len=3 dtype=i32 child=vortex.test.stack-child len=3 dtype=i32 source=session[0] kernel=execute_parent_fn outcome=declined + child_execute_parent attempt slot=0 parent=vortex.test.stack-parent len=3 dtype=i32 child=vortex.test.stack-child len=3 dtype=i32 source=session[1] kernel=execute_parent_fn outcome=declined + child_execute_parent none current=vortex.test.stack-parent len=3 dtype=i32 + execute encoding=vortex.test.stack-parent len=3 dtype=i32 + request ExecuteSlot slot=0 target=Primitive parent=vortex.test.stack-parent len=3 dtype=i32 + ExecuteSlot slot=0 parent=vortex.test.stack-parent len=3 dtype=i32 child=vortex.test.stack-child len=3 dtype=i32 + iter 2 current=vortex.test.stack-child len=3 dtype=i32 stack_parent=vortex.test.stack-parent len=3 dtype=i32 slot=0 builder_active=false + done_check target=false canonical=false + stack_execute_parent attempt slot=0 parent=vortex.test.stack-parent len=3 dtype=i32 child=vortex.test.stack-child len=3 dtype=i32 source=session[0] kernel=execute_parent_fn outcome=declined + stack_execute_parent applied slot=0 parent=vortex.test.stack-parent len=3 dtype=i32 child=vortex.test.stack-child len=3 dtype=i32 source=session[1] kernel=execute_parent_fn output=vortex.primitive len=3 dtype=i32 + optimize root=vortex.primitive len=3 dtype=i32 session=true + loop input=vortex.primitive len=3 dtype=i32 + reduce none array=vortex.primitive len=3 dtype=i32 + reduce_parent none array=vortex.primitive len=3 dtype=i32 + done output=vortex.primitive len=3 dtype=i32 changed=false + optimize_ctx input=vortex.primitive len=3 dtype=i32 output=vortex.primitive len=3 dtype=i32 changed=false + iter 3 current=vortex.primitive len=3 dtype=i32 builder_active=false + done_check target=true canonical=true + return output=vortex.primitive len=3 dtype=i32 + "); + + 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_array(|| { + 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 len=5 dtype=i32 + iter 1 current=vortex.chunked len=5 dtype=i32 builder_active=false + builder start array=vortex.chunked len=5 dtype=i32 + AppendChild slot=1 parent=vortex.chunked len=5 dtype=i32 child=vortex.primitive len=2 dtype=i32 + builder append child=vortex.primitive len=2 dtype=i32 + iter 2 current=vortex.chunked len=5 dtype=i32 builder_active=true + AppendChild slot=2 parent=vortex.chunked len=5 dtype=i32 child=vortex.primitive len=1 dtype=i32 + builder append child=vortex.primitive len=1 dtype=i32 + iter 3 current=vortex.chunked len=5 dtype=i32 builder_active=true + AppendChild slot=3 parent=vortex.chunked len=5 dtype=i32 child=vortex.primitive len=2 dtype=i32 + builder append child=vortex.primitive len=2 dtype=i32 + iter 4 current=vortex.chunked len=5 dtype=i32 builder_active=true + Done array=vortex.primitive len=0 dtype=i32 + builder finish output=vortex.primitive len=5 dtype=i32 + iter 5 current=vortex.primitive len=5 dtype=i32 builder_active=false + return output=vortex.primitive len=5 dtype=i32 + "); + + Ok(()) + } +} diff --git a/vortex-array/src/test_harness/trace_arrays.rs b/vortex-array/src/test_harness/trace_arrays.rs new file mode 100644 index 00000000000..7a8d0fc2ad3 --- /dev/null +++ b/vortex-array/src/test_harness/trace_arrays.rs @@ -0,0 +1,350 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Small encodings used by trace snapshot tests. + +use std::fmt::Display; +use std::fmt::Formatter; +use std::hash::Hasher; + +use smallvec::smallvec; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::ArrayEq; +use crate::EqMode; +use crate::ArrayHash; +use crate::ArrayRef; +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::Primitive; +use crate::arrays::PrimitiveArray; +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::kernels::ArrayKernelsExt; +use crate::session::ArraySession; +use crate::serde::ArrayChildren; +use crate::validity::Validity; + +/// Create a `StackParent(StackChild)` fixture. +/// +/// `StackParent` requests `ExecuteSlot` until its child is `Primitive`. `StackChild` has one +/// declined parent kernel followed by one successful parent kernel, so strict trace snapshots can +/// assert that stack parent kernels run before the child decodes itself. +pub 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. +pub 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 +} + +/// Create the child encoding used by [`stack_parent_fixture`]. +pub fn stack_child() -> VortexResult { + Ok( + Array::try_from_parts(ArrayParts::new(StackChild, test_dtype(), 3, StackChildData))? + .into_array(), + ) +} + +/// Wrap `child` in the parent encoding used by [`stack_parent_fixture`]. +pub 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 with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + with_empty_buffers(self, array, buffers) + } + + 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 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 with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + with_empty_buffers(self, array, buffers) + } + + 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 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) + } +} diff --git a/vortex-array/src/trace_macros.rs b/vortex-array/src/trace_macros.rs new file mode 100644 index 00000000000..21764b1e1e4 --- /dev/null +++ b/vortex-array/src/trace_macros.rs @@ -0,0 +1,47 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#[cfg(any(test, feature = "_test-harness"))] +macro_rules! trace_array { + ($($event:tt)*) => { + $crate::test_harness::trace::if_active( + || $crate::test_harness::trace::$($event)*, + || {}, + ) + }; +} + +#[cfg(not(any(test, feature = "_test-harness")))] +macro_rules! trace_array { + ($($event:tt)*) => {{}}; +} + +#[cfg(any(test, feature = "_test-harness"))] +macro_rules! trace_array_value { + ($enabled:expr, $disabled:expr) => { + $crate::test_harness::trace::if_active(|| $enabled, || $disabled) + }; +} + +#[cfg(not(any(test, feature = "_test-harness")))] +macro_rules! trace_array_value { + ($enabled:expr, $disabled:expr) => { + $disabled + }; +} + +#[cfg(any(test, feature = "_test-harness"))] +macro_rules! trace_array_use { + ($($value:expr),* $(,)?) => {{}}; +} + +#[cfg(not(any(test, feature = "_test-harness")))] +macro_rules! trace_array_use { + ($($value:expr),* $(,)?) => { + let _ = ($(&$value),*); + }; +} + +pub(crate) use trace_array; +pub(crate) use trace_array_use; +pub(crate) use trace_array_value; From 9523008b1778809e3d540526392f7f4b3bcca947 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Wed, 6 May 2026 21:43:49 +0100 Subject: [PATCH 02/15] Maybe better perf? Signed-off-by: Adam Gutglick --- vortex-array/src/test_harness/trace.rs | 54 ++++++++++++++++++-------- vortex-array/src/trace_macros.rs | 13 ++++--- 2 files changed, 45 insertions(+), 22 deletions(-) diff --git a/vortex-array/src/test_harness/trace.rs b/vortex-array/src/test_harness/trace.rs index c8e60a6457f..becfa326f61 100644 --- a/vortex-array/src/test_harness/trace.rs +++ b/vortex-array/src/test_harness/trace.rs @@ -6,12 +6,16 @@ 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::ArrayId; use crate::ArrayRef; +use crate::dtype::DType; /// Controls how much rule and kernel resolution detail is captured. #[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] @@ -138,6 +142,7 @@ pub fn trace_array_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() { @@ -146,9 +151,13 @@ pub fn trace_array_with( *active = Some(TraceRecorder::new(options)); Ok(()) })?; - TRACE_INTEREST.with(|interest| interest.set(TraceInterest::from(options.resolution))); + 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; + let guard = ActiveTraceGuard { interest }; let output = f(); let recorder = ACTIVE_TRACE.with(|active| { active @@ -165,18 +174,19 @@ pub fn trace_array_with( } /// Returns true when the current thread has an active trace recorder. -#[inline] +#[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] -pub(crate) fn if_active(enabled: impl FnOnce() -> R, disabled: impl FnOnce() -> R) -> R { - if is_active() { enabled() } else { disabled() } -} - -#[inline] +#[inline(always)] fn attempts_enabled() -> bool { + if ATTEMPTS_TRACE_COUNT.load(Ordering::Relaxed) == 0 { + return false; + } TRACE_INTEREST.with(|interest| interest.get() == TraceInterest::Attempts) } @@ -236,17 +246,17 @@ impl From for TraceInterest { #[derive(Clone, Debug, Eq, PartialEq)] pub(crate) struct ArraySummary { - encoding: String, + encoding: ArrayId, len: usize, - dtype: String, + dtype: DType, } impl ArraySummary { pub(crate) fn new(array: &ArrayRef) -> Self { Self { - encoding: array.encoding_id().to_string(), + encoding: array.encoding_id(), len: array.len(), - dtype: array.dtype().to_string(), + dtype: array.dtype().clone(), } } } @@ -614,10 +624,11 @@ fn record(event: TraceEvent) { } fn record_attempt(event: TraceEvent) { + if !attempts_enabled() { + return; + } ACTIVE_TRACE.with(|active| { - if let Some(recorder) = active.borrow_mut().as_mut() - && recorder.options.resolution == TraceResolution::Attempts - { + if let Some(recorder) = active.borrow_mut().as_mut() { recorder.events.push(event); } }); @@ -655,10 +666,19 @@ thread_local! { static ACTIVE_TRACE: RefCell> = const { RefCell::new(None) }; } -struct ActiveTraceGuard; +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(); diff --git a/vortex-array/src/trace_macros.rs b/vortex-array/src/trace_macros.rs index 21764b1e1e4..8692125cbd7 100644 --- a/vortex-array/src/trace_macros.rs +++ b/vortex-array/src/trace_macros.rs @@ -4,10 +4,9 @@ #[cfg(any(test, feature = "_test-harness"))] macro_rules! trace_array { ($($event:tt)*) => { - $crate::test_harness::trace::if_active( - || $crate::test_harness::trace::$($event)*, - || {}, - ) + if $crate::test_harness::trace::is_active() { + $crate::test_harness::trace::$($event)* + } }; } @@ -19,7 +18,11 @@ macro_rules! trace_array { #[cfg(any(test, feature = "_test-harness"))] macro_rules! trace_array_value { ($enabled:expr, $disabled:expr) => { - $crate::test_harness::trace::if_active(|| $enabled, || $disabled) + if $crate::test_harness::trace::is_active() { + $enabled + } else { + $disabled + } }; } From 1a39c41af40c7f05029e706f218f290ede5e55c1 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Thu, 7 May 2026 12:00:45 +0100 Subject: [PATCH 03/15] No codspeed Signed-off-by: Adam Gutglick --- vortex-array/src/test_harness.rs | 1 + vortex-array/src/trace_macros.rs | 12 ++++++------ 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/vortex-array/src/test_harness.rs b/vortex-array/src/test_harness.rs index 04b9ad8d888..74c7d380074 100644 --- a/vortex-array/src/test_harness.rs +++ b/vortex-array/src/test_harness.rs @@ -12,6 +12,7 @@ use crate::ExecutionCtx; use crate::arrays::BoolArray; use crate::arrays::bool::BoolArrayExt; +#[cfg(not(codspeed))] pub mod trace; pub mod trace_arrays; diff --git a/vortex-array/src/trace_macros.rs b/vortex-array/src/trace_macros.rs index 8692125cbd7..f5d37a92ab6 100644 --- a/vortex-array/src/trace_macros.rs +++ b/vortex-array/src/trace_macros.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -#[cfg(any(test, feature = "_test-harness"))] +#[cfg(all(any(test, feature = "_test-harness"), not(codspeed)))] macro_rules! trace_array { ($($event:tt)*) => { if $crate::test_harness::trace::is_active() { @@ -10,12 +10,12 @@ macro_rules! trace_array { }; } -#[cfg(not(any(test, feature = "_test-harness")))] +#[cfg(any(not(any(test, feature = "_test-harness")), codspeed))] macro_rules! trace_array { ($($event:tt)*) => {{}}; } -#[cfg(any(test, feature = "_test-harness"))] +#[cfg(all(any(test, feature = "_test-harness"), not(codspeed)))] macro_rules! trace_array_value { ($enabled:expr, $disabled:expr) => { if $crate::test_harness::trace::is_active() { @@ -26,19 +26,19 @@ macro_rules! trace_array_value { }; } -#[cfg(not(any(test, feature = "_test-harness")))] +#[cfg(any(not(any(test, feature = "_test-harness")), codspeed))] macro_rules! trace_array_value { ($enabled:expr, $disabled:expr) => { $disabled }; } -#[cfg(any(test, feature = "_test-harness"))] +#[cfg(all(any(test, feature = "_test-harness"), not(codspeed)))] macro_rules! trace_array_use { ($($value:expr),* $(,)?) => {{}}; } -#[cfg(not(any(test, feature = "_test-harness")))] +#[cfg(any(not(any(test, feature = "_test-harness")), codspeed))] macro_rules! trace_array_use { ($($value:expr),* $(,)?) => { let _ = ($(&$value),*); From a93ed2afeadbda1dc73f0cbcf5d7af1031deb757 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Thu, 7 May 2026 18:38:38 +0100 Subject: [PATCH 04/15] less stuff Signed-off-by: Adam Gutglick --- vortex-array/src/executor.rs | 13 ++++--- vortex-array/src/lib.rs | 2 - vortex-array/src/optimizer/mod.rs | 2 +- vortex-array/src/trace_macros.rs | 62 +++++++++++++++---------------- 4 files changed, 37 insertions(+), 42 deletions(-) diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 62b29d472c1..66112752d37 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -163,6 +163,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. + #[expect(clippy::cognitive_complexity)] pub fn execute_until(self, ctx: &mut ExecutionCtx) -> VortexResult { let mut current_array = self; let mut current_builder: Option> = None; @@ -174,7 +175,7 @@ impl ArrayRef { crate::trace_array!(record_execute_until_start::(¤t_array)); for iteration in 0..max_iterations { - crate::trace_array_use!(iteration); + crate::trace_array!(use(iteration)); crate::trace_array!(record_execute_until_iteration( iteration + 1, ¤t_array, @@ -203,17 +204,17 @@ impl ArrayRef { return Ok(current_array); } Some(frame) => { - let trace_pop_frame = crate::trace_array_value!( + let trace_pop_frame = crate::trace_array!(value( Some(( frame.parent_array.clone(), current_array.clone(), frame.slot_idx )), None::<(ArrayRef, ArrayRef, usize)> - ); + )); (current_array, current_builder) = pop_frame(frame, current_array)?; if let Some((parent_before, child_before, slot_idx)) = trace_pop_frame { - crate::trace_array_use!(parent_before, child_before, slot_idx,); + crate::trace_array!(use(parent_before, child_before, slot_idx,)); crate::trace_array!(record_execute_until_pop_frame( &parent_before, slot_idx, @@ -633,11 +634,11 @@ fn execute_parent_for_child( kernels: &ParentExecutionKernels, ctx: &mut ExecutionCtx, ) -> VortexResult> { - crate::trace_array_use!(phase); + crate::trace_array!(use(phase)); let key = execute_parent_key(parent.encoding_id(), child.encoding_id()); if let Some(plugins) = kernels.get(&key) { for (plugin_idx, plugin) in plugins.as_ref().iter().enumerate() { - crate::trace_array_use!(plugin_idx); + crate::trace_array!(use(plugin_idx)); if let Some(result) = plugin.execute_parent(child, parent, slot_idx, ctx)? { if cfg!(debug_assertions) { vortex_ensure!( diff --git a/vortex-array/src/lib.rs b/vortex-array/src/lib.rs index a860884dffb..56701b1fe2d 100644 --- a/vortex-array/src/lib.rs +++ b/vortex-array/src/lib.rs @@ -107,8 +107,6 @@ pub mod aggregate_fn; pub mod aliases; mod trace_macros; pub(crate) use trace_macros::trace_array; -pub(crate) use trace_macros::trace_array_use; -pub(crate) use trace_macros::trace_array_value; 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 e09ece28a76..202c096fe20 100644 --- a/vortex-array/src/optimizer/mod.rs +++ b/vortex-array/src/optimizer/mod.rs @@ -100,7 +100,7 @@ fn try_optimize( array_ref.find_reduce_parent(current_array.encoding_id(), child.encoding_id()) { for (plugin_idx, plugin) in plugins.as_ref().iter().enumerate() { - crate::trace_array_use!(plugin_idx); + crate::trace_array!(use(plugin_idx)); if let Some(new_array) = plugin(child, ¤t_array, slot_idx)? { crate::trace_array!(record_parent_reduce_applied( ¤t_array, diff --git a/vortex-array/src/trace_macros.rs b/vortex-array/src/trace_macros.rs index f5d37a92ab6..e7827ed5d03 100644 --- a/vortex-array/src/trace_macros.rs +++ b/vortex-array/src/trace_macros.rs @@ -1,50 +1,46 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -#[cfg(all(any(test, feature = "_test-harness"), not(codspeed)))] macro_rules! trace_array { - ($($event:tt)*) => { - if $crate::test_harness::trace::is_active() { - $crate::test_harness::trace::$($event)* + (@when_enabled { $($enabled:tt)* } else { $($disabled:tt)* }) => {{ + #[cfg(all(any(test, feature = "_test-harness"), not(codspeed)))] + { + $($enabled)* } - }; -} -#[cfg(any(not(any(test, feature = "_test-harness")), codspeed))] -macro_rules! trace_array { - ($($event:tt)*) => {{}}; -} + #[cfg(any(not(any(test, feature = "_test-harness")), codspeed))] + { + $($disabled)* + } + }}; -#[cfg(all(any(test, feature = "_test-harness"), not(codspeed)))] -macro_rules! trace_array_value { - ($enabled:expr, $disabled:expr) => { - if $crate::test_harness::trace::is_active() { - $enabled + (@if_active { $($enabled:tt)* } else { $($disabled:tt)* }) => { + $crate::trace_array!(@when_enabled { + if $crate::test_harness::trace::is_active() { + $($enabled)* + } else { + $($disabled)* + } } else { - $disabled - } + $($disabled)* + }) }; -} -#[cfg(any(not(any(test, feature = "_test-harness")), codspeed))] -macro_rules! trace_array_value { - ($enabled:expr, $disabled:expr) => { - $disabled + (use($($value:expr),* $(,)?)) => { + $crate::trace_array!(@when_enabled {} else { + let _ = ($(&$value),*); + }) }; -} -#[cfg(all(any(test, feature = "_test-harness"), not(codspeed)))] -macro_rules! trace_array_use { - ($($value:expr),* $(,)?) => {{}}; -} + (value($enabled:expr, $disabled:expr)) => { + $crate::trace_array!(@if_active { $enabled } else { $disabled }) + }; -#[cfg(any(not(any(test, feature = "_test-harness")), codspeed))] -macro_rules! trace_array_use { - ($($value:expr),* $(,)?) => { - let _ = ($(&$value),*); + ($($event:tt)*) => { + $crate::trace_array!(@if_active { + $crate::test_harness::trace::$($event)* + } else {}) }; } pub(crate) use trace_array; -pub(crate) use trace_array_use; -pub(crate) use trace_array_value; From 158cfa1f047cb0905933c8924d43afd77e991e71 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Mon, 11 May 2026 10:47:24 +0100 Subject: [PATCH 05/15] More fixes Signed-off-by: Adam Gutglick --- vortex-array/src/executor.rs | 5 ++--- vortex-array/src/test_harness/trace.rs | 17 +++-------------- 2 files changed, 5 insertions(+), 17 deletions(-) diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 66112752d37..3e11810a9c4 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -174,10 +174,10 @@ impl ArrayRef { crate::trace_array!(record_execute_until_start::(¤t_array)); - for iteration in 0..max_iterations { + for iteration in 1..=max_iterations { crate::trace_array!(use(iteration)); crate::trace_array!(record_execute_until_iteration( - iteration + 1, + iteration, ¤t_array, stack .last() @@ -280,7 +280,6 @@ impl ArrayRef { )); } - // execute step let expected_len = current_array.len(); let expected_dtype = current_array.dtype().clone(); let stats = current_array.statistics().to_array_stats(); diff --git a/vortex-array/src/test_harness/trace.rs b/vortex-array/src/test_harness/trace.rs index becfa326f61..218173afac5 100644 --- a/vortex-array/src/test_harness/trace.rs +++ b/vortex-array/src/test_harness/trace.rs @@ -339,7 +339,7 @@ pub(crate) fn record_reduce_attempt(array: &ArrayRef, rule: &dyn Debug, outcome: if !attempts_enabled() { return; } - record_attempt(TraceEvent::ReduceAttempt { + record(TraceEvent::ReduceAttempt { array: ArraySummary::new(array), rule: compact_label(rule), outcome, @@ -365,7 +365,7 @@ pub(crate) fn record_parent_reduce_attempt( if !attempts_enabled() { return; } - record_attempt(TraceEvent::ParentReduceAttempt { + record(TraceEvent::ParentReduceAttempt { parent: ArraySummary::new(parent), child: ArraySummary::new(child), slot_idx, @@ -453,7 +453,7 @@ pub(crate) fn record_execute_parent_attempt( if !attempts_enabled() { return; } - record_attempt(TraceEvent::ExecuteParentAttempt { + record(TraceEvent::ExecuteParentAttempt { phase, parent: ArraySummary::new(parent), child: ArraySummary::new(child), @@ -623,17 +623,6 @@ fn record(event: TraceEvent) { }); } -fn record_attempt(event: TraceEvent) { - if !attempts_enabled() { - return; - } - ACTIVE_TRACE.with(|active| { - if let Some(recorder) = active.borrow_mut().as_mut() { - recorder.events.push(event); - } - }); -} - pub(crate) fn compact_label(value: &dyn Debug) -> String { let label = format!("{value:?}"); if let Some(label) = adapter_field(&label, "rule") { From 5a489a95f938211a17e6b1c26270330a358fa838 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Mon, 11 May 2026 11:31:22 +0100 Subject: [PATCH 06/15] Nicer tracing Signed-off-by: Adam Gutglick --- vortex-array/src/test_harness/trace.rs | 29 +++++++++++++------------- 1 file changed, 15 insertions(+), 14 deletions(-) diff --git a/vortex-array/src/test_harness/trace.rs b/vortex-array/src/test_harness/trace.rs index 218173afac5..1d399538f3d 100644 --- a/vortex-array/src/test_harness/trace.rs +++ b/vortex-array/src/test_harness/trace.rs @@ -13,9 +13,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_err; -use crate::ArrayId; use crate::ArrayRef; -use crate::dtype::DType; /// Controls how much rule and kernel resolution detail is captured. #[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] @@ -244,26 +242,29 @@ impl From for TraceInterest { } } -#[derive(Clone, Debug, Eq, PartialEq)] -pub(crate) struct ArraySummary { - encoding: ArrayId, - len: usize, - dtype: DType, -} +/// Snapshot-friendly wrapper around [`ArrayRef`] that renders just the encoding, length, and +/// dtype using the trace format (`vortex.primitive len=4 dtype=i32`). +/// +/// Carries a clone of the [`ArrayRef`] (a cheap [`Arc`] bump) instead of cloning individual +/// fields, so trace events stay small and don't duplicate the existing array metadata. +#[derive(Clone, Debug)] +pub(crate) struct ArraySummary(ArrayRef); impl ArraySummary { pub(crate) fn new(array: &ArrayRef) -> Self { - Self { - encoding: array.encoding_id(), - len: array.len(), - dtype: array.dtype().clone(), - } + Self(array.clone()) } } impl Display for ArraySummary { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "{} len={} dtype={}", self.encoding, self.len, self.dtype) + write!( + f, + "{} len={} dtype={}", + self.0.encoding_id(), + self.0.len(), + self.0.dtype(), + ) } } From d8e10a7173ed851ef57d6312d11d3c6bc8cf61c7 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Mon, 11 May 2026 11:38:25 +0100 Subject: [PATCH 07/15] More consistet format Signed-off-by: Adam Gutglick --- vortex-array/src/test_harness/trace.rs | 121 ++++++++++++------------- 1 file changed, 58 insertions(+), 63 deletions(-) diff --git a/vortex-array/src/test_harness/trace.rs b/vortex-array/src/test_harness/trace.rs index 1d399538f3d..586c77e374b 100644 --- a/vortex-array/src/test_harness/trace.rs +++ b/vortex-array/src/test_harness/trace.rs @@ -242,11 +242,12 @@ impl From for TraceInterest { } } -/// Snapshot-friendly wrapper around [`ArrayRef`] that renders just the encoding, length, and -/// dtype using the trace format (`vortex.primitive len=4 dtype=i32`). +/// 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 cloning individual -/// fields, so trace events stay small and don't duplicate the existing array metadata. +/// 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); @@ -258,13 +259,7 @@ impl ArraySummary { impl Display for ArraySummary { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!( - f, - "{} len={} dtype={}", - self.0.encoding_id(), - self.0.len(), - self.0.dtype(), - ) + Display::fmt(&self.0, f) } } @@ -1145,9 +1140,9 @@ mod tests { assert!(traced.output.is::()); assert_arrays_eq!(traced.output, values, &mut ctx); insta::assert_snapshot!(traced.trace.to_string(), @r" -optimize root=vortex.filter len=4 dtype=i32 session=false - reduce TrivialFilterRule: vortex.filter len=4 dtype=i32 -> vortex.primitive len=4 dtype=i32 - done output=vortex.primitive len=4 dtype=i32 +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(()) @@ -1176,9 +1171,9 @@ optimize root=vortex.filter len=4 dtype=i32 session=false assert!(optimized_filter.child().is::()); assert_arrays_eq!(traced.output, PrimitiveArray::from_iter([2i32, 3]), &mut ctx); insta::assert_snapshot!(traced.trace.to_string(), @" - optimize root=vortex.filter len=2 dtype=i32 session=false - reduce_parent static:FilterReduceAdaptor(Filter) slot=0 parent=vortex.filter len=2 dtype=i32 child=vortex.filter len=4 dtype=i32 -> vortex.filter len=2 dtype=i32 - done output=vortex.filter len=2 dtype=i32 + 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 traced = trace_array_with( @@ -1189,17 +1184,17 @@ optimize root=vortex.filter len=4 dtype=i32 session=false )?; insta::assert_snapshot!(traced.trace.to_string(), @r" -execute_until target=AnyCanonical root=vortex.filter len=2 dtype=i32 - iter 1 current=vortex.filter len=2 dtype=i32 builder_active=false - ExecuteSlot slot=0 parent=vortex.filter len=2 dtype=i32 child=vortex.filter len=4 dtype=i32 - iter 2 current=vortex.filter len=4 dtype=i32 stack_parent=vortex.filter len=2 dtype=i32 slot=0 builder_active=false - Done array=vortex.primitive len=4 dtype=i32 - iter 3 current=vortex.primitive len=4 dtype=i32 stack_parent=vortex.filter len=2 dtype=i32 slot=0 builder_active=false - pop_frame slot=0 parent=vortex.filter len=2 dtype=i32 child=vortex.primitive len=4 dtype=i32 output=vortex.filter len=2 dtype=i32 - iter 4 current=vortex.filter len=2 dtype=i32 builder_active=false - Done array=vortex.primitive len=2 dtype=i32 - iter 5 current=vortex.primitive len=2 dtype=i32 builder_active=false - return output=vortex.primitive len=2 dtype=i32 +execute_until target=AnyCanonical root=vortex.filter(i32, len=2) + iter 1 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 2 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 3 current=vortex.primitive(i32, len=4) stack_parent=vortex.filter(i32, len=2) slot=0 builder_active=false + pop_frame slot=0 parent=vortex.filter(i32, len=2) child=vortex.primitive(i32, len=4) output=vortex.filter(i32, len=2) + iter 4 current=vortex.filter(i32, len=2) builder_active=false + Done array=vortex.primitive(i32, len=2) + iter 5 current=vortex.primitive(i32, len=2) builder_active=false + return output=vortex.primitive(i32, len=2) "); Ok(()) @@ -1219,28 +1214,28 @@ execute_until target=AnyCanonical root=vortex.filter len=2 dtype=i32 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 len=3 dtype=i32 - iter 1 current=vortex.test.stack-parent len=3 dtype=i32 builder_active=false + execute_until target=AnyCanonical root=vortex.test.stack-parent(i32, len=3) + iter 1 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 len=3 dtype=i32 child=vortex.test.stack-child len=3 dtype=i32 source=session[0] kernel=execute_parent_fn outcome=declined - child_execute_parent attempt slot=0 parent=vortex.test.stack-parent len=3 dtype=i32 child=vortex.test.stack-child len=3 dtype=i32 source=session[1] kernel=execute_parent_fn outcome=declined - child_execute_parent none current=vortex.test.stack-parent len=3 dtype=i32 - execute encoding=vortex.test.stack-parent len=3 dtype=i32 - request ExecuteSlot slot=0 target=Primitive parent=vortex.test.stack-parent len=3 dtype=i32 - ExecuteSlot slot=0 parent=vortex.test.stack-parent len=3 dtype=i32 child=vortex.test.stack-child len=3 dtype=i32 - iter 2 current=vortex.test.stack-child len=3 dtype=i32 stack_parent=vortex.test.stack-parent len=3 dtype=i32 slot=0 builder_active=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) + request ExecuteSlot slot=0 target=Primitive parent=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 2 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 len=3 dtype=i32 child=vortex.test.stack-child len=3 dtype=i32 source=session[0] kernel=execute_parent_fn outcome=declined - stack_execute_parent applied slot=0 parent=vortex.test.stack-parent len=3 dtype=i32 child=vortex.test.stack-child len=3 dtype=i32 source=session[1] kernel=execute_parent_fn output=vortex.primitive len=3 dtype=i32 - optimize root=vortex.primitive len=3 dtype=i32 session=true - loop input=vortex.primitive len=3 dtype=i32 - reduce none array=vortex.primitive len=3 dtype=i32 - reduce_parent none array=vortex.primitive len=3 dtype=i32 - done output=vortex.primitive len=3 dtype=i32 changed=false - optimize_ctx input=vortex.primitive len=3 dtype=i32 output=vortex.primitive len=3 dtype=i32 changed=false - iter 3 current=vortex.primitive len=3 dtype=i32 builder_active=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 3 current=vortex.primitive(i32, len=3) builder_active=false done_check target=true canonical=true - return output=vortex.primitive len=3 dtype=i32 + return output=vortex.primitive(i32, len=3) "); Ok(()) @@ -1265,22 +1260,22 @@ execute_until target=AnyCanonical root=vortex.filter len=2 dtype=i32 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 len=5 dtype=i32 - iter 1 current=vortex.chunked len=5 dtype=i32 builder_active=false - builder start array=vortex.chunked len=5 dtype=i32 - AppendChild slot=1 parent=vortex.chunked len=5 dtype=i32 child=vortex.primitive len=2 dtype=i32 - builder append child=vortex.primitive len=2 dtype=i32 - iter 2 current=vortex.chunked len=5 dtype=i32 builder_active=true - AppendChild slot=2 parent=vortex.chunked len=5 dtype=i32 child=vortex.primitive len=1 dtype=i32 - builder append child=vortex.primitive len=1 dtype=i32 - iter 3 current=vortex.chunked len=5 dtype=i32 builder_active=true - AppendChild slot=3 parent=vortex.chunked len=5 dtype=i32 child=vortex.primitive len=2 dtype=i32 - builder append child=vortex.primitive len=2 dtype=i32 - iter 4 current=vortex.chunked len=5 dtype=i32 builder_active=true - Done array=vortex.primitive len=0 dtype=i32 - builder finish output=vortex.primitive len=5 dtype=i32 - iter 5 current=vortex.primitive len=5 dtype=i32 builder_active=false - return output=vortex.primitive len=5 dtype=i32 + execute_until target=AnyCanonical root=vortex.chunked(i32, len=5) + iter 1 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 2 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 3 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 4 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 5 current=vortex.primitive(i32, len=5) builder_active=false + return output=vortex.primitive(i32, len=5) "); Ok(()) From 059d9e61007343242ecb8bbd35da9ebf32cb94e3 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Mon, 11 May 2026 12:09:30 +0100 Subject: [PATCH 08/15] more thing Signed-off-by: Adam Gutglick --- vortex-array/src/executor.rs | 78 +++++++-------- vortex-array/src/lib.rs | 2 +- vortex-array/src/optimizer/mod.rs | 26 ++--- vortex-array/src/optimizer/rules.rs | 10 +- vortex-array/src/test_harness/trace.rs | 125 +++++++++++++++++++++---- vortex-array/src/trace_macros.rs | 53 +++++++++-- 6 files changed, 213 insertions(+), 81 deletions(-) diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 3e11810a9c4..a4ca5db8f8b 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -163,7 +163,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. - #[expect(clippy::cognitive_complexity)] + #[allow(clippy::cognitive_complexity)] pub fn execute_until(self, ctx: &mut ExecutionCtx) -> VortexResult { let mut current_array = self; let mut current_builder: Option> = None; @@ -172,11 +172,11 @@ impl ArrayRef { let kernels = execute_parent_kernels.as_ref(); let max_iterations = max_iterations(); - crate::trace_array!(record_execute_until_start::(¤t_array)); + crate::trace_op!(record_execute_until_start::(¤t_array)); for iteration in 1..=max_iterations { - crate::trace_array!(use(iteration)); - crate::trace_array!(record_execute_until_iteration( + crate::trace_op!(use(iteration)); + crate::trace_op!(record_execute_until_iteration( iteration, ¤t_array, stack @@ -191,7 +191,7 @@ impl ArrayRef { let done_target = is_done(¤t_array); let done_canonical = AnyCanonical::matches(¤t_array); - crate::trace_array!(record_execute_until_done_check(done_target, done_canonical)); + crate::trace_op!(record_execute_until_done_check(done_target, done_canonical)); if done_target || done_canonical { match stack.pop() { @@ -200,11 +200,11 @@ impl ArrayRef { current_builder.is_none(), "root activation should not retain a builder" ); - crate::trace_array!(record_execute_until_return(¤t_array)); + crate::trace_op!(record_execute_until_return(¤t_array)); return Ok(current_array); } Some(frame) => { - let trace_pop_frame = crate::trace_array!(value( + let trace_pop_frame = crate::trace_op!(value( Some(( frame.parent_array.clone(), current_array.clone(), @@ -214,8 +214,8 @@ impl ArrayRef { )); (current_array, current_builder) = pop_frame(frame, current_array)?; if let Some((parent_before, child_before, slot_idx)) = trace_pop_frame { - crate::trace_array!(use(parent_before, child_before, slot_idx,)); - crate::trace_array!(record_execute_until_pop_frame( + crate::trace_op!(use(parent_before, child_before, slot_idx,)); + crate::trace_op!(record_execute_until_pop_frame( &parent_before, slot_idx, &child_before, @@ -252,13 +252,13 @@ impl ArrayRef { { let frame = stack.pop().vortex_expect("just peeked"); let optimized = result.optimize_ctx(ctx.session())?; - crate::trace_array!(record_execute_optimized(&result, &optimized)); + crate::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() { - crate::trace_array!(record_execute_parent_none( + crate::trace_op!(record_execute_parent_none( "stack_execute_parent", ¤t_array, )); @@ -269,12 +269,12 @@ impl ArrayRef { && let Some(rewritten) = try_execute_parent(¤t_array, kernels, ctx)? { let optimized = rewritten.optimize_ctx(ctx.session())?; - crate::trace_array!(record_execute_optimized(&rewritten, &optimized)); + crate::trace_op!(record_execute_optimized(&rewritten, &optimized)); current_array = optimized; continue; } if current_builder.is_none() { - crate::trace_array!(record_execute_parent_none( + crate::trace_op!(record_execute_parent_none( "child_execute_parent", ¤t_array, )); @@ -284,14 +284,14 @@ impl ArrayRef { let expected_dtype = current_array.dtype().clone(); let stats = current_array.statistics().to_array_stats(); let encoding_id = current_array.encoding_id(); - crate::trace_array!(record_execute_encoding(¤t_array)); + crate::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) }?; - crate::trace_array!(record_execute_slot(i, &parent, &child)); + crate::trace_op!(record_execute_slot(i, &parent, &child)); stack.push(StackFrame { parent_array: parent, parent_builder: current_builder.take(), @@ -305,7 +305,7 @@ impl ArrayRef { } ExecutionStep::AppendChild(i) => { if current_builder.is_none() { - crate::trace_array!(record_builder_start(&array)); + crate::trace_op!(record_builder_start(&array)); current_builder = Some(builder_with_capacity_in( ctx.allocator(), array.dtype(), @@ -314,8 +314,8 @@ impl ArrayRef { } let (parent, child) = unsafe { array.take_slot_unchecked(i) }?; - crate::trace_array!(record_append_child(i, &parent, &child)); - crate::trace_array!(record_builder_append(&child)); + crate::trace_op!(record_append_child(i, &parent, &child)); + crate::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. @@ -329,7 +329,7 @@ impl ArrayRef { } ExecutionStep::Done => { let had_builder = current_builder.is_some(); - crate::trace_array!(record_execute_done(&array)); + crate::trace_op!(record_execute_done(&array)); (current_array, current_builder) = finalize_done( array, current_builder, @@ -339,7 +339,7 @@ impl ArrayRef { encoding_id, )?; if had_builder { - crate::trace_array!(record_builder_finish(¤t_array)); + crate::trace_op!(record_builder_finish(¤t_array)); } } } @@ -465,27 +465,27 @@ impl Drop for ExecutionCtx { /// `AppendChild` is returned. impl Executable for ArrayRef { fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - crate::trace_array!(record_single_step_start(&array)); + crate::trace_op!(record_single_step_start(&array)); if let Some(canonical) = array.as_opt::() { let output = Canonical::from(canonical).into_array(); - crate::trace_array!(record_single_step_applied("canonical", &array, &output)); + crate::trace_op!(record_single_step_applied("canonical", &array, &output)); return Ok(output); } - crate::trace_array!(record_single_step_phase_none("canonical", &array)); + crate::trace_op!(record_single_step_phase_none("canonical", &array)); if let Some(reduced) = array.reduce()? { reduced.statistics().inherit_from(array.statistics()); - crate::trace_array!(record_single_step_applied("reduce", &array, &reduced)); + crate::trace_op!(record_single_step_applied("reduce", &array, &reduced)); return Ok(reduced); } - crate::trace_array!(record_single_step_phase_none("reduce", &array)); + crate::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)? { reduced_parent.statistics().inherit_from(array.statistics()); - crate::trace_array!(record_single_step_applied( + crate::trace_op!(record_single_step_applied( "reduce_parent", &array, &reduced_parent, @@ -493,7 +493,7 @@ impl Executable for ArrayRef { return Ok(reduced_parent); } } - crate::trace_array!(record_single_step_phase_none("reduce_parent", &array)); + crate::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(); @@ -518,7 +518,7 @@ impl Executable for ArrayRef { executed_parent .statistics() .inherit_from(array.statistics()); - crate::trace_array!(record_single_step_applied( + crate::trace_op!(record_single_step_applied( "execute_parent", &array, &executed_parent, @@ -526,14 +526,14 @@ impl Executable for ArrayRef { return Ok(executed_parent); } } - crate::trace_array!(record_single_step_phase_none("execute_parent", &array)); - crate::trace_array!(record_execute_encoding(&array)); + crate::trace_op!(record_single_step_phase_none("execute_parent", &array)); + crate::trace_op!(record_execute_encoding(&array)); let result = array.execute_encoding(ctx)?; let (array, step) = result.into_parts(); match step { ExecutionStep::Done => { - crate::trace_array!(record_execute_done(&array)); + crate::trace_op!(record_execute_done(&array)); Ok(array) } ExecutionStep::ExecuteSlot(i, _) => { @@ -545,11 +545,11 @@ impl Executable for ArrayRef { } ExecutionStep::AppendChild(_) => { // Single-step: build the entire parent via the builder path. - crate::trace_array!(record_builder_start(&array)); + crate::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)?; let output = builder.finish(); - crate::trace_array!(record_builder_finish(&output)); + crate::trace_op!(record_builder_finish(&output)); Ok(output) } } @@ -633,11 +633,11 @@ fn execute_parent_for_child( kernels: &ParentExecutionKernels, ctx: &mut ExecutionCtx, ) -> VortexResult> { - crate::trace_array!(use(phase)); + crate::trace_op!(use(phase)); let key = execute_parent_key(parent.encoding_id(), child.encoding_id()); if let Some(plugins) = kernels.get(&key) { for (plugin_idx, plugin) in plugins.as_ref().iter().enumerate() { - crate::trace_array!(use(plugin_idx)); + crate::trace_op!(use(plugin_idx)); if let Some(result) = plugin.execute_parent(child, parent, slot_idx, ctx)? { if cfg!(debug_assertions) { vortex_ensure!( @@ -649,7 +649,7 @@ fn execute_parent_for_child( "Executed parent canonical dtype mismatch" ); } - crate::trace_array!(record_execute_parent_applied( + crate::trace_op!(record_execute_parent_applied( phase, parent, child, @@ -660,7 +660,7 @@ fn execute_parent_for_child( )); return Ok(Some(result)); } - crate::trace_array!(record_execute_parent_attempt( + crate::trace_op!(record_execute_parent_attempt( phase, parent, child, @@ -800,7 +800,7 @@ 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(); - crate::trace_array!(record_execute_step_request::(&array, slot_idx)); + crate::trace_op!(record_execute_step_request::(&array, slot_idx)); Self { array, step: ExecutionStep::ExecuteSlot(slot_idx, M::matches), @@ -812,7 +812,7 @@ impl ExecutionResult { /// `current_array`. pub fn append_child(array: impl IntoArray, slot_idx: usize) -> Self { let array = array.into_array(); - crate::trace_array!(record_append_child_request(&array, slot_idx)); + crate::trace_op!(record_append_child_request(&array, slot_idx)); Self { array, step: ExecutionStep::AppendChild(slot_idx), diff --git a/vortex-array/src/lib.rs b/vortex-array/src/lib.rs index 56701b1fe2d..9439dc7dd1b 100644 --- a/vortex-array/src/lib.rs +++ b/vortex-array/src/lib.rs @@ -106,7 +106,7 @@ pub mod aggregate_fn; #[doc(hidden)] pub mod aliases; mod trace_macros; -pub(crate) use trace_macros::trace_array; +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 202c096fe20..0d0171fbf5c 100644 --- a/vortex-array/src/optimizer/mod.rs +++ b/vortex-array/src/optimizer/mod.rs @@ -73,20 +73,20 @@ fn try_optimize( let mut any_optimizations = false; let array_ref = session.map(|s| s.kernels()); - crate::trace_array!(record_optimize_start(array, session.is_some())); + crate::trace_op!(record_optimize_start(array, session.is_some())); // Apply reduction rules to the current array until no more rules apply. for _ in 0..=100 { - crate::trace_array!(record_optimize_loop_start(¤t_array)); + crate::trace_op!(record_optimize_loop_start(¤t_array)); if let Some(new_array) = current_array.reduce()? { current_array = new_array; any_optimizations = true; - crate::trace_array!(record_optimize_loop_end()); + crate::trace_op!(record_optimize_loop_end()); continue; } - crate::trace_array!(record_optimize_reduce_none(¤t_array)); + crate::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. @@ -100,9 +100,9 @@ fn try_optimize( array_ref.find_reduce_parent(current_array.encoding_id(), child.encoding_id()) { for (plugin_idx, plugin) in plugins.as_ref().iter().enumerate() { - crate::trace_array!(use(plugin_idx)); + crate::trace_op!(use(plugin_idx)); if let Some(new_array) = plugin(child, ¤t_array, slot_idx)? { - crate::trace_array!(record_parent_reduce_applied( + crate::trace_op!(record_parent_reduce_applied( ¤t_array, child, slot_idx, @@ -113,7 +113,7 @@ fn try_optimize( parent_reduced = Some(new_array); break; } - crate::trace_array!(record_parent_reduce_attempt( + crate::trace_op!(record_parent_reduce_attempt( ¤t_array, child, slot_idx, @@ -137,15 +137,15 @@ fn try_optimize( // If the parent was replaced, then we attempt to reduce it again. current_array = new_array; any_optimizations = true; - crate::trace_array!(record_optimize_loop_end()); + crate::trace_op!(record_optimize_loop_end()); continue; } - crate::trace_array!(record_optimize_parent_reduce_none(¤t_array)); - crate::trace_array!(record_optimize_loop_end()); + crate::trace_op!(record_optimize_parent_reduce_none(¤t_array)); + crate::trace_op!(record_optimize_loop_end()); // No more optimizations can be applied - crate::trace_array!(record_optimize_done(¤t_array, any_optimizations)); + crate::trace_op!(record_optimize_done(¤t_array, any_optimizations)); if any_optimizations { return Ok(Some(current_array)); @@ -164,7 +164,7 @@ fn try_optimize_recursive( let mut current_array = array.clone(); let mut any_optimizations = false; - crate::trace_array!(record_optimize_recursive_start(array)); + crate::trace_op!(record_optimize_recursive_start(array)); if let Some(new_array) = try_optimize(¤t_array, Some(session))? { current_array = new_array; @@ -177,7 +177,7 @@ fn try_optimize_recursive( match slot { Some(child) => { if let Some(new_child) = try_optimize_recursive(child, session)? { - crate::trace_array!(record_optimize_recursive_slot( + crate::trace_op!(record_optimize_recursive_slot( new_slots.len(), child, &new_child, diff --git a/vortex-array/src/optimizer/rules.rs b/vortex-array/src/optimizer/rules.rs index 5ccd48e10c3..6f77f38f6cf 100644 --- a/vortex-array/src/optimizer/rules.rs +++ b/vortex-array/src/optimizer/rules.rs @@ -133,10 +133,10 @@ impl ReduceRuleSet { pub fn evaluate(&self, array: ArrayView<'_, V>) -> VortexResult> { for rule in self.rules.iter() { if let Some(reduced) = rule.reduce(array)? { - crate::trace_array!(record_reduce_applied(array.array(), *rule, &reduced)); + crate::trace_op!(record_reduce_applied(array.array(), *rule, &reduced)); return Ok(Some(reduced)); } - crate::trace_array!(record_reduce_attempt( + crate::trace_op!(record_reduce_attempt( array.array(), *rule, crate::test_harness::trace::AttemptOutcome::Declined, @@ -182,7 +182,7 @@ impl ParentRuleSet { ) -> VortexResult> { for rule in self.rules.iter() { if !rule.matches(parent) { - crate::trace_array!(record_parent_reduce_attempt( + crate::trace_op!(record_parent_reduce_attempt( parent, child.array(), child_idx, @@ -212,7 +212,7 @@ impl ParentRuleSet { ); } - crate::trace_array!(record_parent_reduce_applied( + crate::trace_op!(record_parent_reduce_applied( parent, child.array(), child_idx, @@ -222,7 +222,7 @@ impl ParentRuleSet { )); return Ok(Some(reduced)); } - crate::trace_array!(record_parent_reduce_attempt( + crate::trace_op!(record_parent_reduce_attempt( parent, child.array(), child_idx, diff --git a/vortex-array/src/test_harness/trace.rs b/vortex-array/src/test_harness/trace.rs index 586c77e374b..ab17016ecd0 100644 --- a/vortex-array/src/test_harness/trace.rs +++ b/vortex-array/src/test_harness/trace.rs @@ -1,6 +1,68 @@ // 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!`][crate::trace_op] macro inside the optimizer +//! ([`optimizer`][crate::optimizer]) and executor ([`executor`][crate::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`][crate::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; @@ -25,7 +87,7 @@ pub enum TraceResolution { Attempts, } -/// Options for [`trace_array_with`]. +/// 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. @@ -123,20 +185,49 @@ fn write_indent(f: &mut fmt::Formatter<'_>, depth: usize) -> fmt::Result { Ok(()) } -/// Run `f` while capturing default trace output. +/// 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 +/// [`VortexSessionExecute::execute`][crate::VortexSessionExecute::execute]. While `f` runs, +/// the optimizer and executor emit structured events via the [`trace_op!`][crate::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 /// -/// The default resolution records the rule rewrites, parent kernels, execution steps, and builder -/// activity that actually executed. Use [`trace_array_with`] and [`TraceResolution::Attempts`] -/// when a test needs to assert on every declined rule or kernel attempt. -pub fn trace_array(f: impl FnOnce() -> VortexResult) -> VortexResult> { - trace_array_with(TraceOptions::default(), f) +/// ```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 trace output using `options`. +/// 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_array_with( +pub fn trace_op_with( options: TraceOptions, f: impl FnOnce() -> VortexResult, ) -> VortexResult> { @@ -144,7 +235,7 @@ pub fn trace_array_with( ACTIVE_TRACE.with(|active| { let mut active = active.borrow_mut(); if active.is_some() { - return Err(vortex_err!("trace_array captures cannot be nested")); + return Err(vortex_err!("trace_op captures cannot be nested")); } *active = Some(TraceRecorder::new(options)); Ok(()) @@ -1123,8 +1214,8 @@ mod tests { use crate::session::ArraySession; use crate::test_harness::trace::TraceOptions; use crate::test_harness::trace::TraceResolution; - use crate::test_harness::trace::trace_array; - use crate::test_harness::trace::trace_array_with; + use crate::test_harness::trace::trace_op; + use crate::test_harness::trace::trace_op_with; use crate::test_harness::trace_arrays::stack_parent_fixture; use crate::test_harness::trace_arrays::stack_parent_session; @@ -1135,7 +1226,7 @@ mod tests { let filter = FilterArray::try_new(values.clone(), Mask::new_true(values.len()))?.into_array(); - let traced = trace_array(|| filter.optimize())?; + let traced = trace_op(|| filter.optimize())?; assert!(traced.output.is::()); assert_arrays_eq!(traced.output, values, &mut ctx); @@ -1160,7 +1251,7 @@ optimize root=vortex.filter(i32, len=4) session=false let outer = FilterArray::try_new(inner, Mask::from_iter([false, true, true, false]))?.into_array(); - let traced = trace_array_with( + let traced = trace_op_with( TraceOptions { resolution: TraceResolution::ExecutedOnly, }, @@ -1176,7 +1267,7 @@ optimize root=vortex.filter(i32, len=4) session=false done output=vortex.filter(i32, len=2) "); - let traced = trace_array_with( + let traced = trace_op_with( TraceOptions { resolution: TraceResolution::ExecutedOnly, }, @@ -1205,7 +1296,7 @@ execute_until target=AnyCanonical root=vortex.filter(i32, len=2) let mut ctx = ExecutionCtx::new(stack_parent_session()); let parent = stack_parent_fixture()?; - let traced = trace_array_with( + let traced = trace_op_with( TraceOptions { resolution: TraceResolution::Attempts, }, @@ -1252,7 +1343,7 @@ execute_until target=AnyCanonical root=vortex.filter(i32, len=2) let chunked = ChunkedArray::try_new(chunks, dtype)?.into_array(); let mut ctx = ExecutionCtx::new(VortexSession::empty().with::()); - let traced = trace_array(|| { + let traced = trace_op(|| { chunked .execute::(&mut ctx) .map(IntoArray::into_array) diff --git a/vortex-array/src/trace_macros.rs b/vortex-array/src/trace_macros.rs index e7827ed5d03..7b8503ae086 100644 --- a/vortex-array/src/trace_macros.rs +++ b/vortex-array/src/trace_macros.rs @@ -1,7 +1,48 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -macro_rules! trace_array { +//! 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]. +/// +/// # Forms +/// +/// ```ignore +/// // 1. Event dispatch — the common form. Calls `test_harness::trace::(...)` only when +/// // a recorder is installed on this thread. Otherwise compiles to nothing. +/// crate::trace_op!(record_execute_until_start::(&array)); +/// +/// // 2. `use(...)` — keep otherwise-unused bindings alive in non-test builds so the +/// // surrounding code still type-checks when the trace call is compiled out. +/// crate::trace_op!(use(plugin_idx, slot_idx)); +/// +/// // 3. `value(, )` — pick between two expressions depending on +/// // whether a recorder is active. Useful when the active branch needs to clone state +/// // just for the trace. +/// let snapshot = crate::trace_op!(value(Some(parent.clone()), None)); +/// ``` +/// +/// # 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 `()`. Captured names are bound +/// through the `use(...)` form so the compiler does not warn about unused variables. +/// +/// 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 { (@when_enabled { $($enabled:tt)* } else { $($disabled:tt)* }) => {{ #[cfg(all(any(test, feature = "_test-harness"), not(codspeed)))] { @@ -15,7 +56,7 @@ macro_rules! trace_array { }}; (@if_active { $($enabled:tt)* } else { $($disabled:tt)* }) => { - $crate::trace_array!(@when_enabled { + $crate::trace_op!(@when_enabled { if $crate::test_harness::trace::is_active() { $($enabled)* } else { @@ -27,20 +68,20 @@ macro_rules! trace_array { }; (use($($value:expr),* $(,)?)) => { - $crate::trace_array!(@when_enabled {} else { + $crate::trace_op!(@when_enabled {} else { let _ = ($(&$value),*); }) }; (value($enabled:expr, $disabled:expr)) => { - $crate::trace_array!(@if_active { $enabled } else { $disabled }) + $crate::trace_op!(@if_active { $enabled } else { $disabled }) }; ($($event:tt)*) => { - $crate::trace_array!(@if_active { + $crate::trace_op!(@if_active { $crate::test_harness::trace::$($event)* } else {}) }; } -pub(crate) use trace_array; +pub(crate) use trace_op; From cc5b375990e837b95fdd01bc7a296351048f89bc Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Mon, 11 May 2026 12:22:13 +0100 Subject: [PATCH 09/15] fix rustdoc Signed-off-by: Adam Gutglick --- vortex-array/src/test_harness/trace.rs | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/vortex-array/src/test_harness/trace.rs b/vortex-array/src/test_harness/trace.rs index ab17016ecd0..890956ba508 100644 --- a/vortex-array/src/test_harness/trace.rs +++ b/vortex-array/src/test_harness/trace.rs @@ -6,8 +6,7 @@ //! # What this records //! //! [`trace_op`] runs a closure with a thread-local recorder installed. While the recorder is -//! active, calls to the [`trace_op!`][crate::trace_op] macro inside the optimizer -//! ([`optimizer`][crate::optimizer]) and executor ([`executor`][crate::executor]) push +//! 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. //! @@ -47,7 +46,7 @@ //! - 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`][crate::trace_op] for the gating rules. +//! `trace_op!` for the gating rules. //! //! # Example //! @@ -189,8 +188,8 @@ fn write_indent(f: &mut fmt::Formatter<'_>, depth: usize) -> fmt::Result { /// /// `f` typically invokes an operation that drives the executor or optimizer, such as /// [`ArrayOptimizer::optimize`][crate::optimizer::ArrayOptimizer::optimize] or -/// [`VortexSessionExecute::execute`][crate::VortexSessionExecute::execute]. While `f` runs, -/// the optimizer and executor emit structured events via the [`trace_op!`][crate::trace_op] +/// [`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`]. /// From 7033c6d21ce4875d97389235cffbb9b3a3f050cb Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Mon, 8 Jun 2026 11:21:12 +0100 Subject: [PATCH 10/15] some updates Signed-off-by: Adam Gutglick --- vortex-array/src/test_harness/trace.rs | 2 +- vortex-array/src/test_harness/trace_arrays.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/vortex-array/src/test_harness/trace.rs b/vortex-array/src/test_harness/trace.rs index 890956ba508..5f0c125e2a8 100644 --- a/vortex-array/src/test_harness/trace.rs +++ b/vortex-array/src/test_harness/trace.rs @@ -1260,7 +1260,7 @@ optimize root=vortex.filter(i32, len=4) session=false let optimized_filter = traced.output.as_::(); assert!(optimized_filter.child().is::()); assert_arrays_eq!(traced.output, PrimitiveArray::from_iter([2i32, 3]), &mut ctx); - insta::assert_snapshot!(traced.trace.to_string(), @" + 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) diff --git a/vortex-array/src/test_harness/trace_arrays.rs b/vortex-array/src/test_harness/trace_arrays.rs index 7a8d0fc2ad3..e95df7f6aa4 100644 --- a/vortex-array/src/test_harness/trace_arrays.rs +++ b/vortex-array/src/test_harness/trace_arrays.rs @@ -16,9 +16,9 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::ArrayEq; -use crate::EqMode; use crate::ArrayHash; use crate::ArrayRef; +use crate::EqMode; use crate::ExecutionCtx; use crate::ExecutionResult; use crate::IntoArray; From 7dc3012e798e22fb4a9b48394dd97b8e88686767 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Mon, 8 Jun 2026 11:26:32 +0100 Subject: [PATCH 11/15] less files Signed-off-by: Adam Gutglick --- vortex-array/src/executor.rs | 2 +- .../{test_harness.rs => test_harness/mod.rs} | 1 - vortex-array/src/test_harness/trace.rs | 387 ++++++++++++++++-- vortex-array/src/test_harness/trace_arrays.rs | 350 ---------------- 4 files changed, 362 insertions(+), 378 deletions(-) rename vortex-array/src/{test_harness.rs => test_harness/mod.rs} (98%) delete mode 100644 vortex-array/src/test_harness/trace_arrays.rs diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index a4ca5db8f8b..ae80153300b 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -174,7 +174,7 @@ impl ArrayRef { crate::trace_op!(record_execute_until_start::(¤t_array)); - for iteration in 1..=max_iterations { + for iteration in 0..max_iterations { crate::trace_op!(use(iteration)); crate::trace_op!(record_execute_until_iteration( iteration, diff --git a/vortex-array/src/test_harness.rs b/vortex-array/src/test_harness/mod.rs similarity index 98% rename from vortex-array/src/test_harness.rs rename to vortex-array/src/test_harness/mod.rs index 74c7d380074..d3a6b829e80 100644 --- a/vortex-array/src/test_harness.rs +++ b/vortex-array/src/test_harness/mod.rs @@ -14,7 +14,6 @@ use crate::arrays::bool::BoolArrayExt; #[cfg(not(codspeed))] pub mod trace; -pub mod trace_arrays; /// Check that a named metadata matches its previous versioning. /// diff --git a/vortex-array/src/test_harness/trace.rs b/vortex-array/src/test_harness/trace.rs index 5f0c125e2a8..eeea5d11e5e 100644 --- a/vortex-array/src/test_harness/trace.rs +++ b/vortex-array/src/test_harness/trace.rs @@ -1195,13 +1195,37 @@ impl TraceEvent { #[cfg(test)] mod tests { + 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::ChunkedArray; use crate::arrays::Filter; use crate::arrays::FilterArray; @@ -1209,14 +1233,322 @@ mod tests { use crate::arrays::PrimitiveArray; use crate::arrays::filter::FilterArraySlotsExt; 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::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::test_harness::trace_arrays::stack_parent_fixture; - use crate::test_harness::trace_arrays::stack_parent_session; + 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<()> { @@ -1274,26 +1606,29 @@ optimize root=vortex.filter(i32, len=4) session=false )?; insta::assert_snapshot!(traced.trace.to_string(), @r" -execute_until target=AnyCanonical root=vortex.filter(i32, len=2) - iter 1 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 2 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 3 current=vortex.primitive(i32, len=4) stack_parent=vortex.filter(i32, len=2) slot=0 builder_active=false - pop_frame slot=0 parent=vortex.filter(i32, len=2) child=vortex.primitive(i32, len=4) output=vortex.filter(i32, len=2) - iter 4 current=vortex.filter(i32, len=2) builder_active=false - Done array=vortex.primitive(i32, len=2) - iter 5 current=vortex.primitive(i32, len=2) builder_active=false - return output=vortex.primitive(i32, len=2) -"); + 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 parent=vortex.filter(i32, len=2) child=vortex.primitive(i32, len=4) 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(()) } - #[test] - fn trace_execution_stack_parent_kernel_attempts() -> VortexResult<()> { - let mut ctx = ExecutionCtx::new(stack_parent_session()); - let parent = stack_parent_fixture()?; + #[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 { @@ -1305,7 +1640,7 @@ execute_until target=AnyCanonical root=vortex.filter(i32, len=2) 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 1 current=vortex.test.stack-parent(i32, len=3) builder_active=false + 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 @@ -1313,7 +1648,7 @@ execute_until target=AnyCanonical root=vortex.filter(i32, len=2) execute encoding=vortex.test.stack-parent(i32, len=3) request ExecuteSlot slot=0 target=Primitive parent=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 2 current=vortex.test.stack-child(i32, len=3) stack_parent=vortex.test.stack-parent(i32, len=3) slot=0 builder_active=false + 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) @@ -1323,7 +1658,7 @@ execute_until target=AnyCanonical root=vortex.filter(i32, len=2) 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 3 current=vortex.primitive(i32, len=3) builder_active=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) "); @@ -1351,20 +1686,20 @@ execute_until target=AnyCanonical root=vortex.filter(i32, len=2) 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 1 current=vortex.chunked(i32, len=5) builder_active=false + 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 2 current=vortex.chunked(i32, len=5) builder_active=true + 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 3 current=vortex.chunked(i32, len=5) builder_active=true + 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 4 current=vortex.chunked(i32, len=5) builder_active=true + 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 5 current=vortex.primitive(i32, len=5) builder_active=false + iter 4 current=vortex.primitive(i32, len=5) builder_active=false return output=vortex.primitive(i32, len=5) "); diff --git a/vortex-array/src/test_harness/trace_arrays.rs b/vortex-array/src/test_harness/trace_arrays.rs deleted file mode 100644 index e95df7f6aa4..00000000000 --- a/vortex-array/src/test_harness/trace_arrays.rs +++ /dev/null @@ -1,350 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Small encodings used by trace snapshot tests. - -use std::fmt::Display; -use std::fmt::Formatter; -use std::hash::Hasher; - -use smallvec::smallvec; -use vortex_error::VortexResult; -use vortex_error::vortex_bail; -use vortex_error::vortex_ensure; -use vortex_error::vortex_panic; -use vortex_session::VortexSession; -use vortex_session::registry::CachedId; - -use crate::ArrayEq; -use crate::ArrayHash; -use crate::ArrayRef; -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::Primitive; -use crate::arrays::PrimitiveArray; -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::kernels::ArrayKernelsExt; -use crate::session::ArraySession; -use crate::serde::ArrayChildren; -use crate::validity::Validity; - -/// Create a `StackParent(StackChild)` fixture. -/// -/// `StackParent` requests `ExecuteSlot` until its child is `Primitive`. `StackChild` has one -/// declined parent kernel followed by one successful parent kernel, so strict trace snapshots can -/// assert that stack parent kernels run before the child decodes itself. -pub 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. -pub 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 -} - -/// Create the child encoding used by [`stack_parent_fixture`]. -pub fn stack_child() -> VortexResult { - Ok( - Array::try_from_parts(ArrayParts::new(StackChild, test_dtype(), 3, StackChildData))? - .into_array(), - ) -} - -/// Wrap `child` in the parent encoding used by [`stack_parent_fixture`]. -pub 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 with_buffers( - &self, - array: ArrayView<'_, Self>, - buffers: &[BufferHandle], - ) -> VortexResult> { - with_empty_buffers(self, array, buffers) - } - - 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 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 with_buffers( - &self, - array: ArrayView<'_, Self>, - buffers: &[BufferHandle], - ) -> VortexResult> { - with_empty_buffers(self, array, buffers) - } - - 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 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) - } -} From d225445e3cadc7427b16673b9c3c45e29a69802a Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Mon, 8 Jun 2026 11:42:17 +0100 Subject: [PATCH 12/15] fmt stuff Signed-off-by: Adam Gutglick --- vortex-array/src/executor.rs | 77 +++++++++++++++-------------- vortex-array/src/optimizer/mod.rs | 27 +++++----- vortex-array/src/optimizer/rules.rs | 11 +++-- 3 files changed, 59 insertions(+), 56 deletions(-) diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index ae80153300b..9dae03bb0ea 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. @@ -172,11 +173,11 @@ impl ArrayRef { let kernels = execute_parent_kernels.as_ref(); let max_iterations = max_iterations(); - crate::trace_op!(record_execute_until_start::(¤t_array)); + trace_op!(record_execute_until_start::(¤t_array)); for iteration in 0..max_iterations { - crate::trace_op!(use(iteration)); - crate::trace_op!(record_execute_until_iteration( + trace_op!(use(iteration)); + trace_op!(record_execute_until_iteration( iteration, ¤t_array, stack @@ -191,7 +192,7 @@ impl ArrayRef { let done_target = is_done(¤t_array); let done_canonical = AnyCanonical::matches(¤t_array); - crate::trace_op!(record_execute_until_done_check(done_target, done_canonical)); + trace_op!(record_execute_until_done_check(done_target, done_canonical)); if done_target || done_canonical { match stack.pop() { @@ -200,11 +201,11 @@ impl ArrayRef { current_builder.is_none(), "root activation should not retain a builder" ); - crate::trace_op!(record_execute_until_return(¤t_array)); + trace_op!(record_execute_until_return(¤t_array)); return Ok(current_array); } Some(frame) => { - let trace_pop_frame = crate::trace_op!(value( + let trace_pop_frame = trace_op!(value( Some(( frame.parent_array.clone(), current_array.clone(), @@ -214,8 +215,8 @@ impl ArrayRef { )); (current_array, current_builder) = pop_frame(frame, current_array)?; if let Some((parent_before, child_before, slot_idx)) = trace_pop_frame { - crate::trace_op!(use(parent_before, child_before, slot_idx,)); - crate::trace_op!(record_execute_until_pop_frame( + trace_op!(use(parent_before, child_before, slot_idx,)); + trace_op!(record_execute_until_pop_frame( &parent_before, slot_idx, &child_before, @@ -252,13 +253,13 @@ impl ArrayRef { { let frame = stack.pop().vortex_expect("just peeked"); let optimized = result.optimize_ctx(ctx.session())?; - crate::trace_op!(record_execute_optimized(&result, &optimized)); + 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() { - crate::trace_op!(record_execute_parent_none( + trace_op!(record_execute_parent_none( "stack_execute_parent", ¤t_array, )); @@ -269,12 +270,12 @@ impl ArrayRef { && let Some(rewritten) = try_execute_parent(¤t_array, kernels, ctx)? { let optimized = rewritten.optimize_ctx(ctx.session())?; - crate::trace_op!(record_execute_optimized(&rewritten, &optimized)); + trace_op!(record_execute_optimized(&rewritten, &optimized)); current_array = optimized; continue; } if current_builder.is_none() { - crate::trace_op!(record_execute_parent_none( + trace_op!(record_execute_parent_none( "child_execute_parent", ¤t_array, )); @@ -284,14 +285,14 @@ impl ArrayRef { let expected_dtype = current_array.dtype().clone(); let stats = current_array.statistics().to_array_stats(); let encoding_id = current_array.encoding_id(); - crate::trace_op!(record_execute_encoding(¤t_array)); + 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) }?; - crate::trace_op!(record_execute_slot(i, &parent, &child)); + trace_op!(record_execute_slot(i, &parent, &child)); stack.push(StackFrame { parent_array: parent, parent_builder: current_builder.take(), @@ -305,7 +306,7 @@ impl ArrayRef { } ExecutionStep::AppendChild(i) => { if current_builder.is_none() { - crate::trace_op!(record_builder_start(&array)); + trace_op!(record_builder_start(&array)); current_builder = Some(builder_with_capacity_in( ctx.allocator(), array.dtype(), @@ -314,8 +315,8 @@ impl ArrayRef { } let (parent, child) = unsafe { array.take_slot_unchecked(i) }?; - crate::trace_op!(record_append_child(i, &parent, &child)); - crate::trace_op!(record_builder_append(&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. @@ -329,7 +330,7 @@ impl ArrayRef { } ExecutionStep::Done => { let had_builder = current_builder.is_some(); - crate::trace_op!(record_execute_done(&array)); + trace_op!(record_execute_done(&array)); (current_array, current_builder) = finalize_done( array, current_builder, @@ -339,7 +340,7 @@ impl ArrayRef { encoding_id, )?; if had_builder { - crate::trace_op!(record_builder_finish(¤t_array)); + trace_op!(record_builder_finish(¤t_array)); } } } @@ -465,27 +466,27 @@ impl Drop for ExecutionCtx { /// `AppendChild` is returned. impl Executable for ArrayRef { fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - crate::trace_op!(record_single_step_start(&array)); + trace_op!(record_single_step_start(&array)); if let Some(canonical) = array.as_opt::() { let output = Canonical::from(canonical).into_array(); - crate::trace_op!(record_single_step_applied("canonical", &array, &output)); + trace_op!(record_single_step_applied("canonical", &array, &output)); return Ok(output); } - crate::trace_op!(record_single_step_phase_none("canonical", &array)); + trace_op!(record_single_step_phase_none("canonical", &array)); if let Some(reduced) = array.reduce()? { reduced.statistics().inherit_from(array.statistics()); - crate::trace_op!(record_single_step_applied("reduce", &array, &reduced)); + trace_op!(record_single_step_applied("reduce", &array, &reduced)); return Ok(reduced); } - crate::trace_op!(record_single_step_phase_none("reduce", &array)); + 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)? { reduced_parent.statistics().inherit_from(array.statistics()); - crate::trace_op!(record_single_step_applied( + trace_op!(record_single_step_applied( "reduce_parent", &array, &reduced_parent, @@ -493,7 +494,7 @@ impl Executable for ArrayRef { return Ok(reduced_parent); } } - crate::trace_op!(record_single_step_phase_none("reduce_parent", &array)); + 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(); @@ -518,7 +519,7 @@ impl Executable for ArrayRef { executed_parent .statistics() .inherit_from(array.statistics()); - crate::trace_op!(record_single_step_applied( + trace_op!(record_single_step_applied( "execute_parent", &array, &executed_parent, @@ -526,14 +527,14 @@ impl Executable for ArrayRef { return Ok(executed_parent); } } - crate::trace_op!(record_single_step_phase_none("execute_parent", &array)); - crate::trace_op!(record_execute_encoding(&array)); + trace_op!(record_single_step_phase_none("execute_parent", &array)); + trace_op!(record_execute_encoding(&array)); let result = array.execute_encoding(ctx)?; let (array, step) = result.into_parts(); match step { ExecutionStep::Done => { - crate::trace_op!(record_execute_done(&array)); + trace_op!(record_execute_done(&array)); Ok(array) } ExecutionStep::ExecuteSlot(i, _) => { @@ -545,11 +546,11 @@ impl Executable for ArrayRef { } ExecutionStep::AppendChild(_) => { // Single-step: build the entire parent via the builder path. - crate::trace_op!(record_builder_start(&array)); + 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)?; let output = builder.finish(); - crate::trace_op!(record_builder_finish(&output)); + trace_op!(record_builder_finish(&output)); Ok(output) } } @@ -633,11 +634,11 @@ fn execute_parent_for_child( kernels: &ParentExecutionKernels, ctx: &mut ExecutionCtx, ) -> VortexResult> { - crate::trace_op!(use(phase)); + trace_op!(use(phase)); let key = execute_parent_key(parent.encoding_id(), child.encoding_id()); if let Some(plugins) = kernels.get(&key) { for (plugin_idx, plugin) in plugins.as_ref().iter().enumerate() { - crate::trace_op!(use(plugin_idx)); + trace_op!(use(plugin_idx)); if let Some(result) = plugin.execute_parent(child, parent, slot_idx, ctx)? { if cfg!(debug_assertions) { vortex_ensure!( @@ -649,7 +650,7 @@ fn execute_parent_for_child( "Executed parent canonical dtype mismatch" ); } - crate::trace_op!(record_execute_parent_applied( + trace_op!(record_execute_parent_applied( phase, parent, child, @@ -660,7 +661,7 @@ fn execute_parent_for_child( )); return Ok(Some(result)); } - crate::trace_op!(record_execute_parent_attempt( + trace_op!(record_execute_parent_attempt( phase, parent, child, @@ -800,7 +801,7 @@ 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(); - crate::trace_op!(record_execute_step_request::(&array, slot_idx)); + trace_op!(record_execute_step_request::(&array, slot_idx)); Self { array, step: ExecutionStep::ExecuteSlot(slot_idx, M::matches), @@ -812,7 +813,7 @@ impl ExecutionResult { /// `current_array`. pub fn append_child(array: impl IntoArray, slot_idx: usize) -> Self { let array = array.into_array(); - crate::trace_op!(record_append_child_request(&array, slot_idx)); + trace_op!(record_append_child_request(&array, slot_idx)); Self { array, step: ExecutionStep::AppendChild(slot_idx), diff --git a/vortex-array/src/optimizer/mod.rs b/vortex-array/src/optimizer/mod.rs index 0d0171fbf5c..10c10ba4ac6 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,20 +74,20 @@ fn try_optimize( let mut any_optimizations = false; let array_ref = session.map(|s| s.kernels()); - crate::trace_op!(record_optimize_start(array, session.is_some())); + trace_op!(record_optimize_start(array, session.is_some())); // Apply reduction rules to the current array until no more rules apply. for _ in 0..=100 { - crate::trace_op!(record_optimize_loop_start(¤t_array)); + trace_op!(record_optimize_loop_start(¤t_array)); if let Some(new_array) = current_array.reduce()? { current_array = new_array; any_optimizations = true; - crate::trace_op!(record_optimize_loop_end()); + trace_op!(record_optimize_loop_end()); continue; } - crate::trace_op!(record_optimize_reduce_none(¤t_array)); + 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. @@ -100,9 +101,9 @@ fn try_optimize( array_ref.find_reduce_parent(current_array.encoding_id(), child.encoding_id()) { for (plugin_idx, plugin) in plugins.as_ref().iter().enumerate() { - crate::trace_op!(use(plugin_idx)); + trace_op!(use(plugin_idx)); if let Some(new_array) = plugin(child, ¤t_array, slot_idx)? { - crate::trace_op!(record_parent_reduce_applied( + trace_op!(record_parent_reduce_applied( ¤t_array, child, slot_idx, @@ -113,7 +114,7 @@ fn try_optimize( parent_reduced = Some(new_array); break; } - crate::trace_op!(record_parent_reduce_attempt( + trace_op!(record_parent_reduce_attempt( ¤t_array, child, slot_idx, @@ -137,15 +138,15 @@ fn try_optimize( // If the parent was replaced, then we attempt to reduce it again. current_array = new_array; any_optimizations = true; - crate::trace_op!(record_optimize_loop_end()); + trace_op!(record_optimize_loop_end()); continue; } - crate::trace_op!(record_optimize_parent_reduce_none(¤t_array)); - crate::trace_op!(record_optimize_loop_end()); + trace_op!(record_optimize_parent_reduce_none(¤t_array)); + trace_op!(record_optimize_loop_end()); // No more optimizations can be applied - crate::trace_op!(record_optimize_done(¤t_array, any_optimizations)); + trace_op!(record_optimize_done(¤t_array, any_optimizations)); if any_optimizations { return Ok(Some(current_array)); @@ -164,7 +165,7 @@ fn try_optimize_recursive( let mut current_array = array.clone(); let mut any_optimizations = false; - crate::trace_op!(record_optimize_recursive_start(array)); + trace_op!(record_optimize_recursive_start(array)); if let Some(new_array) = try_optimize(¤t_array, Some(session))? { current_array = new_array; @@ -177,7 +178,7 @@ fn try_optimize_recursive( match slot { Some(child) => { if let Some(new_child) = try_optimize_recursive(child, session)? { - crate::trace_op!(record_optimize_recursive_slot( + trace_op!(record_optimize_recursive_slot( new_slots.len(), child, &new_child, diff --git a/vortex-array/src/optimizer/rules.rs b/vortex-array/src/optimizer/rules.rs index 6f77f38f6cf..4358d1daeb5 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,10 +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)? { - crate::trace_op!(record_reduce_applied(array.array(), *rule, &reduced)); + trace_op!(record_reduce_applied(array.array(), *rule, &reduced)); return Ok(Some(reduced)); } - crate::trace_op!(record_reduce_attempt( + trace_op!(record_reduce_attempt( array.array(), *rule, crate::test_harness::trace::AttemptOutcome::Declined, @@ -182,7 +183,7 @@ impl ParentRuleSet { ) -> VortexResult> { for rule in self.rules.iter() { if !rule.matches(parent) { - crate::trace_op!(record_parent_reduce_attempt( + trace_op!(record_parent_reduce_attempt( parent, child.array(), child_idx, @@ -212,7 +213,7 @@ impl ParentRuleSet { ); } - crate::trace_op!(record_parent_reduce_applied( + trace_op!(record_parent_reduce_applied( parent, child.array(), child_idx, @@ -222,7 +223,7 @@ impl ParentRuleSet { )); return Ok(Some(reduced)); } - crate::trace_op!(record_parent_reduce_attempt( + trace_op!(record_parent_reduce_attempt( parent, child.array(), child_idx, From 78eebdc9d2ffc448bba07b23216ebd90d883f500 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Mon, 8 Jun 2026 11:53:26 +0100 Subject: [PATCH 13/15] simplify even more Signed-off-by: Adam Gutglick --- vortex-array/src/executor.rs | 41 +---- vortex-array/src/optimizer/mod.rs | 14 +- vortex-array/src/optimizer/rules.rs | 23 +-- vortex-array/src/test_harness/trace.rs | 217 ++++++++++++++++--------- vortex-array/src/trace_macros.rs | 59 +------ 5 files changed, 164 insertions(+), 190 deletions(-) diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 9dae03bb0ea..1b53de6eeba 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -175,10 +175,9 @@ impl ArrayRef { trace_op!(record_execute_until_start::(¤t_array)); - for iteration in 0..max_iterations { - trace_op!(use(iteration)); + for _iteration in 0..max_iterations { trace_op!(record_execute_until_iteration( - iteration, + _iteration, ¤t_array, stack .last() @@ -205,24 +204,9 @@ impl ArrayRef { return Ok(current_array); } Some(frame) => { - let trace_pop_frame = trace_op!(value( - Some(( - frame.parent_array.clone(), - current_array.clone(), - frame.slot_idx - )), - None::<(ArrayRef, ArrayRef, usize)> - )); + let _slot_idx = frame.slot_idx; (current_array, current_builder) = pop_frame(frame, current_array)?; - if let Some((parent_before, child_before, slot_idx)) = trace_pop_frame { - trace_op!(use(parent_before, child_before, slot_idx,)); - trace_op!(record_execute_until_pop_frame( - &parent_before, - slot_idx, - &child_before, - ¤t_array, - )); - } + trace_op!(record_execute_until_pop_frame(_slot_idx, ¤t_array)); continue; } } @@ -634,11 +618,9 @@ fn execute_parent_for_child( kernels: &ParentExecutionKernels, ctx: &mut ExecutionCtx, ) -> VortexResult> { - trace_op!(use(phase)); let key = execute_parent_key(parent.encoding_id(), child.encoding_id()); if let Some(plugins) = kernels.get(&key) { - for (plugin_idx, plugin) in plugins.as_ref().iter().enumerate() { - trace_op!(use(plugin_idx)); + 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!( @@ -650,25 +632,22 @@ fn execute_parent_for_child( "Executed parent canonical dtype mismatch" ); } - trace_op!(record_execute_parent_applied( + trace_op!(record_session_execute_parent_applied( phase, parent, child, slot_idx, - crate::test_harness::trace::TraceSource::Session(plugin_idx), - "execute_parent_fn", + _plugin_idx, &result, )); return Ok(Some(result)); } - trace_op!(record_execute_parent_attempt( + trace_op!(record_session_execute_parent_declined( phase, parent, child, slot_idx, - crate::test_harness::trace::TraceSource::Session(plugin_idx), - "execute_parent_fn", - crate::test_harness::trace::AttemptOutcome::Declined, + _plugin_idx, )); } } @@ -801,7 +780,6 @@ 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(); - trace_op!(record_execute_step_request::(&array, slot_idx)); Self { array, step: ExecutionStep::ExecuteSlot(slot_idx, M::matches), @@ -813,7 +791,6 @@ impl ExecutionResult { /// `current_array`. pub fn append_child(array: impl IntoArray, slot_idx: usize) -> Self { let array = array.into_array(); - trace_op!(record_append_child_request(&array, slot_idx)); Self { array, step: ExecutionStep::AppendChild(slot_idx), diff --git a/vortex-array/src/optimizer/mod.rs b/vortex-array/src/optimizer/mod.rs index 10c10ba4ac6..01c01013a68 100644 --- a/vortex-array/src/optimizer/mod.rs +++ b/vortex-array/src/optimizer/mod.rs @@ -100,27 +100,23 @@ fn try_optimize( && let Some(plugins) = array_ref.find_reduce_parent(current_array.encoding_id(), child.encoding_id()) { - for (plugin_idx, plugin) in plugins.as_ref().iter().enumerate() { - trace_op!(use(plugin_idx)); + for (_plugin_idx, plugin) in plugins.as_ref().iter().enumerate() { if let Some(new_array) = plugin(child, ¤t_array, slot_idx)? { - trace_op!(record_parent_reduce_applied( + trace_op!(record_session_parent_reduce_applied( ¤t_array, child, slot_idx, - crate::test_harness::trace::TraceSource::Session(plugin_idx), - "reduce_parent_fn", + _plugin_idx, &new_array, )); parent_reduced = Some(new_array); break; } - trace_op!(record_parent_reduce_attempt( + trace_op!(record_session_parent_reduce_declined( ¤t_array, child, slot_idx, - crate::test_harness::trace::TraceSource::Session(plugin_idx), - "reduce_parent_fn", - crate::test_harness::trace::AttemptOutcome::Declined, + _plugin_idx, )); } if parent_reduced.is_some() { diff --git a/vortex-array/src/optimizer/rules.rs b/vortex-array/src/optimizer/rules.rs index 4358d1daeb5..8539e39aea2 100644 --- a/vortex-array/src/optimizer/rules.rs +++ b/vortex-array/src/optimizer/rules.rs @@ -137,11 +137,7 @@ impl ReduceRuleSet { trace_op!(record_reduce_applied(array.array(), *rule, &reduced)); return Ok(Some(reduced)); } - trace_op!(record_reduce_attempt( - array.array(), - *rule, - crate::test_harness::trace::AttemptOutcome::Declined, - )); + trace_op!(record_reduce_declined(array.array(), *rule)); } Ok(None) } @@ -183,13 +179,11 @@ impl ParentRuleSet { ) -> VortexResult> { for rule in self.rules.iter() { if !rule.matches(parent) { - trace_op!(record_parent_reduce_attempt( + trace_op!(record_static_parent_reduce_no_match( parent, child.array(), child_idx, - crate::test_harness::trace::TraceSource::Static, - crate::test_harness::trace::compact_label(*rule), - crate::test_harness::trace::AttemptOutcome::NoMatch, + *rule, )); continue; } @@ -213,23 +207,20 @@ impl ParentRuleSet { ); } - trace_op!(record_parent_reduce_applied( + trace_op!(record_static_parent_reduce_applied( parent, child.array(), child_idx, - crate::test_harness::trace::TraceSource::Static, - crate::test_harness::trace::compact_label(*rule), + *rule, &reduced, )); return Ok(Some(reduced)); } - trace_op!(record_parent_reduce_attempt( + trace_op!(record_static_parent_reduce_declined( parent, child.array(), child_idx, - crate::test_harness::trace::TraceSource::Static, - crate::test_harness::trace::compact_label(*rule), - crate::test_harness::trace::AttemptOutcome::Declined, + *rule, )); } Ok(None) diff --git a/vortex-array/src/test_harness/trace.rs b/vortex-array/src/test_harness/trace.rs index eeea5d11e5e..31a8daf834e 100644 --- a/vortex-array/src/test_harness/trace.rs +++ b/vortex-array/src/test_harness/trace.rs @@ -279,7 +279,7 @@ fn attempts_enabled() -> bool { } #[derive(Clone, Copy, Debug)] -pub(crate) enum TraceSource { +enum TraceSource { Static, Session(usize), } @@ -294,7 +294,7 @@ impl Display for TraceSource { } #[derive(Clone, Copy, Debug)] -pub(crate) enum AttemptOutcome { +enum AttemptOutcome { Declined, NoMatch, } @@ -421,26 +421,108 @@ pub(crate) fn record_optimize_recursive_slot(slot_idx: usize, input: &ArrayRef, }); } -pub(crate) fn record_reduce_attempt(array: &ArrayRef, rule: &dyn Debug, outcome: AttemptOutcome) { +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, + outcome: AttemptOutcome::Declined, }); } -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_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_parent_reduce_attempt( +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, @@ -461,7 +543,7 @@ pub(crate) fn record_parent_reduce_attempt( }); } -pub(crate) fn record_parent_reduce_applied( +fn record_parent_reduce_applied( parent: &ArrayRef, child: &ArrayRef, slot_idx: usize, @@ -513,21 +595,51 @@ pub(crate) fn record_execute_until_return(output: &ArrayRef) { }); } -pub(crate) fn record_execute_until_pop_frame( +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, - slot_idx: usize, child: &ArrayRef, + slot_idx: usize, + plugin_idx: usize, output: &ArrayRef, ) { - record(TraceEvent::ExecuteUntilPopFrame { - parent: ArraySummary::new(parent), + record_execute_parent_applied( + phase, + parent, + child, slot_idx, - child: ArraySummary::new(child), - output: ArraySummary::new(output), - }); + 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, + ); } -pub(crate) fn record_execute_parent_attempt( +fn record_execute_parent_attempt( phase: &'static str, parent: &ArrayRef, child: &ArrayRef, @@ -550,7 +662,7 @@ pub(crate) fn record_execute_parent_attempt( }); } -pub(crate) fn record_execute_parent_applied( +fn record_execute_parent_applied( phase: &'static str, parent: &ArrayRef, child: &ArrayRef, @@ -603,30 +715,6 @@ pub(crate) fn record_execute_encoding(array: &ArrayRef) { }); } -pub(crate) fn record_execute_step_request(array: &ArrayRef, slot_idx: usize) { - if !attempts_enabled() { - return; - } - record(TraceEvent::ExecutionRequest { - step: "ExecuteSlot", - parent: ArraySummary::new(array), - slot_idx, - target: Some(short_type_name::()), - }); -} - -pub(crate) fn record_append_child_request(array: &ArrayRef, slot_idx: usize) { - if !attempts_enabled() { - return; - } - record(TraceEvent::ExecutionRequest { - step: "AppendChild", - parent: ArraySummary::new(array), - slot_idx, - target: None, - }); -} - pub(crate) fn record_execute_slot(slot_idx: usize, parent: &ArrayRef, child: &ArrayRef) { record(TraceEvent::SlotTransition { step: "ExecuteSlot", @@ -709,7 +797,7 @@ fn record(event: TraceEvent) { }); } -pub(crate) fn compact_label(value: &dyn Debug) -> String { +fn compact_label(value: &dyn Debug) -> String { let label = format!("{value:?}"); if let Some(label) = adapter_field(&label, "rule") { return label.to_string(); @@ -849,9 +937,7 @@ enum TraceEvent { output: ArraySummary, }, ExecuteUntilPopFrame { - parent: ArraySummary, slot_idx: usize, - child: ArraySummary, output: ArraySummary, }, ExecuteParentAttempt { @@ -886,12 +972,6 @@ enum TraceEvent { ExecuteEncoding { array: ArraySummary, }, - ExecutionRequest { - step: &'static str, - parent: ArraySummary, - slot_idx: usize, - target: Option, - }, SlotTransition { step: &'static str, slot_idx: usize, @@ -927,7 +1007,6 @@ impl TraceEvent { | TraceEvent::PhaseNone { .. } | TraceEvent::ExecuteUntilDoneCheck { .. } | TraceEvent::ExecuteEncoding { .. } - | TraceEvent::ExecutionRequest { .. } | TraceEvent::ExecuteOptimized { changed: false, .. } | TraceEvent::ExecuteParentAttempt { .. } | TraceEvent::ReduceAttempt { .. } @@ -981,7 +1060,6 @@ impl TraceEvent { | TraceEvent::ExecuteParentApplied { .. } | TraceEvent::ExecuteOptimized { .. } | TraceEvent::ExecuteEncoding { .. } - | TraceEvent::ExecutionRequest { .. } | TraceEvent::SlotTransition { .. } | TraceEvent::BuilderEvent { .. } | TraceEvent::ExecuteDone { .. } => 2, @@ -1087,15 +1165,9 @@ impl TraceEvent { TraceEvent::ExecuteUntilReturn { output } => { write!(f, "return output={output}") } - TraceEvent::ExecuteUntilPopFrame { - parent, - slot_idx, - child, - output, - } => write!( - f, - "pop_frame slot={slot_idx} parent={parent} child={child} output={output}" - ), + TraceEvent::ExecuteUntilPopFrame { slot_idx, output } => { + write!(f, "pop_frame slot={slot_idx} output={output}") + } TraceEvent::ExecuteParentAttempt { phase, parent, @@ -1148,18 +1220,6 @@ impl TraceEvent { TraceEvent::ExecuteEncoding { array } => { write!(f, "execute encoding={array}") } - TraceEvent::ExecutionRequest { - step, - parent, - slot_idx, - target, - } => { - write!(f, "request {step} slot={slot_idx}")?; - if let Some(target) = target { - write!(f, " target={target}")?; - } - write!(f, " parent={parent}") - } TraceEvent::SlotTransition { step, slot_idx, @@ -1612,7 +1672,7 @@ optimize root=vortex.filter(i32, len=4) session=false 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 parent=vortex.filter(i32, len=2) child=vortex.primitive(i32, len=4) output=vortex.filter(i32, len=2) + 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 @@ -1646,7 +1706,6 @@ optimize root=vortex.filter(i32, len=4) session=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[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) - request ExecuteSlot slot=0 target=Primitive parent=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 diff --git a/vortex-array/src/trace_macros.rs b/vortex-array/src/trace_macros.rs index 7b8503ae086..fd66f38b201 100644 --- a/vortex-array/src/trace_macros.rs +++ b/vortex-array/src/trace_macros.rs @@ -13,75 +13,26 @@ /// while debug-test builds with an active recorder route the call into /// [`test_harness::trace`][crate::test_harness::trace]. /// -/// # Forms -/// -/// ```ignore -/// // 1. Event dispatch — the common form. Calls `test_harness::trace::(...)` only when -/// // a recorder is installed on this thread. Otherwise compiles to nothing. -/// crate::trace_op!(record_execute_until_start::(&array)); -/// -/// // 2. `use(...)` — keep otherwise-unused bindings alive in non-test builds so the -/// // surrounding code still type-checks when the trace call is compiled out. -/// crate::trace_op!(use(plugin_idx, slot_idx)); -/// -/// // 3. `value(, )` — pick between two expressions depending on -/// // whether a recorder is active. Useful when the active branch needs to clone state -/// // just for the trace. -/// let snapshot = crate::trace_op!(value(Some(parent.clone()), None)); -/// ``` -/// /// # 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 `()`. Captured names are bound -/// through the `use(...)` form so the compiler does not warn about unused variables. +/// - 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 { - (@when_enabled { $($enabled:tt)* } else { $($disabled:tt)* }) => {{ + ($event:ident $(::<$($generic:ty),*>)?($($arg:expr),* $(,)?)) => {{ #[cfg(all(any(test, feature = "_test-harness"), not(codspeed)))] { - $($enabled)* - } - - #[cfg(any(not(any(test, feature = "_test-harness")), codspeed))] - { - $($disabled)* - } - }}; - - (@if_active { $($enabled:tt)* } else { $($disabled:tt)* }) => { - $crate::trace_op!(@when_enabled { if $crate::test_harness::trace::is_active() { - $($enabled)* - } else { - $($disabled)* + $crate::test_harness::trace::$event $(::<$($generic),*>)?($($arg),*); } - } else { - $($disabled)* - }) - }; - - (use($($value:expr),* $(,)?)) => { - $crate::trace_op!(@when_enabled {} else { - let _ = ($(&$value),*); - }) - }; - - (value($enabled:expr, $disabled:expr)) => { - $crate::trace_op!(@if_active { $enabled } else { $disabled }) - }; - - ($($event:tt)*) => { - $crate::trace_op!(@if_active { - $crate::test_harness::trace::$($event)* - } else {}) - }; + } + }}; } pub(crate) use trace_op; From 1f773126b67435e561a509549d9c0174b40540e9 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Mon, 8 Jun 2026 11:58:31 +0100 Subject: [PATCH 14/15] fix clippy Signed-off-by: Adam Gutglick --- vortex-array/src/executor.rs | 1 + vortex-array/src/optimizer/mod.rs | 1 + 2 files changed, 2 insertions(+) diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 1b53de6eeba..3e73d52b350 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -620,6 +620,7 @@ fn execute_parent_for_child( ) -> VortexResult> { let key = execute_parent_key(parent.encoding_id(), child.encoding_id()); if let Some(plugins) = kernels.get(&key) { + #[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) { diff --git a/vortex-array/src/optimizer/mod.rs b/vortex-array/src/optimizer/mod.rs index 01c01013a68..41a6bb69b25 100644 --- a/vortex-array/src/optimizer/mod.rs +++ b/vortex-array/src/optimizer/mod.rs @@ -100,6 +100,7 @@ fn try_optimize( && let Some(plugins) = array_ref.find_reduce_parent(current_array.encoding_id(), child.encoding_id()) { + #[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)? { trace_op!(record_session_parent_reduce_applied( From 8ead1e82ee9588c22e4064a2642c16e00ed79da7 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 30 Jul 2026 15:47:11 +0100 Subject: [PATCH 15/15] test: trace reductions and executions of scan ops over complex and compressed arrays (#8828) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Stacked on #7814 (the `trace_op` execution/optimization tracing harness). Builds out trace-test coverage for how take/filter/compare/like reduce and execute over complex arrays, using the harness from the base PR: **vortex-array** (promotes `test_harness/trace.rs` to a directory module `trace/mod.rs` + `trace/tests.rs`; harness source unchanged): - take (`DictArray`) over a chunked array — no reduce rule fires; execution canonicalizes the chunked values chunk-by-chunk before the dict take kernel runs. - filter over a struct with complex children (dict-of-strings + chunked fields) — `FilterStructRule` pushes the filter into each field, `FilterReduceAdaptor(Dict)` rewrites the dict field at optimize time. - compare (`Binary` Eq vs constant) over a dict — `DictionaryScalarFnValuesPushDownRule` reduces the compare into the dictionary values. - like over a dict-of-strings — `LikeReduceAdaptor(Dict)` reduces into the dictionary values. **vortex-runend** (new `trace_tests.rs`, `insta` added as dev-dependency): - compare vs constant reduces via `RunEndScalarFnRule` at optimize time; filter and take execute via the `FilterExecuteAdaptor`/`TakeExecuteAdaptor` parent kernels. **vortex-btrblocks** (new `trace_tests.rs`; `insta`, `tpchgen`, `tpchgen-arrow`, `arrow-array` added as dev-dependencies): generates TPC-H lineitem (SF 0.001, 4096 rows, deterministic — tpchgen and the compressor's sampling seed are both fixed), compresses it with `BtrBlocksCompressor`, and traces TPC-H-style scan predicates over the resulting encodings: - `l_shipdate >= const` over `ext(date) → for → bitpacked` — extension compare kernel at execution time, with `CastReduceAdaptor(FoR)/(BitPacked)` reductions inside the arrow fallback. - `l_quantity < const` over `decimal_byte_parts → dict → bitpacked` — the decimal_byte_parts compare kernel pushes into the byte-parts dictionary, where `DictionaryScalarFnValuesPushDownRule` reduces the compare to the 50 dictionary values before decode. - `l_shipmode = const` over dict-of-FSST — dict pushdown at optimize time, FSST compare kernel + dict decode at execution time. - `l_comment LIKE` over FSST — the FSST like kernel matches in compressed space. - filter over a struct of compressed columns — `FilterStructRule` plus per-encoding pushdown via `DecimalBytePartsFilterPushDownRule`, `ExtensionFilterPushDownRule`, `FoRFilterPushDownRule`, and `FilterReduceAdaptor(Dict)`, leaving execution a no-op. - take over a struct of compressed columns — `TakeReduceAdaptor(Struct)` absorbs the take entirely at optimize time. Each btrblocks test also asserts the canonical result matches running the same operation over the uncompressed column. A 1MiB TPC-H text pool replaces the spec-default 300MiB one, cutting per-test setup from ~9s to ~0.5s without changing the traces. This PR is tests-only; no production code changes. - `cargo nextest run -p vortex-array` — 2961 passed (includes the 8 trace tests, run in parallel to confirm the thread-local recorder is deterministic). - `cargo nextest run -p vortex-runend` — 61 passed. - `cargo nextest run -p vortex-btrblocks` — 43 passed. - `cargo clippy --all-targets` on the three touched crates — clean. - `cargo +nightly fmt --all` and `git diff --check` — clean. 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_01RMBcCiwpC3dbGLXRbK7UuR --- _Generated by [Claude Code](https://claude.ai/code/session_01RMBcCiwpC3dbGLXRbK7UuR)_ --------- Co-authored-by: Claude Fable 5 --- Cargo.lock | 6 + encodings/runend/Cargo.toml | 1 + encodings/runend/src/lib.rs | 3 + encodings/runend/src/trace_tests.rs | 148 ++++ vortex-array/src/executor.rs | 17 +- .../test_harness/{trace.rs => trace/mod.rs} | 512 +----------- vortex-array/src/test_harness/trace/tests.rs | 771 ++++++++++++++++++ vortex-btrblocks/Cargo.toml | 5 + vortex-btrblocks/src/lib.rs | 3 + vortex-btrblocks/src/trace_tests.rs | 448 ++++++++++ 10 files changed, 1392 insertions(+), 522 deletions(-) create mode 100644 encodings/runend/src/trace_tests.rs rename vortex-array/src/test_harness/{trace.rs => trace/mod.rs} (66%) create mode 100644 vortex-array/src/test_harness/trace/tests.rs create mode 100644 vortex-btrblocks/src/trace_tests.rs 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/executor.rs b/vortex-array/src/executor.rs index 3e73d52b350..4e87ff5e7dd 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -611,7 +611,7 @@ fn finalize_done( } fn execute_parent_for_child( - phase: &'static str, + _phase: &'static str, parent: &ArrayRef, child: &ArrayRef, slot_idx: usize, @@ -634,7 +634,7 @@ fn execute_parent_for_child( ); } trace_op!(record_session_execute_parent_applied( - phase, + _phase, parent, child, slot_idx, @@ -644,7 +644,7 @@ fn execute_parent_for_child( return Ok(Some(result)); } trace_op!(record_session_execute_parent_declined( - phase, + _phase, parent, child, slot_idx, @@ -664,14 +664,9 @@ fn try_execute_parent( ) -> VortexResult> { for (slot_idx, slot) in array.slots().iter().enumerate() { let Some(child) = slot else { continue }; - if let Some(executed_parent) = execute_parent_for_child( - "child_execute_parent", - array, - child, - slot_idx, - kernels, - ctx, - )? { + if let Some(executed_parent) = + execute_parent_for_child("child_execute_parent", array, child, slot_idx, kernels, ctx)? + { ctx.log(format_args!( "execute_parent: slot[{}]({}) rewrote {} -> {}", slot_idx, diff --git a/vortex-array/src/test_harness/trace.rs b/vortex-array/src/test_harness/trace/mod.rs similarity index 66% rename from vortex-array/src/test_harness/trace.rs rename to vortex-array/src/test_harness/trace/mod.rs index 31a8daf834e..709259f0bcb 100644 --- a/vortex-array/src/test_harness/trace.rs +++ b/vortex-array/src/test_harness/trace/mod.rs @@ -1254,514 +1254,4 @@ impl TraceEvent { } #[cfg(test)] -mod tests { - 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::ChunkedArray; - use crate::arrays::Filter; - use crate::arrays::FilterArray; - use crate::arrays::Primitive; - use crate::arrays::PrimitiveArray; - use crate::arrays::filter::FilterArraySlotsExt; - 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::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 mut ctx = ExecutionCtx::new(VortexSession::empty().with::()); - 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 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 mut ctx = ExecutionCtx::new(VortexSession::empty().with::()); - 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 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 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(()) - } -} +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-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(()) +}