feat: iterative execution for DateTimePartsArray (#7698)

Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: joseph-isaacs

encodings/bytebool/src/array.rs

@@ -73,7 +73,7 @@ impl VTable for ByteBool {
         len: usize,
         slots: &[Option<ArrayRef>],
     ) -> VortexResult<()> {
-        let validity = child_to_validity(&slots[VALIDITY_SLOT], dtype.nullability());
+        let validity = child_to_validity(slots[VALIDITY_SLOT].as_ref(), dtype.nullability());
         ByteBoolData::validate(data.buffer(), &validity, dtype, len)
     }
 
@@ -186,7 +186,7 @@ impl Display for ByteBoolData {
 pub trait ByteBoolArrayExt: TypedArrayRef<ByteBool> {
     fn validity(&self) -> Validity {
         child_to_validity(
-            &self.as_ref().slots()[VALIDITY_SLOT],
+            self.as_ref().slots()[VALIDITY_SLOT].as_ref(),
             self.as_ref().dtype().nullability(),
         )
     }

encodings/datetime-parts/public-api.lock

@@ -152,6 +152,62 @@ pub fn vortex_datetime_parts::DateTimePartsMetadata::clear(&mut self)
 
 pub fn vortex_datetime_parts::DateTimePartsMetadata::encoded_len(&self) -> usize
 
+pub struct vortex_datetime_parts::DateTimePartsParts
+
+pub vortex_datetime_parts::DateTimePartsParts::days: vortex_array::array::erased::ArrayRef
+
+pub vortex_datetime_parts::DateTimePartsParts::seconds: vortex_array::array::erased::ArrayRef
+
+pub vortex_datetime_parts::DateTimePartsParts::subseconds: vortex_array::array::erased::ArrayRef
+
+pub struct vortex_datetime_parts::DateTimePartsSlots
+
+pub vortex_datetime_parts::DateTimePartsSlots::days: vortex_array::array::erased::ArrayRef
+
+pub vortex_datetime_parts::DateTimePartsSlots::seconds: vortex_array::array::erased::ArrayRef
+
+pub vortex_datetime_parts::DateTimePartsSlots::subseconds: vortex_array::array::erased::ArrayRef
+
+impl vortex_datetime_parts::DateTimePartsSlots
+
+pub const vortex_datetime_parts::DateTimePartsSlots::COUNT: usize
+
+pub const vortex_datetime_parts::DateTimePartsSlots::DAYS: usize
+
+pub const vortex_datetime_parts::DateTimePartsSlots::NAMES: [&'static str; 3]
+
+pub const vortex_datetime_parts::DateTimePartsSlots::SECONDS: usize
+
+pub const vortex_datetime_parts::DateTimePartsSlots::SUBSECONDS: usize
+
+pub fn vortex_datetime_parts::DateTimePartsSlots::from_slots(slots: alloc::vec::Vec<core::option::Option<vortex_array::array::erased::ArrayRef>>) -> Self
+
+pub fn vortex_datetime_parts::DateTimePartsSlots::into_slots(self) -> alloc::vec::Vec<core::option::Option<vortex_array::array::erased::ArrayRef>>
+
+pub struct vortex_datetime_parts::DateTimePartsSlotsView<'a>
+
+pub vortex_datetime_parts::DateTimePartsSlotsView::days: &'a vortex_array::array::erased::ArrayRef
+
+pub vortex_datetime_parts::DateTimePartsSlotsView::seconds: &'a vortex_array::array::erased::ArrayRef
+
+pub vortex_datetime_parts::DateTimePartsSlotsView::subseconds: &'a vortex_array::array::erased::ArrayRef
+
+impl<'a> vortex_datetime_parts::DateTimePartsSlotsView<'a>
+
+pub fn vortex_datetime_parts::DateTimePartsSlotsView<'a>::from_slots(slots: &'a [core::option::Option<vortex_array::array::erased::ArrayRef>]) -> Self
+
+pub fn vortex_datetime_parts::DateTimePartsSlotsView<'a>::to_owned(&self) -> vortex_datetime_parts::DateTimePartsSlots
+
+impl<'a> core::clone::Clone for vortex_datetime_parts::DateTimePartsSlotsView<'a>
+
+pub fn vortex_datetime_parts::DateTimePartsSlotsView<'a>::clone(&self) -> vortex_datetime_parts::DateTimePartsSlotsView<'a>
+
+impl<'a> core::fmt::Debug for vortex_datetime_parts::DateTimePartsSlotsView<'a>
+
+pub fn vortex_datetime_parts::DateTimePartsSlotsView<'a>::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+impl<'a> core::marker::Copy for vortex_datetime_parts::DateTimePartsSlotsView<'a>
+
 pub struct vortex_datetime_parts::TemporalParts
 
 pub vortex_datetime_parts::TemporalParts::days: vortex_array::array::erased::ArrayRef
@@ -160,22 +216,34 @@ pub vortex_datetime_parts::TemporalParts::seconds: vortex_array::array::erased::
 
 pub vortex_datetime_parts::TemporalParts::subseconds: vortex_array::array::erased::ArrayRef
 
-pub trait vortex_datetime_parts::DateTimePartsArrayExt: vortex_array::array::typed::TypedArrayRef<vortex_datetime_parts::DateTimeParts>
+pub trait vortex_datetime_parts::DateTimePartsArraySlotsExt: vortex_array::array::typed::TypedArrayRef<vortex_datetime_parts::DateTimeParts>
 
-pub fn vortex_datetime_parts::DateTimePartsArrayExt::days(&self) -> &vortex_array::array::erased::ArrayRef
+pub fn vortex_datetime_parts::DateTimePartsArraySlotsExt::days(&self) -> &vortex_array::array::erased::ArrayRef
 
-pub fn vortex_datetime_parts::DateTimePartsArrayExt::seconds(&self) -> &vortex_array::array::erased::ArrayRef
+pub fn vortex_datetime_parts::DateTimePartsArraySlotsExt::seconds(&self) -> &vortex_array::array::erased::ArrayRef
 
-pub fn vortex_datetime_parts::DateTimePartsArrayExt::subseconds(&self) -> &vortex_array::array::erased::ArrayRef
+pub fn vortex_datetime_parts::DateTimePartsArraySlotsExt::slots_view(&self) -> vortex_datetime_parts::DateTimePartsSlotsView<'_>
 
-impl<T: vortex_array::array::typed::TypedArrayRef<vortex_datetime_parts::DateTimeParts>> vortex_datetime_parts::DateTimePartsArrayExt for T
+pub fn vortex_datetime_parts::DateTimePartsArraySlotsExt::subseconds(&self) -> &vortex_array::array::erased::ArrayRef
+
+impl<T: vortex_array::array::typed::TypedArrayRef<vortex_datetime_parts::DateTimeParts>> vortex_datetime_parts::DateTimePartsArraySlotsExt for T
 
 pub fn T::days(&self) -> &vortex_array::array::erased::ArrayRef
 
 pub fn T::seconds(&self) -> &vortex_array::array::erased::ArrayRef
 
+pub fn T::slots_view(&self) -> vortex_datetime_parts::DateTimePartsSlotsView<'_>
+
 pub fn T::subseconds(&self) -> &vortex_array::array::erased::ArrayRef
 
+pub trait vortex_datetime_parts::DateTimePartsOwnedExt
+
+pub fn vortex_datetime_parts::DateTimePartsOwnedExt::into_parts(self) -> vortex_datetime_parts::DateTimePartsParts
+
+impl vortex_datetime_parts::DateTimePartsOwnedExt for vortex_array::array::typed::Array<vortex_datetime_parts::DateTimeParts>
+
+pub fn vortex_array::array::typed::Array<vortex_datetime_parts::DateTimeParts>::into_parts(self) -> vortex_datetime_parts::DateTimePartsParts
+
 pub fn vortex_datetime_parts::initialize(session: &vortex_session::VortexSession)
 
 pub fn vortex_datetime_parts::split_temporal(array: vortex_array::arrays::datetime::TemporalArray, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_datetime_parts::TemporalParts>

encodings/datetime-parts/src/array.rs

@@ -7,6 +7,7 @@ use std::fmt::Formatter;
 use std::hash::Hasher;
 
 use prost::Message;
+use vortex_array::AnyCanonical;
 use vortex_array::Array;
 use vortex_array::ArrayEq;
 use vortex_array::ArrayHash;
@@ -18,17 +19,18 @@ use vortex_array::ExecutionCtx;
 use vortex_array::ExecutionResult;
 use vortex_array::IntoArray;
 use vortex_array::Precision;
-use vortex_array::TypedArrayRef;
+use vortex_array::array_slots;
+use vortex_array::arrays::Primitive;
 use vortex_array::arrays::TemporalArray;
 use vortex_array::buffer::BufferHandle;
 use vortex_array::dtype::DType;
 use vortex_array::dtype::Nullability;
 use vortex_array::dtype::PType;
+use vortex_array::require_child;
 use vortex_array::serde::ArrayChildren;
 use vortex_array::vtable::VTable;
 use vortex_array::vtable::ValidityChild;
 use vortex_array::vtable::ValidityVTableFromChild;
-use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_error::vortex_ensure;
@@ -104,16 +106,8 @@ impl VTable for DateTimeParts {
         len: usize,
         slots: &[Option<ArrayRef>],
     ) -> VortexResult<()> {
-        let days = slots[DAYS_SLOT]
-            .as_ref()
-            .vortex_expect("DateTimePartsArray days slot");
-        let seconds = slots[SECONDS_SLOT]
-            .as_ref()
-            .vortex_expect("DateTimePartsArray seconds slot");
-        let subseconds = slots[SUBSECONDS_SLOT]
-            .as_ref()
-            .vortex_expect("DateTimePartsArray subseconds slot");
-        DateTimePartsData::validate(dtype, days, seconds, subseconds, len)
+        let slots = DateTimePartsSlotsView::from_slots(slots);
+        DateTimePartsData::validate(dtype, slots.days, slots.seconds, slots.subseconds, len)
     }
 
     fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
@@ -181,12 +175,20 @@ impl VTable for DateTimeParts {
     }
 
     fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
-        SLOT_NAMES[idx].to_string()
+        DateTimePartsSlots::NAMES[idx].to_string()
     }
 
     fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
+        let array = require_child!(array, array.days(), DateTimePartsSlots::DAYS => Primitive);
+        let array =
+            require_child!(array, array.seconds(), DateTimePartsSlots::SECONDS => AnyCanonical);
+        let array = require_child!(array, array.subseconds(), DateTimePartsSlots::SUBSECONDS => AnyCanonical);
+
+        let dtype = array.dtype().clone();
+        let parts = array.into_parts();
+
         Ok(ExecutionResult::done(
-            decode_to_temporal(&array, ctx)?.into_array(),
+            decode_to_temporal(parts, &dtype, ctx)?.into_array(),
         ))
     }
 
@@ -208,46 +210,58 @@ impl VTable for DateTimeParts {
     }
 }
 
-/// The days component of the datetime, stored as an integer array.
-pub(super) const DAYS_SLOT: usize = 0;
-/// The seconds component of the datetime (within the day).
-pub(super) const SECONDS_SLOT: usize = 1;
-/// The sub-second component of the datetime.
-pub(super) const SUBSECONDS_SLOT: usize = 2;
-pub(super) const NUM_SLOTS: usize = 3;
-pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["days", "seconds", "subseconds"];
+#[array_slots(DateTimeParts)]
+pub struct DateTimePartsSlots {
+    /// The days component of the datetime, stored as an integer array.
+    pub days: ArrayRef,
+    /// The seconds component of the datetime (within the day).
+    pub seconds: ArrayRef,
+    /// The sub-second component of the datetime.
+    pub subseconds: ArrayRef,
+}
 
 #[derive(Clone, Debug)]
 pub struct DateTimePartsData {}
 
-impl Display for DateTimePartsData {
-    fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result {
-        Ok(())
-    }
+pub struct DateTimePartsParts {
+    pub days: ArrayRef,
+    pub seconds: ArrayRef,
+    pub subseconds: ArrayRef,
 }
 
-pub trait DateTimePartsArrayExt: TypedArrayRef<DateTimeParts> {
-    fn days(&self) -> &ArrayRef {
-        self.as_ref().slots()[DAYS_SLOT]
-            .as_ref()
-            .vortex_expect("DateTimePartsArray days slot")
-    }
+pub trait DateTimePartsOwnedExt {
+    fn into_parts(self) -> DateTimePartsParts;
+}
 
-    fn seconds(&self) -> &ArrayRef {
-        self.as_ref().slots()[SECONDS_SLOT]
-            .as_ref()
-            .vortex_expect("DateTimePartsArray seconds slot")
+impl DateTimePartsOwnedExt for Array<DateTimeParts> {
+    fn into_parts(self) -> DateTimePartsParts {
+        match self.try_into_parts() {
+            Ok(parts) => {
+                let slots = DateTimePartsSlots::from_slots(parts.slots);
+                DateTimePartsParts {
+                    days: slots.days,
+                    seconds: slots.seconds,
+                    subseconds: slots.subseconds,
+                }
+            }
+            Err(array) => {
+                let view = DateTimePartsSlotsView::from_slots(array.as_ref().slots());
+                DateTimePartsParts {
+                    days: view.days.clone(),
+                    seconds: view.seconds.clone(),
+                    subseconds: view.subseconds.clone(),
+                }
+            }
+        }
     }
+}
 
-    fn subseconds(&self) -> &ArrayRef {
-        self.as_ref().slots()[SUBSECONDS_SLOT]
-            .as_ref()
-            .vortex_expect("DateTimePartsArray subseconds slot")
+impl Display for DateTimePartsData {
+    fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result {
+        Ok(())
     }
 }
 
-impl<T: TypedArrayRef<DateTimeParts>> DateTimePartsArrayExt for T {}
-
 #[derive(Clone, Debug)]
 pub struct DateTimeParts;