diff --git a/encodings/alp/src/alp_rd/array.rs b/encodings/alp/src/alp_rd/array.rs index 3bd7cdc6812..05364bbe006 100644 --- a/encodings/alp/src/alp_rd/array.rs +++ b/encodings/alp/src/alp_rd/array.rs @@ -16,7 +16,9 @@ use vortex_array::IntoArray; use vortex_array::Precision; use vortex_array::ProstMetadata; use vortex_array::SerializeMetadata; +use vortex_array::arrays::ConstantVTable; use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::PrimitiveVTable; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -41,7 +43,6 @@ use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_error::vortex_panic; -use vortex_mask::Mask; use vortex_session::VortexSession; use crate::alp_rd::kernel::PARENT_KERNELS; @@ -297,17 +298,27 @@ impl VTable for ALPRDVTable { } fn execute(array: &Self::Array, ctx: &mut ExecutionCtx) -> VortexResult { - let left_parts = array.left_parts().clone().execute::(ctx)?; - let right_parts = array.right_parts().clone().execute::(ctx)?; + // Ensure left_parts (child 0) is a PrimitiveArray. + let left_parts = if let Some(primitive) = array.left_parts().as_opt::() { + primitive.clone() + } else if array.left_parts().is::() { + array.left_parts().to_canonical()?.into_primitive() + } else { + return Ok(ExecutionStep::ExecuteChild(0)); + }; + + // Ensure right_parts (child 1) is a PrimitiveArray. + let right_parts = if let Some(primitive) = array.right_parts().as_opt::() { + primitive.clone() + } else if array.right_parts().is::() { + array.right_parts().to_canonical()?.into_primitive() + } else { + return Ok(ExecutionStep::ExecuteChild(1)); + }; // Decode the left_parts using our builtin dictionary. let left_parts_dict = array.left_parts_dictionary(); - - let validity = array - .left_parts() - .validity()? - .to_array(array.len()) - .execute::(ctx)?; + let validity = left_parts.validity_mask()?; let decoded_array = if array.is_f32() { PrimitiveArray::new( diff --git a/encodings/decimal-byte-parts/public-api.lock b/encodings/decimal-byte-parts/public-api.lock index 41a2a23a086..aab5e297326 100644 --- a/encodings/decimal-byte-parts/public-api.lock +++ b/encodings/decimal-byte-parts/public-api.lock @@ -108,7 +108,7 @@ pub fn vortex_decimal_byte_parts::DecimalBytePartsVTable::deserialize(bytes: &[u pub fn vortex_decimal_byte_parts::DecimalBytePartsVTable::dtype(array: &vortex_decimal_byte_parts::DecimalBytePartsArray) -> &vortex_array::dtype::DType -pub fn vortex_decimal_byte_parts::DecimalBytePartsVTable::execute(array: &Self::Array, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult +pub fn vortex_decimal_byte_parts::DecimalBytePartsVTable::execute(array: &Self::Array, _ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult pub fn vortex_decimal_byte_parts::DecimalBytePartsVTable::execute_parent(array: &Self::Array, parent: &vortex_array::array::ArrayRef, child_idx: usize, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult> diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs index 9bfbd9c21de..edbd78f90de 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -18,8 +18,9 @@ use vortex_array::IntoArray; use vortex_array::Precision; use vortex_array::ProstMetadata; use vortex_array::SerializeMetadata; +use vortex_array::arrays::ConstantVTable; use vortex_array::arrays::DecimalArray; -use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::PrimitiveVTable; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; @@ -190,8 +191,31 @@ impl VTable for DecimalBytePartsVTable { PARENT_RULES.evaluate(array, parent, child_idx) } - fn execute(array: &Self::Array, ctx: &mut ExecutionCtx) -> VortexResult { - to_canonical_decimal(array, ctx).map(ExecutionStep::Done) + fn execute(array: &Self::Array, _ctx: &mut ExecutionCtx) -> VortexResult { + // Ensure msp (child 0) is a PrimitiveArray. + let prim = if let Some(primitive) = array.msp.as_opt::() { + primitive.clone() + } else if array.msp.is::() { + array.msp.to_canonical()?.into_primitive() + } else { + return Ok(ExecutionStep::ExecuteChild(0)); + }; + + Ok(ExecutionStep::Done(match_each_signed_integer_ptype!( + prim.ptype(), + |P| { + // SAFETY: The primitive array's buffer is already validated with correct type. + // The decimal dtype matches the array's dtype, and validity is preserved. + unsafe { + DecimalArray::new_unchecked( + prim.to_buffer::

(), + *array.decimal_dtype(), + prim.validity().clone(), + ) + } + .into_array() + } + ))) } fn execute_parent( @@ -277,30 +301,6 @@ impl DecimalBytePartsVTable { pub const ID: ArrayId = ArrayId::new_ref("vortex.decimal_byte_parts"); } -/// Converts a DecimalBytePartsArray to its canonical DecimalArray representation. -fn to_canonical_decimal( - array: &DecimalBytePartsArray, - ctx: &mut ExecutionCtx, -) -> VortexResult { - // TODO(joe): support parts len != 1 - let prim = array.msp.clone().execute::(ctx)?; - // Depending on the decimal type and the min/max of the primitive array we can choose - // the correct buffer size - - Ok(match_each_signed_integer_ptype!(prim.ptype(), |P| { - // SAFETY: The primitive array's buffer is already validated with correct type. - // The decimal dtype matches the array's dtype, and validity is preserved. - unsafe { - DecimalArray::new_unchecked( - prim.to_buffer::

(), - *array.decimal_dtype(), - prim.validity().clone(), - ) - } - .into_array() - })) -} - impl OperationsVTable for DecimalBytePartsVTable { fn scalar_at(array: &DecimalBytePartsArray, index: usize) -> VortexResult { // TODO(joe): support parts len != 1 diff --git a/encodings/fastlanes/src/for/array/for_compress.rs b/encodings/fastlanes/src/for/array/for_compress.rs index 9990df5fb6a..161d2b04439 100644 --- a/encodings/fastlanes/src/for/array/for_compress.rs +++ b/encodings/fastlanes/src/for/array/for_compress.rs @@ -68,8 +68,9 @@ mod test { use super::*; use crate::BitPackedArray; - use crate::r#for::array::for_decompress::decompress; + use crate::r#for::array::for_decompress::apply_reference; use crate::r#for::array::for_decompress::fused_decompress; + use crate::r#for::array::for_decompress::try_fused_decompress; static SESSION: LazyLock = LazyLock::new(|| VortexSession::empty().with::()); @@ -169,7 +170,13 @@ mod test { let expected_unsigned = PrimitiveArray::from_iter(unsigned); assert_arrays_eq!(encoded, expected_unsigned); - let decompressed = decompress(&compressed, &mut SESSION.create_execution_ctx())?; + let mut ctx = SESSION.create_execution_ctx(); + let decompressed = if let Some(result) = try_fused_decompress(&compressed, &mut ctx)? { + result + } else { + let encoded = compressed.encoded().to_primitive(); + apply_reference(&compressed, encoded) + }; array .as_slice::() .iter() diff --git a/encodings/fastlanes/src/for/array/for_decompress.rs b/encodings/fastlanes/src/for/array/for_decompress.rs index f7292a5a06e..b462fa0646d 100644 --- a/encodings/fastlanes/src/for/array/for_decompress.rs +++ b/encodings/fastlanes/src/for/array/for_decompress.rs @@ -45,23 +45,28 @@ impl + FoR> UnpackStrategy for FoRStrategy } } -pub fn decompress(array: &FoRArray, ctx: &mut ExecutionCtx) -> VortexResult { - let ptype = array.ptype(); - - // Try to do fused unpack. +/// Try the fused BitPacked decompression path. Returns `None` if the child is not BitPacked +/// or the reference type is not unsigned. +pub fn try_fused_decompress( + array: &FoRArray, + ctx: &mut ExecutionCtx, +) -> VortexResult> { if array.reference_scalar().dtype().is_unsigned_int() && let Some(bp) = array.encoded().as_opt::() { return match_each_unsigned_integer_ptype!(array.ptype(), |T| { - fused_decompress::(array, bp, ctx) + fused_decompress::(array, bp, ctx).map(Some) }); } + Ok(None) +} - // TODO(ngates): Do we need this to be into_encoded() somehow? - let encoded = array.encoded().clone().execute::(ctx)?; +/// Apply the FoR reference value to an already-decoded PrimitiveArray. +pub fn apply_reference(array: &FoRArray, encoded: PrimitiveArray) -> PrimitiveArray { + let ptype = array.ptype(); let validity = encoded.validity().clone(); - Ok(match_each_integer_ptype!(ptype, |T| { + match_each_integer_ptype!(ptype, |T| { let min = array .reference_scalar() .as_primitive() @@ -75,7 +80,7 @@ pub fn decompress(array: &FoRArray, ctx: &mut ExecutionCtx) -> VortexResult VortexResult { - Ok(ExecutionStep::Done(decompress(array, ctx)?.into_array())) + // Try fused decompress with BitPacked child (no child execution needed). + if let Some(result) = try_fused_decompress(array, ctx)? { + return Ok(ExecutionStep::Done(result.into_array())); + } + + // If child is already a PrimitiveArray, add the reference value. + if array.encoded().is::() { + let encoded = array.encoded().as_::().clone(); + return Ok(ExecutionStep::Done( + apply_reference(array, encoded).into_array(), + )); + } + + // If child is a constant, compute the result as a constant. + if let Some(constant) = array.encoded().as_opt::() { + let scalar = constant.scalar(); + if scalar.is_null() { + return Ok(ExecutionStep::Done( + ConstantArray::new(Scalar::null(array.dtype().clone()), array.len()) + .into_array(), + )); + } + return Ok(ExecutionStep::Done(match_each_integer_ptype!( + array.ptype(), + |T| { + let enc_val = scalar + .as_primitive() + .typed_value::() + .vortex_expect("constant must be non-null after check"); + let ref_val = array + .reference_scalar() + .as_primitive() + .typed_value::() + .vortex_expect("reference must be non-null"); + ConstantArray::new( + Scalar::primitive( + enc_val.wrapping_add(ref_val), + scalar.dtype().nullability(), + ), + array.len(), + ) + .into_array() + } + ))); + } + + // Otherwise, ask the scheduler to execute the child first. + Ok(ExecutionStep::ExecuteChild(0)) } fn execute_parent( diff --git a/encodings/zigzag/public-api.lock b/encodings/zigzag/public-api.lock index 1f9fca6c735..45307c66270 100644 --- a/encodings/zigzag/public-api.lock +++ b/encodings/zigzag/public-api.lock @@ -96,7 +96,7 @@ pub fn vortex_zigzag::ZigZagVTable::deserialize(_bytes: &[u8], _dtype: &vortex_a pub fn vortex_zigzag::ZigZagVTable::dtype(array: &vortex_zigzag::ZigZagArray) -> &vortex_array::dtype::DType -pub fn vortex_zigzag::ZigZagVTable::execute(array: &Self::Array, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult +pub fn vortex_zigzag::ZigZagVTable::execute(array: &Self::Array, _ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult pub fn vortex_zigzag::ZigZagVTable::execute_parent(array: &Self::Array, parent: &vortex_array::array::ArrayRef, child_idx: usize, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult> diff --git a/encodings/zigzag/src/array.rs b/encodings/zigzag/src/array.rs index b7229031151..104ab79ca76 100644 --- a/encodings/zigzag/src/array.rs +++ b/encodings/zigzag/src/array.rs @@ -12,6 +12,9 @@ use vortex_array::ExecutionCtx; use vortex_array::ExecutionStep; use vortex_array::IntoArray; use vortex_array::Precision; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::ConstantVTable; +use vortex_array::arrays::PrimitiveVTable; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::PType; @@ -149,10 +152,39 @@ impl VTable for ZigZagVTable { Ok(()) } - fn execute(array: &Self::Array, ctx: &mut ExecutionCtx) -> VortexResult { - Ok(ExecutionStep::Done( - zigzag_decode(array.encoded().clone().execute(ctx)?).into_array(), - )) + fn execute(array: &Self::Array, _ctx: &mut ExecutionCtx) -> VortexResult { + // If child is already a PrimitiveArray, decode it directly. + if array.encoded().is::() { + let encoded = array.encoded().as_::().clone(); + return Ok(ExecutionStep::Done(zigzag_decode(encoded).into_array())); + } + + // If child is a constant, decode the scalar value. + if let Some(constant) = array.encoded().as_opt::() { + let scalar = constant.scalar(); + if scalar.is_null() { + return Ok(ExecutionStep::Done( + ConstantArray::new(Scalar::null(array.dtype().clone()), array.len()) + .into_array(), + )); + } + let result = match_each_unsigned_integer_ptype!(scalar.as_primitive().ptype(), |P| { + let val = scalar + .as_primitive() + .typed_value::

() + .vortex_expect("constant must be non-null after check"); + Scalar::primitive( + <

::Int>::decode(val), + array.dtype().nullability(), + ) + }); + return Ok(ExecutionStep::Done( + ConstantArray::new(result, array.len()).into_array(), + )); + } + + // Otherwise, ask the scheduler to execute the child first. + Ok(ExecutionStep::ExecuteChild(0)) } fn reduce_parent( diff --git a/encodings/zstd/public-api.lock b/encodings/zstd/public-api.lock index 9f351777a4c..c863b07ac90 100644 --- a/encodings/zstd/public-api.lock +++ b/encodings/zstd/public-api.lock @@ -196,7 +196,7 @@ pub fn vortex_zstd::ZstdVTable::deserialize(bytes: &[u8], _dtype: &vortex_array: pub fn vortex_zstd::ZstdVTable::dtype(array: &vortex_zstd::ZstdArray) -> &vortex_array::dtype::DType -pub fn vortex_zstd::ZstdVTable::execute(array: &Self::Array, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult +pub fn vortex_zstd::ZstdVTable::execute(array: &Self::Array, _ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult pub fn vortex_zstd::ZstdVTable::id(_array: &Self::Array) -> vortex_array::vtable::dyn_::ArrayId diff --git a/encodings/zstd/src/array.rs b/encodings/zstd/src/array.rs index f7263939e4d..9578ebcecb8 100644 --- a/encodings/zstd/src/array.rs +++ b/encodings/zstd/src/array.rs @@ -272,11 +272,8 @@ impl VTable for ZstdVTable { Ok(()) } - fn execute(array: &Self::Array, ctx: &mut ExecutionCtx) -> VortexResult { - array - .decompress()? - .execute::(ctx) - .map(ExecutionStep::Done) + fn execute(array: &Self::Array, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(ExecutionStep::Done(array.decompress()?)) } fn reduce_parent( diff --git a/encodings/zstd/src/zstd_buffers.rs b/encodings/zstd/src/zstd_buffers.rs index ff733475a95..3fd2e84ce3e 100644 --- a/encodings/zstd/src/zstd_buffers.rs +++ b/encodings/zstd/src/zstd_buffers.rs @@ -470,9 +470,7 @@ impl VTable for ZstdBuffersVTable { fn execute(array: &Self::Array, ctx: &mut ExecutionCtx) -> VortexResult { let session = ctx.session(); let inner_array = array.decompress_and_build_inner(session)?; - inner_array - .execute::(ctx) - .map(ExecutionStep::Done) + Ok(ExecutionStep::Done(inner_array)) } } diff --git a/vortex-array/public-api.lock b/vortex-array/public-api.lock index b885c08284c..a0e6c5cbb50 100644 --- a/vortex-array/public-api.lock +++ b/vortex-array/public-api.lock @@ -2978,7 +2978,7 @@ pub fn vortex_array::arrays::SharedVTable::deserialize(_bytes: &[u8], _dtype: &v pub fn vortex_array::arrays::SharedVTable::dtype(array: &vortex_array::arrays::SharedArray) -> &vortex_array::dtype::DType -pub fn vortex_array::arrays::SharedVTable::execute(array: &Self::Array, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult +pub fn vortex_array::arrays::SharedVTable::execute(array: &Self::Array, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult pub fn vortex_array::arrays::SharedVTable::execute_parent(array: &Self::Array, parent: &vortex_array::ArrayRef, child_idx: usize, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult> @@ -3138,7 +3138,7 @@ pub fn vortex_array::arrays::SliceVTable::deserialize(_bytes: &[u8], _dtype: &vo pub fn vortex_array::arrays::SliceVTable::dtype(array: &vortex_array::arrays::SliceArray) -> &vortex_array::dtype::DType -pub fn vortex_array::arrays::SliceVTable::execute(array: &Self::Array, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult +pub fn vortex_array::arrays::SliceVTable::execute(array: &Self::Array, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult pub fn vortex_array::arrays::SliceVTable::execute_parent(array: &Self::Array, parent: &vortex_array::ArrayRef, child_idx: usize, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult> @@ -17348,7 +17348,7 @@ pub fn vortex_array::arrays::SharedVTable::deserialize(_bytes: &[u8], _dtype: &v pub fn vortex_array::arrays::SharedVTable::dtype(array: &vortex_array::arrays::SharedArray) -> &vortex_array::dtype::DType -pub fn vortex_array::arrays::SharedVTable::execute(array: &Self::Array, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult +pub fn vortex_array::arrays::SharedVTable::execute(array: &Self::Array, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult pub fn vortex_array::arrays::SharedVTable::execute_parent(array: &Self::Array, parent: &vortex_array::ArrayRef, child_idx: usize, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult> @@ -17402,7 +17402,7 @@ pub fn vortex_array::arrays::SliceVTable::deserialize(_bytes: &[u8], _dtype: &vo pub fn vortex_array::arrays::SliceVTable::dtype(array: &vortex_array::arrays::SliceArray) -> &vortex_array::dtype::DType -pub fn vortex_array::arrays::SliceVTable::execute(array: &Self::Array, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult +pub fn vortex_array::arrays::SliceVTable::execute(array: &Self::Array, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult pub fn vortex_array::arrays::SliceVTable::execute_parent(array: &Self::Array, parent: &vortex_array::ArrayRef, child_idx: usize, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult> diff --git a/vortex-array/src/arrays/dict/vtable/mod.rs b/vortex-array/src/arrays/dict/vtable/mod.rs index 37d88aaf7ee..b8e0dc51c15 100644 --- a/vortex-array/src/arrays/dict/vtable/mod.rs +++ b/vortex-array/src/arrays/dict/vtable/mod.rs @@ -14,6 +14,7 @@ use vortex_session::VortexSession; use super::DictArray; use super::DictMetadata; use super::take_canonical; +use crate::AnyCanonical; use crate::Array; use crate::ArrayRef; use crate::Canonical; @@ -23,6 +24,8 @@ use crate::Precision; use crate::ProstMetadata; use crate::SerializeMetadata; use crate::arrays::ConstantArray; +use crate::arrays::ConstantVTable; +use crate::arrays::constant::constant_canonicalize; use crate::arrays::dict::compute::rules::PARENT_RULES; use crate::buffer::BufferHandle; use crate::dtype::DType; @@ -196,18 +199,24 @@ impl VTable for DictVTable { return Ok(ExecutionStep::Done(canonical)); } - // TODO(joe): if the values are constant return a constant - let values = array.values().clone().execute::(ctx)?; - let codes = array - .codes() - .clone() - .execute::(ctx)? - .into_primitive(); - - // TODO(ngates): if indices are sorted and unique (strict-sorted), then we should delegate to - // the filter function since they're typically optimised for this case. - // TODO(ngates): if indices min is quite high, we could slice self and offset the indices - // such that canonicalize does less work. + // Ensure codes (child 0) are in canonical form. + let codes = if let Some(canonical) = array.codes().as_opt::() { + Canonical::from(canonical) + } else if let Some(constant) = array.codes().as_opt::() { + constant_canonicalize(constant)? + } else { + return Ok(ExecutionStep::ExecuteChild(0)); + }; + let codes = codes.into_primitive(); + + // Ensure values (child 1) are in canonical form. + let values = if let Some(canonical) = array.values().as_opt::() { + Canonical::from(canonical) + } else if let Some(constant) = array.values().as_opt::() { + constant_canonicalize(constant)? + } else { + return Ok(ExecutionStep::ExecuteChild(1)); + }; Ok(ExecutionStep::Done( take_canonical(values, &codes, ctx)?.into_array(), diff --git a/vortex-array/src/arrays/filter/vtable.rs b/vortex-array/src/arrays/filter/vtable.rs index fff8c90b992..ef329007ecf 100644 --- a/vortex-array/src/arrays/filter/vtable.rs +++ b/vortex-array/src/arrays/filter/vtable.rs @@ -13,12 +13,16 @@ use vortex_error::vortex_panic; use vortex_mask::Mask; use vortex_session::VortexSession; +use crate::AnyCanonical; use crate::Array; use crate::ArrayEq; use crate::ArrayHash; use crate::ArrayRef; +use crate::Canonical; use crate::IntoArray; use crate::Precision; +use crate::arrays::ConstantArray; +use crate::arrays::ConstantVTable; use crate::arrays::filter::array::FilterArray; use crate::arrays::filter::execute::execute_filter; use crate::arrays::filter::execute::execute_filter_fast_paths; @@ -156,18 +160,30 @@ impl VTable for FilterVTable { } fn execute(array: &Self::Array, ctx: &mut ExecutionCtx) -> VortexResult { - if let Some(canonical) = execute_filter_fast_paths(array, ctx)? { - return Ok(ExecutionStep::Done(canonical)); + if let Some(result) = execute_filter_fast_paths(array, ctx)? { + return Ok(ExecutionStep::Done(result)); } let Mask::Values(mask_values) = &array.mask else { unreachable!("`execute_filter_fast_paths` handles AllTrue and AllFalse") }; - // We rely on the optimization pass that runs prior to this execution for filter pushdown, - // so now we can just execute the filter without worrying. - Ok(ExecutionStep::Done( - execute_filter(array.child.clone().execute(ctx)?, mask_values).into_array(), - )) + // If child is already canonical, filter it directly. + if let Some(canonical) = array.child.as_opt::() { + return Ok(ExecutionStep::Done( + execute_filter(Canonical::from(canonical), mask_values).into_array(), + )); + } + + // If child is a constant, filtering just changes the length. + if let Some(constant) = array.child.as_opt::() { + return Ok(ExecutionStep::Done( + ConstantArray::new(constant.scalar().clone(), mask_values.true_count()) + .into_array(), + )); + } + + // Otherwise, ask the scheduler to execute the child first. + Ok(ExecutionStep::ExecuteChild(0)) } fn reduce_parent( diff --git a/vortex-array/src/arrays/masked/vtable/mod.rs b/vortex-array/src/arrays/masked/vtable/mod.rs index c63a0687196..b2a8099bf1b 100644 --- a/vortex-array/src/arrays/masked/vtable/mod.rs +++ b/vortex-array/src/arrays/masked/vtable/mod.rs @@ -13,12 +13,15 @@ use vortex_error::vortex_ensure; use vortex_error::vortex_panic; use vortex_session::VortexSession; +use crate::AnyCanonical; use crate::ArrayRef; use crate::Canonical; use crate::EmptyMetadata; use crate::IntoArray; use crate::Precision; use crate::arrays::ConstantArray; +use crate::arrays::ConstantVTable; +use crate::arrays::constant::constant_canonicalize; use crate::arrays::masked::MaskedArray; use crate::arrays::masked::compute::rules::PARENT_RULES; use crate::arrays::masked::mask_validity_canonical; @@ -178,10 +181,24 @@ impl VTable for MaskedVTable { // While we could manually convert the dtype, `mask_validity_canonical` is already O(1) for // `AllTrue` masks (no data copying), so there's no benefit. - let child = array.child().clone().execute::(ctx)?; - Ok(ExecutionStep::Done( - mask_validity_canonical(child, &validity_mask, ctx)?.into_array(), - )) + // If child is already canonical, apply the validity mask directly. + if let Some(canonical) = array.child().as_opt::() { + return Ok(ExecutionStep::Done( + mask_validity_canonical(Canonical::from(canonical), &validity_mask, ctx)? + .into_array(), + )); + } + + // If child is a constant, expand to canonical then apply the validity mask. + if let Some(constant) = array.child().as_opt::() { + let canonical = constant_canonicalize(constant)?; + return Ok(ExecutionStep::Done( + mask_validity_canonical(canonical, &validity_mask, ctx)?.into_array(), + )); + } + + // Otherwise, ask the scheduler to execute the child first. + Ok(ExecutionStep::ExecuteChild(0)) } fn reduce_parent( diff --git a/vortex-array/src/arrays/shared/vtable.rs b/vortex-array/src/arrays/shared/vtable.rs index 9a7ca38b263..8e1829511b1 100644 --- a/vortex-array/src/arrays/shared/vtable.rs +++ b/vortex-array/src/arrays/shared/vtable.rs @@ -8,12 +8,15 @@ use vortex_error::VortexResult; use vortex_error::vortex_panic; use vortex_session::VortexSession; +use crate::AnyCanonical; use crate::ArrayRef; use crate::Canonical; use crate::EmptyMetadata; use crate::ExecutionCtx; use crate::ExecutionStep; +use crate::IntoArray; use crate::Precision; +use crate::arrays::ConstantVTable; use crate::arrays::shared::SharedArray; use crate::buffer::BufferHandle; use crate::dtype::DType; @@ -145,10 +148,15 @@ impl VTable for SharedVTable { Ok(()) } - fn execute(array: &Self::Array, ctx: &mut ExecutionCtx) -> VortexResult { - array - .get_or_compute(|source| source.clone().execute::(ctx)) - .map(ExecutionStep::Done) + fn execute(array: &Self::Array, _ctx: &mut ExecutionCtx) -> VortexResult { + let current = array.current_array_ref(); + if let Some(canonical) = current.as_opt::() { + return Ok(ExecutionStep::Done(Canonical::from(canonical).into_array())); + } + if current.as_opt::().is_some() { + return Ok(ExecutionStep::Done(current.clone())); + } + Ok(ExecutionStep::ExecuteChild(0)) } } impl OperationsVTable for SharedVTable { diff --git a/vortex-array/src/arrays/slice/vtable.rs b/vortex-array/src/arrays/slice/vtable.rs index 26cc932a8ea..09fb50caf88 100644 --- a/vortex-array/src/arrays/slice/vtable.rs +++ b/vortex-array/src/arrays/slice/vtable.rs @@ -20,7 +20,10 @@ use crate::ArrayEq; use crate::ArrayHash; use crate::ArrayRef; use crate::Canonical; +use crate::IntoArray; use crate::Precision; +use crate::arrays::ConstantArray; +use crate::arrays::ConstantVTable; use crate::arrays::slice::array::SliceArray; use crate::arrays::slice::rules::PARENT_RULES; use crate::buffer::BufferHandle; @@ -155,23 +158,25 @@ impl VTable for SliceVTable { Ok(()) } - fn execute(array: &Self::Array, ctx: &mut ExecutionCtx) -> VortexResult { - // Execute the child to get canonical form, then slice it - let Some(canonical) = array.child.as_opt::() else { - // If the child is not canonical, recurse. - return array - .child - .clone() - .execute::(ctx)? - .slice(array.slice_range().clone()) + fn execute(array: &Self::Array, _ctx: &mut ExecutionCtx) -> VortexResult { + // If child is already canonical, slice it directly. + if let Some(canonical) = array.child.as_opt::() { + // TODO(ngates): we should inline canonical slice logic here. + return Canonical::from(canonical) + .as_ref() + .slice(array.range.clone()) .map(ExecutionStep::Done); - }; + } + + // If child is a constant, slicing just changes the length. + if let Some(constant) = array.child.as_opt::() { + return Ok(ExecutionStep::Done( + ConstantArray::new(constant.scalar().clone(), array.range.len()).into_array(), + )); + } - // TODO(ngates): we should inline canonical slice logic here. - Canonical::from(canonical) - .as_ref() - .slice(array.range.clone()) - .map(ExecutionStep::Done) + // Otherwise, ask the scheduler to execute the child first. + Ok(ExecutionStep::ExecuteChild(0)) } fn reduce_parent(