Skip to content

Commit ca47b51

Browse files
oleggatoroleggator
committed
sql: add truncate statement support
1 parent d6ff3db commit ca47b51

File tree

9 files changed

+93
-14
lines changed

9 files changed

+93
-14
lines changed

‎datafusion/core/src/physical_planner.rs

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1227,6 +1227,14 @@ impl DefaultPhysicalPlanner {
12271227
return internal_err!(
12281228
"Unsupported logical plan: Analyze must be root of the plan"
12291229
)
1230+
},
1231+
LogicalPlan::Truncate(_) => {
1232+
// There is no default plan for "TRUNCATE" -- it must be
1233+
// handled at a higher level (so that the appropriate
1234+
// table can be truncated)
1235+
return internal_err!(
1236+
"Unsupported logical plan: Truncate"
1237+
)
12301238
}
12311239
};
12321240
Ok(exec_node)

‎datafusion/expr/src/logical_plan/display.rs

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -654,7 +654,12 @@ impl<'a, 'b> PgJsonVisitor<'a, 'b> {
654654
"ListColumn": expr_vec_fmt!(list_type_columns),
655655
"StructColumn": expr_vec_fmt!(struct_type_columns),
656656
})
657-
}
657+
},
658+
LogicalPlan::Truncate(_) => {
659+
json!({
660+
"Node Type": "Truncate"
661+
})
662+
},
658663
}
659664
}
660665
}

‎datafusion/expr/src/logical_plan/dml.rs

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,31 @@ impl Hash for CopyTo {
7171
}
7272
}
7373

74+
/// Operator that truncates the content of a table
75+
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
76+
pub struct Truncate {
77+
/// The table name
78+
pub table_name: TableReference,
79+
/// The schema of the output relation
80+
pub output_schema: DFSchemaRef,
81+
}
82+
83+
impl Truncate {
84+
/// Creates a new truncate statement with the output schema set empty.
85+
pub fn new(table_name: TableReference) -> Self {
86+
Self {
87+
table_name,
88+
89+
// The output schema is always empty
90+
output_schema: make_empty_schema(),
91+
}
92+
}
93+
}
94+
95+
fn make_empty_schema() -> DFSchemaRef {
96+
Arc::new(Schema::empty().try_into().unwrap())
97+
}
98+
7499
/// The operator that modifies the content of a database (adapted from
75100
/// substrait WriteRel)
76101
#[derive(Debug, Clone, PartialEq, Eq, Hash)]

‎datafusion/expr/src/logical_plan/plan.rs

Lines changed: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ use std::sync::Arc;
2525
use super::dml::CopyTo;
2626
use super::DdlStatement;
2727
use crate::builder::{change_redundant_column, unnest_with_options};
28+
use crate::dml::Truncate;
2829
use crate::expr::{Placeholder, Sort as SortExpr, WindowFunction};
2930
use crate::expr_rewriter::{
3031
create_col_from_scalar_expr, normalize_cols, normalize_sorts, NamePreserver,
@@ -280,6 +281,8 @@ pub enum LogicalPlan {
280281
Unnest(Unnest),
281282
/// A variadic query (e.g. "Recursive CTEs")
282283
RecursiveQuery(RecursiveQuery),
284+
/// Truncate a table
285+
Truncate(Truncate),
283286
}
284287

285288
impl Default for LogicalPlan {
@@ -329,7 +332,8 @@ impl LogicalPlan {
329332
LogicalPlan::RecursiveQuery(RecursiveQuery { static_term, .. }) => {
330333
// we take the schema of the static term as the schema of the entire recursive query
331334
static_term.schema()
332-
}
335+
},
336+
LogicalPlan::Truncate(Truncate{output_schema, ..}) => output_schema,
333337
}
334338
}
335339

@@ -481,7 +485,8 @@ impl LogicalPlan {
481485
| LogicalPlan::Statement { .. }
482486
| LogicalPlan::EmptyRelation { .. }
483487
| LogicalPlan::Values { .. }
484-
| LogicalPlan::DescribeTable(_) => vec![],
488+
| LogicalPlan::DescribeTable(_)
489+
| LogicalPlan::Truncate(_) => vec![],
485490
}
486491
}
487492

@@ -598,7 +603,8 @@ impl LogicalPlan {
598603
| LogicalPlan::Copy(_)
599604
| LogicalPlan::Ddl(_)
600605
| LogicalPlan::DescribeTable(_)
601-
| LogicalPlan::Unnest(_) => Ok(None),
606+
| LogicalPlan::Unnest(_)
607+
| LogicalPlan::Truncate(_) => Ok(None),
602608
}
603609
}
604610

@@ -767,7 +773,8 @@ impl LogicalPlan {
767773
}) => {
768774
// Update schema with unnested column type.
769775
unnest_with_options(Arc::unwrap_or_clone(input), exec_columns, options)
770-
}
776+
},
777+
LogicalPlan::Truncate(_) => Ok(self),
771778
}
772779
}
773780

@@ -1117,7 +1124,8 @@ impl LogicalPlan {
11171124
LogicalPlan::EmptyRelation(_)
11181125
| LogicalPlan::Ddl(_)
11191126
| LogicalPlan::Statement(_)
1120-
| LogicalPlan::DescribeTable(_) => {
1127+
| LogicalPlan::DescribeTable(_)
1128+
| LogicalPlan::Truncate(_) => {
11211129
// All of these plan types have no inputs / exprs so should not be called
11221130
self.assert_no_expressions(expr)?;
11231131
self.assert_no_inputs(inputs)?;
@@ -1367,7 +1375,8 @@ impl LogicalPlan {
13671375
| LogicalPlan::DescribeTable(_)
13681376
| LogicalPlan::Prepare(_)
13691377
| LogicalPlan::Statement(_)
1370-
| LogicalPlan::Extension(_) => None,
1378+
| LogicalPlan::Extension(_)
1379+
| LogicalPlan::Truncate(_) => None,
13711380
}
13721381
}
13731382

@@ -1990,6 +1999,9 @@ impl LogicalPlan {
19901999
write!(f, "Unnest: lists[{}] structs[{}]",
19912000
expr_vec_fmt!(list_type_columns),
19922001
expr_vec_fmt!(struct_type_columns))
2002+
},
2003+
LogicalPlan::Truncate(_) => {
2004+
write!(f, "Truncate")
19932005
}
19942006
}
19952007
}

‎datafusion/expr/src/logical_plan/tree_node.rs

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -374,7 +374,8 @@ impl TreeNode for LogicalPlan {
374374
| LogicalPlan::Statement { .. }
375375
| LogicalPlan::EmptyRelation { .. }
376376
| LogicalPlan::Values { .. }
377-
| LogicalPlan::DescribeTable(_) => Transformed::no(self),
377+
| LogicalPlan::DescribeTable(_)
378+
| LogicalPlan::Truncate(_) => Transformed::no(self),
378379
})
379380
}
380381
}
@@ -527,7 +528,8 @@ impl LogicalPlan {
527528
| LogicalPlan::Ddl(_)
528529
| LogicalPlan::Copy(_)
529530
| LogicalPlan::DescribeTable(_)
530-
| LogicalPlan::Prepare(_) => Ok(TreeNodeRecursion::Continue),
531+
| LogicalPlan::Prepare(_)
532+
| LogicalPlan::Truncate(_) => Ok(TreeNodeRecursion::Continue),
531533
}
532534
}
533535

@@ -739,7 +741,8 @@ impl LogicalPlan {
739741
| LogicalPlan::Ddl(_)
740742
| LogicalPlan::Copy(_)
741743
| LogicalPlan::DescribeTable(_)
742-
| LogicalPlan::Prepare(_) => Transformed::no(self),
744+
| LogicalPlan::Prepare(_)
745+
| LogicalPlan::Truncate(_) => Transformed::no(self),
743746
})
744747
}
745748

‎datafusion/optimizer/src/common_subexpr_eliminate.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -803,7 +803,8 @@ impl OptimizerRule for CommonSubexprEliminate {
803803
| LogicalPlan::Copy(_)
804804
| LogicalPlan::Unnest(_)
805805
| LogicalPlan::RecursiveQuery(_)
806-
| LogicalPlan::Prepare(_) => {
806+
| LogicalPlan::Prepare(_)
807+
| LogicalPlan::Truncate(_) => {
807808
// This rule handles recursion itself in a `ApplyOrder::TopDown` like
808809
// manner.
809810
plan.map_children(|c| self.rewrite(c, config))?

‎datafusion/optimizer/src/optimize_projections/mod.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -349,7 +349,8 @@ fn optimize_projections(
349349
| LogicalPlan::RecursiveQuery(_)
350350
| LogicalPlan::Statement(_)
351351
| LogicalPlan::Values(_)
352-
| LogicalPlan::DescribeTable(_) => {
352+
| LogicalPlan::DescribeTable(_)
353+
| LogicalPlan::Truncate(_) => {
353354
// These operators have no inputs, so stop the optimization process.
354355
return Ok(Transformed::no(plan));
355356
}

‎datafusion/sql/src/statement.rs

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@ use datafusion_common::{
3737
DataFusionError, Result, ScalarValue, SchemaError, SchemaReference, TableReference,
3838
ToDFSchema,
3939
};
40-
use datafusion_expr::dml::CopyTo;
40+
use datafusion_expr::dml::{CopyTo, Truncate};
4141
use datafusion_expr::expr_rewriter::normalize_col_with_schemas_and_ambiguity_check;
4242
use datafusion_expr::logical_plan::builder::project;
4343
use datafusion_expr::logical_plan::DdlStatement;
@@ -542,7 +542,21 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
542542
}
543543
self.update_to_plan(table, assignments, from, selection)
544544
}
545+
Statement::Truncate {
546+
table_name,
547+
partitions,
548+
table,
549+
} => {
550+
if !table {
551+
plan_err!("Truncate of non-tables not yet supported")?;
552+
}
545553

554+
if partitions.is_some() {
555+
plan_err!("Partition clause not supported")?;
556+
}
557+
558+
self.truncate_to_plan(table_name)
559+
}
546560
Statement::Delete(Delete {
547561
tables,
548562
using,
@@ -1499,6 +1513,15 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
14991513
Ok(plan)
15001514
}
15011515

1516+
fn truncate_to_plan(&self, table_name: ObjectName) -> Result<LogicalPlan> {
1517+
// Do a table lookup to verify the table exists
1518+
let table_ref = self.object_name_to_table_reference(table_name.clone())?;
1519+
let _ = self.context_provider.get_table_source(table_ref.clone())?;
1520+
1521+
let plan = LogicalPlan::Truncate(Truncate::new(table_ref));
1522+
Ok(plan)
1523+
}
1524+
15021525
fn show_columns_to_plan(
15031526
&self,
15041527
extended: bool,

‎datafusion/sql/src/unparser/plan.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -112,7 +112,8 @@ impl Unparser<'_> {
112112
| LogicalPlan::Copy(_)
113113
| LogicalPlan::DescribeTable(_)
114114
| LogicalPlan::RecursiveQuery(_)
115-
| LogicalPlan::Unnest(_) => not_impl_err!("Unsupported plan: {plan:?}"),
115+
| LogicalPlan::Unnest(_)
116+
| LogicalPlan::Truncate(_) => not_impl_err!("Unsupported plan: {plan:?}"),
116117
}
117118
}
118119

0 commit comments

Comments
 (0)