diff --git a/vortex-array/src/arrays/list/array.rs b/vortex-array/src/arrays/list/array.rs index dbe2943fc51..95c6260d786 100644 --- a/vortex-array/src/arrays/list/array.rs +++ b/vortex-array/src/arrays/list/array.rs @@ -7,6 +7,7 @@ use std::sync::Arc; use num_traits::AsPrimitive; use smallvec::smallvec; +use vortex_array_macros::array_slots; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -38,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 list elements are non-null. + pub validity: ArrayRef, +} /// A list array that stores variable-length lists of elements, similar to `Vec>`. /// @@ -279,20 +281,20 @@ pub trait ListArrayExt: TypedArrayRef { } fn elements(&self) -> &ArrayRef { - self.as_ref().slots()[ELEMENTS_SLOT] + self.as_ref().slots()[ListSlots::ELEMENTS] .as_ref() .vortex_expect("ListArray elements slot") } fn offsets(&self) -> &ArrayRef { - self.as_ref().slots()[OFFSETS_SLOT] + self.as_ref().slots()[ListSlots::OFFSETS] .as_ref() .vortex_expect("ListArray offsets slot") } fn list_validity(&self) -> Validity { child_to_validity( - self.as_ref().slots()[VALIDITY_SLOT].as_ref(), + self.as_ref().slots()[ListSlots::VALIDITY].as_ref(), self.nullability(), ) } @@ -405,17 +407,12 @@ impl Array { } pub fn into_data_parts(self) -> ListDataParts { + let slots = ListSlotsView::from_slots(self.slots()); 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(); ListDataParts { - elements, - offsets, + elements: slots.elements.clone(), + offsets: slots.offsets.clone(), validity, dtype, } diff --git a/vortex-array/src/arrays/list/mod.rs b/vortex-array/src/arrays/list/mod.rs index 422a9523ff9..28962e70c03 100644 --- a/vortex-array/src/arrays/list/mod.rs +++ b/vortex-array/src/arrays/list/mod.rs @@ -4,6 +4,7 @@ mod array; pub use array::ListArrayExt; pub use array::ListData; +pub use array::ListSlots; pub use array::ListDataParts; pub use vtable::ListArray; diff --git a/vortex-array/src/arrays/list/vtable/mod.rs b/vortex-array/src/arrays/list/vtable/mod.rs index 55f7753c995..255b0c817f9 100644 --- a/vortex-array/src/arrays/list/vtable/mod.rs +++ b/vortex-array/src/arrays/list/vtable/mod.rs @@ -4,8 +4,8 @@ use std::hash::Hasher; use std::sync::Arc; +use super::array::ListSlotsView; use prost::Message; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -28,10 +28,7 @@ use crate::array::VTable; use crate::array::with_empty_buffers; use crate::arrays::list::ListArrayExt; 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::compute::rules::PARENT_RULES; use crate::arrays::listview::list_view_from_list; use crate::buffer::BufferHandle; @@ -122,17 +119,9 @@ impl VTable for List { len: usize, slots: &[Option], ) -> VortexResult<()> { - vortex_ensure!( - slots.len() == NUM_SLOTS, - "ListArray expected {NUM_SLOTS} slots, found {}", - 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 view = ListSlotsView::from_slots(slots); + let elements = view.elements; + let offsets = view.offsets; vortex_ensure!( offsets.len().saturating_sub(1) == len, "ListArray length {} does not match outer length {}", @@ -192,7 +181,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-ffi/src/array.rs b/vortex-ffi/src/array.rs index 61a8ab00a08..cceceb8bff2 100644 --- a/vortex-ffi/src/array.rs +++ b/vortex-ffi/src/array.rs @@ -23,6 +23,8 @@ use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::StructArray; use vortex::array::arrays::VarBinView; use vortex::array::arrays::bool::BoolArrayExt; +use vortex::array::arrays::dict::DictSlots; +use vortex::array::arrays::list::ListSlots; use vortex::array::arrays::struct_::StructArrayExt; use vortex::array::arrow::FromArrowArray; use vortex::array::legacy_session; @@ -611,6 +613,41 @@ pub unsafe extern "C" fn vx_array_apply( }) } +/// Return array's encoding name +#[unsafe(no_mangle)] +pub unsafe extern "C" fn vx_array_encoding(array: *const vx_array) -> vx_view { + let array = vx_array::as_ref(array); + vx_view::from_bytes(array.encoding_id().as_str().as_bytes()) +} + +/// The codes array mapping each element to a dictionary entry. +pub const VX_SLOT_DICT_CODES: usize = DictSlots::CODES; +/// The dictionary values array containing the unique values. +pub const VX_SLOT_DICT_VALUES: usize = DictSlots::VALUES; + +/// The elements data array containing all list elements concatenated together. +pub const VX_SLOT_LIST_ELEMENTS: usize = ListSlots::ELEMENTS; +/// The offsets array defining the start/end of each list within the elements array. +pub const VX_SLOT_LIST_OFFSETS: usize = ListSlots::OFFSETS; +/// The validity bitmap indicating which list elements are non-null. +pub const VX_SLOT_LIST_VALIDITY: usize = ListSlots::VALIDITY; + +/// Return array's slot from VX_SLOT_* constants. Returns NULL if +/// index is out of bounds. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn vx_array_slot(array: *const vx_array, index: usize) -> *const vx_array { + let array = vx_array::as_ref(array); + let slots = array.slots(); + if index >= slots.len() { + return ptr::null(); + } + if let Some(slot) = &slots[index] { + vx_array::new(Arc::new(slot.clone())) + } else { + ptr::null() + } +} + #[cfg(test)] mod tests { use std::ptr; @@ -625,12 +662,15 @@ mod tests { use vortex::array::arrays::VarBinViewArray; use vortex::array::arrays::bool::BoolArrayExt; use vortex::array::validity::Validity; + use vortex::buffer::BitBuffer; use vortex::buffer::buffer; #[cfg(not(miri))] use vortex::dtype::half::f16; use vortex::expr::eq; use vortex::expr::lit; use vortex::expr::root; + use vortex_array::arrays::DictArray; + use vortex_array::assert_arrays_eq; use crate::array::*; use crate::dtype::vx_dtype_free; @@ -1189,4 +1229,42 @@ mod tests { vx_array_free(array); } } + + #[test] + #[cfg_attr(miri, ignore)] + fn test_dict_slots() { + let mut ctx = legacy_session().create_execution_ctx(); + let codes_buffer = buffer![0u32, 1, 2, 2, 1]; + let codes_validity = Validity::from(BitBuffer::from(vec![true, false, true, false, true])); + let codes = PrimitiveArray::new(codes_buffer, codes_validity).into_array(); + let values = PrimitiveArray::new(buffer![3, 6, 9], Validity::AllValid).into_array(); + let dict = DictArray::try_new(codes.clone(), values.clone()) + .unwrap() + .into_array(); + + unsafe { + let array = vx_array::new(Arc::new(dict)); + let encoding_view = vx_array_encoding(array); + assert_eq!( + "vortex.dict".as_bytes(), + from_raw_parts(encoding_view.ptr, encoding_view.len) + ); + let codes_slot = vx_array_slot(array, VX_SLOT_DICT_CODES); + { + let codes_slot = vx_array::as_ref(codes_slot); + assert_arrays_eq!(codes, codes_slot, &mut ctx); + } + + let values_slot = vx_array_slot(array, VX_SLOT_DICT_VALUES); + { + let values_slot = vx_array::as_ref(values_slot); + assert_arrays_eq!(values, values_slot, &mut ctx); + } + + let invalid_slot = vx_array_slot(array, 999); + assert_eq!(invalid_slot, ptr::null()); + + vx_array_free(array); + } + } }