From 6fe50222cffaa334b7eccbf2e5c63e52ea16eae1 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 18 Oct 2026 00:07:22 +0000 Subject: [PATCH 1/3] Validate MaskedArray with an explicit execution context `MaskedArray::try_new` checked that the child has no nulls using the hidden global `legacy_session()`. Take an `ExecutionCtx` instead, and validate during deserialization with a context from the session. Add `MaskedArray::new_unchecked` and use it in the Masked filter, slice and take rules and in Constant take, where the invariants hold by construction. Constant take no longer executes the indices' validity. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Y1zHpE1K1Y9HvyU4kt24Tq --- .../fastlanes/src/rle/array/rle_compress.rs | 1 + encodings/fsst/src/array.rs | 33 ++------ encodings/fsst/src/canonical.rs | 3 - encodings/fsst/src/compress.rs | 7 +- encodings/fsst/src/compute/filter.rs | 1 - encodings/fsst/src/compute/mod.rs | 1 - encodings/runend/src/array.rs | 83 ++++++++++++------- encodings/zstd/src/zstd_buffers.rs | 11 +-- lang/cpp/examples/reader.cpp | 2 +- lang/cpp/include/vortex/array.hpp | 6 +- lang/cpp/src/array.cpp | 12 +-- lang/cpp/tests/array.cpp | 18 ++-- lang/cpp/tests/arrow.cpp | 2 +- lang/cpp/tests/scan.cpp | 6 +- .../src/arrays/constant/compute/take.rs | 21 ++--- vortex-array/src/arrays/masked/array.rs | 37 +++++++-- .../src/arrays/masked/compute/filter.rs | 14 ++-- .../src/arrays/masked/compute/mask.rs | 9 +- .../src/arrays/masked/compute/slice.rs | 6 +- .../src/arrays/masked/compute/take.rs | 14 ++-- vortex-array/src/arrays/masked/execute.rs | 2 +- vortex-array/src/arrays/masked/tests.rs | 30 +++++-- .../src/arrays/masked/vtable/canonical.rs | 8 ++ vortex-array/src/arrays/masked/vtable/mod.rs | 23 +++-- vortex-array/src/arrays/mod.rs | 1 + vortex-array/src/arrays/null/compute/rules.rs | 2 - vortex-array/src/arrays/null/compute/take.rs | 23 ++--- vortex-array/src/arrays/null/mod.rs | 10 +++ vortex-array/src/arrays/varbinview/array.rs | 40 ++++++--- vortex-array/src/display/mod.rs | 26 ++++-- .../unstable/row/types/element/tuple/tests.rs | 5 +- vortex-btrblocks/src/schemes/string/fsst.rs | 1 - .../src/compressor/constant/mod.rs | 2 +- vortex-cuda/benches/fsst_cuda.rs | 1 - vortex-cuda/src/arrow/canonical.rs | 1 - vortex-cuda/src/kernel/arrays/masked.rs | 2 +- vortex-ffi/cinclude/vortex.h | 5 +- vortex-ffi/src/array.rs | 24 +++--- vortex-ffi/test/array.cpp | 2 +- vortex-ffi/test/scan.cpp | 8 +- vortex-spatial/benches/area.rs | 2 + vortex-spatial/benches/binary_predicates.rs | 15 +++- vortex-spatial/benches/convex_hull.rs | 2 + vortex-spatial/benches/distance.rs | 4 +- vortex-spatial/benches/length.rs | 3 +- vortex-spatial/src/scalar_fn/collect.rs | 11 ++- vortex-spatial/src/scalar_fn/convex_hull.rs | 6 +- vortex-spatial/src/scalar_fn/length.rs | 3 + vortex-tensor/benches/inner_product.rs | 6 +- vortex-tensor/benches/l2_norm.rs | 6 +- .../src/scalar_fns/cosine_similarity.rs | 5 +- vortex-tensor/src/scalar_fns/inner_product.rs | 5 +- vortex-tensor/src/scalar_fns/l2_norm.rs | 5 +- vortex-tensor/src/scalar_fns/l2_normalize.rs | 9 +- .../common_encoding_tree_throughput.rs | 1 - 55 files changed, 346 insertions(+), 240 deletions(-) diff --git a/encodings/fastlanes/src/rle/array/rle_compress.rs b/encodings/fastlanes/src/rle/array/rle_compress.rs index ea288d95571..596e1641658 100644 --- a/encodings/fastlanes/src/rle/array/rle_compress.rs +++ b/encodings/fastlanes/src/rle/array/rle_compress.rs @@ -372,6 +372,7 @@ mod tests { let masked_indices = MaskedArray::try_new( ConstantArray::new(1u16, indices_prim.len()).into_array(), indices_prim.validity()?, + ctx, )? .into_array(); RLE::try_new( diff --git a/encodings/fsst/src/array.rs b/encodings/fsst/src/array.rs index 3f8bd2e5541..4ce0504cd88 100644 --- a/encodings/fsst/src/array.rs +++ b/encodings/fsst/src/array.rs @@ -26,7 +26,6 @@ use vortex_array::EqMode; use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; use vortex_array::TypedArrayRef; -use vortex_array::VortexSessionExecute; use vortex_array::array_slots; use vortex_array::arrays::Primitive; use vortex_array::arrays::VarBin; @@ -40,7 +39,6 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::OffsetBuilderPType; use vortex_array::dtype::PType; -use vortex_array::legacy_session; use vortex_array::match_each_integer_ptype; use vortex_array::match_each_varbin_builder; use vortex_array::require_child; @@ -128,7 +126,6 @@ impl VTable for FSST { *ID } - #[allow(clippy::disallowed_methods)] fn validate( &self, data: &Self::TypedArrayData, @@ -136,9 +133,7 @@ impl VTable for FSST { len: usize, slots: &[Option], ) -> VortexResult<()> { - // TODO(ctx): trait fixes - VTable::validate has a fixed signature. - let mut ctx = legacy_session().create_execution_ctx(); - data.validate(dtype, len, slots, &mut ctx) + data.validate(dtype, len, slots) } fn nbuffers(_array: ArrayView<'_, Self>) -> usize { @@ -236,13 +231,12 @@ impl VTable for FSST { metadata: &[u8], buffers: &[BufferHandle], children: &dyn ArrayChildren, - session: &VortexSession, + _session: &VortexSession, ) -> VortexResult> { let metadata = FSSTMetadata::decode(metadata)?; let symbols = Buffer::::from_byte_buffer(buffers[0].clone().try_to_host_sync()?); let symbol_lengths = Buffer::::from_byte_buffer(buffers[1].clone().try_to_host_sync()?); - let mut ctx = session.create_execution_ctx(); if buffers.len() == 2 { return Self::deserialize_legacy( self, @@ -252,7 +246,6 @@ impl VTable for FSST { &symbols, &symbol_lengths, children, - &mut ctx, ); } @@ -295,7 +288,6 @@ impl VTable for FSST { &uncompressed_lengths, dtype, len, - &mut ctx, )?; let slots = FSSTSlots { uncompressed_lengths, @@ -614,7 +606,6 @@ impl FSST { symbol_lengths: Buffer, codes: VarBinArray, uncompressed_lengths: ArrayRef, - ctx: &mut ExecutionCtx, ) -> VortexResult { let len = codes.len(); FSSTData::validate_parts_from_codes( @@ -624,7 +615,6 @@ impl FSST { &uncompressed_lengths, &dtype, len, - ctx, )?; let slots = FSSTData::make_slots(&codes, &uncompressed_lengths); let codes_bytes = codes.bytes_handle().clone(); @@ -637,7 +627,6 @@ impl FSST { symbol_table: Arc, codes: VarBinArray, uncompressed_lengths: ArrayRef, - ctx: &mut ExecutionCtx, ) -> VortexResult { let len = codes.len(); FSSTData::validate_parts_from_codes( @@ -647,7 +636,6 @@ impl FSST { &uncompressed_lengths, &dtype, len, - ctx, )?; let slots = FSSTData::make_slots(&codes, &uncompressed_lengths); let codes_bytes = codes.bytes_handle().clone(); @@ -668,7 +656,6 @@ impl FSST { symbols: &Buffer, symbol_lengths: &Buffer, children: &dyn ArrayChildren, - ctx: &mut ExecutionCtx, ) -> VortexResult> { if children.len() != 2 { vortex_bail!(InvalidArgument: "Expected 2 children, got {}", children.len()); @@ -699,7 +686,6 @@ impl FSST { &uncompressed_lengths, dtype, len, - ctx, )?; let slots = FSSTData::make_slots(&codes, &uncompressed_lengths); let codes_bytes = codes.bytes_handle().clone(); @@ -792,7 +778,6 @@ impl FSSTData { dtype: &DType, len: usize, slots: &[Option], - ctx: &mut ExecutionCtx, ) -> VortexResult<()> { let fsst_slots = FSSTSlotsView::from_slots(slots); Self::validate_parts( @@ -804,7 +789,6 @@ impl FSSTData { fsst_slots.uncompressed_lengths, dtype, len, - ctx, ) } @@ -819,7 +803,6 @@ impl FSSTData { uncompressed_lengths: &ArrayRef, dtype: &DType, len: usize, - ctx: &mut ExecutionCtx, ) -> VortexResult<()> { vortex_ensure!( matches!(dtype, DType::Binary(_) | DType::Utf8(_)), @@ -864,14 +847,12 @@ impl FSSTData { // checked by the decode plan instead, once `execute` has made it primitive. if codes_bytes.is_on_host() && codes_offsets.is_host() - && codes_offsets.is::() + && let Some(codes_offsets) = codes_offsets.as_opt::() && !codes_offsets.is_empty() { - let last_offset: usize = (&codes_offsets - .execute_scalar(codes_offsets.len() - 1, ctx) - .vortex_expect("offsets must support scalar_at")) - .try_into() - .vortex_expect("Failed to convert offset to usize"); + let last_offset: usize = match_each_integer_ptype!(codes_offsets.ptype(), |O| { + codes_offsets.as_slice::()[codes_offsets.len() - 1].as_() + }); vortex_ensure!( last_offset <= codes_bytes.len(), InvalidArgument: "Last codes offset {} exceeds codes bytes length {}", @@ -917,7 +898,6 @@ impl FSSTData { uncompressed_lengths: &ArrayRef, dtype: &DType, len: usize, - ctx: &mut ExecutionCtx, ) -> VortexResult<()> { Self::validate_parts( symbols, @@ -928,7 +908,6 @@ impl FSSTData { uncompressed_lengths, dtype, len, - ctx, ) } diff --git a/encodings/fsst/src/canonical.rs b/encodings/fsst/src/canonical.rs index 464f0bd4f5a..57856e2dc03 100644 --- a/encodings/fsst/src/canonical.rs +++ b/encodings/fsst/src/canonical.rs @@ -364,7 +364,6 @@ mod tests { encoded.symbol_table(), codes, encoded.uncompressed_lengths().slice(1..3)?, - &mut ctx, )?; let mut builder = VarBinViewBuilder::with_capacity_in(input.dtype().clone(), 2, ctx.allocator().clone()); @@ -389,7 +388,6 @@ mod tests { encoded.symbol_table(), encoded.codes(), PrimitiveArray::from_iter([4i32]).into_array(), - &mut ctx, )?; assert!(fsst_decode_bytes(invalid.as_view(), &mut ctx).is_err()); @@ -413,7 +411,6 @@ mod tests { encoded.symbol_table(), encoded.codes(), lengths, - &mut ctx, )?; assert!(fsst_decode_bytes(invalid.as_view(), &mut ctx).is_err()); } diff --git a/encodings/fsst/src/compress.rs b/encodings/fsst/src/compress.rs index 2cc28adcf11..957e319cd3c 100644 --- a/encodings/fsst/src/compress.rs +++ b/encodings/fsst/src/compress.rs @@ -228,7 +228,7 @@ where } } } - sink.finish(strings.dtype().clone(), ctx) + sink.finish(strings.dtype().clone()) } fn compress_varbin( @@ -252,7 +252,7 @@ where let off = offsets.as_slice::(); for_each_varbin_row(off, bytes, mask, |row| sink.emit(row)); }); - sink.finish(strings.dtype().clone(), ctx) + sink.finish(strings.dtype().clone()) } /// Call `f` once per row of a `VarBinArray` with the row bytes or `None`. @@ -333,7 +333,7 @@ impl<'c, O: OffsetBuilderPType + 'static> FsstSink<'c, O> { self.builder.append_value(&self.buffer); } - fn finish(mut self, dtype: DType, ctx: &mut ExecutionCtx) -> VortexResult { + fn finish(mut self, dtype: DType) -> VortexResult { let codes = self.builder.finish_into_varbin(); // Pad the symbol table here so that the array can hand it straight to a `Decompressor` // without copying. @@ -346,7 +346,6 @@ impl<'c, O: OffsetBuilderPType + 'static> FsstSink<'c, O> { )?), codes, self.uncompressed_lengths.into_array(), - ctx, ) } } diff --git a/encodings/fsst/src/compute/filter.rs b/encodings/fsst/src/compute/filter.rs index 45a95520aca..bb1f9ec6a65 100644 --- a/encodings/fsst/src/compute/filter.rs +++ b/encodings/fsst/src/compute/filter.rs @@ -37,7 +37,6 @@ impl FilterKernel for FSST { array.symbol_table(), filtered_codes, array.uncompressed_lengths().filter(mask.clone())?, - ctx, )? .into_array(), )) diff --git a/encodings/fsst/src/compute/mod.rs b/encodings/fsst/src/compute/mod.rs index f606072488c..501479e1944 100644 --- a/encodings/fsst/src/compute/mod.rs +++ b/encodings/fsst/src/compute/mod.rs @@ -77,7 +77,6 @@ impl TakeExecute for FSST { .fill_null(Scalar::zero_value( &array.uncompressed_lengths_dtype().clone(), ))?, - ctx, )? .into_array(), )) diff --git a/encodings/runend/src/array.rs b/encodings/runend/src/array.rs index 4ce0350a37a..0913ec01457 100644 --- a/encodings/runend/src/array.rs +++ b/encodings/runend/src/array.rs @@ -33,7 +33,6 @@ use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; -use vortex_array::legacy_session; use vortex_array::require_child; use vortex_array::serde::ArrayChildren; use vortex_array::validity::Validity; @@ -93,7 +92,6 @@ impl VTable for RunEnd { *ID } - #[allow(clippy::disallowed_methods)] fn validate( &self, data: &Self::TypedArrayData, @@ -104,9 +102,7 @@ impl VTable for RunEnd { let run_end_slots = RunEndSlotsView::from_slots(slots); let ends = run_end_slots.ends; let values = run_end_slots.values; - // TODO(ctx): trait fixes - VTable::validate has a fixed signature. - let mut ctx = legacy_session().create_execution_ctx(); - RunEndData::validate_parts(ends, values, data.offset, len, &mut ctx)?; + RunEndData::validate_shape(ends, values, data.offset, len)?; vortex_ensure_eq!(values.dtype(), dtype); Ok(()) } @@ -153,7 +149,7 @@ impl VTable for RunEnd { metadata: &[u8], _buffers: &[BufferHandle], children: &dyn ArrayChildren, - _session: &VortexSession, + session: &VortexSession, ) -> VortexResult> { let metadata = RunEndMetadata::decode(metadata)?; let ends_dtype = DType::Primitive(metadata.ends_ptype(), Nullability::NonNullable); @@ -162,6 +158,13 @@ impl VTable for RunEnd { let values = children.get(1, dtype, runs)?; let offset = usize::try_from(metadata.offset).vortex_expect("Offset must be a valid usize"); + RunEndData::validate_parts( + &ends, + &values, + offset, + len, + &mut session.create_execution_ctx(), + )?; let slots = RunEndSlots { ends, values }.into_slots(); let data = RunEndData::new(offset); Ok(ArrayParts::new( @@ -327,22 +330,7 @@ impl RunEndData { length: usize, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - // DType validation - vortex_ensure!( - ends.dtype().is_unsigned_int(), - "run ends must be unsigned integers, was {}", - ends.dtype(), - ); - vortex_ensure_eq!(ends.len(), values.len()); - - // Handle empty run-ends - if ends.is_empty() { - vortex_ensure_eq!(offset, 0, "non-zero offset provided for empty RunEndArray"); - return Ok(()); - } - - // Zero-length logical slices may retain run metadata from the source array. - if length == 0 { + if !Self::validate_shape(ends, values, offset, length)? { return Ok(()); } @@ -367,15 +355,52 @@ impl RunEndData { return Ok(()); } - // 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.execute_scalar(0, ctx)?)?; - if first_run_end < offset { - vortex_bail!("First run end {first_run_end} must be >= offset {offset}"); - } + let first_run_end = if offset == 0 { + 0 + } else { + usize::try_from(&ends.execute_scalar(0, ctx)?)? + }; + let last_run_end = usize::try_from(&ends.execute_scalar(ends.len() - 1, ctx)?)?; + Self::validate_run_end_bounds(first_run_end, last_run_end, offset, length) + } + + /// Validates the dtypes and lengths of the parts. Returns `false` when there are no runs to + /// check against `offset..offset + length`. + fn validate_shape( + ends: &ArrayRef, + values: &ArrayRef, + offset: usize, + length: usize, + ) -> VortexResult { + // DType validation + vortex_ensure!( + ends.dtype().is_unsigned_int(), + "run ends must be unsigned integers, was {}", + ends.dtype(), + ); + vortex_ensure_eq!(ends.len(), values.len()); + + // Handle empty run-ends + if ends.is_empty() { + vortex_ensure_eq!(offset, 0, "non-zero offset provided for empty RunEndArray"); + return Ok(false); + } + + // Zero-length logical slices may retain run metadata from the source array. + Ok(length != 0) + } + + /// Validates that the first and last run ends cover `offset..offset + length`. + fn validate_run_end_bounds( + first_run_end: usize, + last_run_end: usize, + offset: usize, + length: usize, + ) -> VortexResult<()> { + 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.execute_scalar(ends.len() - 1, 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/zstd/src/zstd_buffers.rs b/encodings/zstd/src/zstd_buffers.rs index 53423510bcb..548f5bc2bff 100644 --- a/encodings/zstd/src/zstd_buffers.rs +++ b/encodings/zstd/src/zstd_buffers.rs @@ -98,6 +98,7 @@ impl ZstdBuffers { compressed_buffers, uncompressed_sizes, buffer_alignments, + session: session.clone(), }; let slots: ArraySlots = serialization.children.into_iter().map(Some).collect(); let compressed = Array::try_from_parts(ArrayParts::new( @@ -160,6 +161,8 @@ pub struct ZstdBuffersData { compressed_buffers: Vec, uncompressed_sizes: Vec, buffer_alignments: Vec, + /// The session used to rebuild the wrapped array, which resolves its encoding. + session: VortexSession, } impl Display for ZstdBuffersData { @@ -513,6 +516,7 @@ impl VTable for ZstdBuffers { compressed_buffers, uncompressed_sizes: metadata.uncompressed_sizes.clone(), buffer_alignments: metadata.buffer_alignments.clone(), + session: session.clone(), }; data.validate()?; @@ -554,16 +558,13 @@ impl OperationsVTable for ZstdBuffers { } impl ValidityVTable for ZstdBuffers { - #[allow(clippy::disallowed_methods)] fn validity(array: ArrayView<'_, ZstdBuffers>) -> VortexResult { if !array.dtype().is_nullable() { return Ok(Validity::NonNullable); } - let inner_array = ZstdBuffers::decompress_and_build_inner( - &array.into_owned(), - vortex_array::legacy_session(), - )?; + let session = array.data().session.clone(); + let inner_array = ZstdBuffers::decompress_and_build_inner(&array.into_owned(), &session)?; inner_array.validity() } } diff --git a/lang/cpp/examples/reader.cpp b/lang/cpp/examples/reader.cpp index a447ec6917b..a5706631bf5 100644 --- a/lang/cpp/examples/reader.cpp +++ b/lang/cpp/examples/reader.cpp @@ -21,7 +21,7 @@ int main() { for (Partition &partition : scan.partitions()) { for (Array &array : partition.batches()) { - const Array age = array.field("age"); + const Array age = array.field(session, "age"); const PrimitiveView age_view = age.values(session); const std::span age_values = age_view.values(); for (uint8_t value : age_values) { diff --git a/lang/cpp/include/vortex/array.hpp b/lang/cpp/include/vortex/array.hpp index c60dde2a8bd..15ff95a4919 100644 --- a/lang/cpp/include/vortex/array.hpp +++ b/lang/cpp/include/vortex/array.hpp @@ -196,19 +196,19 @@ class Array { Validity validity() const; // Number of null/invalid elements in Array - size_t null_count() const; + size_t null_count(const Session &session) const; /** * Get a Struct field by index. Throws if Array is not a Struct or if index * is out of bounds. */ - Array field(size_t index) const; + Array field(const Session &session, size_t index) const; /** * Get a Struct field by name. Throws if Array is not a Struct or doesn't * have this named field. */ - Array field(std::string_view name) const; + Array field(const Session &session, std::string_view name) const; /* * Create a new Array slicing [begin; end) rows from original. diff --git a/lang/cpp/src/array.cpp b/lang/cpp/src/array.cpp index 7364c9abc34..d1c398dfee2 100644 --- a/lang/cpp/src/array.cpp +++ b/lang/cpp/src/array.cpp @@ -228,21 +228,21 @@ Validity Array::validity() const { return Access::adopt(static_cast(raw.type), raw.array); } -size_t Array::null_count() const { +size_t Array::null_count(const Session &session) const { vx_error *error = nullptr; - const size_t count = vx_array_invalid_count(handle_.get(), &error); + const size_t count = vx_array_invalid_count(Access::c_ptr(session), handle_.get(), &error); throw_on_error(error); return count; } -Array Array::field(size_t index) const { +Array Array::field(const Session &session, size_t index) const { vx_error *error = nullptr; - const vx_array *out = vx_array_get_field(handle_.get(), index, &error); + const vx_array *out = vx_array_get_field(Access::c_ptr(session), handle_.get(), index, &error); throw_on_error(error); return Access::adopt(out); } -Array Array::field(std::string_view name) const { +Array Array::field(const Session &session, std::string_view name) const { const DataType dt = dtype(); const std::unique_ptr fields( struct_fields_or_throw(detail::Access::c_ptr(dt)), @@ -251,7 +251,7 @@ Array Array::field(std::string_view name) const { for (uint64_t i = 0; i < fields_size; ++i) { const vx_view field = vx_struct_fields_field_name(fields.get(), i); if (std::string_view {field.ptr, field.len} == name) { - return this->field(i); + return this->field(session, i); } } throw VortexException("no field named \"" + std::string(name) + "\"", ErrorCode::InvalidArgument); diff --git a/lang/cpp/tests/array.cpp b/lang/cpp/tests/array.cpp index 8e12bd6ccb2..5a3b14a9062 100644 --- a/lang/cpp/tests/array.cpp +++ b/lang/cpp/tests/array.cpp @@ -28,7 +28,7 @@ TEST_CASE("Empty array", "[array]") { REQUIRE(empty.size() == 0); REQUIRE(empty.is_primitive(I32)); - REQUIRE(empty.null_count() == 0); + REQUIRE(empty.null_count(session) == 0); auto view = empty.values(session); REQUIRE(view.size() == 0); @@ -41,7 +41,7 @@ void test_primitive_array(Array array, const int32_t *begin) { REQUIRE(array.size() == 3); REQUIRE(array.is_primitive(I32)); REQUIRE_FALSE(array.nullable()); - REQUIRE(array.null_count() == 0); + REQUIRE(array.null_count(session) == 0); auto view = array.values(session); REQUIRE(view.size() == 3); @@ -80,7 +80,7 @@ TEST_CASE("Validity from a boolean mask", "[array]") { Array a = Array::primitive(std::span(data), Validity::from_array(mask)); REQUIRE(a.nullable()); - REQUIRE(a.null_count() == 1); + REQUIRE(a.null_count(session) == 1); auto view = a.values(session); REQUIRE_FALSE(view.is_null(0)); @@ -104,7 +104,7 @@ TEST_CASE("AllInvalid", "[array]") { Session session; std::vector data = {1, 2}; Array a = Array::primitive(std::span(data), AllInvalid); - REQUIRE(a.null_count() == 2); + REQUIRE(a.null_count(session) == 2); auto view = a.values(session); REQUIRE(view.is_null(0)); REQUIRE(view.is_null(1)); @@ -128,17 +128,17 @@ TEST_CASE("make_struct and fields", "[array]") { REQUIRE(s.has_dtype(DataTypeVariant::Struct)); REQUIRE(s.dtype().fields().size() == 2); - Array by_index = s.field(0); + Session session; + Array by_index = s.field(session, 0); REQUIRE(by_index.is_primitive(U8)); - Session session; - Array by_name = s.field("height"); + Array by_name = s.field(session, "height"); REQUIRE(by_name.is_primitive(U16)); auto view = by_name.values(session); REQUIRE(view.values()[2] == 170); - REQUIRE_THROWS_AS(s.field(2), VortexException); - REQUIRE_THROWS_AS(s.field("nope"), VortexException); + REQUIRE_THROWS_AS(s.field(session, 2), VortexException); + REQUIRE_THROWS_AS(s.field(session, "nope"), VortexException); std::vector fields_vec; fields_vec.emplace_back("age", Array::primitive(ages)); diff --git a/lang/cpp/tests/arrow.cpp b/lang/cpp/tests/arrow.cpp index bfd287fcecb..a73d82de3f7 100644 --- a/lang/cpp/tests/arrow.cpp +++ b/lang/cpp/tests/arrow.cpp @@ -77,7 +77,7 @@ TEST_CASE("Import Arrow array as Vortex array", "[arrow]") { REQUIRE(vx.size() == 3); REQUIRE(vx.has_dtype(DataTypeVariant::Struct)); - Array a = vx.field(0); + Array a = vx.field(session, 0); REQUIRE(a.is_primitive(I32)); auto view = a.values(session); REQUIRE(view.values()[0] == 10); diff --git a/lang/cpp/tests/scan.cpp b/lang/cpp/tests/scan.cpp index 82854e11f2a..8e700b4b9a3 100644 --- a/lang/cpp/tests/scan.cpp +++ b/lang/cpp/tests/scan.cpp @@ -66,16 +66,16 @@ void verify_age_field(const Session &session, void verify_sample_array(const Session &session, const Array &array) { REQUIRE(array.size() == SAMPLE_ROWS); REQUIRE(array.has_dtype(DataTypeVariant::Struct)); - verify_age_field(session, array.field(0)); + verify_age_field(session, array.field(session, 0)); - Array height = array.field(1); + Array height = array.field(session, 1); REQUIRE(height.is_primitive(U16)); auto view = height.values(session); for (size_t i = 0; i < SAMPLE_ROWS; ++i) { REQUIRE(view.values()[i] == (i + 1) % 200); } - REQUIRE_THROWS_AS(array.field(2), VortexException); + REQUIRE_THROWS_AS(array.field(session, 2), VortexException); } TEST_CASE("Basic scan", "[scan]") { diff --git a/vortex-array/src/arrays/constant/compute/take.rs b/vortex-array/src/arrays/constant/compute/take.rs index ecd9db452ce..5fd36fb9e4c 100644 --- a/vortex-array/src/arrays/constant/compute/take.rs +++ b/vortex-array/src/arrays/constant/compute/take.rs @@ -2,32 +2,23 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; -use vortex_mask::AllOr; use crate::ArrayRef; use crate::IntoArray; -use crate::VortexSessionExecute; use crate::array::ArrayView; use crate::arrays::Constant; use crate::arrays::ConstantArray; use crate::arrays::MaskedArray; use crate::arrays::dict::TakeReduce; use crate::arrays::dict::TakeReduceAdaptor; -use crate::legacy_session; use crate::optimizer::rules::ParentRuleSet; use crate::scalar::Scalar; use crate::validity::Validity; impl TakeReduce for Constant { - #[allow(clippy::disallowed_methods)] fn take(array: ArrayView<'_, Constant>, indices: &ArrayRef) -> VortexResult> { - let mut ctx = legacy_session().create_execution_ctx(); - let result = match indices - .validity()? - .execute_mask(indices.len(), &mut ctx)? - .bit_buffer() - { - AllOr::All => { + let result = match indices.validity()? { + Validity::NonNullable | Validity::AllValid => { let scalar = Scalar::try_new( array .scalar() @@ -37,7 +28,7 @@ impl TakeReduce for Constant { )?; ConstantArray::new(scalar, indices.len()).into_array() } - AllOr::None => ConstantArray::new( + Validity::AllInvalid => ConstantArray::new( Scalar::null( array .dtype() @@ -46,14 +37,16 @@ impl TakeReduce for Constant { indices.len(), ) .into_array(), - AllOr::Some(v) => { + validity @ Validity::Array(_) => { let arr = ConstantArray::new(array.scalar().clone(), indices.len()).into_array(); if array.scalar().is_null() { return Ok(Some(arr)); } - MaskedArray::try_new(arr, Validity::from(v.clone()))?.into_array() + // SAFETY: `arr` repeats a non-null scalar, and the validity comes from the indices, + // which have the same length as `arr`. + unsafe { MaskedArray::new_unchecked(arr, validity) }.into_array() } }; Ok(Some(result)) diff --git a/vortex-array/src/arrays/masked/array.rs b/vortex-array/src/arrays/masked/array.rs index a83d5ab5701..f1243c8f1a9 100644 --- a/vortex-array/src/arrays/masked/array.rs +++ b/vortex-array/src/arrays/masked/array.rs @@ -9,7 +9,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use crate::ArrayRef; -use crate::VortexSessionExecute; +use crate::ExecutionCtx; use crate::array::Array; use crate::array::ArrayParts; use crate::array::TypedArrayRef; @@ -17,7 +17,6 @@ use crate::array::child_to_validity; use crate::array::validity_to_child; use crate::array_slots; use crate::arrays::Masked; -use crate::legacy_session; use crate::validity::Validity; #[array_slots(Masked)] @@ -77,16 +76,15 @@ impl MaskedData { impl Array { /// Constructs a new `MaskedArray`. - #[allow(clippy::disallowed_methods)] - pub fn try_new(child: ArrayRef, validity: Validity) -> VortexResult { + pub fn try_new( + child: ArrayRef, + validity: Validity, + ctx: &mut ExecutionCtx, + ) -> VortexResult { let dtype = child.dtype().as_nullable(); let len = child.len(); let validity_slot = validity_to_child(&validity, len); - let data = MaskedData::try_new( - len, - child.all_valid(&mut legacy_session().create_execution_ctx())?, - validity, - )?; + let data = MaskedData::try_new(len, child.all_valid(ctx)?, validity)?; Ok(unsafe { Array::from_parts_unchecked(ArrayParts::new( Masked, @@ -97,4 +95,25 @@ impl Array { )) }) } + + /// Constructs a new `MaskedArray` without validating the child or validity. + /// + /// # Safety + /// + /// The caller must ensure that `child` contains no nulls, that `validity` is nullable, and + /// that an array-backed `validity` has the same length as `child`. + pub unsafe fn new_unchecked(child: ArrayRef, validity: Validity) -> Self { + let dtype = child.dtype().as_nullable(); + let len = child.len(); + let validity_slot = validity_to_child(&validity, len); + unsafe { + Array::from_parts_unchecked(ArrayParts::new( + Masked, + dtype, + len, + MaskedData, + smallvec![Some(child), validity_slot], + )) + } + } } diff --git a/vortex-array/src/arrays/masked/compute/filter.rs b/vortex-array/src/arrays/masked/compute/filter.rs index 287a996a12b..9d0fdfe8b23 100644 --- a/vortex-array/src/arrays/masked/compute/filter.rs +++ b/vortex-array/src/arrays/masked/compute/filter.rs @@ -21,9 +21,10 @@ impl FilterReduce for Masked { // The child is guaranteed to have no nulls, so filtering it is straightforward let filtered_child = array.child().filter(mask.clone())?; - // Construct new MaskedArray + // SAFETY: filtering a null-free child keeps it null-free, and the child and validity are + // filtered by the same mask, so their lengths match. Ok(Some( - MaskedArray::try_new(filtered_child, filtered_validity)?.into_array(), + unsafe { MaskedArray::new_unchecked(filtered_child, filtered_validity) }.into_array(), )) } } @@ -44,19 +45,22 @@ mod tests { #[case( MaskedArray::try_new( PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]).into_array(), - Validity::from_iter([true, true, false, true, false]) + Validity::from_iter([true, true, false, true, false]), + &mut array_session().create_execution_ctx(), ).unwrap() )] #[case( MaskedArray::try_new( PrimitiveArray::from_iter([10i32, 20, 30]).into_array(), - Validity::AllValid + Validity::AllValid, + &mut array_session().create_execution_ctx(), ).unwrap() )] #[case( MaskedArray::try_new( PrimitiveArray::from_iter(0..100).into_array(), - Validity::from_iter((0..100).map(|i| i % 3 != 0)) + Validity::from_iter((0..100).map(|i| i % 3 != 0)), + &mut array_session().create_execution_ctx(), ).unwrap() )] fn test_filter_masked_conformance(#[case] array: MaskedArray) { diff --git a/vortex-array/src/arrays/masked/compute/mask.rs b/vortex-array/src/arrays/masked/compute/mask.rs index ba901086459..ed9d17039a1 100644 --- a/vortex-array/src/arrays/masked/compute/mask.rs +++ b/vortex-array/src/arrays/masked/compute/mask.rs @@ -40,19 +40,22 @@ mod tests { #[case( MaskedArray::try_new( PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]).into_array(), - Validity::from_iter([true, true, false, true, false]) + Validity::from_iter([true, true, false, true, false]), + &mut array_session().create_execution_ctx(), ).unwrap() )] #[case( MaskedArray::try_new( PrimitiveArray::from_iter([10i32, 20, 30]).into_array(), - Validity::AllValid + Validity::AllValid, + &mut array_session().create_execution_ctx(), ).unwrap() )] #[case( MaskedArray::try_new( PrimitiveArray::from_iter(0..100).into_array(), - Validity::from_iter((0..100).map(|i| i % 3 != 0)) + Validity::from_iter((0..100).map(|i| i % 3 != 0)), + &mut array_session().create_execution_ctx(), ).unwrap() )] fn test_mask_masked_conformance(#[case] array: MaskedArray) { diff --git a/vortex-array/src/arrays/masked/compute/slice.rs b/vortex-array/src/arrays/masked/compute/slice.rs index 5a51ca9f681..d9d3f197018 100644 --- a/vortex-array/src/arrays/masked/compute/slice.rs +++ b/vortex-array/src/arrays/masked/compute/slice.rs @@ -18,6 +18,10 @@ impl SliceReduce for Masked { let child = array.child().slice(range.clone())?; let validity = array.validity()?.slice(range)?; - Ok(Some(MaskedArray::try_new(child, validity)?.into_array())) + // SAFETY: slicing a null-free child keeps it null-free, and the child and validity are + // sliced by the same range, so their lengths match. + Ok(Some( + unsafe { MaskedArray::new_unchecked(child, validity) }.into_array(), + )) } } diff --git a/vortex-array/src/arrays/masked/compute/take.rs b/vortex-array/src/arrays/masked/compute/take.rs index 49ad0fc9fe6..cf621aaa377 100644 --- a/vortex-array/src/arrays/masked/compute/take.rs +++ b/vortex-array/src/arrays/masked/compute/take.rs @@ -32,9 +32,10 @@ impl TakeReduce for Masked { // Compute the new validity by taking from array's validity and merging with indices validity let taken_validity = array.validity()?.take(indices)?; - // Construct new MaskedArray + // SAFETY: the child is null-free and is taken with non-null indices, so the taken child is + // null-free. The child and validity are taken by the same indices, so their lengths match. Ok(Some( - MaskedArray::try_new(taken_child, taken_validity)?.into_array(), + unsafe { MaskedArray::new_unchecked(taken_child, taken_validity) }.into_array(), )) } } @@ -55,19 +56,22 @@ mod tests { #[case( MaskedArray::try_new( PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]).into_array(), - Validity::from_iter([true, true, false, true, false]) + Validity::from_iter([true, true, false, true, false]), + &mut array_session().create_execution_ctx(), ).unwrap() )] #[case( MaskedArray::try_new( PrimitiveArray::from_iter([10i32, 20, 30]).into_array(), - Validity::AllValid + Validity::AllValid, + &mut array_session().create_execution_ctx(), ).unwrap() )] #[case( MaskedArray::try_new( PrimitiveArray::from_iter(0..100).into_array(), - Validity::from_iter((0..100).map(|i| i % 3 != 0)) + Validity::from_iter((0..100).map(|i| i % 3 != 0)), + &mut array_session().create_execution_ctx(), ).unwrap() )] fn test_take_masked_conformance(#[case] array: MaskedArray) { diff --git a/vortex-array/src/arrays/masked/execute.rs b/vortex-array/src/arrays/masked/execute.rs index eebcf1a2335..b81c16a5971 100644 --- a/vortex-array/src/arrays/masked/execute.rs +++ b/vortex-array/src/arrays/masked/execute.rs @@ -205,7 +205,7 @@ fn mask_validity_variant( let masked_core_storage = match core_validity { Validity::NonNullable | Validity::AllValid => { // Core storage has no nulls, so wrap it in MaskedArray to apply the mask. - MaskedArray::try_new(core_storage, validity)?.into_array() + MaskedArray::try_new(core_storage, validity, ctx)?.into_array() } Validity::AllInvalid => { // Already all-null, ANDing with any mask is still all-null. diff --git a/vortex-array/src/arrays/masked/tests.rs b/vortex-array/src/arrays/masked/tests.rs index 93a61e9f06f..b12ffb5f265 100644 --- a/vortex-array/src/arrays/masked/tests.rs +++ b/vortex-array/src/arrays/masked/tests.rs @@ -27,6 +27,7 @@ fn test_repeated_probe_preserves_masked_nulls() -> VortexResult<()> { let array = MaskedArray::try_new( PrimitiveArray::from_iter([1i32, 2, 3]).into_array(), Validity::from_iter([true, false, true]), + &mut ctx, )? .into_array(); let expected = [Some(1i32), None, Some(3)]; @@ -48,8 +49,9 @@ fn test_repeated_probe_preserves_masked_nulls() -> VortexResult<()> { #[case(Validity::AllValid, Nullability::Nullable)] #[case(Validity::from_iter([true, false, true]), Nullability::Nullable)] fn test_dtype_nullability(#[case] validity: Validity, #[case] expected: Nullability) { + let mut ctx = array_session().create_execution_ctx(); let child = PrimitiveArray::from_iter([1i32, 2, 3]).into_array(); - let array = MaskedArray::try_new(child, validity).unwrap(); + let array = MaskedArray::try_new(child, validity, &mut ctx).unwrap(); assert_eq!( array.dtype(), @@ -80,7 +82,7 @@ fn test_empty_child_with_array_validity() -> VortexResult<()> { assert!(child.all_valid(&mut ctx)?); assert!(child.all_invalid(&mut ctx)?); - let array = MaskedArray::try_new(child, validity)?; + let array = MaskedArray::try_new(child, validity, &mut ctx)?; assert!(array.is_empty()); Ok(()) @@ -88,9 +90,10 @@ fn test_empty_child_with_array_validity() -> VortexResult<()> { #[test] fn test_canonical_dtype_matches_array_dtype() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // The canonical form should have the same nullability as the array's dtype. let child = PrimitiveArray::from_iter([1i32, 2, 3]).into_array(); - let array = MaskedArray::try_new(child, Validity::AllValid)?; + let array = MaskedArray::try_new(child, Validity::AllValid, &mut ctx)?; let canonical = array .clone() @@ -102,10 +105,15 @@ fn test_canonical_dtype_matches_array_dtype() -> VortexResult<()> { #[test] fn test_masked_child_with_validity() { + let mut ctx = array_session().create_execution_ctx(); // When validity has nulls, masked_child should apply inverted mask. let child = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]).into_array(); - let array = - MaskedArray::try_new(child, Validity::from_iter([true, false, true, false, true])).unwrap(); + let array = MaskedArray::try_new( + child, + Validity::from_iter([true, false, true, false, true]), + &mut ctx, + ) + .unwrap(); // Positions where validity is false should be null in masked_child. let mut ctx = array_session().create_execution_ctx(); @@ -144,7 +152,7 @@ fn test_masked_child_all_valid() { let mut ctx = array_session().create_execution_ctx(); // When validity is AllValid, masked_child should invert to AllInvalid. let child = PrimitiveArray::from_iter([10i32, 20, 30]).into_array(); - let array = MaskedArray::try_new(child, Validity::AllValid).unwrap(); + let array = MaskedArray::try_new(child, Validity::AllValid, &mut ctx).unwrap(); assert_eq!(array.len(), 3); assert_eq!( @@ -166,6 +174,7 @@ fn test_masked_child_all_valid() { #[case(Validity::from_iter([false, false, false]))] #[case(Validity::from_iter([true, false, true, false]))] fn test_masked_child_preserves_length(#[case] validity: Validity) { + let mut ctx = array_session().create_execution_ctx(); let len = match &validity { Validity::Array(arr) => arr.len(), _ => 3, @@ -173,7 +182,7 @@ fn test_masked_child_preserves_length(#[case] validity: Validity) { #[expect(clippy::cast_possible_truncation)] let child = PrimitiveArray::from_iter(0..len as i32).into_array(); - let array = MaskedArray::try_new(child, validity.clone()).unwrap(); + let array = MaskedArray::try_new(child, validity.clone(), &mut ctx).unwrap(); assert_eq!(array.len(), len); @@ -189,6 +198,7 @@ fn test_masked_child_preserves_length(#[case] validity: Validity) { #[test] fn masked_listview_execute_preserves_zctl_true() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // Masking only intersects validity, so the zero-copy-to-list survives // execution. let elements = PrimitiveArray::from_iter([1i32, 2, 3]).into_array(); @@ -199,7 +209,11 @@ fn masked_listview_execute_preserves_zctl_true() -> VortexResult<()> { .with_zero_copy_to_list(true) }; - let masked = MaskedArray::try_new(list_view.into_array(), Validity::from_iter([true, false]))?; + let masked = MaskedArray::try_new( + list_view.into_array(), + Validity::from_iter([true, false]), + &mut ctx, + )?; let canonical = masked .into_array() .execute::(&mut array_session().create_execution_ctx())?; diff --git a/vortex-array/src/arrays/masked/vtable/canonical.rs b/vortex-array/src/arrays/masked/vtable/canonical.rs index 32a5e334afa..41e755c406c 100644 --- a/vortex-array/src/arrays/masked/vtable/canonical.rs +++ b/vortex-array/src/arrays/masked/vtable/canonical.rs @@ -16,17 +16,21 @@ mod tests { use crate::validity::Validity; fn masked_all_valid() -> MaskedArray { + let mut ctx = array_session().create_execution_ctx(); MaskedArray::try_new( PrimitiveArray::from_iter([1i32, 2, 3]).into_array(), Validity::AllValid, + &mut ctx, ) .expect("valid masked array") } fn masked_with_nulls() -> MaskedArray { + let mut ctx = array_session().create_execution_ctx(); MaskedArray::try_new( PrimitiveArray::from_iter([1i32, 2, 3]).into_array(), Validity::from_iter([true, false, true]), + &mut ctx, ) .expect("valid masked array") } @@ -49,9 +53,11 @@ mod tests { #[test] fn test_canonical_with_nulls() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let array = MaskedArray::try_new( PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]).into_array(), Validity::from_iter([true, false, true, false, true]), + &mut ctx, )?; let canonical = array @@ -72,9 +78,11 @@ mod tests { #[test] fn test_canonical_all_valid() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let array = MaskedArray::try_new( PrimitiveArray::from_iter([10i32, 20, 30]).into_array(), Validity::AllValid, + &mut ctx, )?; let canonical = array diff --git a/vortex-array/src/arrays/masked/vtable/mod.rs b/vortex-array/src/arrays/masked/vtable/mod.rs index 1bfbc74e5f1..5896f323066 100644 --- a/vortex-array/src/arrays/masked/vtable/mod.rs +++ b/vortex-array/src/arrays/masked/vtable/mod.rs @@ -42,7 +42,6 @@ use crate::buffer::BufferHandle; use crate::dtype::DType; use crate::executor::ExecutionCtx; use crate::executor::ExecutionResult; -use crate::legacy_session; use crate::require_child; use crate::scalar::Scalar; use crate::serde::ArrayChildren; @@ -74,7 +73,6 @@ impl VTable for Masked { *ID } - #[expect(clippy::disallowed_methods)] fn validate( &self, _data: &MaskedData, @@ -95,10 +93,6 @@ impl VTable for Masked { *dtype, "MaskedArray dtype does not match child and validity", ); - vortex_ensure!( - child.all_valid(&mut legacy_session().create_execution_ctx())?, - "MaskedArray children must not have nulls", - ); Ok(()) } @@ -129,7 +123,6 @@ impl VTable for Masked { Ok(Some(vec![])) } - #[allow(clippy::disallowed_methods)] fn deserialize( &self, dtype: &DType, @@ -138,7 +131,7 @@ impl VTable for Masked { buffers: &[BufferHandle], children: &dyn ArrayChildren, - _session: &VortexSession, + session: &VortexSession, ) -> VortexResult> { if !metadata.is_empty() { vortex_bail!( @@ -168,7 +161,7 @@ impl VTable for Masked { let validity_slot = validity_to_child(&validity, len); let data = MaskedData::try_new( len, - child.all_valid(&mut legacy_session().create_execution_ctx())?, + child.all_valid(&mut session.create_execution_ctx())?, validity, )?; Ok(ArrayParts::new( @@ -241,19 +234,22 @@ mod tests { #[case( MaskedArray::try_new( PrimitiveArray::from_iter([1i32, 2, 3]).into_array(), - Validity::AllValid + Validity::AllValid, + &mut array_session().create_execution_ctx(), ).unwrap() )] #[case( MaskedArray::try_new( PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]).into_array(), - Validity::from_iter([true, true, false, true, false]) + Validity::from_iter([true, true, false, true, false]), + &mut array_session().create_execution_ctx(), ).unwrap() )] #[case( MaskedArray::try_new( PrimitiveArray::from_iter(0..100).into_array(), - Validity::from_iter((0..100).map(|i| i % 3 != 0)) + Validity::from_iter((0..100).map(|i| i % 3 != 0)), + &mut array_session().create_execution_ctx(), ).unwrap() )] fn test_serde_roundtrip(#[case] array: MaskedArray) { @@ -298,13 +294,14 @@ mod tests { /// MaskedArray always has Nullable dtype. #[test] fn test_execute_with_all_valid_preserves_nullable_dtype() -> Result<(), VortexError> { + let mut ctx = array_session().create_execution_ctx(); // Create a MaskedArray with AllValid validity. // Child has NonNullable dtype, but MaskedArray's dtype is Nullable. let child = PrimitiveArray::from_iter([1i32, 2, 3]).into_array(); assert_eq!(child.dtype().nullability(), Nullability::NonNullable); - let array = MaskedArray::try_new(child, Validity::AllValid)?; + let array = MaskedArray::try_new(child, Validity::AllValid, &mut ctx)?; assert_eq!(array.dtype().nullability(), Nullability::Nullable); // Execute the array. This should produce a Canonical with Nullable dtype. diff --git a/vortex-array/src/arrays/mod.rs b/vortex-array/src/arrays/mod.rs index 8a8b6a234f7..ed6e272b2b1 100644 --- a/vortex-array/src/arrays/mod.rs +++ b/vortex-array/src/arrays/mod.rs @@ -142,6 +142,7 @@ pub(crate) fn initialize(session: &VortexSession) { list::initialize(session); listview::initialize(session); map::initialize(session); + null::initialize(session); patched::initialize(session); primitive::initialize(session); struct_::initialize(session); diff --git a/vortex-array/src/arrays/null/compute/rules.rs b/vortex-array/src/arrays/null/compute/rules.rs index f63c904667a..b2ee44a19dd 100644 --- a/vortex-array/src/arrays/null/compute/rules.rs +++ b/vortex-array/src/arrays/null/compute/rules.rs @@ -2,7 +2,6 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use crate::arrays::Null; -use crate::arrays::dict::TakeReduceAdaptor; use crate::arrays::filter::FilterReduceAdaptor; use crate::arrays::slice::SliceReduceAdaptor; use crate::optimizer::rules::ParentRuleSet; @@ -14,5 +13,4 @@ pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&CastReduceAdaptor(Null)), ParentRuleSet::lift(&MaskReduceAdaptor(Null)), ParentRuleSet::lift(&SliceReduceAdaptor(Null)), - ParentRuleSet::lift(&TakeReduceAdaptor(Null)), ]); diff --git a/vortex-array/src/arrays/null/compute/take.rs b/vortex-array/src/arrays/null/compute/take.rs index dc328614743..6a60c94a54c 100644 --- a/vortex-array/src/arrays/null/compute/take.rs +++ b/vortex-array/src/arrays/null/compute/take.rs @@ -5,23 +5,23 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use crate::ArrayRef; +use crate::ExecutionCtx; use crate::IntoArray; -use crate::VortexSessionExecute; use crate::array::ArrayView; use crate::arrays::Null; use crate::arrays::NullArray; use crate::arrays::PrimitiveArray; -use crate::arrays::dict::TakeReduce; -use crate::arrays::dict::TakeReduceAdaptor; +use crate::arrays::dict::TakeExecute; use crate::match_each_integer_ptype; -use crate::optimizer::rules::ParentRuleSet; -impl TakeReduce for Null { +impl TakeExecute for Null { #[expect(clippy::cast_possible_truncation)] - fn take(array: ArrayView<'_, Null>, indices: &ArrayRef) -> VortexResult> { - #[allow(clippy::disallowed_methods)] - let mut ctx = crate::legacy_session().create_execution_ctx(); - let indices = indices.clone().execute::(&mut ctx)?; + fn take( + array: ArrayView<'_, Null>, + indices: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let indices = indices.clone().execute::(ctx)?; // Enforce all indices are valid match_each_integer_ptype!(indices.ptype(), |T| { @@ -35,8 +35,3 @@ impl TakeReduce for Null { Ok(Some(NullArray::new(indices.len()).into_array())) } } - -impl Null { - pub const TAKE_RULES: ParentRuleSet = - ParentRuleSet::new(&[ParentRuleSet::lift(&TakeReduceAdaptor::(Self))]); -} diff --git a/vortex-array/src/arrays/null/mod.rs b/vortex-array/src/arrays/null/mod.rs index 62deab08eb8..a4db271bf8a 100644 --- a/vortex-array/src/arrays/null/mod.rs +++ b/vortex-array/src/arrays/null/mod.rs @@ -11,6 +11,7 @@ use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::ArraySlots; +use crate::ArrayVTable; use crate::ExecutionCtx; use crate::ExecutionResult; use crate::array::Array; @@ -22,11 +23,14 @@ use crate::array::OperationsVTable; use crate::array::VTable; use crate::array::ValidityVTable; use crate::array::with_empty_buffers; +use crate::arrays::Dict; +use crate::arrays::dict::TakeExecuteAdaptor; use crate::arrays::null::compute::rules::PARENT_RULES; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::builders::NullBuilder; use crate::dtype::DType; +use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar::Scalar; use crate::serde::ArrayChildren; use crate::validity::Validity; @@ -36,6 +40,12 @@ pub(crate) mod compute; /// A [`Null`]-encoded Vortex array. pub type NullArray = Array; +pub(crate) fn initialize(session: &VortexSession) { + session + .kernels() + .register_execute_parent_kernel(Dict.id(), Null, TakeExecuteAdaptor(Null)); +} + impl VTable for Null { type TypedArrayData = EmptyArrayData; diff --git a/vortex-array/src/arrays/varbinview/array.rs b/vortex-array/src/arrays/varbinview/array.rs index 8327f91db2b..180704da9f2 100644 --- a/vortex-array/src/arrays/varbinview/array.rs +++ b/vortex-array/src/arrays/varbinview/array.rs @@ -25,14 +25,17 @@ use crate::ArrayRef; use crate::ArraySlots; use crate::ExecutionCtx; #[cfg(debug_assertions)] -use crate::VortexSessionExecute; use crate::array::Array; use crate::array::ArrayParts; use crate::array::TypedArrayRef; use crate::array::child_to_validity; use crate::array::validity_to_child; use crate::array_slots; +#[cfg(debug_assertions)] +use crate::arrays::Bool; use crate::arrays::VarBinView; +#[cfg(debug_assertions)] +use crate::arrays::bool::BoolArrayExt; use crate::arrays::varbinview::BinaryView; use crate::arrays::varbinview::ResolvedViews; use crate::buffer::BufferHandle; @@ -41,7 +44,6 @@ use crate::builders::VarBinViewBuilder; use crate::dtype::DType; use crate::dtype::Nullability; #[cfg(debug_assertions)] -use crate::legacy_session; use crate::validity::Validity; #[array_slots(VarBinView)] @@ -313,15 +315,7 @@ impl VarBinViewData { validity: Validity, ) -> Self { #[cfg(debug_assertions)] - #[expect(clippy::disallowed_methods)] - Self::validate( - &views, - &buffers, - &dtype, - &validity, - &mut legacy_session().create_execution_ctx(), - ) - .vortex_expect("[Debug Assertion]: Invalid `VarBinViewArray` parameters"); + Self::debug_validate(&views, &buffers, &dtype, &validity); let handles: Vec = buffers .iter() @@ -366,6 +360,30 @@ impl VarBinViewData { Self::validate_views(views, &mask, buffers, check_utf8) } + /// Debug-only validation for unchecked construction. + /// + /// This never executes the validity, so the views are only checked when the validity is + /// constant or a canonical host-resident bitmap. + #[cfg(debug_assertions)] + fn debug_validate( + views: &Buffer, + buffers: &Arc<[ByteBuffer]>, + dtype: &DType, + validity: &Validity, + ) { + let mask = match validity { + Validity::NonNullable | Validity::AllValid => Mask::AllTrue(views.len()), + Validity::AllInvalid => Mask::AllFalse(views.len()), + Validity::Array(array) => match array.as_opt::() { + Some(bools) if array.is_host() => Mask::from_buffer(bools.to_bit_buffer()), + _ => return, + }, + }; + Self::check_nullability(dtype, validity) + .and_then(|check_utf8| Self::validate_views(views, &mask, buffers, check_utf8)) + .vortex_expect("[Debug Assertion]: Invalid `VarBinViewArray` parameters"); + } + /// Validates components like validate() and replaces views at null slots to empty views pub(crate) fn validate_and_fix( views: Buffer, diff --git a/vortex-array/src/display/mod.rs b/vortex-array/src/display/mod.rs index 780140af2ac..2fe4e6a4f36 100644 --- a/vortex-array/src/display/mod.rs +++ b/vortex-array/src/display/mod.rs @@ -21,7 +21,7 @@ pub use tree_display::TreeDisplay; use crate::ArrayRef; use crate::VortexSessionExecute; -use crate::legacy_session; +use crate::array_session; /// Describe how to convert an array to a string. /// @@ -522,7 +522,6 @@ impl ArrayRef { DisplayArrayAs(self, DisplayOptions::TableDisplay) } - #[allow(clippy::disallowed_methods)] fn fmt_as(&self, f: &mut std::fmt::Formatter, options: &DisplayOptions) -> std::fmt::Result { match options { DisplayOptions::MetadataOnly => EncodingSummaryExtractor::write(self, f), @@ -537,17 +536,25 @@ impl ArrayRef { let limit = self.len().min(f.precision().unwrap_or(DISPLAY_LIMIT)); let is_truncated = self.len() > limit; - let fmt_scalar = |i| { - self.execute_scalar(i, &mut legacy_session().create_execution_ctx()) + // Formatting cannot take an execution context, so the values are read with a + // default array session. + let mut ctx = array_session().create_execution_ctx(); + let mut fmt_scalar = |i| { + self.execute_scalar(i, &mut ctx) .map_or_else(|e| format!(""), |s| s.to_string()) }; + let head = (0..limit.saturating_sub(3)) + .map(&mut fmt_scalar) + .collect::>(); + let tail = (self.len().saturating_sub(3)..self.len()) + .map(&mut fmt_scalar) + .collect::>(); write!( f, "{opening_brace}{}{closing_brace}", - (0..limit.saturating_sub(3)) - .map(fmt_scalar) + head.into_iter() .chain(repeat_n("...".to_string(), is_truncated as usize)) - .chain((self.len().saturating_sub(3)..self.len()).map(fmt_scalar)) + .chain(tail) .format(sep) ) } @@ -585,8 +592,9 @@ impl ArrayRef { use crate::dtype::DType; let mut builder = tabled::builder::Builder::default(); - // Reuse a single execution context across all per-row accesses below. - let mut ctx = legacy_session().create_execution_ctx(); + // Formatting cannot take an execution context, so the rows are read with a default + // array session. Reuse a single execution context across all per-row accesses. + let mut ctx = array_session().create_execution_ctx(); // Special logic for struct arrays. let DType::Struct(sf, _) = self.dtype() else { diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs index f2f4293e539..acf36415af6 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs @@ -144,9 +144,10 @@ fn test_batch_constant_uses_constant_decoder() -> VortexResult<()> { #[test] fn test_batch_const_unwraps_filtered_masked_constant() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let child = ConstantArray::new(7_i64, 3).into_array(); - let masked = - MaskedArray::try_new(child, Validity::from_iter([true, false, true]))?.into_array(); + let masked = MaskedArray::try_new(child, Validity::from_iter([true, false, true]), &mut ctx)? + .into_array(); let filtered = masked.filter(Mask::from_iter([true, true, false]))?; let Some(const_array) = batch_const(&filtered) else { diff --git a/vortex-btrblocks/src/schemes/string/fsst.rs b/vortex-btrblocks/src/schemes/string/fsst.rs index 7cb87154d3d..27f4fad3197 100644 --- a/vortex-btrblocks/src/schemes/string/fsst.rs +++ b/vortex-btrblocks/src/schemes/string/fsst.rs @@ -122,7 +122,6 @@ impl Scheme for FSSTScheme { )?), compressed_codes, compressed_original_lengths, - exec_ctx, )?; Ok(fsst.into_array()) diff --git a/vortex-compressor/src/compressor/constant/mod.rs b/vortex-compressor/src/compressor/constant/mod.rs index d452a56e1f2..ca28f72a094 100644 --- a/vortex-compressor/src/compressor/constant/mod.rs +++ b/vortex-compressor/src/compressor/constant/mod.rs @@ -188,7 +188,7 @@ fn compress_constant(source: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult< if mask.all_true() { Ok(const_arr) } else { - Ok(MaskedArray::try_new(const_arr, validity)?.into_array()) + Ok(MaskedArray::try_new(const_arr, validity, ctx)?.into_array()) } } diff --git a/vortex-cuda/benches/fsst_cuda.rs b/vortex-cuda/benches/fsst_cuda.rs index 784856c0585..6924b56f79c 100644 --- a/vortex-cuda/benches/fsst_cuda.rs +++ b/vortex-cuda/benches/fsst_cuda.rs @@ -89,7 +89,6 @@ fn make_fixture(n: usize) -> FSSTBenchFixture { ), fsst.codes(), fsst.uncompressed_lengths().clone(), - setup_ctx.execution_ctx(), ) .vortex_expect("rebuild FSST fixture with Binary dtype") .into_array(); diff --git a/vortex-cuda/src/arrow/canonical.rs b/vortex-cuda/src/arrow/canonical.rs index fb06f4c48ab..78d4cc7204c 100644 --- a/vortex-cuda/src/arrow/canonical.rs +++ b/vortex-cuda/src/arrow/canonical.rs @@ -2676,7 +2676,6 @@ mod tests { fsst.symbol_table(), fsst.codes(), PrimitiveArray::from_iter([i32::MAX, i32::MAX]).into_array(), - ctx.execution_ctx(), )? .into_array(); let schema = arrow_schema_for_array(&oversized, &mut ctx)?; diff --git a/vortex-cuda/src/kernel/arrays/masked.rs b/vortex-cuda/src/kernel/arrays/masked.rs index b739ed2bf60..1f8ddaf1390 100644 --- a/vortex-cuda/src/kernel/arrays/masked.rs +++ b/vortex-cuda/src/kernel/arrays/masked.rs @@ -82,7 +82,7 @@ mod tests { let validity = Validity::Array( BoolArray::from_iter([true, false, true, true].into_iter()).into_array(), ); - let masked = MaskedArray::try_new(child, validity)?; + let masked = MaskedArray::try_new(child, validity, cuda_ctx.execution_ctx())?; let gpu_result = MaskedExecutor .execute(masked.clone().into_array(), &mut cuda_ctx) diff --git a/vortex-ffi/cinclude/vortex.h b/vortex-ffi/cinclude/vortex.h index 384dc29305c..c6b250fc9dc 100644 --- a/vortex-ffi/cinclude/vortex.h +++ b/vortex-ffi/cinclude/vortex.h @@ -677,7 +677,8 @@ size_t vx_array_len(const vx_array *array); */ const vx_dtype *vx_array_dtype(const vx_array *array); -const vx_array *vx_array_get_field(const vx_array *array, size_t index, vx_error **error_out); +const vx_array * +vx_array_get_field(const vx_session *session, const vx_array *array, size_t index, vx_error **error_out); const vx_array *vx_array_slice(const vx_array *array, size_t start, size_t stop, vx_error **error_out); @@ -694,7 +695,7 @@ bool vx_array_element_is_invalid(const vx_session *session, /** * Check how many items in the array are invalid (null). */ -size_t vx_array_invalid_count(const vx_array *array, vx_error **error_out); +size_t vx_array_invalid_count(const vx_session *session, const vx_array *array, vx_error **error_out); /** * Increase reference count on vx_array diff --git a/vortex-ffi/src/array.rs b/vortex-ffi/src/array.rs index df938879ba1..b033d3e1137 100644 --- a/vortex-ffi/src/array.rs +++ b/vortex-ffi/src/array.rs @@ -24,7 +24,6 @@ use vortex::array::arrays::StructArray; use vortex::array::arrays::VarBinView; use vortex::array::arrays::bool::BoolArrayExt; use vortex::array::arrays::struct_::StructArrayExt; -use vortex::array::legacy_session; use vortex::array::validity::Validity; use vortex::buffer::BitBuffer; use vortex::buffer::Buffer; @@ -267,16 +266,17 @@ pub unsafe extern "C-unwind" fn vx_array_dtype(array: *const vx_array) -> *const // Returns NULL and sets error_out if index is out of bounds or array doesn't // have dtype DTYPE_STRUCT. #[unsafe(no_mangle)] -#[allow(clippy::disallowed_methods)] pub unsafe extern "C-unwind" fn vx_array_get_field( + session: *const vx_session, array: *const vx_array, index: usize, error_out: *mut *mut vx_error, ) -> *const vx_array { try_or_default(error_out, || { + let session = unsafe { vx_session_ref(session) }?; let array = vx_array::as_ref(array); - let mut ctx = legacy_session().create_execution_ctx(); + let mut ctx = session.create_execution_ctx(); let struct_array = array.clone().execute::(&mut ctx)?; let field_array = struct_array .unmasked_field_opt(index) @@ -305,7 +305,6 @@ pub unsafe extern "C-unwind" fn vx_array_slice( /// validity array. Sets error if index is out of bounds or underlying validity /// array is corrupted. #[unsafe(no_mangle)] -#[allow(clippy::disallowed_methods)] pub unsafe extern "C-unwind" fn vx_array_element_is_invalid( session: *const vx_session, array: *const vx_array, @@ -320,15 +319,16 @@ pub unsafe extern "C-unwind" fn vx_array_element_is_invalid( /// Check how many items in the array are invalid (null). #[unsafe(no_mangle)] -#[allow(clippy::disallowed_methods)] pub unsafe extern "C-unwind" fn vx_array_invalid_count( + session: *const vx_session, array: *const vx_array, error_out: *mut *mut vx_error, ) -> usize { try_or_default(error_out, || { + let session = unsafe { vx_session_ref(session) }?; vortex_ensure!(!array.is_null()); let array = vx_array::as_ref(array); - array.invalid_count(&mut legacy_session().create_execution_ctx()) + array.invalid_count(&mut session.create_execution_ctx()) }) } @@ -878,7 +878,7 @@ mod tests { )); assert_no_error(error); - let null_count = vx_array_invalid_count(ffi_array, &raw mut error); + let null_count = vx_array_invalid_count(session, ffi_array, &raw mut error); assert_no_error(error); assert_eq!(null_count, 1); @@ -905,11 +905,11 @@ mod tests { let ffi_array = vx_array::new(struct_array.into_array()); let mut error = ptr::null_mut(); - let field0 = vx_array_get_field(ffi_array, 0, &raw mut error); + let field0 = vx_array_get_field(session, ffi_array, 0, &raw mut error); assert_no_error(error); assert_eq!(vx_array_len(field0), 3); - let field1 = vx_array_get_field(ffi_array, 1, &raw mut error); + let field1 = vx_array_get_field(session, ffi_array, 1, &raw mut error); assert_no_error(error); assert_eq!(vx_array_len(field1), 3); assert_eq!(get_u8(session, field1, 0), 30); @@ -917,7 +917,7 @@ mod tests { assert_eq!(get_u8(session, field1, 2), 35); // Test out of bounds - let field_oob = vx_array_get_field(ffi_array, 2, &raw mut error); + let field_oob = vx_array_get_field(session, ffi_array, 2, &raw mut error); assert!(!error.is_null()); assert!(field_oob.is_null()); vx_error_free(error); @@ -1098,14 +1098,14 @@ mod tests { assert_eq!(vx_array_len(vx), 3); assert!(!vx_array_is_nullable(vx)); - let a = vx_array_get_field(vx, 0, &raw mut error); + let a = vx_array_get_field(session, vx, 0, &raw mut error); assert_no_error(error); assert!(vx_array_is_primitive(a, vx_ptype::PTYPE_I32)); assert_eq!(get_i32(session, a, 0), 1); assert_eq!(get_i32(session, a, 2), 3); vx_array_free(a); - let b = vx_array_get_field(vx, 1, &raw mut error); + let b = vx_array_get_field(session, vx, 1, &raw mut error); assert_no_error(error); assert!(vx_array_has_dtype(b, vx_dtype_variant::DTYPE_UTF8)); assert!(vx_array_element_is_invalid(session, b, 1, &raw mut error)); diff --git a/vortex-ffi/test/array.cpp b/vortex-ffi/test/array.cpp index dc409e18c24..0b80fa84621 100644 --- a/vortex-ffi/test/array.cpp +++ b/vortex-ffi/test/array.cpp @@ -156,7 +156,7 @@ TEST_CASE("Array with validity", "[array]") { REQUIRE(vx_array_element_is_invalid(session, array, i, &error) == !valid[i]); require_no_error(error); } - REQUIRE(vx_array_invalid_count(array, &error) == 2); + REQUIRE(vx_array_invalid_count(session, array, &error) == 2); require_no_error(error); validity = {}; diff --git a/vortex-ffi/test/scan.cpp b/vortex-ffi/test/scan.cpp index f1334105ca6..5eb402fb480 100644 --- a/vortex-ffi/test/scan.cpp +++ b/vortex-ffi/test/scan.cpp @@ -361,17 +361,17 @@ void verify_sample_array(vx_session *session, const vx_array *array) { require_no_error(error); REQUIRE(validity.type == VX_VALIDITY_NON_NULLABLE); - const vx_array *age_field = vx_array_get_field(array, 0, &error); + const vx_array *age_field = vx_array_get_field(session, array, 0, &error); require_no_error(error); verify_age_field(session, age_field); vx_array_free(age_field); - const vx_array *height_field = vx_array_get_field(array, 1, &error); + const vx_array *height_field = vx_array_get_field(session, array, 1, &error); require_no_error(error); verify_height_field(session, height_field); vx_array_free(height_field); - REQUIRE(vx_array_get_field(array, 2, &error) == nullptr); + REQUIRE(vx_array_get_field(session, array, 2, &error) == nullptr); REQUIRE(error != nullptr); vx_error_free(error); } @@ -751,7 +751,7 @@ TEST_CASE("Filter with literal expression", "[filter]") { REQUIRE(vx_array_len(array) == SAMPLE_ROWS - threshold); vx_error *error = nullptr; - const vx_array *filtered_age = vx_array_get_field(array, 0, &error); + const vx_array *filtered_age = vx_array_get_field(session, array, 0, &error); require_no_error(error); REQUIRE(filtered_age != nullptr); defer { diff --git a/vortex-spatial/benches/area.rs b/vortex-spatial/benches/area.rs index bef930b53ba..904dd7d8fd2 100644 --- a/vortex-spatial/benches/area.rs +++ b/vortex-spatial/benches/area.rs @@ -128,9 +128,11 @@ fn multipolygon(bencher: Bencher) { #[divan::bench] fn nullable_polygon(bencher: Bencher) { + let mut ctx = SESSION.create_execution_ctx(); let geometry = MaskedArray::try_new( simple_polygons(), Validity::from_iter((0..ROWS).map(|row| !row.is_multiple_of(8))), + &mut ctx, ) .unwrap() .into_array(); diff --git a/vortex-spatial/benches/binary_predicates.rs b/vortex-spatial/benches/binary_predicates.rs index 8cc25888e59..90c6479efb2 100644 --- a/vortex-spatial/benches/binary_predicates.rs +++ b/vortex-spatial/benches/binary_predicates.rs @@ -148,21 +148,30 @@ fn execute(array: VortexResult, ctx: &mut ExecutionCtx) -> Array /// Marks one row in `null_every` null without changing the geometry storage. fn nullable_every(array: ArrayRef, null_every: usize) -> ArrayRef { + let mut ctx = SESSION.create_execution_ctx(); let validity = Validity::from_iter((0..array.len()).map(|i| !i.is_multiple_of(null_every))); - MaskedArray::try_new(array, validity).unwrap().into_array() + MaskedArray::try_new(array, validity, &mut ctx) + .unwrap() + .into_array() } /// A deterministic 90%-null validity pattern. fn ninety_percent_null(array: ArrayRef) -> ArrayRef { + let mut ctx = SESSION.create_execution_ctx(); let validity = Validity::from_iter((0..array.len()).map(|i| i.is_multiple_of(10))); - MaskedArray::try_new(array, validity).unwrap().into_array() + MaskedArray::try_new(array, validity, &mut ctx) + .unwrap() + .into_array() } /// A deterministic 50% validity pattern with a distinct phase per operand. Adjacent phases make /// the two columns jointly cover every valid/null combination once per four rows. fn half_valid(array: ArrayRef, phase: usize) -> ArrayRef { + let mut ctx = SESSION.create_execution_ctx(); let validity = Validity::from_iter((0..array.len()).map(|i| (i + phase) % 4 < 2)); - MaskedArray::try_new(array, validity).unwrap().into_array() + MaskedArray::try_new(array, validity, &mut ctx) + .unwrap() + .into_array() } mod contains { diff --git a/vortex-spatial/benches/convex_hull.rs b/vortex-spatial/benches/convex_hull.rs index b39207430f1..8145cd8af33 100644 --- a/vortex-spatial/benches/convex_hull.rs +++ b/vortex-spatial/benches/convex_hull.rs @@ -92,9 +92,11 @@ fn sixty_four_points(bencher: Bencher) { #[divan::bench] fn nullable_eight_points(bencher: Bencher) { + let mut ctx = SESSION.create_execution_ctx(); let input = MaskedArray::try_new( multipoints(ROWS, 8), Validity::from_iter((0..ROWS).map(|row| !row.is_multiple_of(8))), + &mut ctx, ) .unwrap() .into_array(); diff --git a/vortex-spatial/benches/distance.rs b/vortex-spatial/benches/distance.rs index 50714841e1e..c3edfd8502f 100644 --- a/vortex-spatial/benches/distance.rs +++ b/vortex-spatial/benches/distance.rs @@ -131,11 +131,11 @@ fn polygon_column_x_constant_polygon(bencher: Bencher) { #[divan::bench] fn nullable_point_column_x_constant_point(bencher: Bencher) { + let mut ctx = SESSION.create_execution_ctx(); let validity = Validity::from_iter((0..POINT_ROWS).map(|i| !i.is_multiple_of(8))); - let points = MaskedArray::try_new(points(POINT_ROWS, 0), validity) + let points = MaskedArray::try_new(points(POINT_ROWS, 0), validity, &mut ctx) .unwrap() .into_array(); - let mut ctx = SESSION.create_execution_ctx(); let point = point_constant(0.0, 0.0, POINT_ROWS, &mut ctx); bench_distance(bencher, points, point); } diff --git a/vortex-spatial/benches/length.rs b/vortex-spatial/benches/length.rs index cdee06526c2..24a0108e851 100644 --- a/vortex-spatial/benches/length.rs +++ b/vortex-spatial/benches/length.rs @@ -89,13 +89,14 @@ fn sixteen_vertex_lines(bencher: Bencher) { #[divan::bench] fn nullable_two_vertex_lines(bencher: Bencher) { + let mut ctx = SESSION.create_execution_ctx(); let lines = MaskedArray::try_new( linestrings(2), Validity::from_iter((0..ROWS).map(|i| !i.is_multiple_of(8))), + &mut ctx, ) .unwrap() .into_array(); - let mut ctx = SESSION.create_execution_ctx(); bencher .counter(ItemsCount::new(ROWS)) .bench_local(|| lengths(&lines, &mut ctx)); diff --git a/vortex-spatial/src/scalar_fn/collect.rs b/vortex-spatial/src/scalar_fn/collect.rs index d41bd6c489c..54ca7eb58dd 100644 --- a/vortex-spatial/src/scalar_fn/collect.rs +++ b/vortex-spatial/src/scalar_fn/collect.rs @@ -319,6 +319,7 @@ mod tests { use vortex_array::Columnar; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::ListArray; @@ -365,7 +366,7 @@ mod tests { /// Assert that `ST_Collect` of `input` equals `expected`. fn assert_collects(input: ArrayRef, expected: ArrayRef) -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let result = SpatialCollect::try_new(input)?.into_array(); assert_arrays_eq!(result, expected, &mut ctx); @@ -409,7 +410,7 @@ mod tests { #[test] fn all_valid_collect_reuses_geometry_storage() -> VortexResult<()> { let points = point_column(vec![0.0, 1.0, 2.0], vec![3.0, 4.0, 5.0])?; - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let point_storage = points .clone() .execute::(&mut ctx)? @@ -440,7 +441,7 @@ mod tests { } else { point_column(vec![0.0, 1.0, 2.0], vec![3.0, 4.0, 5.0])? }; - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let input = list(points, &[0, 2, 3])?; assert!( input @@ -499,7 +500,7 @@ mod tests { nullable_point_column(vec![Some((0.0, 2.0)), None, Some((1.0, 3.0))])?, &[0, 3], )?; - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let scalar = input.execute_scalar(0, &mut ctx)?; let input = ConstantArray::new(scalar, 3).into_array(); @@ -538,6 +539,7 @@ mod tests { #[test] fn propagates_null_list_rows() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let input = list_with_validity( point_column(vec![0.0, 1.0], vec![2.0, 3.0])?, &[0, 1, 2], @@ -546,6 +548,7 @@ mod tests { let expected = MaskedArray::try_new( multipoint_column(vec![vec![(0.0, 2.0)], vec![(1.0, 3.0)]])?, Validity::from_iter([true, false]), + &mut ctx, )? .into_array(); diff --git a/vortex-spatial/src/scalar_fn/convex_hull.rs b/vortex-spatial/src/scalar_fn/convex_hull.rs index 840a606e878..b8266d48622 100644 --- a/vortex-spatial/src/scalar_fn/convex_hull.rs +++ b/vortex-spatial/src/scalar_fn/convex_hull.rs @@ -176,12 +176,14 @@ mod tests { #[test] fn propagates_nulls() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); let input = MaskedArray::try_new( multipoint_column(vec![ vec![(0.0, 0.0), (1.0, 0.0), (0.0, 1.0)], vec![(2.0, 2.0)], ])?, Validity::from_iter([true, false]), + &mut ctx, )? .into_array(); let expected = MaskedArray::try_new( @@ -190,10 +192,10 @@ mod tests { vec![vec![(2.0, 2.0), (2.0, 2.0)]], ])?, Validity::from_iter([true, false]), + &mut ctx, )? .into_array(); let result = SpatialConvexHull::try_new(input)?.into_array(); - let mut ctx = vortex_array::array_session().create_execution_ctx(); assert_arrays_eq!(result, expected, &mut ctx); Ok(()) @@ -221,6 +223,7 @@ mod tests { #[test] fn collect_hull_area_pipeline() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); let points = point_column( vec![0.0, 2.0, 2.0, 0.0, 1.0, 0.0, 1.0, 2.0], vec![0.0, 0.0, 2.0, 2.0, 1.0, 0.0, 1.0, 2.0], @@ -236,7 +239,6 @@ mod tests { let hulls = SpatialConvexHull::try_new(collected)?.into_array(); let areas = SpatialArea::try_new(hulls)?.into_array(); let expected = PrimitiveArray::from_iter([4.0_f64, 0.0]).into_array(); - let mut ctx = vortex_array::array_session().create_execution_ctx(); assert_arrays_eq!(areas, expected, &mut ctx); Ok(()) diff --git a/vortex-spatial/src/scalar_fn/length.rs b/vortex-spatial/src/scalar_fn/length.rs index 09b5a1bd274..a539f9f0c63 100644 --- a/vortex-spatial/src/scalar_fn/length.rs +++ b/vortex-spatial/src/scalar_fn/length.rs @@ -297,6 +297,7 @@ mod tests { vec![vec![(0.0, 0.0), (0.0, 5.0)]], ])?, Validity::from_iter([true, false, true]), + &mut ctx, )? .into_array(); @@ -359,6 +360,7 @@ mod tests { vec![(0.0, 0.0), (0.0, 4.0)], ])?, Validity::from_iter([true, false, true]), + &mut ctx, )? .into_array(); @@ -382,6 +384,7 @@ mod tests { vec![(0.0, 0.0), (0.0, 4.0)], ])?, Validity::from_iter([false, false]), + &mut ctx, )? .into_array(); let expected = diff --git a/vortex-tensor/benches/inner_product.rs b/vortex-tensor/benches/inner_product.rs index e6f67b7d814..9d6905d5861 100644 --- a/vortex-tensor/benches/inner_product.rs +++ b/vortex-tensor/benches/inner_product.rs @@ -17,6 +17,7 @@ use mimalloc::MiMalloc; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::array_session; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::MaskedArray; use vortex_array::arrays::PrimitiveArray; @@ -57,7 +58,7 @@ fn vectors(width: usize, seed: usize) -> ArrayRef { } fn bench_inner_product(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef) { - let session = vortex_array::array_session(); + let session = array_session(); bencher .counter(ItemsCount::new(lhs.len())) .with_inputs(|| { @@ -78,8 +79,9 @@ fn non_nullable(bencher: Bencher, width: usize) { #[divan::bench(args = WIDTHS)] fn nullable(bencher: Bencher, width: usize) { + let mut ctx = array_session().create_execution_ctx(); let validity = Validity::from_iter((0..ELEMENTS / width).map(|i| i % 8 != 0)); - let lhs = MaskedArray::try_new(vectors(width, 0), validity) + let lhs = MaskedArray::try_new(vectors(width, 0), validity, &mut ctx) .unwrap() .into_array(); bench_inner_product(bencher, lhs, vectors(width, 31)); diff --git a/vortex-tensor/benches/l2_norm.rs b/vortex-tensor/benches/l2_norm.rs index bb5dc0dd0ab..a7c3d9bae83 100644 --- a/vortex-tensor/benches/l2_norm.rs +++ b/vortex-tensor/benches/l2_norm.rs @@ -17,6 +17,7 @@ use mimalloc::MiMalloc; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::array_session; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::MaskedArray; use vortex_array::arrays::PrimitiveArray; @@ -55,7 +56,7 @@ fn vectors(width: usize) -> ArrayRef { } fn bench_l2_norm(bencher: Bencher, input: ArrayRef) { - let session = vortex_array::array_session(); + let session = array_session(); bencher .counter(ItemsCount::new(input.len())) .with_inputs(|| { @@ -74,8 +75,9 @@ fn non_nullable(bencher: Bencher, width: usize) { #[divan::bench(args = WIDTHS)] fn nullable(bencher: Bencher, width: usize) { + let mut ctx = array_session().create_execution_ctx(); let validity = Validity::from_iter((0..ELEMENTS / width).map(|i| i % 8 != 0)); - let input = MaskedArray::try_new(vectors(width), validity) + let input = MaskedArray::try_new(vectors(width), validity, &mut ctx) .unwrap() .into_array(); bench_l2_norm(bencher, input); diff --git a/vortex-tensor/src/scalar_fns/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index d2629ec5b28..069a8b3663a 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -393,13 +393,14 @@ mod tests { #[test] fn null_input_row() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); // 2 rows of dim-2 vectors. Row 1 of rhs is masked as null. let lhs = tensor_array(&[2], &[3.0, 4.0, 1.0, 0.0])?; let rhs = tensor_array(&[2], &[3.0, 4.0, 0.0, 1.0])?; - let rhs = MaskedArray::try_new(rhs, Validity::from_iter([true, false]))?.into_array(); + let rhs = + MaskedArray::try_new(rhs, Validity::from_iter([true, false]), &mut ctx)?.into_array(); let result = CosineSimilarity::try_new(lhs, rhs)?; - let mut ctx = SESSION.create_execution_ctx(); let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; // Row 0: self-similarity = 1.0, row 1: null. diff --git a/vortex-tensor/src/scalar_fns/inner_product.rs b/vortex-tensor/src/scalar_fns/inner_product.rs index 1a843a1fd27..35f7e32e343 100644 --- a/vortex-tensor/src/scalar_fns/inner_product.rs +++ b/vortex-tensor/src/scalar_fns/inner_product.rs @@ -274,13 +274,14 @@ mod tests { #[test] fn null_input_row() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); // 3 rows of dim-2 vectors. Row 1 of lhs is masked as null. let lhs = tensor_array(&[2], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0])?; let rhs = tensor_array(&[2], &[7.0, 8.0, 9.0, 10.0, 11.0, 12.0])?; - let lhs = MaskedArray::try_new(lhs, Validity::from_iter([true, false, true]))?.into_array(); + let lhs = MaskedArray::try_new(lhs, Validity::from_iter([true, false, true]), &mut ctx)? + .into_array(); let result = InnerProduct::try_new(lhs, rhs)?; - let mut ctx = SESSION.create_execution_ctx(); let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; // Row 0: 1*7 + 2*8 = 23, row 1: null, row 2: 5*11 + 6*12 = 127. diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index 2c5fb5ccea8..ea5bbefd5ae 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -241,12 +241,13 @@ mod tests { #[test] fn null_input_row() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); // 2 rows of dim-2 vectors. Row 1 is masked as null. let arr = tensor_array(&[2], &[3.0, 4.0, 0.0, 0.0])?; - let arr = MaskedArray::try_new(arr, Validity::from_iter([true, false]))?.into_array(); + let arr = + MaskedArray::try_new(arr, Validity::from_iter([true, false]), &mut ctx)?.into_array(); let result = L2Norm::try_new(arr)?; - let mut ctx = SESSION.create_execution_ctx(); let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; // Row 0: norm = 5.0, row 1: null. diff --git a/vortex-tensor/src/scalar_fns/l2_normalize.rs b/vortex-tensor/src/scalar_fns/l2_normalize.rs index 73ef976726f..0922ee26e3c 100644 --- a/vortex-tensor/src/scalar_fns/l2_normalize.rs +++ b/vortex-tensor/src/scalar_fns/l2_normalize.rs @@ -254,6 +254,7 @@ mod tests { use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::MaskedArray; @@ -364,6 +365,7 @@ mod tests { #[test] fn test_input_nulls_produce_null_struct_rows() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let input = vector_array( 2, &[ @@ -371,9 +373,9 @@ mod tests { 1.0, 0.0, // The null row's placeholder values. ], )?; - let input = MaskedArray::try_new(input, Validity::from_iter([true, false]))?.into_array(); + let input = + MaskedArray::try_new(input, Validity::from_iter([true, false]), &mut ctx)?.into_array(); let output = evaluate(input)?; - let mut ctx = SESSION.create_execution_ctx(); assert!(output.is_valid(0, &mut ctx)?); assert!(!output.is_valid(1, &mut ctx)?); @@ -460,8 +462,9 @@ mod tests { #[test] fn test_return_dtype_has_nonnullable_fields() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let input = vector_array(2, &[3.0f64, 4.0])?; - let input = MaskedArray::try_new(input, Validity::AllValid)?.into_array(); + let input = MaskedArray::try_new(input, Validity::AllValid, &mut ctx)?.into_array(); let output = L2Normalize::try_new(input)?.into_array(); let DType::Struct(fields, Nullability::Nullable) = output.dtype() else { diff --git a/vortex/benches/common_encoding_tree_throughput.rs b/vortex/benches/common_encoding_tree_throughput.rs index 8dcc482bf8f..0077e0006c0 100644 --- a/vortex/benches/common_encoding_tree_throughput.rs +++ b/vortex/benches/common_encoding_tree_throughput.rs @@ -326,7 +326,6 @@ mod setup { ), compressed_codes, fsst.uncompressed_lengths().clone(), - &mut ctx, ) .unwrap(); From e1b5a060106cb636da4c2c29e3dc3cba5bd2b4fc Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 18 Oct 2026 11:41:21 +0000 Subject: [PATCH 2/3] Drop changes that belong to the other legacy_session PRs The masked commit also carried the Display, VarBinView, Null take, ZstdBuffers, RunEnd, FSST and FFI changes, which have their own PRs. Restore those files to develop so this PR only changes MaskedArray and its callers. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Y1zHpE1K1Y9HvyU4kt24Tq --- encodings/fsst/src/array.rs | 33 ++++++-- encodings/fsst/src/canonical.rs | 3 + encodings/fsst/src/compress.rs | 7 +- encodings/fsst/src/compute/filter.rs | 1 + encodings/fsst/src/compute/mod.rs | 1 + encodings/runend/src/array.rs | 83 +++++++------------ encodings/zstd/src/zstd_buffers.rs | 11 ++- lang/cpp/examples/reader.cpp | 2 +- lang/cpp/include/vortex/array.hpp | 6 +- lang/cpp/src/array.cpp | 12 +-- lang/cpp/tests/array.cpp | 20 ++--- lang/cpp/tests/arrow.cpp | 2 +- lang/cpp/tests/scan.cpp | 16 ++-- vortex-array/src/arrays/mod.rs | 1 - vortex-array/src/arrays/null/compute/rules.rs | 2 + vortex-array/src/arrays/null/compute/take.rs | 23 +++-- vortex-array/src/arrays/null/mod.rs | 10 --- vortex-array/src/arrays/varbinview/array.rs | 40 +++------ vortex-array/src/display/mod.rs | 26 ++---- vortex-btrblocks/src/schemes/string/fsst.rs | 1 + vortex-cuda/benches/fsst_cuda.rs | 1 + vortex-cuda/src/arrow/canonical.rs | 1 + vortex-ffi/cinclude/vortex.h | 5 +- vortex-ffi/src/array.rs | 24 +++--- vortex-ffi/test/array.cpp | 8 +- vortex-ffi/test/scan.cpp | 8 +- .../common_encoding_tree_throughput.rs | 1 + 27 files changed, 159 insertions(+), 189 deletions(-) diff --git a/encodings/fsst/src/array.rs b/encodings/fsst/src/array.rs index 4ce0504cd88..3f8bd2e5541 100644 --- a/encodings/fsst/src/array.rs +++ b/encodings/fsst/src/array.rs @@ -26,6 +26,7 @@ use vortex_array::EqMode; use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; use vortex_array::TypedArrayRef; +use vortex_array::VortexSessionExecute; use vortex_array::array_slots; use vortex_array::arrays::Primitive; use vortex_array::arrays::VarBin; @@ -39,6 +40,7 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::OffsetBuilderPType; use vortex_array::dtype::PType; +use vortex_array::legacy_session; use vortex_array::match_each_integer_ptype; use vortex_array::match_each_varbin_builder; use vortex_array::require_child; @@ -126,6 +128,7 @@ impl VTable for FSST { *ID } + #[allow(clippy::disallowed_methods)] fn validate( &self, data: &Self::TypedArrayData, @@ -133,7 +136,9 @@ impl VTable for FSST { len: usize, slots: &[Option], ) -> VortexResult<()> { - data.validate(dtype, len, slots) + // TODO(ctx): trait fixes - VTable::validate has a fixed signature. + let mut ctx = legacy_session().create_execution_ctx(); + data.validate(dtype, len, slots, &mut ctx) } fn nbuffers(_array: ArrayView<'_, Self>) -> usize { @@ -231,12 +236,13 @@ impl VTable for FSST { metadata: &[u8], buffers: &[BufferHandle], children: &dyn ArrayChildren, - _session: &VortexSession, + session: &VortexSession, ) -> VortexResult> { let metadata = FSSTMetadata::decode(metadata)?; let symbols = Buffer::::from_byte_buffer(buffers[0].clone().try_to_host_sync()?); let symbol_lengths = Buffer::::from_byte_buffer(buffers[1].clone().try_to_host_sync()?); + let mut ctx = session.create_execution_ctx(); if buffers.len() == 2 { return Self::deserialize_legacy( self, @@ -246,6 +252,7 @@ impl VTable for FSST { &symbols, &symbol_lengths, children, + &mut ctx, ); } @@ -288,6 +295,7 @@ impl VTable for FSST { &uncompressed_lengths, dtype, len, + &mut ctx, )?; let slots = FSSTSlots { uncompressed_lengths, @@ -606,6 +614,7 @@ impl FSST { symbol_lengths: Buffer, codes: VarBinArray, uncompressed_lengths: ArrayRef, + ctx: &mut ExecutionCtx, ) -> VortexResult { let len = codes.len(); FSSTData::validate_parts_from_codes( @@ -615,6 +624,7 @@ impl FSST { &uncompressed_lengths, &dtype, len, + ctx, )?; let slots = FSSTData::make_slots(&codes, &uncompressed_lengths); let codes_bytes = codes.bytes_handle().clone(); @@ -627,6 +637,7 @@ impl FSST { symbol_table: Arc, codes: VarBinArray, uncompressed_lengths: ArrayRef, + ctx: &mut ExecutionCtx, ) -> VortexResult { let len = codes.len(); FSSTData::validate_parts_from_codes( @@ -636,6 +647,7 @@ impl FSST { &uncompressed_lengths, &dtype, len, + ctx, )?; let slots = FSSTData::make_slots(&codes, &uncompressed_lengths); let codes_bytes = codes.bytes_handle().clone(); @@ -656,6 +668,7 @@ impl FSST { symbols: &Buffer, symbol_lengths: &Buffer, children: &dyn ArrayChildren, + ctx: &mut ExecutionCtx, ) -> VortexResult> { if children.len() != 2 { vortex_bail!(InvalidArgument: "Expected 2 children, got {}", children.len()); @@ -686,6 +699,7 @@ impl FSST { &uncompressed_lengths, dtype, len, + ctx, )?; let slots = FSSTData::make_slots(&codes, &uncompressed_lengths); let codes_bytes = codes.bytes_handle().clone(); @@ -778,6 +792,7 @@ impl FSSTData { dtype: &DType, len: usize, slots: &[Option], + ctx: &mut ExecutionCtx, ) -> VortexResult<()> { let fsst_slots = FSSTSlotsView::from_slots(slots); Self::validate_parts( @@ -789,6 +804,7 @@ impl FSSTData { fsst_slots.uncompressed_lengths, dtype, len, + ctx, ) } @@ -803,6 +819,7 @@ impl FSSTData { uncompressed_lengths: &ArrayRef, dtype: &DType, len: usize, + ctx: &mut ExecutionCtx, ) -> VortexResult<()> { vortex_ensure!( matches!(dtype, DType::Binary(_) | DType::Utf8(_)), @@ -847,12 +864,14 @@ impl FSSTData { // checked by the decode plan instead, once `execute` has made it primitive. if codes_bytes.is_on_host() && codes_offsets.is_host() - && let Some(codes_offsets) = codes_offsets.as_opt::() + && codes_offsets.is::() && !codes_offsets.is_empty() { - let last_offset: usize = match_each_integer_ptype!(codes_offsets.ptype(), |O| { - codes_offsets.as_slice::()[codes_offsets.len() - 1].as_() - }); + let last_offset: usize = (&codes_offsets + .execute_scalar(codes_offsets.len() - 1, ctx) + .vortex_expect("offsets must support scalar_at")) + .try_into() + .vortex_expect("Failed to convert offset to usize"); vortex_ensure!( last_offset <= codes_bytes.len(), InvalidArgument: "Last codes offset {} exceeds codes bytes length {}", @@ -898,6 +917,7 @@ impl FSSTData { uncompressed_lengths: &ArrayRef, dtype: &DType, len: usize, + ctx: &mut ExecutionCtx, ) -> VortexResult<()> { Self::validate_parts( symbols, @@ -908,6 +928,7 @@ impl FSSTData { uncompressed_lengths, dtype, len, + ctx, ) } diff --git a/encodings/fsst/src/canonical.rs b/encodings/fsst/src/canonical.rs index 18b29e4aab3..9087b21cba9 100644 --- a/encodings/fsst/src/canonical.rs +++ b/encodings/fsst/src/canonical.rs @@ -374,6 +374,7 @@ mod tests { encoded.symbol_table(), codes, encoded.uncompressed_lengths().slice(1..3)?, + &mut ctx, )?; let mut builder = VarBinViewBuilder::with_capacity_in(input.dtype().clone(), 2, ctx.allocator().clone()); @@ -398,6 +399,7 @@ mod tests { encoded.symbol_table(), encoded.codes(), PrimitiveArray::from_iter([4i32]).into_array(), + &mut ctx, )?; assert!(fsst_decode_bytes(invalid.as_view(), &mut ctx).is_err()); @@ -421,6 +423,7 @@ mod tests { encoded.symbol_table(), encoded.codes(), lengths, + &mut ctx, )?; assert!(fsst_decode_bytes(invalid.as_view(), &mut ctx).is_err()); } diff --git a/encodings/fsst/src/compress.rs b/encodings/fsst/src/compress.rs index f116605808b..e931d515b5c 100644 --- a/encodings/fsst/src/compress.rs +++ b/encodings/fsst/src/compress.rs @@ -228,7 +228,7 @@ where } } } - sink.finish(strings.dtype().clone()) + sink.finish(strings.dtype().clone(), ctx) } fn compress_varbin( @@ -252,7 +252,7 @@ where let off = offsets.as_slice::(); for_each_varbin_row(off, bytes, mask, |row| sink.emit(row)); }); - sink.finish(strings.dtype().clone()) + sink.finish(strings.dtype().clone(), ctx) } /// Call `f` once per row of a `VarBinArray` with the row bytes or `None`. @@ -333,7 +333,7 @@ impl<'c, O: OffsetBuilderPType + 'static> FsstSink<'c, O> { self.builder.append_value(&self.buffer); } - fn finish(mut self, dtype: DType) -> VortexResult { + fn finish(mut self, dtype: DType, ctx: &mut ExecutionCtx) -> VortexResult { let codes = self.builder.finish_into_varbin(); // Pad the symbol table here so that the array can hand it straight to a `Decompressor` // without copying. @@ -346,6 +346,7 @@ impl<'c, O: OffsetBuilderPType + 'static> FsstSink<'c, O> { )?), codes, self.uncompressed_lengths.into_array(), + ctx, ) } } diff --git a/encodings/fsst/src/compute/filter.rs b/encodings/fsst/src/compute/filter.rs index bb1f9ec6a65..45a95520aca 100644 --- a/encodings/fsst/src/compute/filter.rs +++ b/encodings/fsst/src/compute/filter.rs @@ -37,6 +37,7 @@ impl FilterKernel for FSST { array.symbol_table(), filtered_codes, array.uncompressed_lengths().filter(mask.clone())?, + ctx, )? .into_array(), )) diff --git a/encodings/fsst/src/compute/mod.rs b/encodings/fsst/src/compute/mod.rs index e1564cffe9a..0b95b75e1e7 100644 --- a/encodings/fsst/src/compute/mod.rs +++ b/encodings/fsst/src/compute/mod.rs @@ -77,6 +77,7 @@ impl TakeExecute for FSST { .fill_null(Scalar::zero_value( &array.uncompressed_lengths_dtype().clone(), ))?, + ctx, )? .into_array(), )) diff --git a/encodings/runend/src/array.rs b/encodings/runend/src/array.rs index 01bed165e25..bfb60c9233e 100644 --- a/encodings/runend/src/array.rs +++ b/encodings/runend/src/array.rs @@ -33,6 +33,7 @@ use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; +use vortex_array::legacy_session; use vortex_array::require_child; use vortex_array::serde::ArrayChildren; use vortex_array::validity::Validity; @@ -92,6 +93,7 @@ impl VTable for RunEnd { *ID } + #[allow(clippy::disallowed_methods)] fn validate( &self, data: &Self::TypedArrayData, @@ -102,7 +104,9 @@ impl VTable for RunEnd { let run_end_slots = RunEndSlotsView::from_slots(slots); let ends = run_end_slots.ends; let values = run_end_slots.values; - RunEndData::validate_shape(ends, values, data.offset, len)?; + // TODO(ctx): trait fixes - VTable::validate has a fixed signature. + let mut ctx = legacy_session().create_execution_ctx(); + RunEndData::validate_parts(ends, values, data.offset, len, &mut ctx)?; vortex_ensure_eq!(values.dtype(), dtype); Ok(()) } @@ -149,7 +153,7 @@ impl VTable for RunEnd { metadata: &[u8], _buffers: &[BufferHandle], children: &dyn ArrayChildren, - session: &VortexSession, + _session: &VortexSession, ) -> VortexResult> { let metadata = RunEndMetadata::decode(metadata)?; let ends_dtype = DType::Primitive(metadata.ends_ptype(), Nullability::NonNullable); @@ -158,13 +162,6 @@ impl VTable for RunEnd { let values = children.get(1, dtype, runs)?; let offset = usize::try_from(metadata.offset).vortex_expect("Offset must be a valid usize"); - RunEndData::validate_parts( - &ends, - &values, - offset, - len, - &mut session.create_execution_ctx(), - )?; let slots = RunEndSlots { ends, values }.into_slots(); let data = RunEndData::new(offset); Ok(ArrayParts::new( @@ -330,7 +327,22 @@ impl RunEndData { length: usize, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - if !Self::validate_shape(ends, values, offset, length)? { + // DType validation + vortex_ensure!( + ends.dtype().is_unsigned_int(), + "run ends must be unsigned integers, was {}", + ends.dtype(), + ); + vortex_ensure_eq!(ends.len(), values.len()); + + // Handle empty run-ends + if ends.is_empty() { + vortex_ensure_eq!(offset, 0, "non-zero offset provided for empty RunEndArray"); + return Ok(()); + } + + // Zero-length logical slices may retain run metadata from the source array. + if length == 0 { return Ok(()); } @@ -355,52 +367,15 @@ impl RunEndData { return Ok(()); } - let first_run_end = if offset == 0 { - 0 - } else { - usize::try_from(&ends.execute_scalar(0, ctx)?)? - }; - let last_run_end = usize::try_from(&ends.execute_scalar(ends.len() - 1, ctx)?)?; - Self::validate_run_end_bounds(first_run_end, last_run_end, offset, length) - } - - /// Validates the dtypes and lengths of the parts. Returns `false` when there are no runs to - /// check against `offset..offset + length`. - fn validate_shape( - ends: &ArrayRef, - values: &ArrayRef, - offset: usize, - length: usize, - ) -> VortexResult { - // DType validation - vortex_ensure!( - ends.dtype().is_unsigned_int(), - "run ends must be unsigned integers, was {}", - ends.dtype(), - ); - vortex_ensure_eq!(ends.len(), values.len()); - - // Handle empty run-ends - if ends.is_empty() { - vortex_ensure_eq!(offset, 0, "non-zero offset provided for empty RunEndArray"); - return Ok(false); - } - - // Zero-length logical slices may retain run metadata from the source array. - Ok(length != 0) - } - - /// Validates that the first and last run ends cover `offset..offset + length`. - fn validate_run_end_bounds( - first_run_end: usize, - last_run_end: usize, - offset: usize, - length: usize, - ) -> VortexResult<()> { - if first_run_end < offset { - vortex_bail!("First run end {first_run_end} must be >= offset {offset}"); + // 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.execute_scalar(0, 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.execute_scalar(ends.len() - 1, 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/zstd/src/zstd_buffers.rs b/encodings/zstd/src/zstd_buffers.rs index 9b94137676b..a390271cab1 100644 --- a/encodings/zstd/src/zstd_buffers.rs +++ b/encodings/zstd/src/zstd_buffers.rs @@ -98,7 +98,6 @@ impl ZstdBuffers { compressed_buffers, uncompressed_sizes, buffer_alignments, - session: session.clone(), }; let slots: ArraySlots = serialization.children.into_iter().map(Some).collect(); let compressed = Array::try_from_parts(ArrayParts::new( @@ -161,8 +160,6 @@ pub struct ZstdBuffersData { compressed_buffers: Vec, uncompressed_sizes: Vec, buffer_alignments: Vec, - /// The session used to rebuild the wrapped array, which resolves its encoding. - session: VortexSession, } impl Display for ZstdBuffersData { @@ -516,7 +513,6 @@ impl VTable for ZstdBuffers { compressed_buffers, uncompressed_sizes: metadata.uncompressed_sizes.clone(), buffer_alignments: metadata.buffer_alignments.clone(), - session: session.clone(), }; data.validate()?; @@ -558,13 +554,16 @@ impl OperationsVTable for ZstdBuffers { } impl ValidityVTable for ZstdBuffers { + #[allow(clippy::disallowed_methods)] fn validity(array: ArrayView<'_, ZstdBuffers>) -> VortexResult { if !array.dtype().is_nullable() { return Ok(Validity::NonNullable); } - let session = array.data().session.clone(); - let inner_array = ZstdBuffers::decompress_and_build_inner(&array.into_owned(), &session)?; + let inner_array = ZstdBuffers::decompress_and_build_inner( + &array.into_owned(), + vortex_array::legacy_session(), + )?; inner_array.validity() } } diff --git a/lang/cpp/examples/reader.cpp b/lang/cpp/examples/reader.cpp index a5706631bf5..a447ec6917b 100644 --- a/lang/cpp/examples/reader.cpp +++ b/lang/cpp/examples/reader.cpp @@ -21,7 +21,7 @@ int main() { for (Partition &partition : scan.partitions()) { for (Array &array : partition.batches()) { - const Array age = array.field(session, "age"); + const Array age = array.field("age"); const PrimitiveView age_view = age.values(session); const std::span age_values = age_view.values(); for (uint8_t value : age_values) { diff --git a/lang/cpp/include/vortex/array.hpp b/lang/cpp/include/vortex/array.hpp index 15ff95a4919..c60dde2a8bd 100644 --- a/lang/cpp/include/vortex/array.hpp +++ b/lang/cpp/include/vortex/array.hpp @@ -196,19 +196,19 @@ class Array { Validity validity() const; // Number of null/invalid elements in Array - size_t null_count(const Session &session) const; + size_t null_count() const; /** * Get a Struct field by index. Throws if Array is not a Struct or if index * is out of bounds. */ - Array field(const Session &session, size_t index) const; + Array field(size_t index) const; /** * Get a Struct field by name. Throws if Array is not a Struct or doesn't * have this named field. */ - Array field(const Session &session, std::string_view name) const; + Array field(std::string_view name) const; /* * Create a new Array slicing [begin; end) rows from original. diff --git a/lang/cpp/src/array.cpp b/lang/cpp/src/array.cpp index d1c398dfee2..7364c9abc34 100644 --- a/lang/cpp/src/array.cpp +++ b/lang/cpp/src/array.cpp @@ -228,21 +228,21 @@ Validity Array::validity() const { return Access::adopt(static_cast(raw.type), raw.array); } -size_t Array::null_count(const Session &session) const { +size_t Array::null_count() const { vx_error *error = nullptr; - const size_t count = vx_array_invalid_count(Access::c_ptr(session), handle_.get(), &error); + const size_t count = vx_array_invalid_count(handle_.get(), &error); throw_on_error(error); return count; } -Array Array::field(const Session &session, size_t index) const { +Array Array::field(size_t index) const { vx_error *error = nullptr; - const vx_array *out = vx_array_get_field(Access::c_ptr(session), handle_.get(), index, &error); + const vx_array *out = vx_array_get_field(handle_.get(), index, &error); throw_on_error(error); return Access::adopt(out); } -Array Array::field(const Session &session, std::string_view name) const { +Array Array::field(std::string_view name) const { const DataType dt = dtype(); const std::unique_ptr fields( struct_fields_or_throw(detail::Access::c_ptr(dt)), @@ -251,7 +251,7 @@ Array Array::field(const Session &session, std::string_view name) const { for (uint64_t i = 0; i < fields_size; ++i) { const vx_view field = vx_struct_fields_field_name(fields.get(), i); if (std::string_view {field.ptr, field.len} == name) { - return this->field(session, i); + return this->field(i); } } throw VortexException("no field named \"" + std::string(name) + "\"", ErrorCode::InvalidArgument); diff --git a/lang/cpp/tests/array.cpp b/lang/cpp/tests/array.cpp index fee9fd19b37..db6ea6b1eda 100644 --- a/lang/cpp/tests/array.cpp +++ b/lang/cpp/tests/array.cpp @@ -29,7 +29,7 @@ TEST_CASE("Empty array", "[array]") { REQUIRE(empty.size() == 0); REQUIRE(empty.is_primitive(I32)); - REQUIRE(empty.null_count(session) == 0); + REQUIRE(empty.null_count() == 0); auto view = empty.values(session); REQUIRE(view.size() == 0); @@ -42,7 +42,7 @@ void test_primitive_array(Array array, const int32_t *begin) { REQUIRE(array.size() == 3); REQUIRE(array.is_primitive(I32)); REQUIRE_FALSE(array.nullable()); - REQUIRE(array.null_count(session) == 0); + REQUIRE(array.null_count() == 0); auto view = array.values(session); REQUIRE(view.size() == 3); @@ -81,7 +81,7 @@ TEST_CASE("Validity from a boolean mask", "[array]") { Array a = Array::primitive(std::span(data), Validity::from_array(mask)); REQUIRE(a.nullable()); - REQUIRE(a.null_count(session) == 1); + REQUIRE(a.null_count() == 1); auto view = a.values(session); REQUIRE_FALSE(view.is_null(0)); @@ -105,7 +105,7 @@ TEST_CASE("AllInvalid", "[array]") { Session session; std::vector data = {1, 2}; Array a = Array::primitive(std::span(data), AllInvalid); - REQUIRE(a.null_count(session) == 2); + REQUIRE(a.null_count() == 2); auto view = a.values(session); REQUIRE(view.is_null(0)); REQUIRE(view.is_null(1)); @@ -129,25 +129,25 @@ TEST_CASE("make_struct and fields", "[array]") { REQUIRE(s.has_dtype(DataTypeVariant::Struct)); REQUIRE(s.dtype().fields().size() == 2); - Session session; - Array by_index = s.field(session, 0); + Array by_index = s.field(0); REQUIRE(by_index.is_primitive(U8)); - Array by_name = s.field(session, "height"); + Session session; + Array by_name = s.field("height"); REQUIRE(by_name.is_primitive(U16)); auto view = by_name.values(session); REQUIRE(std::ranges::equal(view.values(), heights)); REQUIRE(std::ranges::equal(by_index.values(session).values(), ages)); - REQUIRE_THROWS_AS(s.field(session, 2), VortexException); - REQUIRE_THROWS_AS(s.field(session, "nope"), VortexException); + REQUIRE_THROWS_AS(s.field(2), VortexException); + REQUIRE_THROWS_AS(s.field("nope"), VortexException); std::vector fields_vec; fields_vec.emplace_back("age", Array::primitive(ages)); Array other = make_struct(fields_vec); REQUIRE(other.size() == 3); REQUIRE(other.has_dtype(DataTypeVariant::Struct)); - REQUIRE(std::ranges::equal(other.field(session, "age").values(session).values(), ages)); + REQUIRE(std::ranges::equal(other.field("age").values(session).values(), ages)); } TEST_CASE("Mismatched field length", "[array]") { diff --git a/lang/cpp/tests/arrow.cpp b/lang/cpp/tests/arrow.cpp index ba2d7a60210..2d5c2a0246f 100644 --- a/lang/cpp/tests/arrow.cpp +++ b/lang/cpp/tests/arrow.cpp @@ -83,7 +83,7 @@ TEST_CASE("Import Arrow array as Vortex array", "[arrow]") { REQUIRE(vx.size() == 3); REQUIRE(vx.has_dtype(DataTypeVariant::Struct)); - Array a = vx.field(session, 0); + Array a = vx.field(0); REQUIRE(a.is_primitive(I32)); auto view = a.values(session); REQUIRE(view.size() == 3); diff --git a/lang/cpp/tests/scan.cpp b/lang/cpp/tests/scan.cpp index 18f5b519fb7..0082418e7ae 100644 --- a/lang/cpp/tests/scan.cpp +++ b/lang/cpp/tests/scan.cpp @@ -69,7 +69,7 @@ void verify_sample_array(const Session &session, size_t rows = SAMPLE_ROWS) { REQUIRE(array.size() == rows); REQUIRE(array.has_dtype(DataTypeVariant::Struct)); - Array age = array.field(session, 0); + Array age = array.field(0); REQUIRE(age.is_primitive(U8)); auto ages = age.values(session); for (size_t i = 0; i < rows; ++i) { @@ -77,7 +77,7 @@ void verify_sample_array(const Session &session, REQUIRE(ages.values()[i] == (first + i) % SAMPLE_ROWS); } - Array height = array.field(session, 1); + Array height = array.field(1); REQUIRE(height.is_primitive(U16)); auto view = height.values(session); for (size_t i = 0; i < rows; ++i) { @@ -85,7 +85,7 @@ void verify_sample_array(const Session &session, REQUIRE(view.values()[i] == ((first + i) % SAMPLE_ROWS + 1) % 200); } - REQUIRE_THROWS_AS(array.field(session, 2), VortexException); + REQUIRE_THROWS_AS(array.field(2), VortexException); } TEST_CASE("Basic scan", "[scan]") { @@ -264,7 +264,7 @@ TEST_CASE("Filter age", "[filter]") { REQUIRE(array.has_value()); REQUIRE(array->size() == SAMPLE_ROWS - threshold); - verify_age_field(session, array->field(session, 0), threshold, SAMPLE_ROWS - threshold); + verify_age_field(session, array->field(0), threshold, SAMPLE_ROWS - threshold); } TEST_CASE("Filter invalid values", "[filter]") { @@ -300,7 +300,7 @@ TEST_CASE("Filter with operators", "[filter]") { auto array = partition->next(); REQUIRE(array.has_value()); REQUIRE(array->size() == 5); - verify_age_field(session, array->field(session, "age"), 90, 5); + verify_age_field(session, array->field("age"), 90, 5); } } @@ -332,7 +332,7 @@ TEST_CASE("Row range and limit", "[scan]") { REQUIRE(array.has_value()); REQUIRE(array->size() == 5); - verify_age_field(session, array->field(session, 0), 10, 5); + verify_age_field(session, array->field(0), 10, 5); } TEST_CASE("Selection", "[scan]") { @@ -350,7 +350,7 @@ TEST_CASE("Selection", "[scan]") { REQUIRE(array.has_value()); REQUIRE(array->size() == 3); - auto ages = array->field(session, 0).values(session); + auto ages = array->field(0).values(session); REQUIRE(ages.values()[0] == 3); REQUIRE(ages.values()[1] == 5); REQUIRE(ages.values()[2] == 8); @@ -364,6 +364,6 @@ TEST_CASE("Selection", "[scan]") { REQUIRE(array.has_value()); REQUIRE(array->size() == 97); - verify_age_field(session, array->field(session, 0), 3, 97); + verify_age_field(session, array->field(0), 3, 97); } } // namespace diff --git a/vortex-array/src/arrays/mod.rs b/vortex-array/src/arrays/mod.rs index ed6e272b2b1..8a8b6a234f7 100644 --- a/vortex-array/src/arrays/mod.rs +++ b/vortex-array/src/arrays/mod.rs @@ -142,7 +142,6 @@ pub(crate) fn initialize(session: &VortexSession) { list::initialize(session); listview::initialize(session); map::initialize(session); - null::initialize(session); patched::initialize(session); primitive::initialize(session); struct_::initialize(session); diff --git a/vortex-array/src/arrays/null/compute/rules.rs b/vortex-array/src/arrays/null/compute/rules.rs index b2ee44a19dd..f63c904667a 100644 --- a/vortex-array/src/arrays/null/compute/rules.rs +++ b/vortex-array/src/arrays/null/compute/rules.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use crate::arrays::Null; +use crate::arrays::dict::TakeReduceAdaptor; use crate::arrays::filter::FilterReduceAdaptor; use crate::arrays::slice::SliceReduceAdaptor; use crate::optimizer::rules::ParentRuleSet; @@ -13,4 +14,5 @@ pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&CastReduceAdaptor(Null)), ParentRuleSet::lift(&MaskReduceAdaptor(Null)), ParentRuleSet::lift(&SliceReduceAdaptor(Null)), + ParentRuleSet::lift(&TakeReduceAdaptor(Null)), ]); diff --git a/vortex-array/src/arrays/null/compute/take.rs b/vortex-array/src/arrays/null/compute/take.rs index 6a60c94a54c..dc328614743 100644 --- a/vortex-array/src/arrays/null/compute/take.rs +++ b/vortex-array/src/arrays/null/compute/take.rs @@ -5,23 +5,23 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use crate::ArrayRef; -use crate::ExecutionCtx; use crate::IntoArray; +use crate::VortexSessionExecute; use crate::array::ArrayView; use crate::arrays::Null; use crate::arrays::NullArray; use crate::arrays::PrimitiveArray; -use crate::arrays::dict::TakeExecute; +use crate::arrays::dict::TakeReduce; +use crate::arrays::dict::TakeReduceAdaptor; use crate::match_each_integer_ptype; +use crate::optimizer::rules::ParentRuleSet; -impl TakeExecute for Null { +impl TakeReduce for Null { #[expect(clippy::cast_possible_truncation)] - fn take( - array: ArrayView<'_, Null>, - indices: &ArrayRef, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - let indices = indices.clone().execute::(ctx)?; + fn take(array: ArrayView<'_, Null>, indices: &ArrayRef) -> VortexResult> { + #[allow(clippy::disallowed_methods)] + let mut ctx = crate::legacy_session().create_execution_ctx(); + let indices = indices.clone().execute::(&mut ctx)?; // Enforce all indices are valid match_each_integer_ptype!(indices.ptype(), |T| { @@ -35,3 +35,8 @@ impl TakeExecute for Null { Ok(Some(NullArray::new(indices.len()).into_array())) } } + +impl Null { + pub const TAKE_RULES: ParentRuleSet = + ParentRuleSet::new(&[ParentRuleSet::lift(&TakeReduceAdaptor::(Self))]); +} diff --git a/vortex-array/src/arrays/null/mod.rs b/vortex-array/src/arrays/null/mod.rs index a4db271bf8a..62deab08eb8 100644 --- a/vortex-array/src/arrays/null/mod.rs +++ b/vortex-array/src/arrays/null/mod.rs @@ -11,7 +11,6 @@ use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::ArraySlots; -use crate::ArrayVTable; use crate::ExecutionCtx; use crate::ExecutionResult; use crate::array::Array; @@ -23,14 +22,11 @@ use crate::array::OperationsVTable; use crate::array::VTable; use crate::array::ValidityVTable; use crate::array::with_empty_buffers; -use crate::arrays::Dict; -use crate::arrays::dict::TakeExecuteAdaptor; use crate::arrays::null::compute::rules::PARENT_RULES; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::builders::NullBuilder; use crate::dtype::DType; -use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar::Scalar; use crate::serde::ArrayChildren; use crate::validity::Validity; @@ -40,12 +36,6 @@ pub(crate) mod compute; /// A [`Null`]-encoded Vortex array. pub type NullArray = Array; -pub(crate) fn initialize(session: &VortexSession) { - session - .kernels() - .register_execute_parent_kernel(Dict.id(), Null, TakeExecuteAdaptor(Null)); -} - impl VTable for Null { type TypedArrayData = EmptyArrayData; diff --git a/vortex-array/src/arrays/varbinview/array.rs b/vortex-array/src/arrays/varbinview/array.rs index 180704da9f2..8327f91db2b 100644 --- a/vortex-array/src/arrays/varbinview/array.rs +++ b/vortex-array/src/arrays/varbinview/array.rs @@ -25,17 +25,14 @@ use crate::ArrayRef; use crate::ArraySlots; use crate::ExecutionCtx; #[cfg(debug_assertions)] +use crate::VortexSessionExecute; use crate::array::Array; use crate::array::ArrayParts; use crate::array::TypedArrayRef; use crate::array::child_to_validity; use crate::array::validity_to_child; use crate::array_slots; -#[cfg(debug_assertions)] -use crate::arrays::Bool; use crate::arrays::VarBinView; -#[cfg(debug_assertions)] -use crate::arrays::bool::BoolArrayExt; use crate::arrays::varbinview::BinaryView; use crate::arrays::varbinview::ResolvedViews; use crate::buffer::BufferHandle; @@ -44,6 +41,7 @@ use crate::builders::VarBinViewBuilder; use crate::dtype::DType; use crate::dtype::Nullability; #[cfg(debug_assertions)] +use crate::legacy_session; use crate::validity::Validity; #[array_slots(VarBinView)] @@ -315,7 +313,15 @@ impl VarBinViewData { validity: Validity, ) -> Self { #[cfg(debug_assertions)] - Self::debug_validate(&views, &buffers, &dtype, &validity); + #[expect(clippy::disallowed_methods)] + Self::validate( + &views, + &buffers, + &dtype, + &validity, + &mut legacy_session().create_execution_ctx(), + ) + .vortex_expect("[Debug Assertion]: Invalid `VarBinViewArray` parameters"); let handles: Vec = buffers .iter() @@ -360,30 +366,6 @@ impl VarBinViewData { Self::validate_views(views, &mask, buffers, check_utf8) } - /// Debug-only validation for unchecked construction. - /// - /// This never executes the validity, so the views are only checked when the validity is - /// constant or a canonical host-resident bitmap. - #[cfg(debug_assertions)] - fn debug_validate( - views: &Buffer, - buffers: &Arc<[ByteBuffer]>, - dtype: &DType, - validity: &Validity, - ) { - let mask = match validity { - Validity::NonNullable | Validity::AllValid => Mask::AllTrue(views.len()), - Validity::AllInvalid => Mask::AllFalse(views.len()), - Validity::Array(array) => match array.as_opt::() { - Some(bools) if array.is_host() => Mask::from_buffer(bools.to_bit_buffer()), - _ => return, - }, - }; - Self::check_nullability(dtype, validity) - .and_then(|check_utf8| Self::validate_views(views, &mask, buffers, check_utf8)) - .vortex_expect("[Debug Assertion]: Invalid `VarBinViewArray` parameters"); - } - /// Validates components like validate() and replaces views at null slots to empty views pub(crate) fn validate_and_fix( views: Buffer, diff --git a/vortex-array/src/display/mod.rs b/vortex-array/src/display/mod.rs index 2fe4e6a4f36..780140af2ac 100644 --- a/vortex-array/src/display/mod.rs +++ b/vortex-array/src/display/mod.rs @@ -21,7 +21,7 @@ pub use tree_display::TreeDisplay; use crate::ArrayRef; use crate::VortexSessionExecute; -use crate::array_session; +use crate::legacy_session; /// Describe how to convert an array to a string. /// @@ -522,6 +522,7 @@ impl ArrayRef { DisplayArrayAs(self, DisplayOptions::TableDisplay) } + #[allow(clippy::disallowed_methods)] fn fmt_as(&self, f: &mut std::fmt::Formatter, options: &DisplayOptions) -> std::fmt::Result { match options { DisplayOptions::MetadataOnly => EncodingSummaryExtractor::write(self, f), @@ -536,25 +537,17 @@ impl ArrayRef { let limit = self.len().min(f.precision().unwrap_or(DISPLAY_LIMIT)); let is_truncated = self.len() > limit; - // Formatting cannot take an execution context, so the values are read with a - // default array session. - let mut ctx = array_session().create_execution_ctx(); - let mut fmt_scalar = |i| { - self.execute_scalar(i, &mut ctx) + let fmt_scalar = |i| { + self.execute_scalar(i, &mut legacy_session().create_execution_ctx()) .map_or_else(|e| format!(""), |s| s.to_string()) }; - let head = (0..limit.saturating_sub(3)) - .map(&mut fmt_scalar) - .collect::>(); - let tail = (self.len().saturating_sub(3)..self.len()) - .map(&mut fmt_scalar) - .collect::>(); write!( f, "{opening_brace}{}{closing_brace}", - head.into_iter() + (0..limit.saturating_sub(3)) + .map(fmt_scalar) .chain(repeat_n("...".to_string(), is_truncated as usize)) - .chain(tail) + .chain((self.len().saturating_sub(3)..self.len()).map(fmt_scalar)) .format(sep) ) } @@ -592,9 +585,8 @@ impl ArrayRef { use crate::dtype::DType; let mut builder = tabled::builder::Builder::default(); - // Formatting cannot take an execution context, so the rows are read with a default - // array session. Reuse a single execution context across all per-row accesses. - let mut ctx = array_session().create_execution_ctx(); + // Reuse a single execution context across all per-row accesses below. + let mut ctx = legacy_session().create_execution_ctx(); // Special logic for struct arrays. let DType::Struct(sf, _) = self.dtype() else { diff --git a/vortex-btrblocks/src/schemes/string/fsst.rs b/vortex-btrblocks/src/schemes/string/fsst.rs index 27f4fad3197..7cb87154d3d 100644 --- a/vortex-btrblocks/src/schemes/string/fsst.rs +++ b/vortex-btrblocks/src/schemes/string/fsst.rs @@ -122,6 +122,7 @@ impl Scheme for FSSTScheme { )?), compressed_codes, compressed_original_lengths, + exec_ctx, )?; Ok(fsst.into_array()) diff --git a/vortex-cuda/benches/fsst_cuda.rs b/vortex-cuda/benches/fsst_cuda.rs index 6924b56f79c..784856c0585 100644 --- a/vortex-cuda/benches/fsst_cuda.rs +++ b/vortex-cuda/benches/fsst_cuda.rs @@ -89,6 +89,7 @@ fn make_fixture(n: usize) -> FSSTBenchFixture { ), fsst.codes(), fsst.uncompressed_lengths().clone(), + setup_ctx.execution_ctx(), ) .vortex_expect("rebuild FSST fixture with Binary dtype") .into_array(); diff --git a/vortex-cuda/src/arrow/canonical.rs b/vortex-cuda/src/arrow/canonical.rs index 5f68ed512e9..5bea5a9921b 100644 --- a/vortex-cuda/src/arrow/canonical.rs +++ b/vortex-cuda/src/arrow/canonical.rs @@ -2729,6 +2729,7 @@ mod tests { fsst.symbol_table(), fsst.codes(), PrimitiveArray::from_iter([i32::MAX, i32::MAX]).into_array(), + ctx.execution_ctx(), )? .into_array(); let schema = arrow_schema_for_array(&oversized, &mut ctx)?; diff --git a/vortex-ffi/cinclude/vortex.h b/vortex-ffi/cinclude/vortex.h index c6b250fc9dc..384dc29305c 100644 --- a/vortex-ffi/cinclude/vortex.h +++ b/vortex-ffi/cinclude/vortex.h @@ -677,8 +677,7 @@ size_t vx_array_len(const vx_array *array); */ const vx_dtype *vx_array_dtype(const vx_array *array); -const vx_array * -vx_array_get_field(const vx_session *session, const vx_array *array, size_t index, vx_error **error_out); +const vx_array *vx_array_get_field(const vx_array *array, size_t index, vx_error **error_out); const vx_array *vx_array_slice(const vx_array *array, size_t start, size_t stop, vx_error **error_out); @@ -695,7 +694,7 @@ bool vx_array_element_is_invalid(const vx_session *session, /** * Check how many items in the array are invalid (null). */ -size_t vx_array_invalid_count(const vx_session *session, const vx_array *array, vx_error **error_out); +size_t vx_array_invalid_count(const vx_array *array, vx_error **error_out); /** * Increase reference count on vx_array diff --git a/vortex-ffi/src/array.rs b/vortex-ffi/src/array.rs index a47bf01e3a9..cf37e3cce37 100644 --- a/vortex-ffi/src/array.rs +++ b/vortex-ffi/src/array.rs @@ -24,6 +24,7 @@ use vortex::array::arrays::StructArray; use vortex::array::arrays::VarBinView; use vortex::array::arrays::bool::BoolArrayExt; use vortex::array::arrays::struct_::StructArrayExt; +use vortex::array::legacy_session; use vortex::array::validity::Validity; use vortex::buffer::BitBuffer; use vortex::buffer::Buffer; @@ -266,17 +267,16 @@ pub unsafe extern "C-unwind" fn vx_array_dtype(array: *const vx_array) -> *const // Returns NULL and sets error_out if index is out of bounds or array doesn't // have dtype DTYPE_STRUCT. #[unsafe(no_mangle)] +#[allow(clippy::disallowed_methods)] pub unsafe extern "C-unwind" fn vx_array_get_field( - session: *const vx_session, array: *const vx_array, index: usize, error_out: *mut *mut vx_error, ) -> *const vx_array { try_or_default(error_out, || { - let session = unsafe { vx_session_ref(session) }?; let array = vx_array::as_ref(array); - let mut ctx = session.create_execution_ctx(); + let mut ctx = legacy_session().create_execution_ctx(); let struct_array = array.clone().execute::(&mut ctx)?; let field_array = struct_array .unmasked_field_opt(index) @@ -305,6 +305,7 @@ pub unsafe extern "C-unwind" fn vx_array_slice( /// validity array. Sets error if index is out of bounds or underlying validity /// array is corrupted. #[unsafe(no_mangle)] +#[allow(clippy::disallowed_methods)] pub unsafe extern "C-unwind" fn vx_array_element_is_invalid( session: *const vx_session, array: *const vx_array, @@ -319,16 +320,15 @@ pub unsafe extern "C-unwind" fn vx_array_element_is_invalid( /// Check how many items in the array are invalid (null). #[unsafe(no_mangle)] +#[allow(clippy::disallowed_methods)] pub unsafe extern "C-unwind" fn vx_array_invalid_count( - session: *const vx_session, array: *const vx_array, error_out: *mut *mut vx_error, ) -> usize { try_or_default(error_out, || { - let session = unsafe { vx_session_ref(session) }?; vortex_ensure!(!array.is_null()); let array = vx_array::as_ref(array); - array.invalid_count(&mut session.create_execution_ctx()) + array.invalid_count(&mut legacy_session().create_execution_ctx()) }) } @@ -879,7 +879,7 @@ mod tests { )); assert_no_error(error); - let null_count = vx_array_invalid_count(session, ffi_array, &raw mut error); + let null_count = vx_array_invalid_count(ffi_array, &raw mut error); assert_no_error(error); assert_eq!(null_count, 1); @@ -906,7 +906,7 @@ mod tests { let ffi_array = vx_array::new(struct_array.into_array()); let mut error = ptr::null_mut(); - let field0 = vx_array_get_field(session, ffi_array, 0, &raw mut error); + let field0 = vx_array_get_field(ffi_array, 0, &raw mut error); assert_no_error(error); assert_eq!(vx_array_len(field0), 3); for (index, expected) in ["Alice", "Bob", "Charlie"].into_iter().enumerate() { @@ -915,7 +915,7 @@ mod tests { assert_eq!(actual.as_str().unwrap(), expected); } - let field1 = vx_array_get_field(session, ffi_array, 1, &raw mut error); + let field1 = vx_array_get_field(ffi_array, 1, &raw mut error); assert_no_error(error); assert_eq!(vx_array_len(field1), 3); assert_eq!(get_u8(session, field1, 0), 30); @@ -923,7 +923,7 @@ mod tests { assert_eq!(get_u8(session, field1, 2), 35); // Test out of bounds - let field_oob = vx_array_get_field(session, ffi_array, 2, &raw mut error); + let field_oob = vx_array_get_field(ffi_array, 2, &raw mut error); assert!(!error.is_null()); assert!(field_oob.is_null()); vx_error_free(error); @@ -1117,14 +1117,14 @@ mod tests { assert_eq!(vx_array_len(vx), 3); assert!(!vx_array_is_nullable(vx)); - let a = vx_array_get_field(session, vx, 0, &raw mut error); + let a = vx_array_get_field(vx, 0, &raw mut error); assert_no_error(error); assert!(vx_array_is_primitive(a, vx_ptype::PTYPE_I32)); assert_eq!(get_i32(session, a, 0), 1); assert_eq!(get_i32(session, a, 2), 3); vx_array_free(a); - let b = vx_array_get_field(session, vx, 1, &raw mut error); + let b = vx_array_get_field(vx, 1, &raw mut error); assert_no_error(error); assert!(vx_array_has_dtype(b, vx_dtype_variant::DTYPE_UTF8)); assert!(vx_array_element_is_invalid(session, b, 1, &raw mut error)); diff --git a/vortex-ffi/test/array.cpp b/vortex-ffi/test/array.cpp index b792b168764..e49e94de10d 100644 --- a/vortex-ffi/test/array.cpp +++ b/vortex-ffi/test/array.cpp @@ -159,7 +159,7 @@ TEST_CASE("Array with validity", "[array]") { REQUIRE(array_get_u8(session, array, i) == 7); } } - REQUIRE(vx_array_invalid_count(session, array, &error) == 2); + REQUIRE(vx_array_invalid_count(array, &error) == 2); require_no_error(error); validity = {}; @@ -174,10 +174,6 @@ TEST_CASE("Array with validity", "[array]") { } TEST_CASE("Struct array creation", "[array]") { - vx_session *session = vx_session_new(); - defer { - vx_session_free(session); - }; vx_error *error = nullptr; vx_validity validity = {}; @@ -209,7 +205,7 @@ TEST_CASE("Struct array creation", "[array]") { REQUIRE(to_string_view(vx_struct_fields_field_name(fields, 0)) == "age"); vx_struct_fields_free(fields); vx_dtype_free(dtype); - const vx_array *age = vx_array_get_field(session, struct_array, 0, &error); + const vx_array *age = vx_array_get_field(struct_array, 0, &error); require_no_error(error); REQUIRE(vx_array_has_dtype(age, DTYPE_NULL)); REQUIRE(vx_array_len(age) == 5); diff --git a/vortex-ffi/test/scan.cpp b/vortex-ffi/test/scan.cpp index 427672e62ab..6e0756d603e 100644 --- a/vortex-ffi/test/scan.cpp +++ b/vortex-ffi/test/scan.cpp @@ -357,17 +357,17 @@ void verify_sample_array(vx_session *session, const vx_array *array) { require_no_error(error); REQUIRE(validity.type == VX_VALIDITY_NON_NULLABLE); - const vx_array *age_field = vx_array_get_field(session, array, 0, &error); + const vx_array *age_field = vx_array_get_field(array, 0, &error); require_no_error(error); verify_age_field(session, age_field); vx_array_free(age_field); - const vx_array *height_field = vx_array_get_field(session, array, 1, &error); + const vx_array *height_field = vx_array_get_field(array, 1, &error); require_no_error(error); verify_height_field(session, height_field); vx_array_free(height_field); - REQUIRE(vx_array_get_field(session, array, 2, &error) == nullptr); + REQUIRE(vx_array_get_field(array, 2, &error) == nullptr); REQUIRE(error != nullptr); vx_error_free(error); } @@ -747,7 +747,7 @@ TEST_CASE("Filter with literal expression", "[filter]") { REQUIRE(vx_array_len(array) == SAMPLE_ROWS - threshold); vx_error *error = nullptr; - const vx_array *filtered_age = vx_array_get_field(session, array, 0, &error); + const vx_array *filtered_age = vx_array_get_field(array, 0, &error); require_no_error(error); REQUIRE(filtered_age != nullptr); defer { diff --git a/vortex/benches/common_encoding_tree_throughput.rs b/vortex/benches/common_encoding_tree_throughput.rs index 0077e0006c0..8dcc482bf8f 100644 --- a/vortex/benches/common_encoding_tree_throughput.rs +++ b/vortex/benches/common_encoding_tree_throughput.rs @@ -326,6 +326,7 @@ mod setup { ), compressed_codes, fsst.uncompressed_lengths().clone(), + &mut ctx, ) .unwrap(); From b4736716bfe28ddb7b19462ba4926c2834bf0351 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 19 Oct 2026 15:21:48 +0000 Subject: [PATCH 3/3] Thread execution contexts through Patches `Patches::new` validated, and `slice`, `chunk_offset_at`, `min_index` and `max_index` read, the patch indices with the hidden global `legacy_session()`. Take an `ExecutionCtx` instead, as do Patched `lane_range` and `Sparse::try_new`. Sparse, ALP, ALP-RD and BitPacked deserialization validate their patches with a context from the session. The BitPacked and ALP-RD encoders and the Sparse patch reshaping build valid patches by construction and use `Patches::new_unchecked`. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Y1zHpE1K1Y9HvyU4kt24Tq --- encodings/alp/src/alp/array.rs | 19 ++- encodings/alp/src/alp/compress.rs | 1 + encodings/alp/src/alp/compute/slice.rs | 4 +- encodings/alp/src/alp_rd/array.rs | 5 +- encodings/alp/src/alp_rd/mod.rs | 22 ++- encodings/alp/src/alp_rd/slice.rs | 4 +- .../array/bitpack_compress/blocked.rs | 19 ++- .../array/bitpack_compress/global.rs | 19 ++- .../array/bitpack_decompress/blocked.rs | 1 + .../fastlanes/src/bitpacking/compute/slice.rs | 4 +- .../fastlanes/src/bitpacking/plugin/mod.rs | 6 +- .../fastlanes/src/bitpacking/plugin/v1.rs | 16 +- .../fastlanes/src/bitpacking/plugin/v2.rs | 16 +- .../src/bitpacking/vtable/operations.rs | 3 + encodings/sparse/benches/sparse_canonical.rs | 26 ++- encodings/sparse/benches/sparse_pushdown.rs | 15 +- encodings/sparse/src/canonical.rs | 80 ++++++--- encodings/sparse/src/compute/between.rs | 4 +- encodings/sparse/src/compute/cast.rs | 17 +- encodings/sparse/src/compute/compare.rs | 3 +- encodings/sparse/src/compute/fill_null.rs | 2 + encodings/sparse/src/compute/filter.rs | 10 ++ encodings/sparse/src/compute/is_constant.rs | 5 + encodings/sparse/src/compute/like.rs | 6 +- encodings/sparse/src/compute/min_max.rs | 8 +- encodings/sparse/src/compute/mod.rs | 46 ++++-- encodings/sparse/src/compute/nan_count.rs | 6 +- encodings/sparse/src/compute/null_count.rs | 4 +- encodings/sparse/src/compute/sum.rs | 9 + encodings/sparse/src/compute/take.rs | 10 ++ encodings/sparse/src/lib.rs | 90 ++++++---- encodings/sparse/src/ops.rs | 2 +- encodings/sparse/src/slice.rs | 4 +- vortex-array/benches/patches_lookup.rs | 2 + vortex-array/benches/take_patches.rs | 5 +- vortex-array/src/arrays/bool/array.rs | 2 + vortex-array/src/arrays/patched/array.rs | 24 +-- .../src/arrays/patched/compute/compare.rs | 17 +- .../src/arrays/patched/compute/filter.rs | 14 +- .../src/arrays/patched/compute/take.rs | 2 + vortex-array/src/arrays/patched/vtable/mod.rs | 13 +- .../src/arrays/patched/vtable/operations.rs | 8 +- .../src/arrays/patched/vtable/slice.rs | 14 +- .../src/compute/conformance/search_sorted.rs | 24 ++- vortex-array/src/patches.rs | 154 +++++++++++------- vortex-array/src/validity.rs | 1 + vortex-btrblocks/src/schemes/float/sparse.rs | 1 + .../src/schemes/integer/sparse.rs | 1 + vortex-btrblocks/src/schemes/patches.rs | 1 + vortex-btrblocks/src/schemes/string/sparse.rs | 1 + vortex-cuda/src/hybrid_dispatch/mod.rs | 1 + vortex-cuda/src/kernel/encodings/alp.rs | 2 + vortex-cuda/src/kernel/patches/types.rs | 10 +- .../arrays/synthetic/encodings/patched.rs | 1 + vortex/benches/take_union.rs | 3 +- 55 files changed, 546 insertions(+), 241 deletions(-) diff --git a/encodings/alp/src/alp/array.rs b/encodings/alp/src/alp/array.rs index 757da49f998..3f73318ea10 100644 --- a/encodings/alp/src/alp/array.rs +++ b/encodings/alp/src/alp/array.rs @@ -21,6 +21,7 @@ use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; use vortex_array::IntoArray; use vortex_array::TypedArrayRef; +use vortex_array::VortexSessionExecute; use vortex_array::array_slots; use vortex_array::arrays::Primitive; use vortex_array::buffer::BufferHandle; @@ -137,7 +138,7 @@ impl VTable for ALP { metadata: &[u8], _buffers: &[BufferHandle], children: &dyn ArrayChildren, - _session: &VortexSession, + session: &VortexSession, ) -> VortexResult> { let metadata = ALPMetadata::decode(metadata)?; let encoded_ptype = match &dtype { @@ -157,7 +158,14 @@ impl VTable for ALP { .map(|dtype| children.get(3, &dtype, usize::try_from(p.chunk_offsets_len())?)) .transpose()?; - Patches::new(len, p.offset()?, indices, values, chunk_offsets) + Patches::new( + len, + p.offset()?, + indices, + values, + chunk_offsets, + &mut session.create_execution_ctx(), + ) }) .transpose()?; @@ -693,6 +701,7 @@ mod tests { /// be preserved, or when building ALPArrays manually without chunk_offsets. #[test] fn test_execute_decompress_with_patches_no_chunk_offsets_regression_5948() { + let mut ctx = SESSION.create_execution_ctx(); // Create an array with values that will produce patches. PI doesn't encode cleanly. let values: Vec = vec![1.0, 2.0, PI, 4.0, 5.0]; let original = PrimitiveArray::from_iter(values); @@ -722,6 +731,7 @@ mod tests { original_patches.indices().clone(), original_patches.values().clone(), None, // NO chunk_offsets - this triggers the bug! + &mut ctx, ) .unwrap(); @@ -748,10 +758,7 @@ mod tests { ); // The execute path has the bug - it drops patches when chunk_offsets is None. - let result_execute = { - let mut ctx = SESSION.create_execution_ctx(); - execute_decompress(alp_without_chunk_offsets, &mut ctx).unwrap() - }; + let result_execute = { execute_decompress(alp_without_chunk_offsets, &mut ctx).unwrap() }; let execute_slice = result_execute.as_slice::(); // This assertion FAILS until the bug is fixed because execute_decompress drops patches. diff --git a/encodings/alp/src/alp/compress.rs b/encodings/alp/src/alp/compress.rs index acfefc1c7fe..6dd122080cf 100644 --- a/encodings/alp/src/alp/compress.rs +++ b/encodings/alp/src/alp/compress.rs @@ -143,6 +143,7 @@ where Buffer::from(valid_exceptional_positions).into_array(), valid_exceptional_values, Some(Buffer::from(chunk_offsets).into_array()), + ctx, )?) }; Ok((exponents, encoded_array, patches)) diff --git a/encodings/alp/src/alp/compute/slice.rs b/encodings/alp/src/alp/compute/slice.rs index cb8c122fdc6..014d17e40bd 100644 --- a/encodings/alp/src/alp/compute/slice.rs +++ b/encodings/alp/src/alp/compute/slice.rs @@ -18,14 +18,14 @@ impl SliceKernel for ALP { fn slice( array: ArrayView<'_, Self>, range: Range, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult> { let sliced_alp = ALP::new( array.encoded().slice(range.clone())?, array.exponents(), array .patches() - .map(|p| p.slice(range)) + .map(|p| p.slice(range, ctx)) .transpose()? .flatten(), ) diff --git a/encodings/alp/src/alp_rd/array.rs b/encodings/alp/src/alp_rd/array.rs index fc1aeb5fb1b..ca7421e788f 100644 --- a/encodings/alp/src/alp_rd/array.rs +++ b/encodings/alp/src/alp_rd/array.rs @@ -22,6 +22,7 @@ use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; use vortex_array::IntoArray; use vortex_array::TypedArrayRef; +use vortex_array::VortexSessionExecute; use vortex_array::array_slots; use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; @@ -161,7 +162,6 @@ impl VTable for ALPRD { )) } - #[allow(clippy::disallowed_methods)] fn deserialize( &self, dtype: &DType, @@ -169,7 +169,7 @@ impl VTable for ALPRD { metadata: &[u8], _buffers: &[BufferHandle], children: &dyn ArrayChildren, - _session: &VortexSession, + session: &VortexSession, ) -> VortexResult> { let metadata = ALPRDMetadata::decode(metadata)?; if children.len() < 2 { @@ -214,6 +214,7 @@ impl VTable for ALPRD { values, // TODO(0ax1): handle chunk offsets None, + &mut session.create_execution_ctx(), ) }) .transpose()?; diff --git a/encodings/alp/src/alp_rd/mod.rs b/encodings/alp/src/alp_rd/mod.rs index 2cdeacd207d..8b0cd8f2f39 100644 --- a/encodings/alp/src/alp_rd/mod.rs +++ b/encodings/alp/src/alp_rd/mod.rs @@ -104,15 +104,19 @@ where .into_array() }; - Patches::new( - len, - 0, - packed_pos, - Buffer::from(exceptions).into_array(), - // TODO(0ax1): handle chunk offsets - None, - ) - .vortex_expect("Patches construction in encode") + // SAFETY: the exception positions are collected in increasing order from positions in + // the input, with one exception value each. + unsafe { + Patches::new_unchecked( + len, + 0, + packed_pos, + Buffer::from(exceptions).into_array(), + // TODO(0ax1): handle chunk offsets + None, + None, + ) + } }); ALPRD::try_new( diff --git a/encodings/alp/src/alp_rd/slice.rs b/encodings/alp/src/alp_rd/slice.rs index 5f81bdc7b5d..613ff0873ff 100644 --- a/encodings/alp/src/alp_rd/slice.rs +++ b/encodings/alp/src/alp_rd/slice.rs @@ -18,11 +18,11 @@ impl SliceKernel for ALPRD { fn slice( array: ArrayView<'_, Self>, range: Range, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult> { let left_parts_exceptions = array .left_parts_patches() - .map(|patches| patches.slice(range.clone())) + .map(|patches| patches.slice(range.clone(), ctx)) .transpose()? .flatten(); diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_compress/blocked.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_compress/blocked.rs index 9a2ba8803c6..cc0105317a2 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_compress/blocked.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_compress/blocked.rs @@ -254,13 +254,18 @@ where if indices.is_empty() { return Ok(None); } - Ok(Some(Patches::new( - values.len(), - 0, - indices.into_array(), - PrimitiveArray::new(patch_values, patch_validity).into_array(), - Some(chunk_offsets.into_array()), - )?)) + // SAFETY: the indices are collected in increasing order from positions in `values`, with one + // patch value each. + Ok(Some(unsafe { + Patches::new_unchecked( + values.len(), + 0, + indices.into_array(), + PrimitiveArray::new(patch_values, patch_validity).into_array(), + Some(chunk_offsets.into_array()), + Some(0), + ) + })) } /// Byte boundaries of blocks packed at `widths`, in the narrowest unsigned type that holds them. diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_compress/global.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_compress/global.rs index 70cec66325b..7e61427c612 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_compress/global.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_compress/global.rs @@ -279,13 +279,18 @@ where if indices.is_empty() { Ok(None) } else { - Ok(Some(Patches::new( - data.len(), - 0, - indices.into_array(), - PrimitiveArray::new(values, patch_validity).into_array(), - Some(chunk_offsets.into_array()), - )?)) + // SAFETY: the indices are collected in increasing order from positions in `data`, with + // one patch value each. + Ok(Some(unsafe { + Patches::new_unchecked( + data.len(), + 0, + indices.into_array(), + PrimitiveArray::new(values, patch_validity).into_array(), + Some(chunk_offsets.into_array()), + Some(0), + ) + })) } } diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress/blocked.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress/blocked.rs index cff33674731..6dcb3177928 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress/blocked.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress/blocked.rs @@ -454,6 +454,7 @@ mod tests { PrimitiveArray::from_iter(indices.map(|i| i + 113)).into_array(), PrimitiveArray::new(buffer![999u32; 6], Validity::AllValid).into_array(), None, + &mut ctx, )?; let BitWidthsView::Blocked(block_offsets) = array.bit_widths() else { vortex_bail!("expected block offsets"); diff --git a/encodings/fastlanes/src/bitpacking/compute/slice.rs b/encodings/fastlanes/src/bitpacking/compute/slice.rs index c7855e4d412..93cb829da97 100644 --- a/encodings/fastlanes/src/bitpacking/compute/slice.rs +++ b/encodings/fastlanes/src/bitpacking/compute/slice.rs @@ -35,14 +35,14 @@ impl SliceKernel for BitPacked { fn slice( array: ArrayView<'_, Self>, range: Range, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult> { let BitWidthsView::Global(bit_width) = array.bit_widths() else { return Ok(None); }; let patches = array .patches() - .map(|p| p.slice(range.clone())) + .map(|p| p.slice(range.clone(), ctx)) .transpose()? .flatten(); diff --git a/encodings/fastlanes/src/bitpacking/plugin/mod.rs b/encodings/fastlanes/src/bitpacking/plugin/mod.rs index 56f7ca89697..31b1f06ec90 100644 --- a/encodings/fastlanes/src/bitpacking/plugin/mod.rs +++ b/encodings/fastlanes/src/bitpacking/plugin/mod.rs @@ -86,12 +86,12 @@ impl ArrayPlugin for BitPackedPlugin { fn deserialize( &self, parts: ArrayDeserialization<'_>, - _session: &VortexSession, + session: &VortexSession, ) -> VortexResult { if parts.serialized_id == bitpacked_v1_id() { - v1::deserialize(parts) + v1::deserialize(parts, session) } else if parts.serialized_id == bitpacked_v2_id() { - v2::deserialize(parts) + v2::deserialize(parts, session) } else { vortex_bail!( "BitPacked plugin does not recognize serialized ID {}", diff --git a/encodings/fastlanes/src/bitpacking/plugin/v1.rs b/encodings/fastlanes/src/bitpacking/plugin/v1.rs index d1a19d43f48..78da347f2c5 100644 --- a/encodings/fastlanes/src/bitpacking/plugin/v1.rs +++ b/encodings/fastlanes/src/bitpacking/plugin/v1.rs @@ -12,6 +12,7 @@ use vortex_array::ArraySerialization; use vortex_array::ArraySlots; use vortex_array::ArrayView; use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; use vortex_array::patches::Patches; use vortex_array::patches::PatchesData; use vortex_array::patches::PatchesMetadata; @@ -20,6 +21,7 @@ use vortex_array::vtable::validity_to_child; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; +use vortex_session::VortexSession; use super::bitpacked_v1_id; use crate::BitPacked; @@ -58,7 +60,10 @@ pub(super) fn serialize( )) } -pub(super) fn deserialize(parts: ArrayDeserialization<'_>) -> VortexResult { +pub(super) fn deserialize( + parts: ArrayDeserialization<'_>, + session: &VortexSession, +) -> VortexResult { let ArrayDeserialization { dtype, len, @@ -108,7 +113,14 @@ pub(super) fn deserialize(parts: ArrayDeserialization<'_>) -> VortexResult) -> VortexResult { +pub(super) fn deserialize( + parts: ArrayDeserialization<'_>, + session: &VortexSession, +) -> VortexResult { let ArrayDeserialization { dtype, len, @@ -127,7 +132,14 @@ pub(super) fn deserialize(parts: ArrayDeserialization<'_>) -> VortexResult ArrayRef { + let mut ctx = SESSION.create_execution_ctx(); let patch_indices: Buffer = (0..len).step_by(patch_stride).map(|i| i as u32).collect(); let n_patches = patch_indices.len(); @@ -70,12 +71,19 @@ fn make_sparse_list(len: usize, patch_stride: usize, list_size: usize) -> ArrayR NonNullable, ); - Sparse::try_new(patch_indices.into_array(), patch_values, len, fill_value) - .vortex_expect("sparse list input should be valid") - .into_array() + Sparse::try_new( + patch_indices.into_array(), + patch_values, + len, + fill_value, + &mut ctx, + ) + .vortex_expect("sparse list input should be valid") + .into_array() } fn make_sparse_fixed_size_list(len: usize, patch_stride: usize, list_size: u32) -> ArrayRef { + let mut ctx = SESSION.create_execution_ctx(); let patch_indices: Buffer = (0..len).step_by(patch_stride).map(|i| i as u32).collect(); let n_patches = patch_indices.len(); @@ -91,9 +99,15 @@ fn make_sparse_fixed_size_list(len: usize, patch_stride: usize, list_size: u32) NonNullable, ); - Sparse::try_new(patch_indices.into_array(), patch_values, len, fill_value) - .vortex_expect("sparse fixed-size-list input should be valid") - .into_array() + Sparse::try_new( + patch_indices.into_array(), + patch_values, + len, + fill_value, + &mut ctx, + ) + .vortex_expect("sparse fixed-size-list input should be valid") + .into_array() } #[divan::bench(args = LIST_ARGS)] diff --git a/encodings/sparse/benches/sparse_pushdown.rs b/encodings/sparse/benches/sparse_pushdown.rs index dea4f94a832..0a28d634898 100644 --- a/encodings/sparse/benches/sparse_pushdown.rs +++ b/encodings/sparse/benches/sparse_pushdown.rs @@ -58,6 +58,7 @@ static SESSION: LazyLock = LazyLock::new(|| { /// fill value 1. When `constant` is true every patch also equals 1, so the whole array /// is constant (the worst case for `is_constant`: it must scan all patches to confirm). fn make_sparse(num_patches: usize, constant: bool) -> ArrayRef { + let mut ctx = SESSION.create_execution_ctx(); let stride = LEN / num_patches; let indices: Buffer = (0..num_patches).map(|i| (i * stride) as u32).collect(); let values: Buffer = (0..num_patches) @@ -68,6 +69,7 @@ fn make_sparse(num_patches: usize, constant: bool) -> ArrayRef { values.into_array(), LEN, Scalar::from(1i32), + &mut ctx, ) .vortex_expect("valid sparse") .into_array() @@ -76,6 +78,7 @@ fn make_sparse(num_patches: usize, constant: bool) -> ArrayRef { /// Build a sparse `i32` array of `LEN` with a null fill and `num_patches` nullable patches /// (every third patch null), so `null_count` does real `O(P)` work over the patch validity. fn make_sparse_nullable(num_patches: usize) -> ArrayRef { + let mut ctx = SESSION.create_execution_ctx(); let stride = LEN / num_patches; let indices: Buffer = (0..num_patches).map(|i| (i * stride) as u32).collect(); let values = PrimitiveArray::from_option_iter( @@ -83,9 +86,15 @@ fn make_sparse_nullable(num_patches: usize) -> ArrayRef { ) .into_array(); let nullable = DType::Primitive(PType::I32, Nullability::Nullable); - Sparse::try_new(indices.into_array(), values, LEN, Scalar::null(nullable)) - .vortex_expect("valid sparse") - .into_array() + Sparse::try_new( + indices.into_array(), + values, + LEN, + Scalar::null(nullable), + &mut ctx, + ) + .vortex_expect("valid sparse") + .into_array() } #[divan::bench] diff --git a/encodings/sparse/src/canonical.rs b/encodings/sparse/src/canonical.rs index afededfeab0..4fca16be926 100644 --- a/encodings/sparse/src/canonical.rs +++ b/encodings/sparse/src/canonical.rs @@ -927,7 +927,8 @@ mod test { let mut ctx = SESSION.create_execution_ctx(); let indices = buffer![0u64, 1, 7].into_array(); let values = BoolArray::from_iter([Some(true), None, Some(false)]).into_array(); - let sparse_bools = Sparse::try_new(indices, values, 10, Scalar::from(fill_value)).unwrap(); + let sparse_bools = + Sparse::try_new(indices, values, 10, Scalar::from(fill_value), &mut ctx).unwrap(); let expected = BoolArray::from_iter([ Some(true), @@ -953,7 +954,8 @@ mod test { let mut ctx = SESSION.create_execution_ctx(); let indices = buffer![0u64, 1, 7].into_array(); let values = PrimitiveArray::from_option_iter([Some(0i32), None, Some(1)]).into_array(); - let sparse_ints = Sparse::try_new(indices, values, 10, Scalar::from(fill_value)).unwrap(); + let sparse_ints = + Sparse::try_new(indices, values, 10, Scalar::from(fill_value), &mut ctx).unwrap(); assert_eq!(*sparse_ints.dtype(), DType::Primitive(PType::I32, Nullable)); let expected = PrimitiveArray::from_option_iter([ @@ -1004,7 +1006,8 @@ mod test { vec![Scalar::from(Some(-10i32)), Scalar::from(Some(-1i32))], ); let len = 10; - let sparse_struct = Sparse::try_new(indices, patch_values, len, fill_scalar).unwrap(); + let sparse_struct = + Sparse::try_new(indices, patch_values, len, fill_scalar, &mut ctx).unwrap(); let expected_a = PrimitiveArray::from_option_iter((0..len).map(|i| { if i == 0 { @@ -1071,7 +1074,8 @@ mod test { let fill_scalar = Scalar::null(struct_dtype); let len = 10; - let sparse_struct = Sparse::try_new(indices, patch_values, len, fill_scalar).unwrap(); + let sparse_struct = + Sparse::try_new(indices, patch_values, len, fill_scalar, &mut ctx).unwrap(); let expected_a = PrimitiveArray::from_option_iter((0..len).map(|i| { if i == 0 { @@ -1122,7 +1126,8 @@ mod test { .into_array(); let len = 10; let fill_scalar = Scalar::decimal(DecimalValue::I32(123), decimal_dtype, Nullable); - let sparse_struct = Sparse::try_new(indices, patch_values, len, fill_scalar).unwrap(); + let sparse_struct = + Sparse::try_new(indices, patch_values, len, fill_scalar, &mut ctx).unwrap(); let expected = SESSION .arrow() @@ -1176,6 +1181,7 @@ mod test { strings, 12, Scalar::from(Some("123".to_owned())), + &mut ctx, ) .unwrap(); @@ -1217,6 +1223,7 @@ mod test { strings, 12, Scalar::null(DType::Utf8(Nullable)), + &mut ctx, ) .unwrap(); @@ -1251,6 +1258,7 @@ mod test { strings, 9, Scalar::from("123".to_owned()), + &mut ctx, ) .unwrap(); @@ -1281,6 +1289,7 @@ mod test { strings, 12, Scalar::null(DType::Utf8(Nullable)), + &mut ctx, ) .unwrap(); @@ -1322,6 +1331,7 @@ mod test { binaries, 12, Scalar::from(Some(ByteBuffer::from(b"123".to_vec()))), + &mut ctx, ) .unwrap(); @@ -1364,7 +1374,7 @@ mod test { let indices = buffer![0u8, 3u8, 4u8, 5u8].into_array(); let fill_value = Scalar::null(lists.dtype().clone()); - let sparse = Sparse::try_new(indices, lists, 6, fill_value)?.into_array(); + let sparse = Sparse::try_new(indices, lists, 6, fill_value, &mut ctx)?.into_array(); let expected = ListArray::from_iter_opt_slow::( [ @@ -1400,7 +1410,7 @@ mod test { let indices = buffer![0u8, 3u8, 4u8, 5u8].into_array(); let fill_value = Scalar::null(lists.dtype().clone()); - let sparse = Sparse::try_new(indices, lists, 6, fill_value) + let sparse = Sparse::try_new(indices, lists, 6, fill_value, &mut ctx) .unwrap() .into_array(); @@ -1434,7 +1444,7 @@ mod test { let indices = buffer![0u8, 3u8, 4u8, 5u8].into_array(); let fill_value = Scalar::from(Some(vec![5i32, 6, 7, 8])); - let sparse = Sparse::try_new(indices, lists, 6, fill_value)?.into_array(); + let sparse = Sparse::try_new(indices, lists, 6, fill_value, &mut ctx)?.into_array(); let actual = sparse.execute::(&mut ctx)?.into_array(); let result_listview = actual.execute::(&mut ctx)?; @@ -1512,8 +1522,14 @@ mod test { // Patches at 10 and 20 of 10,000 rows, so the fill covers three gaps. let indices = buffer![10u32, 20].into_array(); let fill = vec![7i32, 8, 9]; - let sparse = - Sparse::try_new(indices, lists, 10_000, Scalar::from(Some(fill.clone())))?.into_array(); + let sparse = Sparse::try_new( + indices, + lists, + 10_000, + Scalar::from(Some(fill.clone())), + &mut ctx, + )? + .into_array(); let actual = sparse.execute::(&mut ctx)?; assert_eq!(actual.len(), 10_000); @@ -1569,7 +1585,8 @@ mod test { let indices = PrimitiveArray::from_iter(0..patches_u32).into_array(); let fill = Scalar::from(Some(vec![-1i32])); - let sparse = Sparse::try_new(indices, patch_values, 2 * PATCHES, fill)?.into_array(); + let sparse = + Sparse::try_new(indices, patch_values, 2 * PATCHES, fill, &mut ctx)?.into_array(); let actual = sparse.execute::(&mut ctx)?; let chunks = actual @@ -1618,7 +1635,8 @@ mod test { // Rows 0..3 are one run of patches, row 5 is another, and rows 3-4 and 6 are gaps. let indices = buffer![0u8, 1, 2, 5].into_array(); let fill = Scalar::from(Some(vec![7i32, 8])); - let sparse = Sparse::try_new(indices, patches.into_array(), 7, fill)?.into_array(); + let sparse = + Sparse::try_new(indices, patches.into_array(), 7, fill, &mut ctx)?.into_array(); let actual = sparse.execute::(&mut ctx)?; @@ -1656,6 +1674,7 @@ mod test { strings, 12, Scalar::null(DType::Binary(Nullable)), + &mut ctx, ) .unwrap(); @@ -1680,6 +1699,7 @@ mod test { #[test] fn test_sparse_fixed_size_list_null_fill() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); // Create a FixedSizeListArray with 3 lists of size 3. let elements = buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9].into_array(); let fsl = FixedSizeListArray::try_new(elements, 3, Validity::AllValid, 3)?.into_array(); @@ -1690,7 +1710,7 @@ mod test { 3, Nullable, )); - let sparse = Sparse::try_new(indices, fsl, 5, fill_value)?.into_array(); + let sparse = Sparse::try_new(indices, fsl, 5, fill_value, &mut ctx)?.into_array(); // Expected: [1,2,3], null, [4,5,6], [7,8,9], null. let expected_elements = @@ -1709,6 +1729,7 @@ mod test { #[test] fn test_sparse_fixed_size_list_non_null_fill() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let elements = buffer![1i32, 2, 3, 4, 5, 6].into_array(); // Non-nullable values to match the non-nullable fill below. let fsl = FixedSizeListArray::try_new(elements, 2, Validity::NonNullable, 3)?.into_array(); @@ -1722,7 +1743,7 @@ mod test { ], NonNullable, ); - let sparse = Sparse::try_new(indices, fsl, 6, fill_value)?.into_array(); + let sparse = Sparse::try_new(indices, fsl, 6, fill_value, &mut ctx)?.into_array(); // Expected: [1,2], [99,88], [3,4], [99,88], [5,6], [99,88]. let expected_elements = buffer![1i32, 2, 99, 88, 3, 4, 99, 88, 5, 6, 99, 88].into_array(); @@ -1735,6 +1756,7 @@ mod test { #[test] fn test_sparse_fixed_size_list_with_validity() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); // Create FSL values with some nulls. let elements = buffer![10i32, 20, 30, 40, 50, 60].into_array(); let fsl = FixedSizeListArray::try_new( @@ -1754,7 +1776,7 @@ mod test { ], Nullable, ); - let sparse = Sparse::try_new(indices, fsl, 6, fill_value)?.into_array(); + let sparse = Sparse::try_new(indices, fsl, 6, fill_value, &mut ctx)?.into_array(); // Expected validity: [true, true, true, false, true, true]. // Expected elements: [7,8], [10,20], [7,8], [30,40], [50,60], [7,8]. @@ -1775,6 +1797,7 @@ mod test { #[test] fn test_sparse_fixed_size_list_truly_sparse() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); // Test with a truly sparse array where most values are the fill value. // This demonstrates the compression benefit of sparse encoding. @@ -1796,7 +1819,7 @@ mod test { NonNullable, ); - let sparse = Sparse::try_new(indices, fsl, 100, fill_value)?.into_array(); + let sparse = Sparse::try_new(indices, fsl, 100, fill_value, &mut ctx)?.into_array(); // Build expected: 97 copies of [99,99] with patches at positions 5, 50, 95. let mut expected_elements_vec = Vec::with_capacity(200); @@ -1833,6 +1856,7 @@ mod test { #[test] fn test_sparse_fixed_size_list_single_element() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); // Test with a single element FSL array. let elements = buffer![42i32, 43].into_array(); // Non-nullable values to match the non-nullable fill below. @@ -1847,7 +1871,7 @@ mod test { ], NonNullable, ); - let sparse = Sparse::try_new(indices, fsl, 1, fill_value)?.into_array(); + let sparse = Sparse::try_new(indices, fsl, 1, fill_value, &mut ctx)?.into_array(); // Expected: just [42, 43]. let expected_elements = buffer![42i32, 43].into_array(); @@ -1867,7 +1891,7 @@ mod test { let indices = buffer![0u8, 1u8, 2u8, 3u8].into_array(); let fill_value = Scalar::from(Some(vec![42i32; 252])); // 252 + 4 elements = 256 > u8::MAX - let sparse = Sparse::try_new(indices, lists, 5, fill_value)?.into_array(); + let sparse = Sparse::try_new(indices, lists, 5, fill_value, &mut ctx)?.into_array(); let actual = sparse.execute::(&mut ctx)?.into_array(); let mut expected_elements = buffer_mut![1, 2, 1, 2]; @@ -1939,6 +1963,7 @@ mod test { list_view.into_array(), 10, Scalar::null(list_dtype), + &mut ctx, )?; let expected = ListArray::from_iter_opt_slow::( @@ -2003,7 +2028,13 @@ mod test { let indices = buffer![0u8, 1].into_array(); // Extract only the values we need from the sliced array let values = sliced.slice(0..2)?; - let sparse = Sparse::try_new(indices, values, 5, Scalar::null(sliced.dtype().clone()))?; + let sparse = Sparse::try_new( + indices, + values, + 5, + Scalar::null(sliced.dtype().clone()), + &mut ctx, + )?; let expected = ListArray::from_iter_opt_slow::( [Some(vec![2i32, 3, 4]), Some(vec![5]), None, None, None], @@ -2031,8 +2062,14 @@ mod test { Some(value) => Scalar::from(value.to_owned()).into_nullable(), None => Scalar::null(DType::Utf8(Nullable)), }; - let array = Sparse::try_new(buffer![1u16, 4, 8].into_array(), strings, 10, fill_scalar)? - .into_array(); + let array = Sparse::try_new( + buffer![1u16, 4, 8].into_array(), + strings, + 10, + fill_scalar, + &mut ctx, + )? + .into_array(); let expected_values = [ fill, @@ -2082,6 +2119,7 @@ mod test { strings.into_array(), 6, Scalar::from("123".to_owned()), + &mut ctx, )?; let actual = array diff --git a/encodings/sparse/src/compute/between.rs b/encodings/sparse/src/compute/between.rs index 382d6c65405..35f5230d5ca 100644 --- a/encodings/sparse/src/compute/between.rs +++ b/encodings/sparse/src/compute/between.rs @@ -96,11 +96,13 @@ mod tests { #[case] lower_strict: StrictComparison, #[case] upper_strict: StrictComparison, ) { + let mut ctx = SESSION.create_execution_ctx(); let array = Sparse::try_new( buffer![1u64, 3, 5].into_array(), buffer![10i32, 20, 30].into_array(), 8, Scalar::from(1i32), + &mut ctx, ) .unwrap() .into_array(); @@ -113,8 +115,6 @@ mod tests { let lower = ConstantArray::new(Scalar::from(lo), len).into_array(); let upper = ConstantArray::new(Scalar::from(hi), len).into_array(); - let mut ctx = SESSION.create_execution_ctx(); - // Kernel path: between pushes through the Sparse encoding. let kernel = array .clone() diff --git a/encodings/sparse/src/compute/cast.rs b/encodings/sparse/src/compute/cast.rs index f07fe514a50..d59dcf8b143 100644 --- a/encodings/sparse/src/compute/cast.rs +++ b/encodings/sparse/src/compute/cast.rs @@ -62,11 +62,13 @@ mod tests { #[test] fn test_cast_sparse_with_null_fill() { + let mut ctx = SESSION.create_execution_ctx(); let sparse = Sparse::try_new( buffer![1u64, 3, 5].into_array(), PrimitiveArray::from_option_iter([Some(42i32), Some(84), Some(126)]).into_array(), 8, Scalar::null_native::(), + &mut ctx, ) .unwrap(); @@ -85,25 +87,29 @@ mod tests { buffer![2u64, 5, 8].into_array(), buffer![100i32, 200, 300].into_array(), 10, - Scalar::from(0i32) + Scalar::from(0i32), + &mut SESSION.create_execution_ctx(), ).unwrap())] #[case(Sparse::try_new( buffer![0u64, 4, 9].into_array(), buffer![1.5f32, 2.5, 3.5].into_array(), 10, - Scalar::from(0.0f32) + Scalar::from(0.0f32), + &mut SESSION.create_execution_ctx(), ).unwrap())] #[case(Sparse::try_new( buffer![1u64, 3, 7].into_array(), PrimitiveArray::from_option_iter([Some(100i32), None, Some(300)]).into_array(), 10, - Scalar::null_native::() + Scalar::null_native::(), + &mut SESSION.create_execution_ctx(), ).unwrap())] #[case(Sparse::try_new( buffer![5u64].into_array(), buffer![42u8].into_array(), 10, - Scalar::from(0u8) + Scalar::from(0u8), + &mut SESSION.create_execution_ctx(), ).unwrap())] fn test_cast_sparse_conformance(#[case] array: SparseArray) { test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); @@ -124,6 +130,7 @@ mod tests { .into_array(), 5, Scalar::null_native::(), + &mut ctx, )?; let casted = sparse @@ -142,6 +149,7 @@ mod tests { #[test] fn test_fill_null_sparse_with_null_fill() -> vortex_error::VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); // Regression test for https://github.com/vortex-data/vortex/issues/6932 // fill_null on a sparse array with null fill triggers an internal cast to // non-nullable, which must not panic. @@ -151,6 +159,7 @@ mod tests { PrimitiveArray::from_option_iter([Some(10u64), Some(20)]).into_array(), 5, Scalar::null_native::(), + &mut ctx, )?; let filled = sparse.into_array().fill_null(Scalar::from(0u64))?; diff --git a/encodings/sparse/src/compute/compare.rs b/encodings/sparse/src/compute/compare.rs index 3a4d9c5d5ce..c3a83d164ee 100644 --- a/encodings/sparse/src/compute/compare.rs +++ b/encodings/sparse/src/compute/compare.rs @@ -85,16 +85,17 @@ mod tests { #[case::lte(Scalar::from(10i32), Operator::Lte)] #[case::neq(Scalar::from(1i32), Operator::NotEq)] fn compare_matches_canonical(#[case] rhs: Scalar, #[case] op: Operator) { + let mut ctx = SESSION.create_execution_ctx(); let array: SparseArray = Sparse::try_new( buffer![1u64, 3, 5].into_array(), buffer![10i32, 20, 30].into_array(), 8, Scalar::from(1i32), + &mut ctx, ) .unwrap(); let arr = array.into_array(); let len = arr.len(); - let mut ctx = SESSION.create_execution_ctx(); // Kernel path: compare pushes through the Sparse encoding. let kernel_bool = arr diff --git a/encodings/sparse/src/compute/fill_null.rs b/encodings/sparse/src/compute/fill_null.rs index 24c9ea4c10d..8af5645aa10 100644 --- a/encodings/sparse/src/compute/fill_null.rs +++ b/encodings/sparse/src/compute/fill_null.rs @@ -78,6 +78,7 @@ mod tests { PrimitiveArray::from_option_iter([Some(10i32), None, Some(30)]).into_array().cast(nullable_i32()).unwrap(), 8, Scalar::null(nullable_i32()), + &mut SESSION.create_execution_ctx(), ).unwrap().into_array())] // non-null fill, nullable patches with a null #[case(Sparse::try_new( @@ -85,6 +86,7 @@ mod tests { PrimitiveArray::from_option_iter([Some(7i32), None]).into_array().cast(nullable_i32()).unwrap(), 4, Scalar::from(1i32).cast(&nullable_i32()).unwrap(), + &mut SESSION.create_execution_ctx(), ).unwrap().into_array())] fn fill_null_matches_canonical(#[case] array: vortex_array::ArrayRef) { let mut ctx = SESSION.create_execution_ctx(); diff --git a/encodings/sparse/src/compute/filter.rs b/encodings/sparse/src/compute/filter.rs index c1818262b9b..7a924155ff1 100644 --- a/encodings/sparse/src/compute/filter.rs +++ b/encodings/sparse/src/compute/filter.rs @@ -64,11 +64,13 @@ mod tests { #[fixture] fn array() -> ArrayRef { + let mut ctx = SESSION.create_execution_ctx(); Sparse::try_new( buffer![2u64, 9, 15].into_array(), PrimitiveArray::new(buffer![33_i32, 44, 55], Validity::AllValid).into_array(), 20, Scalar::null_native::(), + &mut ctx, ) .unwrap() .into_array() @@ -76,6 +78,7 @@ mod tests { #[rstest] fn test_filter(array: ArrayRef) { + let mut ctx = SESSION.create_execution_ctx(); let mut predicate = vec![false, false, true]; predicate.extend_from_slice(&[false; 17]); let mask = Mask::from_iter(predicate); @@ -89,6 +92,7 @@ mod tests { PrimitiveArray::new(buffer![33_i32], Validity::AllValid).into_array(), 1, Scalar::null_native::(), + &mut ctx, ) .unwrap(); @@ -101,12 +105,14 @@ mod tests { #[test] fn true_fill_value() { + let mut ctx = SESSION.create_execution_ctx(); let mask = Mask::from_iter([false, true, false, true, false, true, true]); let array = Sparse::try_new( buffer![0_u64, 3, 6].into_array(), PrimitiveArray::new(buffer![33_i32, 44, 55], Validity::AllValid).into_array(), 7, Scalar::null_native::(), + &mut ctx, ) .unwrap() .into_array(); @@ -122,6 +128,7 @@ mod tests { PrimitiveArray::new(buffer![44_i32, 55], Validity::AllValid).into_array(), 4, Scalar::null_native::(), + &mut ctx, ) .unwrap(); @@ -134,6 +141,7 @@ mod tests { #[test] fn test_filter_sparse_array() { + let mut ctx = SESSION.create_execution_ctx(); let null_fill_value = Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)); test_filter_conformance( &Sparse::try_new( @@ -144,6 +152,7 @@ mod tests { .unwrap(), 5, null_fill_value, + &mut ctx, ) .unwrap() .into_array(), @@ -157,6 +166,7 @@ mod tests { buffer![100i32, 200, 300].into_array(), 5, ten_fill_value, + &mut ctx, ) .unwrap() .into_array(), diff --git a/encodings/sparse/src/compute/is_constant.rs b/encodings/sparse/src/compute/is_constant.rs index eee8683df0a..71e3130aea6 100644 --- a/encodings/sparse/src/compute/is_constant.rs +++ b/encodings/sparse/src/compute/is_constant.rs @@ -114,6 +114,7 @@ mod tests { buffer![7i32, 7, 7].into_array(), 10, Scalar::from(7i32), + &mut SESSION.create_execution_ctx(), ).unwrap(), true, )] @@ -123,6 +124,7 @@ mod tests { buffer![9i32, 9].into_array(), 5, Scalar::from(7i32), + &mut SESSION.create_execution_ctx(), ).unwrap(), false, )] @@ -132,6 +134,7 @@ mod tests { buffer![3i32].into_array(), 5, Scalar::from(7i32), + &mut SESSION.create_execution_ctx(), ).unwrap(), false, )] @@ -141,6 +144,7 @@ mod tests { buffer![5i32, 5, 5, 5].into_array(), 4, Scalar::from(99i32), // fill is unreachable + &mut SESSION.create_execution_ctx(), ).unwrap(), true, )] @@ -150,6 +154,7 @@ mod tests { buffer![1i32, 2, 3].into_array(), 3, Scalar::from(99i32), + &mut SESSION.create_execution_ctx(), ).unwrap(), false, )] diff --git a/encodings/sparse/src/compute/like.rs b/encodings/sparse/src/compute/like.rs index b46a4ce61e0..e5a2333b190 100644 --- a/encodings/sparse/src/compute/like.rs +++ b/encodings/sparse/src/compute/like.rs @@ -101,12 +101,12 @@ mod tests { #[case] fill: Scalar, #[case] options: LikeOptions, ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let values = VarBinViewArray::from_iter_str(["world", "help", "foo"]).into_array(); let values = values.cast(fill.dtype().clone())?; - let array = - Sparse::try_new(buffer![1u64, 3, 5].into_array(), values, 8, fill)?.into_array(); + let array = Sparse::try_new(buffer![1u64, 3, 5].into_array(), values, 8, fill, &mut ctx)? + .into_array(); let pattern = ConstantArray::new(pattern, array.len()).into_array(); - let mut ctx = SESSION.create_execution_ctx(); // Kernel path: like pushes through the Sparse encoding. let kernel = Like::try_new(array.clone(), pattern.clone(), options)? diff --git a/encodings/sparse/src/compute/min_max.rs b/encodings/sparse/src/compute/min_max.rs index 553b8eaaee0..3e5037d32bc 100644 --- a/encodings/sparse/src/compute/min_max.rs +++ b/encodings/sparse/src/compute/min_max.rs @@ -91,13 +91,13 @@ mod tests { #[rstest] // fill below all patches - #[case(Sparse::try_new(buffer![1u64, 3, 5].into_array(), buffer![10i32, 20, 30].into_array(), 8, Scalar::from(1i32)).unwrap())] + #[case(Sparse::try_new(buffer![1u64, 3, 5].into_array(), buffer![10i32, 20, 30].into_array(), 8, Scalar::from(1i32), &mut SESSION.create_execution_ctx()).unwrap())] // fill above all patches - #[case(Sparse::try_new(buffer![1u64, 3, 5].into_array(), buffer![10i32, 20, 30].into_array(), 8, Scalar::from(99i32)).unwrap())] + #[case(Sparse::try_new(buffer![1u64, 3, 5].into_array(), buffer![10i32, 20, 30].into_array(), 8, Scalar::from(99i32), &mut SESSION.create_execution_ctx()).unwrap())] // fill in the middle - #[case(Sparse::try_new(buffer![1u64, 3, 5].into_array(), buffer![10i32, 20, 30].into_array(), 8, Scalar::from(15i32)).unwrap())] + #[case(Sparse::try_new(buffer![1u64, 3, 5].into_array(), buffer![10i32, 20, 30].into_array(), 8, Scalar::from(15i32), &mut SESSION.create_execution_ctx()).unwrap())] // every position patched (fill unreachable) - #[case(Sparse::try_new(buffer![0u64, 1, 2].into_array(), buffer![7i32, 3, 9].into_array(), 3, Scalar::from(99i32)).unwrap())] + #[case(Sparse::try_new(buffer![0u64, 1, 2].into_array(), buffer![7i32, 3, 9].into_array(), 3, Scalar::from(99i32), &mut SESSION.create_execution_ctx()).unwrap())] fn min_max_matches_canonical(#[case] array: SparseArray) { let arr = array.into_array(); let kernel: Option = min_max( diff --git a/encodings/sparse/src/compute/mod.rs b/encodings/sparse/src/compute/mod.rs index 6d565b221aa..a47b1ac9431 100644 --- a/encodings/sparse/src/compute/mod.rs +++ b/encodings/sparse/src/compute/mod.rs @@ -49,11 +49,13 @@ mod tests { #[fixture] fn array() -> ArrayRef { + let mut ctx = SESSION.create_execution_ctx(); Sparse::try_new( buffer![2u64, 9, 15].into_array(), PrimitiveArray::new(buffer![33_i32, 44, 55], Validity::AllValid).into_array(), 20, Scalar::null_native::(), + &mut ctx, ) .unwrap() .into_array() @@ -66,6 +68,7 @@ mod tests { #[test] fn test_mask_sparse_array() { + let mut ctx = SESSION.create_execution_ctx(); let null_fill_value = Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)); test_mask_conformance( &Sparse::try_new( @@ -76,6 +79,7 @@ mod tests { .unwrap(), 5, null_fill_value, + &mut ctx, ) .unwrap() .into_array(), @@ -89,6 +93,7 @@ mod tests { buffer![100i32, 200, 300].into_array(), 5, ten_fill_value, + &mut ctx, ) .unwrap() .into_array(), @@ -102,46 +107,53 @@ mod tests { buffer![2u64, 5, 8].into_array(), PrimitiveArray::from_option_iter([Some(100i32), Some(200), Some(300)]).into_array(), 10, - Scalar::null_native::() + Scalar::null_native::(), + &mut SESSION.create_execution_ctx(), ).unwrap())] #[case::sparse_i32_value_fill(Sparse::try_new( buffer![1u64, 3, 7].into_array(), buffer![42i32, 84, 126].into_array(), 10, - Scalar::from(0i32) + Scalar::from(0i32), + &mut SESSION.create_execution_ctx(), ).unwrap())] // Different types #[case::sparse_u64(Sparse::try_new( buffer![0u64, 4, 9].into_array(), buffer![1000u64, 2000, 3000].into_array(), 10, - Scalar::from(999u64) + Scalar::from(999u64), + &mut SESSION.create_execution_ctx(), ).unwrap())] #[case::sparse_f32(Sparse::try_new( buffer![2u64, 6].into_array(), buffer![f32::consts::PI, f32::consts::E].into_array(), 8, - Scalar::from(0.0f32) + Scalar::from(0.0f32), + &mut SESSION.create_execution_ctx(), ).unwrap())] // Edge cases #[case::sparse_single_patch(Sparse::try_new( buffer![5u64].into_array(), buffer![42i32].into_array(), 10, - Scalar::from(-1i32) + Scalar::from(-1i32), + &mut SESSION.create_execution_ctx(), ).unwrap())] #[case::sparse_dense_patches(Sparse::try_new( buffer![0u64, 1, 2, 3, 4].into_array(), PrimitiveArray::from_option_iter([Some(10i32), Some(20), Some(30), Some(40), Some(50)]).into_array(), 5, - Scalar::null_native::() + Scalar::null_native::(), + &mut SESSION.create_execution_ctx(), ).unwrap())] // Large sparse arrays #[case::sparse_large(Sparse::try_new( buffer![100u64, 500, 900, 1500, 1999].into_array(), buffer![111i32, 222, 333, 444, 555].into_array(), 2000, - Scalar::from(0i32) + Scalar::from(0i32), + &mut SESSION.create_execution_ctx(), ).unwrap())] // Nullable patches #[case::sparse_nullable_patches({ @@ -153,7 +165,8 @@ mod tests { .cast(null_fill_value.dtype().clone()) .unwrap(), 10, - null_fill_value + null_fill_value, + &mut SESSION.create_execution_ctx(), ).unwrap() })] fn test_sparse_consistency(#[case] array: SparseArray) { @@ -165,38 +178,43 @@ mod tests { buffer![2u64, 5, 8].into_array(), buffer![100i32, 200, 300].into_array(), 10, - Scalar::from(0i32) + Scalar::from(0i32), + &mut SESSION.create_execution_ctx(), ).unwrap())] #[case::sparse_u32_basic(Sparse::try_new( buffer![1u64, 3, 7].into_array(), buffer![1000u32, 2000, 3000].into_array(), 10, - Scalar::from(100u32) + Scalar::from(100u32), + &mut SESSION.create_execution_ctx(), ).unwrap())] #[case::sparse_i64_basic(Sparse::try_new( buffer![0u64, 4, 9].into_array(), buffer![5000i64, 6000, 7000].into_array(), 10, - Scalar::from(1000i64) + Scalar::from(1000i64), + &mut SESSION.create_execution_ctx(), ).unwrap())] #[case::sparse_f32_basic(Sparse::try_new( buffer![2u64, 6].into_array(), buffer![1.5f32, 2.5].into_array(), 8, - Scalar::from(0.5f32) + Scalar::from(0.5f32), + &mut SESSION.create_execution_ctx(), ).unwrap())] #[case::sparse_f64_basic(Sparse::try_new( buffer![1u64, 5, 9].into_array(), buffer![10.1f64, 20.2, 30.3].into_array(), 10, - Scalar::from(5.0f64) + Scalar::from(5.0f64), + &mut SESSION.create_execution_ctx(), ).unwrap())] #[case::sparse_i32_large(Sparse::try_new( buffer![10u64, 50, 90, 150, 199].into_array(), buffer![111i32, 222, 333, 444, 555].into_array(), 200, Scalar::from(0i32) - ).unwrap())] + , &mut SESSION.create_execution_ctx()).unwrap())] fn test_sparse_binary_numeric(#[case] array: SparseArray) { test_binary_numeric_array(&array.into_array(), &mut SESSION.create_execution_ctx()); } diff --git a/encodings/sparse/src/compute/nan_count.rs b/encodings/sparse/src/compute/nan_count.rs index 9a8a3a871c1..3d3d20a41e4 100644 --- a/encodings/sparse/src/compute/nan_count.rs +++ b/encodings/sparse/src/compute/nan_count.rs @@ -99,11 +99,11 @@ mod tests { #[rstest] // NaN fill value → all unpatched positions are NaN - #[case(Sparse::try_new(buffer![1u64, 3].into_array(), buffer![1.0f32, 2.0].into_array(), 6, Scalar::from(f32::NAN)).unwrap())] + #[case(Sparse::try_new(buffer![1u64, 3].into_array(), buffer![1.0f32, 2.0].into_array(), 6, Scalar::from(f32::NAN), &mut SESSION.create_execution_ctx()).unwrap())] // NaN patch values, finite fill - #[case(Sparse::try_new(buffer![1u64, 3].into_array(), buffer![f32::NAN, 2.0].into_array(), 6, Scalar::from(0.0f32)).unwrap())] + #[case(Sparse::try_new(buffer![1u64, 3].into_array(), buffer![f32::NAN, 2.0].into_array(), 6, Scalar::from(0.0f32), &mut SESSION.create_execution_ctx()).unwrap())] // no NaNs anywhere - #[case(Sparse::try_new(buffer![1u64, 3].into_array(), buffer![1.0f32, 2.0].into_array(), 6, Scalar::from(0.0f32)).unwrap())] + #[case(Sparse::try_new(buffer![1u64, 3].into_array(), buffer![1.0f32, 2.0].into_array(), 6, Scalar::from(0.0f32), &mut SESSION.create_execution_ctx()).unwrap())] fn nan_count_matches_canonical(#[case] array: SparseArray) { let arr = array.into_array(); let kernel = nan_count(&arr, &mut SESSION.create_execution_ctx()).unwrap(); diff --git a/encodings/sparse/src/compute/null_count.rs b/encodings/sparse/src/compute/null_count.rs index f63dad80205..205c00b528e 100644 --- a/encodings/sparse/src/compute/null_count.rs +++ b/encodings/sparse/src/compute/null_count.rs @@ -93,13 +93,14 @@ mod tests { #[rstest] // non-null fill, no null patches → 0 - #[case(Sparse::try_new(buffer![1u64, 3].into_array(), buffer![10i32, 20].into_array(), 5, Scalar::from(1i32)).unwrap())] + #[case(Sparse::try_new(buffer![1u64, 3].into_array(), buffer![10i32, 20].into_array(), 5, Scalar::from(1i32), &mut SESSION.create_execution_ctx()).unwrap())] // null fill (8 - 2 = 6 fill nulls), patches non-null #[case(Sparse::try_new( buffer![1u64, 3].into_array(), PrimitiveArray::from_option_iter([Some(10i32), Some(20)]).into_array().cast(nullable_i32()).unwrap(), 8, Scalar::null(nullable_i32()), + &mut SESSION.create_execution_ctx(), ).unwrap())] // null fill + some null patches #[case(Sparse::try_new( @@ -107,6 +108,7 @@ mod tests { PrimitiveArray::from_option_iter([Some(10i32), None, Some(30)]).into_array().cast(nullable_i32()).unwrap(), 6, Scalar::null(nullable_i32()), + &mut SESSION.create_execution_ctx(), ).unwrap())] fn null_count_matches_canonical(#[case] array: SparseArray) { let arr = array.into_array(); diff --git a/encodings/sparse/src/compute/sum.rs b/encodings/sparse/src/compute/sum.rs index 31f33e169eb..02c36ad36f3 100644 --- a/encodings/sparse/src/compute/sum.rs +++ b/encodings/sparse/src/compute/sum.rs @@ -147,6 +147,7 @@ mod tests { buffer![10i32, 20].into_array(), 5, Scalar::from(1i32), + &mut SESSION.create_execution_ctx(), ).unwrap(), // 10 + 1 + 20 + 1 + 1 = 33 33i64, @@ -157,6 +158,7 @@ mod tests { buffer![7i32, 8].into_array(), 10, Scalar::from(0i32), + &mut SESSION.create_execution_ctx(), ).unwrap(), 15i64, )] @@ -176,6 +178,7 @@ mod tests { vortex_array::dtype::PType::I64, vortex_array::dtype::Nullability::Nullable, )), + &mut SESSION.create_execution_ctx(), ).unwrap(), 16i64, )] @@ -186,11 +189,13 @@ mod tests { #[test] fn sum_v2_kernel_nonempty() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let array = Sparse::try_new( buffer![0u64, 2].into_array(), buffer![10i32, 20].into_array(), 5, Scalar::from(1i32), + &mut ctx, )?; let result = check_v2(array)?; assert_eq!(result.as_primitive().typed_value::(), Some(33)); @@ -199,12 +204,14 @@ mod tests { #[test] fn sum_v2_kernel_all_null() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let array = Sparse::try_new( buffer![0u64, 3].into_array(), vortex_array::arrays::PrimitiveArray::from_option_iter([None::, None]) .into_array(), 6, Scalar::null_native::(), + &mut ctx, )?; assert!(check_v2(array)?.is_null()); Ok(()) @@ -212,11 +219,13 @@ mod tests { #[test] fn sum_v2_kernel_preserves_options() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let array = Sparse::try_new( buffer![0u64].into_array(), buffer![1.0f64].into_array(), 3, Scalar::from(f64::NAN), + &mut ctx, )? .into_array(); let result = run_v2_kernel(&array, NumericalAggregateOpts::include_nans())?; diff --git a/encodings/sparse/src/compute/take.rs b/encodings/sparse/src/compute/take.rs index 3babc19e57b..c59c2c61533 100644 --- a/encodings/sparse/src/compute/take.rs +++ b/encodings/sparse/src/compute/take.rs @@ -86,11 +86,13 @@ mod test { } fn sparse_array() -> ArrayRef { + let mut ctx = SESSION.create_execution_ctx(); Sparse::try_new( buffer![0u64, 37, 47, 99].into_array(), PrimitiveArray::new(buffer![1.23f64, 0.47, 9.99, 3.5], Validity::AllValid).into_array(), 100, test_array_fill_value(), + &mut ctx, ) .unwrap() .into_array() @@ -151,11 +153,13 @@ mod test { #[test] fn nullable_take() { + let mut ctx = SESSION.create_execution_ctx(); let arr = Sparse::try_new( buffer![1u32].into_array(), buffer![10].into_array(), 10, Scalar::primitive(1, Nullability::NonNullable), + &mut ctx, ) .unwrap(); @@ -176,11 +180,13 @@ mod test { #[test] fn nullable_take_with_many_patches() { + let mut ctx = SESSION.create_execution_ctx(); let arr = Sparse::try_new( buffer![1u32, 3, 7, 8, 9].into_array(), buffer![10, 8, 3, 2, 1].into_array(), 10, Scalar::primitive(1, Nullability::NonNullable), + &mut ctx, ) .unwrap(); @@ -205,12 +211,14 @@ mod test { PrimitiveArray::new(buffer![1.23f64, 0.47, 9.99, 3.5], Validity::AllValid).into_array(), 100, Scalar::null_native::(), + &mut SESSION.create_execution_ctx(), ).unwrap())] #[case(Sparse::try_new( buffer![1u32, 3, 7, 8, 9].into_array(), buffer![10, 8, 3, 2, 1].into_array(), 10, Scalar::from(0i32), + &mut SESSION.create_execution_ctx(), ).unwrap())] #[case({ let nullable_values = PrimitiveArray::from_option_iter([Some(100i64), None, Some(300)]); @@ -219,6 +227,7 @@ mod test { nullable_values.into_array(), 10, Scalar::null_native::(), + &mut SESSION.create_execution_ctx(), ).unwrap() })] #[case(Sparse::try_new( @@ -226,6 +235,7 @@ mod test { buffer![999i32].into_array(), 20, Scalar::from(-1i32), + &mut SESSION.create_execution_ctx(), ).unwrap())] fn test_take_sparse_conformance(#[case] sparse: SparseArray) { test_take_conformance(&sparse.into_array(), &mut SESSION.create_execution_ctx()); diff --git a/encodings/sparse/src/lib.rs b/encodings/sparse/src/lib.rs index f499524ac4e..f1ec2ddc32c 100644 --- a/encodings/sparse/src/lib.rs +++ b/encodings/sparse/src/lib.rs @@ -22,6 +22,7 @@ use vortex_array::EqMode; use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::Primitive; @@ -148,17 +149,23 @@ pub(crate) trait SparseOwnedExt { impl SparseOwnedExt for Array { fn into_parts(self) -> VortexResult { - let patches = Patches::new( - self.len(), - self.patches().offset(), - self.slots()[SparseSlots::PATCH_INDICES] - .clone() - .vortex_expect("indices"), - self.slots()[SparseSlots::PATCH_VALUES] - .clone() - .vortex_expect("values"), - self.slots()[SparseSlots::PATCH_CHUNK_OFFSETS].clone(), - )?; + let chunk_offsets = self.slots()[SparseSlots::PATCH_CHUNK_OFFSETS].clone(); + let offset_within_chunk = chunk_offsets.as_ref().map(|_| 0); + // SAFETY: the patch slots come from a valid sparse array. + let patches = unsafe { + Patches::new_unchecked( + self.len(), + self.patches().offset(), + self.slots()[SparseSlots::PATCH_INDICES] + .clone() + .vortex_expect("indices"), + self.slots()[SparseSlots::PATCH_VALUES] + .clone() + .vortex_expect("values"), + chunk_offsets, + offset_within_chunk, + ) + }; Ok(SparseParts { patches, fill_value: self.fill_scalar().clone(), @@ -290,6 +297,7 @@ impl VTable for Sparse { patch_indices, patch_values, None, + &mut session.create_execution_ctx(), )?; let slots = SparseData::make_slots(&patches); let data = SparseData::from_patches(&patches, fill_value)?; @@ -328,7 +336,7 @@ impl VTable for Sparse { ConstantArray::new(offset_scalar, indices.len()).into_array(), Operator::Sub, )?; - let patches = Patches::new(len, 0, resolved_indices.clone(), values, None)?; + let patches = Patches::new(len, 0, resolved_indices.clone(), values, None, ctx)?; // Decompose, update in place, and reassemble without re-validation. match array.try_into_parts() { Ok(mut parts) => { @@ -423,6 +431,7 @@ impl Sparse { values: ArrayRef, len: usize, fill_value: Scalar, + ctx: &mut ExecutionCtx, ) -> VortexResult { let dtype = fill_value.dtype().clone(); vortex_ensure_eq!( @@ -430,7 +439,7 @@ impl Sparse { &dtype, "sparse values dtype must match fill value dtype" ); - let patches = Patches::new(len, 0, indices, values, None)?; + let patches = Patches::new(len, 0, indices, values, None, ctx)?; let slots = SparseData::make_slots(&patches); let data = SparseData::from_patches(&patches, fill_value)?; Ok( @@ -601,6 +610,7 @@ impl SparseData { non_null_values, array.len(), Scalar::null(array.dtype().clone()), + ctx, ) .map(IntoArray::into_array); } @@ -641,7 +651,7 @@ impl SparseData { Mask::Values(values) => values.indices().iter().map(|v| *v as u64).collect(), }; - Sparse::try_new(indices.into_array(), non_top_values, array.len(), fill) + Sparse::try_new(indices.into_array(), non_top_values, array.len(), fill, ctx) .map(IntoArray::into_array) } } @@ -662,14 +672,19 @@ pub trait SparseExt { Operator::Sub, )?; - Patches::new( - patches.array_len(), - 0, - indices, - patches.values().clone(), - // TODO(0ax1): handle chunk offsets - None, - ) + // SAFETY: subtracting the offset from valid patch indices keeps them sorted and within + // the array length. + Ok(unsafe { + Patches::new_unchecked( + patches.array_len(), + 0, + indices, + patches.values().clone(), + // TODO(0ax1): handle chunk offsets + None, + None, + ) + }) } } @@ -748,13 +763,20 @@ mod test { } fn sparse_array(fill_value: Scalar) -> ArrayRef { + let mut ctx = SESSION.create_execution_ctx(); // merged array: [null, null, 100, null, null, 200, null, null, 300, null] let mut values = buffer![100i32, 200, 300].into_array(); values = values.cast(fill_value.dtype().clone()).unwrap(); - Sparse::try_new(buffer![2u64, 5, 8].into_array(), values, 10, fill_value) - .unwrap() - .into_array() + Sparse::try_new( + buffer![2u64, 5, 8].into_array(), + values, + 10, + fill_value, + &mut ctx, + ) + .unwrap() + .into_array() } #[test] @@ -792,11 +814,13 @@ mod test { #[test] pub fn test_scalar_at_again() { + let mut ctx = SESSION.create_execution_ctx(); let arr = Sparse::try_new( ConstantArray::new(10u32, 1).into_array(), ConstantArray::new(Scalar::primitive(1234u32, Nullability::Nullable), 1).into_array(), 100, Scalar::null(DType::Primitive(PType::U32, Nullability::Nullable)), + &mut ctx, ) .unwrap(); @@ -848,6 +872,7 @@ mod test { #[test] pub fn validity_mask_sliced_nonnull_fill() { + let mut ctx = SESSION.create_execution_ctx(); let sliced = Sparse::try_new( buffer![2u64, 5, 8].into_array(), ConstantArray::new( @@ -857,6 +882,7 @@ mod test { .into_array(), 10, Scalar::primitive(1.0f32, Nullability::Nullable), + &mut ctx, ) .unwrap() .slice(2..7) @@ -934,7 +960,8 @@ mod test { let values = VarBinViewArray::from_iter_nullable_str([Some("patched"), None, Some("last")]) .into_array(); let fill = Scalar::null(values.dtype().clone()); - let array = Sparse::try_new(buffer![1u8, 3, 5].into_array(), values, 6, fill).unwrap(); + let array = + Sparse::try_new(buffer![1u8, 3, 5].into_array(), values, 6, fill, &mut ctx).unwrap(); let expected = VarBinViewArray::from_iter_nullable_str([ None, Some("patched"), @@ -962,6 +989,7 @@ mod test { values, 2, Scalar::utf8("fill".to_string(), Nullability::NonNullable), + &mut ctx, ) .unwrap(); let expected = VarBinViewArray::from_iter_str(["fill", "second"]).into_array(); @@ -977,18 +1005,20 @@ mod test { #[test] #[should_panic] fn test_invalid_length() { + let mut ctx = SESSION.create_execution_ctx(); let values = buffer![15_u32, 135, 13531, 42].into_array(); let indices = buffer![10_u64, 11, 50, 100].into_array(); - Sparse::try_new(indices, values, 100, 0_u32.into()).unwrap(); + Sparse::try_new(indices, values, 100, 0_u32.into(), &mut ctx).unwrap(); } #[test] fn test_valid_length() { + let mut ctx = SESSION.create_execution_ctx(); let values = buffer![15_u32, 135, 13531, 42].into_array(); let indices = buffer![10_u64, 11, 50, 100].into_array(); - Sparse::try_new(indices, values, 101, 0_u32.into()).unwrap(); + Sparse::try_new(indices, values, 101, 0_u32.into(), &mut ctx).unwrap(); } #[test] @@ -1017,10 +1047,12 @@ mod test { #[test] fn validity_mask_includes_null_values_when_fill_is_null() { + let mut ctx = SESSION.create_execution_ctx(); let indices = buffer![0u8, 2, 4, 6, 8].into_array(); let values = PrimitiveArray::from_option_iter([Some(0i16), Some(1), None, None, Some(4)]) .into_array(); - let array = Sparse::try_new(indices, values, 10, Scalar::null_native::()).unwrap(); + let array = + Sparse::try_new(indices, values, 10, Scalar::null_native::(), &mut ctx).unwrap(); let actual = array .validity() .unwrap() diff --git a/encodings/sparse/src/ops.rs b/encodings/sparse/src/ops.rs index 10d71f6b082..06616b7e6f8 100644 --- a/encodings/sparse/src/ops.rs +++ b/encodings/sparse/src/ops.rs @@ -50,7 +50,7 @@ mod tests { let values = buffer![0u64].into_array(); let indices = buffer![0u8].into_array(); - let sparse = Sparse::try_new(indices, values, 1000, 999u64.into()).unwrap(); + let sparse = Sparse::try_new(indices, values, 1000, 999u64.into(), &mut ctx).unwrap(); let sliced = sparse.slice(0..1000).unwrap(); let mut expected = vec![999u64; 1000]; expected[0] = 0; diff --git a/encodings/sparse/src/slice.rs b/encodings/sparse/src/slice.rs index 02d48b91d80..badfc1e1796 100644 --- a/encodings/sparse/src/slice.rs +++ b/encodings/sparse/src/slice.rs @@ -18,9 +18,9 @@ impl SliceKernel for Sparse { fn slice( array: ArrayView<'_, Self>, range: Range, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult> { - let Some(new_patches) = array.patches().slice(range.clone())? else { + let Some(new_patches) = array.patches().slice(range.clone(), ctx)? else { return Ok(Some( ConstantArray::new(array.fill_scalar().clone(), range.len()).into_array(), )); diff --git a/vortex-array/benches/patches_lookup.rs b/vortex-array/benches/patches_lookup.rs index 47a1340f511..c81fdaa573e 100644 --- a/vortex-array/benches/patches_lookup.rs +++ b/vortex-array/benches/patches_lookup.rs @@ -37,6 +37,7 @@ const PATCH_LOW: usize = 100_000; const PATCH_HIGH: usize = 110_000; fn patches_from_indices(index_iter: impl Iterator, chunked: bool) -> Patches { + let mut ctx = SESSION.create_execution_ctx(); let mut indices: Vec = index_iter.collect(); indices.sort(); indices.dedup(); @@ -56,6 +57,7 @@ fn patches_from_indices(index_iter: impl Iterator, chunked: bool) -> Buffer::from(indices).into_array(), values.into_array(), chunk_offsets, + &mut ctx, ) .unwrap(); if chunked { diff --git a/vortex-array/benches/take_patches.rs b/vortex-array/benches/take_patches.rs index 395c275c7d3..1f91e873264 100644 --- a/vortex-array/benches/take_patches.rs +++ b/vortex-array/benches/take_patches.rs @@ -114,16 +114,18 @@ fn take_map(bencher: Bencher, (patches_sparsity, index_multiple): (f64, f64)) { } fn fixture(len: usize, sparsity: f64, rng: &mut StdRng) -> Patches { + let mut ctx = SESSION.create_execution_ctx(); let indices = (0..len) .filter(|_| rng.random_bool(sparsity)) .map(|x| x as u64) .collect::>(); let sparse_len = indices.len(); let values = Buffer::from_iter((0..sparse_len).map(|x| x as u64)).into_array(); - Patches::new(len, 0, indices.into_array(), values, None).unwrap() + Patches::new(len, 0, indices.into_array(), values, None, &mut ctx).unwrap() } fn fixture_with_chunk_offsets(len: usize, sparsity: f64, rng: &mut StdRng) -> Patches { + let mut ctx = SESSION.create_execution_ctx(); let patch_indices = (0..len) .filter(|_| rng.random_bool(sparsity)) .map(|x| x as u64) @@ -146,6 +148,7 @@ fn fixture_with_chunk_offsets(len: usize, sparsity: f64, rng: &mut StdRng) -> Pa Buffer::from(patch_indices).into_array(), values, Some(Buffer::from(chunk_offsets).into_array()), + &mut ctx, ) .unwrap() } diff --git a/vortex-array/src/arrays/bool/array.rs b/vortex-array/src/arrays/bool/array.rs index 0d9d40782c4..0ece3025322 100644 --- a/vortex-array/src/arrays/bool/array.rs +++ b/vortex-array/src/arrays/bool/array.rs @@ -453,6 +453,7 @@ mod tests { buffer![4u32].into_array(), BoolArray::from(BitBuffer::new_unset(1)).into_array(), None, + &mut ctx, ) .unwrap(); let arr = arr.patch(&patches, &mut ctx).unwrap(); @@ -503,6 +504,7 @@ mod tests { PrimitiveArray::new(buffer![0u32], Validity::NonNullable).into_array(), BoolArray::from(BitBuffer::new_unset(1)).into_array(), None, + &mut ctx, ) .unwrap(); let arr = arr.patch(&patches, &mut ctx).unwrap(); diff --git a/vortex-array/src/arrays/patched/array.rs b/vortex-array/src/arrays/patched/array.rs index eda6907c935..48e679efb86 100644 --- a/vortex-array/src/arrays/patched/array.rs +++ b/vortex-array/src/arrays/patched/array.rs @@ -17,7 +17,6 @@ use crate::ArraySlots; use crate::Canonical; use crate::ExecutionCtx; use crate::IntoArray; -use crate::VortexSessionExecute; use crate::array::Array; use crate::array::ArrayParts; use crate::array::TypedArrayRef; @@ -31,7 +30,6 @@ use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::dtype::NativePType; use crate::dtype::PType; -use crate::legacy_session; use crate::match_each_native_ptype; use crate::match_each_unsigned_integer_ptype; use crate::patches::Patches; @@ -112,19 +110,21 @@ pub trait PatchedArrayExt: PatchedArraySlotsExt { } #[inline] - #[allow(clippy::disallowed_methods)] - fn lane_range(&self, chunk: usize, lane: usize) -> VortexResult> { + fn lane_range( + &self, + chunk: usize, + lane: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { assert!(chunk * 1024 <= self.len() + self.offset()); assert!(lane < self.n_lanes()); - let start = self.lane_offsets().execute_scalar( - chunk * self.n_lanes() + lane, - &mut legacy_session().create_execution_ctx(), - )?; - let stop = self.lane_offsets().execute_scalar( - chunk * self.n_lanes() + lane + 1, - &mut legacy_session().create_execution_ctx(), - )?; + let start = self + .lane_offsets() + .execute_scalar(chunk * self.n_lanes() + lane, ctx)?; + let stop = self + .lane_offsets() + .execute_scalar(chunk * self.n_lanes() + lane + 1, ctx)?; let start = start .as_primitive() diff --git a/vortex-array/src/arrays/patched/compute/compare.rs b/vortex-array/src/arrays/patched/compute/compare.rs index 81d30c2bf99..0587d82f57f 100644 --- a/vortex-array/src/arrays/patched/compute/compare.rs +++ b/vortex-array/src/arrays/patched/compute/compare.rs @@ -165,6 +165,7 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::Patched; @@ -178,6 +179,7 @@ mod tests { #[test] fn test_basic() { + let mut ctx = array_session().create_execution_ctx(); let lhs = PrimitiveArray::from_iter(0u32..512).into_array(); let patches = Patches::new( 512, @@ -185,10 +187,11 @@ mod tests { buffer![509u16, 510, 511].into_array(), buffer![u32::MAX; 3].into_array(), None, + &mut ctx, ) .unwrap(); - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let lhs = Patched::from_array_and_patches(lhs, &patches, &mut ctx) .unwrap() @@ -211,6 +214,7 @@ mod tests { #[test] fn test_with_offset() { + let mut ctx = array_session().create_execution_ctx(); let lhs = PrimitiveArray::from_iter(0u32..512).into_array(); let patches = Patches::new( 512, @@ -218,10 +222,11 @@ mod tests { buffer![5u16, 510, 511].into_array(), buffer![u32::MAX; 3].into_array(), None, + &mut ctx, ) .unwrap(); - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let lhs = Patched::from_array_and_patches(lhs, &patches, &mut ctx).unwrap(); // Slice the array so that the first patch should be skipped. @@ -244,6 +249,7 @@ mod tests { #[test] fn test_subnormal_f32() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // Subnormal f32 values are smaller than f32::MIN_POSITIVE but greater than 0 let subnormal: f32 = f32::MIN_POSITIVE / 2.0; assert!(subnormal > 0.0 && subnormal < f32::MIN_POSITIVE); @@ -256,9 +262,10 @@ mod tests { buffer![509u16, 510, 511].into_array(), buffer![f32::NAN, subnormal, f32::NEG_INFINITY].into_array(), None, + &mut ctx, )?; - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let lhs = Patched::from_array_and_patches(lhs, &patches, &mut ctx)? .into_array() .try_downcast::() @@ -282,6 +289,7 @@ mod tests { #[test] fn test_pos_neg_zero() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let lhs = PrimitiveArray::from_iter([-0.0f32; 10]).into_array(); let patches = Patches::new( @@ -290,9 +298,10 @@ mod tests { buffer![5u16, 6, 7, 8, 9].into_array(), buffer![f32::NAN, f32::NEG_INFINITY, 0f32, -0.0f32, f32::INFINITY].into_array(), None, + &mut ctx, )?; - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let lhs = Patched::from_array_and_patches(lhs, &patches, &mut ctx)? .into_array() .try_downcast::() diff --git a/vortex-array/src/arrays/patched/compute/filter.rs b/vortex-array/src/arrays/patched/compute/filter.rs index 8db401afb9b..85705e9c710 100644 --- a/vortex-array/src/arrays/patched/compute/filter.rs +++ b/vortex-array/src/arrays/patched/compute/filter.rs @@ -84,7 +84,7 @@ mod tests { let patched_indices = buffer![3u8, 4].into_array(); let patched_values = buffer![u16::MAX; 2].into_array(); - let patches = Patches::new(5, 0, patched_indices, patched_values, None)?; + let patches = Patches::new(5, 0, patched_indices, patched_values, None, &mut ctx)?; let array = Patched::from_array_and_patches(array, &patches, &mut ctx)?.into_array(); @@ -108,7 +108,7 @@ mod tests { let patched_indices = buffer![5u16, 1030].into_array(); let patched_values = buffer![u16::MAX; 2].into_array(); - let patches = Patches::new(4096, 0, patched_indices, patched_values, None)?; + let patches = Patches::new(4096, 0, patched_indices, patched_values, None, &mut ctx)?; let array = Patched::from_array_and_patches(array, &patches, &mut ctx)?.into_array(); @@ -133,7 +133,7 @@ mod tests { let patched_indices = buffer![1024u16, 1025].into_array(); let patched_values = buffer![u16::MAX, u16::MAX].into_array(); - let patches = Patches::new(4096, 0, patched_indices, patched_values, None)?; + let patches = Patches::new(4096, 0, patched_indices, patched_values, None, &mut ctx)?; let array = Patched::from_array_and_patches(array, &patches, &mut ctx)?.into_array(); @@ -157,7 +157,7 @@ mod tests { let patched_indices = buffer![1024u16, 1025].into_array(); let patched_values = buffer![u16::MAX, u16::MAX].into_array(); - let patches = Patches::new(4096, 1, patched_indices, patched_values, None)?; + let patches = Patches::new(4096, 1, patched_indices, patched_values, None, &mut ctx)?; let array = Patched::from_array_and_patches(array, &patches, &mut ctx)?.into_array(); @@ -184,7 +184,7 @@ mod tests { let patched_indices = buffer![2048u16, 2049].into_array(); let patched_values = buffer![u16::MAX, u16::MAX].into_array(); - let patches = Patches::new(6144, 0, patched_indices, patched_values, None)?; + let patches = Patches::new(6144, 0, patched_indices, patched_values, None, &mut ctx)?; let patched = Patched::from_array_and_patches(array, &patches, &mut ctx)?; @@ -220,7 +220,7 @@ mod tests { let patched_indices = buffer![5u16, 1030].into_array(); let patched_values = buffer![9999u16, 8888].into_array(); - let patches = Patches::new(4096, 0, patched_indices, patched_values, None)?; + let patches = Patches::new(4096, 0, patched_indices, patched_values, None, &mut ctx)?; let array = Patched::from_array_and_patches(array, &patches, &mut ctx)?.into_array(); // Slice to create offset > 0. @@ -252,7 +252,7 @@ mod tests { let patched_indices = buffer![5000u16, 6000].into_array(); let patched_values = buffer![u16::MAX; 2].into_array(); - let patches = Patches::new(6144, 0, patched_indices, patched_values, None)?; + let patches = Patches::new(6144, 0, patched_indices, patched_values, None, &mut ctx)?; let patched = Patched::from_array_and_patches(array, &patches, &mut ctx)?; diff --git a/vortex-array/src/arrays/patched/compute/take.rs b/vortex-array/src/arrays/patched/compute/take.rs index 07b67280920..48725596ce2 100644 --- a/vortex-array/src/arrays/patched/compute/take.rs +++ b/vortex-array/src/arrays/patched/compute/take.rs @@ -150,6 +150,7 @@ mod tests { patch_values: &[u16], slice: Range, ) -> VortexResult { + let mut ctx = array_session().create_execution_ctx(); let values = PrimitiveArray::from_iter(base.iter().copied()).into_array(); let patches = Patches::new( base.len(), @@ -157,6 +158,7 @@ mod tests { PrimitiveArray::from_iter(patch_indices.iter().copied()).into_array(), PrimitiveArray::from_iter(patch_values.iter().copied()).into_array(), None, + &mut ctx, )?; let session = VortexSession::empty(); diff --git a/vortex-array/src/arrays/patched/vtable/mod.rs b/vortex-array/src/arrays/patched/vtable/mod.rs index ec332a40566..c153e2b8c00 100644 --- a/vortex-array/src/arrays/patched/vtable/mod.rs +++ b/vortex-array/src/arrays/patched/vtable/mod.rs @@ -379,6 +379,7 @@ mod tests { #[test] fn test_execute() { + let mut ctx = array_session().create_execution_ctx(); let values = buffer![0u16; 1024].into_array(); let patches = Patches::new( 1024, @@ -386,6 +387,7 @@ mod tests { buffer![1u32, 2, 3].into_array(), buffer![1u16; 3].into_array(), None, + &mut ctx, ) .unwrap(); @@ -406,6 +408,7 @@ mod tests { #[test] fn test_execute_sliced() { + let mut ctx = array_session().create_execution_ctx(); let values = buffer![0u16; 1024].into_array(); let patches = Patches::new( 1024, @@ -413,6 +416,7 @@ mod tests { buffer![1u32, 2, 3].into_array(), buffer![1u16; 3].into_array(), None, + &mut ctx, ) .unwrap(); @@ -433,6 +437,7 @@ mod tests { #[test] fn test_append_to_builder_non_nullable() { + let mut ctx = array_session().create_execution_ctx(); let values = PrimitiveArray::new(buffer![0u16; 1024], Validity::NonNullable).into_array(); let patches = Patches::new( 1024, @@ -440,6 +445,7 @@ mod tests { buffer![1u32, 2, 3].into_array(), buffer![10u16, 20, 30].into_array(), None, + &mut ctx, ) .unwrap(); @@ -470,6 +476,7 @@ mod tests { #[test] fn test_append_to_builder_sliced() { + let mut ctx = array_session().create_execution_ctx(); let values = PrimitiveArray::new(buffer![0u16; 1024], Validity::NonNullable).into_array(); let patches = Patches::new( 1024, @@ -477,6 +484,7 @@ mod tests { buffer![1u32, 2, 3].into_array(), buffer![10u16, 20, 30].into_array(), None, + &mut ctx, ) .unwrap(); @@ -507,6 +515,7 @@ mod tests { #[test] fn test_append_to_builder_with_validity() { + let mut ctx = array_session().create_execution_ctx(); // Create inner array with nulls at indices 0 and 5. let validity = Validity::from_iter((0..10).map(|i| i != 0 && i != 5)); let values = PrimitiveArray::new(buffer![0u16; 10], validity).into_array(); @@ -518,6 +527,7 @@ mod tests { buffer![1u32, 2, 3].into_array(), buffer![10u16, 20, 30].into_array(), None, + &mut ctx, ) .unwrap(); @@ -560,6 +570,7 @@ mod tests { patch_indices: &[u32], patch_values: &[u16], ) -> VortexResult { + let mut ctx = array_session().create_execution_ctx(); let values: Vec = inner.into_iter().collect(); let len = values.len(); let array = PrimitiveArray::from_iter(values).into_array(); @@ -567,7 +578,7 @@ mod tests { let indices = PrimitiveArray::from_iter(patch_indices.iter().copied()).into_array(); let patch_vals = PrimitiveArray::from_iter(patch_values.iter().copied()).into_array(); - let patches = Patches::new(len, 0, indices, patch_vals, None)?; + let patches = Patches::new(len, 0, indices, patch_vals, None, &mut ctx)?; let session = array_session(); let mut ctx = session.create_execution_ctx(); diff --git a/vortex-array/src/arrays/patched/vtable/operations.rs b/vortex-array/src/arrays/patched/vtable/operations.rs index bd7440965ab..ae115cd8cba 100644 --- a/vortex-array/src/arrays/patched/vtable/operations.rs +++ b/vortex-array/src/arrays/patched/vtable/operations.rs @@ -31,7 +31,7 @@ impl OperationsVTable for Patched { let lane = (index + array.offset()) % array.n_lanes(); - let range = array.lane_range(chunk, lane)?; + let range = array.lane_range(chunk, lane, ctx)?; // Get the range of indices corresponding to the lane, potentially decoding them to avoid // the overhead of repeated scalar_at calls. @@ -73,6 +73,7 @@ mod tests { #[test] fn test_simple() { + let mut ctx = array_session().create_execution_ctx(); let values = buffer![0u16; 1024].into_array(); let patches = Patches::new( 1024, @@ -80,6 +81,7 @@ mod tests { buffer![1u32, 2, 3].into_array(), buffer![1u16; 3].into_array(), None, + &mut ctx, ) .unwrap(); @@ -118,6 +120,7 @@ mod tests { #[test] fn test_multi_chunk() { + let mut ctx = array_session().create_execution_ctx(); let values = buffer![0u16; 4096].into_array(); let patches = Patches::new( 4096, @@ -125,6 +128,7 @@ mod tests { buffer![1u32, 2, 3].into_array(), buffer![1u16; 3].into_array(), None, + &mut ctx, ) .unwrap(); @@ -150,6 +154,7 @@ mod tests { #[test] fn test_multi_chunk_sliced() { + let mut ctx = array_session().create_execution_ctx(); let values = buffer![0u16; 4096].into_array(); let patches = Patches::new( 4096, @@ -157,6 +162,7 @@ mod tests { buffer![1u32, 2, 3].into_array(), buffer![1u16; 3].into_array(), None, + &mut ctx, ) .unwrap(); diff --git a/vortex-array/src/arrays/patched/vtable/slice.rs b/vortex-array/src/arrays/patched/vtable/slice.rs index 99d60ddae26..70f82decbe7 100644 --- a/vortex-array/src/arrays/patched/vtable/slice.rs +++ b/vortex-array/src/arrays/patched/vtable/slice.rs @@ -67,6 +67,7 @@ mod tests { use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::Patched; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -75,12 +76,13 @@ mod tests { #[test] fn test_reduce() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let values = buffer![0u16; 512].into_array(); let patch_indices = buffer![1u32, 8, 30].into_array(); let patch_values = buffer![u16::MAX; 3].into_array(); - let patches = Patches::new(512, 0, patch_indices, patch_values, None)?; + let patches = Patches::new(512, 0, patch_indices, patch_values, None, &mut ctx)?; - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let patched_array = Patched::from_array_and_patches(values, &patches, &mut ctx)?; @@ -129,10 +131,11 @@ mod tests { patch_indices.into_array(), patch_values.into_array(), None, + &mut array_session().create_execution_ctx(), ) .unwrap(); - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let patched_array = Patched::from_array_and_patches(inner.into_array(), &patches, &mut ctx) .unwrap() @@ -145,13 +148,14 @@ mod tests { #[test] fn test_stacked_slices() { + let mut ctx = array_session().create_execution_ctx(); let values = PrimitiveArray::from_iter(0u64..10_000).into_array(); let patched_indices = buffer![1u32, 2, 1024, 2048, 3072, 3088].into_array(); let patched_values = buffer![0u64, 1, 2, 3, 4, 5].into_array(); - let patches = Patches::new(10_000, 0, patched_indices, patched_values, None).unwrap(); - let mut ctx = crate::array_session().create_execution_ctx(); + let patches = + Patches::new(10_000, 0, patched_indices, patched_values, None, &mut ctx).unwrap(); let patched_array = Patched::from_array_and_patches(values, &patches, &mut ctx) .unwrap() diff --git a/vortex-array/src/compute/conformance/search_sorted.rs b/vortex-array/src/compute/conformance/search_sorted.rs index 0d0cf9d8cc8..fc32af9ddbb 100644 --- a/vortex-array/src/compute/conformance/search_sorted.rs +++ b/vortex-array/src/compute/conformance/search_sorted.rs @@ -22,6 +22,7 @@ pub fn all_null() -> ArrayRef { } pub fn sparse_high_null_fill() -> ArrayRef { + let mut ctx = array_session().create_execution_ctx(); PrimitiveArray::new(buffer![0; 20], Validity::AllInvalid) .patch( &Patches::new( @@ -30,15 +31,17 @@ pub fn sparse_high_null_fill() -> ArrayRef { buffer![17u64, 18, 19].into_array(), PrimitiveArray::new(buffer![33_i32, 44, 55], Validity::AllValid).into_array(), None, + &mut ctx, ) .unwrap(), - &mut array_session().create_execution_ctx(), + &mut ctx, ) .unwrap() .into_array() } pub fn sparse_high_non_null_fill() -> ArrayRef { + let mut ctx = array_session().create_execution_ctx(); PrimitiveArray::new(buffer![22; 20], Validity::NonNullable) .patch( &Patches::new( @@ -47,15 +50,17 @@ pub fn sparse_high_non_null_fill() -> ArrayRef { buffer![17u64, 18, 19].into_array(), buffer![33_i32, 44, 55].into_array(), None, + &mut ctx, ) .unwrap(), - &mut array_session().create_execution_ctx(), + &mut ctx, ) .unwrap() .into_array() } pub fn sparse_low() -> ArrayRef { + let mut ctx = array_session().create_execution_ctx(); PrimitiveArray::new(buffer![60; 20], Validity::NonNullable) .patch( &Patches::new( @@ -64,15 +69,17 @@ pub fn sparse_low() -> ArrayRef { buffer![0u64, 1, 2].into_array(), buffer![33i32, 44, 55].into_array(), None, + &mut ctx, ) .unwrap(), - &mut array_session().create_execution_ctx(), + &mut ctx, ) .unwrap() .into_array() } pub fn sparse_low_high() -> ArrayRef { + let mut ctx = array_session().create_execution_ctx(); PrimitiveArray::new(buffer![30; 20], Validity::NonNullable) .patch( &Patches::new( @@ -81,15 +88,17 @@ pub fn sparse_low_high() -> ArrayRef { buffer![0u64, 1, 17, 18, 19].into_array(), buffer![11i32, 22, 33, 44, 55].into_array(), None, + &mut ctx, ) .unwrap(), - &mut array_session().create_execution_ctx(), + &mut ctx, ) .unwrap() .into_array() } pub fn sparse_edge_patch_high() -> ArrayRef { + let mut ctx = array_session().create_execution_ctx(); PrimitiveArray::new(buffer![33; 20], Validity::NonNullable) .patch( &Patches::new( @@ -98,15 +107,17 @@ pub fn sparse_edge_patch_high() -> ArrayRef { buffer![0u64, 1, 2, 19].into_array(), buffer![11i32, 22, 23, 55].into_array(), None, + &mut ctx, ) .unwrap(), - &mut array_session().create_execution_ctx(), + &mut ctx, ) .unwrap() .into_array() } pub fn sparse_edge_patch_low() -> ArrayRef { + let mut ctx = array_session().create_execution_ctx(); PrimitiveArray::new(buffer![22; 20], Validity::NonNullable) .patch( &Patches::new( @@ -115,9 +126,10 @@ pub fn sparse_edge_patch_low() -> ArrayRef { buffer![0u64, 17, 18, 19].into_array(), buffer![11i32, 33, 44, 55].into_array(), None, + &mut ctx, ) .unwrap(), - &mut array_session().create_execution_ctx(), + &mut ctx, ) .unwrap() .into_array() diff --git a/vortex-array/src/patches.rs b/vortex-array/src/patches.rs index 64fc57da038..853a0dff4bb 100644 --- a/vortex-array/src/patches.rs +++ b/vortex-array/src/patches.rs @@ -25,7 +25,6 @@ use crate::ArrayRef; use crate::ArraySlots; use crate::ExecutionCtx; use crate::IntoArray; -use crate::VortexSessionExecute; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::primitive::PrimitiveArrayExt; @@ -37,7 +36,6 @@ use crate::dtype::Nullability; use crate::dtype::Nullability::NonNullable; use crate::dtype::PType; use crate::dtype::UnsignedPType; -use crate::legacy_session; use crate::match_each_unsigned_integer_ptype; use crate::scalar::Scalar; use crate::search_sorted::SearchResult; @@ -239,13 +237,13 @@ pub struct Patches { } impl Patches { - #[allow(clippy::disallowed_methods)] pub fn new( array_len: usize, offset: usize, indices: ArrayRef, values: ArrayRef, chunk_offsets: Option, + ctx: &mut ExecutionCtx, ) -> VortexResult { vortex_ensure_eq!(indices.len(), values.len()); vortex_ensure!( @@ -263,11 +261,8 @@ 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.execute_scalar( - indices.len() - 1, - &mut legacy_session().create_execution_ctx(), - )?) - .map_err(|_| vortex_err!("indices must be a number"))?; + let max = usize::try_from(&indices.execute_scalar(indices.len() - 1, ctx)?) + .map_err(|_| vortex_err!("indices must be a number"))?; vortex_ensure!( max - offset < array_len, "Patch indices {max:?}, offset {offset} are longer than the array length {array_len}" @@ -276,9 +271,8 @@ impl Patches { #[cfg(debug_assertions)] { use crate::aggregate_fn::fns::is_sorted::is_sorted; - let mut ctx = legacy_session().create_execution_ctx(); assert!( - is_sorted(&indices, &mut ctx).unwrap_or(false), + is_sorted(&indices, ctx).unwrap_or(false), "Patch indices must be sorted" ); } @@ -379,14 +373,13 @@ impl Patches { } #[inline] - #[allow(clippy::disallowed_methods)] - pub fn chunk_offset_at(&self, idx: usize) -> VortexResult { + pub fn chunk_offset_at(&self, idx: usize, ctx: &mut ExecutionCtx) -> VortexResult { let Some(chunk_offsets) = &self.chunk_offsets else { vortex_bail!("chunk_offsets must be set to retrieve offset at index") }; chunk_offsets - .execute_scalar(idx, &mut legacy_session().create_execution_ctx())? + .execute_scalar(idx, ctx)? .as_primitive() .as_::() .ok_or_else(|| vortex_err!("chunk offset does not fit in usize")) @@ -615,11 +608,10 @@ impl Patches { } /// Returns the minimum patch index - #[allow(clippy::disallowed_methods)] - pub fn min_index(&self) -> VortexResult { + pub fn min_index(&self, ctx: &mut ExecutionCtx) -> VortexResult { let first = self .indices - .execute_scalar(0, &mut legacy_session().create_execution_ctx())? + .execute_scalar(0, ctx)? .as_primitive() .as_::() .ok_or_else(|| vortex_err!("index does not fit in usize"))?; @@ -627,14 +619,10 @@ impl Patches { } /// Returns the maximum patch index - #[allow(clippy::disallowed_methods)] - pub fn max_index(&self) -> VortexResult { + pub fn max_index(&self, ctx: &mut ExecutionCtx) -> VortexResult { let last = self .indices - .execute_scalar( - self.indices.len() - 1, - &mut legacy_session().create_execution_ctx(), - )? + .execute_scalar(self.indices.len() - 1, ctx)? .as_primitive() .as_::() .ok_or_else(|| vortex_err!("index does not fit in usize"))?; @@ -740,11 +728,9 @@ impl Patches { } /// Slice the patches by a range of the patched array. - #[allow(clippy::disallowed_methods)] - pub fn slice(&self, range: Range) -> VortexResult> { - let mut ctx = legacy_session().create_execution_ctx(); - let slice_start_idx = self.search_index(range.start, &mut ctx)?.to_index(); - let slice_end_idx = self.search_index(range.end, &mut ctx)?.to_index(); + pub fn slice(&self, range: Range, ctx: &mut ExecutionCtx) -> VortexResult> { + let slice_start_idx = self.search_index(range.start, ctx)?.to_index(); + let slice_end_idx = self.search_index(range.end, ctx)?.to_index(); if slice_start_idx == slice_end_idx { return Ok(None); @@ -768,11 +754,11 @@ impl Patches { .as_ref() .map(|new_chunk_offsets| -> VortexResult { let new_chunk_base = new_chunk_offsets - .execute_scalar(0, &mut ctx)? + .execute_scalar(0, ctx)? .as_primitive() .as_::() .ok_or_else(|| vortex_err!("chunk offset does not fit in usize"))?; - let parent_chunk_base = self.chunk_offset_at(0)?; + let parent_chunk_base = self.chunk_offset_at(0, ctx)?; let parent_within = self.offset_within_chunk.unwrap_or(0); Ok(parent_chunk_base + parent_within + slice_start_idx - new_chunk_base) }) @@ -941,8 +927,8 @@ impl Patches { let take_indices_unsigned = take_indices.reinterpret_cast(take_indices.ptype().to_unsigned()); - let min_index = self.min_index()?; - let max_index = self.max_index()?; + let min_index = self.min_index(ctx)?; + let max_index = self.max_index(ctx)?; let Some((new_sparse_indices, value_indices)) = match_each_unsigned_integer_ptype!(indices.ptype(), |Indices| { @@ -1391,6 +1377,7 @@ mod test { buffer![10u32, 11, 20].into_array(), buffer![100, 110, 200].into_array(), None, + &mut ctx, ) .unwrap(); @@ -1420,6 +1407,7 @@ mod test { buffer![2u64, 9, 15].into_array(), PrimitiveArray::new(buffer![33_i32, 44, 55], Validity::AllValid).into_array(), None, + &mut ctx, ) .unwrap(); @@ -1463,6 +1451,7 @@ mod test { buffer![2u64, 9, 15].into_array(), buffer![33_i32, 44, 55].into_array(), Some(buffer![0u64].into_array()), + &mut ctx, ) .unwrap(); @@ -1507,6 +1496,7 @@ mod test { buffer![100u64, 500, 1200, 1800].into_array(), buffer![10_i32, 20, 30, 40].into_array(), Some(buffer![0u64, 2].into_array()), + &mut ctx, ) .unwrap(); @@ -1541,6 +1531,7 @@ mod test { buffer![2u64, 9, 15].into_array(), buffer![33_i32, 44, 55].into_array(), Some(buffer![0u64].into_array()), + &mut ctx, ) .unwrap(); @@ -1564,6 +1555,7 @@ mod test { buffer![5u64, 10, 20, 25].into_array(), buffer![100_i32, 200, 300, 400].into_array(), Some(buffer![0u64].into_array()), + &mut ctx, ) .unwrap(); @@ -1598,6 +1590,7 @@ mod test { BufferMut::from_iter(0..1500u64).into_array(), BufferMut::from_iter(0..1500i32).into_array(), Some(buffer![0u64, 1024u64].into_array()), + &mut ctx, ) .unwrap(); @@ -1629,9 +1622,9 @@ mod test { let values = buffer![15_u32, 135, 13531, 42].into_array(); let indices = buffer![10_u64, 11, 50, 100].into_array(); - let patches = Patches::new(101, 0, indices, values, None).unwrap(); + let patches = Patches::new(101, 0, indices, values, None, &mut ctx).unwrap(); - let sliced = patches.slice(15..100).unwrap().unwrap(); + let sliced = patches.slice(15..100, &mut ctx).unwrap().unwrap(); assert_eq!(sliced.array_len(), 100 - 15); assert_arrays_eq!( sliced.values(), @@ -1645,7 +1638,7 @@ mod test { PrimitiveArray::from_iter([50u64]), &mut ctx ); - let doubly_sliced = sliced.slice(35..36).unwrap().unwrap(); + let doubly_sliced = sliced.slice(35..36, &mut ctx).unwrap().unwrap(); assert_eq!(doubly_sliced.array_len(), 1); assert_eq!(doubly_sliced.offset(), 50); assert_eq!( @@ -1668,6 +1661,7 @@ mod test { buffer![2u64, 5, 8].into_array(), buffer![100i32, 200, 300].into_array(), None, + &mut ctx, ) .unwrap(); @@ -1685,6 +1679,7 @@ mod test { buffer![2u64, 5, 8].into_array(), buffer![100i32, 200, 300].into_array(), None, + &mut ctx, ) .unwrap(); @@ -1719,6 +1714,7 @@ mod test { buffer![2u64, 5, 8].into_array(), buffer![100i32, 200, 300].into_array(), None, + &mut ctx, ) .unwrap(); @@ -1753,6 +1749,7 @@ mod test { buffer![7u64, 10, 13].into_array(), // actual indices are 2, 5, 8 buffer![100i32, 200, 300].into_array(), None, + &mut ctx, ) .unwrap(); @@ -1785,6 +1782,7 @@ mod test { buffer![2u64, 5, 8].into_array(), PrimitiveArray::from_option_iter([Some(100i32), None, Some(300)]).into_array(), None, + &mut ctx, ) .unwrap(); @@ -1846,6 +1844,7 @@ mod test { .collect::>() .into_array(), None, + &mut ctx, )?; let mask = if cache_mask_indices { @@ -1904,6 +1903,7 @@ mod test { (0..=u8::MAX).collect::>().into_array(), (0..256i32).collect::>().into_array(), None, + &mut ctx, )?; let mask = Mask::from_indices(1000, [3, 255, 256, 900]); @@ -1934,6 +1934,7 @@ mod test { buffer![2u64, 5, 8].into_array(), buffer![100i32, 200, 300].into_array(), None, + &mut ctx, ) .unwrap(); @@ -1962,6 +1963,7 @@ mod test { buffer![2u64, 5, 8].into_array(), buffer![100i32, 200, 300].into_array(), None, + &mut ctx, ) .unwrap(); @@ -1980,6 +1982,7 @@ mod test { buffer![2u64, 5, 8].into_array(), buffer![100i32, 200, 300].into_array(), None, + &mut ctx, ) .unwrap(); @@ -2008,10 +2011,11 @@ mod test { buffer![2u64, 5, 8].into_array(), buffer![100i32, 200, 300].into_array(), None, + &mut ctx, ) .unwrap(); - let sliced = patches.slice(0..10).unwrap().unwrap(); + let sliced = patches.slice(0..10, &mut ctx).unwrap().unwrap(); assert_arrays_eq!( sliced.indices(), @@ -2034,11 +2038,12 @@ mod test { buffer![2u64, 5, 8].into_array(), buffer![100i32, 200, 300].into_array(), None, + &mut ctx, ) .unwrap(); // Slice from 3 to 8 (includes patch at 5) - let sliced = patches.slice(3..8).unwrap().unwrap(); + let sliced = patches.slice(3..8, &mut ctx).unwrap().unwrap(); assert_arrays_eq!( sliced.indices(), @@ -2056,17 +2061,19 @@ mod test { #[test] fn test_slice_no_patches() { + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 10, 0, buffer![2u64, 5, 8].into_array(), buffer![100i32, 200, 300].into_array(), None, + &mut ctx, ) .unwrap(); // Slice from 6 to 7 (no patches in this range) - let sliced = patches.slice(6..7).unwrap(); + let sliced = patches.slice(6..7, &mut ctx).unwrap(); assert!(sliced.is_none()); } @@ -2079,11 +2086,12 @@ mod test { buffer![7u64, 10, 13].into_array(), // actual indices are 2, 5, 8 buffer![100i32, 200, 300].into_array(), None, + &mut ctx, ) .unwrap(); // Slice from 3 to 8 (includes patch at actual index 5) - let sliced = patches.slice(3..8).unwrap().unwrap(); + let sliced = patches.slice(3..8, &mut ctx).unwrap().unwrap(); assert_arrays_eq!( sliced.indices(), @@ -2100,17 +2108,19 @@ mod test { #[test] fn test_indices_range() { + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 10, 0, buffer![2u64, 5, 8].into_array(), buffer![100i32, 200, 300].into_array(), None, + &mut ctx, ) .unwrap(); - assert_eq!(patches.min_index().unwrap(), 2); - assert_eq!(patches.max_index().unwrap(), 8); + assert_eq!(patches.min_index(&mut ctx).unwrap(), 2); + assert_eq!(patches.max_index(&mut ctx).unwrap(), 8); } #[test] @@ -2122,6 +2132,7 @@ mod test { buffer![2u64, 5, 8].into_array(), buffer![100i32, 200, 300].into_array(), None, + &mut ctx, ) .unwrap(); @@ -2168,6 +2179,7 @@ mod test { buffer![0u64, 9].into_array(), buffer![100i32, 200].into_array(), None, + &mut ctx, ) .unwrap(); @@ -2199,6 +2211,7 @@ mod test { buffer![2u64, 5, 8].into_array(), buffer![100i32, 200, 300].into_array(), None, + &mut ctx, ) .unwrap(); @@ -2220,6 +2233,7 @@ mod test { buffer![2u64, 5, 8].into_array(), buffer![100i32, 200, 300].into_array(), None, + &mut ctx, ) .unwrap(); @@ -2251,6 +2265,7 @@ mod test { buffer![2u64].into_array(), buffer![42i32].into_array(), None, + &mut ctx, ) .unwrap(); @@ -2284,6 +2299,7 @@ mod test { buffer![3u64, 4, 5, 6].into_array(), buffer![100i32, 200, 300, 400].into_array(), None, + &mut ctx, ) .unwrap(); @@ -2315,6 +2331,7 @@ mod test { buffer![16u64, 17, 19].into_array(), // actual indices are 1, 2, 4 buffer![100i32, 200, 300].into_array(), None, + &mut ctx, ) .unwrap(); @@ -2347,6 +2364,7 @@ mod test { buffer![2u64, 5, 8].into_array(), buffer![100i32, 200, 300].into_array(), None, + &mut ctx, ) .unwrap(); @@ -2361,7 +2379,8 @@ mod test { let indices = buffer![100u64, 200, 3000, 3100].into_array(); let values = buffer![10i32, 20, 30, 40].into_array(); let chunk_offsets = buffer![0u64, 2, 2, 3].into_array(); - let patches = Patches::new(4096, 0, indices, values, Some(chunk_offsets)).unwrap(); + let patches = + Patches::new(4096, 0, indices, values, Some(chunk_offsets), &mut ctx).unwrap(); assert!(patches.chunk_offsets.is_some()); @@ -2407,9 +2426,10 @@ mod test { let indices = buffer![100u64, 500, 1200, 1300, 1500, 1800, 2100, 2500].into_array(); let values = buffer![10i32, 20, 30, 35, 40, 45, 50, 60].into_array(); let chunk_offsets = buffer![0u64, 2, 6].into_array(); - let patches = Patches::new(3000, 0, indices, values, Some(chunk_offsets)).unwrap(); + let patches = + Patches::new(3000, 0, indices, values, Some(chunk_offsets), &mut ctx).unwrap(); - let sliced = patches.slice(1000..2200).unwrap().unwrap(); + let sliced = patches.slice(1000..2200, &mut ctx).unwrap().unwrap(); assert!(sliced.chunk_offsets.is_some()); assert_eq!(sliced.offset(), 1000); @@ -2444,9 +2464,10 @@ mod test { let indices = buffer![100u64, 500, 1200, 1300, 1500, 1800, 2100, 2500].into_array(); let values = buffer![10i32, 20, 30, 35, 40, 45, 50, 60].into_array(); let chunk_offsets = buffer![0u64, 2, 6].into_array(); - let patches = Patches::new(3000, 0, indices, values, Some(chunk_offsets)).unwrap(); + let patches = + Patches::new(3000, 0, indices, values, Some(chunk_offsets), &mut ctx).unwrap(); - let sliced = patches.slice(1300..2200).unwrap().unwrap(); + let sliced = patches.slice(1300..2200, &mut ctx).unwrap().unwrap(); assert!(sliced.chunk_offsets.is_some()); assert_eq!(sliced.offset(), 1300); @@ -2480,12 +2501,14 @@ mod test { #[test] fn test_chunk_offsets_slice_empty_result() { + let mut ctx = array_session().create_execution_ctx(); let indices = buffer![100u64, 200, 3000, 3100].into_array(); let values = buffer![10i32, 20, 30, 40].into_array(); let chunk_offsets = buffer![0u64, 2, 2, 3].into_array(); - let patches = Patches::new(4000, 0, indices, values, Some(chunk_offsets)).unwrap(); + let patches = + Patches::new(4000, 0, indices, values, Some(chunk_offsets), &mut ctx).unwrap(); - let sliced = patches.slice(1000..2000).unwrap(); + let sliced = patches.slice(1000..2000, &mut ctx).unwrap(); assert!(sliced.is_none()); } @@ -2495,9 +2518,10 @@ mod test { let indices = buffer![100u64, 1200, 1300, 2500].into_array(); let values = buffer![10i32, 20, 30, 40].into_array(); let chunk_offsets = buffer![0u64, 1, 3].into_array(); - let patches = Patches::new(3000, 0, indices, values, Some(chunk_offsets)).unwrap(); + let patches = + Patches::new(3000, 0, indices, values, Some(chunk_offsets), &mut ctx).unwrap(); - let sliced = patches.slice(1100..1250).unwrap().unwrap(); + let sliced = patches.slice(1100..1250, &mut ctx).unwrap().unwrap(); assert_eq!(sliced.num_patches(), 1); assert_eq!(sliced.offset(), 1100); @@ -2526,9 +2550,10 @@ mod test { let indices = buffer![100u64, 200, 1100, 1200, 2100, 2200].into_array(); let values = buffer![10i32, 20, 30, 40, 50, 60].into_array(); let chunk_offsets = buffer![0u64, 2, 4].into_array(); - let patches = Patches::new(3000, 0, indices, values, Some(chunk_offsets)).unwrap(); + let patches = + Patches::new(3000, 0, indices, values, Some(chunk_offsets), &mut ctx).unwrap(); - let sliced = patches.slice(150..2150).unwrap().unwrap(); + let sliced = patches.slice(150..2150, &mut ctx).unwrap().unwrap(); assert_eq!(sliced.num_patches(), 4); assert_eq!(sliced.offset(), 150); @@ -2562,7 +2587,8 @@ mod test { let indices = buffer![1023u64, 1024, 1025, 2047, 2048].into_array(); let values = buffer![10i32, 20, 30, 40, 50].into_array(); let chunk_offsets = buffer![0u64, 1, 4].into_array(); - let patches = Patches::new(3000, 0, indices, values, Some(chunk_offsets)).unwrap(); + let patches = + Patches::new(3000, 0, indices, values, Some(chunk_offsets), &mut ctx).unwrap(); assert_eq!( patches.search_index(1023, &mut ctx).unwrap(), @@ -2601,10 +2627,11 @@ mod test { let indices = buffer![0u64, 1, 1023, 1024, 2047, 2048].into_array(); let values = buffer![10i32, 20, 30, 40, 50, 60].into_array(); let chunk_offsets = buffer![0u64, 3, 5].into_array(); - let patches = Patches::new(3000, 0, indices, values, Some(chunk_offsets)).unwrap(); + let patches = + Patches::new(3000, 0, indices, values, Some(chunk_offsets), &mut ctx).unwrap(); // Slice at the very beginning - let sliced = patches.slice(0..10).unwrap().unwrap(); + let sliced = patches.slice(0..10, &mut ctx).unwrap().unwrap(); assert_eq!(sliced.num_patches(), 2); assert_eq!( sliced.search_index(0, &mut ctx).unwrap(), @@ -2622,7 +2649,7 @@ mod test { ); // Slice at the very end - let sliced = patches.slice(2040..3000).unwrap().unwrap(); + let sliced = patches.slice(2040..3000, &mut ctx).unwrap().unwrap(); assert_eq!(sliced.num_patches(), 2); // patches at 2047 and 2048 assert_eq!( sliced.search_index(7, &mut ctx).unwrap(), @@ -2645,12 +2672,13 @@ mod test { let indices = buffer![100u64, 200, 300, 400, 500, 600].into_array(); let values = buffer![10i32, 20, 30, 40, 50, 60].into_array(); let chunk_offsets = buffer![0u64].into_array(); - let patches = Patches::new(1000, 0, indices, values, Some(chunk_offsets)).unwrap(); + let patches = + Patches::new(1000, 0, indices, values, Some(chunk_offsets), &mut ctx).unwrap(); - let sliced1 = patches.slice(150..550).unwrap().unwrap(); + let sliced1 = patches.slice(150..550, &mut ctx).unwrap().unwrap(); assert_eq!(sliced1.num_patches(), 4); // 200, 300, 400, 500 - let sliced2 = sliced1.slice(100..250).unwrap().unwrap(); + let sliced2 = sliced1.slice(100..250, &mut ctx).unwrap().unwrap(); assert_eq!(sliced2.num_patches(), 1); // 300 assert_eq!(sliced2.offset(), 250); @@ -2681,11 +2709,12 @@ mod test { let indices = buffer![0u64, 1024].into_array(); let values = buffer![1i32, 2].into_array(); let chunk_offsets = buffer![0u64, 1].into_array(); - let patches = Patches::new(2048, 0, indices, values, Some(chunk_offsets)).unwrap(); + let patches = + Patches::new(2048, 0, indices, values, Some(chunk_offsets), &mut ctx).unwrap(); // Drop chunk 0, then re-slice the result. - let dropped_first = patches.slice(1024..2048).unwrap().unwrap(); - let resliced = dropped_first.slice(0..1024).unwrap().unwrap(); + let dropped_first = patches.slice(1024..2048, &mut ctx).unwrap().unwrap(); + let resliced = dropped_first.slice(0..1024, &mut ctx).unwrap().unwrap(); assert_eq!(resliced.num_patches(), 1); assert_eq!(resliced.array_len(), 1024); assert_eq!(resliced.offset(), 1024); @@ -2706,7 +2735,8 @@ mod test { let chunk_offsets = buffer![0u64].into_array(); let indices = buffer![1023u64].into_array(); let values = buffer![42i32].into_array(); - let patches = Patches::new(1024, 0, indices, values, Some(chunk_offsets)).unwrap(); + let patches = + Patches::new(1024, 0, indices, values, Some(chunk_offsets), &mut ctx).unwrap(); assert_eq!( patches.search_index(2048, &mut ctx).unwrap(), SearchResult::NotFound(1) diff --git a/vortex-array/src/validity.rs b/vortex-array/src/validity.rs index 83b3fefe08b..71b36a522b1 100644 --- a/vortex-array/src/validity.rs +++ b/vortex-array/src/validity.rs @@ -382,6 +382,7 @@ impl Validity { patch_values.into_array(), // TODO(0ax1): chunk offsets None, + ctx, )?; Ok(Self::Array(source.patch(&patches, ctx)?.into_array())) diff --git a/vortex-btrblocks/src/schemes/float/sparse.rs b/vortex-btrblocks/src/schemes/float/sparse.rs index 3d9c25b18e4..b23ca87a321 100644 --- a/vortex-btrblocks/src/schemes/float/sparse.rs +++ b/vortex-btrblocks/src/schemes/float/sparse.rs @@ -112,6 +112,7 @@ impl Scheme for NullDominatedSparseScheme { sparse.patches().values().clone(), sparse.len(), sparse.fill_scalar().clone(), + exec_ctx, ) .map(|a| a.into_array()) } else { diff --git a/vortex-btrblocks/src/schemes/integer/sparse.rs b/vortex-btrblocks/src/schemes/integer/sparse.rs index 429ff5c1a31..438f1084cda 100644 --- a/vortex-btrblocks/src/schemes/integer/sparse.rs +++ b/vortex-btrblocks/src/schemes/integer/sparse.rs @@ -204,6 +204,7 @@ impl Scheme for SparseScheme { compressed_values, sparse.len(), sparse.fill_scalar().clone(), + exec_ctx, ) .map(|a| a.into_array()) } else { diff --git a/vortex-btrblocks/src/schemes/patches.rs b/vortex-btrblocks/src/schemes/patches.rs index 69ca8450f12..0ed60a9e112 100644 --- a/vortex-btrblocks/src/schemes/patches.rs +++ b/vortex-btrblocks/src/schemes/patches.rs @@ -51,5 +51,6 @@ pub fn compress_patches(patches: Patches, ctx: &mut ExecutionCtx) -> VortexResul indices, values, chunk_offsets, + ctx, ) } diff --git a/vortex-btrblocks/src/schemes/string/sparse.rs b/vortex-btrblocks/src/schemes/string/sparse.rs index 8620c366f77..604018f98ba 100644 --- a/vortex-btrblocks/src/schemes/string/sparse.rs +++ b/vortex-btrblocks/src/schemes/string/sparse.rs @@ -120,6 +120,7 @@ impl Scheme for NullDominatedSparseScheme { sparse.patches().values().clone(), sparse.len(), sparse.fill_scalar().clone(), + exec_ctx, ) .map(|a| a.into_array()) } else { diff --git a/vortex-cuda/src/hybrid_dispatch/mod.rs b/vortex-cuda/src/hybrid_dispatch/mod.rs index ab078198a79..eecd8adb12c 100644 --- a/vortex-cuda/src/hybrid_dispatch/mod.rs +++ b/vortex-cuda/src/hybrid_dispatch/mod.rs @@ -278,6 +278,7 @@ mod tests { PrimitiveArray::new(buffer![0u32, 1024u32], NN).into_array(), PrimitiveArray::new(buffer![99.9f32, 88.8f32], NN).into_array(), None, + ctx.execution_ctx(), ) .unwrap(); let arr = ALP::try_new(encoded, Exponents { e: 0, f: 2 }, Some(patches))?; diff --git a/vortex-cuda/src/kernel/encodings/alp.rs b/vortex-cuda/src/kernel/encodings/alp.rs index 3595d265d89..3c364646af2 100644 --- a/vortex-cuda/src/kernel/encodings/alp.rs +++ b/vortex-cuda/src/kernel/encodings/alp.rs @@ -188,6 +188,7 @@ mod tests { PrimitiveArray::new(buffer![0u32, 4u32], Validity::NonNullable).into_array(), PrimitiveArray::new(buffer![0.0f32, 999f32], Validity::NonNullable).into_array(), Some(PrimitiveArray::new(buffer![0u32], Validity::NonNullable).into_array()), + cuda_ctx.execution_ctx(), )?; let alp_array = ALP::try_new( @@ -427,6 +428,7 @@ mod tests { PrimitiveArray::new(patch_indices, Validity::NonNullable).into_array(), PrimitiveArray::new(patch_values, Validity::NonNullable).into_array(), Some(PrimitiveArray::new(buffer![0u32], Validity::NonNullable).into_array()), + cuda_ctx.execution_ctx(), )?; let alp_array = ALP::try_new( diff --git a/vortex-cuda/src/kernel/patches/types.rs b/vortex-cuda/src/kernel/patches/types.rs index 18ae7a0249d..7ed4308eb1a 100644 --- a/vortex-cuda/src/kernel/patches/types.rs +++ b/vortex-cuda/src/kernel/patches/types.rs @@ -239,6 +239,7 @@ mod tests { #[test] fn test_patches_with_chunk_offsets() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); // Test creating patches with pre-computed chunk_offsets let indices = PrimitiveArray::from_iter([0u32, 500, 1024, 2000]); let values = PrimitiveArray::from_iter([10u32, 20, 30, 40]); @@ -250,11 +251,12 @@ mod tests { indices.into_array(), values.into_array(), Some(chunk_offsets.into_array()), + &mut ctx, )?; assert!(patches.chunk_offsets().is_some()); - assert_eq!(patches.chunk_offset_at(0)?, 0); - assert_eq!(patches.chunk_offset_at(1)?, 2); - assert_eq!(patches.chunk_offset_at(2)?, 4); + assert_eq!(patches.chunk_offset_at(0, &mut ctx)?, 0); + assert_eq!(patches.chunk_offset_at(1, &mut ctx)?, 2); + assert_eq!(patches.chunk_offset_at(2, &mut ctx)?, 4); Ok(()) } @@ -284,6 +286,7 @@ mod tests { indices.into_array(), values.into_array(), Some(chunk_offsets.into_array()), + cuda_ctx.execution_ctx(), )?; let mut device_patches = load_device_patches(&patches, &mut cuda_ctx).await?; @@ -328,6 +331,7 @@ mod tests { indices.into_array(), values.into_array(), Some(chunk_offsets), + cuda_ctx.execution_ctx(), )?; let mut device_patches = load_device_patches(&patches, &mut cuda_ctx).await?; diff --git a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/patched.rs b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/patched.rs index 62b9ea29115..ea363c09c37 100644 --- a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/patched.rs +++ b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/patched.rs @@ -36,6 +36,7 @@ impl FlatLayoutFixture for PatchedFixture { PrimitiveArray::from_iter([0u32, 1024, 1025, 1026]).into_array(), PrimitiveArray::from_option_iter([Some(1u64), Some(2), Some(3), Some(4)]).into_array(), None, + ctx, )?; Ok(Patched::from_array_and_patches(array, &patches, ctx)?.into_array()) diff --git a/vortex/benches/take_union.rs b/vortex/benches/take_union.rs index ff0f1e2a25a..b7088601d13 100644 --- a/vortex/benches/take_union.rs +++ b/vortex/benches/take_union.rs @@ -167,6 +167,7 @@ fn dense_children_union(rng: &mut StdRng) -> ArrayRef { /// The dominant integer child stays canonical. At 98% density a sparse child would store a patch /// index for nearly every row and buy nothing. fn compressed_children_union(rng: &mut StdRng) -> ArrayRef { + let mut ctx = SESSION.create_execution_ctx(); union_array( (0..VARIANT_NAMES.len()) .map(|variant| { @@ -178,7 +179,7 @@ fn compressed_children_union(rng: &mut StdRng) -> ArrayRef { let values = variant_child(variant, indices.len(), rng); let fill = Scalar::default_value(values.dtype()); - Sparse::try_new(indices.into_array(), values, ARRAY_SIZE, fill) + Sparse::try_new(indices.into_array(), values, ARRAY_SIZE, fill, &mut ctx) .unwrap() .into_array() })