diff --git a/Cargo.lock b/Cargo.lock index 781932651ca..d2b06dc3b9c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9777,7 +9777,9 @@ dependencies = [ name = "vortex-btrblocks" version = "0.1.0" dependencies = [ + "arrow-array", "codspeed-divan-compat", + "insta", "itertools 0.14.0", "num-traits", "pco", @@ -9785,6 +9787,8 @@ dependencies = [ "rstest", "rustc-hash", "test-with", + "tpchgen", + "tpchgen-arrow", "tracing", "vortex-alp", "vortex-array", @@ -10501,6 +10505,7 @@ dependencies = [ "arrow-array", "arrow-schema", "codspeed-divan-compat", + "insta", "itertools 0.14.0", "num-traits", "prost 0.14.4", diff --git a/encodings/runend/Cargo.toml b/encodings/runend/Cargo.toml index 01a5b8d7a3e..7d923ffeeee 100644 --- a/encodings/runend/Cargo.toml +++ b/encodings/runend/Cargo.toml @@ -32,6 +32,7 @@ workspace = true arrow-array = { workspace = true } arrow-schema = { workspace = true } divan = { workspace = true } +insta = { workspace = true } itertools = { workspace = true } rand = { workspace = true } rstest = { workspace = true } diff --git a/encodings/runend/src/lib.rs b/encodings/runend/src/lib.rs index 8770dbbf58e..13dd44581ad 100644 --- a/encodings/runend/src/lib.rs +++ b/encodings/runend/src/lib.rs @@ -18,6 +18,9 @@ mod iter; mod kernel; mod ops; mod rules; +#[cfg(test)] +#[cfg(not(codspeed))] +mod trace_tests; #[doc(hidden)] pub mod _benchmarking { diff --git a/encodings/runend/src/trace_tests.rs b/encodings/runend/src/trace_tests.rs new file mode 100644 index 00000000000..cbcd0137d8f --- /dev/null +++ b/encodings/runend/src/trace_tests.rs @@ -0,0 +1,149 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Snapshot tests tracing how reductions and executions flow through run-end arrays. + +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::DictArray; +use vortex_array::arrays::FilterArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::assert_arrays_eq; +use vortex_array::optimizer::ArrayOptimizer; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::fns::binary::Binary; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::session::ArraySession; +use vortex_array::test_harness::trace::trace_op; +use vortex_error::VortexResult; +use vortex_mask::Mask; +use vortex_session::VortexSession; + +use crate::RunEnd; + +fn execution_ctx() -> ExecutionCtx { + ExecutionCtx::new(VortexSession::empty().with::()) +} + +/// 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]) + ); + insta::assert_snapshot!(traced.trace.to_string(), @" + execute_until target=AnyCanonical root=vortex.runend(bool, len=9) + iter 0 current=vortex.runend(bool, len=9) builder_active=false + execute_until target=AnyCanonical root=vortex.binary(bool, len=3) + iter 0 current=vortex.binary(bool, len=3) builder_active=false + optimize root=vortex.cast(i32?, len=3) session=false + reduce_parent static:CastReduceAdaptor(Primitive) slot=0 parent=vortex.cast(i32?, len=3) child=vortex.primitive(i32, len=3) -> vortex.primitive(i32?, len=3) + done output=vortex.primitive(i32?, len=3) + optimize root=vortex.slice(i32, len=1) session=false + reduce_parent static:SliceReduceAdaptor(Constant) slot=0 parent=vortex.slice(i32, len=1) child=vortex.constant(i32, len=3) -> vortex.constant(i32, len=1) + done output=vortex.constant(i32, len=1) + optimize root=vortex.cast(i32?, len=1) session=false + reduce_parent static:CastReduceAdaptor(Constant) slot=0 parent=vortex.cast(i32?, len=1) child=vortex.constant(i32, len=1) -> vortex.constant(i32?, len=1) + done output=vortex.constant(i32?, len=1) + execute_until target=AnyCanonical root=vortex.constant(i32?, len=1) + iter 0 current=vortex.constant(i32?, len=1) builder_active=false + Done array=vortex.primitive(i32?, len=1) + iter 1 current=vortex.primitive(i32?, len=1) builder_active=false + return output=vortex.primitive(i32?, len=1) + Done array=vortex.bool(bool, len=3) + iter 1 current=vortex.bool(bool, len=3) builder_active=false + return output=vortex.bool(bool, len=3) + Done array=vortex.bool(bool, len=9) + iter 1 current=vortex.bool(bool, len=9) builder_active=false + return output=vortex.bool(bool, len=9) + "); + + Ok(()) +} + +#[test] +fn trace_filter_on_runend() -> VortexResult<()> { + let runend = runend_array()?; + let filtered = FilterArray::try_new( + runend, + Mask::from_iter([true, false, false, true, true, false, false, true, false]), + )? + .into_array(); + + let traced = trace_op(|| { + filtered + .execute::(&mut execution_ctx()) + .map(IntoArray::into_array) + })?; + + assert_arrays_eq!(traced.output, PrimitiveArray::from_iter([1i32, 2, 2, 3])); + insta::assert_snapshot!(traced.trace.to_string(), @" + execute_until target=AnyCanonical root=vortex.filter(i32, len=4) + iter 0 current=vortex.filter(i32, len=4) builder_active=false + child_execute_parent static:kernel[2] slot=0 parent=vortex.filter(i32, len=4) child=vortex.runend(i32, len=9) -> vortex.dict(i32, len=4) + iter 1 current=vortex.dict(i32, len=4) builder_active=false + child_execute_parent static:kernel[3] slot=1 parent=vortex.dict(i32, len=4) child=vortex.primitive(i32, len=3) -> vortex.primitive(i32, len=4) + iter 2 current=vortex.primitive(i32, len=4) builder_active=false + return output=vortex.primitive(i32, len=4) + "); + + Ok(()) +} + +#[test] +fn trace_take_on_runend() -> VortexResult<()> { + let runend = runend_array()?; + // A take is expressed as a `DictArray` whose codes are the take indices. + let indices = PrimitiveArray::from_iter([8u64, 0, 4, 4]).into_array(); + let take = DictArray::try_new(indices, runend)?.into_array(); + + let traced = trace_op(|| { + take.execute::(&mut execution_ctx()) + .map(IntoArray::into_array) + })?; + + assert_arrays_eq!(traced.output, PrimitiveArray::from_iter([3i32, 1, 2, 2])); + insta::assert_snapshot!(traced.trace.to_string(), @" + execute_until target=AnyCanonical root=vortex.dict(i32, len=4) + iter 0 current=vortex.dict(i32, len=4) builder_active=false + child_execute_parent static:kernel[3] slot=1 parent=vortex.dict(i32, len=4) child=vortex.runend(i32, len=9) -> vortex.dict(i32, len=4) + iter 1 current=vortex.dict(i32, len=4) builder_active=false + child_execute_parent static:kernel[3] slot=1 parent=vortex.dict(i32, len=4) child=vortex.primitive(i32, len=3) -> vortex.primitive(i32, len=4) + iter 2 current=vortex.primitive(i32, len=4) builder_active=false + return output=vortex.primitive(i32, len=4) + "); + + Ok(()) +} diff --git a/vortex-array/src/test_harness/trace.rs b/vortex-array/src/test_harness/trace/mod.rs similarity index 67% rename from vortex-array/src/test_harness/trace.rs rename to vortex-array/src/test_harness/trace/mod.rs index 367588007df..2e0bf7f9ea2 100644 --- a/vortex-array/src/test_harness/trace.rs +++ b/vortex-array/src/test_harness/trace/mod.rs @@ -1331,513 +1331,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::arrays::ChunkedArray; - use crate::arrays::Filter; - use crate::arrays::FilterArray; - use crate::arrays::Primitive; - use crate::arrays::PrimitiveArray; - use crate::arrays::filter::FilterArrayExt; - 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::kernel::ParentKernelSet; - use crate::matcher::Matcher; - use crate::optimizer::ArrayOptimizer; - 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()?) - } - - 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 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 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, - ]))) - } - - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - STACK_CHILD_PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } - } - - const STACK_CHILD_PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&StackDeclineKernel), - ParentKernelSet::lift(&StackParentKernel), - ]); - - #[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); - 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])); - 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, - ) -> VortexResult<()> { - let mut ctx = ExecutionCtx::new(VortexSession::empty().with::()); - 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])); - insta::assert_snapshot!(traced.trace.to_string(), @r" - 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=static kernel=kernel[0] 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=static kernel=kernel[1] 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=static kernel=kernel[0] 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=static kernel=kernel[1] 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])); - insta::assert_snapshot!(traced.trace.to_string(), @r" - 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) - execute_until target=AnyCanonical root=vortex.primitive(i32, len=2) - iter 0 current=vortex.primitive(i32, len=2) builder_active=false - return output=vortex.primitive(i32, len=2) - execute_until target=AnyCanonical root=vortex.primitive(i32, len=2) - iter 0 current=vortex.primitive(i32, len=2) builder_active=false - return output=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) - execute_until target=AnyCanonical root=vortex.primitive(i32, len=1) - iter 0 current=vortex.primitive(i32, len=1) builder_active=false - return output=vortex.primitive(i32, len=1) - execute_until target=AnyCanonical root=vortex.primitive(i32, len=1) - iter 0 current=vortex.primitive(i32, len=1) builder_active=false - return output=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) - execute_until target=AnyCanonical root=vortex.primitive(i32, len=2) - iter 0 current=vortex.primitive(i32, len=2) builder_active=false - return output=vortex.primitive(i32, len=2) - execute_until target=AnyCanonical root=vortex.primitive(i32, len=2) - iter 0 current=vortex.primitive(i32, len=2) builder_active=false - return output=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..416266408fd --- /dev/null +++ b/vortex-array/src/test_harness/trace/tests.rs @@ -0,0 +1,790 @@ +// 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::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::FilterArrayExt; +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::kernel::ParentKernelSet; +use crate::matcher::Matcher; +use crate::optimizer::ArrayOptimizer; +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()?) +} + +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 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 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, + ]))) + } + + fn execute_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + STACK_CHILD_PARENT_KERNELS.execute(array, parent, child_idx, ctx) + } +} + +const STACK_CHILD_PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ + ParentKernelSet::lift(&StackDeclineKernel), + ParentKernelSet::lift(&StackParentKernel), +]); + +#[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); + 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])); + 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, +) -> VortexResult<()> { + let mut ctx = ExecutionCtx::new(VortexSession::empty().with::()); + 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])); + insta::assert_snapshot!(traced.trace.to_string(), @r" + 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=static kernel=kernel[0] 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=static kernel=kernel[1] 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=static kernel=kernel[0] 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=static kernel=kernel[1] 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])); + insta::assert_snapshot!(traced.trace.to_string(), @r" + 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) + execute_until target=AnyCanonical root=vortex.primitive(i32, len=2) + iter 0 current=vortex.primitive(i32, len=2) builder_active=false + return output=vortex.primitive(i32, len=2) + execute_until target=AnyCanonical root=vortex.primitive(i32, len=2) + iter 0 current=vortex.primitive(i32, len=2) builder_active=false + return output=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) + execute_until target=AnyCanonical root=vortex.primitive(i32, len=1) + iter 0 current=vortex.primitive(i32, len=1) builder_active=false + return output=vortex.primitive(i32, len=1) + execute_until target=AnyCanonical root=vortex.primitive(i32, len=1) + iter 0 current=vortex.primitive(i32, len=1) builder_active=false + return output=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) + execute_until target=AnyCanonical root=vortex.primitive(i32, len=2) + iter 0 current=vortex.primitive(i32, len=2) builder_active=false + return output=vortex.primitive(i32, len=2) + execute_until target=AnyCanonical root=vortex.primitive(i32, len=2) + iter 0 current=vortex.primitive(i32, len=2) builder_active=false + return output=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]) + ); + 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) + execute_until target=AnyCanonical root=vortex.primitive(i32, len=1) + iter 0 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) + execute_until target=AnyCanonical root=vortex.primitive(i32, len=2) + iter 0 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 static:kernel[3] 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 static:kernel[3] slot=1 parent=vortex.dict(i32, len=4) child=vortex.primitive(i32, len=3) -> vortex.primitive(i32, len=4) + iter 2 current=vortex.primitive(i32, len=4) builder_active=false + return output=vortex.primitive(i32, len=4) + "); + + Ok(()) +} + +#[test] +fn trace_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); + 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]) + ); + 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 + optimize root=vortex.cast(i32?, len=3) session=false + reduce_parent static:CastReduceAdaptor(Primitive) slot=0 parent=vortex.cast(i32?, len=3) child=vortex.primitive(i32, len=3) -> vortex.primitive(i32?, len=3) + done output=vortex.primitive(i32?, len=3) + optimize root=vortex.slice(i32, len=1) session=false + reduce_parent static:SliceReduceAdaptor(Constant) slot=0 parent=vortex.slice(i32, len=1) child=vortex.constant(i32, len=3) -> vortex.constant(i32, len=1) + done output=vortex.constant(i32, len=1) + optimize root=vortex.cast(i32?, len=1) session=false + reduce_parent static:CastReduceAdaptor(Constant) slot=0 parent=vortex.cast(i32?, len=1) child=vortex.constant(i32, len=1) -> vortex.constant(i32?, len=1) + done output=vortex.constant(i32?, len=1) + execute_until target=AnyCanonical root=vortex.constant(i32?, len=1) + iter 0 current=vortex.constant(i32?, len=1) builder_active=false + Done array=vortex.primitive(i32?, len=1) + iter 1 current=vortex.primitive(i32?, len=1) builder_active=false + return output=vortex.primitive(i32?, len=1) + Done array=vortex.bool(bool, len=3) + iter 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 static:kernel[2] 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]) + ); + 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 + optimize root=vortex.cast(utf8?, len=3) session=false + reduce_parent static:CastReduceAdaptor(VarBinView) slot=0 parent=vortex.cast(utf8?, len=3) child=vortex.varbinview(utf8, len=3) -> vortex.varbinview(utf8?, len=3) + done output=vortex.varbinview(utf8?, len=3) + optimize root=vortex.slice(utf8, len=1) session=false + reduce_parent static:SliceReduceAdaptor(Constant) slot=0 parent=vortex.slice(utf8, len=1) child=vortex.constant(utf8, len=3) -> vortex.constant(utf8, len=1) + done output=vortex.constant(utf8, len=1) + optimize root=vortex.cast(utf8?, len=1) session=false + reduce_parent static:CastReduceAdaptor(Constant) slot=0 parent=vortex.cast(utf8?, len=1) child=vortex.constant(utf8, len=1) -> vortex.constant(utf8?, len=1) + done output=vortex.constant(utf8?, len=1) + execute_until target=AnyCanonical root=vortex.constant(utf8?, len=1) + iter 0 current=vortex.constant(utf8?, len=1) builder_active=false + Done array=vortex.varbinview(utf8?, len=1) + iter 1 current=vortex.varbinview(utf8?, len=1) builder_active=false + return output=vortex.varbinview(utf8?, len=1) + 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 static:kernel[2] 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 1adb6508828..bbea747656a 100644 --- a/vortex-btrblocks/Cargo.toml +++ b/vortex-btrblocks/Cargo.toml @@ -41,9 +41,13 @@ 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-session = { workspace = true } diff --git a/vortex-btrblocks/src/lib.rs b/vortex-btrblocks/src/lib.rs index 39db05246a6..97065c22b8c 100644 --- a/vortex-btrblocks/src/lib.rs +++ b/vortex-btrblocks/src/lib.rs @@ -58,6 +58,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..e0141f310c6 --- /dev/null +++ b/vortex-btrblocks/src/trace_tests.rs @@ -0,0 +1,435 @@ +// 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::PrimitiveArray; +use vortex_array::arrays::Struct; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::arrow::FromArrowArray; +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::test_harness::trace::Traced; +use vortex_array::test_harness::trace::trace_op; +use vortex_error::VortexResult; +use vortex_mask::Mask; +use vortex_session::VortexSession; + +use crate::BtrBlocksCompressor; + +fn execution_ctx() -> ExecutionCtx { + ExecutionCtx::new(VortexSession::empty().with::()) +} + +/// 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); + 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 static:kernel[0] 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 + optimize root=vortex.cast(i32?, len=4096) session=false + optimize root=vortex.cast(i32?, len=4096) session=false + reduce_parent static:CastReduceAdaptor(BitPacked) slot=0 parent=vortex.cast(i32?, len=4096) child=fastlanes.bitpacked(i32, len=4096) -> fastlanes.bitpacked(i32?, len=4096) + done output=fastlanes.bitpacked(i32?, len=4096) + reduce_parent static:CastReduceAdaptor(FoR) slot=0 parent=vortex.cast(i32?, len=4096) child=fastlanes.for(i32, len=4096) -> fastlanes.for(i32?, len=4096) + done output=fastlanes.for(i32?, len=4096) + 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) + optimize root=vortex.slice(i32, len=1) session=false + reduce_parent static:SliceReduceAdaptor(Constant) slot=0 parent=vortex.slice(i32, len=1) child=vortex.constant(i32, len=4096) -> vortex.constant(i32, len=1) + done output=vortex.constant(i32, len=1) + optimize root=vortex.cast(i32?, len=1) session=false + reduce_parent static:CastReduceAdaptor(Constant) slot=0 parent=vortex.cast(i32?, len=1) child=vortex.constant(i32, len=1) -> vortex.constant(i32?, len=1) + done output=vortex.constant(i32?, len=1) + execute_until target=AnyCanonical root=vortex.constant(i32?, len=1) + iter 0 current=vortex.constant(i32?, len=1) builder_active=false + Done array=vortex.primitive(i32?, len=1) + iter 1 current=vortex.primitive(i32?, len=1) builder_active=false + return output=vortex.primitive(i32?, len=1) + Done array=vortex.bool(bool, len=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 static:kernel[0] 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 + optimize root=vortex.cast(i16?, len=50) session=false + reduce_parent static:CastReduceAdaptor(Sequence) slot=0 parent=vortex.cast(i16?, len=50) child=vortex.sequence(i16, len=50) -> vortex.sequence(i16?, len=50) + done output=vortex.sequence(i16?, len=50) + 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) + optimize root=vortex.slice(i16, len=1) session=false + reduce_parent static:SliceReduceAdaptor(Constant) slot=0 parent=vortex.slice(i16, len=1) child=vortex.constant(i16, len=50) -> vortex.constant(i16, len=1) + done output=vortex.constant(i16, len=1) + optimize root=vortex.cast(i16?, len=1) session=false + reduce_parent static:CastReduceAdaptor(Constant) slot=0 parent=vortex.cast(i16?, len=1) child=vortex.constant(i16, len=1) -> vortex.constant(i16?, len=1) + done output=vortex.constant(i16?, len=1) + execute_until target=AnyCanonical root=vortex.constant(i16?, len=1) + iter 0 current=vortex.constant(i16?, len=1) builder_active=false + Done array=vortex.primitive(i16?, len=1) + iter 1 current=vortex.primitive(i16?, len=1) builder_active=false + return output=vortex.primitive(i16?, len=1) + 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 static:kernel[2] 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 static:kernel[1] 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 + execute_until target=AnyCanonical root=vortex.cast(i32, len=8) + iter 0 current=vortex.cast(i32, len=8) builder_active=false + child_execute_parent static:kernel[1] slot=0 parent=vortex.cast(i32, len=8) child=vortex.primitive(u8, len=8) -> vortex.primitive(i32, len=8) + iter 1 current=vortex.primitive(i32, len=8) builder_active=false + return output=vortex.primitive(i32, len=8) + child_execute_parent static:kernel[1] 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 static:kernel[2] 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 static:kernel[4] 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 + 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(()) +}