feat: support prepare statement #4490

NGA-TRAN · 2022-12-02T22:02:56Z

Which issue does this PR close?

This PR generates a LogicalPlan for PREPARE SELECT ...
EXECUTE statement will be done in a follow-on PR

Which prepare statements supported in this PR:

Params in WHERE clause including both column > $1 and column in [$1, $2, ...]
Params in HAVING clause including both AGG(column) > $1 and AGG(column) in [$1, $2, ...]
Params in VALUES table in SELECT. Eg: SELECT * FROM (VALUES(1, $1), (2, $2)) AS t (num, letter);

Those are all I can think of for SELECT. Let me know if you think of other ones and I will add them in the tests. This PR does NOT support params in UPDATE and INSERT

What changes are included in this PR?

Add LogicalPlan::Prepare
Support Expr::Placeholder as it is a part of Prepare stmt
Generate LogicalPlan::Prepare for PREPARE stmt
Generate corresponding protobuf messages for LogicalPlan::Prepare
Map given data types of the PREPARE stmt to their corresponding Placeholders in the stmt

Are these changes tested?

Yes

Are there any user-facing changes?

Yes, they can use PREPARE stmt now

datafusion/sql/src/planner.rs

NGA-TRAN · 2022-12-02T22:12:12Z

@alamb, @thinkharderdev @metesynnada I will add more tests but this PR has implemented the PREPARE stmt to create a logical plan from a sql string (see the test)

Can you have a quick look to see if it is what you expect while I add more tests?

I will implement the EXECUTE stmt either in the following PR (preferred) or in this PR if you think better.

datafusion/proto/src/generated/datafusion.rs

alamb

This is looking exactly along the right track @NGA-TRAN -- thank you !

datafusion/core/src/physical_plan/planner.rs

datafusion/expr/src/expr_schema.rs

datafusion/expr/src/logical_plan/builder.rs

datafusion/expr/src/logical_plan/plan.rs

datafusion/proto/src/generated/datafusion.rs

datafusion/sql/src/planner.rs

…s expected

…ion to get data types of the params , , ...

…PlannerContext PR

NGA-TRAN · 2022-12-06T20:56:09Z

@alamb : the PR is ready for review. The supported queries are described in the description

alamb · 2022-12-06T21:15:39Z

I plan to review this carefully tomorrow morning

metesynnada · 2022-12-07T10:05:34Z

LGTM. As far as I can see, we do not store the plans, just use them immediately after we prepare. I hope I am not missing anything.

alamb

Thanks @NGA-TRAN this is looking quite good

My biggest concern about the code in this PR is the introduction of the _with_params variants. I worry that both the variants means that some codepaths will support parameters and some will not and we'll be chasing down gaps / inconsistencies.

What do you think?

I tried making running a PREPARE statement in datafusion-cli and it didn't seem to work:

cd datafusion-cli
cargo run

And I got an error

❯ PREPARE my_plan(INT) AS SELECT 1 + $1;
Internal("Placehoder $1 does not exist in the parameter list: []")

I would recommend adding some basic end to end tests (perhaps make a prepared.slt file, following the model of https://github.com/apache/arrow-datafusion/blob/master/datafusion/core/tests/sqllogictests/test_files/aggregate.slt, which can be run like cargo run -p datafusion --test sqllogictest)

datafusion/expr/src/expr.rs

datafusion/expr/src/expr_schema.rs

datafusion/expr/src/logical_plan/builder.rs

mingmwang · 2022-12-07T11:44:28Z

I will take a look at this PR carefully tomorrow.
Generally, Prepared Statement and Plan Cache support is quite complex in RDBMS.

datafusion/expr/src/expr.rs

NGA-TRAN · 2022-12-07T18:44:30Z

I am working on adding more tests which I think the last piece of this PR

NGA-TRAN · 2022-12-07T19:53:28Z

@alamb : I have addressed all your comments. Can you re-review it?

Notes: If you run queries in DataFusion-cli, you will get error Internal("Unsupported logical plan: Prepare"). I think this means we are successfully creating a LogicalPlan::Prepare (the purpose of this PR) but since we have not stored it anywhere and have not implemented physical plan for it, we get this error naturally.

Example:

cd datafusion-cli 
cargo run
DataFusion CLI v15.0.0
❯ PREPARE my_plan(INT) AS SELECT $1
;
Internal("Unsupported logical plan: Prepare")
❯

NGA-TRAN · 2022-12-07T19:55:16Z

datafusion/core/tests/sqllogictests/test_files/prepare.slt

+# PREPARE my_plan(STRING, STRING) AS SELECT * FROM (VALUES(1, $1), (2, $2)) AS t (num, letter);
+
+# And then we may want to add test_prepare_statement* here
+


Since we do not store the prepare logical plan anywhere and still throw error when we try to generate physical plan for it, this .stl tests are not that useful yet but they will be valuable after those are done

Perhaps I can add a subtask to implement it under #4539 -- then in DataFusion we normally leave a link to the ticket as a comment

All corresponding prepare tests are added here. They will be available for testing (mean many statement error will become statement ok) after we store the prepare logical plan and not throw error

datafusion/core/tests/sqllogictests/test_files/prepare.slt

alamb

Looks great to me -- thanks @NGA-TRAN

I plan to wait for @mingmwang 's review prior to merging this PR but I think it could go in now

datafusion/core/tests/sqllogictests/test_files/prepare.slt

alamb · 2022-12-07T20:29:12Z

datafusion/core/tests/sqllogictests/test_files/prepare.slt

+# PREPARE my_plan(STRING, STRING) AS SELECT * FROM (VALUES(1, $1), (2, $2)) AS t (num, letter);
+
+# And then we may want to add test_prepare_statement* here
+


Perhaps I can add a subtask to implement it under #4539 -- then in DataFusion we normally leave a link to the ticket as a comment

alamb · 2022-12-07T20:29:24Z

datafusion/expr/src/expr.rs

+    /// (e.g. `$foo` or `$1`)
+    Placeholder {
+        /// The identifier of the parameter (e.g, $1 or $foo)
+        id: String,


datafusion/expr/src/logical_plan/builder.rs

alamb · 2022-12-07T20:32:32Z

datafusion/sql/src/planner.rs

@@ -16,7 +16,7 @@
 // under the License.

 //! SQL Query Planner (produces logical plan from SQL AST)
-
+use log::debug;


👍 I have wanted some debug logging in the planner as well

datafusion/sql/src/planner.rs

alamb · 2022-12-07T20:35:28Z

datafusion/sql/src/planner.rs

+        prepare_stmt_quick_test(sql, expected_plan, expected_dt);
+    }
+
+    #[test]


👍 very nice tests

alamb · 2022-12-07T20:36:20Z

datafusion/sql/src/planner.rs

+
+        let expected_plan = "Prepare: \"my_plan\" [Int32, Float64] \
+        \n  Projection: Int64(1) + $1 + $2\n    EmptyRelation";
+        let expected_dt = "[Int32, Float64]";


… message per review comment

mingmwang · 2022-12-08T09:32:20Z

datafusion/expr/src/logical_plan/plan.rs

+    pub data_types: Vec<DataType>,
+    /// The logical plan of the statements
+    pub input: Arc<LogicalPlan>,
+}
+


I assume the data types Vec size is the same with the place holders in the input plan, but is there any check for this?

The data types of the placeholders are from the data types of the Vec here so they match. We do check if the Vec contains enough params, too. However, there are flexibility:

The length of the Vec can be longer than the number of the placeholders which will be fine and we have tests for this.

The data type of the Vec can be anything and we will use them for the placeholders, we do not check if the data types are compatible with the variables in the expression because: (i) we allow data type casting, and (2) when we reach the placeholders, we no longer have the context which variable/column/expression the place holder is used for. We do not want to add more context to backtrack which will cost compile time as well as complicated implementation.

However, I am working on #4550 that convert Prepare Logical Plan to a logical plan with all placeholders replaced with actual values. There, I will throw error if the data types provided do not work. We follow the same behavior of Postgres

mingmwang · 2022-12-08T09:38:38Z

datafusion/expr/src/utils.rs

+            name: name.clone(),
+            data_types: data_types.clone(),
+            input: Arc::new(inputs[0].clone()),
+        })),


Is it possible and allowed here that the method passed in a totally different input plan ?

With the code we implement, the answer is no unless there are bugs

mingmwang · 2022-12-08T09:56:21Z

datafusion/sql/src/planner.rs

+                        SQLExpr::Value(Value::Placeholder(param)) => {
+                            Self::create_placeholder_expr(param, param_data_types)
+                        }


For SQLValues with Placeholder, should we expect that the param_data_types is the same data type ?

I do not get the question but I think my answer for this https://github.com/apache/arrow-datafusion/pull/4490/files#r1043122807 probably answer this question, too

datafusion/sql/src/planner.rs

mingmwang · 2022-12-08T10:02:03Z

@NGA-TRAN
Overall, this PR LGTM, thanks!

mingmwang · 2022-12-08T10:20:45Z

BTW, do we support this case ?

Plain SQL:

SELECT * FROM tab WHERE col_a is null;

Prepared Stmt:

REPARE my_plan_1(INT) AS 
   SELECT * FROM tab WHERE col_a is $1;

NGA-TRAN · 2022-12-08T15:13:00Z

@mingmwang
I have tried and will add this test in my next commit. We will get error. Thanks

REPARE my_plan_1(INT) AS 
   SELECT * FROM tab WHERE col_a is $1;

value: SQL(ParserError("Expected [NOT] NULL or TRUE|FALSE or [NOT] DISTINCT FROM after IS, found: $1"))

NGA-TRAN

@alamb The PR is ready to master merge. I will resolve conflicts shortly but more conflicts may come if we do not merge soon 😄

datafusion/sql/src/planner.rs

alamb · 2022-12-08T15:52:38Z

I plan to merge this as soon as the CI checks complete

* feat: support prepare statement * fix: typo * chore: address preliminary review comments * fix: put Placeholder last to have the expression comparison to work as expected * test: add more tests and starting to pass param_data_types to expression to get data types of the params , , ... * test: one more test and a bit of refactor while waiting for the CTEs/PlannerContext PR * feat: use prepare stmt's param data types in the planner context * chore: cleanup * refactor: address review comments * chore: cleanup * test: more prepare statement tests * chore: cleanup * chore: fix typos and add tests into the sqllogicaltests * docs: add docstring * chore: update test panic message due to recent change to have clearer message per review comment * chore: add a test and a doc string per review comments * fix: output of a test after master merge

feat: support prepare statement

f7408fc

NGA-TRAN marked this pull request as draft December 2, 2022 22:03

github-actions bot added core Core DataFusion crate logical-expr Logical plan and expressions optimizer Optimizer rules sql SQL Planner labels Dec 2, 2022

fix: typo

ff0dc10

NGA-TRAN commented Dec 2, 2022

View reviewed changes

datafusion/sql/src/planner.rs Outdated Show resolved Hide resolved

NGA-TRAN commented Dec 2, 2022

View reviewed changes

datafusion/proto/src/generated/datafusion.rs Outdated Show resolved Hide resolved

alamb reviewed Dec 3, 2022

View reviewed changes

NGA-TRAN added 3 commits December 5, 2022 09:37

chore: address preliminary review comments

5a50288

fix: put Placeholder last to have the expression comparison to work a…

e5ec877

…s expected

test: add more tests and starting to pass param_data_types to express…

678a30b

…ion to get data types of the params , , ...

NGA-TRAN mentioned this pull request Dec 5, 2022

refactor: make Ctes a struct to also store data types provided by prepare stmt #4520

Merged

NGA-TRAN added 4 commits December 6, 2022 09:52

test: one more test and a bit of refactor while waiting for the CTEs/…

e04da97

…PlannerContext PR

chore: merge master to branch

860f85d

feat: use prepare stmt's param data types in the planner context

ba8ad0f

chore: cleanup

67cbfdd

NGA-TRAN marked this pull request as ready for review December 6, 2022 20:48

alamb mentioned this pull request Dec 7, 2022

[Epic] Prepared Statement Support #4539

Closed

5 tasks

alamb reviewed Dec 7, 2022

View reviewed changes

datafusion/expr/src/expr.rs Outdated Show resolved Hide resolved

NGA-TRAN added 2 commits December 7, 2022 11:34

refactor: address review comments

001251c

chore: cleanup

5d78097

test: more prepare statement tests

1344395

github-actions bot added the sqllogictest SQL Logic Tests (.slt) label Dec 7, 2022

chore: cleanup

87787da

NGA-TRAN commented Dec 7, 2022

View reviewed changes

chore: merge master to branch

199b657

NGA-TRAN commented Dec 7, 2022

View reviewed changes

datafusion/core/tests/sqllogictests/test_files/prepare.slt Outdated Show resolved Hide resolved

alamb approved these changes Dec 7, 2022

View reviewed changes

NGA-TRAN added 3 commits December 7, 2022 16:13

chore: fix typos and add tests into the sqllogicaltests

cdb328c

docs: add docstring

58a3117

chore: update test panic message due to recent change to have clearer…

30705d9

… message per review comment

mingmwang reviewed Dec 8, 2022

View reviewed changes

datafusion/sql/src/planner.rs Show resolved Hide resolved

chore: add a test and a doc string per review comments

87a3523

NGA-TRAN commented Dec 8, 2022

View reviewed changes

datafusion/sql/src/planner.rs Show resolved Hide resolved

chore: merge master to branch

91e104a

fix: output of a test after master merge

dd6c3e0

alamb merged commit 4538912 into apache:master Dec 8, 2022

jonahgao mentioned this pull request May 16, 2023

Fix nullability calculation for boolean expressions #6365

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: support prepare statement #4490

feat: support prepare statement #4490

NGA-TRAN commented Dec 2, 2022 •

edited by alamb

Loading

NGA-TRAN commented Dec 2, 2022

alamb left a comment

NGA-TRAN commented Dec 6, 2022

alamb commented Dec 6, 2022

metesynnada commented Dec 7, 2022

alamb left a comment

mingmwang commented Dec 7, 2022

NGA-TRAN commented Dec 7, 2022

NGA-TRAN commented Dec 7, 2022 •

edited

Loading

NGA-TRAN Dec 7, 2022

alamb Dec 7, 2022

NGA-TRAN Dec 7, 2022

alamb left a comment

alamb Dec 7, 2022

alamb Dec 7, 2022

alamb Dec 7, 2022

alamb Dec 7, 2022

alamb Dec 7, 2022

mingmwang Dec 8, 2022

NGA-TRAN Dec 8, 2022

mingmwang Dec 8, 2022

NGA-TRAN Dec 8, 2022

mingmwang Dec 8, 2022

NGA-TRAN Dec 8, 2022

mingmwang commented Dec 8, 2022

mingmwang commented Dec 8, 2022

NGA-TRAN commented Dec 8, 2022

NGA-TRAN left a comment

alamb commented Dec 8, 2022

		# PREPARE my_plan(STRING, STRING) AS SELECT * FROM (VALUES(1, $1), (2, $2)) AS t (num, letter);

		# And then we may want to add test_prepare_statement* here

feat: support prepare statement #4490

feat: support prepare statement #4490

Conversation

NGA-TRAN commented Dec 2, 2022 • edited by alamb Loading

Which issue does this PR close?

Which prepare statements supported in this PR:

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

NGA-TRAN commented Dec 2, 2022

alamb left a comment

Choose a reason for hiding this comment

NGA-TRAN commented Dec 6, 2022

alamb commented Dec 6, 2022

metesynnada commented Dec 7, 2022

alamb left a comment

Choose a reason for hiding this comment

mingmwang commented Dec 7, 2022

NGA-TRAN commented Dec 7, 2022

NGA-TRAN commented Dec 7, 2022 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mingmwang commented Dec 8, 2022

mingmwang commented Dec 8, 2022

NGA-TRAN commented Dec 8, 2022

NGA-TRAN left a comment

Choose a reason for hiding this comment

alamb commented Dec 8, 2022

NGA-TRAN commented Dec 2, 2022 •

edited by alamb

Loading

NGA-TRAN commented Dec 7, 2022 •

edited

Loading