alamb commented on code in PR #13242:
URL: https://github.com/apache/datafusion/pull/13242#discussion_r1829802769
##########
datafusion/core/src/execution/context/mod.rs:
##########
@@ -1088,6 +1113,49 @@ impl SessionContext {
}
}
+ fn execute_prepared(&self, execute: Execute) -> Result<DataFrame> {
+ let Execute {
+ name, parameters, ..
+ } = execute;
+ let prepared = self.state.read().get_prepared(&name).ok_or_else(|| {
+ exec_datafusion_err!("Prepared statement '{}' does not exist",
name)
+ })?;
+
+ // Only allow literals as parameters for now.
+ let mut params: Vec<ScalarValue> = parameters
+ .into_iter()
+ .map(|e| match e {
+ Expr::Literal(scalar) => Ok(scalar),
+ _ => not_impl_err!("Unsupported parameter type: {}", e),
+ })
+ .collect::<Result<_>>()?;
+
+ // If the prepared statement provides data types, cast the params to
those types.
+ if !prepared.data_types.is_empty() {
+ if params.len() != prepared.data_types.len() {
+ return exec_err!(
+ "Prepared statement '{}' expects {} parameters, but {}
provided",
+ name,
+ prepared.data_types.len(),
+ params.len()
+ );
+ }
+ params = params
+ .into_iter()
+ .zip(prepared.data_types.iter())
+ .map(|(e, dt)| e.cast_to(dt))
+ .collect::<Result<_>>()?;
+ }
+
+ let params = ParamValues::List(params);
Review Comment:
nice
##########
datafusion/core/src/execution/context/mod.rs:
##########
@@ -687,7 +688,31 @@ impl SessionContext {
LogicalPlan::Statement(Statement::SetVariable(stmt)) => {
self.set_variable(stmt).await
}
-
+ LogicalPlan::Prepare(Prepare {
Review Comment:
I am worried that `LogicalPlan::Prepare` seems to be runnable even if we run
`sql_with_options` and disable statements.
https://github.com/apache/datafusion/blob/aad77446f54c89cdaaa4f61909a3e0230b8cc326/datafusion/core/src/execution/context/mod.rs#L605-L614
I realize that this PR does not change if Prepare is a statement or not, but
it is the first that actually makes `Prepare` do something
Could you please add a test like this (but for `PREPARE` statements) and
make sure the statement can't be executed if statements are disabled?
https://github.com/apache/datafusion/blob/6692382f22f04542534bba0183cf0682fd932da1/datafusion/core/tests/sql/sql_api.rs#L99-L114
Perhaps (as a follow on PR) we can make `LogicalPlan::Prepare` a statement
https://docs.rs/datafusion/latest/datafusion/logical_expr/enum.Statement.html
##########
datafusion/sqllogictest/test_files/prepare.slt:
##########
@@ -30,56 +30,175 @@ select * from person;
# Error due to syntax and semantic violation
# Syntax error: no name specified after the keyword prepare
-statement error
+statement error DataFusion error: SQL error: ParserError
Review Comment:
❤️
##########
datafusion/sqllogictest/test_files/prepare.slt:
##########
@@ -30,56 +30,175 @@ select * from person;
# Error due to syntax and semantic violation
# Syntax error: no name specified after the keyword prepare
-statement error
+statement error DataFusion error: SQL error: ParserError
PREPARE AS SELECT id, age FROM person WHERE age = $foo;
# param following a non-number, $foo, not supported
-statement error
+statement error Invalid placeholder, not a number: \$foo
PREPARE my_plan(INT) AS SELECT id, age FROM person WHERE age = $foo;
# not specify table hence cannot specify columns
-statement error
+statement error Schema error: No field named id
PREPARE my_plan(INT) AS SELECT id + $1;
# not specify data types for all params
-statement error
+statement error Prepare specifies 1 data types but query has 2 parameters
PREPARE my_plan(INT) AS SELECT 1 + $1 + $2;
+# sepecify too many data types for params
+statement error Prepare specifies 2 data types but query has 1 parameters
+PREPARE my_plan(INT, INT) AS SELECT 1 + $1;
+
# cannot use IS param
-statement error
+statement error SQL error: ParserError
PREPARE my_plan(INT) AS SELECT id, age FROM person WHERE age is $1;
+# TODO: allow prepare without specifying data types
+statement error Placeholder type could not be resolved
+PREPARE my_plan AS SELECT $1;
+
# #######################
-# TODO: all the errors below should work ok after we store the prepare logical
plan somewhere
-statement error
+# Test prepare and execute statements
+
+# execute a non-existing plan
+statement error Prepared statement \'my_plan\' does not exist
+EXECUTE my_plan('Foo', 'Bar');
+
+statement ok
PREPARE my_plan(STRING, STRING) AS SELECT * FROM (VALUES(1, $1), (2, $2)) AS t
(num, letter);
-statement error
+query IT
+EXECUTE my_plan('Foo', 'Bar');
+----
+1 Foo
+2 Bar
+
+# duplicate prepare statement
+statement error Prepared statement \'my_plan\' already exists
+PREPARE my_plan(STRING, STRING) AS SELECT * FROM (VALUES(1, $1), (2, $2)) AS t
(num, letter);
+
+statement error Prepare specifies 1 data types but query has 0 parameters
PREPARE my_plan(INT) AS SELECT id, age FROM person WHERE age = 10;
-statement error
-PREPARE my_plan AS SELECT id, age FROM person WHERE age = 10;
+# prepare statement has no params
+statement ok
+PREPARE my_plan2 AS SELECT id, age FROM person WHERE age = 20;
+
+query II
+EXECUTE my_plan2;
+----
+1 20
+
+statement ok
+PREPARE my_plan3(INT) AS SELECT $1;
-statement error
-PREPARE my_plan(INT) AS SELECT $1;
+query I
+EXECUTE my_plan3(10);
+----
+10
-statement error
-PREPARE my_plan(INT) AS SELECT 1 + $1;
+statement ok
+PREPARE my_plan4(INT) AS SELECT 1 + $1;
-statement error
-PREPARE my_plan(INT, DOUBLE) AS SELECT 1 + $1 + $2;
+query I
+EXECUTE my_plan4(10);
+----
+11
-statement error
-PREPARE my_plan(INT) AS SELECT id, age FROM person WHERE age = $1;
+statement ok
+PREPARE my_plan5(INT, DOUBLE) AS SELECT 1 + $1 + $2;
-statement error
-PREPARE my_plan(INT, STRING, DOUBLE, INT, DOUBLE, STRING) AS SELECT id, age,
$6 FROM person WHERE age IN ($1, $4) AND salary > $3 and salary < $5 OR
first_name < $2";
+query R
+EXECUTE my_plan5(10, 20.5);
+----
+31.5
-statement error
-PREPARE my_plan(INT, DOUBLE, DOUBLE, DOUBLE) AS SELECT id, SUM(age) FROM
person WHERE salary > $2 GROUP BY id HAVING sum(age) < $1 AND SUM(age) > 10 OR
SUM(age) in ($3, $4);
+statement ok
+PREPARE my_plan6(INT) AS SELECT id, age FROM person WHERE age = $1;
+
+query II
+EXECUTE my_plan6(20);
+----
+1 20
+
+# EXECUTE param is a different type but compatible
+query II
+EXECUTE my_plan6('20');
+----
+1 20
+
+query II
+EXECUTE my_plan6(20.0);
+----
+1 20
+
+# invalid execute param
+statement error Cast error: Cannot cast string 'foo' to value of Int32 type
+EXECUTE my_plan6('foo');
+
+# TODO: support non-literal expressions
+statement error Unsupported parameter type
+EXECUTE my_plan6(10 + 20);
+
+statement ok
+PREPARE my_plan7(INT, STRING, DOUBLE, INT, DOUBLE, STRING)
+ AS
+SELECT id, age, $6 FROM person WHERE age IN ($1, $4) AND salary > $3 and
salary < $5 OR first_name < $2;
+
+query IIT
+EXECUTE my_plan7(10, 'jane', 99999.45, 20, 200000.45, 'foo');
+----
+1 20 foo
+
+statement ok
+PREPARE my_plan8(INT, DOUBLE, DOUBLE, DOUBLE)
+ AS
+SELECT id, SUM(age) FROM person WHERE salary > $2 GROUP BY id
+ HAVING sum(age) < $1 AND SUM(age) > 10 OR SUM(age) in ($3, $4);
+
+query II
+EXECUTE my_plan8(100000, 99999.45, 100000.45, 200000.45);
+----
+1 20
+
+statement ok
+PREPARE my_plan9(STRING, STRING) AS SELECT * FROM (VALUES(1, $1), (2, $2)) AS
t (num, letter);
+
+query IT
+EXECUTE my_plan9('Foo', 'Bar');
+----
+1 Foo
+2 Bar
+
+
+# Test issue: https://github.com/apache/datafusion/issues/12294
Review Comment:
😍
##########
datafusion/core/tests/sql/select.rs:
##########
@@ -106,9 +105,9 @@ async fn test_prepare_statement() -> Result<()> {
let ctx = create_ctx_with_partition(&tmp_dir, partition_count).await?;
// sql to statement then to prepare logical plan with parameters
- // c1 defined as UINT32, c2 defined as UInt64 but the params are Int32 and
Float64
- let dataframe =
- ctx.sql("PREPARE my_plan(INT, DOUBLE) AS SELECT c1, c2 FROM test WHERE
c1 > $2 AND c1 < $1").await?;
Review Comment:
I agree -- PREPARE is a statement in my mind (and has no results)
##########
datafusion/core/src/execution/session_state.rs:
##########
@@ -906,6 +910,29 @@ impl SessionState {
let udtf = self.table_functions.remove(name);
Ok(udtf.map(|x| x.function().clone()))
}
+
+ /// Store the logical plan and the parameter types of a prepared statement.
+ pub(crate) fn store_prepared(
+ &mut self,
+ name: String,
+ data_types: Vec<DataType>,
+ plan: Arc<LogicalPlan>,
+ ) -> datafusion_common::Result<()> {
+ match self.prepared_plans.entry(name) {
+ Entry::Vacant(e) => {
+ e.insert(Arc::new(PreparedPlan { data_types, plan }));
+ Ok(())
+ }
+ Entry::Occupied(e) => {
+ exec_err!("Prepared statement '{}' already exists", e.key())
Review Comment:
Since this errors if an existing prepared statement should we add some way
to erase / clear prepared plans? Now there is no way to avoid accumulating them
over time
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For additional commands, e-mail: github-h...@datafusion.apache.org
