opt(sbroad): fold constants under casts

Close #1228

sbroad/sbroad-core/src/executor.rs

@@ -57,6 +57,7 @@ impl Plan {
     /// # Errors
     /// - Failed to optimize the plan.
     pub fn optimize(&mut self) -> Result<(), SbroadError> {
+        self.cast_constants()?;
         self.replace_in_operator()?;
         self.push_down_not()?;
         self.split_columns()?;

sbroad/sbroad-core/src/executor/tests/concat.rs

@@ -5,7 +5,7 @@ use crate::ir::value::Value;
 fn concat1_test() {
     broadcast_check(
         r#"SELECT CAST('1' as string) || 'hello' FROM "t1""#,
-        r#"SELECT (CAST (? as string)) || (?) as "col_1" FROM "t1""#,
+        r#"SELECT (?) || (?) as "col_1" FROM "t1""#,
         vec![Value::from("1"), Value::from("hello")],
     );
 }

sbroad/sbroad-core/src/ir/explain/tests.rs

@@ -585,3 +585,6 @@ mod delete;
 
 #[cfg(test)]
 mod query_explain;
+
+#[cfg(test)]
+mod cast_constants;

sbroad/sbroad-core/src/ir/explain/tests/cast_constants.rs

+use crate::executor::{engine::mock::RouterRuntimeMock, Query};
+use pretty_assertions::assert_eq;
+
+fn assert_expain_matches(sql: &str, expected_explain: &str) {
+    let metadata = &RouterRuntimeMock::new();
+    let mut query = Query::new(metadata, sql, vec![]).unwrap();
+    let actual_explain = query.to_explain().unwrap();
+    assert_eq!(actual_explain.as_str(), expected_explain);
+}
+
+#[test]
+fn select_values_rows() {
+    assert_expain_matches(
+        "SELECT * FROM (VALUES (1::int, 2::decimal::unsigned, 'txt'::text::text::text))",
+        r#"projection ("COLUMN_1"::integer -> "COLUMN_1", "COLUMN_2"::unsigned -> "COLUMN_2", "COLUMN_3"::string -> "COLUMN_3")
+    scan
+        values
+            value row (data=ROW(1::integer, 2::unsigned, 'txt'::string))
+execution options:
+    vdbe_max_steps = 45000
+    vtable_max_rows = 5000
+buckets = any
+"#,
+    );
+}
+
+#[test]
+fn insert_values_rows() {
+    assert_expain_matches(
+        "INSERT INTO t1 VALUES ('txt'::text::text::text, 2::decimal::unsigned::double::integer)",
+        r#"insert "t1" on conflict: fail
+    motion [policy: segment([ref("COLUMN_1"), ref("COLUMN_2")])]
+        values
+            value row (data=ROW('txt'::string, 2::integer))
+execution options:
+    vdbe_max_steps = 45000
+    vtable_max_rows = 5000
+buckets = unknown
+"#,
+    );
+}
+
+#[test]
+fn select_selection() {
+    assert_expain_matches(
+        "SELECT * FROM t3 WHERE a = 'kek'::text::text::text",
+        r#"projection ("t3"."a"::string -> "a", "t3"."b"::integer -> "b")
+    selection ROW("t3"."a"::string) = ROW('kek'::string)
+        scan "t3"
+execution options:
+    vdbe_max_steps = 45000
+    vtable_max_rows = 5000
+buckets = [1610]
+"#,
+    );
+}
+
+#[test]
+fn update_selection() {
+    assert_expain_matches(
+        "UPDATE t SET c = 2 WHERE a = 1::int::int and b = 2::unsigned::decimal",
+        r#"update "t"
+"c" = "col_0"
+    motion [policy: local]
+        projection (2::unsigned -> "col_0", "t"."b"::unsigned -> "col_1")
+            selection ROW("t"."a"::unsigned) = ROW(1::integer) and ROW("t"."b"::unsigned) = ROW(2::decimal)
+                scan "t"
+execution options:
+    vdbe_max_steps = 45000
+    vtable_max_rows = 5000
+buckets = [550]
+"#,
+    );
+}
+
+#[test]
+fn delete_selection() {
+    assert_expain_matches(
+        r#"DELETE FROM "t2" where "e" = 3::unsigned and "f" = 2::decimal"#,
+        r#"delete "t2"
+    motion [policy: local]
+        projection ("t2"."g"::unsigned -> "pk_col_0", "t2"."h"::unsigned -> "pk_col_1")
+            selection ROW("t2"."e"::unsigned) = ROW(3::unsigned) and ROW("t2"."f"::unsigned) = ROW(2::decimal)
+                scan "t2"
+execution options:
+    vdbe_max_steps = 45000
+    vtable_max_rows = 5000
+buckets = [9374]
+"#,
+    );
+}

sbroad/sbroad-core/src/ir/explain/tests/concat.rs

@@ -6,7 +6,7 @@ fn concat1_test() {
         r#"SELECT CAST('1' as string) || 'hello' FROM "t1""#,
         &format!(
             "{}\n{}\n{}\n{}\n{}\n",
-            r#"projection (ROW('1'::string::string) || ROW('hello'::string) -> "col_1")"#,
+            r#"projection (ROW('1'::string) || ROW('hello'::string) -> "col_1")"#,
             r#"    scan "t1""#,
             r#"execution options:"#,
             r#"    vdbe_max_steps = 45000"#,
@@ -22,7 +22,7 @@ fn concat2_test() {
         &format!(
             "{}\n{}\n{}\n{}\n{}\n{}\n",
             r#"projection ("t1"."a"::string -> "a")"#,
-            r#"    selection ROW(ROW(ROW('1'::string::string) || ROW("func"(('hello'::string))::integer)) || ROW('2'::string)) = ROW(42::unsigned)"#,
+            r#"    selection ROW(ROW(ROW('1'::string) || ROW("func"(('hello'::string))::integer)) || ROW('2'::string)) = ROW(42::unsigned)"#,
             r#"        scan "t1""#,
             r#"execution options:"#,
             r#"    vdbe_max_steps = 45000"#,

sbroad/sbroad-core/src/ir/transformation.rs

@@ -3,6 +3,7 @@
 //! Contains rule-based transformations.
 
 pub mod bool_in;
+pub mod cast_constants;
 pub mod dnf;
 pub mod equality_propagation;
 pub mod merge_tuples;

sbroad/sbroad-core/src/ir/transformation/cast_constants.rs

+use crate::{
+    errors::SbroadError,
+    ir::{
+        node::{
+            expression::{Expression, MutExpression},
+            Cast, Constant, Node, NodeId, Row,
+        },
+        relation::Type,
+        tree::traversal::{LevelNode, PostOrderWithFilter},
+        value::Value,
+        Plan,
+    },
+};
+
+fn apply_cast(plan: &Plan, child_id: NodeId, target_type: Type) -> Option<Value> {
+    match plan.get_expression_node(child_id).ok()? {
+        Expression::Constant(Constant { value }) => {
+            let value = value.clone();
+            value.cast(target_type).ok()
+        }
+        Expression::Cast(Cast {
+            child: cast_child,
+            to: cast_type,
+        }) => {
+            let cast_type = cast_type.as_relation_type();
+            let value = apply_cast(plan, *cast_child, cast_type);
+            // Note: We don't throw errors if casting fails.
+            // It's possible that some type and value combinations are missing,
+            // but in such cases, we simply skip this evaluation and continue with other casts.
+            // An optimization failure should not prevent the execution of the plan.
+            value.and_then(|x| x.cast(target_type).ok())
+        }
+        _ => None,
+    }
+}
+
+impl Plan {
+    /// Evaluates cast constant expressions and replaces them with actual values in the plan.
+    ///
+    /// This function focuses on simplifying the plan by eliminating unnecessary casts in selection
+    /// expressions, enabling bucket filtering and in value rows, enabling local materialization.
+    pub fn cast_constants(&mut self) -> Result<(), SbroadError> {
+        // For simplicity, we only evaluate constants wrapped in Row,
+        // e.g., Row(Cast(Constant), Cast(Cast(Constant))).
+        // This approach includes the target cases from the function comment
+        // (selection expressions and values rows).
+        let rows_filter = |node_id| {
+            matches!(
+                self.get_node(node_id),
+                Ok(Node::Expression(Expression::Row(_)))
+            )
+        };
+        let mut subtree = PostOrderWithFilter::with_capacity(
+            |node| self.subtree_iter(node, false),
+            self.nodes.len(),
+            Box::new(rows_filter),
+        );
+
+        let top_id = self.get_top()?;
+        subtree.populate_nodes(top_id);
+        let row_ids = subtree.take_nodes();
+        drop(subtree);
+
+        let mut new_list = Vec::new();
+        for LevelNode(_, row_id) in row_ids {
+            // Clone row children list to overcome borrow checker.
+            new_list.clear();
+            if let Expression::Row(Row { list, .. }) = self.get_expression_node(row_id)? {
+                new_list.clone_from(list);
+            }
+
+            // Try to apply cast to constants, push new values in the plan and remember ids in a
+            // copy if row children list.
+            for row_child in new_list.iter_mut() {
+                if let Expression::Cast(Cast {
+                    child: cast_child,
+                    to,
+                }) = self.get_expression_node(*row_child)?
+                {
+                    let to = to.as_relation_type();
+                    if let Some(value) = apply_cast(self, *cast_child, to) {
+                        *row_child = self.add_const(value);
+                    }
+                }
+            }
+
+            // Change row children to the new ones with casts applied.
+            if let MutExpression::Row(Row { ref mut list, .. }) =
+                self.get_mut_expression_node(row_id)?
+            {
+                new_list.clone_into(list);
+            }
+        }
+
+        Ok(())
+    }
+}