feat(cubesql): Filter push down for date_part('year') = ?y & date_part('quarter') = ?q

Author: ovr

rust/cubesql/cubesql/src/compile/mod.rs

@@ -9462,6 +9462,84 @@ ORDER BY "source"."str0" ASC
         )
     }
 
+    #[tokio::test]
+    async fn test_filter_extract_by_year_and_quarter() {
+        init_testing_logger();
+
+        let logical_plan = convert_select_to_query_plan(
+            r#"
+            SELECT
+                COUNT(*) AS "count",
+                EXTRACT(YEAR FROM "KibanaSampleDataEcommerce"."order_date") AS "yr:completedAt:ok"
+            FROM "public"."KibanaSampleDataEcommerce" "KibanaSampleDataEcommerce"
+            WHERE EXTRACT(YEAR FROM "KibanaSampleDataEcommerce"."order_date") = 2019 AND EXTRACT(QUARTER FROM "KibanaSampleDataEcommerce"."order_date") = 2
+            GROUP BY 2
+            ;"#
+                .to_string(),
+            DatabaseProtocol::PostgreSQL,
+        )
+            .await
+            .as_logical_plan();
+
+        assert_eq!(
+            logical_plan.find_cube_scan().request,
+            V1LoadRequestQuery {
+                measures: Some(vec!["KibanaSampleDataEcommerce.count".to_string()]),
+                dimensions: Some(vec![]),
+                segments: Some(vec![]),
+                time_dimensions: Some(vec![V1LoadRequestQueryTimeDimension {
+                    dimension: "KibanaSampleDataEcommerce.order_date".to_string(),
+                    granularity: Some("year".to_string()),
+                    date_range: Some(json!(vec![
+                        "2019-04-01".to_string(),
+                        "2019-06-31".to_string(),
+                    ])),
+                },]),
+                order: Some(vec![]),
+                ..Default::default()
+            }
+        )
+    }
+
+    #[tokio::test]
+    async fn test_filter_extract_by_year_and_month() {
+        init_testing_logger();
+
+        let logical_plan = convert_select_to_query_plan(
+            r#"
+            SELECT
+                COUNT(*) AS "count",
+                EXTRACT(YEAR FROM "KibanaSampleDataEcommerce"."order_date") AS "yr:completedAt:ok"
+            FROM "public"."KibanaSampleDataEcommerce" "KibanaSampleDataEcommerce"
+            WHERE EXTRACT(YEAR FROM "KibanaSampleDataEcommerce"."order_date") = 2019 AND EXTRACT(MONTH FROM "KibanaSampleDataEcommerce"."order_date") = 2
+            GROUP BY 2
+            ;"#
+                .to_string(),
+            DatabaseProtocol::PostgreSQL,
+        )
+            .await
+            .as_logical_plan();
+
+        assert_eq!(
+            logical_plan.find_cube_scan().request,
+            V1LoadRequestQuery {
+                measures: Some(vec!["KibanaSampleDataEcommerce.count".to_string()]),
+                dimensions: Some(vec![]),
+                segments: Some(vec![]),
+                time_dimensions: Some(vec![V1LoadRequestQueryTimeDimension {
+                    dimension: "KibanaSampleDataEcommerce.order_date".to_string(),
+                    granularity: Some("year".to_string()),
+                    date_range: Some(json!(vec![
+                        "2019-02-01".to_string(),
+                        "2019-02-28".to_string(),
+                    ])),
+                },]),
+                order: Some(vec![]),
+                ..Default::default()
+            }
+        )
+    }
+
     #[tokio::test]
     async fn test_tableau_filter_extract_by_year() {
         init_testing_logger();

rust/cubesql/cubesql/src/compile/rewrite/rules/filters.rs

@@ -1726,6 +1726,44 @@ impl RewriteRules for FilterRules {
                     "?filter_aliases",
                 ),
             ),
+            // EXTRACT(YEAR FROM "KibanaSampleDataEcommerce"."order_date") = 2019
+            // AND EXTRACT(MONTH FROM "KibanaSampleDataEcommerce"."order_date") = 3
+            transforming_rewrite(
+                "extract-date-range-and-gran-equals",
+                filter_op(
+                    filter_op_filters(
+                        filter_member("?member", "FilterMemberOp:inDateRange", "?values"),
+                        filter_replacer(
+                            binary_expr(
+                                self.fun_expr(
+                                    "DatePart",
+                                    vec![literal_expr("?granularity"), column_expr("?column")],
+                                ),
+                                "=",
+                                literal_expr("?value"),
+                            ),
+                            "?alias_to_cube",
+                            "?members",
+                            "?filter_aliases",
+                        ),
+                    ),
+                    "FilterOpOp:and",
+                ),
+                filter_member("?member", "FilterMemberOp:inDateRange", "?new_values"),
+                self.transform_filter_extract_date_range_and_trunc_gran_equals(
+                    "?member",
+                    "?values",
+                    "?granularity",
+                    "?column",
+                    "?value",
+                    "?alias_to_cube",
+                    "?members",
+                    "?filter_aliases",
+                    "?new_values",
+                ),
+            ),
+            // TODO: Introduce rule to unwrap TRUNC(EXTRACT(?granularity FROM ?column_expr))
+            //
             // TRUNC(EXTRACT(YEAR FROM "KibanaSampleDataEcommerce"."order_date")) = 2019
             // AND TRUNC(EXTRACT(MONTH FROM "KibanaSampleDataEcommerce"."order_date")) = 3
             transforming_rewrite(
@@ -1765,6 +1803,7 @@ impl RewriteRules for FilterRules {
                     "?new_values",
                 ),
             ),
+            // TODO: Introduce rule to unwrap TRUNC(EXTRACT(?granularity FROM ?column_expr))
             // When the filter set above is paired with other filters, it needs to be
             // regrouped for the above rewrite rule to match
             rewrite(
@@ -1829,6 +1868,7 @@ impl RewriteRules for FilterRules {
                     "FilterOpOp:and",
                 ),
             ),
+            // TODO: Introduce rule to unwrap TRUNC(EXTRACT(?granularity FROM ?column_expr))
             // The filter set above may be inverted, let's account for that as well
             rewrite(
                 "extract-date-range-and-trunc-reverse",
@@ -1877,6 +1917,7 @@ impl RewriteRules for FilterRules {
                     "FilterOpOp:and",
                 ),
             ),
+            // TODO: Introduce rule to unwrap TRUNC(EXTRACT(?granularity FROM ?column_expr))
             rewrite(
                 "extract-date-range-and-trunc-reverse-nested",
                 filter_op(
@@ -3991,6 +4032,7 @@ impl FilterRules {
                     if start_date_year != end_date.year() {
                         return false;
                     }
+
                     // Month value must be valid
                     if !(1..=12).contains(&value) {
                         return false;
@@ -4022,6 +4064,48 @@ impl FilterRules {
                         new_end_date.format("%Y-%m-%d").to_string(),
                     ]
                 }
+                "quarter" | "qtr" => {
+                    // Check that the range only covers one year
+                    let start_date_year = start_date.year();
+                    if start_date_year != end_date.year() {
+                        return false;
+                    }
+
+                    // Quarter value must be valid (1-4)
+                    if !(1..=4).contains(&value) {
+                        return false;
+                    }
+
+                    let quarter_start_month = (value - 1) * 3 + 1;
+
+                    // Obtain the new range
+                    let Some(new_start_date) =
+                        NaiveDate::from_ymd_opt(start_date_year, quarter_start_month as u32, 1)
+                    else {
+                        return false;
+                    };
+
+                    let Some(new_end_date) = new_start_date
+                        .checked_add_months(Months::new(3))
+                        .and_then(|date| date.checked_sub_days(Days::new(1)))
+                    else {
+                        return false;
+                    };
+
+                    // Paranoid check, If the resulting range is outside of the original range, we can't merge
+                    // the filters
+                    if new_start_date > end_date || new_end_date < start_date {
+                        return false;
+                    }
+
+                    let new_start_date = max(new_start_date, start_date);
+                    let new_end_date = min(new_end_date, end_date);
+
+                    vec![
+                        new_start_date.format("%Y-%m-%d").to_string(),
+                        new_end_date.format("%Y-%m-%d").to_string(),
+                    ]
+                }
                 // TODO: handle more granularities
                 _ => return false,
             };