diff --git a/encodings/alp/src/alp/ops.rs b/encodings/alp/src/alp/ops.rs index dcd8cecf1d5..acb40ba9422 100644 --- a/encodings/alp/src/alp/ops.rs +++ b/encodings/alp/src/alp/ops.rs @@ -18,7 +18,7 @@ impl OperationsVTable for ALP { fn scalar_at( array: ArrayView<'_, ALP>, index: usize, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { if let Some(patches) = array.patches() && let Some(patch) = patches.get_patched(index)? @@ -26,7 +26,7 @@ impl OperationsVTable for ALP { return patch.cast(array.dtype()); } - let encoded_val = array.encoded().scalar_at(index)?; + let encoded_val = array.encoded().scalar_at(index, ctx)?; Ok(match_each_alp_float_ptype!(array.dtype().as_ptype(), |T| { let encoded_val: ::ALPInt = diff --git a/encodings/alp/src/alp_rd/ops.rs b/encodings/alp/src/alp_rd/ops.rs index 90f1307f59f..0bac27d157a 100644 --- a/encodings/alp/src/alp_rd/ops.rs +++ b/encodings/alp/src/alp_rd/ops.rs @@ -16,7 +16,7 @@ impl OperationsVTable for ALPRD { fn scalar_at( array: ArrayView<'_, ALPRD>, index: usize, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { // The left value can either be a direct value, or an exception. // The exceptions array represents exception positions with non-null values. @@ -32,7 +32,7 @@ impl OperationsVTable for ALPRD { _ => { let left_code: u16 = array .left_parts() - .scalar_at(index)? + .scalar_at(index, ctx)? .as_primitive() .as_::() .vortex_expect("left_code must be non-null"); @@ -44,7 +44,7 @@ impl OperationsVTable for ALPRD { Ok(if array.dtype().as_ptype() == PType::F32 { let right: u32 = array .right_parts() - .scalar_at(index)? + .scalar_at(index, ctx)? .as_primitive() .as_::() .vortex_expect("non-null"); @@ -53,7 +53,7 @@ impl OperationsVTable for ALPRD { } else { let right: u64 = array .right_parts() - .scalar_at(index)? + .scalar_at(index, ctx)? .as_primitive() .as_::() .vortex_expect("non-null"); diff --git a/encodings/datetime-parts/src/compute/is_constant.rs b/encodings/datetime-parts/src/compute/is_constant.rs index 94e6806e785..d84e294f7aa 100644 --- a/encodings/datetime-parts/src/compute/is_constant.rs +++ b/encodings/datetime-parts/src/compute/is_constant.rs @@ -37,6 +37,6 @@ impl DynAggregateKernel for DateTimePartsIsConstantKernel { let result = is_constant(array.days(), ctx)? && is_constant(array.seconds(), ctx)? && is_constant(array.subseconds(), ctx)?; - Ok(Some(IsConstant::make_partial(batch, result)?)) + Ok(Some(IsConstant::make_partial(batch, result, ctx)?)) } } diff --git a/encodings/datetime-parts/src/ops.rs b/encodings/datetime-parts/src/ops.rs index 2c224b078d5..edcec365f01 100644 --- a/encodings/datetime-parts/src/ops.rs +++ b/encodings/datetime-parts/src/ops.rs @@ -20,7 +20,7 @@ impl OperationsVTable for DateTimeParts { fn scalar_at( array: ArrayView<'_, DateTimeParts>, index: usize, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { let DType::Extension(ext) = array.dtype().clone() else { vortex_panic!( @@ -33,25 +33,25 @@ impl OperationsVTable for DateTimeParts { vortex_panic!(Compute: "must decode TemporalMetadata from extension metadata"); }; - if !array.as_ref().is_valid(index)? { + if !array.as_ref().is_valid(index, ctx)? { return Ok(Scalar::null(DType::Extension(ext))); } let days: i64 = array .days() - .scalar_at(index)? + .scalar_at(index, ctx)? .as_primitive() .as_::() .vortex_expect("days fits in i64"); let seconds: i64 = array .seconds() - .scalar_at(index)? + .scalar_at(index, ctx)? .as_primitive() .as_::() .vortex_expect("seconds fits in i64"); let subseconds: i64 = array .subseconds() - .scalar_at(index)? + .scalar_at(index, ctx)? .as_primitive() .as_::() .vortex_expect("subseconds fits in i64"); diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/is_constant.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/is_constant.rs index 4d271eb80b3..f4da528a9b3 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/is_constant.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/is_constant.rs @@ -35,6 +35,6 @@ impl DynAggregateKernel for DecimalBytePartsIsConstantKernel { }; let result = is_constant(array.msp(), ctx)?; - Ok(Some(IsConstant::make_partial(batch, result)?)) + Ok(Some(IsConstant::make_partial(batch, result, ctx)?)) } } 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 ea90bf575e0..bccc767657d 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -303,10 +303,10 @@ impl OperationsVTable for DecimalByteParts { fn scalar_at( array: ArrayView<'_, DecimalByteParts>, index: usize, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { // TODO(joe): support parts len != 1 - let scalar = array.msp().scalar_at(index)?; + let scalar = array.msp().scalar_at(index, ctx)?; // Note. values in msp, can only be signed integers upto size i64. let primitive_scalar = scalar.as_primitive(); diff --git a/encodings/fastlanes/src/bitpacking/compute/is_constant.rs b/encodings/fastlanes/src/bitpacking/compute/is_constant.rs index 900284e926d..94f3feea46f 100644 --- a/encodings/fastlanes/src/bitpacking/compute/is_constant.rs +++ b/encodings/fastlanes/src/bitpacking/compute/is_constant.rs @@ -34,7 +34,7 @@ impl DynAggregateKernel for BitPackedIsConstantKernel { &self, aggregate_fn: &AggregateFnRef, batch: &ArrayRef, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult> { if !aggregate_fn.is::() { return Ok(None); @@ -48,7 +48,7 @@ impl DynAggregateKernel for BitPackedIsConstantKernel { bitpacked_is_constant::() }>(array)? }); - Ok(Some(IsConstant::make_partial(batch, result)?)) + Ok(Some(IsConstant::make_partial(batch, result, ctx)?)) } } diff --git a/encodings/fastlanes/src/delta/vtable/operations.rs b/encodings/fastlanes/src/delta/vtable/operations.rs index 731588a67bb..9b387ec7c83 100644 --- a/encodings/fastlanes/src/delta/vtable/operations.rs +++ b/encodings/fastlanes/src/delta/vtable/operations.rs @@ -14,10 +14,10 @@ impl OperationsVTable for Delta { fn scalar_at( array: ArrayView<'_, Delta>, index: usize, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { let decompressed = array.array().slice(index..index + 1)?.to_primitive(); - decompressed.into_array().scalar_at(0) + decompressed.into_array().scalar_at(0, ctx) } } diff --git a/encodings/fastlanes/src/for/compute/is_constant.rs b/encodings/fastlanes/src/for/compute/is_constant.rs index 246a9c93226..29393884e95 100644 --- a/encodings/fastlanes/src/for/compute/is_constant.rs +++ b/encodings/fastlanes/src/for/compute/is_constant.rs @@ -35,6 +35,6 @@ impl DynAggregateKernel for FoRIsConstantKernel { }; let result = is_constant(array.encoded(), ctx)?; - Ok(Some(IsConstant::make_partial(batch, result)?)) + Ok(Some(IsConstant::make_partial(batch, result, ctx)?)) } } diff --git a/encodings/fastlanes/src/for/compute/is_sorted.rs b/encodings/fastlanes/src/for/compute/is_sorted.rs index 3b81c45c03d..92cd8b36888 100644 --- a/encodings/fastlanes/src/for/compute/is_sorted.rs +++ b/encodings/fastlanes/src/for/compute/is_sorted.rs @@ -49,7 +49,12 @@ impl DynAggregateKernel for FoRIsSortedKernel { is_sorted(&unsigned_array, ctx)? }; - Ok(Some(IsSorted::make_partial(batch, result, options.strict)?)) + Ok(Some(IsSorted::make_partial( + batch, + result, + options.strict, + ctx, + )?)) } } diff --git a/encodings/fastlanes/src/for/vtable/operations.rs b/encodings/fastlanes/src/for/vtable/operations.rs index e90ff7f4ad2..de6c07e7e47 100644 --- a/encodings/fastlanes/src/for/vtable/operations.rs +++ b/encodings/fastlanes/src/for/vtable/operations.rs @@ -15,9 +15,9 @@ impl OperationsVTable for FoR { fn scalar_at( array: ArrayView<'_, FoR>, index: usize, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { - let encoded_pvalue = array.encoded().scalar_at(index)?; + let encoded_pvalue = array.encoded().scalar_at(index, ctx)?; let encoded_pvalue = encoded_pvalue.as_primitive(); let reference = array.reference_scalar(); let reference = reference.as_primitive(); diff --git a/encodings/fastlanes/src/rle/array/mod.rs b/encodings/fastlanes/src/rle/array/mod.rs index a44ae69ddbd..d23a0e21409 100644 --- a/encodings/fastlanes/src/rle/array/mod.rs +++ b/encodings/fastlanes/src/rle/array/mod.rs @@ -5,7 +5,9 @@ use std::fmt::Display; use std::fmt::Formatter; use vortex_array::ArrayRef; +use vortex_array::LEGACY_SESSION; use vortex_array::TypedArrayRef; +use vortex_array::VortexSessionExecute as _; use vortex_error::VortexExpect as _; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -110,15 +112,17 @@ pub trait RLEArrayExt: TypedArrayRef { reason = "expect is safe here as scalar_at returns a valid primitive" )] fn values_idx_offset(&self, chunk_idx: usize) -> usize { + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + self.values_idx_offsets() - .scalar_at(chunk_idx) + .scalar_at(chunk_idx, &mut ctx) .expect("index must be in bounds") .as_primitive() .as_::() .expect("index must be of type usize") - self .values_idx_offsets() - .scalar_at(0) + .scalar_at(0, &mut ctx) .expect("index must be in bounds") .as_primitive() .as_::() diff --git a/encodings/fastlanes/src/rle/vtable/operations.rs b/encodings/fastlanes/src/rle/vtable/operations.rs index 4d54cd9cf32..c125a433c69 100644 --- a/encodings/fastlanes/src/rle/vtable/operations.rs +++ b/encodings/fastlanes/src/rle/vtable/operations.rs @@ -16,10 +16,10 @@ impl OperationsVTable for RLE { fn scalar_at( array: ArrayView<'_, RLE>, index: usize, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { let offset_in_chunk = array.offset(); - let chunk_relative_idx = array.indices().scalar_at(offset_in_chunk + index)?; + let chunk_relative_idx = array.indices().scalar_at(offset_in_chunk + index, ctx)?; let chunk_relative_idx = chunk_relative_idx .as_primitive() @@ -31,7 +31,7 @@ impl OperationsVTable for RLE { let scalar = array .values() - .scalar_at(value_idx_offset + chunk_relative_idx)?; + .scalar_at(value_idx_offset + chunk_relative_idx, ctx)?; Scalar::try_new(array.dtype().clone(), scalar.into_value()) } diff --git a/encodings/fsst/src/array.rs b/encodings/fsst/src/array.rs index 5526276aeb4..f069eddfd56 100644 --- a/encodings/fsst/src/array.rs +++ b/encodings/fsst/src/array.rs @@ -23,8 +23,10 @@ use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; use vortex_array::IntoArray; +use vortex_array::LEGACY_SESSION; use vortex_array::Precision; use vortex_array::TypedArrayRef; +use vortex_array::VortexSessionExecute as _; use vortex_array::arrays::VarBin; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::varbin::VarBinArrayExt; @@ -570,10 +572,12 @@ impl FSSTData { vortex_bail!(InvalidArgument: "codes nullability must match outer dtype nullability"); } + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + // Validate that last offset doesn't exceed bytes length (when host-resident). if codes_bytes.is_on_host() && codes_offsets.is_host() && !codes_offsets.is_empty() { let last_offset: usize = (&codes_offsets - .scalar_at(codes_offsets.len() - 1) + .scalar_at(codes_offsets.len() - 1, &mut ctx) .vortex_expect("offsets must support scalar_at")) .try_into() .vortex_expect("Failed to convert offset to usize"); diff --git a/encodings/fsst/src/ops.rs b/encodings/fsst/src/ops.rs index 86172377f7e..afd13704dd4 100644 --- a/encodings/fsst/src/ops.rs +++ b/encodings/fsst/src/ops.rs @@ -17,9 +17,9 @@ impl OperationsVTable for FSST { fn scalar_at( array: ArrayView<'_, FSST>, index: usize, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { - let compressed = array.codes().scalar_at(index)?; + let compressed = array.codes().scalar_at(index, ctx)?; let binary_datum = compressed.as_binary().value().vortex_expect("non-null"); let decoded_buffer = ByteBuffer::from(array.decompressor().decompress(binary_datum)); diff --git a/encodings/parquet-variant/src/operations.rs b/encodings/parquet-variant/src/operations.rs index 2efad44ae6e..f2d057ba302 100644 --- a/encodings/parquet-variant/src/operations.rs +++ b/encodings/parquet-variant/src/operations.rs @@ -8,6 +8,8 @@ use parquet_variant::Variant as PqVariant; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; +use vortex_array::LEGACY_SESSION; +use vortex_array::VortexSessionExecute as _; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; use vortex_array::dtype::FieldName; @@ -41,15 +43,15 @@ impl OperationsVTable for ParquetVariant { fn scalar_at( array: ArrayView<'_, ParquetVariant>, index: usize, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { - if array.validity()?.is_null(index)? { + if array.validity()?.is_null(index, ctx)? { return Ok(Scalar::null(DType::Variant(Nullability::Nullable))); } let metadata = array .metadata_array() - .scalar_at(index)? + .scalar_at(index, ctx)? .as_binary() .value() .cloned() @@ -74,16 +76,17 @@ fn scalar_from_variant_storage( typed_value: Option<&ArrayRef>, index: usize, ) -> VortexResult { + let mut ctx = LEGACY_SESSION.create_execution_ctx(); if let Some(typed_value) = typed_value - && typed_value.is_valid(index)? + && typed_value.is_valid(index, &mut ctx)? { return scalar_from_typed_value_array(metadata, value, typed_value, index); } if let Some(value) = value - && value.is_valid(index)? + && value.is_valid(index, &mut ctx)? { - return scalar_from_unshredded_value(metadata, &value.scalar_at(index)?); + return scalar_from_unshredded_value(metadata, &value.scalar_at(index, &mut ctx)?); } Ok(Scalar::null(DType::Null)) @@ -95,11 +98,17 @@ fn scalar_from_typed_value_array( typed_value: &ArrayRef, index: usize, ) -> VortexResult { + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let value_scalar = match value { - Some(value) if value.is_valid(index)? => Some(value.scalar_at(index)?), + Some(value) if value.is_valid(index, &mut ctx)? => Some(value.scalar_at(index, &mut ctx)?), _ => None, }; - scalar_from_typed_value_scalar(metadata, value_scalar, typed_value.scalar_at(index)?) + scalar_from_typed_value_scalar( + metadata, + value_scalar, + typed_value.scalar_at(index, &mut ctx)?, + ) } fn scalar_from_typed_value_scalar( diff --git a/encodings/pco/src/array.rs b/encodings/pco/src/array.rs index 8f83dbd8fa8..b77ef0873af 100644 --- a/encodings/pco/src/array.rs +++ b/encodings/pco/src/array.rs @@ -28,10 +28,8 @@ use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; use vortex_array::Precision; use vortex_array::ToCanonical; -use vortex_array::VortexSessionExecute; use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; use vortex_array::buffer::BufferHandle; @@ -629,15 +627,14 @@ impl OperationsVTable for Pco { fn scalar_at( array: ArrayView<'_, Pco>, index: usize, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); let unsliced_validity = child_to_validity(&array.slots()[0], array.dtype().nullability()); array ._slice(index, index + 1) - .decompress(&unsliced_validity, &mut ctx)? + .decompress(&unsliced_validity, ctx)? .into_array() - .scalar_at(0) + .scalar_at(0, ctx) } } diff --git a/encodings/runend/src/array.rs b/encodings/runend/src/array.rs index f714d09fc70..f87db639375 100644 --- a/encodings/runend/src/array.rs +++ b/encodings/runend/src/array.rs @@ -18,8 +18,10 @@ use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; use vortex_array::IntoArray; +use vortex_array::LEGACY_SESSION; use vortex_array::Precision; use vortex_array::TypedArrayRef; +use vortex_array::VortexSessionExecute as _; use vortex_array::arrays::Primitive; use vortex_array::arrays::VarBinViewArray; use vortex_array::buffer::BufferHandle; @@ -313,7 +315,8 @@ impl RunEndData { if ends.is_empty() { Ok(0) } else { - usize::try_from(&ends.scalar_at(ends.len() - 1)?) + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + usize::try_from(&ends.scalar_at(ends.len() - 1, &mut ctx)?) } } @@ -370,15 +373,17 @@ impl RunEndData { return Ok(()); } + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + // Validate the offset and length are valid for the given ends and values if offset != 0 && length != 0 { - let first_run_end = usize::try_from(&ends.scalar_at(0)?)?; + let first_run_end = usize::try_from(&ends.scalar_at(0, &mut ctx)?)?; if first_run_end < offset { vortex_bail!("First run end {first_run_end} must be >= offset {offset}"); } } - let last_run_end = usize::try_from(&ends.scalar_at(ends.len() - 1)?)?; + let last_run_end = usize::try_from(&ends.scalar_at(ends.len() - 1, &mut ctx)?)?; let min_required_end = offset + length; if last_run_end < min_required_end { vortex_bail!("Last run end {last_run_end} must be >= offset+length {min_required_end}"); diff --git a/encodings/runend/src/compute/is_constant.rs b/encodings/runend/src/compute/is_constant.rs index 53418f3ef50..d21f2334018 100644 --- a/encodings/runend/src/compute/is_constant.rs +++ b/encodings/runend/src/compute/is_constant.rs @@ -35,6 +35,6 @@ impl DynAggregateKernel for RunEndIsConstantKernel { }; let result = is_constant(array.values(), ctx)?; - Ok(Some(IsConstant::make_partial(batch, result)?)) + Ok(Some(IsConstant::make_partial(batch, result, ctx)?)) } } diff --git a/encodings/runend/src/compute/is_sorted.rs b/encodings/runend/src/compute/is_sorted.rs index e26a3d8cb57..85a86990102 100644 --- a/encodings/runend/src/compute/is_sorted.rs +++ b/encodings/runend/src/compute/is_sorted.rs @@ -45,6 +45,11 @@ impl DynAggregateKernel for RunEndIsSortedKernel { is_sorted(array.values(), ctx)? }; - Ok(Some(IsSorted::make_partial(batch, result, options.strict)?)) + Ok(Some(IsSorted::make_partial( + batch, + result, + options.strict, + ctx, + )?)) } } diff --git a/encodings/runend/src/kernel.rs b/encodings/runend/src/kernel.rs index 432453ee609..daad05e4353 100644 --- a/encodings/runend/src/kernel.rs +++ b/encodings/runend/src/kernel.rs @@ -7,6 +7,8 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::LEGACY_SESSION; +use vortex_array::VortexSessionExecute as _; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::Slice; use vortex_array::arrays::dict::TakeExecuteAdaptor; @@ -55,9 +57,11 @@ fn slice(array: ArrayView<'_, RunEnd>, range: Range) -> VortexResult for RunEnd { fn scalar_at( array: ArrayView<'_, RunEnd>, index: usize, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { - array.values().scalar_at(array.find_physical_index(index)?) + array + .values() + .scalar_at(array.find_physical_index(index)?, ctx) } } diff --git a/encodings/sequence/src/compute/is_sorted.rs b/encodings/sequence/src/compute/is_sorted.rs index c930246508f..7df01aded17 100644 --- a/encodings/sequence/src/compute/is_sorted.rs +++ b/encodings/sequence/src/compute/is_sorted.rs @@ -27,7 +27,6 @@ impl DynAggregateKernel for SequenceIsSortedKernel { batch: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { - let _ = ctx; let Some(options) = aggregate_fn.as_opt::() else { return Ok(None); }; @@ -47,6 +46,11 @@ impl DynAggregateKernel for SequenceIsSortedKernel { }) })?; - Ok(Some(IsSorted::make_partial(batch, result, options.strict)?)) + Ok(Some(IsSorted::make_partial( + batch, + result, + options.strict, + ctx, + )?)) } } diff --git a/encodings/sparse/src/canonical.rs b/encodings/sparse/src/canonical.rs index c53fd8989e8..793a5f5cb09 100644 --- a/encodings/sparse/src/canonical.rs +++ b/encodings/sparse/src/canonical.rs @@ -8,6 +8,8 @@ use num_traits::NumCast; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::LEGACY_SESSION; +use vortex_array::VortexSessionExecute as _; use vortex_array::arrays::BoolArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::ListViewArray; @@ -185,6 +187,8 @@ fn execute_sparse_lists_inner( let mut patch_idx = 0; + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + // Loop over the patch indices and set them to the corresponding scalar values. For positions // that are not patched, use the fill value. for position in 0..len { @@ -199,7 +203,7 @@ fn execute_sparse_lists_inner( builder .append_value( patch_values - .scalar_at(patch_idx) + .scalar_at(patch_idx, &mut ctx) .vortex_expect("scalar_at") .as_list(), ) @@ -271,6 +275,8 @@ fn execute_sparse_fixed_size_list_inner( .iter() .map(|x| (*x).to_usize().vortex_expect("index must fit in usize")); + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + for (patch_idx, sparse_idx) in indices.enumerate() { // Fill gap before this patch with fill values. append_n_lists( @@ -284,7 +290,7 @@ fn execute_sparse_fixed_size_list_inner( if values .validity() .vortex_expect("sparse fixed-size-list validity should be derivable") - .is_valid(patch_idx) + .is_valid(patch_idx, &mut ctx) .vortex_expect("is_valid") { let patch_list = values @@ -292,7 +298,7 @@ fn execute_sparse_fixed_size_list_inner( .vortex_expect("fixed_size_list_elements_at"); for i in 0..list_size as usize { builder - .append_scalar(&patch_list.scalar_at(i).vortex_expect("scalar_at")) + .append_scalar(&patch_list.scalar_at(i, &mut ctx).vortex_expect("scalar_at")) .vortex_expect("element dtype must match"); } } else { diff --git a/encodings/zigzag/src/array.rs b/encodings/zigzag/src/array.rs index 851b40b438b..c33d6105e96 100644 --- a/encodings/zigzag/src/array.rs +++ b/encodings/zigzag/src/array.rs @@ -240,9 +240,9 @@ impl OperationsVTable for ZigZag { fn scalar_at( array: ArrayView<'_, ZigZag>, index: usize, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { - let scalar = array.encoded().scalar_at(index)?; + let scalar = array.encoded().scalar_at(index, ctx)?; if scalar.is_null() { return scalar.primitive_reinterpret_cast(ZigZagArrayExt::ptype(&array)); } diff --git a/encodings/zstd/src/array.rs b/encodings/zstd/src/array.rs index a11d48defb6..c032ccdcd46 100644 --- a/encodings/zstd/src/array.rs +++ b/encodings/zstd/src/array.rs @@ -21,10 +21,8 @@ use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; use vortex_array::Precision; use vortex_array::ToCanonical; -use vortex_array::VortexSessionExecute; use vortex_array::accessor::ArrayAccessor; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; @@ -1023,14 +1021,13 @@ impl OperationsVTable for Zstd { fn scalar_at( array: ArrayView<'_, Zstd>, index: usize, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); let unsliced_validity = child_to_validity(&array.slots()[0], array.dtype().nullability()); let sliced = array.data().with_slice(index, index + 1); sliced - .decompress(array.dtype(), &unsliced_validity, &mut ctx)? - .scalar_at(0) + .decompress(array.dtype(), &unsliced_validity, ctx)? + .scalar_at(0, ctx) } } diff --git a/fuzz/src/array/compare.rs b/fuzz/src/array/compare.rs index 849a7fdd694..a45260ede95 100644 --- a/fuzz/src/array/compare.rs +++ b/fuzz/src/array/compare.rs @@ -3,7 +3,9 @@ use vortex_array::ArrayRef; use vortex_array::IntoArray; +use vortex_array::LEGACY_SESSION; use vortex_array::ToCanonical; +use vortex_array::VortexSessionExecute as _; use vortex_array::accessor::ArrayAccessor; use vortex_array::arrays::BoolArray; use vortex_array::arrays::bool::BoolArrayExt; @@ -131,8 +133,9 @@ pub fn compare_canonical_array( ) }), DType::Struct(..) | DType::List(..) | DType::FixedSizeList(..) => { + let mut ctx = LEGACY_SESSION.create_execution_ctx(); let scalar_vals: Vec = (0..array.len()) - .map(|i| array.scalar_at(i).vortex_expect("scalar_at")) + .map(|i| array.scalar_at(i, &mut ctx).vortex_expect("scalar_at")) .collect(); BoolArray::from_iter(scalar_vals.iter().map(|v| { scalar_cmp(v, value, operator) diff --git a/fuzz/src/array/fill_null.rs b/fuzz/src/array/fill_null.rs index 21430ce2b4d..b8d4e027093 100644 --- a/fuzz/src/array/fill_null.rs +++ b/fuzz/src/array/fill_null.rs @@ -6,7 +6,9 @@ use std::sync::Arc; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray; +use vortex_array::LEGACY_SESSION; use vortex_array::ToCanonical; +use vortex_array::VortexSessionExecute as _; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::DecimalArray; @@ -199,6 +201,7 @@ fn fill_varbinview_array( match array.dtype() { DType::Utf8(_) => { + let mut ctx = LEGACY_SESSION.create_execution_ctx(); let fill_str = fill_value .as_utf8() .value() @@ -207,7 +210,7 @@ fn fill_varbinview_array( .map(|i| { if validity_bits.value(i) { array_ref - .scalar_at(i) + .scalar_at(i, &mut ctx) .vortex_expect("scalar_at") .as_utf8() .value() @@ -233,6 +236,7 @@ fn fill_varbinview_array( } } DType::Binary(_) => { + let mut ctx = LEGACY_SESSION.create_execution_ctx(); let fill_bytes = fill_value .as_binary() .value() @@ -241,7 +245,7 @@ fn fill_varbinview_array( .map(|i| { if validity_bits.value(i) { array_ref - .scalar_at(i) + .scalar_at(i, &mut ctx) .vortex_expect("scalar_at") .as_binary() .value() diff --git a/fuzz/src/array/mod.rs b/fuzz/src/array/mod.rs index 4576eb7b167..c65f0d46c1a 100644 --- a/fuzz/src/array/mod.rs +++ b/fuzz/src/array/mod.rs @@ -13,6 +13,7 @@ pub(crate) use slice::*; pub use sort::sort_canonical_array; pub(crate) use sum::*; pub(crate) use take::*; +use vortex_array::LEGACY_SESSION; mod cast; mod compare; @@ -250,7 +251,7 @@ impl<'a> Arbitrary<'a> for FuzzArrayAction { let scalar = if u.arbitrary()? { current_array - .scalar_at(u.choose_index(current_array.len())?) + .scalar_at(u.choose_index(current_array.len())?, &mut ctx) .vortex_expect("scalar_at") } else { random_scalar(u, current_array.dtype())? @@ -291,7 +292,7 @@ impl<'a> Arbitrary<'a> for FuzzArrayAction { ActionType::Compare => { let scalar = if u.arbitrary()? { current_array - .scalar_at(u.choose_index(current_array.len())?) + .scalar_at(u.choose_index(current_array.len())?, &mut ctx) .vortex_expect("scalar_at") } else { // We can compare arrays with different nullability @@ -354,7 +355,7 @@ impl<'a> Arbitrary<'a> for FuzzArrayAction { } let fill_value = if u.arbitrary()? && !current_array.is_empty() { current_array - .scalar_at(u.choose_index(current_array.len())?) + .scalar_at(u.choose_index(current_array.len())?, &mut ctx) .vortex_expect("scalar_at") } else { random_scalar( @@ -669,7 +670,9 @@ pub fn run_fuzz_action(fuzz_action: FuzzArrayAction) -> VortexFuzzResult { Action::ScalarAt(indices) => { let expected_scalars = expected.scalar_vec(); for (j, &idx) in indices.iter().enumerate() { - let scalar = current_array.scalar_at(idx).vortex_expect("scalar_at"); + let scalar = current_array + .scalar_at(idx, &mut ctx) + .vortex_expect("scalar_at"); assert_scalar_eq(&expected_scalars[j], &scalar, i)?; } } @@ -726,9 +729,12 @@ pub fn assert_array_eq(lhs: &ArrayRef, rhs: &ArrayRef, step: usize) -> VortexFuz Backtrace::capture(), )); } + + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + for idx in 0..lhs.len() { - let l = lhs.scalar_at(idx).vortex_expect("scalar_at"); - let r = rhs.scalar_at(idx).vortex_expect("scalar_at"); + let l = lhs.scalar_at(idx, &mut ctx).vortex_expect("scalar_at"); + let r = rhs.scalar_at(idx, &mut ctx).vortex_expect("scalar_at"); if l != r { return Err(VortexFuzzError::ArrayNotEqual( diff --git a/fuzz/src/array/scalar_at.rs b/fuzz/src/array/scalar_at.rs index 91898390df6..c8e33f84b49 100644 --- a/fuzz/src/array/scalar_at.rs +++ b/fuzz/src/array/scalar_at.rs @@ -5,6 +5,8 @@ use std::sync::Arc; use vortex_array::Canonical; use vortex_array::IntoArray; +use vortex_array::LEGACY_SESSION; +use vortex_array::VortexSessionExecute as _; use vortex_array::arrays::bool::BoolArrayExt; use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; @@ -23,8 +25,9 @@ use vortex_error::VortexResult; /// This implementation manually extracts the scalar value from each canonical type /// without using the scalar_at method, to serve as an independent baseline for testing. pub fn scalar_at_canonical_array(canonical: Canonical, index: usize) -> VortexResult { + let mut ctx = LEGACY_SESSION.create_execution_ctx(); let canonical_ref = canonical.clone().into_array(); - if canonical_ref.is_invalid(index)? { + if canonical_ref.is_invalid(index, &mut ctx)? { return Ok(Scalar::null(canonical_ref.dtype().clone())); } Ok(match canonical { diff --git a/fuzz/src/array/search_sorted.rs b/fuzz/src/array/search_sorted.rs index d22b583a39b..44f4f4efd3e 100644 --- a/fuzz/src/array/search_sorted.rs +++ b/fuzz/src/array/search_sorted.rs @@ -5,7 +5,9 @@ use std::cmp::Ordering; use std::fmt::Debug; use vortex_array::ArrayRef; +use vortex_array::LEGACY_SESSION; use vortex_array::ToCanonical; +use vortex_array::VortexSessionExecute as _; use vortex_array::accessor::ArrayAccessor; use vortex_array::arrays::bool::BoolArrayExt; use vortex_array::dtype::DType; @@ -61,6 +63,7 @@ pub fn search_sorted_canonical_array( scalar: &Scalar, side: SearchSortedSide, ) -> VortexResult { + let mut ctx = LEGACY_SESSION.create_execution_ctx(); match array.dtype() { DType::Bool(_) => { let bool_array = array.to_bool(); @@ -127,7 +130,7 @@ pub fn search_sorted_canonical_array( } DType::Struct(..) | DType::List(..) | DType::FixedSizeList(..) => { let scalar_vals = (0..array.len()) - .map(|i| array.scalar_at(i)) + .map(|i| array.scalar_at(i, &mut ctx)) .collect::>>()?; scalar_vals.search_sorted(&scalar.cast(array.dtype())?, side) } diff --git a/fuzz/src/array/sort.rs b/fuzz/src/array/sort.rs index 6c07c76c20d..699260ef823 100644 --- a/fuzz/src/array/sort.rs +++ b/fuzz/src/array/sort.rs @@ -5,7 +5,9 @@ use std::cmp::Ordering; use vortex_array::ArrayRef; use vortex_array::IntoArray; +use vortex_array::LEGACY_SESSION; use vortex_array::ToCanonical; +use vortex_array::VortexSessionExecute as _; use vortex_array::accessor::ArrayAccessor; use vortex_array::arrays::BoolArray; use vortex_array::arrays::DecimalArray; @@ -71,12 +73,13 @@ pub fn sort_canonical_array(array: &ArrayRef) -> VortexResult { Ok(VarBinViewArray::from_iter(opt_values, array.dtype().clone()).into_array()) } DType::Struct(..) | DType::List(..) | DType::FixedSizeList(..) => { + let mut ctx = LEGACY_SESSION.create_execution_ctx(); let mut sort_indices = (0..array.len()).collect::>(); sort_indices.sort_by(|a, b| { array - .scalar_at(*a) + .scalar_at(*a, &mut ctx) .vortex_expect("scalar_at") - .partial_cmp(&array.scalar_at(*b).vortex_expect("scalar_at")) + .partial_cmp(&array.scalar_at(*b, &mut ctx).vortex_expect("scalar_at")) .vortex_expect("must be a valid comparison") }); take_canonical_array_non_nullable_indices(array, &sort_indices) diff --git a/fuzz/src/array/take.rs b/fuzz/src/array/take.rs index a113829077b..d34139ed985 100644 --- a/fuzz/src/array/take.rs +++ b/fuzz/src/array/take.rs @@ -3,7 +3,9 @@ use vortex_array::ArrayRef; use vortex_array::IntoArray; +use vortex_array::LEGACY_SESSION; use vortex_array::ToCanonical; +use vortex_array::VortexSessionExecute; use vortex_array::accessor::ArrayAccessor; use vortex_array::arrays::BoolArray; use vortex_array::arrays::DecimalArray; @@ -40,6 +42,7 @@ pub fn take_canonical_array_non_nullable_indices( } pub fn take_canonical_array(array: &ArrayRef, indices: &[Option]) -> VortexResult { + let mut ctx = LEGACY_SESSION.create_execution_ctx(); let nullable: Nullability = indices.contains(&None).into(); let validity = if array.dtype().is_nullable() || nullable == Nullability::Nullable { @@ -128,7 +131,7 @@ pub fn take_canonical_array(array: &ArrayRef, indices: &[Option]) -> Vort if let Some(idx) = idx { builder.append_scalar( &array - .scalar_at(*idx)? + .scalar_at(*idx, &mut ctx)? .cast(&array.dtype().union_nullability(nullable)) .vortex_expect("cannot cast scalar nullability"), )?; diff --git a/vortex-array/src/aggregate_fn/fns/first/mod.rs b/vortex-array/src/aggregate_fn/fns/first/mod.rs index 38e55607279..34bad773322 100644 --- a/vortex-array/src/aggregate_fn/fns/first/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/first/mod.rs @@ -94,13 +94,13 @@ impl AggregateFnVTable for First { &self, partial: &mut Self::Partial, batch: &ArrayRef, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { if partial.value.is_some() { return Ok(true); } if let Some(idx) = batch.validity_mask()?.first() { - let scalar = batch.scalar_at(idx)?; + let scalar = batch.scalar_at(idx, ctx)?; partial.value = Some(scalar.into_nullable()); } Ok(true) diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs b/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs index 4df47b770de..7d7a1173bf3 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs @@ -186,13 +186,13 @@ impl IsConstant { /// Kernels that compute `is_constant` by delegating to child arrays can call this /// to package the boolean result into the partial struct format expected by the /// accumulator, avoiding duplicated boilerplate. - pub fn make_partial(batch: &ArrayRef, is_constant: bool) -> VortexResult { + pub fn make_partial(batch: &ArrayRef, is_constant: bool, ctx: &mut ExecutionCtx) -> VortexResult { let partial_dtype = make_is_constant_partial_dtype(batch.dtype()); if is_constant { if batch.is_empty() { return Ok(Scalar::null(partial_dtype)); } - let first_value = batch.scalar_at(0)?.into_nullable(); + let first_value = batch.scalar_at(0, ctx)?.into_nullable(); Ok(Scalar::struct_( partial_dtype, vec![Scalar::bool(true, Nullability::NonNullable), first_value], @@ -386,7 +386,7 @@ impl AggregateFnVTable for IsConstant { // All valid from here. Check batch-level constancy. if c.len() == 1 { - partial.check_value(array_ref.scalar_at(0)?.into_nullable()); + partial.check_value(array_ref.scalar_at(0, ctx)?.into_nullable()); return Ok(()); } @@ -410,7 +410,7 @@ impl AggregateFnVTable for IsConstant { return Ok(()); } - partial.check_value(array_ref.scalar_at(0)?.into_nullable()); + partial.check_value(array_ref.scalar_at(0, ctx)?.into_nullable()); Ok(()) } } diff --git a/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs b/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs index 780fc492baa..c66a199d0c6 100644 --- a/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs @@ -114,7 +114,7 @@ fn is_sorted_impl(array: &ArrayRef, strict: bool, ctx: &mut ExecutionCtx) -> Vor 0 => {} // If we have a potential null value - it has to be the first one. 1 => { - if !array.is_invalid(0)? { + if !array.is_invalid(0, ctx)? { cache_is_sorted(array, strict, false); return Ok(false); } @@ -171,13 +171,13 @@ impl IsSorted { /// Kernels that compute `is_sorted` by delegating to child arrays can call this /// to package the boolean result into the partial struct format expected by the /// accumulator, avoiding duplicated boilerplate. - pub fn make_partial(batch: &ArrayRef, is_sorted: bool, strict: bool) -> VortexResult { + pub fn make_partial(batch: &ArrayRef, is_sorted: bool, strict: bool, ctx: &mut ExecutionCtx) -> VortexResult { let partial_dtype = make_is_sorted_partial_dtype(batch.dtype()); if batch.is_empty() { return Ok(Scalar::null(partial_dtype)); } - let first_value = batch.scalar_at(0)?.into_nullable(); - let last_value = batch.scalar_at(batch.len() - 1)?.into_nullable(); + let first_value = batch.scalar_at(0, ctx)?.into_nullable(); + let last_value = batch.scalar_at(batch.len() - 1, ctx)?.into_nullable(); // SAFETY: We constructed partial_dtype and the children match its field dtypes exactly. Ok(unsafe { Scalar::struct_unchecked( @@ -430,7 +430,7 @@ impl AggregateFnVTable for IsSorted { let array_ref = c.clone().into_array(); // Check boundary with previous chunk. - let first_value = array_ref.scalar_at(0)?.into_nullable(); + let first_value = array_ref.scalar_at(0, ctx)?.into_nullable(); if let Some(ref self_last) = partial.last_value { if !self_last.is_null() && !first_value.is_null() { let boundary_ok = if partial.strict { @@ -441,7 +441,7 @@ impl AggregateFnVTable for IsSorted { if !boundary_ok { partial.is_sorted = false; partial.last_value = - Some(array_ref.scalar_at(array_ref.len() - 1)?.into_nullable()); + Some(array_ref.scalar_at(array_ref.len() - 1, ctx)?.into_nullable()); if partial.first_value.is_none() { partial.first_value = Some(first_value); } @@ -452,7 +452,7 @@ impl AggregateFnVTable for IsSorted { { partial.is_sorted = false; partial.last_value = - Some(array_ref.scalar_at(array_ref.len() - 1)?.into_nullable()); + Some(array_ref.scalar_at(array_ref.len() - 1, ctx)?.into_nullable()); if partial.first_value.is_none() { partial.first_value = Some(first_value); } @@ -480,7 +480,7 @@ impl AggregateFnVTable for IsSorted { partial.first_value = Some(first_value); } partial.last_value = - Some(array_ref.scalar_at(array_ref.len() - 1)?.into_nullable()); + Some(array_ref.scalar_at(array_ref.len() - 1, ctx)?.into_nullable()); Ok(()) } } diff --git a/vortex-array/src/aggregate_fn/fns/last/mod.rs b/vortex-array/src/aggregate_fn/fns/last/mod.rs index 009b487ad2d..e0c2b344ecf 100644 --- a/vortex-array/src/aggregate_fn/fns/last/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/last/mod.rs @@ -95,10 +95,10 @@ impl AggregateFnVTable for Last { &self, partial: &mut Self::Partial, batch: &ArrayRef, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { if let Some(idx) = batch.validity_mask()?.last() { - let scalar = batch.scalar_at(idx)?; + let scalar = batch.scalar_at(idx, ctx)?; partial.value = Some(scalar.into_nullable()); } Ok(true) diff --git a/vortex-array/src/aggregate_fn/vtable.rs b/vortex-array/src/aggregate_fn/vtable.rs index 73dd2604700..2b458460b90 100644 --- a/vortex-array/src/aggregate_fn/vtable.rs +++ b/vortex-array/src/aggregate_fn/vtable.rs @@ -15,6 +15,8 @@ use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; use crate::IntoArray; +use crate::LEGACY_SESSION; +use crate::VortexSessionExecute; use crate::aggregate_fn::AggregateFn; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnRef; @@ -142,7 +144,7 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync { let scalar = self.to_scalar(partial)?; let array = ConstantArray::new(scalar, 1).into_array(); let result = self.finalize(array)?; - result.scalar_at(0) + result.scalar_at(0, &mut LEGACY_SESSION.create_execution_ctx()) } } diff --git a/vortex-array/src/array/erased.rs b/vortex-array/src/array/erased.rs index cfefe316ca4..e1473b6cbf4 100644 --- a/vortex-array/src/array/erased.rs +++ b/vortex-array/src/array/erased.rs @@ -209,24 +209,24 @@ impl ArrayRef { } /// Fetch the scalar at the given index. - pub fn scalar_at(&self, index: usize) -> VortexResult { + pub fn scalar_at(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult { vortex_ensure!(index < self.len(), OutOfBounds: index, 0, self.len()); - if self.is_invalid(index)? { + if self.is_invalid(index, ctx)? { return Ok(Scalar::null(self.dtype().clone())); } - let scalar = self.0.scalar_at(self, index)?; + let scalar = self.0.scalar_at(self, index, ctx)?; vortex_ensure!(self.dtype() == scalar.dtype(), "Scalar dtype mismatch"); Ok(scalar) } /// Returns whether the item at `index` is valid. - pub fn is_valid(&self, index: usize) -> VortexResult { + pub fn is_valid(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult { vortex_ensure!(index < self.len(), OutOfBounds: index, 0, self.len()); match self.validity()? { Validity::NonNullable | Validity::AllValid => Ok(true), Validity::AllInvalid => Ok(false), Validity::Array(a) => a - .scalar_at(index)? + .scalar_at(index, ctx)? .as_bool() .value() .ok_or_else(|| vortex_err!("validity value at index {} is null", index)), @@ -234,8 +234,8 @@ impl ArrayRef { } /// Returns whether the item at `index` is invalid. - pub fn is_invalid(&self, index: usize) -> VortexResult { - Ok(!self.is_valid(index)?) + pub fn is_invalid(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult { + Ok(!self.is_valid(index, ctx)?) } /// Returns whether all items in the array are valid. diff --git a/vortex-array/src/array/mod.rs b/vortex-array/src/array/mod.rs index 9ef3e57cab2..59ec25bfa44 100644 --- a/vortex-array/src/array/mod.rs +++ b/vortex-array/src/array/mod.rs @@ -17,8 +17,6 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use crate::ExecutionCtx; -use crate::LEGACY_SESSION; -use crate::VortexSessionExecute; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::dtype::DType; @@ -77,7 +75,12 @@ pub(crate) trait DynArray: 'static + private::Sealed + Send + Sync + Debug { /// Fetch the scalar at the given index. /// /// This method panics if the index is out of bounds for the array. - fn scalar_at(&self, this: &ArrayRef, index: usize) -> VortexResult; + fn scalar_at( + &self, + this: &ArrayRef, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult; /// Returns the [`Validity`] of the array. fn validity(&self, this: &ArrayRef) -> VortexResult; @@ -218,13 +221,14 @@ impl DynArray for ArrayInner { self.vtable.id() } - fn scalar_at(&self, this: &ArrayRef, index: usize) -> VortexResult { + fn scalar_at( + &self, + this: &ArrayRef, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { let view = unsafe { ArrayView::new_unchecked(this, &self.data) }; - >::scalar_at( - view, - index, - &mut LEGACY_SESSION.create_execution_ctx(), - ) + >::scalar_at(view, index, ctx) } fn validity(&self, this: &ArrayRef) -> VortexResult { diff --git a/vortex-array/src/array/typed.rs b/vortex-array/src/array/typed.rs index 55656409128..914d5510ac4 100644 --- a/vortex-array/src/array/typed.rs +++ b/vortex-array/src/array/typed.rs @@ -15,6 +15,7 @@ use std::sync::Arc; use vortex_error::VortexResult; use crate::ArrayRef; +use crate::ExecutionCtx; use crate::IntoArray; use crate::array::ArrayId; use crate::array::ArrayView; @@ -344,8 +345,12 @@ impl Array { self.inner.slice(range) } - pub fn scalar_at(&self, index: usize) -> VortexResult { - self.inner.scalar_at(index) + pub fn scalar_at( + &self, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + self.inner.scalar_at(index, ctx) } pub fn filter(&self, mask: vortex_mask::Mask) -> VortexResult { @@ -360,12 +365,12 @@ impl Array { self.inner.validity() } - pub fn is_valid(&self, index: usize) -> VortexResult { - self.inner.is_valid(index) + pub fn is_valid(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult { + self.inner.is_valid(index, ctx) } - pub fn is_invalid(&self, index: usize) -> VortexResult { - self.inner.is_invalid(index) + pub fn is_invalid(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult { + self.inner.is_invalid(index, ctx) } pub fn all_valid(&self) -> VortexResult { @@ -403,7 +408,7 @@ impl Array { pub fn append_to_builder( &self, builder: &mut dyn crate::builders::ArrayBuilder, - ctx: &mut crate::ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult<()> { self.inner.append_to_builder(builder, ctx) } diff --git a/vortex-array/src/arrays/assertions.rs b/vortex-array/src/arrays/assertions.rs index 2dab3a23c44..d6bf260faad 100644 --- a/vortex-array/src/arrays/assertions.rs +++ b/vortex-array/src/arrays/assertions.rs @@ -30,8 +30,9 @@ fn execute_to_canonical(array: ArrayRef, ctx: &mut ExecutionCtx) -> ArrayRef { #[expect(clippy::unwrap_used)] fn find_mismatched_indices(left: &ArrayRef, right: &ArrayRef) -> Vec { assert_eq!(left.len(), right.len()); + let mut ctx = LEGACY_SESSION.create_execution_ctx(); (0..left.len()) - .filter(|i| left.scalar_at(*i).unwrap() != right.scalar_at(*i).unwrap()) + .filter(|i| left.scalar_at(*i, &mut ctx).unwrap() != right.scalar_at(*i, &mut ctx).unwrap()) .collect() } diff --git a/vortex-array/src/arrays/chunked/vtable/operations.rs b/vortex-array/src/arrays/chunked/vtable/operations.rs index 7c7095a05c7..51ecfeafe02 100644 --- a/vortex-array/src/arrays/chunked/vtable/operations.rs +++ b/vortex-array/src/arrays/chunked/vtable/operations.rs @@ -14,10 +14,10 @@ impl OperationsVTable for Chunked { fn scalar_at( array: ArrayView<'_, Chunked>, index: usize, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { let (chunk_index, chunk_offset) = array.find_chunk_idx(index)?; - array.chunk(chunk_index).scalar_at(chunk_offset) + array.chunk(chunk_index).scalar_at(chunk_offset, ctx) } } diff --git a/vortex-array/src/arrays/constant/vtable/canonical.rs b/vortex-array/src/arrays/constant/vtable/canonical.rs index 137f48fd8eb..c8859d28135 100644 --- a/vortex-array/src/arrays/constant/vtable/canonical.rs +++ b/vortex-array/src/arrays/constant/vtable/canonical.rs @@ -338,6 +338,8 @@ mod tests { use crate::dtype::half::f16; use crate::expr::stats::Stat; use crate::expr::stats::StatsProvider; + use crate::LEGACY_SESSION; + use crate::VortexSessionExecute; use crate::scalar::Scalar; use crate::validity::Validity; @@ -659,14 +661,14 @@ mod tests { let element_validity = elements .validity() .vortex_expect("constant canonical element validity should be derivable"); - assert!(element_validity.is_valid(0).unwrap()); - assert!(!element_validity.is_valid(1).unwrap()); - assert!(element_validity.is_valid(2).unwrap()); + assert!(element_validity.is_valid(0, &mut LEGACY_SESSION.create_execution_ctx()).unwrap()); + assert!(!element_validity.is_valid(1, &mut LEGACY_SESSION.create_execution_ctx()).unwrap()); + assert!(element_validity.is_valid(2, &mut LEGACY_SESSION.create_execution_ctx()).unwrap()); // Pattern should repeat. - assert!(element_validity.is_valid(3).unwrap()); - assert!(!element_validity.is_valid(4).unwrap()); - assert!(element_validity.is_valid(5).unwrap()); + assert!(element_validity.is_valid(3, &mut LEGACY_SESSION.create_execution_ctx()).unwrap()); + assert!(!element_validity.is_valid(4, &mut LEGACY_SESSION.create_execution_ctx()).unwrap()); + assert!(element_validity.is_valid(5, &mut LEGACY_SESSION.create_execution_ctx()).unwrap()); } #[test] diff --git a/vortex-array/src/arrays/dict/compute/is_constant.rs b/vortex-array/src/arrays/dict/compute/is_constant.rs index f5a903a126f..82518963106 100644 --- a/vortex-array/src/arrays/dict/compute/is_constant.rs +++ b/vortex-array/src/arrays/dict/compute/is_constant.rs @@ -40,7 +40,7 @@ impl DynAggregateKernel for DictIsConstantKernel { // If codes are constant, only one dictionary value is referenced → constant. if is_constant(dict.codes(), ctx)? { - return Ok(Some(IsConstant::make_partial(batch, true)?)); + return Ok(Some(IsConstant::make_partial(batch, true, ctx)?)); } // Otherwise, check the values array. Filter to only referenced values if needed. @@ -53,6 +53,6 @@ impl DynAggregateKernel for DictIsConstantKernel { is_constant(&filtered_values, ctx)? }; - Ok(Some(IsConstant::make_partial(batch, result)?)) + Ok(Some(IsConstant::make_partial(batch, result, ctx)?)) } } diff --git a/vortex-array/src/arrays/dict/compute/is_sorted.rs b/vortex-array/src/arrays/dict/compute/is_sorted.rs index cb20287ae83..6ab5d5727d5 100644 --- a/vortex-array/src/arrays/dict/compute/is_sorted.rs +++ b/vortex-array/src/arrays/dict/compute/is_sorted.rs @@ -44,7 +44,7 @@ impl DynAggregateKernel for DictIsSortedKernel { }; if result { - Ok(Some(IsSorted::make_partial(batch, true, strict)?)) + Ok(Some(IsSorted::make_partial(batch, true, strict, ctx)?)) } else { // We can't definitively say it's NOT sorted without canonicalizing, // so return None to let the accumulator handle it. diff --git a/vortex-array/src/arrays/dict/vtable/operations.rs b/vortex-array/src/arrays/dict/vtable/operations.rs index 9c75e4f2924..5b7f119d26b 100644 --- a/vortex-array/src/arrays/dict/vtable/operations.rs +++ b/vortex-array/src/arrays/dict/vtable/operations.rs @@ -15,11 +15,11 @@ impl OperationsVTable for Dict { fn scalar_at( array: ArrayView<'_, Dict>, index: usize, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { let Some(dict_index) = array .codes() - .scalar_at(index)? + .scalar_at(index, ctx)? .as_primitive() .as_::() else { @@ -28,7 +28,7 @@ impl OperationsVTable for Dict { Ok(array .values() - .scalar_at(dict_index)? + .scalar_at(dict_index, ctx)? .cast(array.dtype()) .vortex_expect("Array dtype will only differ by nullability")) } diff --git a/vortex-array/src/arrays/extension/vtable/operations.rs b/vortex-array/src/arrays/extension/vtable/operations.rs index d86c7599aef..b9ac304fb4c 100644 --- a/vortex-array/src/arrays/extension/vtable/operations.rs +++ b/vortex-array/src/arrays/extension/vtable/operations.rs @@ -14,11 +14,11 @@ impl OperationsVTable for Extension { fn scalar_at( array: ArrayView<'_, Extension>, index: usize, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { Ok(Scalar::extension_ref( array.ext_dtype().clone(), - array.storage_array().scalar_at(index)?, + array.storage_array().scalar_at(index, ctx)?, )) } } diff --git a/vortex-array/src/arrays/filter/vtable.rs b/vortex-array/src/arrays/filter/vtable.rs index 8eb7fb62213..8dd6f9234f2 100644 --- a/vortex-array/src/arrays/filter/vtable.rs +++ b/vortex-array/src/arrays/filter/vtable.rs @@ -181,10 +181,10 @@ impl OperationsVTable for Filter { fn scalar_at( array: ArrayView<'_, Filter>, index: usize, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { let rank_idx = array.mask.rank(index); - array.child().scalar_at(rank_idx) + array.child().scalar_at(rank_idx, ctx) } } diff --git a/vortex-array/src/arrays/fixed_size_list/array.rs b/vortex-array/src/arrays/fixed_size_list/array.rs index dcbee9f8f6f..6f4e45b97a9 100644 --- a/vortex-array/src/arrays/fixed_size_list/array.rs +++ b/vortex-array/src/arrays/fixed_size_list/array.rs @@ -242,11 +242,11 @@ pub trait FixedSizeListArrayExt: TypedArrayRef { index, self.as_ref().len(), ); - debug_assert!( - self.fixed_size_list_validity() - .is_valid(index) - .unwrap_or(false) - ); + // debug_assert!( + // self.fixed_size_list_validity() + // .is_valid(index, &mut LEGACY_SESSION.create_execution_ctx()) + // .unwrap_or(false) + // ); let start = self.list_size() as usize * index; let end = self.list_size() as usize * (index + 1); diff --git a/vortex-array/src/arrays/fixed_size_list/vtable/operations.rs b/vortex-array/src/arrays/fixed_size_list/vtable/operations.rs index ebd2ed463a4..3c60a5fd718 100644 --- a/vortex-array/src/arrays/fixed_size_list/vtable/operations.rs +++ b/vortex-array/src/arrays/fixed_size_list/vtable/operations.rs @@ -14,12 +14,12 @@ impl OperationsVTable for FixedSizeList { fn scalar_at( array: ArrayView<'_, FixedSizeList>, index: usize, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { // By the preconditions we know that the list scalar is not null. let list = array.fixed_size_list_elements_at(index)?; let children_elements: Vec = (0..list.len()) - .map(|i| list.scalar_at(i)) + .map(|i| list.scalar_at(i, ctx)) .collect::>()?; debug_assert_eq!(children_elements.len(), array.list_size() as usize); diff --git a/vortex-array/src/arrays/list/array.rs b/vortex-array/src/arrays/list/array.rs index 75db7d0f08c..fbdb63681af 100644 --- a/vortex-array/src/arrays/list/array.rs +++ b/vortex-array/src/arrays/list/array.rs @@ -305,7 +305,7 @@ pub trait ListArrayExt: TypedArrayRef { })) } else { self.offsets() - .scalar_at(index)? + .scalar_at(index, &mut LEGACY_SESSION.create_execution_ctx())? .as_primitive() .as_::() .ok_or_else(|| vortex_error::vortex_err!("offset value does not fit in usize")) @@ -340,7 +340,7 @@ pub trait ListArrayExt: TypedArrayRef { } let offsets = self.offsets(); - let first_offset = offsets.scalar_at(0)?; + let first_offset = offsets.scalar_at(0, &mut LEGACY_SESSION.create_execution_ctx())?; let adjusted_offsets = offsets.clone().binary( ConstantArray::new(first_offset, offsets.len()).into_array(), Operator::Sub, diff --git a/vortex-array/src/arrays/list/vtable/operations.rs b/vortex-array/src/arrays/list/vtable/operations.rs index c9fc69c7bbe..dacf30f10d3 100644 --- a/vortex-array/src/arrays/list/vtable/operations.rs +++ b/vortex-array/src/arrays/list/vtable/operations.rs @@ -16,12 +16,12 @@ impl OperationsVTable for List { fn scalar_at( array: ArrayView<'_, List>, index: usize, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { // By the preconditions we know that the list scalar is not null. let elems = array.list_elements_at(index)?; let scalars: Vec = (0..elems.len()) - .map(|i| elems.scalar_at(i)) + .map(|i| elems.scalar_at(i, ctx)) .collect::>()?; Ok(Scalar::list( diff --git a/vortex-array/src/arrays/listview/array.rs b/vortex-array/src/arrays/listview/array.rs index 72c376b2f44..026575764e5 100644 --- a/vortex-array/src/arrays/listview/array.rs +++ b/vortex-array/src/arrays/listview/array.rs @@ -13,7 +13,9 @@ use vortex_error::vortex_ensure; use vortex_error::vortex_err; use crate::ArrayRef; +use crate::LEGACY_SESSION; use crate::ToCanonical; +use crate::VortexSessionExecute; use crate::array::Array; use crate::array::ArrayParts; use crate::array::TypedArrayRef; @@ -349,7 +351,7 @@ pub trait ListViewArrayExt: TypedArrayRef { .map(|p| match_each_integer_ptype!(p.ptype(), |P| { p.as_slice::

()[index].as_() })) .unwrap_or_else(|| { self.offsets() - .scalar_at(index) + .scalar_at(index, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("offsets must support scalar_at") .as_primitive() .as_::() @@ -369,7 +371,7 @@ pub trait ListViewArrayExt: TypedArrayRef { .map(|p| match_each_integer_ptype!(p.ptype(), |P| { p.as_slice::

()[index].as_() })) .unwrap_or_else(|| { self.sizes() - .scalar_at(index) + .scalar_at(index, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("sizes must support scalar_at") .as_primitive() .as_::() diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index 12c2c653162..db291ee8f13 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -162,7 +162,7 @@ impl ListViewArray { let mut n_elements = NewOffset::zero(); for index in 0..len { - if !self.validity()?.is_valid(index)? { + if !self.validity()?.is_valid(index, &mut LEGACY_SESSION.create_execution_ctx())? { new_offsets.push(n_elements); new_sizes.push(S::zero()); continue; @@ -230,7 +230,7 @@ impl ListViewArray { let mut n_elements = NewOffset::zero(); for index in 0..len { - if !self.validity()?.is_valid(index)? { + if !self.validity()?.is_valid(index, &mut LEGACY_SESSION.create_execution_ctx())? { // For NULL lists, place them after the previous item's data to maintain the // no-overlap invariant for zero-copy to `ListArray` arrays. new_offsets.push(n_elements); @@ -440,9 +440,9 @@ mod tests { // Verify nullability is preserved assert_eq!(flattened.dtype().nullability(), Nullability::Nullable); - assert!(flattened.validity()?.is_valid(0).unwrap()); - assert!(!flattened.validity()?.is_valid(1).unwrap()); - assert!(flattened.validity()?.is_valid(2).unwrap()); + assert!(flattened.validity()?.is_valid(0, &mut LEGACY_SESSION.create_execution_ctx()).unwrap()); + assert!(!flattened.validity()?.is_valid(1, &mut LEGACY_SESSION.create_execution_ctx()).unwrap()); + assert!(flattened.validity()?.is_valid(2, &mut LEGACY_SESSION.create_execution_ctx()).unwrap()); // Verify valid lists contain correct data assert_arrays_eq!( diff --git a/vortex-array/src/arrays/listview/vtable/operations.rs b/vortex-array/src/arrays/listview/vtable/operations.rs index c899083027a..9f94411e1f3 100644 --- a/vortex-array/src/arrays/listview/vtable/operations.rs +++ b/vortex-array/src/arrays/listview/vtable/operations.rs @@ -16,12 +16,12 @@ impl OperationsVTable for ListView { fn scalar_at( array: ArrayView<'_, ListView>, index: usize, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { // By the preconditions we know that the list scalar is not null. let list = array.list_elements_at(index)?; let children: Vec = (0..list.len()) - .map(|i| list.scalar_at(i)) + .map(|i| list.scalar_at(i, ctx)) .collect::>()?; Ok(Scalar::list( diff --git a/vortex-array/src/arrays/masked/vtable/operations.rs b/vortex-array/src/arrays/masked/vtable/operations.rs index eb3132cf957..13e6d14b566 100644 --- a/vortex-array/src/arrays/masked/vtable/operations.rs +++ b/vortex-array/src/arrays/masked/vtable/operations.rs @@ -14,9 +14,9 @@ impl OperationsVTable for Masked { fn scalar_at( array: ArrayView<'_, Masked>, index: usize, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { // Invalid indices are handled by the entrypoint function. - Ok(array.child().scalar_at(index)?.into_nullable()) + Ok(array.child().scalar_at(index, ctx)?.into_nullable()) } } diff --git a/vortex-array/src/arrays/patched/array.rs b/vortex-array/src/arrays/patched/array.rs index b2be4205b7f..84e1040d823 100644 --- a/vortex-array/src/arrays/patched/array.rs +++ b/vortex-array/src/arrays/patched/array.rs @@ -15,6 +15,8 @@ use crate::ArrayRef; use crate::Canonical; use crate::ExecutionCtx; use crate::IntoArray; +use crate::LEGACY_SESSION; +use crate::VortexSessionExecute; use crate::array::Array; use crate::array::ArrayParts; use crate::array::TypedArrayRef; @@ -113,10 +115,10 @@ pub trait PatchedArrayExt: PatchedArraySlotsExt { let start = self .lane_offsets() - .scalar_at(chunk * self.n_lanes() + lane)?; + .scalar_at(chunk * self.n_lanes() + lane, &mut LEGACY_SESSION.create_execution_ctx())?; let stop = self .lane_offsets() - .scalar_at(chunk * self.n_lanes() + lane + 1)?; + .scalar_at(chunk * self.n_lanes() + lane + 1, &mut LEGACY_SESSION.create_execution_ctx())?; let start = start .as_primitive() diff --git a/vortex-array/src/arrays/patched/vtable/operations.rs b/vortex-array/src/arrays/patched/vtable/operations.rs index 95beeaaf87f..b28dcf35f4b 100644 --- a/vortex-array/src/arrays/patched/vtable/operations.rs +++ b/vortex-array/src/arrays/patched/vtable/operations.rs @@ -43,12 +43,12 @@ impl OperationsVTable for Patched { // be slower. for (&patch_index, idx) in std::iter::zip(patch_indices.as_slice::(), range) { if patch_index == chunk_index { - return array.patch_values().scalar_at(idx)?.cast(array.dtype()); + return array.patch_values().scalar_at(idx, ctx)?.cast(array.dtype()); } } // Otherwise, access the underlying value. - array.inner().scalar_at(index) + array.inner().scalar_at(index, ctx) } } diff --git a/vortex-array/src/arrays/scalar_fn/rules.rs b/vortex-array/src/arrays/scalar_fn/rules.rs index d0cb0efdbf4..7c300fde4f9 100644 --- a/vortex-array/src/arrays/scalar_fn/rules.rs +++ b/vortex-array/src/arrays/scalar_fn/rules.rs @@ -11,6 +11,8 @@ use vortex_error::VortexResult; use crate::ArrayRef; use crate::Canonical; use crate::IntoArray; +use crate::LEGACY_SESSION; +use crate::VortexSessionExecute; use crate::array::ArrayView; use crate::arrays::Constant; use crate::arrays::ConstantArray; @@ -79,7 +81,7 @@ impl ArrayReduceRule for ScalarFnConstantRule { if array.is_empty() { Ok(Some(Canonical::empty(array.dtype()).into_array())) } else { - let result = array.array().scalar_at(0)?; + let result = array.array().scalar_at(0, &mut LEGACY_SESSION.create_execution_ctx())?; Ok(Some(ConstantArray::new(result, array.len()).into_array())) } } diff --git a/vortex-array/src/arrays/scalar_fn/vtable/operations.rs b/vortex-array/src/arrays/scalar_fn/vtable/operations.rs index bcb0368772a..58f6255a098 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/operations.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/operations.rs @@ -20,12 +20,12 @@ impl OperationsVTable for ScalarFnVTable { fn scalar_at( array: ArrayView<'_, ScalarFnVTable>, index: usize, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { let inputs: Vec<_> = array .children() .iter() - .map(|child| Ok(ConstantArray::new(child.scalar_at(index)?, 1).into_array())) + .map(|child| Ok(ConstantArray::new(child.scalar_at(index, ctx)?, 1).into_array())) .collect::>()?; let mut ctx = LEGACY_SESSION.create_execution_ctx(); @@ -38,7 +38,7 @@ impl OperationsVTable for ScalarFnVTable { "Scalar function {} returned non-constant array from execution over all scalar inputs", array.scalar_fn(), ); - arr.into_array().scalar_at(0)? + arr.into_array().scalar_at(0, &mut ctx)? } Columnar::Constant(constant) => constant.scalar().clone(), }; diff --git a/vortex-array/src/arrays/shared/vtable.rs b/vortex-array/src/arrays/shared/vtable.rs index 6be759c83eb..c7bd5da6239 100644 --- a/vortex-array/src/arrays/shared/vtable.rs +++ b/vortex-array/src/arrays/shared/vtable.rs @@ -121,9 +121,9 @@ impl OperationsVTable for Shared { fn scalar_at( array: ArrayView<'_, Shared>, index: usize, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { - array.current_array_ref().scalar_at(index) + array.current_array_ref().scalar_at(index, ctx) } } diff --git a/vortex-array/src/arrays/slice/vtable.rs b/vortex-array/src/arrays/slice/vtable.rs index 2ac41a358b9..f2502ff97cd 100644 --- a/vortex-array/src/arrays/slice/vtable.rs +++ b/vortex-array/src/arrays/slice/vtable.rs @@ -166,9 +166,9 @@ impl OperationsVTable for Slice { fn scalar_at( array: ArrayView<'_, Slice>, index: usize, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { - array.child().scalar_at(array.range.start + index) + array.child().scalar_at(array.range.start + index, ctx) } } diff --git a/vortex-array/src/arrays/struct_/vtable/operations.rs b/vortex-array/src/arrays/struct_/vtable/operations.rs index 021d93f3c2f..79c4ad8c4e5 100644 --- a/vortex-array/src/arrays/struct_/vtable/operations.rs +++ b/vortex-array/src/arrays/struct_/vtable/operations.rs @@ -14,11 +14,11 @@ impl OperationsVTable for Struct { fn scalar_at( array: ArrayView<'_, Struct>, index: usize, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { let field_scalars: VortexResult> = array .iter_unmasked_fields() - .map(|field| field.scalar_at(index)) + .map(|field| field.scalar_at(index, ctx)) .collect(); // SAFETY: The vtable guarantees index is in-bounds and non-null before this is called. // Each field's scalar_at returns a scalar with the field's own dtype. diff --git a/vortex-array/src/arrays/varbin/array.rs b/vortex-array/src/arrays/varbin/array.rs index fb35ee503bc..3b3e3fe1b9e 100644 --- a/vortex-array/src/arrays/varbin/array.rs +++ b/vortex-array/src/arrays/varbin/array.rs @@ -13,7 +13,9 @@ use vortex_error::vortex_err; use vortex_mask::Mask; use crate::ArrayRef; +use crate::LEGACY_SESSION; use crate::ToCanonical; +use crate::VortexSessionExecute; use crate::array::Array; use crate::array::ArrayParts; use crate::array::TypedArrayRef; @@ -211,7 +213,7 @@ impl VarBinData { // Skip host-only validation when offsets/bytes are not host-resident. if offsets.is_host() && bytes.is_on_host() { let last_offset = offsets - .scalar_at(offsets.len() - 1)? + .scalar_at(offsets.len() - 1, &mut LEGACY_SESSION.create_execution_ctx())? .as_primitive() .as_::() .ok_or_else( @@ -249,7 +251,7 @@ impl VarBinData { .map(|o| (o[0].as_(), o[1].as_())) .enumerate() { - if validity.is_null(i)? { + if validity.is_null(i, &mut LEGACY_SESSION.create_execution_ctx())? { continue; } @@ -330,7 +332,7 @@ pub trait VarBinArrayExt: TypedArrayRef { (&self .offsets() - .scalar_at(index) + .scalar_at(index, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("offsets must support scalar_at")) .try_into() .vortex_expect("Failed to convert offset to usize") diff --git a/vortex-array/src/arrays/varbin/compute/filter.rs b/vortex-array/src/arrays/varbin/compute/filter.rs index 8f4b6793a72..3054d23f1eb 100644 --- a/vortex-array/src/arrays/varbin/compute/filter.rs +++ b/vortex-array/src/arrays/varbin/compute/filter.rs @@ -14,6 +14,8 @@ use vortex_mask::MaskIter; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; +use crate::LEGACY_SESSION; +use crate::VortexSessionExecute; use crate::array::ArrayView; use crate::arrays::PrimitiveArray; use crate::arrays::VarBin; @@ -184,7 +186,7 @@ fn filter_select_var_bin_by_index_primitive_offset( ) -> VortexResult { let mut builder = VarBinBuilder::::with_capacity(selection_count); for idx in mask_indices.iter().copied() { - if validity.is_valid(idx)? { + if validity.is_valid(idx, &mut LEGACY_SESSION.create_execution_ctx())? { let (start, end) = ( offsets[idx].to_usize().ok_or_else(|| { vortex_err!("Failed to convert offset to usize: {}", offsets[idx]) diff --git a/vortex-array/src/arrays/varbinview/array.rs b/vortex-array/src/arrays/varbinview/array.rs index 37d055107cf..961b363501d 100644 --- a/vortex-array/src/arrays/varbinview/array.rs +++ b/vortex-array/src/arrays/varbinview/array.rs @@ -18,6 +18,8 @@ use vortex_error::vortex_panic; use vortex_mask::Mask; use crate::ArrayRef; +use crate::LEGACY_SESSION; +use crate::VortexSessionExecute; use crate::array::Array; use crate::array::ArrayParts; use crate::array::TypedArrayRef; @@ -317,7 +319,7 @@ impl VarBinViewData { F: Fn(&[u8]) -> bool, { for (idx, &view) in views.iter().enumerate() { - if validity.is_null(idx)? { + if validity.is_null(idx, &mut LEGACY_SESSION.create_execution_ctx())? { continue; } diff --git a/vortex-array/src/arrays/variant/vtable/operations.rs b/vortex-array/src/arrays/variant/vtable/operations.rs index fb6fc0eb499..2057a7397cb 100644 --- a/vortex-array/src/arrays/variant/vtable/operations.rs +++ b/vortex-array/src/arrays/variant/vtable/operations.rs @@ -14,8 +14,8 @@ impl OperationsVTable for Variant { fn scalar_at( array: ArrayView<'_, Variant>, index: usize, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { - array.child().scalar_at(index) + array.child().scalar_at(index, ctx) } } diff --git a/vortex-array/src/arrow/datum.rs b/vortex-array/src/arrow/datum.rs index ce77c3bb159..c5cf3e9d8f4 100644 --- a/vortex-array/src/arrow/datum.rs +++ b/vortex-array/src/arrow/datum.rs @@ -11,6 +11,8 @@ use vortex_error::vortex_panic; use crate::ArrayRef; use crate::IntoArray; +use crate::LEGACY_SESSION; +use crate::VortexSessionExecute; use crate::arrays::Constant; use crate::arrays::ConstantArray; use crate::arrow::FromArrowArray; @@ -103,7 +105,7 @@ where Ok(ConstantArray::new( array - .scalar_at(0) + .scalar_at(0, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("array of length 1 must support scalar_at(0)"), len, ) diff --git a/vortex-array/src/arrow/executor/list.rs b/vortex-array/src/arrow/executor/list.rs index 285ad733cfb..217b3a2e809 100644 --- a/vortex-array/src/arrow/executor/list.rs +++ b/vortex-array/src/arrow/executor/list.rs @@ -147,7 +147,7 @@ fn list_view_zctl( // For ZCTL, we know that we only care about the final size. assert!(!sizes.is_empty()); let final_size = sizes - .scalar_at(sizes.len() - 1)? + .scalar_at(sizes.len() - 1, ctx)? .cast(&DType::Primitive(O::PTYPE, Nullability::NonNullable))?; let final_size = final_size .as_primitive() diff --git a/vortex-array/src/builders/listview.rs b/vortex-array/src/builders/listview.rs index b339dc430f7..f4136d103a9 100644 --- a/vortex-array/src/builders/listview.rs +++ b/vortex-array/src/builders/listview.rs @@ -20,7 +20,9 @@ use vortex_mask::Mask; use crate::ArrayRef; use crate::Canonical; +use crate::LEGACY_SESSION; use crate::ToCanonical; +use crate::VortexSessionExecute; use crate::array::IntoArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; @@ -301,7 +303,7 @@ impl ArrayBuilder for ListViewBuilder { if !listview.is_zero_copy_to_list() { for i in 0..listview.len() { let list = listview - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at failed in extend_from_array_unchecked"); self.append_scalar(&list) diff --git a/vortex-array/src/compute/conformance/binary_numeric.rs b/vortex-array/src/compute/conformance/binary_numeric.rs index c36367746e7..2b3e926dbb0 100644 --- a/vortex-array/src/compute/conformance/binary_numeric.rs +++ b/vortex-array/src/compute/conformance/binary_numeric.rs @@ -47,7 +47,7 @@ fn to_vec_of_scalar(array: &ArrayRef) -> Vec { (0..array.len()) .map(|index| { array - .scalar_at(index) + .scalar_at(index, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") }) .collect_vec() diff --git a/vortex-array/src/compute/conformance/cast.rs b/vortex-array/src/compute/conformance/cast.rs index cab73acc521..eb63b032e3a 100644 --- a/vortex-array/src/compute/conformance/cast.rs +++ b/vortex-array/src/compute/conformance/cast.rs @@ -74,10 +74,10 @@ fn test_cast_identity(array: &ArrayRef) { for i in 0..array.len().min(10) { assert_eq!( array - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), result - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -109,7 +109,7 @@ fn test_cast_from_null(array: &ArrayRef) { for i in 0..array.len().min(10) { assert!( result - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") .is_null() ); @@ -141,10 +141,10 @@ fn test_cast_to_non_nullable(array: &ArrayRef) { for i in 0..array.len().min(10) { assert_eq!( array - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), non_nullable - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -157,10 +157,10 @@ fn test_cast_to_non_nullable(array: &ArrayRef) { for i in 0..array.len().min(10) { assert_eq!( array - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), back_to_nullable - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -190,10 +190,10 @@ fn test_cast_to_nullable(array: &ArrayRef) { for i in 0..array.len().min(10) { assert_eq!( array - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), nullable - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -206,9 +206,9 @@ fn test_cast_to_nullable(array: &ArrayRef) { for i in 0..array.len().min(10) { assert_eq!( array - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), - back.scalar_at(i) + back.scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -294,10 +294,10 @@ fn test_cast_to_primitive(array: &ArrayRef, target_ptype: PType, test_round_trip ); for i in 0..array.len().min(10) { let original = array - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); let casted = casted - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( original diff --git a/vortex-array/src/compute/conformance/consistency.rs b/vortex-array/src/compute/conformance/consistency.rs index dec1f8d1d61..8e1c0ee3f55 100644 --- a/vortex-array/src/compute/conformance/consistency.rs +++ b/vortex-array/src/compute/conformance/consistency.rs @@ -89,10 +89,10 @@ fn test_filter_take_consistency(array: &ArrayRef) { for i in 0..filtered.len() { let filtered_val = filtered - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); let taken_val = taken - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( filtered_val, taken_val, @@ -163,10 +163,10 @@ fn test_double_mask_consistency(array: &ArrayRef) { for i in 0..double_masked.len() { let double_val = double_masked - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); let direct_val = directly_masked - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( double_val, direct_val, @@ -213,10 +213,10 @@ fn test_filter_identity(array: &ArrayRef) { for i in 0..len { let original_val = array - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); let filtered_val = filtered - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( filtered_val, original_val, @@ -269,10 +269,10 @@ fn test_mask_identity(array: &ArrayRef) { for i in 0..len { let original_val = array - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); let masked_val = masked - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); let expected_val = original_val.clone().into_nullable(); assert_eq!( @@ -328,10 +328,10 @@ fn test_slice_filter_consistency(array: &ArrayRef) { for i in 0..filtered.len() { let filtered_val = filtered - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); let sliced_val = sliced - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( filtered_val, sliced_val, @@ -383,10 +383,10 @@ fn test_take_slice_consistency(array: &ArrayRef) { for i in 0..taken.len() { let taken_val = taken - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); let sliced_val = sliced - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( taken_val, sliced_val, @@ -416,26 +416,26 @@ fn test_filter_preserves_order(array: &ArrayRef) { if len >= 4 { assert_eq!( filtered - .scalar_at(0) + .scalar_at(0, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), array - .scalar_at(0) + .scalar_at(0, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); assert_eq!( filtered - .scalar_at(1) + .scalar_at(1, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), array - .scalar_at(2) + .scalar_at(2, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); assert_eq!( filtered - .scalar_at(2) + .scalar_at(2, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), array - .scalar_at(3) + .scalar_at(3, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -458,10 +458,10 @@ fn test_take_repeated_indices(array: &ArrayRef) { for i in 0..3 { assert_eq!( taken - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), array - .scalar_at(0) + .scalar_at(0, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -498,10 +498,10 @@ fn test_mask_filter_null_consistency(array: &ArrayRef) { for i in 0..filtered.len() { assert_eq!( filtered - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), direct_filtered - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -555,10 +555,10 @@ fn test_take_preserves_properties(array: &ArrayRef) { for i in 0..len { assert_eq!( taken - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), array - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -610,11 +610,11 @@ fn test_nullable_indices_consistency(array: &ArrayRef) { // Check first element (from index 0) let expected_0 = array - .scalar_at(0) + .scalar_at(0, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") .into_nullable(); let actual_0 = taken - .scalar_at(0) + .scalar_at(0, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( actual_0, expected_0, @@ -624,7 +624,7 @@ fn test_nullable_indices_consistency(array: &ArrayRef) { // Check second element (should be null) let actual_1 = taken - .scalar_at(1) + .scalar_at(1, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); assert!( actual_1.is_null(), @@ -633,11 +633,11 @@ fn test_nullable_indices_consistency(array: &ArrayRef) { // Check third element (from index 2) let expected_2 = array - .scalar_at(2) + .scalar_at(2, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") .into_nullable(); let actual_2 = taken - .scalar_at(2) + .scalar_at(2, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( actual_2, expected_2, @@ -672,10 +672,10 @@ fn test_large_array_consistency(array: &ArrayRef) { for i in 0..taken.len() { assert_eq!( taken - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), filtered - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -714,7 +714,7 @@ fn test_comparison_inverse_consistency(array: &ArrayRef) { return; } else { array - .scalar_at(len / 2) + .scalar_at(len / 2, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") }; @@ -740,10 +740,10 @@ fn test_comparison_inverse_consistency(array: &ArrayRef) { for i in 0..inverted_eq.len() { let inv_val = inverted_eq - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); let neq_val = neq_result - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( inv_val, neq_val, @@ -768,10 +768,10 @@ fn test_comparison_inverse_consistency(array: &ArrayRef) { for i in 0..inverted_gt.len() { let inv_val = inverted_gt - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); let lte_val = lte_result - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( inv_val, lte_val, @@ -796,10 +796,10 @@ fn test_comparison_inverse_consistency(array: &ArrayRef) { for i in 0..inverted_lt.len() { let inv_val = inverted_lt - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); let gte_val = gte_result - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( inv_val, gte_val, @@ -842,7 +842,7 @@ fn test_comparison_symmetry_consistency(array: &ArrayRef) { return; } else { array - .scalar_at(len / 2) + .scalar_at(len / 2, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") }; @@ -867,10 +867,10 @@ fn test_comparison_symmetry_consistency(array: &ArrayRef) { for i in 0..arr_gt_scalar.len() { let arr_gt = arr_gt_scalar - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); let scalar_lt = scalar_lt_arr - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( arr_gt, scalar_lt, @@ -889,10 +889,10 @@ fn test_comparison_symmetry_consistency(array: &ArrayRef) { ) { for i in 0..arr_eq_scalar.len() { let arr_eq = arr_eq_scalar - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); let scalar_eq = scalar_eq_arr - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( arr_eq, scalar_eq, @@ -951,10 +951,10 @@ fn test_boolean_demorgan_consistency(array: &ArrayRef) { for i in 0..not_a_and_b.len() { let left = not_a_and_b - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); let right = not_a_or_not_b - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( left, right, @@ -979,10 +979,10 @@ fn test_boolean_demorgan_consistency(array: &ArrayRef) { for i in 0..not_a_or_b.len() { let left = not_a_or_b - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); let right = not_a_and_not_b - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( left, right, @@ -1274,14 +1274,14 @@ fn test_cast_slice_consistency(array: &ArrayRef) { // Compare each value against the canonical form for i in 0..slice_then_cast.len() { let slice_cast_val = slice_then_cast - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); // Get the corresponding value from the canonical array (adjusted for slice offset) let canonical_val = canonical .clone() .into_array() - .scalar_at(start + i) + .scalar_at(start + i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); // Cast the canonical scalar to the target dtype @@ -1327,10 +1327,10 @@ fn test_cast_slice_consistency(array: &ArrayRef) { for i in 0..slice_then_cast.len() { let slice_cast_val = slice_then_cast - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); let cast_slice_val = cast_then_slice - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( slice_cast_val, cast_slice_val, diff --git a/vortex-array/src/compute/conformance/filter.rs b/vortex-array/src/compute/conformance/filter.rs index 6be08dff69a..80093ca045b 100644 --- a/vortex-array/src/compute/conformance/filter.rs +++ b/vortex-array/src/compute/conformance/filter.rs @@ -6,6 +6,8 @@ use vortex_mask::Mask; use crate::ArrayRef; use crate::IntoArray; +use crate::LEGACY_SESSION; +use crate::VortexSessionExecute as _; use crate::assert_arrays_eq; use crate::dtype::DType; @@ -94,14 +96,16 @@ fn test_selective_filter(array: &ArrayRef) { .vortex_expect("filter should succeed in conformance test"); assert_eq!(filtered.len(), expected_count); + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + // Verify correct elements are kept for (filtered_idx, i) in (0..len).step_by(2).enumerate() { assert_eq!( filtered - .scalar_at(filtered_idx) + .scalar_at(filtered_idx, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test"), array - .scalar_at(i) + .scalar_at(i, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -118,18 +122,18 @@ fn test_selective_filter(array: &ArrayRef) { assert_eq!(filtered.len(), 2); assert_eq!( filtered - .scalar_at(0) + .scalar_at(0, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test"), array - .scalar_at(0) + .scalar_at(0, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test") ); assert_eq!( filtered - .scalar_at(1) + .scalar_at(1, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test"), array - .scalar_at(len - 1) + .scalar_at(len - 1, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -141,6 +145,8 @@ fn test_single_element_filter(array: &ArrayRef) { return; } + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + // Test selecting only the first element let mut mask_values = vec![false; len]; mask_values[0] = true; @@ -151,10 +157,10 @@ fn test_single_element_filter(array: &ArrayRef) { assert_eq!(filtered.len(), 1); assert_eq!( filtered - .scalar_at(0) + .scalar_at(0, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test"), array - .scalar_at(0) + .scalar_at(0, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test") ); @@ -169,10 +175,10 @@ fn test_single_element_filter(array: &ArrayRef) { assert_eq!(filtered.len(), 1); assert_eq!( filtered - .scalar_at(0) + .scalar_at(0, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test"), array - .scalar_at(len - 1) + .scalar_at(len - 1, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -229,16 +235,18 @@ fn test_alternating_pattern_filter(array: &ArrayRef) { .vortex_expect("filter should succeed in conformance test"); assert_eq!(filtered.len(), expected_count); + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + // Verify correct elements are kept let mut filtered_idx = 0; for (i, &keep) in pattern.iter().enumerate() { if keep { assert_eq!( filtered - .scalar_at(filtered_idx) + .scalar_at(filtered_idx, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test"), array - .scalar_at(i) + .scalar_at(i, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test") ); filtered_idx += 1; diff --git a/vortex-array/src/compute/conformance/mask.rs b/vortex-array/src/compute/conformance/mask.rs index debdff9a806..ea9c44755b2 100644 --- a/vortex-array/src/compute/conformance/mask.rs +++ b/vortex-array/src/compute/conformance/mask.rs @@ -6,6 +6,8 @@ use vortex_mask::Mask; use crate::ArrayRef; use crate::IntoArray; +use crate::LEGACY_SESSION; +use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::bool::BoolArrayExt; use crate::builtins::ArrayBuiltins; @@ -47,21 +49,23 @@ fn test_heterogenous_mask(array: &ArrayRef) { .vortex_expect("mask should succeed in conformance test"); assert_eq!(masked.len(), array.len()); + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + // Verify masked elements are null and unmasked elements are preserved for (i, &masked_out) in mask_pattern.iter().enumerate() { if masked_out { assert!( !masked - .is_valid(i) + .is_valid(i, &mut ctx) .vortex_expect("is_valid should succeed in conformance test") ); } else { assert_eq!( masked - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), array - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") .into_nullable() ); @@ -85,10 +89,10 @@ fn test_empty_mask(array: &ArrayRef) { for i in 0..len { assert_eq!( masked - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), array - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") .into_nullable() ); @@ -107,11 +111,13 @@ fn test_full_mask(array: &ArrayRef) { .vortex_expect("mask should succeed in conformance test"); assert_eq!(masked.len(), array.len()); + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + // All elements should be null for i in 0..len { assert!( !masked - .is_valid(i) + .is_valid(i, &mut ctx) .vortex_expect("is_valid should succeed in conformance test") ); } @@ -129,20 +135,22 @@ fn test_alternating_mask(array: &ArrayRef) { .vortex_expect("mask should succeed in conformance test"); assert_eq!(masked.len(), array.len()); + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + for i in 0..len { if i % 2 == 0 { assert!( !masked - .is_valid(i) + .is_valid(i, &mut ctx) .vortex_expect("is_valid should succeed in conformance test") ); } else { assert_eq!( masked - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), array - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") .into_nullable() ); @@ -167,11 +175,13 @@ fn test_sparse_mask(array: &ArrayRef) { .vortex_expect("mask should succeed in conformance test"); assert_eq!(masked.len(), array.len()); + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + // Count how many elements are valid after masking let valid_count = (0..len) .filter(|&i| { masked - .is_valid(i) + .is_valid(i, &mut ctx) .vortex_expect("is_valid should succeed in conformance test") }) .count(); @@ -183,7 +193,7 @@ fn test_sparse_mask(array: &ArrayRef) { .filter(|&i| { pattern[i] || !array - .is_valid(i) + .is_valid(i, &mut ctx) .vortex_expect("is_valid should succeed in conformance test") }) .count(); @@ -200,23 +210,25 @@ fn test_single_element_mask(array: &ArrayRef) { pattern[0] = true; let mask_array = Mask::from_iter(pattern); + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let masked = array .clone() .mask((!&mask_array).into_array()) .vortex_expect("mask should succeed in conformance test"); assert!( !masked - .is_valid(0) + .is_valid(0, &mut ctx) .vortex_expect("is_valid should succeed in conformance test") ); for i in 1..len { assert_eq!( masked - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), array - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") .into_nullable() ); @@ -242,21 +254,23 @@ fn test_double_mask(array: &ArrayRef) { .mask((!&mask2).into_array()) .vortex_expect("mask should succeed in conformance test"); + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + // Elements should be null if either mask is true for i in 0..len { if mask1_pattern[i] || mask2_pattern[i] { assert!( !double_masked - .is_valid(i) + .is_valid(i, &mut ctx) .vortex_expect("is_valid should succeed in conformance test") ); } else { assert_eq!( double_masked - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), array - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") .into_nullable() ); @@ -285,21 +299,23 @@ fn test_nullable_mask_input(array: &ArrayRef) { .mask((!&mask_array).into_array()) .vortex_expect("mask should succeed in conformance test"); + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + // Elements are masked only if the mask is true AND valid for i in 0..len { if bool_values[i] && validity_values[i] { assert!( !masked - .is_valid(i) + .is_valid(i, &mut ctx) .vortex_expect("is_valid should succeed in conformance test") ); } else { assert_eq!( masked - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), array - .scalar_at(i) + .scalar_at(i, &mut LEGACY_SESSION.create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") .into_nullable() ); diff --git a/vortex-array/src/compute/conformance/take.rs b/vortex-array/src/compute/conformance/take.rs index 9fec410ac35..9864cbadb33 100644 --- a/vortex-array/src/compute/conformance/take.rs +++ b/vortex-array/src/compute/conformance/take.rs @@ -7,6 +7,8 @@ use vortex_error::VortexExpect; use crate::ArrayRef; use crate::Canonical; use crate::IntoArray as _; +use crate::LEGACY_SESSION; +use crate::VortexSessionExecute as _; use crate::arrays::PrimitiveArray; use crate::dtype::Nullability; @@ -61,6 +63,8 @@ fn test_take_all(array: &ArrayRef) { assert_eq!(result.len(), len); assert_eq!(result.dtype(), array.dtype()); + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + // Verify elements match match ( array @@ -81,10 +85,10 @@ fn test_take_all(array: &ArrayRef) { for i in 0..len { assert_eq!( array - .scalar_at(i) + .scalar_at(i, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test"), result - .scalar_at(i) + .scalar_at(i, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -116,14 +120,16 @@ fn test_take_selective(array: &ArrayRef) { .vortex_expect("take should succeed in conformance test"); assert_eq!(result.len(), expected_len); + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + // Verify the taken elements for (result_idx, &original_idx) in indices.iter().enumerate() { assert_eq!( array - .scalar_at(original_idx as usize) + .scalar_at(original_idx as usize, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test"), result - .scalar_at(result_idx) + .scalar_at(result_idx, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -136,21 +142,23 @@ fn test_take_first_and_last(array: &ArrayRef) { .take(indices.into_array()) .vortex_expect("take should succeed in conformance test"); + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + assert_eq!(result.len(), 2); assert_eq!( array - .scalar_at(0) + .scalar_at(0, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test"), result - .scalar_at(0) + .scalar_at(0, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test") ); assert_eq!( array - .scalar_at(len - 1) + .scalar_at(len - 1, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test"), result - .scalar_at(1) + .scalar_at(1, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -179,22 +187,24 @@ fn test_take_with_nullable_indices(array: &ArrayRef) { &array.dtype().with_nullability(Nullability::Nullable) ); + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + // Verify values for (i, idx_opt) in indices_vec.iter().enumerate() { match idx_opt { Some(idx) => { let expected = array - .scalar_at(*idx as usize) + .scalar_at(*idx as usize, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test"); let actual = result - .scalar_at(i) + .scalar_at(i, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!(expected, actual); } None => { assert!( result - .scalar_at(i) + .scalar_at(i, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test") .is_null() ); @@ -214,14 +224,16 @@ fn test_take_repeated_indices(array: &ArrayRef) { .take(indices) .vortex_expect("take should succeed in conformance test"); + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + assert_eq!(result.len(), 3); let first_elem = array - .scalar_at(0) + .scalar_at(0, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test"); for i in 0..3 { assert_eq!( result - .scalar_at(i) + .scalar_at(i, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test"), first_elem ); @@ -248,14 +260,16 @@ fn test_take_reverse(array: &ArrayRef) { assert_eq!(result.len(), len); + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + // Verify elements are in reverse order for i in 0..len { assert_eq!( array - .scalar_at(len - 1 - i) + .scalar_at(len - 1 - i, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test"), result - .scalar_at(i) + .scalar_at(i, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -270,13 +284,15 @@ fn test_take_single_middle(array: &ArrayRef) { .take(indices.into_array()) .vortex_expect("take should succeed in conformance test"); + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + assert_eq!(result.len(), 1); assert_eq!( array - .scalar_at(middle_idx) + .scalar_at(middle_idx, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test"), result - .scalar_at(0) + .scalar_at(0, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -300,14 +316,16 @@ fn test_take_random_unsorted(array: &ArrayRef) { assert_eq!(result.len(), indices.len()); + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + // Verify elements match for (i, &idx) in indices.iter().enumerate() { assert_eq!( array - .scalar_at(idx as usize) + .scalar_at(idx as usize, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test"), result - .scalar_at(i) + .scalar_at(i, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -326,14 +344,16 @@ fn test_take_contiguous_range(array: &ArrayRef) { assert_eq!(result.len(), end - start); + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + // Verify elements for i in 0..(end - start) { assert_eq!( array - .scalar_at(start + i) + .scalar_at(start + i, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test"), result - .scalar_at(i) + .scalar_at(i, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -362,14 +382,16 @@ fn test_take_mixed_repeated(array: &ArrayRef) { assert_eq!(result.len(), indices.len()); + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + // Verify elements for (i, &idx) in indices.iter().enumerate() { assert_eq!( array - .scalar_at(idx as usize) + .scalar_at(idx as usize, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test"), result - .scalar_at(i) + .scalar_at(i, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -393,15 +415,16 @@ fn test_take_large_indices(array: &ArrayRef) { assert_eq!(result.len(), num_indices); + let mut ctx = LEGACY_SESSION.create_execution_ctx(); // Spot check a few elements for i in (0..num_indices).step_by(1000) { let expected_idx = indices[i] as usize; assert_eq!( array - .scalar_at(expected_idx) + .scalar_at(expected_idx, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test"), result - .scalar_at(i) + .scalar_at(i, &mut ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } diff --git a/vortex-array/src/display/mod.rs b/vortex-array/src/display/mod.rs index f299bb7aeeb..47947691c90 100644 --- a/vortex-array/src/display/mod.rs +++ b/vortex-array/src/display/mod.rs @@ -19,6 +19,9 @@ use itertools::Itertools as _; pub use tree_display::TreeDisplay; use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::LEGACY_SESSION; +use crate::VortexSessionExecute; /// Describe how to convert an array to a string. /// @@ -519,7 +522,13 @@ impl ArrayRef { DisplayArrayAs(self, DisplayOptions::TableDisplay) } + fn fmt_scalar(&self, i: usize, ctx: &mut ExecutionCtx) -> String { + self.scalar_at(i, ctx) + .map_or_else(|e| format!(""), |s| s.to_string()) + } + fn fmt_as(&self, f: &mut std::fmt::Formatter, options: &DisplayOptions) -> std::fmt::Result { + let mut ctx = LEGACY_SESSION.create_execution_ctx(); match options { DisplayOptions::MetadataOnly => EncodingSummaryExtractor::write(self, f), DisplayOptions::CommaSeparatedScalars { @@ -531,24 +540,28 @@ impl ArrayRef { let sep = if *omit_comma_after_space { "," } else { ", " }; let sep = if f.alternate() { ",\n" } else { sep }; let limit = self.len().min(f.precision().unwrap_or(DISPLAY_LIMIT)); + let is_truncated = self.len() > limit; - let fmt_scalar = |i| { - self.scalar_at(i) - .map_or_else(|e| format!(""), |s| s.to_string()) - }; + write!(f, "{opening_brace}")?; write!( f, - "{opening_brace}{}{closing_brace}", + "{}", (0..limit.saturating_sub(3)) - .map(fmt_scalar) - .chain(std::iter::repeat_n( - "...".to_string(), - is_truncated as usize - )) - .chain((self.len().saturating_sub(3)..self.len()).map(fmt_scalar)) + .map(|i| self.fmt_scalar(i, &mut ctx)) + .format(sep) + )?; + if is_truncated { + write!(f, "...")?; + } + write!( + f, + "{}", + (self.len().saturating_sub(3)..self.len()) + .map(|i| self.fmt_scalar(i, &mut ctx)) .format(sep) - ) + )?; + write!(f, "{closing_brace}") } DisplayOptions::TreeDisplay { buffers, @@ -587,7 +600,7 @@ impl ArrayRef { // For non-struct arrays, simply display a single column table without header. for row_idx in 0..self.len() { let value = self - .scalar_at(row_idx) + .scalar_at(row_idx, &mut ctx) .map_or_else(|e| format!(""), |s| s.to_string()); builder.push_record([value]); } @@ -602,7 +615,7 @@ impl ArrayRef { builder.push_record(sf.names().iter().map(|name| name.to_string())); for row_idx in 0..self.len() { - if !self.is_valid(row_idx).unwrap_or(false) { + if !self.is_valid(row_idx, &mut ctx).unwrap_or(false) { let null_row = vec!["null".to_string(); sf.names().len()]; builder.push_record(null_row); } else { @@ -611,7 +624,7 @@ impl ArrayRef { crate::arrays::struct_::StructArrayExt::iter_unmasked_fields(&struct_) { let value = field_array - .scalar_at(row_idx) + .scalar_at(row_idx, &mut ctx) .map_or_else(|e| format!(""), |s| s.to_string()); row.push(value); } @@ -628,7 +641,7 @@ impl ArrayRef { } for row_idx in 0..self.len() { - if !self.is_valid(row_idx).unwrap_or(false) { + if !self.is_valid(row_idx, &mut ctx).unwrap_or(false) { table.modify( (1 + row_idx, 0), tabled::settings::Span::column(sf.names().len() as isize), diff --git a/vortex-array/src/patches.rs b/vortex-array/src/patches.rs index d25d6446ce6..3e475441e1c 100644 --- a/vortex-array/src/patches.rs +++ b/vortex-array/src/patches.rs @@ -21,7 +21,9 @@ use vortex_utils::aliases::hash_map::HashMap; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; +use crate::LEGACY_SESSION; use crate::ToCanonical; +use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; @@ -172,7 +174,7 @@ impl Patches { // Perform validation of components when they are host-resident. // This is not possible to do eagerly when the data is on GPU memory. if indices.is_host() && values.is_host() { - let max = usize::try_from(&indices.scalar_at(indices.len() - 1)?) + let max = usize::try_from(&indices.scalar_at(indices.len() - 1, &mut LEGACY_SESSION.create_execution_ctx())?) .map_err(|_| vortex_err!("indices must be a number"))?; vortex_ensure!( max - offset < array_len, @@ -181,8 +183,7 @@ impl Patches { #[cfg(debug_assertions)] { - use crate::VortexSessionExecute; - let mut ctx = crate::LEGACY_SESSION.create_execution_ctx(); + let mut ctx = LEGACY_SESSION.create_execution_ctx(); assert!( crate::aggregate_fn::fns::is_sorted::is_sorted(&indices, &mut ctx) .unwrap_or(false), @@ -292,7 +293,7 @@ impl Patches { }; chunk_offsets - .scalar_at(idx)? + .scalar_at(idx, &mut LEGACY_SESSION.create_execution_ctx())? .as_primitive() .as_::() .ok_or_else(|| vortex_err!("chunk offset does not fit in usize")) @@ -363,7 +364,7 @@ impl Patches { pub fn get_patched(&self, index: usize) -> VortexResult> { self.search_index(index)? .to_found() - .map(|patch_idx| self.values().scalar_at(patch_idx)) + .map(|patch_idx| self.values().scalar_at(patch_idx, &mut LEGACY_SESSION.create_execution_ctx())) .transpose() } @@ -545,7 +546,7 @@ impl Patches { pub fn min_index(&self) -> VortexResult { let first = self .indices - .scalar_at(0)? + .scalar_at(0, &mut LEGACY_SESSION.create_execution_ctx())? .as_primitive() .as_::() .ok_or_else(|| vortex_err!("index does not fit in usize"))?; @@ -556,7 +557,7 @@ impl Patches { pub fn max_index(&self) -> VortexResult { let last = self .indices - .scalar_at(self.indices.len() - 1)? + .scalar_at(self.indices.len() - 1, &mut LEGACY_SESSION.create_execution_ctx())? .as_primitive() .as_::() .ok_or_else(|| vortex_err!("index does not fit in usize"))?; @@ -665,7 +666,7 @@ impl Patches { .as_ref() .map(|new_chunk_offsets| -> VortexResult { let new_chunk_base = new_chunk_offsets - .scalar_at(0)? + .scalar_at(0, &mut LEGACY_SESSION.create_execution_ctx())? .as_primitive() .as_::() .ok_or_else(|| vortex_err!("chunk offset does not fit in usize"))?; diff --git a/vortex-array/src/scalar_fn/fns/between/mod.rs b/vortex-array/src/scalar_fn/fns/between/mod.rs index aaec8108808..03bad4417ec 100644 --- a/vortex-array/src/scalar_fn/fns/between/mod.rs +++ b/vortex-array/src/scalar_fn/fns/between/mod.rs @@ -17,6 +17,8 @@ use crate::ArrayRef; use crate::Canonical; use crate::ExecutionCtx; use crate::IntoArray; +use crate::LEGACY_SESSION; +use crate::VortexSessionExecute; use crate::arrays::ConstantArray; use crate::arrays::Decimal; use crate::arrays::Primitive; @@ -107,7 +109,7 @@ pub(super) fn precondition( } // A quick check to see if either bound is a null constant array. - if (lower.is_invalid(0)? || upper.is_invalid(0)?) + if (lower.is_invalid(0, &mut LEGACY_SESSION.create_execution_ctx())? || upper.is_invalid(0, &mut LEGACY_SESSION.create_execution_ctx())?) && let (Some(c_lower), Some(c_upper)) = (lower.as_constant(), upper.as_constant()) && (c_lower.is_null() || c_upper.is_null()) { diff --git a/vortex-array/src/scalar_fn/fns/case_when.rs b/vortex-array/src/scalar_fn/fns/case_when.rs index f30a17ee580..0d1fd5943b3 100644 --- a/vortex-array/src/scalar_fn/fns/case_when.rs +++ b/vortex-array/src/scalar_fn/fns/case_when.rs @@ -246,7 +246,7 @@ impl ScalarFnVTable for CaseWhen { return Ok(else_value); } - merge_case_branches(branches, else_value) + merge_case_branches(branches, else_value, ctx) } fn is_null_sensitive(&self, _options: &Self::Options) -> bool { @@ -268,6 +268,7 @@ const SLICE_CROSSOVER_RUN_LEN: usize = 4; fn merge_case_branches( branches: Vec<(Mask, ArrayRef)>, else_value: ArrayRef, + ctx: &mut ExecutionCtx, ) -> VortexResult { if branches.len() == 1 { let (mask, then_value) = &branches[0]; @@ -304,7 +305,7 @@ fn merge_case_branches( let fragmented = spans.len() > else_value.len() / SLICE_CROSSOVER_RUN_LEN; if fragmented { - merge_row_by_row(&branch_arrays, &else_value, &spans, &output_dtype, builder) + merge_row_by_row(&branch_arrays, &else_value, &spans, &output_dtype, builder, ctx) } else { merge_run_by_run(&branch_arrays, &else_value, &spans, &output_dtype, builder) } @@ -318,21 +319,22 @@ fn merge_row_by_row( spans: &[(usize, usize, usize)], output_dtype: &DType, mut builder: Box, + ctx: &mut ExecutionCtx, ) -> VortexResult { let mut pos = 0; for &(start, end, branch_idx) in spans { for row in pos..start { - let scalar = else_value.scalar_at(row)?; + let scalar = else_value.scalar_at(row, ctx)?; builder.append_scalar(&scalar.cast(output_dtype)?)?; } for row in start..end { - let scalar = branch_arrays[branch_idx].scalar_at(row)?; + let scalar = branch_arrays[branch_idx].scalar_at(row, ctx)?; builder.append_scalar(&scalar.cast(output_dtype)?)?; } pos = end; } for row in pos..else_value.len() { - let scalar = else_value.scalar_at(row)?; + let scalar = else_value.scalar_at(row, ctx)?; builder.append_scalar(&scalar.cast(output_dtype)?)?; } diff --git a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs index b4d6c3e6bc4..c9b8a13653f 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs @@ -280,7 +280,7 @@ fn list_contains_scalar( // If the list array is constant, we perform a single comparison. if array.len() > 1 && array.is::() { let contains = list_contains_scalar(&array.slice(0..1)?, value, nullability, ctx)?; - return Ok(ConstantArray::new(contains.scalar_at(0)?, array.len()).into_array()); + return Ok(ConstantArray::new(contains.scalar_at(0, ctx)?, array.len()).into_array()); } let list_array = array.clone().execute::(ctx)?; diff --git a/vortex-array/src/search_sorted.rs b/vortex-array/src/search_sorted.rs index cf096bbb8e2..cb3366b72a3 100644 --- a/vortex-array/src/search_sorted.rs +++ b/vortex-array/src/search_sorted.rs @@ -13,6 +13,8 @@ use std::hint; use vortex_error::VortexResult; use crate::ArrayRef; +use crate::LEGACY_SESSION; +use crate::VortexSessionExecute; use crate::scalar::Scalar; #[derive(Debug, Copy, Clone, Eq, PartialEq)] @@ -264,7 +266,7 @@ fn search_sorted_side_idx VortexResult>( impl IndexOrd for ArrayRef { fn index_cmp(&self, idx: usize, elem: &Scalar) -> VortexResult> { - let scalar_a = self.scalar_at(idx)?; + let scalar_a = self.scalar_at(idx, &mut LEGACY_SESSION.create_execution_ctx())?; Ok(scalar_a.partial_cmp(elem)) } diff --git a/vortex-array/src/validity.rs b/vortex-array/src/validity.rs index a5f6903e712..28cb21d0401 100644 --- a/vortex-array/src/validity.rs +++ b/vortex-array/src/validity.rs @@ -123,12 +123,12 @@ impl Validity { /// Returns whether the `index` item is valid. #[inline] - pub fn is_valid(&self, index: usize) -> VortexResult { + pub fn is_valid(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult { Ok(match self { Self::NonNullable | Self::AllValid => true, Self::AllInvalid => false, Self::Array(a) => a - .scalar_at(index) + .scalar_at(index, ctx) .vortex_expect("Validity array must support scalar_at") .as_bool() .value() @@ -137,8 +137,8 @@ impl Validity { } #[inline] - pub fn is_null(&self, index: usize) -> VortexResult { - Ok(!self.is_valid(index)?) + pub fn is_null(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult { + Ok(!self.is_valid(index, ctx)?) } #[inline] diff --git a/vortex-array/src/variants.rs b/vortex-array/src/variants.rs index f60eea316a9..46bccbd4840 100644 --- a/vortex-array/src/variants.rs +++ b/vortex-array/src/variants.rs @@ -133,7 +133,7 @@ impl PrimitiveTyped<'_> { /// Return the primitive value at the given index. pub fn value(&self, idx: usize) -> VortexResult> { self.0 - .is_valid(idx)? + .is_valid(idx, &mut LEGACY_SESSION.create_execution_ctx())? .then(|| self.value_unchecked(idx)) .transpose() } @@ -142,7 +142,7 @@ impl PrimitiveTyped<'_> { pub fn value_unchecked(&self, idx: usize) -> VortexResult { Ok(self .0 - .scalar_at(idx)? + .scalar_at(idx, &mut LEGACY_SESSION.create_execution_ctx())? .as_primitive() .pvalue() .unwrap_or_else(|| PValue::zero(&self.ptype()))) diff --git a/vortex-btrblocks/src/schemes/patches.rs b/vortex-btrblocks/src/schemes/patches.rs index 8a82058ddb7..a6e2d764ba5 100644 --- a/vortex-btrblocks/src/schemes/patches.rs +++ b/vortex-btrblocks/src/schemes/patches.rs @@ -3,7 +3,9 @@ use vortex_array::ArrayRef; use vortex_array::IntoArray; +use vortex_array::LEGACY_SESSION; use vortex_array::ToCanonical; +use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::patches::Patches; @@ -15,6 +17,8 @@ pub fn compress_patches(patches: Patches) -> VortexResult { // Downscale the patch indices. let indices = patches.indices().to_primitive().narrow()?.into_array(); + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + // Check if the values are constant. let values = patches.values(); let values = if values @@ -22,7 +26,7 @@ pub fn compress_patches(patches: Patches) -> VortexResult { .compute_is_constant() .unwrap_or_default() { - ConstantArray::new(values.scalar_at(0)?, values.len()).into_array() + ConstantArray::new(values.scalar_at(0, &mut ctx)?, values.len()).into_array() } else { values.clone() }; diff --git a/vortex-btrblocks/src/schemes/temporal.rs b/vortex-btrblocks/src/schemes/temporal.rs index 47d9ae81cfb..b4c306c9589 100644 --- a/vortex-btrblocks/src/schemes/temporal.rs +++ b/vortex-btrblocks/src/schemes/temporal.rs @@ -6,7 +6,9 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray; +use vortex_array::LEGACY_SESSION; use vortex_array::ToCanonical; +use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::fns::is_constant::is_constant; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::TemporalArray; @@ -83,7 +85,11 @@ impl Scheme for TemporalScheme { )?; if is_constant { - return Ok(ConstantArray::new(ext_array.scalar_at(0)?, ext_array.len()).into_array()); + let mut exec_ctx = LEGACY_SESSION.create_execution_ctx(); + return Ok( + ConstantArray::new(ext_array.scalar_at(0, &mut exec_ctx)?, ext_array.len()) + .into_array(), + ); } let dtype = temporal_array.dtype().clone(); diff --git a/vortex-compressor/src/builtins/constant/mod.rs b/vortex-compressor/src/builtins/constant/mod.rs index 1b177fc530b..2941c154460 100644 --- a/vortex-compressor/src/builtins/constant/mod.rs +++ b/vortex-compressor/src/builtins/constant/mod.rs @@ -5,6 +5,8 @@ use vortex_array::ArrayRef; use vortex_array::IntoArray; +use vortex_array::LEGACY_SESSION; +use vortex_array::VortexSessionExecute as _; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::MaskedArray; use vortex_array::scalar::Scalar; @@ -49,11 +51,13 @@ fn compress_constant_array_with_validity(source: &ArrayRef) -> VortexResult ColumnExporter for RunEndExporter { if start_run_idx == end_run_idx { // NOTE(ngates): would be great if we could just export and set type == CONSTANT // self.values_exporter.export(start_run_idx, 1, vector, cache); - let constant = self.values.scalar_at(start_run_idx)?; + let constant = self.values.scalar_at(start_run_idx, ctx)?; let value = constant.try_to_duckdb_scalar()?; vector.reference_value(&value); return Ok(()); diff --git a/vortex-ffi/src/array.rs b/vortex-ffi/src/array.rs index 5f684622ee6..8c647b0ec38 100644 --- a/vortex-ffi/src/array.rs +++ b/vortex-ffi/src/array.rs @@ -10,7 +10,9 @@ use std::sync::Arc; use paste::paste; use vortex::array::ArrayRef; use vortex::array::IntoArray; +use vortex::array::LEGACY_SESSION; use vortex::array::ToCanonical; +use vortex::array::VortexSessionExecute as _; use vortex::array::arrays::NullArray; use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::struct_::StructArrayExt; @@ -238,7 +240,8 @@ pub unsafe extern "C-unwind" fn vx_array_element_is_invalid( ) -> bool { try_or_default(error, || { vortex_ensure!(!array.is_null()); - vx_array::as_ref(array).is_invalid(index) + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + vx_array::as_ref(array).is_invalid(index, &mut ctx) }) } @@ -324,9 +327,10 @@ macro_rules! ffiarray_get_ptype { paste! { #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn [](array: *const vx_array, index: usize) -> $ptype { + let mut ctx = LEGACY_SESSION.create_execution_ctx(); let array = vx_array::as_ref(array); // TODO(joe): propagate this error up instead of expecting - let value = array.scalar_at(index).vortex_expect("scalar_at failed"); + let value = array.scalar_at(index, &mut ctx).vortex_expect("scalar_at failed"); // TODO(joe): propagate this error up instead of expecting value.as_primitive() .as_::<$ptype>() @@ -335,9 +339,10 @@ macro_rules! ffiarray_get_ptype { #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn [](array: *const vx_array, index: usize) -> $ptype { + let mut ctx = LEGACY_SESSION.create_execution_ctx(); let array = vx_array::as_ref(array); // TODO(joe): propagate this error up instead of expecting - let value = array.scalar_at(index).vortex_expect("scalar_at failed"); + let value = array.scalar_at(index, &mut ctx).vortex_expect("scalar_at failed"); // TODO(joe): propagate this error up instead of expecting value.as_extension() .to_storage_scalar() @@ -368,10 +373,11 @@ pub unsafe extern "C-unwind" fn vx_array_get_utf8( array: *const vx_array, index: u32, ) -> *const vx_string { + let mut ctx = LEGACY_SESSION.create_execution_ctx(); let array = vx_array::as_ref(array); // TODO(joe): propagate this error up instead of expecting let value = array - .scalar_at(index as usize) + .scalar_at(index as usize, &mut ctx) .vortex_expect("scalar_at failed"); let utf8_scalar = value.as_utf8(); if let Some(buffer) = utf8_scalar.value() { @@ -388,10 +394,11 @@ pub unsafe extern "C-unwind" fn vx_array_get_binary( array: *const vx_array, index: u32, ) -> *const vx_binary { + let mut ctx = LEGACY_SESSION.create_execution_ctx(); let array = vx_array::as_ref(array); // TODO(joe): propagate this error up instead of expecting let value = array - .scalar_at(index as usize) + .scalar_at(index as usize, &mut ctx) .vortex_expect("scalar_at failed"); let binary_scalar = value.as_binary(); if let Some(bytes) = binary_scalar.value() { diff --git a/vortex-file/src/v2/file_stats_reader.rs b/vortex-file/src/v2/file_stats_reader.rs index fb31b28d61f..6369569ca86 100644 --- a/vortex-file/src/v2/file_stats_reader.rs +++ b/vortex-file/src/v2/file_stats_reader.rs @@ -102,7 +102,7 @@ impl FileStatsLayoutReader { .execute::(&mut ctx)? .into_bool() .into_array() - .scalar_at(0)?; + .scalar_at(0, &mut ctx)?; Ok(result.as_bool().value() == Some(true)) } diff --git a/vortex-jni/src/array.rs b/vortex-jni/src/array.rs index c45f232b920..0c12d6cda4b 100644 --- a/vortex-jni/src/array.rs +++ b/vortex-jni/src/array.rs @@ -28,7 +28,9 @@ use jni::sys::jshort; use jni::sys::jstring; use vortex::array::ArrayRef; use vortex::array::ArrayView; +use vortex::array::LEGACY_SESSION; use vortex::array::ToCanonical; +use vortex::array::VortexSessionExecute as _; use vortex::array::arrays::VarBin; use vortex::array::arrays::VarBinView; use vortex::array::arrays::extension::ExtensionArrayExt; @@ -267,7 +269,8 @@ pub extern "system" fn Java_dev_vortex_jni_NativeArrayMethods_getNull( ) -> jboolean { let array_ref = unsafe { NativeArray::from_ptr(array_ptr) }; try_or_throw(&mut env, |_| { - let is_null = array_ref.inner.is_invalid(index as usize)?; + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let is_null = array_ref.inner.is_invalid(index as usize, &mut ctx)?; if is_null { Ok(JNI_TRUE) } else { Ok(JNI_FALSE) } }) } @@ -296,14 +299,15 @@ macro_rules! get_primitive { ) -> $jtype { let array_ref = unsafe { NativeArray::from_ptr(array_ptr) }; try_or_throw(&mut env, |_| { + let mut ctx = LEGACY_SESSION.create_execution_ctx(); let scalar_value = if array_ref.is_extension { array_ref .inner .to_extension() .storage_array() - .scalar_at(index as usize)? + .scalar_at(index as usize, &mut ctx)? } else { - array_ref.inner.scalar_at(index as usize)? + array_ref.inner.scalar_at(index as usize, &mut ctx)? }; Ok(scalar_value @@ -335,14 +339,15 @@ pub extern "system" fn Java_dev_vortex_jni_NativeArrayMethods_getBigDecimal( ) -> jobject { let array_ref = unsafe { NativeArray::from_ptr(array_ptr) }; try_or_throw(&mut env, |env| { + let mut ctx = LEGACY_SESSION.create_execution_ctx(); let scalar_value = if array_ref.is_extension { array_ref .inner .to_extension() .storage_array() - .scalar_at(index as usize)? + .scalar_at(index as usize, &mut ctx)? } else { - array_ref.inner.scalar_at(index as usize)? + array_ref.inner.scalar_at(index as usize, &mut ctx)? }; let decimal_scalar = scalar_value.as_decimal(); @@ -404,7 +409,8 @@ pub extern "system" fn Java_dev_vortex_jni_NativeArrayMethods_getBool( ) -> jboolean { let array_ref = unsafe { NativeArray::from_ptr(array_ptr) }; try_or_throw(&mut env, |_| { - let value = array_ref.inner.scalar_at(index as usize)?; + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let value = array_ref.inner.scalar_at(index as usize, &mut ctx)?; match value.as_bool().value() { None => Ok(JNI_FALSE), Some(b) => { @@ -427,7 +433,8 @@ pub extern "system" fn Java_dev_vortex_jni_NativeArrayMethods_getUTF8<'local>( ) -> jstring { let array_ref = unsafe { NativeArray::from_ptr(array_ptr) }; try_or_throw(&mut env, |env| { - let value = array_ref.inner.scalar_at(index as usize)?; + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let value = array_ref.inner.scalar_at(index as usize, &mut ctx)?; match value.as_utf8().value() { None => Ok(JObject::null().into_raw()), Some(buf_str) => Ok(env.new_string(buf_str.as_str())?.into_raw()), @@ -475,7 +482,8 @@ pub extern "system" fn Java_dev_vortex_jni_NativeArrayMethods_getBinary<'local>( ) -> jbyteArray { let array_ref = unsafe { NativeArray::from_ptr(array_ptr) }; try_or_throw(&mut env, |env| { - let value = array_ref.inner.scalar_at(index as usize)?; + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let value = array_ref.inner.scalar_at(index as usize, &mut ctx)?; match value.as_binary().value() { None => Ok(JObject::null().into_raw()), Some(buf) => Ok(env.byte_array_from_slice(buf.as_slice())?.into_raw()), diff --git a/vortex-python/src/arrays/mod.rs b/vortex-python/src/arrays/mod.rs index f0cc68e2f34..bc50f4623a4 100644 --- a/vortex-python/src/arrays/mod.rs +++ b/vortex-python/src/arrays/mod.rs @@ -24,7 +24,9 @@ use pyo3::types::PyRangeMethods; use pyo3_bytes::PyBytes; use vortex::array::ArrayRef; use vortex::array::IntoArray; +use vortex::array::LEGACY_SESSION; use vortex::array::ToCanonical; +use vortex::array::VortexSessionExecute as _; use vortex::array::arrays::Chunked; use vortex::array::arrays::bool::BoolArrayExt; use vortex::array::arrays::chunked::ChunkedArrayExt; @@ -610,7 +612,8 @@ impl PyArray { )) .into()); } - Ok(PyScalar::init(py, slf.scalar_at(index)?)?) + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + Ok(PyScalar::init(py, slf.scalar_at(index, &mut ctx)?)?) } /// Filter, permute, and/or repeat elements by their index. diff --git a/vortex-python/src/scan.rs b/vortex-python/src/scan.rs index 11773588add..e4c6dc5bbf4 100644 --- a/vortex-python/src/scan.rs +++ b/vortex-python/src/scan.rs @@ -3,7 +3,7 @@ use pyo3::exceptions::PyIndexError; use pyo3::prelude::*; -use vortex::array::ArrayRef; +use vortex::array::{ArrayRef, LEGACY_SESSION, VortexSessionExecute as _}; use vortex::layout::scan::repeated_scan::RepeatedScan; use crate::RUNTIME; @@ -59,6 +59,8 @@ impl PyRepeatedScan { .into()); } + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + for batch in slf .get() .scan @@ -68,7 +70,7 @@ impl PyRepeatedScan { if array.is_empty() { continue; } - let scalar = array.scalar_at(0)?; + let scalar = array.scalar_at(0, &mut ctx)?; return Ok(PyScalar::init(slf.py(), scalar)?); } diff --git a/vortex-tensor/src/scalar_fns/l2_denorm.rs b/vortex-tensor/src/scalar_fns/l2_denorm.rs index c0f0f54fb12..916f3fc7249 100644 --- a/vortex-tensor/src/scalar_fns/l2_denorm.rs +++ b/vortex-tensor/src/scalar_fns/l2_denorm.rs @@ -375,7 +375,7 @@ pub fn normalize_as_l2_denorm( let total_elements = row_count * tensor_flat_size; let mut elements = BufferMut::::with_capacity(total_elements); for i in 0..row_count { - let is_valid = norms_validity.is_valid(i)?; + let is_valid = norms_validity.is_valid(i, ctx)?; let norm = norm_values[i]; // SAFETY: We allocated `row_count * tensor_flat_size` capacity and push exactly @@ -590,7 +590,7 @@ fn validate_l2_normalized_rows_impl( let stored_norms = norms.as_ref().map(|norms| norms.as_slice::()); for i in 0..row_count { - if !combined_validity.is_valid(i)? { + if !combined_validity.is_valid(i, ctx)? { continue; } diff --git a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/constant.rs b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/constant.rs index 64feb357093..345ac98c624 100644 --- a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/constant.rs +++ b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/constant.rs @@ -18,6 +18,8 @@ use vortex::array::scalar::DecimalValue; use vortex::array::scalar::Scalar; use vortex::array::validity::Validity; use vortex::error::VortexResult; +use vortex_array::LEGACY_SESSION; +use vortex_array::VortexSessionExecute as _; use super::N; use crate::fixtures::FlatLayoutFixture; @@ -38,6 +40,7 @@ impl FlatLayoutFixture for ConstantFixture { } fn build(&self) -> VortexResult { + let mut ctx = LEGACY_SESSION.create_execution_ctx(); let const_i32 = ConstantArray::new(42i32, N); let const_f64 = ConstantArray::new(99.99f64, N); let const_bool = ConstantArray::new(true, N); @@ -75,7 +78,7 @@ impl FlatLayoutFixture for ConstantFixture { Some("UTC".into()), ) .into_array() - .scalar_at(0)?; + .scalar_at(0, &mut ctx)?; let const_timestamp = ConstantArray::new(timestamp_scalar, N); let arr = StructArray::try_new( diff --git a/vortex-tui/src/browse/ui/layouts.rs b/vortex-tui/src/browse/ui/layouts.rs index a00c363610e..a840099128d 100644 --- a/vortex-tui/src/browse/ui/layouts.rs +++ b/vortex-tui/src/browse/ui/layouts.rs @@ -27,7 +27,9 @@ use ratatui::widgets::Table; use ratatui::widgets::Widget; use ratatui::widgets::Wrap; use vortex::array::ArrayRef; +use vortex::array::LEGACY_SESSION; use vortex::array::ToCanonical; +use vortex::array::VortexSessionExecute as _; use vortex::array::arrays::struct_::StructArrayExt; use vortex::error::VortexExpect; use vortex::layout::layouts::flat::Flat; @@ -113,6 +115,7 @@ fn render_layout_header(app: &AppState, area: Rect, buf: &mut Buffer) { /// Render the inner Array for a FlatLayout. fn render_array(app: &AppState, area: Rect, buf: &mut Buffer, is_stats_table: bool) { + let mut ctx = LEGACY_SESSION.create_execution_ctx(); // Array data is loaded eagerly when navigating to a FlatLayout (synchronously on // native, asynchronously on WASM) and cached in AppState. The render loop never // performs I/O. @@ -163,7 +166,7 @@ fn render_array(app: &AppState, area: Rect, buf: &mut Buffer, is_stats_table: bo std::iter::once(Cell::from(Text::from(format!("{chunk_id}")))) .chain(field_arrays.iter().map(|arr| { Cell::from(Text::from( - arr.scalar_at(chunk_id) + arr.scalar_at(chunk_id, &mut ctx) .vortex_expect("scalar_at failed") .to_string(), ))