sql: add truncate statement support

Author: oleggator

datafusion/expr/src/logical_plan/dml.rs

@@ -119,6 +119,7 @@ pub enum WriteOp {
     Delete,
     Update,
     Ctas,
+    Truncate,
 }
 
 impl WriteOp {
@@ -130,6 +131,7 @@ impl WriteOp {
             WriteOp::Delete => "Delete",
             WriteOp::Update => "Update",
             WriteOp::Ctas => "Ctas",
+            WriteOp::Truncate => "Truncate",
         }
     }
 }

datafusion/sql/src/statement.rs

@@ -542,7 +542,21 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
                 }
                 self.update_to_plan(table, assignments, from, selection)
             }
+            Statement::Truncate {
+                table_name,
+                partitions,
+                table,
+            } => {
+                if !table {
+                    plan_err!("Truncate of non-tables not yet supported")?;
+                }
+
+                if partitions.is_some() {
+                    plan_err!("Partition clause not supported")?;
+                }
 
+                self.truncate_to_plan(table_name)
+            }
             Statement::Delete(Delete {
                 tables,
                 using,
@@ -1499,6 +1513,23 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
         Ok(plan)
     }
 
+    fn truncate_to_plan(&self, table_name: ObjectName) -> Result<LogicalPlan> {
+        // Do a table lookup to verify the table exists
+        let table_ref = self.object_name_to_table_reference(table_name.clone())?;
+        let table_source = self.context_provider.get_table_source(table_ref.clone())?;
+        let schema = (*table_source.schema()).clone();
+        let schema = DFSchema::try_from(schema)?;
+        let scan = LogicalPlanBuilder::empty(false).build()?;
+
+        let plan = LogicalPlan::Dml(DmlStatement::new(
+            table_ref,
+            schema.into(),
+            WriteOp::Truncate,
+            Arc::new(scan),
+        ));
+        Ok(plan)
+    }
+
     fn show_columns_to_plan(
         &self,
         extended: bool,