fix(cubesql): Enable constant folding for unary minus exprs

cube-js · Jan 31, 2024 · 4e5dba3 · 4e5dba3
1 parent b84a345
commit 4e5dba3
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 2 deletions.
diff --git a/rust/cubesql/cubesql/src/compile/mod.rs b/rust/cubesql/cubesql/src/compile/mod.rs
@@ -21136,4 +21136,38 @@ limit
         assert!(sql.contains("order_date"));
         assert!(sql.contains("EXTRACT(DAY FROM"))
     }
+
+    #[tokio::test]
+    async fn test_unary_minus_constant_folding() {
+        if !Rewriter::sql_push_down_enabled() {
+            return;
+        }
+        init_logger();
+
+        let query_plan = convert_select_to_query_plan(
+            r#"
+            SELECT order_date + (-EXTRACT(YEAR FROM order_date) * INTERVAL '1 day') AS t
+            FROM KibanaSampleDataEcommerce
+            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()
+        );
+
+        let logical_plan = query_plan.as_logical_plan();
+        let sql = logical_plan
+            .find_cube_scan_wrapper()
+            .wrapped_sql
+            .unwrap()
+            .sql;
+        assert!(sql.contains("-(EXTRACT(YEAR FROM"));
+        assert!(sql.contains("* INTERVAL '1 DAY'"));
+    }
 }
diff --git a/rust/cubesql/cubesql/src/compile/rewrite/analysis.rs b/rust/cubesql/cubesql/src/compile/rewrite/analysis.rs
@@ -999,7 +999,7 @@ impl LogicalPlanAnalysis {
                 // TODO In case multiple node variant exists ConstantFolding::List will choose one which contains actual constants.
                 Some(ConstantFolding::List(list))
             }
-            LogicalPlanLanguage::AnyExpr(_) => {
+            LogicalPlanLanguage::AnyExpr(_) | LogicalPlanLanguage::NegativeExpr(_) => {
                 let expr = node_to_expr(
                     enode,
                     &egraph.analysis.cube_context,

diff --git a/rust/cubesql/cubesql/src/compile/test/mod.rs b/rust/cubesql/cubesql/src/compile/test/mod.rs
@@ -301,7 +301,7 @@ pub fn get_test_tenant_ctx_customized(custom_templates: Vec<(String, String)>) -
                         "expressions/column_aliased".to_string(),
                         "{{expr}} {{quoted_alias}}".to_string(),
                     ),
-                    ("expressions/binary".to_string(), "{{ left }} {{ op }} {{ right }}".to_string()),
+                    ("expressions/binary".to_string(), "({{ left }} {{ op }} {{ right }})".to_string()),
                     ("expressions/is_null".to_string(), "{{ expr }} IS {% if negate %}NOT {% endif %}NULL".to_string()),
                     ("expressions/case".to_string(), "CASE{% if expr %}{{ expr }} {% endif %}{% for when, then in when_then %} WHEN {{ when }} THEN {{ then }}{% endfor %}{% if else_expr %} ELSE {{ else_expr }}{% endif %} END".to_string()),
                     ("expressions/sort".to_string(), "{{ expr }} {% if asc %}ASC{% else %}DESC{% endif %}{% if nulls_first %} NULLS FIRST {% endif %}".to_string()),