From cd8dc98843ebdd3b914f33a68f398b56d5fc71f8 Mon Sep 17 00:00:00 2001
From: Kaitmazian Maksim <m.kaitmazian@picodata.io>
Date: Thu, 19 Dec 2024 14:03:12 +0300
Subject: [PATCH] feat(sbroad): cast constant evaluation
Close https://git.picodata.io/core/picodata/-/issues/1228
---
sbroad/sbroad-core/src/executor.rs | 1 +
.../sbroad-core/src/executor/tests/concat.rs | 2 +-
sbroad/sbroad-core/src/ir/explain/tests.rs | 3 +
.../src/ir/explain/tests/cast_constants.rs | 80 +++++++++++++++
.../src/ir/explain/tests/concat.rs | 4 +-
sbroad/sbroad-core/src/ir/transformation.rs | 1 +
.../src/ir/transformation/cast_constants.rs | 97 +++++++++++++++++++
7 files changed, 185 insertions(+), 3 deletions(-)
create mode 100644 sbroad/sbroad-core/src/ir/explain/tests/cast_constants.rs
create mode 100644 sbroad/sbroad-core/src/ir/transformation/cast_constants.rs
diff --git a/sbroad/sbroad-core/src/executor.rs b/sbroad/sbroad-core/src/executor.rs
index 5a572420d9..e4453abe67 100644
--- a/sbroad/sbroad-core/src/executor.rs
+++ b/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()?;
diff --git a/sbroad/sbroad-core/src/executor/tests/concat.rs b/sbroad/sbroad-core/src/executor/tests/concat.rs
index c0a6aa8221..09df269f7d 100644
--- a/sbroad/sbroad-core/src/executor/tests/concat.rs
+++ b/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")],
);
}
diff --git a/sbroad/sbroad-core/src/ir/explain/tests.rs b/sbroad/sbroad-core/src/ir/explain/tests.rs
index 0767a6c4b9..0c931a896e 100644
--- a/sbroad/sbroad-core/src/ir/explain/tests.rs
+++ b/sbroad/sbroad-core/src/ir/explain/tests.rs
@@ -585,3 +585,6 @@ mod delete;
#[cfg(test)]
mod query_explain;
+
+#[cfg(test)]
+mod cast_constants;
diff --git a/sbroad/sbroad-core/src/ir/explain/tests/cast_constants.rs b/sbroad/sbroad-core/src/ir/explain/tests/cast_constants.rs
new file mode 100644
index 0000000000..9539ad2f1a
--- /dev/null
+++ b/sbroad/sbroad-core/src/ir/explain/tests/cast_constants.rs
@@ -0,0 +1,80 @@
+use super::*;
+
+#[test]
+fn select_values_rows() {
+ explain_check(
+ "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
+"#,
+ );
+}
+
+#[test]
+fn insert_values_rows() {
+ explain_check(
+ "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
+"#,
+ );
+}
+
+#[test]
+fn select_selection() {
+ explain_check(
+ "SELECT * FROM t1 WHERE a = 1::int::unsigned",
+ r#"projection ("t1"."a"::string -> "a", "t1"."b"::integer -> "b")
+ selection ROW("t1"."a"::string) = ROW(1::unsigned)
+ scan "t1"
+execution options:
+ vdbe_max_steps = 45000
+ vtable_max_rows = 5000
+"#,
+ );
+}
+
+#[test]
+fn update_selection() {
+ explain_check(
+ "UPDATE t SET c = 2 WHERE a = 1::int::int",
+ 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)
+ scan "t"
+execution options:
+ vdbe_max_steps = 45000
+ vtable_max_rows = 5000
+"#,
+ );
+}
+
+#[test]
+fn delete_selection() {
+ explain_check(
+ r#"DELETE FROM "t1" where "a" = 3::int::decimal"#,
+ &format!(
+ "{}\n{}\n{}\n{}\n{}\n{}\n{}\n{}\n",
+ r#"delete "t1""#,
+ r#" motion [policy: local]"#,
+ r#" projection ("t1"."a"::string -> "pk_col_0", "t1"."b"::integer -> "pk_col_1")"#,
+ r#" selection ROW("t1"."a"::string) = ROW(3::decimal)"#,
+ r#" scan "t1""#,
+ r#"execution options:"#,
+ r#" vdbe_max_steps = 45000"#,
+ r#" vtable_max_rows = 5000"#,
+ ),
+ );
+}
diff --git a/sbroad/sbroad-core/src/ir/explain/tests/concat.rs b/sbroad/sbroad-core/src/ir/explain/tests/concat.rs
index 1e65e6c1ef..cc3026cbbf 100644
--- a/sbroad/sbroad-core/src/ir/explain/tests/concat.rs
+++ b/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"#,
diff --git a/sbroad/sbroad-core/src/ir/transformation.rs b/sbroad/sbroad-core/src/ir/transformation.rs
index 333c762bff..49adc95ca4 100644
--- a/sbroad/sbroad-core/src/ir/transformation.rs
+++ b/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;
diff --git a/sbroad/sbroad-core/src/ir/transformation/cast_constants.rs b/sbroad/sbroad-core/src/ir/transformation/cast_constants.rs
new file mode 100644
index 0000000000..d319ecdc38
--- /dev/null
+++ b/sbroad/sbroad-core/src/ir/transformation/cast_constants.rs
@@ -0,0 +1,97 @@
+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.map(|x| x.cast(target_type).ok()).flatten()
+ }
+ _ => 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(())
+ }
+}
--
GitLab