FFI: array slots

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<Vec<T>>`.
 ///
@@ -279,20 +281,20 @@ pub trait ListArrayExt: TypedArrayRef<List> {
     }
 
     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<List> {
     }
 
     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,
         }

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;
 

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<ArrayRef>],
     ) -> 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<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {

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);
+        }
+    }
 }