This is an automated email from the ASF dual-hosted git repository.
jonah pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new a20b3183b3 feat: support `DEALLOCATE` to remove prepared statements
(#13327)
a20b3183b3 is described below
commit a20b3183b33337a4dc0539393b4630dd32b76232
Author: Jonah Gao <[email protected]>
AuthorDate: Sun Nov 10 22:04:27 2024 +0800
feat: support `DEALLOCATE` to remove prepared statements (#13327)
---
datafusion/core/src/execution/context/mod.rs | 6 +++
datafusion/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 +++++++--
datafusion/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 ff009a0164..e04fe6bdde 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 f65bd2a597..d0bbc95a1b 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<Arc<PreparedPlan>>
{
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 b5bd2e0128..5d613d4e80 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 9ba8170f8e..05e2b1af14 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<Expr>,
}
+
+/// 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 3b32a43757..0ac804b706 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 f149cec96f..c3bd15c086 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;
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
