added truncate table support

Author: Nachiket-Roy

datafusion/catalog/src/table.rs

@@ -353,6 +353,14 @@ pub trait TableProvider: Debug + Sync + Send {
     ) -> Result<Arc<dyn ExecutionPlan>> {
         not_impl_err!("UPDATE not supported for {} table", self.table_type())
     }
+
+    /// Truncate rows
+    async fn truncate(
+        &self,
+        _state: &dyn Session,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        not_impl_err!("TRUNCATE not supported for {}", self.table_type())
+    }
 }
 
 /// Arguments for scanning a table with [`TableProvider::scan_with_args`].

datafusion/core/src/physical_planner.rs

@@ -655,6 +655,31 @@ impl DefaultPhysicalPlanner {
                     );
                 }
             }
+            LogicalPlan::Dml(DmlStatement {
+                table_name,
+                target,
+                op: WriteOp::Truncate,
+                ..
+            }) => {
+                if let Some(provider) =
+                    target.as_any().downcast_ref::<DefaultTableSource>()
+                {
+                    provider
+                        .table_provider
+                        .truncate(session_state)
+                        .await
+                        .map_err(|e| {
+                            e.context(format!(
+                                "TRUNCATE operation on table '{}'",
+                                table_name
+                            ))
+                        })?
+                } else {
+                    return exec_err!(
+                        "Table source can't be downcasted to DefaultTableSource"
+                    );
+                }
+            }
             LogicalPlan::Window(Window { window_expr, .. }) => {
                 assert_or_internal_err!(
                     !window_expr.is_empty(),

datafusion/core/tests/custom_sources_cases/dml_planning.rs

@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-//! Tests for DELETE and UPDATE planning to verify filter and assignment extraction.
+//! Tests for DELETE, UPDATE, and TRUNCATE planning to verify filter and assignment extraction.
 
 use std::any::Any;
 use std::sync::{Arc, Mutex};
@@ -165,6 +165,69 @@ impl TableProvider for CaptureUpdateProvider {
     }
 }
 
+/// A TableProvider that captures whether truncate() was called.
+struct CaptureTruncateProvider {
+    schema: SchemaRef,
+    truncate_called: Arc<Mutex<bool>>,
+}
+
+impl CaptureTruncateProvider {
+    fn new(schema: SchemaRef) -> Self {
+        Self {
+            schema,
+            truncate_called: Arc::new(Mutex::new(false)),
+        }
+    }
+
+    fn was_truncated(&self) -> bool {
+        *self.truncate_called.lock().unwrap()
+    }
+}
+
+impl std::fmt::Debug for CaptureTruncateProvider {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("CaptureTruncateProvider")
+            .field("schema", &self.schema)
+            .finish()
+    }
+}
+
+#[async_trait]
+impl TableProvider for CaptureTruncateProvider {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn schema(&self) -> SchemaRef {
+        Arc::clone(&self.schema)
+    }
+
+    fn table_type(&self) -> TableType {
+        TableType::Base
+    }
+
+    async fn scan(
+        &self,
+        _state: &dyn Session,
+        _projection: Option<&Vec<usize>>,
+        _filters: &[Expr],
+        _limit: Option<usize>,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        Ok(Arc::new(EmptyExec::new(Arc::clone(&self.schema))))
+    }
+
+    async fn truncate(
+        &self,
+        _state: &dyn Session,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        *self.truncate_called.lock().unwrap() = true;
+
+        Ok(Arc::new(EmptyExec::new(Arc::new(Schema::new(vec![
+            Field::new("count", DataType::UInt64, false),
+        ])))))
+    }
+}
+
 fn test_schema() -> SchemaRef {
     Arc::new(Schema::new(vec![
         Field::new("id", DataType::Int32, false),
@@ -269,6 +332,26 @@ async fn test_update_assignments() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn test_truncate_calls_provider() -> Result<()> {
+    let provider = Arc::new(CaptureTruncateProvider::new(test_schema()));
+    let ctx = SessionContext::new();
+
+    ctx.register_table("t", Arc::clone(&provider) as Arc<dyn TableProvider>)?;
+
+    ctx.sql("TRUNCATE TABLE t")
+        .await?
+        .collect()
+        .await?;
+
+    assert!(
+        provider.was_truncated(),
+        "truncate() should be called on the TableProvider"
+    );
+
+    Ok(())
+}
+
 #[tokio::test]
 async fn test_unsupported_table_delete() -> Result<()> {
     let schema = test_schema();
@@ -295,3 +378,18 @@ async fn test_unsupported_table_update() -> Result<()> {
     assert!(result.is_err() || result.unwrap().collect().await.is_err());
     Ok(())
 }
+
+#[tokio::test]
+async fn test_unsupported_table_truncate() -> Result<()> {
+    let schema = test_schema();
+    let ctx = SessionContext::new();
+
+    let empty_table = datafusion::datasource::empty::EmptyTable::new(schema);
+    ctx.register_table("empty_t", Arc::new(empty_table))?;
+
+    let result = ctx.sql("TRUNCATE TABLE empty_t").await;
+
+    assert!(result.is_err() || result.unwrap().collect().await.is_err());
+
+    Ok(())
+}

datafusion/expr/src/logical_plan/dml.rs

@@ -237,6 +237,8 @@ pub enum WriteOp {
     Update,
     /// `CREATE TABLE AS SELECT` operation
     Ctas,
+    /// `TRUNCATE` operation
+    Truncate,
 }
 
 impl WriteOp {
@@ -247,6 +249,7 @@ impl WriteOp {
             WriteOp::Delete => "Delete",
             WriteOp::Update => "Update",
             WriteOp::Ctas => "Ctas",
+            WriteOp::Truncate => "Truncate",
         }
     }
 }

datafusion/proto/src/generated/prost.rs

@@ -239,6 +239,7 @@ impl From<protobuf::dml_node::Type> for WriteOp {
             }
             protobuf::dml_node::Type::InsertReplace => WriteOp::Insert(InsertOp::Replace),
             protobuf::dml_node::Type::Ctas => WriteOp::Ctas,
+            protobuf::dml_node::Type::Truncate => WriteOp::Truncate,
         }
     }
 }

datafusion/proto/src/logical_plan/from_proto.rs

@@ -728,6 +728,7 @@ impl From<&WriteOp> for protobuf::dml_node::Type {
             WriteOp::Delete => protobuf::dml_node::Type::Delete,
             WriteOp::Update => protobuf::dml_node::Type::Update,
             WriteOp::Ctas => protobuf::dml_node::Type::Ctas,
+            WriteOp::Truncate => protobuf::dml_node::Type::Truncate,
         }
     }
 }

datafusion/proto/src/logical_plan/to_proto.rs

@@ -1362,6 +1362,28 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
                     exec_err!("Function name not provided")
                 }
             }
+            Statement::Truncate { table_names, .. } => {
+                if table_names.len() != 1 {
+                    return not_impl_err!("TRUNCATE with multiple tables is not supported");
+                }
+
+                let target = &table_names[0]; // TruncateTableTarget
+                let table = self.object_name_to_table_reference(target.name.clone())?;
+                let source = self.context_provider.get_table_source(table.clone())?;
+
+                Ok(LogicalPlan::Dml(DmlStatement {
+                    table_name: table.clone(),
+                    target: source,
+                    op: WriteOp::Truncate,
+                    input: Arc::new(LogicalPlan::EmptyRelation(
+                        EmptyRelation {
+                            produce_one_row: false,
+                            schema: DFSchemaRef::new(DFSchema::empty()),
+                        },
+                    )),
+                    output_schema: DFSchemaRef::new(DFSchema::empty()),
+                }))
+            }
             Statement::CreateIndex(CreateIndex {
                 name,
                 table_name,

datafusion/sql/src/statement.rs

@@ -0,0 +1,91 @@
+
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+
+#   http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+##########
+## Truncate tests for MemTable
+##########
+
+# Test that TRUNCATE preserves table structure but removes data
+statement ok
+truncate table t1;
+
+query I
+select count(*) from t1;
+----
+0
+
+query TT
+describe t1;
+----
+a:Int32:None
+b:Utf8:None
+c:Float64:None
+d:Int32:None
+
+# Test TRUNCATE then INSERT works correctly
+statement ok
+truncate table t1;
+
+statement ok
+insert into t1 values (10, 'new', 1.23, 100);
+
+query I
+select count(*) from t1;
+----
+1
+
+query I
+select a from t1;
+----
+10
+
+# Test multiple TRUNCATE operations
+statement ok
+insert into t1 values (20, 'another', 4.56, 200);
+
+query I
+select count(*) from t1;
+----
+2
+
+statement ok
+truncate table t1;
+
+query I
+select count(*) from t1;
+----
+0
+
+# Truncate already empty table
+statement ok
+truncate table t1;
+
+query I
+select count(*) from t1;
+----
+0
+
+# Clean up
+statement ok
+drop table t1;
+
+statement ok
+drop table test_schema.t5;
+
+statement ok
+drop schema test_schema;

datafusion/sqllogictest/test_files/dml_truncate.slt

@@ -0,0 +1,52 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+
+#   http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+##########
+## Truncate Tests
+##########
+
+statement ok
+create table t1(a int, b varchar, c double, d int);
+
+statement ok
+insert into t1 values (1, 'abc', 3.14, 4), (2, 'def', 2.71, 5);
+
+# Truncate all rows from table
+query TT
+explain truncate table t1;
+----
+logical_plan
+01)Dml: op=[Truncate] table=[t1]
+physical_plan
+01)CooperativeExec
+02)--DmlResultExec: rows_affected=0
+
+# Test TRUNCATE with fully qualified table name
+statement ok
+create schema test_schema;
+
+statement ok
+create table test_schema.t5(a int);
+
+query TT
+explain truncate table test_schema.t5;
+----
+logical_plan
+01)Dml: op=[Truncate] table=[test_schema.t5]
+physical_plan
+01)CooperativeExec
+02)--DmlResultExec: rows_affected=0