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

Author: mcheshkov

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

@@ -1780,7 +1780,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,11 +2,13 @@ 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,
+            agg_fun_expr, aggregate, alias_expr,
+            analysis::ConstantFolding,
+            binary_expr, case_expr, column_expr, cube_scan_wrapper, grouping_set_expr,
+            literal_null, original_expr_name, rewrite,
             rewriter::{CubeEGraph, CubeRewrite},
             rules::{members::MemberRules, wrapper::WrapperRules},
-            subquery, transforming_chain_rewrite, transforming_rewrite, wrapped_select,
+            subquery, transforming_chain_rewrite, transforming_rewrite, udaf_expr, wrapped_select,
             wrapped_select_aggr_expr_empty_tail, wrapped_select_filter_expr_empty_tail,
             wrapped_select_group_expr_empty_tail, wrapped_select_having_expr_empty_tail,
             wrapped_select_joins_empty_tail, wrapped_select_order_expr_empty_tail,
@@ -319,6 +321,91 @@ impl WrapperRules {
                 "GroupingSetExprMembers",
             );
         }
+
+        // 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"),
+                            "?replace_agg_type".to_string(),
+                            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",
+                    "?measure_column",
+                    "?fun",
+                    "?cube_members",
+                    "?replace_agg_type",
+                    "?out_measure_alias",
+                ),
+            ),
+            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",
+                ),
+            ),
+        ]);
     }
 
     pub fn aggregate_rules_subquery(&self, rules: &mut Vec<CubeRewrite>) {
@@ -1058,4 +1145,128 @@ impl WrapperRules {
             )
         }
     }
+
+    fn transform_filtered_measure(
+        &self,
+        aggr_expr_var: &'static str,
+        literal_var: &'static str,
+        column_var: &'static str,
+        fun_name_var: &'static str,
+        cube_members_var: &'static str,
+        replace_agg_type_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 column_var = var!(column_var);
+        let fun_name_var = var!(fun_name_var);
+        let cube_members_var = var!(cube_members_var);
+        let replace_agg_type_var = var!(replace_agg_type_var);
+        let out_measure_alias_var = var!(out_measure_alias_var);
+
+        let meta = self.meta_context.clone();
+        let disable_strict_agg_type_match = self.config_obj.disable_strict_agg_type_match();
+
+        move |egraph, subst| {
+            match &egraph[subst[literal_var]].data.constant {
+                Some(ConstantFolding::Scalar(_)) => {
+                    // Do nothing
+                }
+                _ => {
+                    return false;
+                }
+            }
+
+            let Some(alias) = original_expr_name(egraph, subst[aggr_expr_var]) else {
+                return false;
+            };
+
+            for fun in var_iter!(egraph[subst[fun_name_var]], AggregateFunctionExprFun)
+                .cloned()
+                .collect::<Vec<_>>()
+            {
+                let call_agg_type = MemberRules::get_agg_type(Some(&fun), false);
+
+                let column_iter = var_iter!(egraph[subst[column_var]], ColumnExprColumn)
+                    .cloned()
+                    .collect::<Vec<_>>();
+
+                if let Some(member_names_to_expr) = &mut egraph
+                    .index_mut(subst[cube_members_var])
+                    .data
+                    .member_name_to_expr
+                {
+                    for column in column_iter {
+                        if let Some((&(Some(ref member), _, _), _)) =
+                            LogicalPlanData::do_find_member_by_alias(
+                                member_names_to_expr,
+                                &column.name,
+                            )
+                        {
+                            if let Some(measure) = meta.find_measure_with_name(member) {
+                                if !measure.allow_add_filter(call_agg_type.as_deref()) {
+                                    continue;
+                                }
+
+                                let Some(call_agg_type) = &call_agg_type else {
+                                    // call_agg_type is None, rewrite as is
+                                    Self::insert_patch_measure(
+                                        egraph,
+                                        subst,
+                                        column,
+                                        None,
+                                        alias,
+                                        None,
+                                        Some(replace_agg_type_var),
+                                        out_measure_alias_var,
+                                    );
+
+                                    return true;
+                                };
+
+                                if measure
+                                    .is_same_agg_type(call_agg_type, disable_strict_agg_type_match)
+                                {
+                                    Self::insert_patch_measure(
+                                        egraph,
+                                        subst,
+                                        column,
+                                        None,
+                                        alias,
+                                        None,
+                                        Some(replace_agg_type_var),
+                                        out_measure_alias_var,
+                                    );
+
+                                    return true;
+                                }
+
+                                if measure.allow_replace_agg_type(
+                                    call_agg_type,
+                                    disable_strict_agg_type_match,
+                                ) {
+                                    Self::insert_patch_measure(
+                                        egraph,
+                                        subst,
+                                        column,
+                                        Some(call_agg_type.clone()),
+                                        alias,
+                                        None,
+                                        Some(replace_agg_type_var),
+                                        out_measure_alias_var,
+                                    );
+
+                                    return true;
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+
+            false
+
+            // TODO share code with Self::pushdown_measure: locate cube and measure, check that ?fun matches measure, etc
+        }
+    }
 }

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

@@ -1801,3 +1801,104 @@ 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(),
+                json!({
+                    "cubeName": "MultiTypeCube",
+                    "alias": "sum_case_multity",
+                    "expr": {
+                        "type": "PatchMeasure",
+                        "sourceMeasure": "MultiTypeCube.maxPrice",
+                        "replaceAggregationType": "sum",
+                        "addFilters": [{
+                            "cubeParams": ["MultiTypeCube"],
+                            "sql": "((${MultiTypeCube.dim_str1} = $0$) = 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()
+        }
+    );
+}

rust/cubesql/cubesql/src/transport/ext.rs

@@ -12,6 +12,8 @@ pub trait V1CubeMetaMeasureExt {
 
     fn allow_replace_agg_type(&self, query_agg_type: &str, disable_strict_match: bool) -> bool;
 
+    fn allow_add_filter(&self, query_agg_type: Option<&str>) -> bool;
+
     fn get_sql_type(&self) -> ColumnType;
 }
 
@@ -72,6 +74,30 @@ impl V1CubeMetaMeasureExt for CubeMetaMeasure {
         }
     }
 
+    // This should be aligned with BaseMeasure.preparePatchedMeasure
+    // See packages/cubejs-schema-compiler/src/adapter/BaseMeasure.ts:16
+    fn allow_add_filter(&self, query_agg_type: Option<&str>) -> bool {
+        let Some(agg_type) = &self.agg_type else {
+            return false;
+        };
+
+        let agg_type = match query_agg_type {
+            Some(query_agg_type) => query_agg_type,
+            None => agg_type,
+        };
+
+        match agg_type {
+            "sum"
+            | "avg"
+            | "min"
+            | "max"
+            | "count"
+            | "count_distinct"
+            | "count_distinct_approx" => true,
+            _ => false,
+        }
+    }
+
     fn get_sql_type(&self) -> ColumnType {
         let from_type = match &self.r#type.to_lowercase().as_str() {
             &"number" => ColumnType::Double,