feat: support DEALLOCATE to remove prepared statements (#13327)

Author: jonahgao

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)),
         }
     }

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

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,
 };
 

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,
+}

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,

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;