wip datetimeparts fixes

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

Author: joseph-isaacs

bench-vortex/Cargo.toml

@@ -73,7 +73,7 @@ tokio-stream = { workspace = true }
 tokio-util = { workspace = true }
 tpchgen = { workspace = true }
 tpchgen-arrow = { workspace = true }
-tracing = { workspace = true, features = ["max_level_debug"] }
+tracing = { workspace = true, features = ["max_level_trace"] }
 tracing-perfetto = { workspace = true }
 tracing-subscriber = { workspace = true, features = [
     "env-filter",

encodings/datetime-parts/src/array.rs

@@ -12,6 +12,7 @@ use vortex_array::ArrayHash;
 use vortex_array::ArrayRef;
 use vortex_array::Canonical;
 use vortex_array::DeserializeMetadata;
+use vortex_array::IntoArray;
 use vortex_array::Precision;
 use vortex_array::ProstMetadata;
 use vortex_array::SerializeMetadata;
@@ -40,6 +41,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 +162,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,65 @@
+// 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)]);
+
+#[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/expr/transform/partition.rs

@@ -71,7 +71,7 @@ where
             Nullability::NonNullable,
         );
 
-        let expr = expr.optimize_recursive(scope)?;
+        // let expr = expr.optimize_recursive(scope)?;
         let expr_dtype = expr.return_dtype(scope)?;
 
         partitions.push(expr);

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,17 @@ impl BtrBlocksCompressor {
                 if let Ok(temporal_array) = TemporalArray::try_from(ext_array.to_array())
                     && let TemporalMetadata::Timestamp(..) = temporal_array.temporal_metadata()
                 {
+                    let temporal_array_ref = temporal_array.clone().into_array();
+                    if temporal_array_ref
+                        .as_ref()
+                        .is_constant_opts(Cost::Canonicalize)
+                    {
+                        return Ok(ConstantArray::new(
+                            temporal_array_ref.scalar_at(0),
+                            ext_array.len(),
+                        )
+                        .into_array());
+                    }
                     return compress_temporal(temporal_array);
                 }