From 951a766746c384f93233f626e507891b3fe3e42f Mon Sep 17 00:00:00 2001 From: Jonah Gao Date: Sun, 10 Nov 2024 22:04:27 +0800 Subject: [PATCH] feat: support `DEALLOCATE` to remove prepared statements (#13327) --- datafusion/core/src/execution/context/mod.rs | 6 +++ .../core/src/execution/session_state.rs | 11 ++++ datafusion/expr/src/logical_plan/mod.rs | 2 +- datafusion/expr/src/logical_plan/statement.rs | 14 +++++ datafusion/sql/src/statement.rs | 17 ++++-- .../sqllogictest/test_files/prepare.slt | 52 +++++++++++++++++++ 6 files changed, 97 insertions(+), 5 deletions(-) diff --git a/datafusion/core/src/execution/context/mod.rs b/datafusion/core/src/execution/context/mod.rs index ff009a01645f2..e04fe6bddec94 100644 --- a/datafusion/core/src/execution/context/mod.rs +++ b/datafusion/core/src/execution/context/mod.rs @@ -715,6 +715,12 @@ impl SessionContext { LogicalPlan::Statement(Statement::Execute(execute)) => { self.execute_prepared(execute) } + LogicalPlan::Statement(Statement::Deallocate(deallocate)) => { + self.state + .write() + .remove_prepared(deallocate.name.as_str())?; + self.return_empty_dataframe() + } plan => Ok(DataFrame::new(self.state(), plan)), } } diff --git a/datafusion/core/src/execution/session_state.rs b/datafusion/core/src/execution/session_state.rs index f65bd2a5977f7..d0bbc95a1b08b 100644 --- a/datafusion/core/src/execution/session_state.rs +++ b/datafusion/core/src/execution/session_state.rs @@ -934,6 +934,17 @@ impl SessionState { pub(crate) fn get_prepared(&self, name: &str) -> Option> { self.prepared_plans.get(name).map(Arc::clone) } + + /// Remove the prepared plan with the given name. + pub(crate) fn remove_prepared( + &mut self, + name: &str, + ) -> datafusion_common::Result<()> { + match self.prepared_plans.remove(name) { + Some(_) => Ok(()), + None => exec_err!("Prepared statement '{}' does not exist", name), + } + } } /// A builder to be used for building [`SessionState`]'s. Defaults will diff --git a/datafusion/expr/src/logical_plan/mod.rs b/datafusion/expr/src/logical_plan/mod.rs index b5bd2e0128717..5d613d4e80dbf 100644 --- a/datafusion/expr/src/logical_plan/mod.rs +++ b/datafusion/expr/src/logical_plan/mod.rs @@ -42,7 +42,7 @@ pub use plan::{ SubqueryAlias, TableScan, ToStringifiedPlan, Union, Unnest, Values, Window, }; pub use statement::{ - Execute, Prepare, SetVariable, Statement, TransactionAccessMode, + Deallocate, Execute, Prepare, SetVariable, Statement, TransactionAccessMode, TransactionConclusion, TransactionEnd, TransactionIsolationLevel, TransactionStart, }; diff --git a/datafusion/expr/src/logical_plan/statement.rs b/datafusion/expr/src/logical_plan/statement.rs index 9ba8170f8eb27..05e2b1af14d3b 100644 --- a/datafusion/expr/src/logical_plan/statement.rs +++ b/datafusion/expr/src/logical_plan/statement.rs @@ -48,6 +48,9 @@ pub enum Statement { Prepare(Prepare), /// Execute a prepared statement. This is used to implement SQL 'EXECUTE'. Execute(Execute), + /// Deallocate a prepared statement. + /// This is used to implement SQL 'DEALLOCATE'. + Deallocate(Deallocate), } impl Statement { @@ -65,6 +68,7 @@ impl Statement { Statement::SetVariable(_) => "SetVariable", Statement::Prepare(_) => "Prepare", Statement::Execute(_) => "Execute", + Statement::Deallocate(_) => "Deallocate", } } @@ -167,6 +171,9 @@ impl Statement { expr_vec_fmt!(parameters) ) } + Statement::Deallocate(Deallocate { name }) => { + write!(f, "Deallocate: {}", name) + } } } } @@ -245,3 +252,10 @@ pub struct Execute { /// The execute parameters pub parameters: Vec, } + +/// Deallocate a prepared statement. +#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Hash)] +pub struct Deallocate { + /// The name of the prepared statement to deallocate + pub name: String, +} diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs index 3b32a437576b8..0ac804b706c8b 100644 --- a/datafusion/sql/src/statement.rs +++ b/datafusion/sql/src/statement.rs @@ -46,10 +46,10 @@ use datafusion_expr::utils::expr_to_columns; use datafusion_expr::{ cast, col, Analyze, CreateCatalog, CreateCatalogSchema, CreateExternalTable as PlanCreateExternalTable, CreateFunction, CreateFunctionBody, - CreateIndex as PlanCreateIndex, CreateMemoryTable, CreateView, DescribeTable, - DmlStatement, DropCatalogSchema, DropFunction, DropTable, DropView, EmptyRelation, - Execute, Explain, Expr, ExprSchemable, Filter, LogicalPlan, LogicalPlanBuilder, - OperateFunctionArg, PlanType, Prepare, SetVariable, SortExpr, + CreateIndex as PlanCreateIndex, CreateMemoryTable, CreateView, Deallocate, + DescribeTable, DmlStatement, DropCatalogSchema, DropFunction, DropTable, DropView, + EmptyRelation, Execute, Explain, Expr, ExprSchemable, Filter, LogicalPlan, + LogicalPlanBuilder, OperateFunctionArg, PlanType, Prepare, SetVariable, SortExpr, Statement as PlanStatement, ToStringifiedPlan, TransactionAccessMode, TransactionConclusion, TransactionEnd, TransactionIsolationLevel, TransactionStart, Volatility, WriteOp, @@ -665,6 +665,15 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { parameters, }))) } + Statement::Deallocate { + name, + // Similar to PostgreSQL, the PREPARE keyword is ignored + prepare: _, + } => Ok(LogicalPlan::Statement(PlanStatement::Deallocate( + Deallocate { + name: ident_to_string(&name), + }, + ))), Statement::ShowTables { extended, diff --git a/datafusion/sqllogictest/test_files/prepare.slt b/datafusion/sqllogictest/test_files/prepare.slt index f149cec96f197..c3bd15c086262 100644 --- a/datafusion/sqllogictest/test_files/prepare.slt +++ b/datafusion/sqllogictest/test_files/prepare.slt @@ -64,6 +64,10 @@ PREPARE my_plan AS SELECT $1; statement error Prepared statement \'my_plan\' does not exist EXECUTE my_plan('Foo', 'Bar'); +# deallocate a non-existing plan +statement error Prepared statement \'my_plan\' does not exist +DEALLOCATE my_plan; + statement ok PREPARE my_plan(STRING, STRING) AS SELECT * FROM (VALUES(1, $1), (2, $2)) AS t (num, letter); @@ -77,6 +81,28 @@ EXECUTE my_plan('Foo', 'Bar'); 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); +# deallocate a plan +statement ok +DEALLOCATE my_plan; + +# can't EXECUTE a deallocated plan +statement error Prepared statement \'my_plan\' does not exist +EXECUTE my_plan('Foo', 'Bar'); + +# re-prepare a deallocated plan +statement ok +PREPARE my_plan(STRING, STRING) AS SELECT * FROM (VALUES(1, $1), (2, $2)) AS t (num, letter); + +query IT +EXECUTE my_plan('Foo', 'Bar'); +---- +1 Foo +2 Bar + +# deallocate with the PREPARE keyword +statement ok +DEALLOCATE PREPARE my_plan; + 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; @@ -89,6 +115,9 @@ EXECUTE my_plan2; ---- 1 20 +statement ok +DEALLOCATE my_plan2; + statement ok PREPARE my_plan3(INT) AS SELECT $1; @@ -97,6 +126,9 @@ EXECUTE my_plan3(10); ---- 10 +statement ok +DEALLOCATE my_plan3; + statement ok PREPARE my_plan4(INT) AS SELECT 1 + $1; @@ -105,6 +137,9 @@ EXECUTE my_plan4(10); ---- 11 +statement ok +DEALLOCATE my_plan4; + statement ok PREPARE my_plan5(INT, DOUBLE) AS SELECT 1 + $1 + $2; @@ -113,6 +148,9 @@ EXECUTE my_plan5(10, 20.5); ---- 31.5 +statement ok +DEALLOCATE my_plan5; + statement ok PREPARE my_plan6(INT) AS SELECT id, age FROM person WHERE age = $1; @@ -140,6 +178,9 @@ EXECUTE my_plan6('foo'); statement error Unsupported parameter type EXECUTE my_plan6(10 + 20); +statement ok +DEALLOCATE my_plan6; + statement ok PREPARE my_plan7(INT, STRING, DOUBLE, INT, DOUBLE, STRING) AS @@ -150,6 +191,9 @@ EXECUTE my_plan7(10, 'jane', 99999.45, 20, 200000.45, 'foo'); ---- 1 20 foo +statement ok +DEALLOCATE my_plan7; + statement ok PREPARE my_plan8(INT, DOUBLE, DOUBLE, DOUBLE) AS @@ -161,6 +205,9 @@ EXECUTE my_plan8(100000, 99999.45, 100000.45, 200000.45); ---- 1 20 +statement ok +DEALLOCATE my_plan8; + statement ok PREPARE my_plan9(STRING, STRING) AS SELECT * FROM (VALUES(1, $1), (2, $2)) AS t (num, letter); @@ -170,6 +217,8 @@ EXECUTE my_plan9('Foo', 'Bar'); 1 Foo 2 Bar +statement ok +DEALLOCATE my_plan9; # Test issue: https://github.com/apache/datafusion/issues/12294 # prepare argument is in the LIMIT clause @@ -196,6 +245,9 @@ EXECUTE get_N_rand_ints_from_last_run(2); 1 1 +statement ok +DEALLOCATE get_N_rand_ints_from_last_run; + statement ok DROP TABLE test;