[WIP] feat(cubesql): Rewrite measure filters in query as PatchMeasure

Author: mcheshkov

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

@@ -1778,7 +1778,6 @@ fn literal_bool(literal_bool: bool) -> String {
     format!("(LiteralExpr LiteralExprValue:b:{})", literal_bool)
 }
 
-#[allow(dead_code)]
 fn literal_null() -> String {
     format!("(LiteralExpr LiteralExprValue:null)")
 }

rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/aggregate.rs

@@ -2,8 +2,9 @@ use crate::{
     compile::{
         engine::udf::{MEASURE_UDAF_NAME, PATCH_MEASURE_UDAF_NAME},
         rewrite::{
-            aggregate, alias_expr, cube_scan_wrapper, grouping_set_expr, original_expr_name,
-            rewrite,
+            aggregate, alias_expr,
+            analysis::ConstantFolding,
+            cube_scan_wrapper, grouping_set_expr, original_expr_name, rewrite,
             rewriter::{CubeEGraph, CubeRewrite},
             rules::{members::MemberRules, wrapper::WrapperRules},
             subquery, transforming_chain_rewrite, transforming_rewrite, wrapped_select,
@@ -319,6 +320,96 @@ impl WrapperRules {
                 "GroupingSetExprMembers",
             );
         }
+
+        use crate::compile::rewrite::{
+            agg_fun_expr, binary_expr, case_expr, column_expr, literal_null, udaf_expr,
+        };
+
+        // incoming structure: agg_fun(?name, case(?cond, (?when_value, measure_column)))
+        // optional "else null" is fine
+        // only single when-then
+        rules.extend(vec![
+            transforming_chain_rewrite(
+                "wrapper-push-down-aggregation-over-filtered-measure",
+                wrapper_pushdown_replacer("?aggr_expr", "?context"),
+                vec![
+                    (
+                        "?aggr_expr",
+                        agg_fun_expr(
+                            "?fun",
+                            vec![case_expr(
+                                Some("?case_expr".to_string()),
+                                vec![("?literal".to_string(), column_expr("?measure_column"))],
+                                // TODO make `ELSE NULL` optional and/or add generic rewrite to normalize it
+                                Some(literal_null()),
+                            )],
+                            "?distinct",
+                        ),
+                    ),
+                    (
+                        "?context",
+                        wrapper_replacer_context(
+                            "?alias_to_cube",
+                            "WrapperReplacerContextPushToCube:true",
+                            "?in_projection",
+                            "?cube_members",
+                            "?grouped_subqueries",
+                            "?ungrouped_scan",
+                        ),
+                    ),
+                ],
+                alias_expr(
+                    udaf_expr(
+                        PATCH_MEASURE_UDAF_NAME,
+                        vec![
+                            column_expr("?measure_column"),
+                            // TODO support doing both: changing agg type and adding filters
+                            literal_null(),
+                            wrapper_pushdown_replacer(
+                                // = is a proper way to filter here:
+                                // CASE NULL WHEN ... will return null
+                                // So NULL in ?case_expr is equivalent to hitting ELSE branch
+                                // TODO add "is not null" to cond? just to make is always boolean
+                                binary_expr("?case_expr", "=", "?literal"),
+                                "?context",
+                            ),
+                        ],
+                    ),
+                    "?out_measure_alias",
+                ),
+                self.transform_filtered_measure("?aggr_expr", "?literal", "?out_measure_alias"),
+            ),
+            transforming_rewrite(
+                "wrapper-pull-up-aggregation-over-filtered-measure",
+                udaf_expr(
+                    PATCH_MEASURE_UDAF_NAME,
+                    vec![
+                        column_expr("?measure_column"),
+                        "?new_agg_type".to_string(),
+                        wrapper_pullup_replacer("?filter_expr", "?context"),
+                    ],
+                ),
+                wrapper_pullup_replacer(
+                    udaf_expr(
+                        PATCH_MEASURE_UDAF_NAME,
+                        vec![
+                            column_expr("?measure_column"),
+                            "?new_agg_type".to_string(),
+                            "?filter_expr".to_string(),
+                        ],
+                    ),
+                    "?context",
+                ),
+                |_egraph, _subst| {
+                    // dbg!("wrapper-pull-up-aggregation-over-filtered-measure call");
+                    // dbg!(&egraph[subst[var!("?measure_column")]]);
+                    // dbg!(&egraph[subst[var!("?filter_expr")]]);
+
+                    // TODO do we need to check something here? like a SQL generator?
+                    true
+                },
+            ),
+        ]);
     }
 
     pub fn aggregate_rules_subquery(&self, rules: &mut Vec<CubeRewrite>) {
@@ -1044,4 +1135,38 @@ impl WrapperRules {
             )
         }
     }
+
+    fn transform_filtered_measure(
+        &self,
+        aggr_expr_var: &'static str,
+        literal_var: &'static str,
+        out_measure_alias_var: &'static str,
+    ) -> impl Fn(&mut CubeEGraph, &mut Subst) -> bool {
+        let aggr_expr_var = var!(aggr_expr_var);
+        let literal_var = var!(literal_var);
+        let out_measure_alias_var = var!(out_measure_alias_var);
+
+        move |egraph, subst| {
+            match &egraph[subst[literal_var]].data.constant {
+                Some(ConstantFolding::Scalar(_)) => {
+                    // Do nothing
+                }
+                _ => {
+                    return false;
+                }
+            }
+
+            // TODO share code with Self::pushdown_measure: locate cube and measure, check that ?fun matches measure, etc
+            // TODO support both changing agg fun and add filter
+
+            let Some(alias) = original_expr_name(egraph, subst[aggr_expr_var]) else {
+                return false;
+            };
+            let alias_expr_alias =
+                egraph.add(LogicalPlanLanguage::AliasExprAlias(AliasExprAlias(alias)));
+            subst.insert(out_measure_alias_var, alias_expr_alias);
+
+            true
+        }
+    }
 }

rust/cubesql/cubesql/src/compile/test/test_wrapper.rs

@@ -1518,3 +1518,88 @@ async fn wrapper_min_from_avg_measure() {
         }
     );
 }
+
+#[tokio::test]
+async fn test_ad_hoc_measure_filter() {
+    if !Rewriter::sql_push_down_enabled() {
+        return;
+    }
+    init_testing_logger();
+
+    let query_plan = convert_select_to_query_plan(
+        // language=PostgreSQL
+        r#"SELECT
+    dim_str0,
+    AVG(
+        CASE (
+            (
+                CAST(TRUNC(EXTRACT(YEAR FROM dim_date0)) AS INTEGER) = 2024
+            )
+            AND
+            (
+                CAST(TRUNC(EXTRACT(MONTH FROM dim_date0)) AS INTEGER) <= 11
+            )
+        )
+        WHEN TRUE
+        THEN avgPrice
+        ELSE NULL
+        END
+    ),
+    SUM(
+        CASE (dim_str1 = 'foo')
+        WHEN TRUE
+        THEN maxPrice
+        ELSE NULL
+        END
+    )
+FROM MultiTypeCube
+GROUP BY
+    1
+;"#
+        .to_string(),
+        DatabaseProtocol::PostgreSQL,
+    )
+    .await;
+
+    let physical_plan = query_plan.as_physical_plan().await.unwrap();
+    println!(
+        "Physical plan: {}",
+        displayable(physical_plan.as_ref()).indent()
+    );
+
+    assert_eq!(
+        query_plan
+            .as_logical_plan()
+            .find_cube_scan_wrapped_sql()
+            .request,
+        TransportLoadRequestQuery {
+            measures: Some(vec![json!({
+                "cubeName": "MultiTypeCube",
+                "alias": "avg_case_cast_tr",
+                "expr": {
+                    "type": "PatchMeasure",
+                    "sourceMeasure": "MultiTypeCube.avgPrice",
+                    "replaceAggregationType": null,
+                    "addFilters": [{
+                        "cubeParams": ["MultiTypeCube"],
+                        "sql": "(((CAST(TRUNC(EXTRACT(YEAR FROM ${MultiTypeCube.dim_date0})) AS INTEGER) = 2024) AND (CAST(TRUNC(EXTRACT(MONTH FROM ${MultiTypeCube.dim_date0})) AS INTEGER) <= 11)) = TRUE)"
+                    }],
+                },
+                "groupingSet": null,
+            }).to_string(),]),
+            dimensions: Some(vec![json!({
+                "cubeName": "MultiTypeCube",
+                "alias": "dim_str0",
+                "expr": {
+                    "type": "SqlFunction",
+                    "cubeParams": ["MultiTypeCube"],
+                    "sql": "${MultiTypeCube.dim_str0}",
+                },
+                "groupingSet": null,
+            }).to_string(),]),
+            segments: Some(vec![]),
+            order: Some(vec![]),
+            ..Default::default()
+        }
+    );
+}