Implemented Multiple DatePart Caching/Retrieving

src/local/src/tests/snapshots/liquid_cache_local__tests__date_optimizer__date_extraction.snap

@@ -4,18 +4,18 @@ expression: stats
 ---
 entries.total: 123
 entries.after_first_run: 123
-entries.memory.arrow: 4
-entries.memory.liquid: 38
-entries.memory.hybrid_liquid: 81
+entries.memory.arrow: 2
+entries.memory.liquid: 26
+entries.memory.hybrid_liquid: 95
 entries.disk.liquid: 0
 entries.disk.arrow: 0
-usage.memory_bytes: 1032554
-usage.disk_bytes: 1081512
+usage.memory_bytes: 976214
+usage.disk_bytes: 1268440
 runtime.get_arrow_array_calls: 0
 runtime.get_with_selection_calls: 123
 runtime.get_with_predicate_calls: 0
 runtime.get_predicate_hybrid_success: 0
 runtime.get_predicate_hybrid_needs_io: 0
 runtime.get_predicate_hybrid_unsupported: 0
 runtime.try_read_liquid_calls: 0
-runtime.hit_date32_expression_calls: 81
+runtime.hit_date32_expression_calls: 95

src/local/src/tests/snapshots/liquid_cache_local__tests__date_optimizer__date_extraction_case2.snap

@@ -4,18 +4,18 @@ expression: stats
 ---
 entries.total: 123
 entries.after_first_run: 123
-entries.memory.arrow: 4
-entries.memory.liquid: 38
-entries.memory.hybrid_liquid: 81
+entries.memory.arrow: 2
+entries.memory.liquid: 26
+entries.memory.hybrid_liquid: 95
 entries.disk.liquid: 0
 entries.disk.arrow: 0
-usage.memory_bytes: 1032554
-usage.disk_bytes: 1081512
+usage.memory_bytes: 976214
+usage.disk_bytes: 1268440
 runtime.get_arrow_array_calls: 0
 runtime.get_with_selection_calls: 246
 runtime.get_with_predicate_calls: 0
 runtime.get_predicate_hybrid_success: 0
 runtime.get_predicate_hybrid_needs_io: 0
 runtime.get_predicate_hybrid_unsupported: 0
 runtime.try_read_liquid_calls: 0
-runtime.hit_date32_expression_calls: 162
+runtime.hit_date32_expression_calls: 190

src/local/src/tests/snapshots/liquid_cache_local__tests__date_optimizer__date_extraction_day.snap

@@ -4,18 +4,18 @@ expression: stats
 ---
 entries.total: 123
 entries.after_first_run: 123
-entries.memory.arrow: 4
-entries.memory.liquid: 38
-entries.memory.hybrid_liquid: 81
+entries.memory.arrow: 2
+entries.memory.liquid: 26
+entries.memory.hybrid_liquid: 95
 entries.disk.liquid: 0
 entries.disk.arrow: 0
-usage.memory_bytes: 1032554
-usage.disk_bytes: 1081512
+usage.memory_bytes: 976214
+usage.disk_bytes: 1268440
 runtime.get_arrow_array_calls: 0
 runtime.get_with_selection_calls: 123
 runtime.get_with_predicate_calls: 0
 runtime.get_predicate_hybrid_success: 0
 runtime.get_predicate_hybrid_needs_io: 0
 runtime.get_predicate_hybrid_unsupported: 0
 runtime.try_read_liquid_calls: 0
-runtime.hit_date32_expression_calls: 81
+runtime.hit_date32_expression_calls: 95

src/local/src/tests/snapshots/liquid_cache_local__tests__date_optimizer__date_extraction_month.snap

@@ -5,17 +5,17 @@ expression: stats
 entries.total: 123
 entries.after_first_run: 123
 entries.memory.arrow: 4
-entries.memory.liquid: 46
-entries.memory.hybrid_liquid: 73
+entries.memory.liquid: 38
+entries.memory.hybrid_liquid: 81
 entries.disk.liquid: 0
 entries.disk.arrow: 0
-usage.memory_bytes: 1023346
-usage.disk_bytes: 974696
+usage.memory_bytes: 1032554
+usage.disk_bytes: 1081512
 runtime.get_arrow_array_calls: 0
 runtime.get_with_selection_calls: 123
 runtime.get_with_predicate_calls: 0
 runtime.get_predicate_hybrid_success: 0
 runtime.get_predicate_hybrid_needs_io: 0
 runtime.get_predicate_hybrid_unsupported: 0
 runtime.try_read_liquid_calls: 0
-runtime.hit_date32_expression_calls: 73
+runtime.hit_date32_expression_calls: 81

src/parquet/src/optimizers/lineage_opt.rs

@@ -2,7 +2,7 @@
 //! It then attaches the metadata to schema adapter, which is then passed to the physical plan.
 //! The physical optimizer will move the metadata to the fields of the schema.
 
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::str::FromStr;
 use std::sync::{Arc, Mutex, OnceLock};
 
@@ -46,7 +46,7 @@ impl SupportedIntervalUnit {
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub(crate) struct DateExtraction {
     pub(crate) column: Column,
-    pub(crate) component: SupportedIntervalUnit,
+    pub(crate) components: HashSet<SupportedIntervalUnit>,
 }
 
 /// Metadata describing a Variant column that participates in a `variant_get`.
@@ -59,10 +59,36 @@ pub(crate) struct VariantExtraction {
 /// Annotation that should be attached to a column in the file schema.
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub(crate) enum ColumnAnnotation {
-    DatePart(SupportedIntervalUnit),
+    DatePart(HashSet<SupportedIntervalUnit>),
     VariantPath(String),
 }
 
+impl ColumnAnnotation {
+    /// Serialize DatePart units to a comma-separated string.
+    /// Returns None if this is not a DatePart annotation.
+    pub(crate) fn serialize_date_part(&self) -> Option<String> {
+        match self {
+            ColumnAnnotation::DatePart(units) => {
+                let mut sorted_units: Vec<&SupportedIntervalUnit> = units.iter().collect();
+                // Sort by a consistent order: Year, Month, Day
+                sorted_units.sort_by_key(|unit| match unit {
+                    SupportedIntervalUnit::Year => 0,
+                    SupportedIntervalUnit::Month => 1,
+                    SupportedIntervalUnit::Day => 2,
+                });
+                Some(
+                    sorted_units
+                        .iter()
+                        .map(|unit| unit.metadata_value())
+                        .collect::<Vec<_>>()
+                        .join(","),
+                )
+            }
+            ColumnAnnotation::VariantPath(_) => None,
+        }
+    }
+}
+
 /// Logical optimizer that analyses the logical plan to detect columns that
 /// are only used via compatible `EXTRACT` or `variant_get` projections.
 #[derive(Debug, Default)]
@@ -499,24 +525,40 @@ impl TableColumnUsage {
         let mut extractions = Vec::new();
         for (key, stats) in self.usage.iter() {
             if matches!(stats.data_type, DataType::Date32) {
-                // Check if every usage's first operation is Extract with the same unit
-                let first_unit = stats.usages.first().and_then(|usage| {
-                    if let Some(Operation::Extract(unit)) = usage.first() {
-                        Some(unit)
-                    } else {
-                        None
+                // Collect all extract units from paths where the first n operations are all extracts
+                let mut all_units = HashSet::new();
+                let mut all_paths_valid = true;
+
+                for usage in &stats.usages {
+                    // Collect all Extract units from the leading sequence of extracts
+                    let mut path_units = HashSet::new();
+                    for op in usage {
+                        match op {
+                            Operation::Extract(unit) => {
+                                path_units.insert(unit);
+                            }
+                            _ => {
+                                // Stop at first non-extract operation
+                                break;
+                            }
+                        }
                     }
-                });
-                if let Some(first_unit) = first_unit {
-                    let all_matches = stats.usages.iter().all(|usage| {
-                        matches!(usage.first(), Some(Operation::Extract(unit)) if unit == first_unit)
-                    });
-                    if all_matches {
-                        extractions.push(DateExtraction {
-                            column: key.to_column(),
-                            component: *first_unit,
-                        });
+
+                    if path_units.is_empty() {
+                        // This path doesn't start with Extract, so skip this column
+                        all_paths_valid = false;
+                        break;
                     }
+
+                    // Union the units from this path into the overall set
+                    all_units.extend(path_units);
+                }
+
+                if all_paths_valid && !all_units.is_empty() {
+                    extractions.push(DateExtraction {
+                        column: key.to_column(),
+                        components: all_units,
+                    });
                 }
             }
         }
@@ -554,7 +596,7 @@ fn build_annotation_map(
     for extraction in date_findings {
         annotations.insert(
             ColumnKey::from_column(&extraction.column),
-            ColumnAnnotation::DatePart(extraction.component),
+            ColumnAnnotation::DatePart(extraction.components.clone()),
         );
     }
     for extraction in variant_findings {
@@ -1034,10 +1076,19 @@ mod tests {
         let expected_field_metadata = expected
             .iter()
             .map(|extraction| {
-                (
-                    extraction.column.name().to_string(),
-                    extraction.component.metadata_value().to_string(),
-                )
+                let mut sorted_units: Vec<&SupportedIntervalUnit> =
+                    extraction.components.iter().collect();
+                sorted_units.sort_by_key(|unit| match unit {
+                    SupportedIntervalUnit::Year => 0,
+                    SupportedIntervalUnit::Month => 1,
+                    SupportedIntervalUnit::Day => 2,
+                });
+                let metadata_value = sorted_units
+                    .iter()
+                    .map(|unit| unit.metadata_value())
+                    .collect::<Vec<_>>()
+                    .join(",");
+                (extraction.column.name().to_string(), metadata_value)
             })
             .collect::<HashMap<String, String>>();
         assert_eq!(field_metadata_map, expected_field_metadata);
@@ -1048,20 +1099,14 @@ mod tests {
         general_test(
             "SELECT EXTRACT(YEAR FROM table_a.date) AS year, EXTRACT(DAY FROM table_b.date) AS day FROM table_a INNER JOIN table_b ON table_a.event_ts = table_b.event_ts",
             vec![
-                DateExtraction { column: Column::new(Some("table_a"), "date"), component: SupportedIntervalUnit::Year },
-                DateExtraction { column: Column::new(Some("table_b"), "date"), component: SupportedIntervalUnit::Day },
-            ],
-        )
-        .await;
-    }
-
-    #[tokio::test]
-    async fn single_table_multiple_extracts() {
-        general_test(
-            "SELECT EXTRACT(YEAR FROM date_copy) AS year, EXTRACT(DAY FROM date) AS day FROM table_a",
-            vec![
-                DateExtraction { column: Column::new(Some("table_a"), "date"), component: SupportedIntervalUnit::Day },
-                DateExtraction { column: Column::new(Some("table_a"), "date_copy"), component: SupportedIntervalUnit::Year },
+                DateExtraction {
+                    column: Column::new(Some("table_a"), "date"), 
+                    components: HashSet::from([SupportedIntervalUnit::Year]),
+                },
+                DateExtraction {
+                    column: Column::new(Some("table_b"), "date"), 
+                    components: HashSet::from([SupportedIntervalUnit::Day]),
+                },
             ],
         )
         .await;
@@ -1079,7 +1124,7 @@ mod tests {
         ];
         let expected = vec![DateExtraction {
             column: Column::new(Some("table_a"), "date"),
-            component: SupportedIntervalUnit::Day,
+            components: HashSet::from([SupportedIntervalUnit::Day]),
         }];
         for sql in statements {
             general_test(sql, expected.clone()).await;
@@ -1125,11 +1170,9 @@ mod tests {
     #[tokio::test]
     async fn inconsistent_extracts_are_ignored() {
         let statements = vec![
-            "SELECT EXTRACT(DAY FROM date) AS day, EXTRACT(MONTH FROM date) AS month FROM table_a",
             "SELECT EXTRACT(DAY FROM date + INTERVAL '1 day') AS day FROM table_a",
             "SELECT date FROM table_a",
             "SELECT EXTRACT(DAY FROM table_a.date) AS day FROM table_a INNER JOIN table_b ON table_a.date = table_b.date",
-            "SELECT (SELECT MAX(EXTRACT(DAY FROM date)) FROM table_a) AS max_day, (SELECT MIN(EXTRACT(Month FROM date)) FROM table_a) AS min_day",
             "SELECT EXTRACT(YEAR FROM event_ts) AS year FROM table_a", // todo: time stamp is not supported yet.
         ];
 
@@ -1138,6 +1181,21 @@ mod tests {
         }
     }
 
+    #[tokio::test]
+    async fn single_table_multiple_extracts() {
+        let statements = vec![
+            "SELECT EXTRACT(DAY FROM date) AS day, EXTRACT(MONTH FROM date) AS month FROM table_a",
+            "SELECT (SELECT MAX(EXTRACT(DAY FROM date)) FROM table_a) AS max_day, (SELECT MIN(EXTRACT(Month FROM date)) FROM table_a) AS min_day",
+        ];
+        let expected = vec![DateExtraction {
+            column: Column::new(Some("table_a"), "date"),
+            components: HashSet::from([SupportedIntervalUnit::Month, SupportedIntervalUnit::Day]),
+        }];
+        for sql in statements {
+            general_test(sql, expected.clone()).await;
+        }
+    }
+
     #[tokio::test]
     async fn variant_get_metadata_is_propagated() {
         let temp_dir = TempDir::new().unwrap();

src/parquet/src/optimizers/mod.rs

@@ -90,9 +90,11 @@ pub fn rewrite_data_source_plan(
                             {
                                 let (metadata_key, metadata_value): (&str, String) =
                                     match annotation {
-                                        ColumnAnnotation::DatePart(unit) => (
+                                        ColumnAnnotation::DatePart(_) => (
                                             DATE_MAPPING_METADATA_KEY,
-                                            unit.metadata_value().to_string(),
+                                            annotation
+                                                .serialize_date_part()
+                                                .expect("DatePart should serialize"),
                                         ),
                                         ColumnAnnotation::VariantPath(path) => {
                                             (VARIANT_MAPPING_METADATA_KEY, path)

src/storage/bench/squeeze_date32.rs

@@ -1,9 +1,10 @@
 use arrow::array::{Array, ArrayRef, cast::AsArray};
+use arrow::compute::DatePart;
 use arrow::datatypes::Date32Type;
 use clap::Parser;
 use datafusion::prelude::*;
 use futures::StreamExt;
-use liquid_cache_storage::liquid_array::{Date32Field, LiquidPrimitiveArray, SqueezedDate32Array};
+use liquid_cache_storage::liquid_array::{LiquidPrimitiveArray, SqueezedDate32Array};
 
 #[global_allocator]
 static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
@@ -76,9 +77,9 @@ async fn run_for_column(ctx: &SessionContext, col: &str, limit: Option<usize>) {
         let liquid = LiquidPrimitiveArray::<Date32Type>::from_arrow_array(prim.clone());
         total_liquid_bytes += liquid.get_array_memory_size();
 
-        let squeezed_year = SqueezedDate32Array::from_liquid_date32(&liquid, Date32Field::Year);
-        let squeezed_month = SqueezedDate32Array::from_liquid_date32(&liquid, Date32Field::Month);
-        let squeezed_day = SqueezedDate32Array::from_liquid_date32(&liquid, Date32Field::Day);
+        let squeezed_year = SqueezedDate32Array::from_liquid_date32(&liquid, DatePart::Year);
+        let squeezed_month = SqueezedDate32Array::from_liquid_date32(&liquid, DatePart::Month);
+        let squeezed_day = SqueezedDate32Array::from_liquid_date32(&liquid, DatePart::Day);
 
         total_year_bytes += squeezed_year.get_array_memory_size();
         total_month_bytes += squeezed_month.get_array_memory_size();