From 15657466a3b074671b45187186165fe41b3092ae Mon Sep 17 00:00:00 2001
From: Alex Qyoun-ae <4062971+MazterQyou@users.noreply.github.com>
Date: Fri, 26 Jan 2024 16:38:26 +0400
Subject: [PATCH] chore(cubesql): Realias `Int32` and `Boolean` on constant
 folding

---
 .../cubesql/src/compile/rewrite/analysis.rs   |  2 +
 .../src/compile/rewrite/rules/filters.rs      | 42 ++++++++++++++-----
 ...ile__tests__test_extract_string_field.snap | 10 ++---
 3 files changed, 39 insertions(+), 15 deletions(-)

diff --git a/rust/cubesql/cubesql/src/compile/rewrite/analysis.rs b/rust/cubesql/cubesql/src/compile/rewrite/analysis.rs
index 80621e613bd2d..3206dfe90ec4b 100644
--- a/rust/cubesql/cubesql/src/compile/rewrite/analysis.rs
+++ b/rust/cubesql/cubesql/src/compile/rewrite/analysis.rs
@@ -1287,10 +1287,12 @@ impl Analysis<LogicalPlanLanguage> for LogicalPlanAnalysis {
                     ScalarValue::Date32(_)
                         | ScalarValue::Date64(_)
                         | ScalarValue::Int64(_)
+                        | ScalarValue::Int32(_)
                         | ScalarValue::Float64(_)
                         | ScalarValue::IntervalYearMonth(_)
                         | ScalarValue::IntervalDayTime(_)
                         | ScalarValue::Utf8(_)
+                        | ScalarValue::Boolean(_)
                 ) {
                 egraph[id]
                     .data
diff --git a/rust/cubesql/cubesql/src/compile/rewrite/rules/filters.rs b/rust/cubesql/cubesql/src/compile/rewrite/rules/filters.rs
index 68bb0f411091f..16d45dc08946e 100644
--- a/rust/cubesql/cubesql/src/compile/rewrite/rules/filters.rs
+++ b/rust/cubesql/cubesql/src/compile/rewrite/rules/filters.rs
@@ -101,7 +101,7 @@ impl RewriteRules for FilterRules {
             transforming_rewrite(
                 "push-down-limit-filter",
                 filter(
-                    literal_expr("?literal"),
+                    "?literal_false",
                     cube_scan(
                         "?source_table_name",
                         "?members",
@@ -132,7 +132,7 @@ impl RewriteRules for FilterRules {
                     ),
                 ),
                 self.push_down_limit_filter(
-                    "?literal",
+                    "?literal_false",
                     "?new_limit",
                     "?new_limit_skip",
                     "?new_limit_fetch",
@@ -141,18 +141,19 @@ impl RewriteRules for FilterRules {
             // Transform Filter: Boolean(true)
             // It's safe to push down filter under projection, next filter-truncate-true will truncate it
             // TODO: Find a better solution how to drop filter node at all once
-            rewrite(
+            transforming_rewrite(
                 "push-down-filter-projection",
                 filter(
-                    literal_bool(true),
+                    "?literal_true",
                     projection("?expr", "?input", "?alias", "?projection_split"),
                 ),
                 projection(
                     "?expr",
-                    filter(literal_bool(true), "?input"),
+                    filter("?literal_true", "?input"),
                     "?alias",
                     "?projection_split",
                 ),
+                self.transform_literal_true("?literal_true"),
             ),
             rewrite(
                 "swap-limit-filter",
@@ -257,26 +258,28 @@ impl RewriteRules for FilterRules {
                 ),
             ),
             // Transform Filter: Boolean(True) same as TRUE = TRUE, which is useless
-            rewrite(
+            transforming_rewrite(
                 "filter-truncate-true",
                 filter_replacer(
-                    literal_bool(true),
+                    "?literal_true",
                     "?alias_to_cube",
                     "?members",
                     "?filter_aliases",
                 ),
                 cube_scan_filters_empty_tail(),
+                self.transform_literal_true("?literal_true"),
             ),
             // We use this rule to transform: (?expr IN (?list..)) = TRUE and ((?expr IN (?list..)) = TRUE) = TRUE
-            rewrite(
+            transforming_rewrite(
                 "filter-truncate-in-list-true",
                 filter_replacer(
-                    binary_expr("?expr", "=", literal_bool(true)),
+                    binary_expr("?expr", "=", "?literal_true"),
                     "?alias_to_cube",
                     "?members",
                     "?filter_aliases",
                 ),
                 filter_replacer("?expr", "?alias_to_cube", "?members", "?filter_aliases"),
+                self.transform_literal_true("?literal_true"),
             ),
             transforming_rewrite(
                 "filter-replacer",
@@ -2686,7 +2689,9 @@ impl FilterRules {
         let new_limit_skip_var = var!(new_limit_skip_var);
         let new_limit_fetch_var = var!(new_limit_fetch_var);
         move |egraph, subst| {
-            for literal_value in var_iter!(egraph[subst[literal_var]], LiteralExprValue) {
+            if let Some(ConstantFolding::Scalar(literal_value)) =
+                &egraph[subst[literal_var]].data.constant
+            {
                 if let ScalarValue::Boolean(Some(false)) = literal_value {
                     subst.insert(
                         new_limit_var,
@@ -2707,6 +2712,23 @@ impl FilterRules {
         }
     }
 
+    fn transform_literal_true(
+        &self,
+        literal_var: &'static str,
+    ) -> impl Fn(&mut EGraph<LogicalPlanLanguage, LogicalPlanAnalysis>, &mut Subst) -> bool {
+        let literal_var = var!(literal_var);
+        move |egraph, subst| {
+            if let Some(ConstantFolding::Scalar(literal)) =
+                &egraph[subst[literal_var]].data.constant
+            {
+                if let ScalarValue::Boolean(Some(true)) = literal {
+                    return true;
+                }
+            }
+            false
+        }
+    }
+
     fn push_down_limit_projection(
         &self,
         input_var: &'static str,
diff --git a/rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__test_extract_string_field.snap b/rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__test_extract_string_field.snap
index 4ba799074a945..1f0526cb3be1a 100644
--- a/rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__test_extract_string_field.snap
+++ b/rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__test_extract_string_field.snap
@@ -2,8 +2,8 @@
 source: cubesql/src/compile/mod.rs
 expression: "execute_query(\"SELECT EXTRACT('YEAR' FROM CAST ('2020-12-25 22:48:48.000' AS timestamptz))\".to_string(),\n            DatabaseProtocol::PostgreSQL).await?"
 ---
-+-------------+
-| Int32(2020) |
-+-------------+
-| 2020        |
-+-------------+
++---------------------------------------------------------------------------------------------+
+| datepart(Utf8("YEAR"),CAST(Utf8("2020-12-25 22:48:48.000") AS Timestamp(Nanosecond, None))) |
++---------------------------------------------------------------------------------------------+
+| 2020                                                                                        |
++---------------------------------------------------------------------------------------------+