From 4e5dba3a64d54a8ee86bf655cbad33c473faa5cd Mon Sep 17 00:00:00 2001
From: Alex Qyoun-ae <4062971+MazterQyou@users.noreply.github.com>
Date: Wed, 31 Jan 2024 20:56:00 +0400
Subject: [PATCH] fix(cubesql): Enable constant folding for unary minus exprs

---
 rust/cubesql/cubesql/src/compile/mod.rs       | 34 +++++++++++++++++++
 .../cubesql/src/compile/rewrite/analysis.rs   |  2 +-
 rust/cubesql/cubesql/src/compile/test/mod.rs  |  2 +-
 3 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/rust/cubesql/cubesql/src/compile/mod.rs b/rust/cubesql/cubesql/src/compile/mod.rs
index 568672370d24e..956fbbb3336e5 100644
--- 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
index 80621e613bd2d..e29b8d6bb1be8 100644
--- 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
index 0c2463e3e3ee3..e78dbffee3728 100644
--- 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()),