From 1c5cc675b478a938ba51f2537126b0941d5cb047 Mon Sep 17 00:00:00 2001 From: Mikhail Kot Date: Thu, 23 Jul 2026 14:28:26 +0100 Subject: [PATCH] Use array_slots macro instead of manual slot indices Signed-off-by: Mikhail Kot --- encodings/alp/src/alp_rd/array.rs | 98 +++++++------------ encodings/alp/src/alp_rd/compute/cast.rs | 1 + encodings/alp/src/alp_rd/compute/filter.rs | 1 + encodings/alp/src/alp_rd/compute/mask.rs | 1 + encodings/alp/src/alp_rd/compute/take.rs | 1 + encodings/alp/src/alp_rd/ops.rs | 1 + encodings/alp/src/alp_rd/slice.rs | 1 + .../src/delta/array/delta_decompress.rs | 1 + encodings/fastlanes/src/delta/array/mod.rs | 29 ++---- encodings/fastlanes/src/delta/compute/cast.rs | 1 + encodings/fastlanes/src/delta/mod.rs | 3 + encodings/fastlanes/src/delta/vtable/mod.rs | 21 ++-- encodings/fastlanes/src/delta/vtable/slice.rs | 1 + .../fastlanes/src/delta/vtable/validity.rs | 1 + encodings/fastlanes/src/rle/array/mod.rs | 59 ++++------- .../fastlanes/src/rle/array/rle_compress.rs | 1 + .../fastlanes/src/rle/array/rle_decompress.rs | 1 + encodings/fastlanes/src/rle/compute/cast.rs | 1 + encodings/fastlanes/src/rle/kernel.rs | 1 + encodings/fastlanes/src/rle/mod.rs | 3 + encodings/fastlanes/src/rle/vtable/mod.rs | 52 +++++----- .../fastlanes/src/rle/vtable/operations.rs | 1 + .../fastlanes/src/rle/vtable/validity.rs | 1 + encodings/fsst/src/array.rs | 87 +++++++--------- encodings/fsst/src/canonical.rs | 1 + encodings/fsst/src/compress.rs | 4 +- encodings/fsst/src/compute/byte_length.rs | 2 +- encodings/fsst/src/compute/cast.rs | 3 +- encodings/fsst/src/compute/compare.rs | 1 + encodings/fsst/src/compute/filter.rs | 1 + encodings/fsst/src/compute/like.rs | 2 +- encodings/fsst/src/compute/mod.rs | 1 + encodings/fsst/src/slice.rs | 1 + encodings/onpair/src/compute/filter.rs | 2 +- encodings/parquet-variant/src/array.rs | 98 ++++++++----------- encodings/parquet-variant/src/arrow.rs | 15 +-- .../src/compute/allnondistinct.rs | 6 +- .../src/json_to_variant_tests.rs | 6 +- encodings/parquet-variant/src/kernel.rs | 45 ++++----- encodings/parquet-variant/src/lib.rs | 3 + encodings/parquet-variant/src/operations.rs | 8 +- encodings/parquet-variant/src/vtable.rs | 64 ++++++------ encodings/runend/src/array.rs | 54 ++++------ encodings/runend/src/compute/cast.rs | 1 + encodings/runend/src/compute/compare.rs | 1 + encodings/runend/src/compute/fill_null.rs | 1 + encodings/runend/src/compute/filter.rs | 1 + encodings/runend/src/compute/is_constant.rs | 2 +- encodings/runend/src/compute/is_sorted.rs | 2 +- encodings/runend/src/compute/min_max.rs | 2 +- encodings/runend/src/compute/take.rs | 1 + encodings/runend/src/compute/take_from.rs | 3 +- encodings/runend/src/kernel.rs | 1 + encodings/runend/src/ops.rs | 1 + encodings/runend/src/rules.rs | 1 + encodings/sparse/src/canonical.rs | 2 + fuzz/src/array/mask.rs | 3 +- fuzz/src/array/slice.rs | 3 +- vortex-array/benches/listview_rebuild.rs | 2 +- vortex-array/benches/take_fsl.rs | 2 +- .../src/aggregate_fn/accumulator_grouped.rs | 3 +- .../fns/all_non_distinct/fixed_size_list.rs | 2 +- .../aggregate_fn/fns/all_non_distinct/list.rs | 2 + .../src/aggregate_fn/fns/is_constant/list.rs | 1 + .../fixed_size_list.rs | 2 +- .../uncompressed_size_in_bytes/list_view.rs | 2 +- .../src/arrays/chunked/vtable/canonical.rs | 4 +- .../src/arrays/constant/vtable/canonical.rs | 3 +- .../arrays/filter/execute/fixed_size_list.rs | 1 + .../src/arrays/filter/execute/listview.rs | 3 +- .../src/arrays/fixed_size_list/array.rs | 53 +++++----- .../arrays/fixed_size_list/compute/cast.rs | 1 + .../arrays/fixed_size_list/compute/mask.rs | 1 + .../arrays/fixed_size_list/compute/slice.rs | 1 + .../arrays/fixed_size_list/compute/take.rs | 1 + .../src/arrays/fixed_size_list/mod.rs | 3 + .../src/arrays/fixed_size_list/tests/basic.rs | 1 + .../fixed_size_list/tests/degenerate.rs | 1 + .../src/arrays/fixed_size_list/vtable/mod.rs | 17 ++-- vortex-array/src/arrays/list/array.rs | 60 ++++-------- vortex-array/src/arrays/list/compute/cast.rs | 2 +- .../src/arrays/list/compute/filter.rs | 1 + vortex-array/src/arrays/list/compute/mask.rs | 2 +- vortex-array/src/arrays/list/compute/slice.rs | 2 +- vortex-array/src/arrays/list/compute/take.rs | 1 + vortex-array/src/arrays/list/mod.rs | 3 + vortex-array/src/arrays/list/vtable/mod.rs | 25 ++--- vortex-array/src/arrays/listview/array.rs | 93 +++++++----------- .../src/arrays/listview/compute/cast.rs | 2 +- .../src/arrays/listview/compute/mask.rs | 2 +- .../src/arrays/listview/compute/rules.rs | 2 +- .../src/arrays/listview/compute/slice.rs | 2 +- .../src/arrays/listview/compute/take.rs | 2 +- .../src/arrays/listview/compute/zip.rs | 4 +- .../src/arrays/listview/conversion.rs | 6 +- vortex-array/src/arrays/listview/mod.rs | 3 + vortex-array/src/arrays/listview/rebuild.rs | 2 + .../src/arrays/listview/tests/density.rs | 1 + .../src/arrays/listview/tests/filter.rs | 1 + .../src/arrays/listview/tests/nullability.rs | 1 + .../src/arrays/listview/tests/operations.rs | 1 + .../src/arrays/listview/tests/take.rs | 1 + .../src/arrays/listview/vtable/mod.rs | 32 +++--- vortex-array/src/arrays/masked/execute.rs | 3 +- vortex-array/src/arrays/varbin/array.rs | 43 +++----- vortex-array/src/arrays/varbin/builder.rs | 2 +- .../src/arrays/varbin/compute/cast.rs | 2 +- .../src/arrays/varbin/compute/compare.rs | 2 +- .../src/arrays/varbin/compute/filter.rs | 1 + .../src/arrays/varbin/compute/mask.rs | 2 +- .../src/arrays/varbin/compute/slice.rs | 2 +- .../src/arrays/varbin/compute/take.rs | 1 + vortex-array/src/arrays/varbin/mod.rs | 3 + vortex-array/src/arrays/varbin/vtable/mod.rs | 18 ++-- vortex-array/src/arrays/variant/mod.rs | 40 +++----- vortex-array/src/arrays/variant/vtable/mod.rs | 36 +++---- vortex-array/src/builders/dict/bytes.rs | 2 +- vortex-array/src/builders/fixed_size_list.rs | 3 +- vortex-array/src/builders/list.rs | 7 +- vortex-array/src/builders/listview.rs | 5 +- .../scalar_fn/fns/binary/compare/nested.rs | 4 +- .../src/scalar_fn/fns/list_contains/mod.rs | 2 +- vortex-array/src/scalar_fn/fns/list_length.rs | 2 + vortex-arrow/src/convert.rs | 5 +- vortex-arrow/src/executor/byte.rs | 2 +- vortex-arrow/src/executor/fixed_size_list.rs | 1 + vortex-arrow/src/executor/list.rs | 4 +- vortex-arrow/src/executor/list_view.rs | 1 + vortex-arrow/src/executor/mod.rs | 4 +- vortex-arrow/src/executor/run_end.rs | 1 + vortex-btrblocks/src/schemes/integer/rle.rs | 1 + vortex-btrblocks/src/schemes/string/fsst.rs | 3 +- vortex-compressor/src/compressor/cascade.rs | 3 +- .../src/compressor/structural.rs | 3 +- vortex-cuda/benches/fsst_cuda.rs | 1 + vortex-cuda/src/arrow/canonical.rs | 2 +- vortex-cuda/src/arrow/mod.rs | 6 +- .../src/dynamic_dispatch/plan_builder.rs | 1 + vortex-cuda/src/kernel/encodings/fsst.rs | 3 +- vortex-cuda/src/kernel/encodings/runend.rs | 1 + vortex-duckdb/src/convert/vector.rs | 1 + vortex-duckdb/src/exporter/list_view.rs | 1 + vortex-duckdb/src/exporter/rle.rs | 1 + vortex-duckdb/src/exporter/run_end.rs | 1 + vortex-geo/src/extension/mod.rs | 4 +- vortex-row/src/codec.rs | 1 + vortex-row/src/tests.rs | 1 + vortex-tensor/src/scalar_fns/l2_denorm.rs | 3 +- vortex-tensor/src/utils.rs | 2 +- .../common_encoding_tree_throughput.rs | 5 +- 150 files changed, 610 insertions(+), 699 deletions(-) diff --git a/encodings/alp/src/alp_rd/array.rs b/encodings/alp/src/alp_rd/array.rs index dcaeb0fde79..3f81248332e 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::array_slots; use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; use vortex_array::buffer::BufferHandle; @@ -105,12 +106,15 @@ impl VTable for ALPRD { len: usize, slots: &[Option], ) -> VortexResult<()> { + let slots_view = ALPRDSlotsView::from_slots(slots); + let patches = + PatchesData::patches_from_slots(data.patches_data.as_ref(), len, slots, LP_PATCH_SLOTS); validate_parts( dtype, len, - left_parts_from_slots(slots), - right_parts_from_slots(slots), - patches_from_slots(slots, data.patches_data.as_ref(), len).as_ref(), + slots_view.left_parts, + slots_view.right_parts, + patches.as_ref(), ) } @@ -230,17 +234,18 @@ impl VTable for ALPRD { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + ALPRDSlots::NAMES[idx].to_string() } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - let array = require_child!(array, array.left_parts(), 0 => Primitive); - let array = require_child!(array, array.right_parts(), 1 => Primitive); + let array = require_child!(array, array.left_parts(), ALPRDSlots::LEFT_PARTS => Primitive); + let array = + require_child!(array, array.right_parts(), ALPRDSlots::RIGHT_PARTS => Primitive); require_patches!( array, - LP_PATCH_INDICES_SLOT, - LP_PATCH_VALUES_SLOT, - LP_PATCH_CHUNK_OFFSETS_SLOT + ALPRDSlots::PATCH_INDICES, + ALPRDSlots::PATCH_VALUES, + ALPRDSlots::PATCH_CHUNK_OFFSETS ); let dtype = array.dtype().clone(); @@ -307,29 +312,25 @@ impl VTable for ALPRD { } } -/// The left (most significant) parts of the real-double encoded values. -pub(super) const LEFT_PARTS_SLOT: usize = 0; -/// The right (least significant) parts of the real-double encoded values. -pub(super) const RIGHT_PARTS_SLOT: usize = 1; -/// The indices of left-parts exception values that could not be dictionary-encoded. -pub(super) const LP_PATCH_SLOTS: PatchSlotIndices = PatchSlotIndices { - indices: LP_PATCH_INDICES_SLOT, - values: LP_PATCH_VALUES_SLOT, - chunk_offsets: LP_PATCH_CHUNK_OFFSETS_SLOT, +#[array_slots(ALPRD)] +pub struct ALPRDSlots { + /// The left (most significant) parts of the real-double encoded values. + pub left_parts: ArrayRef, + /// The right (least significant) parts of the real-double encoded values. + pub right_parts: ArrayRef, + /// The indices of left-parts exception values that could not be dictionary-encoded. + pub patch_indices: Option, + /// The exception values for left-parts that could not be dictionary-encoded. + pub patch_values: Option, + /// Chunk offsets for the left-parts patch indices/values. + pub patch_chunk_offsets: Option, +} + +const LP_PATCH_SLOTS: PatchSlotIndices = PatchSlotIndices { + indices: ALPRDSlots::PATCH_INDICES, + values: ALPRDSlots::PATCH_VALUES, + chunk_offsets: ALPRDSlots::PATCH_CHUNK_OFFSETS, }; -pub(super) const LP_PATCH_INDICES_SLOT: usize = 2; -/// The exception values for left-parts that could not be dictionary-encoded. -pub(super) const LP_PATCH_VALUES_SLOT: usize = 3; -/// Chunk offsets for the left-parts patch indices/values. -pub(super) const LP_PATCH_CHUNK_OFFSETS_SLOT: usize = 4; -pub(super) const NUM_SLOTS: usize = 5; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = [ - "left_parts", - "right_parts", - "patch_indices", - "patch_values", - "patch_chunk_offsets", -]; #[derive(Clone, Debug)] pub struct ALPRDData { @@ -451,26 +452,6 @@ impl ALPRDData { } } -fn left_parts_from_slots(slots: &[Option]) -> &ArrayRef { - slots[LEFT_PARTS_SLOT] - .as_ref() - .vortex_expect("ALPRDArray left_parts slot") -} - -fn right_parts_from_slots(slots: &[Option]) -> &ArrayRef { - slots[RIGHT_PARTS_SLOT] - .as_ref() - .vortex_expect("ALPRDArray right_parts slot") -} - -fn patches_from_slots( - slots: &[Option], - patches_data: Option<&PatchesData>, - len: usize, -) -> Option { - PatchesData::patches_from_slots(patches_data, len, slots, LP_PATCH_SLOTS) -} - #[allow(clippy::disallowed_methods)] fn validate_parts( dtype: &DType, @@ -539,24 +520,17 @@ fn validate_parts( Ok(()) } -pub trait ALPRDArrayExt: TypedArrayRef { - fn left_parts(&self) -> &ArrayRef { - left_parts_from_slots(self.as_ref().slots()) - } - - fn right_parts(&self) -> &ArrayRef { - right_parts_from_slots(self.as_ref().slots()) - } - +pub trait ALPRDArrayExt: ALPRDArraySlotsExt { fn right_bit_width(&self) -> u8 { ALPRDData::right_bit_width(self) } fn left_parts_patches(&self) -> Option { - patches_from_slots( - self.as_ref().slots(), + PatchesData::patches_from_slots( self.patches_data.as_ref(), self.as_ref().len(), + self.as_ref().slots(), + LP_PATCH_SLOTS, ) } diff --git a/encodings/alp/src/alp_rd/compute/cast.rs b/encodings/alp/src/alp_rd/compute/cast.rs index 9a08f00e298..be244048e85 100644 --- a/encodings/alp/src/alp_rd/compute/cast.rs +++ b/encodings/alp/src/alp_rd/compute/cast.rs @@ -10,6 +10,7 @@ use vortex_array::scalar_fn::fns::cast::CastReduce; use vortex_error::VortexResult; use crate::ALPRDArrayExt; +use crate::ALPRDArraySlotsExt; use crate::alp_rd::ALPRD; impl CastReduce for ALPRD { diff --git a/encodings/alp/src/alp_rd/compute/filter.rs b/encodings/alp/src/alp_rd/compute/filter.rs index 69c26c53e3b..9f353e31a21 100644 --- a/encodings/alp/src/alp_rd/compute/filter.rs +++ b/encodings/alp/src/alp_rd/compute/filter.rs @@ -11,6 +11,7 @@ use vortex_mask::Mask; use crate::ALPRD; use crate::ALPRDArrayExt; +use crate::ALPRDArraySlotsExt; impl FilterKernel for ALPRD { fn filter( diff --git a/encodings/alp/src/alp_rd/compute/mask.rs b/encodings/alp/src/alp_rd/compute/mask.rs index 25bdc8f9ab4..7e2906508e2 100644 --- a/encodings/alp/src/alp_rd/compute/mask.rs +++ b/encodings/alp/src/alp_rd/compute/mask.rs @@ -12,6 +12,7 @@ use vortex_error::VortexResult; use crate::ALPRD; use crate::ALPRDArrayExt; +use crate::ALPRDArraySlotsExt; impl MaskReduce for ALPRD { #[allow(clippy::disallowed_methods)] diff --git a/encodings/alp/src/alp_rd/compute/take.rs b/encodings/alp/src/alp_rd/compute/take.rs index 18d5e2b5ad4..3f5b5d9d53f 100644 --- a/encodings/alp/src/alp_rd/compute/take.rs +++ b/encodings/alp/src/alp_rd/compute/take.rs @@ -12,6 +12,7 @@ use vortex_error::VortexResult; use crate::ALPRD; use crate::ALPRDArrayExt; +use crate::ALPRDArraySlotsExt; impl TakeExecute for ALPRD { fn take( diff --git a/encodings/alp/src/alp_rd/ops.rs b/encodings/alp/src/alp_rd/ops.rs index 0fc0dc40fa1..344ee69f532 100644 --- a/encodings/alp/src/alp_rd/ops.rs +++ b/encodings/alp/src/alp_rd/ops.rs @@ -11,6 +11,7 @@ use vortex_error::VortexResult; use crate::ALPRD; use crate::ALPRDArrayExt; +use crate::ALPRDArraySlotsExt; impl OperationsVTable for ALPRD { fn scalar_at( diff --git a/encodings/alp/src/alp_rd/slice.rs b/encodings/alp/src/alp_rd/slice.rs index c36e138ff84..5f81bdc7b5d 100644 --- a/encodings/alp/src/alp_rd/slice.rs +++ b/encodings/alp/src/alp_rd/slice.rs @@ -11,6 +11,7 @@ use vortex_array::arrays::slice::SliceKernel; use vortex_error::VortexResult; use crate::ALPRDArrayExt; +use crate::ALPRDArraySlotsExt; use crate::alp_rd::ALPRD; impl SliceKernel for ALPRD { diff --git a/encodings/fastlanes/src/delta/array/delta_decompress.rs b/encodings/fastlanes/src/delta/array/delta_decompress.rs index 7dcbeb44950..f5d3ed172b3 100644 --- a/encodings/fastlanes/src/delta/array/delta_decompress.rs +++ b/encodings/fastlanes/src/delta/array/delta_decompress.rs @@ -20,6 +20,7 @@ use vortex_error::VortexResult; use crate::DeltaArray; use crate::bit_transpose::untranspose_validity; use crate::delta::array::DeltaArrayExt; +use crate::delta::array::DeltaArraySlotsExt; pub fn delta_decompress( array: &DeltaArray, diff --git a/encodings/fastlanes/src/delta/array/mod.rs b/encodings/fastlanes/src/delta/array/mod.rs index 7754bb37a59..15fb1fbcfc9 100644 --- a/encodings/fastlanes/src/delta/array/mod.rs +++ b/encodings/fastlanes/src/delta/array/mod.rs @@ -7,21 +7,22 @@ use std::fmt::Formatter; use fastlanes::FastLanes; use vortex_array::ArrayRef; use vortex_array::TypedArrayRef; +use vortex_array::array_slots; use vortex_array::dtype::PType; use vortex_array::match_each_unsigned_integer_ptype; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; pub mod delta_compress; pub mod delta_decompress; -/// The base values for each block of deltas. -pub(super) const BASES_SLOT: usize = 0; -/// The delta-encoded values relative to the base values. -pub(super) const DELTAS_SLOT: usize = 1; -pub(super) const NUM_SLOTS: usize = 2; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["bases", "deltas"]; +#[array_slots(crate::Delta)] +pub struct DeltaSlots { + /// The base values for each block of deltas. + pub bases: ArrayRef, + /// The delta-encoded values relative to the base values. + pub deltas: ArrayRef, +} /// A FastLanes-style delta-encoded array of primitive values. /// @@ -88,19 +89,7 @@ impl Display for DeltaData { } } -pub trait DeltaArrayExt: TypedArrayRef { - fn bases(&self) -> &ArrayRef { - self.as_ref().slots()[BASES_SLOT] - .as_ref() - .vortex_expect("DeltaArray bases slot") - } - - fn deltas(&self) -> &ArrayRef { - self.as_ref().slots()[DELTAS_SLOT] - .as_ref() - .vortex_expect("DeltaArray deltas slot") - } - +pub trait DeltaArrayExt: DeltaArraySlotsExt { fn offset(&self) -> usize { self.offset } diff --git a/encodings/fastlanes/src/delta/compute/cast.rs b/encodings/fastlanes/src/delta/compute/cast.rs index c6de602bb0b..5fe84f62b62 100644 --- a/encodings/fastlanes/src/delta/compute/cast.rs +++ b/encodings/fastlanes/src/delta/compute/cast.rs @@ -12,6 +12,7 @@ use vortex_error::VortexResult; use crate::delta::Delta; use crate::delta::array::DeltaArrayExt; +use crate::delta::array::DeltaArraySlotsExt; impl CastReduce for Delta { fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult> { diff --git a/encodings/fastlanes/src/delta/mod.rs b/encodings/fastlanes/src/delta/mod.rs index 52ea9b33574..7a8a3729048 100644 --- a/encodings/fastlanes/src/delta/mod.rs +++ b/encodings/fastlanes/src/delta/mod.rs @@ -2,7 +2,10 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors mod array; +pub use array::DeltaArraySlotsExt; pub use array::DeltaData; +pub use array::DeltaSlots; +pub use array::DeltaSlotsView; pub use array::delta_compress::delta_compress; mod compute; diff --git a/encodings/fastlanes/src/delta/vtable/mod.rs b/encodings/fastlanes/src/delta/vtable/mod.rs index 6dcfedc0bc9..f2ed1ad0019 100644 --- a/encodings/fastlanes/src/delta/vtable/mod.rs +++ b/encodings/fastlanes/src/delta/vtable/mod.rs @@ -21,9 +21,7 @@ use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::PType; use vortex_array::serde::ArrayChildren; -use vortex_array::smallvec::smallvec; use vortex_array::vtable::VTable; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_err; @@ -32,10 +30,10 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::DeltaData; -use crate::delta::array::BASES_SLOT; -use crate::delta::array::DELTAS_SLOT; use crate::delta::array::DeltaArrayExt; -use crate::delta::array::SLOT_NAMES; +use crate::delta::array::DeltaArraySlotsExt; +use crate::delta::array::DeltaSlots; +use crate::delta::array::DeltaSlotsView; use crate::delta::array::delta_decompress::delta_decompress; use crate::delta::array::lane_count; use crate::delta_compress; @@ -87,12 +85,7 @@ impl VTable for Delta { len: usize, slots: &[Option], ) -> VortexResult<()> { - let bases = slots[BASES_SLOT] - .as_ref() - .vortex_expect("DeltaArray bases slot"); - let deltas = slots[DELTAS_SLOT] - .as_ref() - .vortex_expect("DeltaArray deltas slot"); + let DeltaSlotsView { bases, deltas } = DeltaSlotsView::from_slots(slots); validate_parts(bases, deltas, data.offset, dtype, len) } @@ -125,7 +118,7 @@ impl VTable for Delta { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + DeltaSlots::NAMES[idx].to_string() } fn serialize( @@ -175,7 +168,7 @@ impl VTable for Delta { let deltas = children.get(1, dtype, deltas_len)?; let data = DeltaData::try_new(metadata.offset as usize)?; - let slots = smallvec![Some(bases), Some(deltas)]; + let slots = DeltaSlots { bases, deltas }.into_slots(); Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) } @@ -198,7 +191,7 @@ impl Delta { ) -> VortexResult { let dtype = bases.dtype().with_nullability(deltas.dtype().nullability()); let data = DeltaData::try_new(offset)?; - let slots = smallvec![Some(bases), Some(deltas)]; + let slots = DeltaSlots { bases, deltas }.into_slots(); Array::try_from_parts(ArrayParts::new(Delta, dtype, len, data).with_slots(slots)) } diff --git a/encodings/fastlanes/src/delta/vtable/slice.rs b/encodings/fastlanes/src/delta/vtable/slice.rs index 17b088dac92..be028b2f2ca 100644 --- a/encodings/fastlanes/src/delta/vtable/slice.rs +++ b/encodings/fastlanes/src/delta/vtable/slice.rs @@ -11,6 +11,7 @@ use vortex_array::arrays::slice::SliceReduce; use vortex_error::VortexResult; use crate::delta::array::DeltaArrayExt; +use crate::delta::array::DeltaArraySlotsExt; use crate::delta::array::lane_count; use crate::delta::vtable::Delta; diff --git a/encodings/fastlanes/src/delta/vtable/validity.rs b/encodings/fastlanes/src/delta/vtable/validity.rs index c8d8b7b4dc0..2e4d782886c 100644 --- a/encodings/fastlanes/src/delta/vtable/validity.rs +++ b/encodings/fastlanes/src/delta/vtable/validity.rs @@ -11,6 +11,7 @@ use vortex_error::VortexResult; use crate::Delta; use crate::bit_transpose::untranspose_validity; use crate::delta::array::DeltaArrayExt; +use crate::delta::array::DeltaArraySlotsExt; impl ValidityVTable for Delta { #[allow(clippy::disallowed_methods)] diff --git a/encodings/fastlanes/src/rle/array/mod.rs b/encodings/fastlanes/src/rle/array/mod.rs index 7ddff474836..5eb832f62cb 100644 --- a/encodings/fastlanes/src/rle/array/mod.rs +++ b/encodings/fastlanes/src/rle/array/mod.rs @@ -7,29 +7,30 @@ use std::fmt::Formatter; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::TypedArrayRef; -use vortex_error::VortexExpect as _; +use vortex_array::array_slots; use vortex_error::VortexResult; use vortex_error::vortex_ensure; pub mod rle_compress; pub mod rle_decompress; -/// Run values in the dictionary. -pub(super) const VALUES_SLOT: usize = 0; -/// Chunk-local indices from all chunks. The start of each chunk is looked up in `values_idx_offsets`. -pub(super) const INDICES_SLOT: usize = 1; -/// Index start positions of each value chunk. -/// -/// # Example -/// ```text -/// // Chunk 0: [10, 20] (starts at index 0) -/// // Chunk 1: [30, 40] (starts at index 2) -/// let values = [10, 20, 30, 40]; // Global values array -/// let values_idx_offsets = [0, 2]; // Chunk 0 starts at index 0, Chunk 1 starts at index 2 -/// ``` -pub(super) const VALUES_IDX_OFFSETS_SLOT: usize = 2; -pub(super) const NUM_SLOTS: usize = 3; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["values", "indices", "values_idx_offsets"]; +#[array_slots(crate::RLE)] +pub struct RLESlots { + /// Run values in the dictionary. + pub values: ArrayRef, + /// Chunk-local indices from all chunks. The start of each chunk is looked up in `values_idx_offsets`. + pub indices: ArrayRef, + /// Index start positions of each value chunk. + /// + /// # Example + /// ```text + /// // Chunk 0: [10, 20] (starts at index 0) + /// // Chunk 1: [30, 40] (starts at index 2) + /// let values = [10, 20, 30, 40]; // Global values array + /// let values_idx_offsets = [0, 2]; // Chunk 0 starts at index 0, Chunk 1 starts at index 2 + /// ``` + pub values_idx_offsets: ArrayRef, +} #[derive(Clone, Debug)] pub struct RLEData { @@ -79,28 +80,7 @@ impl RLEData { } } -pub trait RLEArrayExt: TypedArrayRef { - #[inline] - fn values(&self) -> &ArrayRef { - self.as_ref().slots()[VALUES_SLOT] - .as_ref() - .vortex_expect("RLEArray values slot must be populated") - } - - #[inline] - fn indices(&self) -> &ArrayRef { - self.as_ref().slots()[INDICES_SLOT] - .as_ref() - .vortex_expect("RLEArray indices slot must be populated") - } - - #[inline] - fn values_idx_offsets(&self) -> &ArrayRef { - self.as_ref().slots()[VALUES_IDX_OFFSETS_SLOT] - .as_ref() - .vortex_expect("RLEArray values_idx_offsets slot must be populated") - } - +pub trait RLEArrayExt: RLEArraySlotsExt { /// Values index offset relative to the first chunk. /// /// Offsets in `values_idx_offsets` are absolute and need to be shifted @@ -160,6 +140,7 @@ mod tests { use crate::RLE; use crate::RLEData; use crate::rle::array::RLEArrayExt; + use crate::rle::array::RLEArraySlotsExt; use crate::test::SESSION; #[test] diff --git a/encodings/fastlanes/src/rle/array/rle_compress.rs b/encodings/fastlanes/src/rle/array/rle_compress.rs index b7e04a13699..5622c3b90c0 100644 --- a/encodings/fastlanes/src/rle/array/rle_compress.rs +++ b/encodings/fastlanes/src/rle/array/rle_compress.rs @@ -171,6 +171,7 @@ mod tests { use super::*; use crate::rle::array::RLEArrayExt; + use crate::rle::array::RLEArraySlotsExt; static SESSION: LazyLock = LazyLock::new(|| { let session = vortex_array::array_session(); diff --git a/encodings/fastlanes/src/rle/array/rle_decompress.rs b/encodings/fastlanes/src/rle/array/rle_decompress.rs index 51766398070..7d5b70fb385 100644 --- a/encodings/fastlanes/src/rle/array/rle_decompress.rs +++ b/encodings/fastlanes/src/rle/array/rle_decompress.rs @@ -17,6 +17,7 @@ use vortex_error::vortex_panic; use crate::FL_CHUNK_SIZE; use crate::RLEArray; use crate::rle::RLEArrayExt; +use crate::rle::RLEArraySlotsExt; /// Decompresses an RLE array back into a primitive array. pub fn rle_decompress(array: &RLEArray, ctx: &mut ExecutionCtx) -> VortexResult { diff --git a/encodings/fastlanes/src/rle/compute/cast.rs b/encodings/fastlanes/src/rle/compute/cast.rs index 38e7927f8e6..a8aa7b199f0 100644 --- a/encodings/fastlanes/src/rle/compute/cast.rs +++ b/encodings/fastlanes/src/rle/compute/cast.rs @@ -12,6 +12,7 @@ use vortex_error::VortexResult; use crate::rle::RLE; use crate::rle::RLEArrayExt; +use crate::rle::RLEArraySlotsExt; impl CastReduce for RLE { fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult> { diff --git a/encodings/fastlanes/src/rle/kernel.rs b/encodings/fastlanes/src/rle/kernel.rs index 9dd05ffbae5..933e6e3d42a 100644 --- a/encodings/fastlanes/src/rle/kernel.rs +++ b/encodings/fastlanes/src/rle/kernel.rs @@ -18,6 +18,7 @@ use vortex_session::VortexSession; use crate::FL_CHUNK_SIZE; use crate::RLE; use crate::rle::RLEArrayExt; +use crate::rle::RLEArraySlotsExt; pub(crate) fn initialize(session: &VortexSession) { let kernels = session.kernels(); diff --git a/encodings/fastlanes/src/rle/mod.rs b/encodings/fastlanes/src/rle/mod.rs index 649d11f307c..d9f2c0499f3 100644 --- a/encodings/fastlanes/src/rle/mod.rs +++ b/encodings/fastlanes/src/rle/mod.rs @@ -3,7 +3,10 @@ mod array; pub use array::RLEArrayExt; +pub use array::RLEArraySlotsExt; pub use array::RLEData; +pub use array::RLESlots; +pub use array::RLESlotsView; mod compute; mod kernel; diff --git a/encodings/fastlanes/src/rle/vtable/mod.rs b/encodings/fastlanes/src/rle/vtable/mod.rs index cea655e53b4..c05806c8428 100644 --- a/encodings/fastlanes/src/rle/vtable/mod.rs +++ b/encodings/fastlanes/src/rle/vtable/mod.rs @@ -22,9 +22,7 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::serde::ArrayChildren; -use vortex_array::smallvec::smallvec; use vortex_array::vtable::VTable; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; @@ -32,11 +30,10 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::RLEData; -use crate::rle::array::INDICES_SLOT; use crate::rle::array::RLEArrayExt; -use crate::rle::array::SLOT_NAMES; -use crate::rle::array::VALUES_IDX_OFFSETS_SLOT; -use crate::rle::array::VALUES_SLOT; +use crate::rle::array::RLEArraySlotsExt; +use crate::rle::array::RLESlots; +use crate::rle::array::RLESlotsView; use crate::rle::array::rle_decompress::rle_decompress; use crate::rle::vtable::rules::RULES; @@ -93,20 +90,12 @@ impl VTable for RLE { len: usize, slots: &[Option], ) -> VortexResult<()> { - validate_parts( - slots[VALUES_SLOT] - .as_ref() - .vortex_expect("RLEArray values slot must be populated"), - slots[INDICES_SLOT] - .as_ref() - .vortex_expect("RLEArray indices slot must be populated"), - slots[VALUES_IDX_OFFSETS_SLOT] - .as_ref() - .vortex_expect("RLEArray values_idx_offsets slot must be populated"), - data.offset, - dtype, - len, - ) + let RLESlotsView { + values, + indices, + values_idx_offsets, + } = RLESlotsView::from_slots(slots); + validate_parts(values, indices, values_idx_offsets, data.offset, dtype, len) } fn nbuffers(_array: ArrayView<'_, Self>) -> usize { @@ -138,7 +127,7 @@ impl VTable for RLE { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + RLESlots::NAMES[idx].to_string() } fn serialize( @@ -195,7 +184,12 @@ impl VTable for RLE { usize::try_from(metadata.values_idx_offsets_len)?, )?; - let slots = smallvec![Some(values), Some(indices), Some(values_idx_offsets)]; + let slots = RLESlots { + values, + indices, + values_idx_offsets, + } + .into_slots(); let data = RLEData::try_new(metadata.offset as usize)?; Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) } @@ -219,7 +213,12 @@ impl RLE { length: usize, ) -> VortexResult { let dtype = DType::Primitive(values.dtype().as_ptype(), indices.dtype().nullability()); - let slots = smallvec![Some(values), Some(indices), Some(values_idx_offsets)]; + let slots = RLESlots { + values, + indices, + values_idx_offsets, + } + .into_slots(); let data = RLEData::try_new(offset)?; Array::try_from_parts(ArrayParts::new(RLE, dtype, length, data).with_slots(slots)) } @@ -236,7 +235,12 @@ impl RLE { length: usize, ) -> RLEArray { let dtype = DType::Primitive(values.dtype().as_ptype(), indices.dtype().nullability()); - let slots = smallvec![Some(values), Some(indices), Some(values_idx_offsets)]; + let slots = RLESlots { + values, + indices, + values_idx_offsets, + } + .into_slots(); let data = unsafe { RLEData::new_unchecked(offset) }; unsafe { Array::from_parts_unchecked(ArrayParts::new(RLE, dtype, length, data).with_slots(slots)) diff --git a/encodings/fastlanes/src/rle/vtable/operations.rs b/encodings/fastlanes/src/rle/vtable/operations.rs index b8f79db95a6..ca4d2d39545 100644 --- a/encodings/fastlanes/src/rle/vtable/operations.rs +++ b/encodings/fastlanes/src/rle/vtable/operations.rs @@ -11,6 +11,7 @@ use vortex_error::VortexResult; use super::RLE; use crate::FL_CHUNK_SIZE; use crate::rle::RLEArrayExt; +use crate::rle::RLEArraySlotsExt; impl OperationsVTable for RLE { fn scalar_at( diff --git a/encodings/fastlanes/src/rle/vtable/validity.rs b/encodings/fastlanes/src/rle/vtable/validity.rs index 917b43da996..603cc9db542 100644 --- a/encodings/fastlanes/src/rle/vtable/validity.rs +++ b/encodings/fastlanes/src/rle/vtable/validity.rs @@ -7,6 +7,7 @@ use vortex_error::VortexResult; use crate::rle::RLE; use crate::rle::RLEArrayExt; +use crate::rle::RLEArraySlotsExt; impl ValidityVTable for RLE { fn validity(array: ArrayView<'_, RLE>) -> VortexResult { diff --git a/encodings/fsst/src/array.rs b/encodings/fsst/src/array.rs index de0cf3157ba..e6b45166994 100644 --- a/encodings/fsst/src/array.rs +++ b/encodings/fsst/src/array.rs @@ -25,11 +25,11 @@ use vortex_array::EqMode; 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::VarBin; use vortex_array::arrays::VarBinArray; -use vortex_array::arrays::varbin::VarBinArrayExt; +use vortex_array::arrays::varbin::VarBinArraySlotsExt; use vortex_array::buffer::BufferHandle; use vortex_array::builders::ArrayBuilder; use vortex_array::builders::VarBinViewBuilder; @@ -38,7 +38,6 @@ use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::legacy_session; use vortex_array::serde::ArrayChildren; -use vortex_array::smallvec::smallvec; use vortex_array::validity::Validity; use vortex_array::vtable::VTable; use vortex_array::vtable::ValidityVTable; @@ -172,9 +171,7 @@ impl VTable for FSST { array: ArrayView<'_, Self>, _session: &VortexSession, ) -> VortexResult>> { - let codes_offsets = array.as_ref().slots()[CODES_OFFSETS_SLOT] - .as_ref() - .vortex_expect("FSSTArray codes_offsets slot"); + let codes_offsets = array.codes_offsets(); Ok(Some( FSSTMetadata { uncompressed_lengths_ptype: array.uncompressed_lengths().dtype().as_ptype().into(), @@ -276,11 +273,12 @@ impl VTable for FSST { len, &mut ctx, )?; - let slots = smallvec![ - Some(uncompressed_lengths), - Some(codes_offsets), - validity_to_child(&codes_validity, len), - ]; + let slots = FSSTSlots { + uncompressed_lengths, + codes_offsets, + codes_validity: validity_to_child(&codes_validity, len), + } + .into_slots(); let data = FSSTData::try_new(symbols, symbol_lengths, codes_bytes, len)?; return Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)); } @@ -292,7 +290,7 @@ impl VTable for FSST { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + FSSTSlots::NAMES[idx].to_string() } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { @@ -339,15 +337,15 @@ impl VTable for FSST { } } -/// Lengths of the original values before compression, can be compressed. -pub(crate) const UNCOMPRESSED_LENGTHS_SLOT: usize = 0; -/// The offsets array for the FSST-compressed codes. -pub(crate) const CODES_OFFSETS_SLOT: usize = 1; -/// The validity bitmap for the compressed codes. -pub(crate) const CODES_VALIDITY_SLOT: usize = 2; -pub(crate) const NUM_SLOTS: usize = 3; -pub(crate) const SLOT_NAMES: [&str; NUM_SLOTS] = - ["uncompressed_lengths", "codes_offsets", "codes_validity"]; +#[array_slots(FSST)] +pub struct FSSTSlots { + /// Lengths of the original values before compression, can be compressed. + pub uncompressed_lengths: ArrayRef, + /// The offsets array for the FSST-compressed codes. + pub codes_offsets: ArrayRef, + /// The validity bitmap for the compressed codes. + pub codes_validity: Option, +} /// The inner data for an FSST-compressed array. /// @@ -545,16 +543,17 @@ impl FSST { impl FSSTData { fn make_slots(codes: &VarBinArray, uncompressed_lengths: &ArrayRef) -> ArraySlots { - smallvec![ - Some(uncompressed_lengths.clone()), - Some(codes.offsets().clone()), - validity_to_child( + FSSTSlots { + uncompressed_lengths: uncompressed_lengths.clone(), + codes_offsets: codes.offsets().clone(), + codes_validity: validity_to_child( &codes .validity() .vortex_expect("FSST codes validity should be derivable"), codes.len(), ), - ] + } + .into_slots() } /// Build FSST data from a set of `symbols`, `symbol_lengths`, and compressed codes bytes. @@ -593,16 +592,14 @@ impl FSSTData { slots: &[Option], ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - let codes_offsets = slots[CODES_OFFSETS_SLOT] - .as_ref() - .vortex_expect("FSSTArray codes_offsets slot"); + let view = FSSTSlotsView::from_slots(slots); Self::validate_parts( &self.symbol_table.symbols, &self.symbol_table.symbol_lengths, &self.codes_bytes, - codes_offsets, + view.codes_offsets, dtype.nullability(), - uncompressed_lengths_from_slots(slots), + view.uncompressed_lengths, dtype, len, ctx, @@ -795,17 +792,7 @@ impl FSSTData { } } -fn uncompressed_lengths_from_slots(slots: &[Option]) -> &ArrayRef { - slots[UNCOMPRESSED_LENGTHS_SLOT] - .as_ref() - .vortex_expect("FSSTArray uncompressed_lengths slot") -} - -pub trait FSSTArrayExt: TypedArrayRef { - fn uncompressed_lengths(&self) -> &ArrayRef { - uncompressed_lengths_from_slots(self.as_ref().slots()) - } - +pub trait FSSTArrayExt: FSSTArraySlotsExt { fn uncompressed_lengths_dtype(&self) -> &DType { self.uncompressed_lengths().dtype() } @@ -813,14 +800,9 @@ pub trait FSSTArrayExt: TypedArrayRef { /// Reconstruct a [`VarBinArray`] for the compressed codes by combining the bytes /// from [`FSSTData`] with the offsets and validity stored in the array's slots. fn codes(&self) -> VarBinArray { - let offsets = self.as_ref().slots()[CODES_OFFSETS_SLOT] - .as_ref() - .vortex_expect("FSSTArray codes_offsets slot") - .clone(); - let validity = child_to_validity( - self.as_ref().slots()[CODES_VALIDITY_SLOT].as_ref(), - self.as_ref().dtype().nullability(), - ); + let offsets = self.codes_offsets().clone(); + let validity = + child_to_validity(self.codes_validity(), self.as_ref().dtype().nullability()); let codes_bytes = self.codes_bytes_handle().clone(); // SAFETY: components were validated at construction time. unsafe { @@ -839,12 +821,12 @@ pub trait FSSTArrayExt: TypedArrayRef { } } -impl> FSSTArrayExt for T {} +impl FSSTArrayExt for T {} impl ValidityVTable for FSST { fn validity(array: ArrayView<'_, FSST>) -> VortexResult { Ok(child_to_validity( - array.slots()[CODES_VALIDITY_SLOT].as_ref(), + array.slots()[FSSTSlots::CODES_VALIDITY].as_ref(), array.dtype().nullability(), )) } @@ -871,6 +853,7 @@ mod test { use crate::FSST; use crate::array::FSSTArrayExt; + use crate::array::FSSTArraySlotsExt; use crate::array::FSSTMetadata; use crate::fsst_compress; diff --git a/encodings/fsst/src/canonical.rs b/encodings/fsst/src/canonical.rs index 5c58dbbd501..b40bbf5d0bb 100644 --- a/encodings/fsst/src/canonical.rs +++ b/encodings/fsst/src/canonical.rs @@ -21,6 +21,7 @@ use vortex_error::VortexResult; use crate::FSST; use crate::FSSTArrayExt; +use crate::FSSTArraySlotsExt; pub(super) fn canonicalize_fsst( array: ArrayView<'_, FSST>, diff --git a/encodings/fsst/src/compress.rs b/encodings/fsst/src/compress.rs index b9f815b474f..e54b6af26f1 100644 --- a/encodings/fsst/src/compress.rs +++ b/encodings/fsst/src/compress.rs @@ -18,7 +18,7 @@ use vortex_array::IntoArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBin; use vortex_array::arrays::VarBinView; -use vortex_array::arrays::varbin::VarBinArrayExt; +use vortex_array::arrays::varbin::VarBinArraySlotsExt; use vortex_array::arrays::varbin::builder::VarBinBuilder; use vortex_array::arrays::varbinview::BinaryView; use vortex_array::buffer::BufferHandle; @@ -328,7 +328,7 @@ mod tests { use vortex_array::VortexSessionExecute; use vortex_array::array_session; use vortex_array::arrays::VarBinViewArray; - use vortex_array::arrays::varbin::VarBinArrayExt; + use vortex_array::arrays::varbin::VarBinArraySlotsExt; use vortex_array::dtype::PType; use vortex_error::VortexResult; diff --git a/encodings/fsst/src/compute/byte_length.rs b/encodings/fsst/src/compute/byte_length.rs index d78bd8be0ec..05505fe95e5 100644 --- a/encodings/fsst/src/compute/byte_length.rs +++ b/encodings/fsst/src/compute/byte_length.rs @@ -15,7 +15,7 @@ use vortex_array::validity::Validity; use vortex_error::VortexResult; use crate::FSST; -use crate::array::FSSTArrayExt; +use crate::array::FSSTArraySlotsExt; impl ByteLengthKernel for FSST { fn byte_length( diff --git a/encodings/fsst/src/compute/cast.rs b/encodings/fsst/src/compute/cast.rs index a1469681590..a3068adc437 100644 --- a/encodings/fsst/src/compute/cast.rs +++ b/encodings/fsst/src/compute/cast.rs @@ -6,7 +6,7 @@ use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::VarBinArray; -use vortex_array::arrays::varbin::VarBinArrayExt; +use vortex_array::arrays::varbin::VarBinArraySlotsExt; use vortex_array::dtype::DType; use vortex_array::scalar_fn::fns::cast::CastKernel; use vortex_array::scalar_fn::fns::cast::CastReduce; @@ -15,6 +15,7 @@ use vortex_error::VortexResult; use crate::FSST; use crate::FSSTArrayExt; +use crate::FSSTArraySlotsExt; fn build_with_codes_validity( array: ArrayView<'_, FSST>, diff --git a/encodings/fsst/src/compute/compare.rs b/encodings/fsst/src/compute/compare.rs index ec63af3735f..08b6681085b 100644 --- a/encodings/fsst/src/compute/compare.rs +++ b/encodings/fsst/src/compute/compare.rs @@ -21,6 +21,7 @@ use vortex_error::vortex_bail; use crate::FSST; use crate::FSSTArrayExt; +use crate::FSSTArraySlotsExt; impl CompareKernel for FSST { fn compare( lhs: ArrayView<'_, Self>, diff --git a/encodings/fsst/src/compute/filter.rs b/encodings/fsst/src/compute/filter.rs index eae72cac8dd..45a95520aca 100644 --- a/encodings/fsst/src/compute/filter.rs +++ b/encodings/fsst/src/compute/filter.rs @@ -13,6 +13,7 @@ use vortex_mask::Mask; use crate::FSST; use crate::FSSTArrayExt; +use crate::FSSTArraySlotsExt; impl FilterKernel for FSST { fn filter( diff --git a/encodings/fsst/src/compute/like.rs b/encodings/fsst/src/compute/like.rs index c922fba088e..d3ff7850fa0 100644 --- a/encodings/fsst/src/compute/like.rs +++ b/encodings/fsst/src/compute/like.rs @@ -7,7 +7,7 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; -use vortex_array::arrays::varbin::VarBinArrayExt; +use vortex_array::arrays::varbin::VarBinArraySlotsExt; use vortex_array::match_each_integer_ptype; use vortex_array::scalar_fn::fns::like::LikeKernel; use vortex_array::scalar_fn::fns::like::LikeOptions; diff --git a/encodings/fsst/src/compute/mod.rs b/encodings/fsst/src/compute/mod.rs index d447b4b4958..c8477119007 100644 --- a/encodings/fsst/src/compute/mod.rs +++ b/encodings/fsst/src/compute/mod.rs @@ -21,6 +21,7 @@ use vortex_error::vortex_err; use crate::FSST; use crate::FSSTArrayExt; +use crate::FSSTArraySlotsExt; impl TakeExecute for FSST { fn take( diff --git a/encodings/fsst/src/slice.rs b/encodings/fsst/src/slice.rs index b98f5fa03b3..c2b2e320d31 100644 --- a/encodings/fsst/src/slice.rs +++ b/encodings/fsst/src/slice.rs @@ -13,6 +13,7 @@ use vortex_error::vortex_err; use crate::FSST; use crate::FSSTArrayExt; +use crate::FSSTArraySlotsExt; impl SliceReduce for FSST { fn slice(array: ArrayView<'_, Self>, range: Range) -> VortexResult> { diff --git a/encodings/onpair/src/compute/filter.rs b/encodings/onpair/src/compute/filter.rs index 1543bff136a..219a5a72907 100644 --- a/encodings/onpair/src/compute/filter.rs +++ b/encodings/onpair/src/compute/filter.rs @@ -20,7 +20,7 @@ use vortex_array::IntoArray; use vortex_array::arrays::List; use vortex_array::arrays::ListArray; use vortex_array::arrays::filter::FilterKernel; -use vortex_array::arrays::list::ListArrayExt; +use vortex_array::arrays::list::ListArraySlotsExt; use vortex_array::validity::Validity; use vortex_error::VortexExpect; use vortex_error::VortexResult; diff --git a/encodings/parquet-variant/src/array.rs b/encodings/parquet-variant/src/array.rs index bc7fd2c22ec..881efec752f 100644 --- a/encodings/parquet-variant/src/array.rs +++ b/encodings/parquet-variant/src/array.rs @@ -14,6 +14,7 @@ use vortex_array::EmptyArrayData; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::TypedArrayRef; +use vortex_array::array_slots; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::List; use vortex_array::arrays::ListArray; @@ -22,6 +23,7 @@ use vortex_array::arrays::StructArray; use vortex_array::arrays::Variant; use vortex_array::arrays::VariantArray; use vortex_array::arrays::list::ListArrayExt; +use vortex_array::arrays::list::ListArraySlotsExt; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::arrays::variant::VariantArrayExt; use vortex_array::builtins::ArrayBuiltins; @@ -30,7 +32,6 @@ use vortex_array::dtype::FieldName; use vortex_array::dtype::FieldNames; use vortex_array::dtype::Nullability; use vortex_array::scalar::Scalar; -use vortex_array::smallvec::smallvec; use vortex_array::validity::Validity; use vortex_array::vtable::child_to_validity; use vortex_array::vtable::validity_to_child; @@ -49,16 +50,17 @@ use vortex_error::vortex_err; use crate::ParquetVariant; use crate::ParquetVariantArray; -/// The validity bitmap indicating which elements are non-null. -pub(crate) const VALIDITY_SLOT: usize = 0; -/// The metadata array for the Parquet variant values. -pub(crate) const METADATA_SLOT: usize = 1; -/// The value array containing the Parquet variant data. -pub(crate) const VALUE_SLOT: usize = 2; -/// The typed value array for strongly-typed Parquet variant data. -pub(crate) const TYPED_VALUE_SLOT: usize = 3; -pub(crate) const NUM_SLOTS: usize = 4; -pub(crate) const SLOT_NAMES: [&str; NUM_SLOTS] = ["validity", "metadata", "value", "typed_value"]; +#[array_slots(ParquetVariant)] +pub struct ParquetVariantSlots { + /// The validity bitmap indicating which elements are non-null. + pub validity: Option, + /// The metadata array for the Parquet variant values. + pub metadata: ArrayRef, + /// The value array containing the Parquet variant data. + pub value: Option, + /// The typed value array for strongly-typed Parquet variant data. + pub typed_value: Option, +} impl ParquetVariant { /// Creates a Parquet Variant array from canonical extension storage slots. @@ -73,12 +75,13 @@ impl ParquetVariant { ) -> VortexResult> { let len = metadata.len(); let dtype = DType::Variant(validity.nullability()); - let slots = smallvec![ - validity_to_child(&validity, len), - Some(metadata), + let slots = ParquetVariantSlots { + validity: validity_to_child(&validity, len), + metadata, value, typed_value, - ]; + } + .into_slots(); Array::try_from_parts( ArrayParts::new(ParquetVariant, dtype, len, EmptyArrayData).with_slots(slots), ) @@ -147,8 +150,8 @@ pub(crate) fn core_storage_without_typed_value( // the outer `VariantArray::shredded` slot and the original had no `value`, synthesize an // all-null `value` so the remaining `ParquetVariant` still satisfies that invariant and // can round-trip back through `to_arrow`. - let value = array.value_array().cloned().or_else(|| { - array.typed_value_array().map(|_| { + let value = array.value().cloned().or_else(|| { + array.typed_value().map(|_| { ConstantArray::new( Scalar::null(DType::Binary(Nullability::Nullable)), array.len(), @@ -157,13 +160,8 @@ pub(crate) fn core_storage_without_typed_value( }) }); - ParquetVariant::try_new( - array.validity()?, - array.metadata_array().clone(), - value, - None, - ) - .map(IntoArray::into_array) + ParquetVariant::try_new(array.validity()?, array.metadata().clone(), value, None) + .map(IntoArray::into_array) } /// Converts a Parquet `typed_value` tree into the storage-agnostic canonical shredded tree. @@ -274,7 +272,7 @@ fn logical_shredded_from_parquet_field( let parquet_field = ParquetVariant::try_new(validity, metadata.clone(), Some(value), Some(typed_value))?; let shredded = parquet_field - .typed_value_array() + .typed_value() .cloned() .map(|typed_value| { logical_shredded_from_parquet_typed_value(metadata, typed_value, ctx) @@ -380,7 +378,7 @@ fn parquet_shredded_field_from_logical( "cannot rebuild Parquet shredded field: nested Variant lacks Parquet Variant core storage" ) })?; - let value = core.value_array().cloned().ok_or_else(|| { + let value = core.value().cloned().ok_or_else(|| { vortex_err!("cannot rebuild Parquet shredded field: partially shredded Variant has no residual value") })?; let typed_value = variant @@ -416,41 +414,26 @@ fn parquet_shredded_field_from_logical( } /// Accessors and Arrow conversion for Parquet Variant storage arrays. -pub trait ParquetVariantArrayExt: TypedArrayRef { - /// Returns the non-nullable Parquet Variant metadata child. - fn metadata_array(&self) -> &ArrayRef { - self.as_ref().slots()[METADATA_SLOT] - .as_ref() - .vortex_expect("ParquetVariantArray metadata slot") - } - +pub trait ParquetVariantArrayExt: + TypedArrayRef + ParquetVariantArraySlotsExt +{ /// Returns the outer row validity for the Variant values. fn validity(&self) -> Validity { child_to_validity( - self.as_ref().slots()[VALIDITY_SLOT].as_ref(), + self.as_ref().slots()[ParquetVariantSlots::VALIDITY].as_ref(), self.as_ref().dtype().nullability(), ) } - /// Returns the optional raw Parquet Variant `value` child. - fn value_array(&self) -> Option<&ArrayRef> { - self.as_ref().slots()[VALUE_SLOT].as_ref() - } - - /// Returns the optional shredded Parquet Variant `typed_value` child. - fn typed_value_array(&self) -> Option<&ArrayRef> { - self.as_ref().slots()[TYPED_VALUE_SLOT].as_ref() - } - /// Converts this storage array to Arrow's canonical Parquet Variant extension storage. #[expect( deprecated, reason = "TODO(aduffy): figure out what to do with Parquet Variant" )] fn to_arrow(&self, ctx: &mut ExecutionCtx) -> VortexResult { - let metadata = self.metadata_array(); + let metadata = self.metadata(); let len = metadata.len(); - let nulls = to_arrow_null_buffer(self.validity(), len, ctx)?; + let nulls = to_arrow_null_buffer(ParquetVariantArrayExt::validity(self), len, ctx)?; let mut fields = Vec::with_capacity(3); let mut arrays: Vec = Vec::with_capacity(3); @@ -463,7 +446,7 @@ pub trait ParquetVariantArrayExt: TypedArrayRef { ))); arrays.push(metadata_arrow); - if let Some(value) = self.value_array() { + if let Some(value) = self.value() { let value_arrow = value.clone().execute_arrow(None, ctx)?; fields.push(Arc::new(Field::new( "value", @@ -473,7 +456,7 @@ pub trait ParquetVariantArrayExt: TypedArrayRef { arrays.push(value_arrow); } - if let Some(typed_value) = self.typed_value_array() { + if let Some(typed_value) = self.typed_value() { let tv_arrow = typed_value.clone().execute_arrow(None, ctx)?; fields.push(Arc::new(Field::new( "typed_value", @@ -528,6 +511,7 @@ mod tests { use crate::ParquetVariant; use crate::array::ParquetVariantArrayExt; + use crate::array::ParquetVariantArraySlotsExt; use crate::array::parquet_typed_value_from_logical_shredded; static SESSION: LazyLock = LazyLock::new(|| { @@ -632,7 +616,7 @@ mod tests { .as_opt::() .ok_or_else(|| vortex_err!("expected parquet variant array"))?; let typed_value = parquet_array - .typed_value_array() + .typed_value() .ok_or_else(|| vortex_err!("expected typed_value child"))? .clone() .execute::(&mut SESSION.create_execution_ctx())?; @@ -716,10 +700,10 @@ mod tests { let parquet_array = vortex_arr .as_opt::() .ok_or_else(|| vortex_err!("expected parquet variant array"))?; - assert!(parquet_array.value_array().is_none()); + assert!(parquet_array.value().is_none()); let typed_value = parquet_array - .typed_value_array() + .typed_value() .ok_or_else(|| vortex_err!("expected typed_value child"))? .clone() .execute::(&mut SESSION.create_execution_ctx())?; @@ -753,8 +737,8 @@ mod tests { let parquet_array = vortex_arr .as_opt::() .ok_or_else(|| vortex_err!("expected parquet variant array"))?; - assert!(parquet_array.value_array().is_some()); - assert!(parquet_array.typed_value_array().is_some()); + assert!(parquet_array.value().is_some()); + assert!(parquet_array.typed_value().is_some()); let mut ctx = SESSION.create_execution_ctx(); let roundtripped = parquet_array.to_arrow(&mut ctx)?; @@ -826,7 +810,7 @@ mod tests { original .as_opt::() .ok_or_else(|| vortex_err!("expected parquet variant"))? - .typed_value_array() + .typed_value() .is_some(), "fixture must be shredded" ); @@ -849,8 +833,8 @@ mod tests { let rebuilt_typed_value = parquet_typed_value_from_logical_shredded(logical, &mut ctx)?; let rebuilt = ParquetVariant::try_new( ParquetVariantArrayExt::validity(&core), - core.metadata_array().clone(), - core.value_array().cloned(), + core.metadata().clone(), + core.value().cloned(), Some(rebuilt_typed_value), )? .into_array(); diff --git a/encodings/parquet-variant/src/arrow.rs b/encodings/parquet-variant/src/arrow.rs index 65efd6c571c..273a282e87d 100644 --- a/encodings/parquet-variant/src/arrow.rs +++ b/encodings/parquet-variant/src/arrow.rs @@ -36,6 +36,7 @@ use vortex_session::registry::Id; use crate::ParquetVariant; use crate::ParquetVariantArrayExt; +use crate::ParquetVariantArraySlotsExt; use crate::array::parquet_typed_value_from_logical_shredded; /// Arrow canonical extension name for Parquet Variant storage. @@ -68,9 +69,9 @@ pub(crate) fn export_storage_to_target( for field in target_fields { let child = match field.name().as_str() { - "metadata" => Some(parquet_array.metadata_array().clone()), - "value" => parquet_array.value_array().cloned(), - "typed_value" => parquet_array.typed_value_array().cloned(), + "metadata" => Some(parquet_array.metadata().clone()), + "value" => parquet_array.value().cloned(), + "typed_value" => parquet_array.typed_value().cloned(), _ => unreachable!("storage fields were validated before export"), }; let Some(child) = child else { @@ -146,8 +147,8 @@ pub(crate) fn parquet_variant_for_export( ParquetVariant::try_new( ParquetVariantArrayExt::validity(&parquet_core), - parquet_core.metadata_array().clone(), - parquet_core.value_array().cloned(), + parquet_core.metadata().clone(), + parquet_core.value().cloned(), Some(typed_value), ) .map(IntoArray::into_array) @@ -195,8 +196,8 @@ impl ArrowExportVTable for ParquetVariant { && let Some((request_has_value, request_has_typed_value)) = parquet_variant_storage_request(fields) { - let has_value = parquet_array.value_array().is_some(); - let has_typed_value = parquet_array.typed_value_array().is_some(); + let has_value = parquet_array.value().is_some(); + let has_typed_value = parquet_array.typed_value().is_some(); if request_has_value && !request_has_typed_value && has_typed_value { return Ok(ArrowExport::Exported(export_unshredded_storage_to_target( diff --git a/encodings/parquet-variant/src/compute/allnondistinct.rs b/encodings/parquet-variant/src/compute/allnondistinct.rs index 4f5568fbcf6..7e6219bc542 100644 --- a/encodings/parquet-variant/src/compute/allnondistinct.rs +++ b/encodings/parquet-variant/src/compute/allnondistinct.rs @@ -16,7 +16,7 @@ use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use crate::ParquetVariant; -use crate::ParquetVariantArrayExt; +use crate::ParquetVariantArraySlotsExt; /// Lets `AllNonDistinct` compare two `ParquetVariant` arrays without canonicalizing them. /// @@ -51,7 +51,7 @@ impl DynAggregateKernel for AllNonDistinctParquetVariant { return Ok(None); }; - let typed_identical = match (lhs.typed_value_array(), rhs.typed_value_array()) { + let typed_identical = match (lhs.typed_value(), rhs.typed_value()) { (Some(lhs_typed), Some(rhs_typed)) => { if lhs_typed.dtype().eq_ignore_nullability(rhs_typed.dtype()) { all_non_distinct(lhs_typed, rhs_typed, ctx)? @@ -63,7 +63,7 @@ impl DynAggregateKernel for AllNonDistinctParquetVariant { }; if typed_identical { - let values_identical = match (lhs.value_array(), rhs.value_array()) { + let values_identical = match (lhs.value(), rhs.value()) { (Some(lhs_value), Some(rhs_value)) => all_non_distinct(lhs_value, rhs_value, ctx)?, (None, None) => true, // Mixed shredding layouts: let the generic canonical path handle it. diff --git a/encodings/parquet-variant/src/json_to_variant_tests.rs b/encodings/parquet-variant/src/json_to_variant_tests.rs index cd3b2f035ee..951277af7e1 100644 --- a/encodings/parquet-variant/src/json_to_variant_tests.rs +++ b/encodings/parquet-variant/src/json_to_variant_tests.rs @@ -36,7 +36,7 @@ use vortex_json::Json; use vortex_session::VortexSession; use crate::ParquetVariant; -use crate::ParquetVariantArrayExt; +use crate::ParquetVariantArraySlotsExt; use crate::ShreddingSpec; use crate::json_to_variant; @@ -200,7 +200,7 @@ fn shredding_produces_typed_value_child() -> VortexResult<()> { let result = execute_json_to_variant(input, shred_field_as_i64("a")?)?; assert!( - result.as_::().typed_value_array().is_some(), + result.as_::().typed_value().is_some(), "expected shredded typed_value child" ); @@ -283,7 +283,7 @@ fn shredding_root_path_shreds_top_level_values() -> VortexResult<()> { let result = execute_json_to_variant(input, spec)?; assert!( - result.as_::().typed_value_array().is_some(), + result.as_::().typed_value().is_some(), "expected shredded typed_value child" ); assert_variant_i64_rows(&result.slice(0..2)?, &[Some(1), Some(2)])?; diff --git a/encodings/parquet-variant/src/kernel.rs b/encodings/parquet-variant/src/kernel.rs index c783e6600ac..82477469c32 100644 --- a/encodings/parquet-variant/src/kernel.rs +++ b/encodings/parquet-variant/src/kernel.rs @@ -57,6 +57,7 @@ use vortex_session::VortexSession; use crate::ParquetVariant; use crate::ParquetVariantArrayExt; +use crate::ParquetVariantArraySlotsExt; use crate::compute::AllNonDistinctParquetVariant; pub(crate) fn initialize(session: &VortexSession) { @@ -229,15 +230,9 @@ impl SliceKernel for ParquetVariant { _ctx: &mut ExecutionCtx, ) -> VortexResult> { let validity = array.validity()?.slice(range.clone())?; - let metadata = array.metadata_array().slice(range.clone())?; - let value = array - .value_array() - .map(|v| v.slice(range.clone())) - .transpose()?; - let typed_value = array - .typed_value_array() - .map(|tv| tv.slice(range)) - .transpose()?; + let metadata = array.metadata().slice(range.clone())?; + let value = array.value().map(|v| v.slice(range.clone())).transpose()?; + let typed_value = array.typed_value().map(|tv| tv.slice(range)).transpose()?; Ok(Some( ParquetVariant::try_new(validity, metadata, value, typed_value)?.into_array(), )) @@ -251,13 +246,10 @@ impl FilterKernel for ParquetVariant { _ctx: &mut ExecutionCtx, ) -> VortexResult> { let validity = array.validity()?.filter(mask)?; - let metadata = array.metadata_array().filter(mask.clone())?; - let value = array - .value_array() - .map(|v| v.filter(mask.clone())) - .transpose()?; + let metadata = array.metadata().filter(mask.clone())?; + let value = array.value().map(|v| v.filter(mask.clone())).transpose()?; let typed_value = array - .typed_value_array() + .typed_value() .map(|tv| tv.filter(mask.clone())) .transpose()?; Ok(Some( @@ -273,13 +265,10 @@ impl TakeExecute for ParquetVariant { _ctx: &mut ExecutionCtx, ) -> VortexResult> { let validity = array.validity()?.take(indices)?; - let metadata = array.metadata_array().take(indices.clone())?; - let value = array - .value_array() - .map(|v| v.take(indices.clone())) - .transpose()?; + let metadata = array.metadata().take(indices.clone())?; + let value = array.value().map(|v| v.take(indices.clone())).transpose()?; let typed_value = array - .typed_value_array() + .typed_value() .map(|tv| tv.take(indices.clone())) .transpose()?; Ok(Some( @@ -338,7 +327,7 @@ mod tests { use vortex_session::VortexSession; use crate::ParquetVariant; - use crate::ParquetVariantArrayExt; + use crate::ParquetVariantArraySlotsExt; static SESSION: LazyLock = LazyLock::new(|| { let session = vortex_array::array_session(); @@ -922,7 +911,7 @@ mod tests { let executed = array.clone().execute::(&mut ctx)?; let typed_value = executed .as_::() - .typed_value_array() + .typed_value() .ok_or_else(|| vortex_err!("expected typed_value child"))? .clone() .execute::(&mut ctx)?; @@ -1004,8 +993,8 @@ mod tests { .core_storage() .as_opt::() .ok_or_else(|| vortex_err!("expected parquet variant core storage"))?; - assert!(core_storage.typed_value_array().is_none()); - assert!(core_storage.value_array().is_some()); + assert!(core_storage.typed_value().is_none()); + assert!(core_storage.value().is_some()); let shredded = canonical_variant .shredded() @@ -1103,7 +1092,7 @@ mod tests { assert!( parquet_array .as_::() - .typed_value_array() + .typed_value() .is_some() ); @@ -1117,8 +1106,8 @@ mod tests { .core_storage() .as_opt::() .ok_or_else(|| vortex_err!("expected parquet variant core storage"))?; - assert!(core_storage.typed_value_array().is_none()); - assert!(core_storage.value_array().is_some()); + assert!(core_storage.typed_value().is_none()); + assert!(core_storage.value().is_some()); Ok(()) } diff --git a/encodings/parquet-variant/src/lib.rs b/encodings/parquet-variant/src/lib.rs index 88125af04fd..5a82702b8f2 100644 --- a/encodings/parquet-variant/src/lib.rs +++ b/encodings/parquet-variant/src/lib.rs @@ -37,6 +37,9 @@ mod vtable; use std::sync::Arc; pub use array::ParquetVariantArrayExt; +pub use array::ParquetVariantArraySlotsExt; +pub use array::ParquetVariantSlots; +pub use array::ParquetVariantSlotsView; use vortex_array::session::ArraySessionExt; use vortex_arrow::ArrowSessionExt; pub use vortex_json::JsonToVariant; diff --git a/encodings/parquet-variant/src/operations.rs b/encodings/parquet-variant/src/operations.rs index 3cc3d2e5b4a..16563dc7ea1 100644 --- a/encodings/parquet-variant/src/operations.rs +++ b/encodings/parquet-variant/src/operations.rs @@ -27,7 +27,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; -use crate::ParquetVariantArrayExt; +use crate::ParquetVariantArraySlotsExt; use crate::vtable::ParquetVariant; impl OperationsVTable for ParquetVariant { @@ -46,7 +46,7 @@ impl OperationsVTable for ParquetVariant { } let metadata = array - .metadata_array() + .metadata() .execute_scalar(index, ctx)? .as_binary() .value() @@ -54,8 +54,8 @@ impl OperationsVTable for ParquetVariant { .vortex_expect("non-null metadata row must have binary value"); let inner = scalar_from_variant_storage( metadata.as_ref(), - array.value_array(), - array.typed_value_array(), + array.value(), + array.typed_value(), index, ctx, )?; diff --git a/encodings/parquet-variant/src/vtable.rs b/encodings/parquet-variant/src/vtable.rs index 4fa3ea10690..2cd2cd35664 100644 --- a/encodings/parquet-variant/src/vtable.rs +++ b/encodings/parquet-variant/src/vtable.rs @@ -16,7 +16,6 @@ use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::serde::ArrayChildren; -use vortex_array::smallvec::smallvec; use vortex_array::validity::Validity; use vortex_array::vtable::VTable; use vortex_array::vtable::child_to_validity; @@ -30,13 +29,9 @@ use vortex_proto::dtype as pb; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::array::METADATA_SLOT; -use crate::array::NUM_SLOTS; -use crate::array::ParquetVariantArrayExt; -use crate::array::SLOT_NAMES; -use crate::array::TYPED_VALUE_SLOT; -use crate::array::VALIDITY_SLOT; -use crate::array::VALUE_SLOT; +use crate::array::ParquetVariantArraySlotsExt; +use crate::array::ParquetVariantSlots; +use crate::array::ParquetVariantSlotsView; use crate::array::core_storage_without_typed_value; use crate::array::logical_shredded_from_parquet_typed_value; @@ -103,16 +98,18 @@ impl VTable for ParquetVariant { slots: &[Option], ) -> VortexResult<()> { vortex_ensure!( - slots.len() == NUM_SLOTS, - "ParquetVariantArray expects {NUM_SLOTS} slots, got {}", + slots.len() == ParquetVariantSlots::COUNT, + "ParquetVariantArray expects {} slots, got {}", + ParquetVariantSlots::COUNT, slots.len() ); - let validity = child_to_validity(slots[VALIDITY_SLOT].as_ref(), dtype.nullability()); - let metadata = slots[METADATA_SLOT] - .as_ref() - .ok_or_else(|| vortex_err!("ParquetVariantArray metadata slot"))?; - let value = slots[VALUE_SLOT].as_ref(); - let typed_value = slots[TYPED_VALUE_SLOT].as_ref(); + let ParquetVariantSlotsView { + validity, + metadata, + value, + typed_value, + } = ParquetVariantSlotsView::from_slots(slots); + let validity = child_to_validity(validity, dtype.nullability()); vortex_ensure!( matches!(dtype, DType::Variant(_)), @@ -180,7 +177,7 @@ impl VTable for ParquetVariant { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + ParquetVariantSlots::NAMES[idx].to_string() } fn serialize( @@ -188,14 +185,14 @@ impl VTable for ParquetVariant { _session: &VortexSession, ) -> VortexResult>> { let typed_value_dtype = array - .typed_value_array() + .typed_value() .map(|tv| tv.dtype().try_into()) .transpose()?; Ok(Some( ParquetVariantMetadataProto { - has_value: array.value_array().is_some(), + has_value: array.value().is_some(), typed_value_dtype, - value_nullable: array.value_array().is_some_and(|v| v.dtype().is_nullable()), + value_nullable: array.value().is_some_and(|v| v.dtype().is_nullable()), } .encode_to_vec(), )) @@ -265,21 +262,22 @@ impl VTable for ParquetVariant { None }; - let slots = smallvec![ - validity_to_child(&validity, len), - Some(variant_metadata), + let slots = ParquetVariantSlots { + validity: validity_to_child(&validity, len), + metadata: variant_metadata, value, typed_value, - ]; + } + .into_slots(); Ok(ArrayParts::new(self.clone(), dtype.clone(), len, EmptyArrayData).with_slots(slots)) } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { let shredded = array - .typed_value_array() + .typed_value() .cloned() .map(|typed_value| { - logical_shredded_from_parquet_typed_value(array.metadata_array(), typed_value, ctx) + logical_shredded_from_parquet_typed_value(array.metadata(), typed_value, ctx) }) .transpose()?; let core_storage = core_storage_without_typed_value(&array)?; @@ -339,7 +337,7 @@ mod tests { use vortex_session::registry::ReadContext; use crate::ParquetVariant; - use crate::array::ParquetVariantArrayExt; + use crate::array::ParquetVariantArraySlotsExt; static SESSION: LazyLock = LazyLock::new(|| { let session = vortex_array::array_session(); @@ -416,7 +414,7 @@ mod tests { let parquet_variant = ParquetVariant::try_new(Validity::NonNullable, metadata, None, Some(typed_value))?; - assert!(parquet_variant.typed_value_array().is_some()); + assert!(parquet_variant.typed_value().is_some()); let mut ctx = SESSION.create_execution_ctx(); let Canonical::Variant(variant) = parquet_variant @@ -430,7 +428,7 @@ mod tests { .core_storage() .as_opt::() .ok_or_else(|| vortex_err!("expected parquet variant core storage"))?; - assert!(core_storage.typed_value_array().is_none()); + assert!(core_storage.typed_value().is_none()); let shredded = variant .shredded() .ok_or_else(|| vortex_err!("expected canonical shredded child"))?; @@ -533,7 +531,7 @@ mod tests { .as_opt::() .ok_or_else(|| vortex_err!("expected parquet variant array"))?; let typed = decoded_pv - .typed_value_array() + .typed_value() .ok_or_else(|| vortex_err!("expected typed_value child"))?; assert_eq!(typed.dtype(), &DType::Variant(Nullability::NonNullable)); Ok(()) @@ -555,8 +553,8 @@ mod tests { let decoded_pv = decoded .as_opt::() .ok_or_else(|| vortex_err!("expected parquet variant array"))?; - assert!(decoded_pv.value_array().is_some()); - assert!(decoded_pv.typed_value_array().is_none()); + assert!(decoded_pv.value().is_some()); + assert!(decoded_pv.typed_value().is_none()); Ok(()) } @@ -580,7 +578,7 @@ mod tests { .as_opt::() .ok_or_else(|| vortex_err!("expected parquet variant array"))?; let typed = decoded_pv - .typed_value_array() + .typed_value() .ok_or_else(|| vortex_err!("expected typed_value child"))?; assert_eq!( typed.dtype(), diff --git a/encodings/runend/src/array.rs b/encodings/runend/src/array.rs index 2a9659efdd8..1e6beb8d06e 100644 --- a/encodings/runend/src/array.rs +++ b/encodings/runend/src/array.rs @@ -21,6 +21,7 @@ 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::VarBinViewArray; use vortex_array::buffer::BufferHandle; @@ -29,7 +30,6 @@ use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::legacy_session; use vortex_array::serde::ArrayChildren; -use vortex_array::smallvec::smallvec; use vortex_array::validity::Validity; use vortex_array::vtable::VTable; use vortex_array::vtable::ValidityVTable; @@ -93,20 +93,15 @@ impl VTable for RunEnd { len: usize, slots: &[Option], ) -> VortexResult<()> { - let ends = slots[ENDS_SLOT] - .as_ref() - .vortex_expect("RunEndArray ends slot"); - let values = slots[VALUES_SLOT] - .as_ref() - .vortex_expect("RunEndArray values slot"); + let s = RunEndSlotsView::from_slots(slots); // 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_parts(s.ends, s.values, data.offset, len, &mut ctx)?; vortex_ensure!( - values.dtype() == dtype, + s.values.dtype() == dtype, "expected dtype {}, got {}", dtype, - values.dtype() + s.values.dtype() ); Ok(()) } @@ -162,13 +157,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"); - let slots = smallvec![Some(ends), Some(values)]; + let slots = RunEndSlots { ends, values }.into_slots(); let data = RunEndData::new(offset); Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + RunEndSlots::NAMES[idx].to_string() } fn reduce_parent( @@ -184,12 +179,13 @@ impl VTable for RunEnd { } } -/// The run-end positions marking where each run terminates. -pub(super) const ENDS_SLOT: usize = 0; -/// The values for each run. -pub(super) const VALUES_SLOT: usize = 1; -pub(super) const NUM_SLOTS: usize = 2; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["ends", "values"]; +#[array_slots(RunEnd)] +pub struct RunEndSlots { + /// The run-end positions marking where each run terminates. + pub ends: ArrayRef, + /// The values for each run. + pub values: ArrayRef, +} #[derive(Clone, Debug)] pub struct RunEndData { @@ -208,23 +204,11 @@ pub struct RunEndDataParts { pub offset: usize, } -pub trait RunEndArrayExt: TypedArrayRef { +pub trait RunEndArrayExt: TypedArrayRef + RunEndArraySlotsExt { fn offset(&self) -> usize { self.offset } - fn ends(&self) -> &ArrayRef { - self.as_ref().slots()[ENDS_SLOT] - .as_ref() - .vortex_expect("RunEndArray ends slot") - } - - fn values(&self) -> &ArrayRef { - self.as_ref().slots()[VALUES_SLOT] - .as_ref() - .vortex_expect("RunEndArray values slot") - } - fn dtype(&self) -> &DType { self.values().dtype() } @@ -255,7 +239,7 @@ impl RunEnd { length: usize, ) -> RunEndArray { let dtype = values.dtype().clone(); - let slots = smallvec![Some(ends), Some(values)]; + let slots = RunEndSlots { ends, values }.into_slots(); let data = unsafe { RunEndData::new_unchecked(offset) }; unsafe { Array::from_parts_unchecked( @@ -273,7 +257,7 @@ impl RunEnd { let len = RunEndData::logical_len_from_ends(&ends, ctx)?; RunEndData::validate_parts(&ends, &values, 0, len, ctx)?; let dtype = values.dtype().clone(); - let slots = smallvec![Some(ends), Some(values)]; + let slots = RunEndSlots { ends, values }.into_slots(); let data = RunEndData::new(0); Array::try_from_parts(ArrayParts::new(RunEnd, dtype, len, data).with_slots(slots)) } @@ -288,7 +272,7 @@ impl RunEnd { ) -> VortexResult { RunEndData::validate_parts(&ends, &values, offset, length, ctx)?; let dtype = values.dtype().clone(); - let slots = smallvec![Some(ends), Some(values)]; + let slots = RunEndSlots { ends, values }.into_slots(); let data = RunEndData::new(offset); Array::try_from_parts(ArrayParts::new(RunEnd, dtype, length, data).with_slots(slots)) } @@ -305,7 +289,7 @@ impl RunEnd { let ends = ends.into_array(); let len = array.len(); let dtype = values.dtype().clone(); - let slots = smallvec![Some(ends), Some(values)]; + let slots = RunEndSlots { ends, values }.into_slots(); let data = unsafe { RunEndData::new_unchecked(0) }; Array::try_from_parts(ArrayParts::new(RunEnd, dtype, len, data).with_slots(slots)) } else { diff --git a/encodings/runend/src/compute/cast.rs b/encodings/runend/src/compute/cast.rs index 4722f9e5549..aa2dc5eb406 100644 --- a/encodings/runend/src/compute/cast.rs +++ b/encodings/runend/src/compute/cast.rs @@ -11,6 +11,7 @@ use vortex_error::VortexResult; use crate::RunEnd; use crate::array::RunEndArrayExt; +use crate::array::RunEndArraySlotsExt; impl CastReduce for RunEnd { fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult> { // Cast the values array to the target type diff --git a/encodings/runend/src/compute/compare.rs b/encodings/runend/src/compute/compare.rs index ecdd03e3686..f3facb1167a 100644 --- a/encodings/runend/src/compute/compare.rs +++ b/encodings/runend/src/compute/compare.rs @@ -16,6 +16,7 @@ use vortex_error::VortexResult; use crate::RunEnd; use crate::array::RunEndArrayExt; +use crate::array::RunEndArraySlotsExt; use crate::decompress_bool::runend_decode_bools; impl CompareKernel for RunEnd { diff --git a/encodings/runend/src/compute/fill_null.rs b/encodings/runend/src/compute/fill_null.rs index ddce31184eb..90284be3e51 100644 --- a/encodings/runend/src/compute/fill_null.rs +++ b/encodings/runend/src/compute/fill_null.rs @@ -11,6 +11,7 @@ use vortex_error::VortexResult; use crate::RunEnd; use crate::array::RunEndArrayExt; +use crate::array::RunEndArraySlotsExt; impl FillNullReduce for RunEnd { fn fill_null( diff --git a/encodings/runend/src/compute/filter.rs b/encodings/runend/src/compute/filter.rs index 6d8d4aab6ea..4575c03076e 100644 --- a/encodings/runend/src/compute/filter.rs +++ b/encodings/runend/src/compute/filter.rs @@ -23,6 +23,7 @@ use vortex_mask::Mask; use crate::RunEnd; use crate::array::RunEndArrayExt; +use crate::array::RunEndArraySlotsExt; use crate::compute::take::take_indices_unchecked; const FILTER_TAKE_THRESHOLD: f64 = 0.1; diff --git a/encodings/runend/src/compute/is_constant.rs b/encodings/runend/src/compute/is_constant.rs index d21f2334018..710fb7a3712 100644 --- a/encodings/runend/src/compute/is_constant.rs +++ b/encodings/runend/src/compute/is_constant.rs @@ -11,7 +11,7 @@ use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use crate::RunEnd; -use crate::array::RunEndArrayExt; +use crate::array::RunEndArraySlotsExt; /// RunEnd-specific is_constant kernel. /// diff --git a/encodings/runend/src/compute/is_sorted.rs b/encodings/runend/src/compute/is_sorted.rs index 65de2a174f2..17ae94272cb 100644 --- a/encodings/runend/src/compute/is_sorted.rs +++ b/encodings/runend/src/compute/is_sorted.rs @@ -14,7 +14,7 @@ use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use crate::RunEnd; -use crate::array::RunEndArrayExt; +use crate::array::RunEndArraySlotsExt; /// RunEnd-specific is_sorted kernel. /// diff --git a/encodings/runend/src/compute/min_max.rs b/encodings/runend/src/compute/min_max.rs index 4dd39edc5f3..d17c5b3c65c 100644 --- a/encodings/runend/src/compute/min_max.rs +++ b/encodings/runend/src/compute/min_max.rs @@ -12,7 +12,7 @@ use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use crate::RunEnd; -use crate::array::RunEndArrayExt; +use crate::array::RunEndArraySlotsExt; /// RunEnd-specific min/max kernel. /// diff --git a/encodings/runend/src/compute/take.rs b/encodings/runend/src/compute/take.rs index 4ec601e71a6..067bebef742 100644 --- a/encodings/runend/src/compute/take.rs +++ b/encodings/runend/src/compute/take.rs @@ -24,6 +24,7 @@ use vortex_mask::Mask; use crate::RunEnd; use crate::array::RunEndArrayExt; +use crate::array::RunEndArraySlotsExt; use crate::iter::trimmed_ends_iter; const SORTED_LINEAR_RUNS_PER_INDEX_THRESHOLD: usize = 16; diff --git a/encodings/runend/src/compute/take_from.rs b/encodings/runend/src/compute/take_from.rs index 0987cd3fc55..e1bd33a3f97 100644 --- a/encodings/runend/src/compute/take_from.rs +++ b/encodings/runend/src/compute/take_from.rs @@ -13,6 +13,7 @@ use vortex_error::VortexResult; use crate::RunEnd; use crate::array::RunEndArrayExt; +use crate::array::RunEndArraySlotsExt; #[derive(Debug)] pub(crate) struct RunEndTakeFrom; @@ -68,7 +69,7 @@ mod tests { use crate::RunEnd; use crate::RunEndArray; - use crate::array::RunEndArrayExt; + use crate::array::RunEndArraySlotsExt; use crate::compute::take_from::RunEndTakeFrom; static SESSION: LazyLock = LazyLock::new(|| { diff --git a/encodings/runend/src/kernel.rs b/encodings/runend/src/kernel.rs index 13f32df07a1..2eed964e3a9 100644 --- a/encodings/runend/src/kernel.rs +++ b/encodings/runend/src/kernel.rs @@ -24,6 +24,7 @@ use vortex_session::VortexSession; use crate::RunEnd; use crate::array::RunEndArrayExt; +use crate::array::RunEndArraySlotsExt; use crate::compute::take_from::RunEndTakeFrom; pub(super) fn initialize(session: &VortexSession) { diff --git a/encodings/runend/src/ops.rs b/encodings/runend/src/ops.rs index 39fb004c3e7..e2c2e3fc99b 100644 --- a/encodings/runend/src/ops.rs +++ b/encodings/runend/src/ops.rs @@ -15,6 +15,7 @@ use vortex_error::VortexResult; use crate::RunEnd; use crate::array::RunEndArrayExt; +use crate::array::RunEndArraySlotsExt; impl OperationsVTable for RunEnd { fn scalar_at( diff --git a/encodings/runend/src/rules.rs b/encodings/runend/src/rules.rs index ac7d6038146..a6f3d6a4ae3 100644 --- a/encodings/runend/src/rules.rs +++ b/encodings/runend/src/rules.rs @@ -19,6 +19,7 @@ use vortex_error::VortexResult; use crate::RunEnd; use crate::array::RunEndArrayExt; +use crate::array::RunEndArraySlotsExt; pub(super) const RULES: ParentRuleSet = ParentRuleSet::new(&[ // CastReduceAdaptor must come before RunEndScalarFnRule so that cast operations are executed diff --git a/encodings/sparse/src/canonical.rs b/encodings/sparse/src/canonical.rs index 0be6d0a4162..47a06ceb7e0 100644 --- a/encodings/sparse/src/canonical.rs +++ b/encodings/sparse/src/canonical.rs @@ -22,6 +22,7 @@ use vortex_array::arrays::VarBinView; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; use vortex_array::arrays::listview::ListViewArrayExt; +use vortex_array::arrays::listview::ListViewArraySlotsExt; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::arrays::varbinview::build_views::BinaryView; @@ -613,6 +614,7 @@ mod test { use vortex_array::arrays::VarBinArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::listview::ListViewArrayExt; + use vortex_array::arrays::listview::ListViewArraySlotsExt; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; diff --git a/fuzz/src/array/mask.rs b/fuzz/src/array/mask.rs index f3cecec2b10..53b1cd6ec67 100644 --- a/fuzz/src/array/mask.rs +++ b/fuzz/src/array/mask.rs @@ -18,7 +18,8 @@ use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::bool::BoolArrayExt; use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; -use vortex_array::arrays::listview::ListViewArrayExt; +use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use vortex_array::arrays::listview::ListViewArraySlotsExt; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::dtype::Nullability; use vortex_array::match_each_decimal_value_type; diff --git a/fuzz/src/array/slice.rs b/fuzz/src/array/slice.rs index 2a151002805..906200f1007 100644 --- a/fuzz/src/array/slice.rs +++ b/fuzz/src/array/slice.rs @@ -13,7 +13,8 @@ use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::bool::BoolArrayExt; use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; -use vortex_array::arrays::listview::ListViewArrayExt; +use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use vortex_array::arrays::listview::ListViewArraySlotsExt; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::dtype::DType; use vortex_array::match_each_decimal_value_type; diff --git a/vortex-array/benches/listview_rebuild.rs b/vortex-array/benches/listview_rebuild.rs index 2e61b969244..c8681dbed66 100644 --- a/vortex-array/benches/listview_rebuild.rs +++ b/vortex-array/benches/listview_rebuild.rs @@ -19,7 +19,7 @@ use vortex_array::arrays::ListViewArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinViewArray; -use vortex_array::arrays::listview::ListViewArrayExt; +use vortex_array::arrays::listview::ListViewArraySlotsExt; use vortex_array::arrays::listview::ListViewRebuildMode; use vortex_array::dtype::FieldNames; use vortex_array::validity::Validity; diff --git a/vortex-array/benches/take_fsl.rs b/vortex-array/benches/take_fsl.rs index 01386cbb9d9..96a6a96b054 100644 --- a/vortex-array/benches/take_fsl.rs +++ b/vortex-array/benches/take_fsl.rs @@ -30,7 +30,7 @@ use vortex_array::arrays::ConstantArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::PiecewiseSequenceArray; use vortex_array::arrays::PrimitiveArray; -use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; +use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use vortex_array::dtype::IntegerPType; use vortex_array::dtype::NativePType; use vortex_array::dtype::half::f16; diff --git a/vortex-array/src/aggregate_fn/accumulator_grouped.rs b/vortex-array/src/aggregate_fn/accumulator_grouped.rs index b87c04ee204..286f4194cd5 100644 --- a/vortex-array/src/aggregate_fn/accumulator_grouped.rs +++ b/vortex-array/src/aggregate_fn/accumulator_grouped.rs @@ -26,7 +26,8 @@ use crate::arrays::ChunkedArray; use crate::arrays::FixedSizeListArray; use crate::arrays::ListViewArray; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; -use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::builders::builder_with_capacity; use crate::builtins::ArrayBuiltins; use crate::columnar::AnyColumnar; diff --git a/vortex-array/src/aggregate_fn/fns/all_non_distinct/fixed_size_list.rs b/vortex-array/src/aggregate_fn/fns/all_non_distinct/fixed_size_list.rs index daaf9fe4cd7..7256d593029 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_distinct/fixed_size_list.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_distinct/fixed_size_list.rs @@ -8,7 +8,7 @@ use super::filter::filter_valid_rows_if_needed; use crate::ExecutionCtx; use crate::IntoArray; use crate::arrays::FixedSizeListArray; -use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; pub(super) fn check_fixed_size_list_identical( lhs: &FixedSizeListArray, diff --git a/vortex-array/src/aggregate_fn/fns/all_non_distinct/list.rs b/vortex-array/src/aggregate_fn/fns/all_non_distinct/list.rs index 2f4e3e52d91..89d1bb7c912 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_distinct/list.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_distinct/list.rs @@ -10,7 +10,9 @@ use crate::IntoArray; use crate::arrays::ListArray; use crate::arrays::ListViewArray; use crate::arrays::list::ListArrayExt; +use crate::arrays::list::ListArraySlotsExt; use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::listview::list_from_list_view; pub(super) fn check_list_identical( diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/list.rs b/vortex-array/src/aggregate_fn/fns/is_constant/list.rs index c33e325d463..ce684389dfd 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/list.rs +++ b/vortex-array/src/aggregate_fn/fns/is_constant/list.rs @@ -8,6 +8,7 @@ use super::is_constant; use crate::ExecutionCtx; use crate::arrays::ListViewArray; use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; /// Check if a list view array is constant by comparing each list's elements. /// diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/fixed_size_list.rs b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/fixed_size_list.rs index 7f8d55b07c1..03512cb8345 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/fixed_size_list.rs +++ b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/fixed_size_list.rs @@ -8,7 +8,7 @@ use super::uncompressed_size_in_bytes_u64; use super::validity_uncompressed_size_in_bytes; use crate::ExecutionCtx; use crate::arrays::FixedSizeListArray; -use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; pub(super) fn fixed_size_list_uncompressed_size_in_bytes( array: &FixedSizeListArray, diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/list_view.rs b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/list_view.rs index f75fcb92025..3021230d076 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/list_view.rs +++ b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/list_view.rs @@ -8,7 +8,7 @@ use super::uncompressed_size_in_bytes_u64; use super::validity_uncompressed_size_in_bytes; use crate::ExecutionCtx; use crate::arrays::ListViewArray; -use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::listview::ListViewRebuildMode; pub(super) fn list_view_uncompressed_size_in_bytes( diff --git a/vortex-array/src/arrays/chunked/vtable/canonical.rs b/vortex-array/src/arrays/chunked/vtable/canonical.rs index 12c3bbcd9e4..6eb968f8773 100644 --- a/vortex-array/src/arrays/chunked/vtable/canonical.rs +++ b/vortex-array/src/arrays/chunked/vtable/canonical.rs @@ -21,8 +21,8 @@ use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::arrays::VariantArray; use crate::arrays::chunked::ChunkedArrayExt; -use crate::arrays::fixed_size_list::FixedSizeListArrayExt; -use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::listview::ListViewRebuildMode; use crate::arrays::variant::VariantArrayExt; use crate::builders::builder_with_capacity_in; diff --git a/vortex-array/src/arrays/constant/vtable/canonical.rs b/vortex-array/src/arrays/constant/vtable/canonical.rs index e160bcc16a9..813657dbbe7 100644 --- a/vortex-array/src/arrays/constant/vtable/canonical.rs +++ b/vortex-array/src/arrays/constant/vtable/canonical.rs @@ -350,7 +350,8 @@ mod tests { use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; - use crate::arrays::listview::ListViewArrayExt; + use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; + use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::listview::ListViewRebuildMode; use crate::arrays::struct_::StructArrayExt; use crate::assert_arrays_eq; diff --git a/vortex-array/src/arrays/filter/execute/fixed_size_list.rs b/vortex-array/src/arrays/filter/execute/fixed_size_list.rs index 6facbe3e38d..33fdf62bfd6 100644 --- a/vortex-array/src/arrays/filter/execute/fixed_size_list.rs +++ b/vortex-array/src/arrays/filter/execute/fixed_size_list.rs @@ -11,6 +11,7 @@ use vortex_mask::MaskValues; use crate::arrays::FixedSizeListArray; use crate::arrays::filter::execute::filter_validity; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; /// Density threshold for choosing between indices and slices representation when expanding masks. /// diff --git a/vortex-array/src/arrays/filter/execute/listview.rs b/vortex-array/src/arrays/filter/execute/listview.rs index 6d1521541bc..db90f1e6680 100644 --- a/vortex-array/src/arrays/filter/execute/listview.rs +++ b/vortex-array/src/arrays/filter/execute/listview.rs @@ -9,7 +9,7 @@ use vortex_mask::MaskValues; use crate::arrays::ListViewArray; use crate::arrays::filter::execute::filter_validity; -use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; /// [`ListViewArray`] filter implementation. /// @@ -71,6 +71,7 @@ mod test { use crate::arrays::PrimitiveArray; use crate::arrays::filter::execute::ConstantArray; use crate::arrays::listview::ListViewArrayExt; + use crate::arrays::listview::ListViewArraySlotsExt; use crate::assert_arrays_eq; use crate::compute::conformance::filter::test_filter_conformance; use crate::validity::Validity; diff --git a/vortex-array/src/arrays/fixed_size_list/array.rs b/vortex-array/src/arrays/fixed_size_list/array.rs index b6328163727..6a17466a53e 100644 --- a/vortex-array/src/arrays/fixed_size_list/array.rs +++ b/vortex-array/src/arrays/fixed_size_list/array.rs @@ -5,7 +5,6 @@ use std::fmt::Display; use std::fmt::Formatter; use std::sync::Arc; -use smallvec::smallvec; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -17,24 +16,26 @@ use crate::array::ArrayParts; use crate::array::TypedArrayRef; use crate::array::child_to_validity; use crate::array::validity_to_child; +use crate::array_slots; use crate::arrays::FixedSizeList; use crate::dtype::DType; use crate::validity::Validity; -/// The `elements` data array, where each fixed-size list scalar is a _slice_ of the `elements` -/// array, and each inner list element is a _scalar_ of the `elements` array. -/// -/// The fixed-size list scalars are contiguous (regardless of nullability for easy lookups), -/// each with equal size in memory. -pub(super) const ELEMENTS_SLOT: usize = 0; -/// The validity / null map of the array. -/// -/// Note that this null map refers to which fixed-size list scalars are null, **not** which -/// sub-elements of fixed-size list scalars are null. The `elements` array will track individual -/// value nullability. -pub(super) const VALIDITY_SLOT: usize = 1; -pub(super) const NUM_SLOTS: usize = 2; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["elements", "validity"]; +#[array_slots(FixedSizeList)] +pub struct FixedSizeListSlots { + /// The `elements` data array, where each fixed-size list scalar is a _slice_ of the `elements` + /// array, and each inner list element is a _scalar_ of the `elements` array. + /// + /// The fixed-size list scalars are contiguous (regardless of nullability for easy lookups), + /// each with equal size in memory. + pub elements: ArrayRef, + /// The validity / null map of the array. + /// + /// Note that this null map refers to which fixed-size list scalars are null, **not** which + /// sub-elements of fixed-size list scalars are null. The `elements` array will track individual + /// value nullability. + pub validity: Option, +} /// The canonical encoding for fixed-size list arrays. /// @@ -110,7 +111,11 @@ pub struct FixedSizeListDataParts { impl FixedSizeListData { pub(crate) fn make_slots(elements: &ArrayRef, validity: &Validity, len: usize) -> ArraySlots { - smallvec![Some(elements.clone()), validity_to_child(validity, len)] + FixedSizeListSlots { + elements: elements.clone(), + validity: validity_to_child(validity, len), + } + .into_slots() } /// Creates a new `FixedSizeListArray`. @@ -203,7 +208,7 @@ impl FixedSizeListData { } } -pub trait FixedSizeListArrayExt: TypedArrayRef { +pub trait FixedSizeListArrayExt: FixedSizeListArraySlotsExt { fn dtype_parts(&self) -> (&DType, u32, crate::dtype::Nullability) { match self.as_ref().dtype() { DType::FixedSizeList(element_dtype, list_size, nullability) => { @@ -213,12 +218,6 @@ pub trait FixedSizeListArrayExt: TypedArrayRef { } } - fn elements(&self) -> &ArrayRef { - self.as_ref().slots()[ELEMENTS_SLOT] - .as_ref() - .vortex_expect("FixedSizeListArray elements slot") - } - fn list_size(&self) -> u32 { let (_, list_size, _) = self.dtype_parts(); list_size @@ -226,7 +225,7 @@ pub trait FixedSizeListArrayExt: TypedArrayRef { fn fixed_size_list_validity(&self) -> Validity { let (_, _, nullability) = self.dtype_parts(); - child_to_validity(self.as_ref().slots()[VALIDITY_SLOT].as_ref(), nullability) + child_to_validity(self.validity(), nullability) } #[allow(clippy::disallowed_methods)] @@ -310,10 +309,10 @@ impl Array { pub fn into_data_parts(self) -> FixedSizeListDataParts { let dtype = self.dtype().clone(); - let elements = self.slots()[ELEMENTS_SLOT] - .clone() - .vortex_expect("FixedSizeListArray elements slot"); let validity = self.fixed_size_list_validity(); + let elements = FixedSizeListSlotsView::from_slots(self.slots()) + .elements + .clone(); FixedSizeListDataParts { elements, validity, diff --git a/vortex-array/src/arrays/fixed_size_list/compute/cast.rs b/vortex-array/src/arrays/fixed_size_list/compute/cast.rs index 4ac3bd1938f..8e813209a5b 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/cast.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/cast.rs @@ -10,6 +10,7 @@ use crate::array::ArrayView; use crate::arrays::FixedSizeList; use crate::arrays::FixedSizeListArray; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::scalar_fn::fns::cast::CastKernel; diff --git a/vortex-array/src/arrays/fixed_size_list/compute/mask.rs b/vortex-array/src/arrays/fixed_size_list/compute/mask.rs index 3df0413e1aa..1db564552e7 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/mask.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/mask.rs @@ -9,6 +9,7 @@ use crate::array::ArrayView; use crate::arrays::FixedSizeList; use crate::arrays::FixedSizeListArray; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use crate::scalar_fn::fns::mask::MaskReduce; use crate::validity::Validity; diff --git a/vortex-array/src/arrays/fixed_size_list/compute/slice.rs b/vortex-array/src/arrays/fixed_size_list/compute/slice.rs index 5146fd2b600..12f67334153 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/slice.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/slice.rs @@ -11,6 +11,7 @@ use crate::array::ArrayView; use crate::arrays::FixedSizeList; use crate::arrays::FixedSizeListArray; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use crate::arrays::slice::SliceReduce; impl SliceReduce for FixedSizeList { diff --git a/vortex-array/src/arrays/fixed_size_list/compute/take.rs b/vortex-array/src/arrays/fixed_size_list/compute/take.rs index 5f2e0a50835..4e41ffd7ad1 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/take.rs @@ -20,6 +20,7 @@ use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::dict::TakeExecute; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use crate::arrays::primitive::PrimitiveArrayExt; use crate::builders::builder_with_capacity; use crate::dtype::DType; diff --git a/vortex-array/src/arrays/fixed_size_list/mod.rs b/vortex-array/src/arrays/fixed_size_list/mod.rs index 3d90d38e6e9..adf5c9bf167 100644 --- a/vortex-array/src/arrays/fixed_size_list/mod.rs +++ b/vortex-array/src/arrays/fixed_size_list/mod.rs @@ -3,8 +3,11 @@ mod array; pub use array::FixedSizeListArrayExt; +pub use array::FixedSizeListArraySlotsExt; pub use array::FixedSizeListData; pub use array::FixedSizeListDataParts; +pub use array::FixedSizeListSlots; +pub use array::FixedSizeListSlotsView; pub use vtable::FixedSizeListArray; pub(crate) mod compute; diff --git a/vortex-array/src/arrays/fixed_size_list/tests/basic.rs b/vortex-array/src/arrays/fixed_size_list/tests/basic.rs index 961801094d8..b471ebb07a0 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/basic.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/basic.rs @@ -10,6 +10,7 @@ use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::FixedSizeListArray; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; diff --git a/vortex-array/src/arrays/fixed_size_list/tests/degenerate.rs b/vortex-array/src/arrays/fixed_size_list/tests/degenerate.rs index f609fd72405..73085b8c8a8 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/degenerate.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/degenerate.rs @@ -11,6 +11,7 @@ use crate::array_session; use crate::arrays::FixedSizeListArray; use crate::arrays::PrimitiveArray; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; diff --git a/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs b/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs index 9393b8563d4..eb2ee610021 100644 --- a/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs +++ b/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs @@ -5,7 +5,6 @@ use std::hash::Hash; use std::hash::Hasher; use std::sync::Arc; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -26,9 +25,8 @@ use crate::array::ArrayView; use crate::array::VTable; use crate::array::with_empty_buffers; use crate::arrays::fixed_size_list::FixedSizeListData; -use crate::arrays::fixed_size_list::array::ELEMENTS_SLOT; -use crate::arrays::fixed_size_list::array::NUM_SLOTS; -use crate::arrays::fixed_size_list::array::SLOT_NAMES; +use crate::arrays::fixed_size_list::array::FixedSizeListSlots; +use crate::arrays::fixed_size_list::array::FixedSizeListSlotsView; use crate::arrays::fixed_size_list::compute::rules::PARENT_RULES; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; @@ -115,16 +113,15 @@ impl VTable for FixedSizeList { slots: &[Option], ) -> VortexResult<()> { vortex_ensure!( - slots.len() == NUM_SLOTS, - "FixedSizeListArray expected {NUM_SLOTS} slots, found {}", + slots.len() == FixedSizeListSlots::COUNT, + "FixedSizeListArray expected {} slots, found {}", + FixedSizeListSlots::COUNT, slots.len() ); let DType::FixedSizeList(_, list_size, nullability) = dtype else { vortex_bail!("Expected `DType::FixedSizeList`, got {dtype:?}"); }; - let elements = slots[ELEMENTS_SLOT] - .as_ref() - .vortex_expect("FixedSizeListArray elements slot"); + let elements = FixedSizeListSlotsView::from_slots(slots).elements; vortex_ensure!( if *list_size == 0 { data.degenerate_len == len @@ -197,7 +194,7 @@ impl VTable for FixedSizeList { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + FixedSizeListSlots::NAMES[idx].to_string() } fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { diff --git a/vortex-array/src/arrays/list/array.rs b/vortex-array/src/arrays/list/array.rs index 8350b1a3558..e64ec9a223f 100644 --- a/vortex-array/src/arrays/list/array.rs +++ b/vortex-array/src/arrays/list/array.rs @@ -6,7 +6,6 @@ use std::fmt::Formatter; use std::sync::Arc; use num_traits::AsPrimitive; -use smallvec::smallvec; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -26,6 +25,7 @@ use crate::array::ArrayParts; use crate::array::TypedArrayRef; use crate::array::child_to_validity; use crate::array::validity_to_child; +use crate::array_slots; use crate::arrays::ConstantArray; use crate::arrays::List; use crate::arrays::ListArray; @@ -39,14 +39,15 @@ use crate::match_each_native_ptype; use crate::scalar_fn::fns::operators::Operator; use crate::validity::Validity; -/// The elements data array containing all list elements concatenated together. -pub(super) const ELEMENTS_SLOT: usize = 0; -/// The offsets array defining the start/end of each list within the elements array. -pub(super) const OFFSETS_SLOT: usize = 1; -/// The validity bitmap indicating which list elements are non-null. -pub(super) const VALIDITY_SLOT: usize = 2; -pub(super) const NUM_SLOTS: usize = 3; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["elements", "offsets", "validity"]; +#[array_slots(List)] +pub struct ListSlots { + /// The elements data array containing all list elements concatenated together. + pub elements: ArrayRef, + /// The offsets array defining the start/end of each list within the elements array. + pub offsets: ArrayRef, + /// The validity bitmap indicating which lists are non-null. + pub validity: Option, +} /// A list array that stores variable-length lists of elements, similar to `Vec>`. /// @@ -124,11 +125,12 @@ impl ListData { validity: &Validity, len: usize, ) -> ArraySlots { - smallvec![ - Some(elements.clone()), - Some(offsets.clone()), - validity_to_child(validity, len), - ] + ListSlots { + elements: elements.clone(), + offsets: offsets.clone(), + validity: validity_to_child(validity, len), + } + .into_slots() } /// Creates a new `ListArray`. @@ -271,7 +273,7 @@ impl ListData { // the offset type and manual subtraction and fast path where `offsets[0] == 0`. } -pub trait ListArrayExt: TypedArrayRef { +pub trait ListArrayExt: ListArraySlotsExt { fn nullability(&self) -> crate::dtype::Nullability { match self.as_ref().dtype() { DType::List(_, nullability) => *nullability, @@ -279,23 +281,8 @@ pub trait ListArrayExt: TypedArrayRef { } } - fn elements(&self) -> &ArrayRef { - self.as_ref().slots()[ELEMENTS_SLOT] - .as_ref() - .vortex_expect("ListArray elements slot") - } - - fn offsets(&self) -> &ArrayRef { - self.as_ref().slots()[OFFSETS_SLOT] - .as_ref() - .vortex_expect("ListArray offsets slot") - } - fn list_validity(&self) -> Validity { - child_to_validity( - self.as_ref().slots()[VALIDITY_SLOT].as_ref(), - self.nullability(), - ) + child_to_validity(self.validity(), self.nullability()) } #[allow(clippy::disallowed_methods)] @@ -407,16 +394,11 @@ impl Array { pub fn into_data_parts(self) -> ListDataParts { let dtype = self.dtype().clone(); - let elements = self.slots()[ELEMENTS_SLOT] - .clone() - .vortex_expect("ListArray elements slot"); - let offsets = self.slots()[OFFSETS_SLOT] - .clone() - .vortex_expect("ListArray offsets slot"); let validity = self.list_validity(); + let view = ListSlotsView::from_slots(self.slots()); ListDataParts { - elements, - offsets, + elements: view.elements.clone(), + offsets: view.offsets.clone(), validity, dtype, } diff --git a/vortex-array/src/arrays/list/compute/cast.rs b/vortex-array/src/arrays/list/compute/cast.rs index 9fd27f731f1..930b6e3b8f4 100644 --- a/vortex-array/src/arrays/list/compute/cast.rs +++ b/vortex-array/src/arrays/list/compute/cast.rs @@ -9,7 +9,7 @@ use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::List; use crate::arrays::ListArray; -use crate::arrays::list::ListArrayExt; +use crate::arrays::list::ListArraySlotsExt; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::scalar_fn::fns::cast::CastKernel; diff --git a/vortex-array/src/arrays/list/compute/filter.rs b/vortex-array/src/arrays/list/compute/filter.rs index 2b98e3c36d1..af02ce23d7b 100644 --- a/vortex-array/src/arrays/list/compute/filter.rs +++ b/vortex-array/src/arrays/list/compute/filter.rs @@ -22,6 +22,7 @@ use crate::arrays::List; use crate::arrays::ListArray; use crate::arrays::filter::FilterKernel; use crate::arrays::list::ListArrayExt; +use crate::arrays::list::ListArraySlotsExt; use crate::dtype::IntegerPType; use crate::match_each_integer_ptype; use crate::validity::Validity; diff --git a/vortex-array/src/arrays/list/compute/mask.rs b/vortex-array/src/arrays/list/compute/mask.rs index e9e895d08f1..f2727247646 100644 --- a/vortex-array/src/arrays/list/compute/mask.rs +++ b/vortex-array/src/arrays/list/compute/mask.rs @@ -8,7 +8,7 @@ use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::List; use crate::arrays::ListArray; -use crate::arrays::list::ListArrayExt; +use crate::arrays::list::ListArraySlotsExt; use crate::scalar_fn::fns::mask::MaskReduce; use crate::validity::Validity; diff --git a/vortex-array/src/arrays/list/compute/slice.rs b/vortex-array/src/arrays/list/compute/slice.rs index d29eaa5184c..f42d490237e 100644 --- a/vortex-array/src/arrays/list/compute/slice.rs +++ b/vortex-array/src/arrays/list/compute/slice.rs @@ -10,7 +10,7 @@ use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::List; use crate::arrays::ListArray; -use crate::arrays::list::ListArrayExt; +use crate::arrays::list::ListArraySlotsExt; use crate::arrays::slice::SliceReduce; impl SliceReduce for List { diff --git a/vortex-array/src/arrays/list/compute/take.rs b/vortex-array/src/arrays/list/compute/take.rs index 6d71d382f0d..c010b0d2732 100644 --- a/vortex-array/src/arrays/list/compute/take.rs +++ b/vortex-array/src/arrays/list/compute/take.rs @@ -22,6 +22,7 @@ use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::dict::TakeExecute; use crate::arrays::list::ListArrayExt; +use crate::arrays::list::ListArraySlotsExt; use crate::arrays::piecewise_sequence::constant_unsigned_usize; use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::arrays::primitive::PrimitiveArrayExt; diff --git a/vortex-array/src/arrays/list/mod.rs b/vortex-array/src/arrays/list/mod.rs index 422a9523ff9..bfd43cdb401 100644 --- a/vortex-array/src/arrays/list/mod.rs +++ b/vortex-array/src/arrays/list/mod.rs @@ -3,8 +3,11 @@ mod array; pub use array::ListArrayExt; +pub use array::ListArraySlotsExt; pub use array::ListData; pub use array::ListDataParts; +pub use array::ListSlots; +pub use array::ListSlotsView; pub use vtable::ListArray; pub(crate) mod compute; diff --git a/vortex-array/src/arrays/list/vtable/mod.rs b/vortex-array/src/arrays/list/vtable/mod.rs index 55f7753c995..cdc35ac6548 100644 --- a/vortex-array/src/arrays/list/vtable/mod.rs +++ b/vortex-array/src/arrays/list/vtable/mod.rs @@ -5,7 +5,6 @@ use std::hash::Hasher; use std::sync::Arc; use prost::Message; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -26,12 +25,10 @@ use crate::array::ArrayParts; use crate::array::ArrayView; use crate::array::VTable; use crate::array::with_empty_buffers; -use crate::arrays::list::ListArrayExt; +use crate::arrays::list::ListArraySlotsExt; use crate::arrays::list::ListData; -use crate::arrays::list::array::ELEMENTS_SLOT; -use crate::arrays::list::array::NUM_SLOTS; -use crate::arrays::list::array::OFFSETS_SLOT; -use crate::arrays::list::array::SLOT_NAMES; +use crate::arrays::list::array::ListSlots; +use crate::arrays::list::array::ListSlotsView; use crate::arrays::list::compute::rules::PARENT_RULES; use crate::arrays::listview::list_view_from_list; use crate::buffer::BufferHandle; @@ -123,16 +120,14 @@ impl VTable for List { slots: &[Option], ) -> VortexResult<()> { vortex_ensure!( - slots.len() == NUM_SLOTS, - "ListArray expected {NUM_SLOTS} slots, found {}", + slots.len() == ListSlots::COUNT, + "ListArray expected {} slots, found {}", + ListSlots::COUNT, slots.len() ); - let elements = slots[ELEMENTS_SLOT] - .as_ref() - .vortex_expect("ListArray elements slot"); - let offsets = slots[OFFSETS_SLOT] - .as_ref() - .vortex_expect("ListArray offsets slot"); + let ListSlotsView { + elements, offsets, .. + } = ListSlotsView::from_slots(slots); vortex_ensure!( offsets.len().saturating_sub(1) == len, "ListArray length {} does not match outer length {}", @@ -192,7 +187,7 @@ impl VTable for List { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + ListSlots::NAMES[idx].to_string() } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { diff --git a/vortex-array/src/arrays/listview/array.rs b/vortex-array/src/arrays/listview/array.rs index daf7100dc59..13f5f620883 100644 --- a/vortex-array/src/arrays/listview/array.rs +++ b/vortex-array/src/arrays/listview/array.rs @@ -6,7 +6,6 @@ use std::fmt::Formatter; use std::sync::Arc; use num_traits::AsPrimitive; -use smallvec::smallvec; use vortex_buffer::BitBufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -26,6 +25,7 @@ use crate::array::ArrayParts; use crate::array::TypedArrayRef; use crate::array::child_to_validity; use crate::array::validity_to_child; +use crate::array_slots; use crate::arrays::ListView; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; @@ -42,23 +42,24 @@ use crate::match_each_unsigned_integer_ptype; use crate::scalar_fn::fns::operators::Operator; use crate::validity::Validity; -/// The `elements` data array, where each list scalar is a _slice_ of the `elements` array, and -/// each inner list element is a _scalar_ of the `elements` array. -pub(super) const ELEMENTS_SLOT: usize = 0; -/// The `offsets` array indicating the start position of each list in elements. -/// -/// Since we also store `sizes`, this `offsets` field is allowed to be stored out-of-order -/// (which is different from [`ListArray`](crate::arrays::ListArray)). -pub(super) const OFFSETS_SLOT: usize = 1; -/// The `sizes` array indicating the length of each list. -/// -/// This field is intended to be paired with a corresponding offset to determine the list scalar -/// we want to access. -pub(super) const SIZES_SLOT: usize = 2; -/// The validity bitmap indicating which list elements are non-null. -pub(super) const VALIDITY_SLOT: usize = 3; -pub(super) const NUM_SLOTS: usize = 4; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["elements", "offsets", "sizes", "validity"]; +#[array_slots(ListView)] +pub struct ListViewSlots { + /// The `elements` data array, where each list scalar is a _slice_ of the `elements` array, and + /// each inner list element is a _scalar_ of the `elements` array. + pub elements: ArrayRef, + /// The `offsets` array indicating the start position of each list in elements. + /// + /// Since we also store `sizes`, this `offsets` field is allowed to be stored out-of-order + /// (which is different from [`ListArray`](crate::arrays::ListArray)). + pub offsets: ArrayRef, + /// The `sizes` array indicating the length of each list. + /// + /// This field is intended to be paired with a corresponding offset to determine the list scalar + /// we want to access. + pub sizes: ArrayRef, + /// The validity bitmap indicating which lists are non-null. + pub validity: Option, +} /// The canonical encoding for variable-length list arrays. /// @@ -166,12 +167,13 @@ impl ListViewData { validity: &Validity, len: usize, ) -> ArraySlots { - smallvec![ - Some(elements.clone()), - Some(offsets.clone()), - Some(sizes.clone()), - validity_to_child(validity, len), - ] + ListViewSlots { + elements: elements.clone(), + offsets: offsets.clone(), + sizes: sizes.clone(), + validity: validity_to_child(validity, len), + } + .into_slots() } /// Creates a new `ListViewArray`. @@ -347,7 +349,7 @@ fn fill_referenced_mask( } } -pub trait ListViewArrayExt: TypedArrayRef { +pub trait ListViewArrayExt: ListViewArraySlotsExt { fn nullability(&self) -> crate::dtype::Nullability { match self.as_ref().dtype() { DType::List(_, nullability) => *nullability, @@ -355,29 +357,8 @@ pub trait ListViewArrayExt: TypedArrayRef { } } - fn elements(&self) -> &ArrayRef { - self.as_ref().slots()[ELEMENTS_SLOT] - .as_ref() - .vortex_expect("ListViewArray elements slot") - } - - fn offsets(&self) -> &ArrayRef { - self.as_ref().slots()[OFFSETS_SLOT] - .as_ref() - .vortex_expect("ListViewArray offsets slot") - } - - fn sizes(&self) -> &ArrayRef { - self.as_ref().slots()[SIZES_SLOT] - .as_ref() - .vortex_expect("ListViewArray sizes slot") - } - fn listview_validity(&self) -> Validity { - child_to_validity( - self.as_ref().slots()[VALIDITY_SLOT].as_ref(), - self.nullability(), - ) + child_to_validity(self.validity(), self.nullability()) } #[allow(clippy::disallowed_methods)] @@ -667,21 +648,13 @@ impl Array { } pub fn into_data_parts(self) -> ListViewDataParts { - let elements = self.slots()[ELEMENTS_SLOT] - .clone() - .vortex_expect("ListViewArray elements slot"); - let offsets = self.slots()[OFFSETS_SLOT] - .clone() - .vortex_expect("ListViewArray offsets slot"); - let sizes = self.slots()[SIZES_SLOT] - .clone() - .vortex_expect("ListViewArray sizes slot"); let validity = self.listview_validity(); + let view = ListViewSlotsView::from_slots(self.slots()); ListViewDataParts { - elements_dtype: Arc::new(elements.dtype().clone()), - elements, - offsets, - sizes, + elements_dtype: Arc::new(view.elements.dtype().clone()), + elements: view.elements.clone(), + offsets: view.offsets.clone(), + sizes: view.sizes.clone(), validity, } } diff --git a/vortex-array/src/arrays/listview/compute/cast.rs b/vortex-array/src/arrays/listview/compute/cast.rs index 7f4c972940f..1c5e1fdfd84 100644 --- a/vortex-array/src/arrays/listview/compute/cast.rs +++ b/vortex-array/src/arrays/listview/compute/cast.rs @@ -9,7 +9,7 @@ use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::ListView; use crate::arrays::ListViewArray; -use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::scalar_fn::fns::cast::CastKernel; diff --git a/vortex-array/src/arrays/listview/compute/mask.rs b/vortex-array/src/arrays/listview/compute/mask.rs index e14233a71d0..2b5a89132fb 100644 --- a/vortex-array/src/arrays/listview/compute/mask.rs +++ b/vortex-array/src/arrays/listview/compute/mask.rs @@ -8,7 +8,7 @@ use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::ListView; use crate::arrays::ListViewArray; -use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::scalar_fn::fns::mask::MaskReduce; use crate::validity::Validity; diff --git a/vortex-array/src/arrays/listview/compute/rules.rs b/vortex-array/src/arrays/listview/compute/rules.rs index 9b9396d7fe7..233533202a8 100644 --- a/vortex-array/src/arrays/listview/compute/rules.rs +++ b/vortex-array/src/arrays/listview/compute/rules.rs @@ -10,7 +10,7 @@ use crate::arrays::Filter; use crate::arrays::ListView; use crate::arrays::ListViewArray; use crate::arrays::dict::TakeReduceAdaptor; -use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::slice::SliceReduceAdaptor; use crate::optimizer::rules::ArrayParentReduceRule; use crate::optimizer::rules::ParentRuleSet; diff --git a/vortex-array/src/arrays/listview/compute/slice.rs b/vortex-array/src/arrays/listview/compute/slice.rs index 2dd29278392..8cda72b871b 100644 --- a/vortex-array/src/arrays/listview/compute/slice.rs +++ b/vortex-array/src/arrays/listview/compute/slice.rs @@ -10,7 +10,7 @@ use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::ListView; use crate::arrays::ListViewArray; -use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::slice::SliceReduce; impl SliceReduce for ListView { diff --git a/vortex-array/src/arrays/listview/compute/take.rs b/vortex-array/src/arrays/listview/compute/take.rs index 2b6c016d2c3..f5ffe321821 100644 --- a/vortex-array/src/arrays/listview/compute/take.rs +++ b/vortex-array/src/arrays/listview/compute/take.rs @@ -12,7 +12,7 @@ use crate::arrays::ListView; use crate::arrays::ListViewArray; use crate::arrays::dict::TakeExecute; use crate::arrays::dict::TakeReduce; -use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::builtins::ArrayBuiltins; use crate::dtype::Nullability; use crate::match_each_integer_ptype; diff --git a/vortex-array/src/arrays/listview/compute/zip.rs b/vortex-array/src/arrays/listview/compute/zip.rs index 71640d4fb21..d5e2a6bf406 100644 --- a/vortex-array/src/arrays/listview/compute/zip.rs +++ b/vortex-array/src/arrays/listview/compute/zip.rs @@ -21,7 +21,7 @@ use crate::arrays::ChunkedArray; use crate::arrays::ListView; use crate::arrays::ListViewArray; use crate::arrays::chunked::ChunkedArrayExt; -use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::Nullability; @@ -251,7 +251,7 @@ mod tests { use crate::arrays::ListView; use crate::arrays::ListViewArray; use crate::arrays::chunked::ChunkedArrayExt; - use crate::arrays::listview::ListViewArrayExt; + use crate::arrays::listview::ListViewArraySlotsExt; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; diff --git a/vortex-array/src/arrays/listview/conversion.rs b/vortex-array/src/arrays/listview/conversion.rs index e85567a9ed9..054f3efecb9 100644 --- a/vortex-array/src/arrays/listview/conversion.rs +++ b/vortex-array/src/arrays/listview/conversion.rs @@ -16,8 +16,11 @@ use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::arrays::extension::ExtensionArrayExt; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use crate::arrays::list::ListArrayExt; +use crate::arrays::list::ListArraySlotsExt; use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::listview::ListViewRebuildMode; use crate::arrays::struct_::StructArrayExt; use crate::builders::PrimitiveBuilder; @@ -312,7 +315,8 @@ mod tests { use crate::arrays::StructArray; use crate::arrays::VarBinViewArray; use crate::arrays::list::ListArrayExt; - use crate::arrays::listview::ListViewArrayExt; + use crate::arrays::list::ListArraySlotsExt; + use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::listview::list_from_list_view; use crate::arrays::listview::list_view_from_list; use crate::assert_arrays_eq; diff --git a/vortex-array/src/arrays/listview/mod.rs b/vortex-array/src/arrays/listview/mod.rs index b240854c6cf..9955ae0f5fd 100644 --- a/vortex-array/src/arrays/listview/mod.rs +++ b/vortex-array/src/arrays/listview/mod.rs @@ -3,8 +3,11 @@ mod array; pub use array::ListViewArrayExt; +pub use array::ListViewArraySlotsExt; pub use array::ListViewData; pub use array::ListViewDataParts; +pub use array::ListViewSlots; +pub use array::ListViewSlotsView; pub use vtable::ListViewArray; pub(crate) mod compute; diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index cee2972ec5d..015a91cb5df 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -15,6 +15,7 @@ use crate::arrays::ListViewArray; use crate::arrays::PiecewiseSequenceArray; use crate::arrays::PrimitiveArray; use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::primitive::PrimitiveArrayExt; use crate::builtins::ArrayBuiltins; use crate::dtype::IntegerPType; @@ -459,6 +460,7 @@ mod tests { use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::arrays::listview::ListViewArrayExt; + use crate::arrays::listview::ListViewArraySlotsExt; use crate::assert_arrays_eq; use crate::dtype::Nullability; use crate::validity::Validity; diff --git a/vortex-array/src/arrays/listview/tests/density.rs b/vortex-array/src/arrays/listview/tests/density.rs index 8538a9fe20c..6bff870a1b6 100644 --- a/vortex-array/src/arrays/listview/tests/density.rs +++ b/vortex-array/src/arrays/listview/tests/density.rs @@ -18,6 +18,7 @@ use crate::IntoArray; use crate::VortexSessionExecute; use crate::arrays::ListViewArray; use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::listview::tests::common::create_empty_elements_listview; use crate::expr::stats::Precision; use crate::expr::stats::Stat; diff --git a/vortex-array/src/arrays/listview/tests/filter.rs b/vortex-array/src/arrays/listview/tests/filter.rs index e8cc7636b30..d8a6cb4f4db 100644 --- a/vortex-array/src/arrays/listview/tests/filter.rs +++ b/vortex-array/src/arrays/listview/tests/filter.rs @@ -19,6 +19,7 @@ use crate::arrays::ConstantArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::assert_arrays_eq; use crate::compute::conformance::filter::test_filter_conformance; use crate::validity::Validity; diff --git a/vortex-array/src/arrays/listview/tests/nullability.rs b/vortex-array/src/arrays/listview/tests/nullability.rs index 4889b73de21..b32e2052435 100644 --- a/vortex-array/src/arrays/listview/tests/nullability.rs +++ b/vortex-array/src/arrays/listview/tests/nullability.rs @@ -13,6 +13,7 @@ use crate::arrays::BoolArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; diff --git a/vortex-array/src/arrays/listview/tests/operations.rs b/vortex-array/src/arrays/listview/tests/operations.rs index 51f03b1abeb..38ac286b2d1 100644 --- a/vortex-array/src/arrays/listview/tests/operations.rs +++ b/vortex-array/src/arrays/listview/tests/operations.rs @@ -22,6 +22,7 @@ use crate::arrays::ListView; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; use crate::compute::conformance::mask::test_mask_conformance; diff --git a/vortex-array/src/arrays/listview/tests/take.rs b/vortex-array/src/arrays/listview/tests/take.rs index 2af397e262f..6f74081a8c5 100644 --- a/vortex-array/src/arrays/listview/tests/take.rs +++ b/vortex-array/src/arrays/listview/tests/take.rs @@ -18,6 +18,7 @@ use crate::arrays::ConstantArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::assert_arrays_eq; use crate::compute::conformance::take::test_take_conformance; use crate::validity::Validity; diff --git a/vortex-array/src/arrays/listview/vtable/mod.rs b/vortex-array/src/arrays/listview/vtable/mod.rs index 546e77898f8..87afcdc87b2 100644 --- a/vortex-array/src/arrays/listview/vtable/mod.rs +++ b/vortex-array/src/arrays/listview/vtable/mod.rs @@ -6,7 +6,6 @@ use std::hash::Hasher; use std::sync::Arc; use prost::Message; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -26,13 +25,10 @@ use crate::array::ArrayId; use crate::array::ArrayView; use crate::array::VTable; use crate::array::with_empty_buffers; -use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::listview::ListViewData; -use crate::arrays::listview::array::ELEMENTS_SLOT; -use crate::arrays::listview::array::NUM_SLOTS; -use crate::arrays::listview::array::OFFSETS_SLOT; -use crate::arrays::listview::array::SIZES_SLOT; -use crate::arrays::listview::array::SLOT_NAMES; +use crate::arrays::listview::array::ListViewSlots; +use crate::arrays::listview::array::ListViewSlotsView; use crate::arrays::listview::compute::rules::PARENT_RULES; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; @@ -128,19 +124,17 @@ impl VTable for ListView { slots: &[Option], ) -> VortexResult<()> { vortex_ensure!( - slots.len() == NUM_SLOTS, - "ListViewArray expected {NUM_SLOTS} slots, found {}", + slots.len() == ListViewSlots::COUNT, + "ListViewArray expected {} slots, found {}", + ListViewSlots::COUNT, slots.len() ); - let elements = slots[ELEMENTS_SLOT] - .as_ref() - .vortex_expect("ListViewArray elements slot"); - let offsets = slots[OFFSETS_SLOT] - .as_ref() - .vortex_expect("ListViewArray offsets slot"); - let sizes = slots[SIZES_SLOT] - .as_ref() - .vortex_expect("ListViewArray sizes slot"); + let ListViewSlotsView { + elements, + offsets, + sizes, + .. + } = ListViewSlotsView::from_slots(slots); vortex_ensure!( offsets.len() == len && sizes.len() == len, "ListViewArray length {} does not match outer length {}", @@ -219,7 +213,7 @@ impl VTable for ListView { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + ListViewSlots::NAMES[idx].to_string() } fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { diff --git a/vortex-array/src/arrays/masked/execute.rs b/vortex-array/src/arrays/masked/execute.rs index 0de8be7321e..401f3b160c7 100644 --- a/vortex-array/src/arrays/masked/execute.rs +++ b/vortex-array/src/arrays/masked/execute.rs @@ -22,7 +22,8 @@ use crate::arrays::VariantArray; use crate::arrays::bool::BoolArrayExt; use crate::arrays::extension::ExtensionArrayExt; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; -use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::struct_::StructArrayExt; use crate::arrays::variant::VariantArrayExt; use crate::builtins::ArrayBuiltins; diff --git a/vortex-array/src/arrays/varbin/array.rs b/vortex-array/src/arrays/varbin/array.rs index 4cb74fe48b4..d0dd44e5d71 100644 --- a/vortex-array/src/arrays/varbin/array.rs +++ b/vortex-array/src/arrays/varbin/array.rs @@ -5,7 +5,6 @@ use std::fmt::Display; use std::fmt::Formatter; use num_traits::AsPrimitive; -use smallvec::smallvec; use vortex_array::arrays::PrimitiveArray; use vortex_buffer::ByteBuffer; use vortex_error::VortexExpect; @@ -21,6 +20,7 @@ use crate::array::ArrayParts; use crate::array::TypedArrayRef; use crate::array::child_to_validity; use crate::array::validity_to_child; +use crate::array_slots; use crate::arrays::VarBin; use crate::arrays::varbin::builder::VarBinBuilder; use crate::buffer::BufferHandle; @@ -31,12 +31,13 @@ use crate::legacy_session; use crate::match_each_integer_ptype; use crate::validity::Validity; -/// The offsets array defining the start/end of each variable-length binary element. -pub(super) const OFFSETS_SLOT: usize = 0; -/// The validity bitmap indicating which elements are non-null. -pub(super) const VALIDITY_SLOT: usize = 1; -pub(super) const NUM_SLOTS: usize = 2; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["offsets", "validity"]; +#[array_slots(VarBin)] +pub struct VarBinSlots { + /// The offsets array defining the start/end of each variable-length binary element. + pub offsets: ArrayRef, + /// The validity bitmap indicating which elements are non-null. + pub validity: Option, +} #[derive(Clone, Debug)] pub struct VarBinData { @@ -83,7 +84,11 @@ impl VarBinData { } pub(crate) fn make_slots(offsets: ArrayRef, validity: &Validity, len: usize) -> ArraySlots { - smallvec![Some(offsets), validity_to_child(validity, len)] + VarBinSlots { + offsets, + validity: validity_to_child(validity, len), + } + .into_slots() } /// Constructs a new `VarBinArray`. @@ -302,17 +307,7 @@ impl VarBinData { } } -pub trait VarBinArrayExt: TypedArrayRef { - fn offsets(&self) -> &ArrayRef { - self.as_ref().slots()[OFFSETS_SLOT] - .as_ref() - .vortex_expect("VarBinArray offsets slot") - } - - fn validity_child(&self) -> Option<&ArrayRef> { - self.as_ref().slots()[VALIDITY_SLOT].as_ref() - } - +pub trait VarBinArrayExt: VarBinArraySlotsExt { fn dtype_parts(&self) -> (bool, Nullability) { match self.as_ref().dtype() { DType::Utf8(nullability) => (true, *nullability), @@ -330,10 +325,7 @@ pub trait VarBinArrayExt: TypedArrayRef { } fn varbin_validity(&self) -> Validity { - child_to_validity( - self.as_ref().slots()[VALIDITY_SLOT].as_ref(), - self.nullability(), - ) + child_to_validity(self.validity(), self.nullability()) } #[allow(clippy::disallowed_methods)] @@ -457,10 +449,7 @@ impl Array { let len = offsets.len().saturating_sub(1); let slots = VarBinData::make_slots(offsets, &validity, len); let data = VarBinData::build( - slots[OFFSETS_SLOT] - .as_ref() - .vortex_expect("VarBinArray offsets slot") - .clone(), + VarBinSlotsView::from_slots(&slots).offsets.clone(), bytes, dtype.clone(), validity, diff --git a/vortex-array/src/arrays/varbin/builder.rs b/vortex-array/src/arrays/varbin/builder.rs index 422139cc851..eaac8e6f600 100644 --- a/vortex-array/src/arrays/varbin/builder.rs +++ b/vortex-array/src/arrays/varbin/builder.rs @@ -133,7 +133,7 @@ mod tests { use crate::VortexSessionExecute; use crate::array_session; - use crate::arrays::varbin::VarBinArrayExt; + use crate::arrays::varbin::VarBinArraySlotsExt; use crate::arrays::varbin::builder::VarBinBuilder; use crate::dtype::DType; use crate::dtype::Nullability::Nullable; diff --git a/vortex-array/src/arrays/varbin/compute/cast.rs b/vortex-array/src/arrays/varbin/compute/cast.rs index b0c4325c774..cc05ec5ac24 100644 --- a/vortex-array/src/arrays/varbin/compute/cast.rs +++ b/vortex-array/src/arrays/varbin/compute/cast.rs @@ -9,7 +9,7 @@ use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::VarBin; use crate::arrays::VarBinArray; -use crate::arrays::varbin::VarBinArrayExt; +use crate::arrays::varbin::VarBinArraySlotsExt; use crate::dtype::DType; use crate::scalar_fn::fns::cast::CastKernel; use crate::scalar_fn::fns::cast::CastReduce; diff --git a/vortex-array/src/arrays/varbin/compute/compare.rs b/vortex-array/src/arrays/varbin/compute/compare.rs index f6fff972c9b..3f5b2b0b01f 100644 --- a/vortex-array/src/arrays/varbin/compute/compare.rs +++ b/vortex-array/src/arrays/varbin/compute/compare.rs @@ -15,7 +15,7 @@ use crate::array::ArrayView; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::VarBin; -use crate::arrays::varbin::VarBinArrayExt; +use crate::arrays::varbin::VarBinArraySlotsExt; use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::match_each_integer_ptype; diff --git a/vortex-array/src/arrays/varbin/compute/filter.rs b/vortex-array/src/arrays/varbin/compute/filter.rs index 400beb9aa0a..97f6a615a07 100644 --- a/vortex-array/src/arrays/varbin/compute/filter.rs +++ b/vortex-array/src/arrays/varbin/compute/filter.rs @@ -20,6 +20,7 @@ use crate::arrays::VarBin; use crate::arrays::VarBinArray; use crate::arrays::filter::FilterKernel; use crate::arrays::varbin::VarBinArrayExt; +use crate::arrays::varbin::VarBinArraySlotsExt; use crate::arrays::varbin::builder::VarBinBuilder; use crate::dtype::DType; use crate::dtype::IntegerPType; diff --git a/vortex-array/src/arrays/varbin/compute/mask.rs b/vortex-array/src/arrays/varbin/compute/mask.rs index 19236cee47b..0d91b6fe3c3 100644 --- a/vortex-array/src/arrays/varbin/compute/mask.rs +++ b/vortex-array/src/arrays/varbin/compute/mask.rs @@ -8,7 +8,7 @@ use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::VarBin; use crate::arrays::VarBinArray; -use crate::arrays::varbin::VarBinArrayExt; +use crate::arrays::varbin::VarBinArraySlotsExt; use crate::scalar_fn::fns::mask::MaskReduce; use crate::validity::Validity; diff --git a/vortex-array/src/arrays/varbin/compute/slice.rs b/vortex-array/src/arrays/varbin/compute/slice.rs index b3ab91cc578..04ed9e52b6e 100644 --- a/vortex-array/src/arrays/varbin/compute/slice.rs +++ b/vortex-array/src/arrays/varbin/compute/slice.rs @@ -11,7 +11,7 @@ use crate::array::ArrayView; use crate::arrays::VarBin; use crate::arrays::VarBinArray; use crate::arrays::slice::SliceReduce; -use crate::arrays::varbin::VarBinArrayExt; +use crate::arrays::varbin::VarBinArraySlotsExt; impl SliceReduce for VarBin { fn slice(array: ArrayView<'_, Self>, range: Range) -> VortexResult> { diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index bc4b36cd221..5777ec1fe50 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -27,6 +27,7 @@ use crate::arrays::piecewise_sequence::constant_unsigned_usize; use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::arrays::primitive::PrimitiveArrayExt; use crate::arrays::varbin::VarBinArrayExt; +use crate::arrays::varbin::VarBinArraySlotsExt; use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::dtype::PType; diff --git a/vortex-array/src/arrays/varbin/mod.rs b/vortex-array/src/arrays/varbin/mod.rs index 4bca67fad81..f7e159fa37c 100644 --- a/vortex-array/src/arrays/varbin/mod.rs +++ b/vortex-array/src/arrays/varbin/mod.rs @@ -3,8 +3,11 @@ mod array; pub use array::VarBinArrayExt; +pub use array::VarBinArraySlotsExt; pub use array::VarBinData; pub use array::VarBinDataParts; +pub use array::VarBinSlots; +pub use array::VarBinSlotsView; pub use vtable::VarBinArray; pub(crate) mod compute; diff --git a/vortex-array/src/arrays/varbin/vtable/mod.rs b/vortex-array/src/arrays/varbin/vtable/mod.rs index 301fdebc084..0f559e13aeb 100644 --- a/vortex-array/src/arrays/varbin/vtable/mod.rs +++ b/vortex-array/src/arrays/varbin/vtable/mod.rs @@ -20,11 +20,10 @@ use crate::array::Array; use crate::array::ArrayId; use crate::array::ArrayView; use crate::array::VTable; -use crate::arrays::varbin::VarBinArrayExt; +use crate::arrays::varbin::VarBinArraySlotsExt; use crate::arrays::varbin::VarBinData; -use crate::arrays::varbin::array::NUM_SLOTS; -use crate::arrays::varbin::array::OFFSETS_SLOT; -use crate::arrays::varbin::array::SLOT_NAMES; +use crate::arrays::varbin::array::VarBinSlots; +use crate::arrays::varbin::array::VarBinSlotsView; use crate::buffer::BufferHandle; use crate::dtype::DType; use crate::dtype::Nullability; @@ -91,13 +90,12 @@ impl VTable for VarBin { slots: &[Option], ) -> VortexResult<()> { vortex_ensure!( - slots.len() == NUM_SLOTS, - "VarBinArray expected {NUM_SLOTS} slots, found {}", + slots.len() == VarBinSlots::COUNT, + "VarBinArray expected {} slots, found {}", + VarBinSlots::COUNT, slots.len() ); - let offsets = slots[OFFSETS_SLOT] - .as_ref() - .vortex_expect("VarBinArray offsets slot"); + let offsets = VarBinSlotsView::from_slots(slots).offsets; vortex_ensure!( offsets.len().saturating_sub(1) == len, "VarBinArray length {} does not match outer length {}", @@ -192,7 +190,7 @@ impl VTable for VarBin { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + VarBinSlots::NAMES[idx].to_string() } fn reduce_parent( diff --git a/vortex-array/src/arrays/variant/mod.rs b/vortex-array/src/arrays/variant/mod.rs index 1a5157e9aa1..b7944598aec 100644 --- a/vortex-array/src/arrays/variant/mod.rs +++ b/vortex-array/src/arrays/variant/mod.rs @@ -5,7 +5,6 @@ mod vtable; pub(crate) mod compute; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -20,13 +19,16 @@ use crate::ArrayRef; use crate::array::Array; use crate::array::ArrayParts; use crate::array::EmptyArrayData; -use crate::array::TypedArrayRef; +use crate::array_slots; use crate::dtype::DType; -pub(super) const CORE_STORAGE_SLOT: usize = 0; -pub(super) const SHREDDED_SLOT: usize = 1; -pub(super) const NUM_SLOTS: usize = 2; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["core_storage", "shredded"]; +#[array_slots(Variant)] +pub struct VariantSlots { + /// The logical variant storage that preserves the full value for every row. + pub core_storage: ArrayRef, + /// The optional row-aligned typed shredded tree for selected variant paths. + pub shredded: Option, +} /// Accessors for canonical variant storage. /// @@ -37,22 +39,7 @@ pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["core_storage", "shredded"]; /// chunked, constant, or otherwise encoded. Callers must use normal array operations instead of /// assuming a particular slot layout. The shredded child may have any dtype; its dtype is recorded /// during serialization and validated by normal child deserialization. -pub trait VariantArrayExt: TypedArrayRef { - /// Returns the logical variant storage that preserves the full value for every row. - fn core_storage(&self) -> &ArrayRef { - self.as_ref().slots()[CORE_STORAGE_SLOT] - .as_ref() - .vortex_expect("validated variant core_storage slot") - } - - /// Returns the optional row-aligned typed shredded tree for selected variant paths. - /// This functions returns `Some` only if the array was canonicalized and the shredded data - /// was pulled out of the underlying variant storage. - fn shredded(&self) -> Option<&ArrayRef> { - self.as_ref().slots()[SHREDDED_SLOT].as_ref() - } -} -impl> VariantArrayExt for T {} +pub use VariantArraySlotsExt as VariantArrayExt; impl Array { /// Creates a new `VariantArray` with logical variant core storage and optional shredded storage. @@ -71,8 +58,13 @@ impl Array { let len = core_storage.len(); let stats = core_storage.statistics().to_owned(); Ok(Array::try_from_parts( - ArrayParts::new(Variant, dtype, len, EmptyArrayData) - .with_slots(vec![Some(core_storage), shredded].into()), + ArrayParts::new(Variant, dtype, len, EmptyArrayData).with_slots( + VariantSlots { + core_storage, + shredded, + } + .into_slots(), + ), )? .with_stats_set(stats)) } diff --git a/vortex-array/src/arrays/variant/vtable/mod.rs b/vortex-array/src/arrays/variant/vtable/mod.rs index 4d84511ae6a..713d2d91045 100644 --- a/vortex-array/src/arrays/variant/vtable/mod.rs +++ b/vortex-array/src/arrays/variant/vtable/mod.rs @@ -6,7 +6,6 @@ mod operations; mod validity; use prost::Message; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; @@ -25,10 +24,8 @@ use crate::array::ArrayView; use crate::array::EmptyArrayData; use crate::array::VTable; use crate::array::with_empty_buffers; -use crate::arrays::variant::CORE_STORAGE_SLOT; -use crate::arrays::variant::NUM_SLOTS; -use crate::arrays::variant::SHREDDED_SLOT; -use crate::arrays::variant::SLOT_NAMES; +use crate::arrays::variant::VariantSlots; +use crate::arrays::variant::VariantSlotsView; use crate::arrays::variant::compute::rules::RULES; use crate::buffer::BufferHandle; use crate::dtype::DType; @@ -76,17 +73,17 @@ impl VTable for Variant { slots: &[Option], ) -> VortexResult<()> { vortex_ensure!( - slots.len() == NUM_SLOTS, - "VariantArray expects {NUM_SLOTS} slots, got {}", + slots.len() == VariantSlots::COUNT, + "VariantArray expects {} slots, got {}", + VariantSlots::COUNT, slots.len() ); vortex_ensure!( - slots[CORE_STORAGE_SLOT].is_some(), + slots[VariantSlots::CORE_STORAGE].is_some(), "VariantArray core_storage slot must be present" ); - let core_storage = slots[CORE_STORAGE_SLOT] - .as_ref() - .vortex_expect("validated core_storage slot presence"); + let view = VariantSlotsView::from_slots(slots); + let core_storage = view.core_storage; vortex_ensure!( matches!(dtype, DType::Variant(_)), "Expected Variant DType, got {dtype}" @@ -103,7 +100,7 @@ impl VTable for Variant { core_storage.len(), len ); - if let Some(shredded) = slots[SHREDDED_SLOT].as_ref() { + if let Some(shredded) = view.shredded { vortex_ensure!( shredded.len() == len, "VariantArray shredded length {} does not match outer length {}", @@ -138,8 +135,8 @@ impl VTable for Variant { array: ArrayView<'_, Self>, _session: &VortexSession, ) -> VortexResult>> { - let shredded_dtype = array.slots()[SHREDDED_SLOT] - .as_ref() + let shredded_dtype = VariantSlotsView::from_slots(array.slots()) + .shredded .map(|shredded| shredded.dtype().try_into()) .transpose()?; Ok(Some( @@ -180,13 +177,18 @@ impl VTable for Variant { .map(|dtype| children.get(1, &dtype, len)) .transpose()?; Ok( - ArrayParts::new(self.clone(), dtype.clone(), len, EmptyArrayData) - .with_slots(vec![Some(core_storage), shredded].into()), + ArrayParts::new(self.clone(), dtype.clone(), len, EmptyArrayData).with_slots( + VariantSlots { + core_storage, + shredded, + } + .into_slots(), + ), ) } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - match SLOT_NAMES.get(idx) { + match VariantSlots::NAMES.get(idx) { Some(name) => (*name).to_string(), None => vortex_panic!("VariantArray slot_name index {idx} out of bounds"), } diff --git a/vortex-array/src/builders/dict/bytes.rs b/vortex-array/src/builders/dict/bytes.rs index 33b2e5225f7..361735536ed 100644 --- a/vortex-array/src/builders/dict/bytes.rs +++ b/vortex-array/src/builders/dict/bytes.rs @@ -31,7 +31,7 @@ use crate::arrays::PrimitiveArray; use crate::arrays::VarBin; use crate::arrays::VarBinView; use crate::arrays::VarBinViewArray; -use crate::arrays::varbin::VarBinArrayExt; +use crate::arrays::varbin::VarBinArraySlotsExt; use crate::arrays::varbinview::build_views::BinaryView; use crate::dtype::DType; use crate::dtype::PType; diff --git a/vortex-array/src/builders/fixed_size_list.rs b/vortex-array/src/builders/fixed_size_list.rs index 9150d37b50b..7cd0ac692f0 100644 --- a/vortex-array/src/builders/fixed_size_list.rs +++ b/vortex-array/src/builders/fixed_size_list.rs @@ -15,7 +15,7 @@ use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; use crate::arrays::FixedSizeListArray; -use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use crate::builders::ArrayBuilder; use crate::builders::DEFAULT_BUILDER_CAPACITY; use crate::builders::LazyBitBufferBuilder; @@ -289,6 +289,7 @@ mod tests { use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; + use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use crate::builders::ArrayBuilder; use crate::builders::fixed_size_list::FixedSizeListArray; use crate::dtype::DType; diff --git a/vortex-array/src/builders/list.rs b/vortex-array/src/builders/list.rs index 130f37534cf..e0b43b98407 100644 --- a/vortex-array/src/builders/list.rs +++ b/vortex-array/src/builders/list.rs @@ -22,8 +22,8 @@ use crate::arrays::ListArray; use crate::arrays::ListView; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; -use crate::arrays::list::ListArrayExt; -use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::list::ListArraySlotsExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::builders::ArrayBuilder; use crate::builders::DEFAULT_BUILDER_CAPACITY; use crate::builders::LazyBitBufferBuilder; @@ -391,8 +391,9 @@ mod tests { use crate::arrays::ChunkedArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; - use crate::arrays::list::ListArrayExt; + use crate::arrays::list::ListArraySlotsExt; use crate::arrays::listview::ListViewArrayExt; + use crate::arrays::listview::ListViewArraySlotsExt; use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; use crate::builders::ListViewBuilder; diff --git a/vortex-array/src/builders/listview.rs b/vortex-array/src/builders/listview.rs index 4e6523bd7d3..a9a234b353a 100644 --- a/vortex-array/src/builders/listview.rs +++ b/vortex-array/src/builders/listview.rs @@ -27,8 +27,8 @@ use crate::arrays::List; use crate::arrays::ListView; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; -use crate::arrays::list::ListArrayExt; -use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::list::ListArraySlotsExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::listview::ListViewRebuildMode; use crate::builders::ArrayBuilder; use crate::builders::DEFAULT_BUILDER_CAPACITY; @@ -517,6 +517,7 @@ mod tests { use crate::arrays::ListArray; use crate::arrays::ListViewArray; use crate::arrays::listview::ListViewArrayExt; + use crate::arrays::listview::ListViewArraySlotsExt; use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; use crate::builders::listview::PrimitiveArray; diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/nested.rs b/vortex-array/src/scalar_fn/fns/binary/compare/nested.rs index 56d8cbf0496..15494127c6f 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/nested.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/nested.rs @@ -32,8 +32,8 @@ use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::arrays::VarBinViewArray; use crate::arrays::extension::ExtensionArrayExt; -use crate::arrays::fixed_size_list::FixedSizeListArrayExt; -use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::struct_::StructArrayExt; use crate::dtype::DType; use crate::dtype::NativePType; diff --git a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs index edf530cb75f..41918a9dffc 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs @@ -26,7 +26,7 @@ use crate::arrays::ConstantArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::arrays::bool::BoolArrayExt; -use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::primitive::PrimitiveArrayExt; use crate::arrays::scalar_fn::ScalarFnFactoryExt; use crate::builtins::ArrayBuiltins; diff --git a/vortex-array/src/scalar_fn/fns/list_length.rs b/vortex-array/src/scalar_fn/fns/list_length.rs index ea88a3b27d6..41180b1ae20 100644 --- a/vortex-array/src/scalar_fn/fns/list_length.rs +++ b/vortex-array/src/scalar_fn/fns/list_length.rs @@ -17,7 +17,9 @@ use crate::arrays::List; use crate::arrays::ListView; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::arrays::list::ListArrayExt; +use crate::arrays::list::ListArraySlotsExt; use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::Nullability; diff --git a/vortex-arrow/src/convert.rs b/vortex-arrow/src/convert.rs index 0bab0937a8a..e227b401633 100644 --- a/vortex-arrow/src/convert.rs +++ b/vortex-arrow/src/convert.rs @@ -716,8 +716,9 @@ mod tests { use vortex_array::arrays::Struct; use vortex_array::arrays::VarBinView; use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; - use vortex_array::arrays::list::ListArrayExt; - use vortex_array::arrays::listview::ListViewArrayExt; + use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; + use vortex_array::arrays::list::ListArraySlotsExt; + use vortex_array::arrays::listview::ListViewArraySlotsExt; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; diff --git a/vortex-arrow/src/executor/byte.rs b/vortex-arrow/src/executor/byte.rs index ae40ee01c26..e4b37f84da3 100644 --- a/vortex-arrow/src/executor/byte.rs +++ b/vortex-arrow/src/executor/byte.rs @@ -14,7 +14,7 @@ use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::arrays::VarBin; use vortex_array::arrays::VarBinViewArray; -use vortex_array::arrays::varbin::VarBinArrayExt; +use vortex_array::arrays::varbin::VarBinArraySlotsExt; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; diff --git a/vortex-arrow/src/executor/fixed_size_list.rs b/vortex-arrow/src/executor/fixed_size_list.rs index 72ce7c57ff0..2d556f38e73 100644 --- a/vortex-arrow/src/executor/fixed_size_list.rs +++ b/vortex-arrow/src/executor/fixed_size_list.rs @@ -9,6 +9,7 @@ use vortex_array::ExecutionCtx; use vortex_array::arrays::FixedSizeList; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; +use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use vortex_error::VortexResult; use vortex_error::vortex_ensure; diff --git a/vortex-arrow/src/executor/list.rs b/vortex-arrow/src/executor/list.rs index 0b0d7f625b5..145c8dfc80c 100644 --- a/vortex-arrow/src/executor/list.rs +++ b/vortex-arrow/src/executor/list.rs @@ -17,8 +17,8 @@ use vortex_array::arrays::ListArray; use vortex_array::arrays::ListView; use vortex_array::arrays::ListViewArray; use vortex_array::arrays::chunked::ChunkedArrayExt; -use vortex_array::arrays::list::ListArrayExt; -use vortex_array::arrays::listview::ListViewArrayExt; +use vortex_array::arrays::list::ListArraySlotsExt; +use vortex_array::arrays::listview::ListViewArraySlotsExt; use vortex_array::arrays::listview::ListViewDataParts; use vortex_array::arrays::listview::ListViewRebuildMode; use vortex_array::builtins::ArrayBuiltins; diff --git a/vortex-arrow/src/executor/list_view.rs b/vortex-arrow/src/executor/list_view.rs index da242eb9719..ed1ba93e556 100644 --- a/vortex-arrow/src/executor/list_view.rs +++ b/vortex-arrow/src/executor/list_view.rs @@ -13,6 +13,7 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::listview::DEFAULT_REBUILD_DENSITY_THRESHOLD; use vortex_array::arrays::listview::DEFAULT_TRIM_ELEMENTS_THRESHOLD; use vortex_array::arrays::listview::ListViewArrayExt; +use vortex_array::arrays::listview::ListViewArraySlotsExt; use vortex_array::arrays::listview::ListViewDataParts; use vortex_array::arrays::listview::ListViewRebuildMode; use vortex_array::builtins::ArrayBuiltins; diff --git a/vortex-arrow/src/executor/mod.rs b/vortex-arrow/src/executor/mod.rs index 875577788b4..14d94ccc033 100644 --- a/vortex-arrow/src/executor/mod.rs +++ b/vortex-arrow/src/executor/mod.rs @@ -34,8 +34,8 @@ use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::arrays::List; use vortex_array::arrays::VarBin; -use vortex_array::arrays::list::ListArrayExt; -use vortex_array::arrays::varbin::VarBinArrayExt; +use vortex_array::arrays::list::ListArraySlotsExt; +use vortex_array::arrays::varbin::VarBinArraySlotsExt; use vortex_array::dtype::DType; use vortex_array::dtype::PType; use vortex_error::VortexResult; diff --git a/vortex-arrow/src/executor/run_end.rs b/vortex-arrow/src/executor/run_end.rs index b6913c1ea1c..a2be95b3b21 100644 --- a/vortex-arrow/src/executor/run_end.rs +++ b/vortex-arrow/src/executor/run_end.rs @@ -23,6 +23,7 @@ use vortex_error::vortex_err; use vortex_runend::RunEnd; use vortex_runend::RunEndArray; use vortex_runend::RunEndArrayExt; +use vortex_runend::RunEndArraySlotsExt; use crate::ArrowArrayExecutor; diff --git a/vortex-btrblocks/src/schemes/integer/rle.rs b/vortex-btrblocks/src/schemes/integer/rle.rs index 5bc4e5296f1..a9faeab54f2 100644 --- a/vortex-btrblocks/src/schemes/integer/rle.rs +++ b/vortex-btrblocks/src/schemes/integer/rle.rs @@ -21,6 +21,7 @@ use vortex_error::VortexResult; use vortex_fastlanes::Delta; use vortex_fastlanes::RLE; use vortex_fastlanes::RLEArrayExt; +use vortex_fastlanes::RLEArraySlotsExt; use super::RUN_LENGTH_THRESHOLD; use crate::ArrayAndStats; diff --git a/vortex-btrblocks/src/schemes/string/fsst.rs b/vortex-btrblocks/src/schemes/string/fsst.rs index bd5bd010396..593fcad802b 100644 --- a/vortex-btrblocks/src/schemes/string/fsst.rs +++ b/vortex-btrblocks/src/schemes/string/fsst.rs @@ -10,12 +10,13 @@ use vortex_array::IntoArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; -use vortex_array::arrays::varbin::VarBinArrayExt; +use vortex_array::arrays::varbin::VarBinArraySlotsExt; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_error::VortexResult; use vortex_fsst::FSST; use vortex_fsst::FSSTArrayExt; +use vortex_fsst::FSSTArraySlotsExt; use vortex_fsst::fsst_compress; use vortex_fsst::fsst_train_compressor; diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index 9dbbc611448..ae52c2f9bde 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -18,7 +18,8 @@ use vortex_array::arrays::Variant; use vortex_array::arrays::VariantArray; use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; -use vortex_array::arrays::listview::ListViewArrayExt; +use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use vortex_array::arrays::listview::ListViewArraySlotsExt; use vortex_array::arrays::listview::list_from_list_view; use vortex_array::arrays::masked::MaskedArraySlotsExt; use vortex_array::arrays::scalar_fn::AnyScalarFn; diff --git a/vortex-compressor/src/compressor/structural.rs b/vortex-compressor/src/compressor/structural.rs index ec867c1c1dc..8b9cd999d15 100644 --- a/vortex-compressor/src/compressor/structural.rs +++ b/vortex-compressor/src/compressor/structural.rs @@ -12,7 +12,8 @@ use vortex_array::arrays::ListArray; use vortex_array::arrays::ListViewArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::list::ListArrayExt; -use vortex_array::arrays::listview::ListViewArrayExt; +use vortex_array::arrays::list::ListArraySlotsExt; +use vortex_array::arrays::listview::ListViewArraySlotsExt; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_error::VortexResult; diff --git a/vortex-cuda/benches/fsst_cuda.rs b/vortex-cuda/benches/fsst_cuda.rs index 49930c9a30a..6ba8c283d37 100644 --- a/vortex-cuda/benches/fsst_cuda.rs +++ b/vortex-cuda/benches/fsst_cuda.rs @@ -34,6 +34,7 @@ use vortex_cuda::arrow::release_device_array; use vortex_cuda::executor::CudaArrayExt; use vortex_cuda_macros::cuda_available; use vortex_cuda_macros::cuda_not_available; +use vortex_fsst::FSSTArraySlotsExt; use vortex_fsst::test_utils::make_fsst_clickbench_urls; use crate::timed_launch_strategy::TimedLaunchStrategy; diff --git a/vortex-cuda/src/arrow/canonical.rs b/vortex-cuda/src/arrow/canonical.rs index c28b0c242ed..3b5946b3391 100644 --- a/vortex-cuda/src/arrow/canonical.rs +++ b/vortex-cuda/src/arrow/canonical.rs @@ -38,7 +38,7 @@ use vortex::array::arrays::extension::ExtensionArrayExt; use vortex::array::arrays::fixed_size_list::FixedSizeListArrayExt; use vortex::array::arrays::fixed_size_list::FixedSizeListDataParts; use vortex::array::arrays::list::ListDataParts; -use vortex::array::arrays::listview::ListViewArrayExt; +use vortex::array::arrays::listview::ListViewArraySlotsExt; use vortex::array::arrays::listview::list_from_list_view; use vortex::array::arrays::primitive::PrimitiveDataParts; use vortex::array::arrays::struct_::StructDataParts; diff --git a/vortex-cuda/src/arrow/mod.rs b/vortex-cuda/src/arrow/mod.rs index 071e22085c4..75dea478717 100644 --- a/vortex-cuda/src/arrow/mod.rs +++ b/vortex-cuda/src/arrow/mod.rs @@ -37,9 +37,9 @@ use vortex::array::arrays::List; use vortex::array::arrays::ListView; use vortex::array::arrays::Struct; use vortex::array::arrays::dict::DictArraySlotsExt; -use vortex::array::arrays::fixed_size_list::FixedSizeListArrayExt; -use vortex::array::arrays::list::ListArrayExt; -use vortex::array::arrays::listview::ListViewArrayExt; +use vortex::array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use vortex::array::arrays::list::ListArraySlotsExt; +use vortex::array::arrays::listview::ListViewArraySlotsExt; use vortex::array::arrays::struct_::StructArrayExt; use vortex::array::buffer::BufferHandle; use vortex::array::stream::SendableArrayStream; diff --git a/vortex-cuda/src/dynamic_dispatch/plan_builder.rs b/vortex-cuda/src/dynamic_dispatch/plan_builder.rs index d7c5a58712d..4461b53212e 100644 --- a/vortex-cuda/src/dynamic_dispatch/plan_builder.rs +++ b/vortex-cuda/src/dynamic_dispatch/plan_builder.rs @@ -34,6 +34,7 @@ use vortex::encodings::fastlanes::FoR; use vortex::encodings::fastlanes::FoRArrayExt; use vortex::encodings::runend::RunEnd; use vortex::encodings::runend::RunEndArrayExt; +use vortex::encodings::runend::RunEndArraySlotsExt; use vortex::encodings::sequence::Sequence; use vortex::encodings::zigzag::ZigZag; use vortex::encodings::zigzag::ZigZagArrayExt; diff --git a/vortex-cuda/src/kernel/encodings/fsst.rs b/vortex-cuda/src/kernel/encodings/fsst.rs index de46d0b8fa1..48470418195 100644 --- a/vortex-cuda/src/kernel/encodings/fsst.rs +++ b/vortex-cuda/src/kernel/encodings/fsst.rs @@ -16,7 +16,7 @@ use vortex::array::Canonical; use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::VarBinViewArray; use vortex::array::arrays::primitive::PrimitiveDataParts; -use vortex::array::arrays::varbin::VarBinArrayExt; +use vortex::array::arrays::varbin::VarBinArraySlotsExt; use vortex::array::arrays::varbinview::build_views::MAX_BUFFER_LEN; use vortex::array::arrays::varbinview::build_views::build_views; use vortex::array::buffer::BufferHandle; @@ -31,6 +31,7 @@ use vortex::dtype::NativePType; use vortex::encodings::fsst::FSST; use vortex::encodings::fsst::FSSTArray; use vortex::encodings::fsst::FSSTArrayExt; +use vortex::encodings::fsst::FSSTArraySlotsExt; use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::error::vortex_err; diff --git a/vortex-cuda/src/kernel/encodings/runend.rs b/vortex-cuda/src/kernel/encodings/runend.rs index 4ac042e2972..c669dc0cf80 100644 --- a/vortex-cuda/src/kernel/encodings/runend.rs +++ b/vortex-cuda/src/kernel/encodings/runend.rs @@ -22,6 +22,7 @@ use vortex::dtype::PType; use vortex::encodings::runend::RunEnd; use vortex::encodings::runend::RunEndArray; use vortex::encodings::runend::RunEndArrayExt; +use vortex::encodings::runend::RunEndArraySlotsExt; use vortex::error::VortexResult; use vortex::error::vortex_bail; use vortex::error::vortex_ensure; diff --git a/vortex-duckdb/src/convert/vector.rs b/vortex-duckdb/src/convert/vector.rs index a3fe098ffbb..a7a0966f0e4 100644 --- a/vortex-duckdb/src/convert/vector.rs +++ b/vortex-duckdb/src/convert/vector.rs @@ -398,6 +398,7 @@ mod tests { use vortex::array::arrays::VarBinViewArray; use vortex::array::arrays::fixed_size_list::FixedSizeListArrayExt; use vortex::array::arrays::listview::ListViewArrayExt; + use vortex::array::arrays::listview::ListViewArraySlotsExt; use vortex::array::arrays::struct_::StructArrayExt; use vortex::array::assert_arrays_eq; use vortex::error::VortexExpect; diff --git a/vortex-duckdb/src/exporter/list_view.rs b/vortex-duckdb/src/exporter/list_view.rs index baf27949e5d..54dc4db9e79 100644 --- a/vortex-duckdb/src/exporter/list_view.rs +++ b/vortex-duckdb/src/exporter/list_view.rs @@ -13,6 +13,7 @@ use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::listview::DEFAULT_REBUILD_DENSITY_THRESHOLD; use vortex::array::arrays::listview::DEFAULT_TRIM_ELEMENTS_THRESHOLD; use vortex::array::arrays::listview::ListViewArrayExt; +use vortex::array::arrays::listview::ListViewArraySlotsExt; use vortex::array::arrays::listview::ListViewDataParts; use vortex::array::arrays::listview::ListViewRebuildMode; use vortex::array::match_each_integer_ptype; diff --git a/vortex-duckdb/src/exporter/rle.rs b/vortex-duckdb/src/exporter/rle.rs index 92e491d5442..aa2f12ecabf 100644 --- a/vortex-duckdb/src/exporter/rle.rs +++ b/vortex-duckdb/src/exporter/rle.rs @@ -13,6 +13,7 @@ use vortex::dtype::IntegerPType; use vortex::encodings::fastlanes::FL_CHUNK_SIZE; use vortex::encodings::fastlanes::RLEArray; use vortex::encodings::fastlanes::RLEArrayExt; +use vortex::encodings::fastlanes::RLEArraySlotsExt; use vortex::error::VortexResult; use crate::duckdb::ReusableDict; diff --git a/vortex-duckdb/src/exporter/run_end.rs b/vortex-duckdb/src/exporter/run_end.rs index 2c499d700a3..40158895f1a 100644 --- a/vortex-duckdb/src/exporter/run_end.rs +++ b/vortex-duckdb/src/exporter/run_end.rs @@ -13,6 +13,7 @@ use vortex::array::search_sorted::SearchSortedSide; use vortex::dtype::IntegerPType; use vortex::encodings::runend::RunEndArray; use vortex::encodings::runend::RunEndArrayExt; +use vortex::encodings::runend::RunEndArraySlotsExt; use vortex::error::VortexExpect; use vortex::error::VortexResult; diff --git a/vortex-geo/src/extension/mod.rs b/vortex-geo/src/extension/mod.rs index 99150833ab5..b05069f8ca5 100644 --- a/vortex-geo/src/extension/mod.rs +++ b/vortex-geo/src/extension/mod.rs @@ -47,8 +47,8 @@ use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::ListViewArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::extension::ExtensionArrayExt; -use vortex_array::arrays::list::ListArrayExt; -use vortex_array::arrays::listview::ListViewArrayExt; +use vortex_array::arrays::list::ListArraySlotsExt; +use vortex_array::arrays::listview::ListViewArraySlotsExt; use vortex_array::arrays::listview::list_from_list_view; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; diff --git a/vortex-row/src/codec.rs b/vortex-row/src/codec.rs index fd14bb254b8..41c04370a65 100644 --- a/vortex-row/src/codec.rs +++ b/vortex-row/src/codec.rs @@ -34,6 +34,7 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; +use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalType; diff --git a/vortex-row/src/tests.rs b/vortex-row/src/tests.rs index c7f9e8a2b49..e13aeb62452 100644 --- a/vortex-row/src/tests.rs +++ b/vortex-row/src/tests.rs @@ -16,6 +16,7 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::listview::ListViewArrayExt; +use vortex_array::arrays::listview::ListViewArraySlotsExt; use vortex_array::dtype::Nullability; use vortex_array::extension::datetime::Date; use vortex_array::extension::datetime::TimeUnit; diff --git a/vortex-tensor/src/scalar_fns/l2_denorm.rs b/vortex-tensor/src/scalar_fns/l2_denorm.rs index a4ac9e6c91e..b879ab3731c 100644 --- a/vortex-tensor/src/scalar_fns/l2_denorm.rs +++ b/vortex-tensor/src/scalar_fns/l2_denorm.rs @@ -20,6 +20,7 @@ use vortex_array::arrays::ScalarFn as ScalarFnArrayEncoding; use vortex_array::arrays::ScalarFnArray; use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; +use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use vortex_array::arrays::scalar_fn::ExactScalarFn; use vortex_array::arrays::scalar_fn::ScalarFnArrayExt; use vortex_array::arrays::scalar_fn::ScalarFnArrayView; @@ -743,7 +744,7 @@ mod tests { use vortex_array::arrays::MaskedArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::extension::ExtensionArrayExt; - use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; + use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use vortex_array::arrays::scalar_fn::ScalarFnArrayExt; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; use vortex_array::dtype::DType; diff --git a/vortex-tensor/src/utils.rs b/vortex-tensor/src/utils.rs index 841cd4329c1..e8dbf16f171 100644 --- a/vortex-tensor/src/utils.rs +++ b/vortex-tensor/src/utils.rs @@ -11,7 +11,7 @@ use vortex_array::arrays::ConstantArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::ScalarFn; -use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; +use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::arrays::scalar_fn::ExactScalarFn; use vortex_array::arrays::scalar_fn::ScalarFnArrayExt; diff --git a/vortex/benches/common_encoding_tree_throughput.rs b/vortex/benches/common_encoding_tree_throughput.rs index 65c2c1bc040..239e6649745 100644 --- a/vortex/benches/common_encoding_tree_throughput.rs +++ b/vortex/benches/common_encoding_tree_throughput.rs @@ -23,7 +23,7 @@ use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::TemporalArray; use vortex::array::arrays::VarBinArray; use vortex::array::arrays::VarBinViewArray; -use vortex::array::arrays::varbin::VarBinArrayExt; +use vortex::array::arrays::varbin::VarBinArraySlotsExt; use vortex::array::builtins::ArrayBuiltins; use vortex::dtype::DType; use vortex::dtype::PType; @@ -38,10 +38,11 @@ use vortex::encodings::fastlanes::FoR; use vortex::encodings::fastlanes::FoRArrayExt; use vortex::encodings::fsst::FSST; use vortex::encodings::fsst::FSSTArrayExt; +use vortex::encodings::fsst::FSSTArraySlotsExt; use vortex::encodings::fsst::fsst_compress; use vortex::encodings::fsst::fsst_train_compressor; use vortex::encodings::runend::RunEnd; -use vortex::encodings::runend::RunEndArrayExt; +use vortex::encodings::runend::RunEndArraySlotsExt; use vortex::error::VortexExpect; use vortex::extension::datetime::TimeUnit; use vortex_session::VortexSession;