wip datetimeparts fixes (#5778)

Signed-off-by: Joe Isaacs <[email protected]>

Author: joseph-isaacs

encodings/datetime-parts/src/array.rs

@@ -40,6 +40,8 @@ use vortex_error::vortex_bail;
 use vortex_error::vortex_ensure;
 use vortex_error::vortex_err;
 
+use crate::compute::rules::PARENT_RULES;
+
 vtable!(DateTimeParts);
 
 #[derive(Clone, prost::Message)]
@@ -159,6 +161,14 @@ impl VTable for DateTimePartsVTable {
 
         Ok(())
     }
+
+    fn reduce_parent(
+        array: &Self::Array,
+        parent: &ArrayRef,
+        child_idx: usize,
+    ) -> VortexResult<Option<ArrayRef>> {
+        PARENT_RULES.evaluate(array, parent, child_idx)
+    }
 }
 
 #[derive(Clone, Debug)]

encodings/datetime-parts/src/compute/mod.rs

@@ -6,6 +6,7 @@ mod compare;
 mod filter;
 mod is_constant;
 mod mask;
+pub(super) mod rules;
 mod take;
 
 #[cfg(test)]

encodings/datetime-parts/src/compute/rules.rs

@@ -0,0 +1,67 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_array::ArrayRef;
+use vortex_array::IntoArray;
+use vortex_array::arrays::ConstantArray;
+use vortex_array::arrays::ConstantVTable;
+use vortex_array::arrays::FilterArray;
+use vortex_array::arrays::FilterVTable;
+use vortex_array::matchers::Exact;
+use vortex_array::optimizer::ArrayOptimizer;
+use vortex_array::optimizer::rules::ArrayParentReduceRule;
+use vortex_array::optimizer::rules::ParentRuleSet;
+use vortex_error::VortexExpect;
+use vortex_error::VortexResult;
+
+use crate::DateTimePartsArray;
+use crate::DateTimePartsVTable;
+
+pub(crate) const PARENT_RULES: ParentRuleSet<DateTimePartsVTable> =
+    ParentRuleSet::new(&[ParentRuleSet::lift(&DTPFilterPushDownRule)]);
+
+/// Push the filter into the days column of a date time parts, we could extend this to other fields
+/// but its less clear if that is beneficial.
+#[derive(Debug)]
+struct DTPFilterPushDownRule;
+
+impl ArrayParentReduceRule<DateTimePartsVTable> for DTPFilterPushDownRule {
+    type Parent = Exact<FilterVTable>;
+
+    fn parent(&self) -> Self::Parent {
+        Exact::from(&FilterVTable)
+    }
+
+    fn reduce_parent(
+        &self,
+        child: &DateTimePartsArray,
+        parent: &FilterArray,
+        child_idx: usize,
+    ) -> VortexResult<Option<ArrayRef>> {
+        if child_idx != 0 {
+            return Ok(None);
+        }
+
+        if !child.seconds().is::<ConstantVTable>() || !child.subseconds().is::<ConstantVTable>() {
+            return Ok(None);
+        }
+
+        DateTimePartsArray::try_new(
+            child.dtype().clone(),
+            FilterArray::new(child.days().clone(), parent.filter_mask().clone())
+                .into_array()
+                .optimize()?,
+            ConstantArray::new(
+                child.seconds().as_constant().vortex_expect("constant"),
+                parent.filter_mask().true_count(),
+            )
+            .into_array(),
+            ConstantArray::new(
+                child.subseconds().as_constant().vortex_expect("constant"),
+                parent.filter_mask().true_count(),
+            )
+            .into_array(),
+        )
+        .map(|x| Some(x.into_array()))
+    }
+}

vortex-array/src/arrays/datetime/mod.rs

@@ -201,6 +201,12 @@ impl TemporalArray {
     }
 }
 
+impl AsRef<dyn Array> for TemporalArray {
+    fn as_ref(&self) -> &dyn Array {
+        self.ext.as_ref()
+    }
+}
+
 impl From<TemporalArray> for ArrayRef {
     fn from(value: TemporalArray) -> Self {
         value.ext.into_array()

vortex-btrblocks/src/lib.rs

@@ -38,12 +38,14 @@ use vortex_array::ArrayRef;
 use vortex_array::Canonical;
 use vortex_array::IntoArray;
 use vortex_array::ToCanonical;
+use vortex_array::arrays::ConstantArray;
 use vortex_array::arrays::ExtensionArray;
 use vortex_array::arrays::FixedSizeListArray;
 use vortex_array::arrays::ListArray;
 use vortex_array::arrays::StructArray;
 use vortex_array::arrays::TemporalArray;
 use vortex_array::arrays::list_from_list_view;
+use vortex_array::compute::Cost;
 use vortex_array::vtable::VTable;
 use vortex_array::vtable::ValidityHelper;
 use vortex_dtype::DType;
@@ -475,6 +477,13 @@ impl BtrBlocksCompressor {
                 if let Ok(temporal_array) = TemporalArray::try_from(ext_array.to_array())
                     && let TemporalMetadata::Timestamp(..) = temporal_array.temporal_metadata()
                 {
+                    if temporal_array.as_ref().is_constant_opts(Cost::Canonicalize) {
+                        return Ok(ConstantArray::new(
+                            temporal_array.as_ref().scalar_at(0),
+                            ext_array.len(),
+                        )
+                        .into_array());
+                    }
                     return compress_temporal(temporal_array);
                 }
 

vortex-duckdb/src/convert/scalar.rs

@@ -201,8 +201,11 @@ impl ToDuckDBScalar for ExtScalar<'_> {
                 _ => vortex_bail!("cannot have TimeUnit {unit}, so represent a day"),
             },
             TemporalMetadata::Timestamp(unit, tz) => {
-                if tz.is_some() {
-                    todo!("timezones to duckdb scalar")
+                if let Some(tz) = tz {
+                    if tz != "UTC" {
+                        todo!()
+                    }
+                    return Ok(Value::new_timestamp_tz(value()?));
                 }
                 match unit {
                     TimeUnit::Nanoseconds => Ok(Value::new_timestamp_ns(value()?)),

vortex-duckdb/src/duckdb/value.rs

@@ -197,6 +197,14 @@ impl Value {
         }
     }
 
+    pub fn new_timestamp_tz(micros: i64) -> Self {
+        unsafe {
+            Self::own(cpp::duckdb_create_timestamp_tz(cpp::duckdb_timestamp {
+                micros,
+            }))
+        }
+    }
+
     pub fn new_timestamp_ns(nanos: i64) -> Self {
         unsafe {
             Self::own(cpp::duckdb_create_timestamp_ns(cpp::duckdb_timestamp_ns {