diff --git a/datafusion/sql/src/query.rs b/datafusion/sql/src/query.rs index e320d2ee6e9c1..a68ef84871e4d 100644 --- a/datafusion/sql/src/query.rs +++ b/datafusion/sql/src/query.rs @@ -20,7 +20,7 @@ use std::sync::Arc; use crate::planner::{ContextProvider, PlannerContext, SqlToRel}; use crate::stack::StackGuard; -use datafusion_common::{Constraints, DFSchema, Result, not_impl_err}; +use datafusion_common::{Constraints, DFSchema, Result, not_impl_err, plan_err}; use datafusion_expr::expr::{Sort, WildcardOptions}; use datafusion_expr::select_expr::SelectExpr; @@ -28,9 +28,9 @@ use datafusion_expr::{ CreateMemoryTable, DdlStatement, Distinct, Expr, LogicalPlan, LogicalPlanBuilder, }; use sqlparser::ast::{ - Expr as SQLExpr, ExprWithAliasAndOrderBy, Ident, LimitClause, Offset, OffsetRows, - OrderBy, OrderByExpr, OrderByKind, PipeOperator, Query, SelectInto, SetExpr, - SetOperator, SetQuantifier, TableAlias, + ExceptSelectItem, Expr as SQLExpr, ExprWithAliasAndOrderBy, Ident, LimitClause, + Offset, OffsetRows, OrderBy, OrderByExpr, OrderByKind, PipeOperator, Query, + SelectInto, SetExpr, SetOperator, SetQuantifier, TableAlias, }; use sqlparser::tokenizer::Span; @@ -215,6 +215,19 @@ impl SqlToRel<'_, S> { PipeOperator::Join(join) => { self.parse_relation_join(plan, join, planner_context) } + PipeOperator::Drop { columns } => { + let Some((first, rest)) = columns.split_first() else { + return plan_err!("DROP requires at least one column"); + }; + let options = WildcardOptions { + except: Some(ExceptSelectItem { + first_element: first.clone(), + additional_elements: rest.to_vec(), + }), + ..WildcardOptions::default() + }; + self.project(plan, vec![SelectExpr::Wildcard(options)]) + } x => not_impl_err!("`{x}` pipe operator is not supported yet"), } diff --git a/datafusion/sqllogictest/test_files/pipe_operator.slt b/datafusion/sqllogictest/test_files/pipe_operator.slt index 406ddafc7bdea..5302a1d4adec1 100644 --- a/datafusion/sqllogictest/test_files/pipe_operator.slt +++ b/datafusion/sqllogictest/test_files/pipe_operator.slt @@ -196,6 +196,32 @@ query TII apples 2 123 bananas 5 NULL +# DROP pipe - drop two columns +query RT +SELECT * FROM test |> DROP a, n +---- +1.1 a +2.2 b +3.3 c + +# DROP pipe - drop single column +query IRT +SELECT * FROM test |> DROP n +---- +1 1.1 a +2 2.2 b +3 3.3 c + +# DROP pipe - chained with other operators +query R +SELECT * FROM test |> WHERE a > 1 |> DROP a, c, n +---- +2.2 +3.3 + +# Error: DROP non-existent column +statement error +SELECT * FROM test |> DROP nonexistent # Config reset statement ok RESET datafusion.sql_parser.dialect; diff --git a/docs/source/user-guide/sql/select.md b/docs/source/user-guide/sql/select.md index 3564884b041ad..c8c4f1d58b9e7 100644 --- a/docs/source/user-guide/sql/select.md +++ b/docs/source/user-guide/sql/select.md @@ -481,6 +481,7 @@ DataFusion currently supports the following pipe operators: - [EXCEPT](#pipe_except) - [AGGREGATE](#pipe_aggregate) - [JOIN](#pipe_join) +- [DROP](#pipe_drop) (pipe_where)= @@ -670,3 +671,19 @@ select * from range(0,3) | bananas| 5 | NULL | +--------+-------+------+ ``` + +(pipe_drop)= + +### DROP + +Removes columns from the output, equivalent to `SELECT * EXCEPT(columns)`. + +```sql +select 1 as a, 2 as b, 3 as c +|> drop b; ++---+---+ +| a | c | ++---+---+ +| 1 | 3 | ++---+---+ +```