Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: EXCEPT in SELECT #10438

Merged
merged 20 commits into from
Jun 25, 2023
Merged
15 changes: 15 additions & 0 deletions e2e_test/batch/basic/query.slt.part
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,21 @@ select count(*) from t3;
----
1

query III
select * except (v1) from t3;
----
2 NULL

query III
select * except (t3.v1) from t3;
----
2 NULL

query III
select * except (v1), * except (v2) from t3;
----
2 NULL 1 NULL

statement error Division by zero
select v1/0 from t3;

Expand Down
14 changes: 14 additions & 0 deletions e2e_test/batch/basic/subquery.slt.part
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,20 @@ NULL 1
NULL 2
NULL NULL

query II
select * except (b,d) from (select t1.x as a, t1.y as b, t2.x as c, t2.y as d from t1 join t2 on t1.x = t2.x where t1.x=1);
----
1 1
1 1
1 1
1 1

query II
select * except (t1.x, t2.y), * except (t1.y, t2.x) from t1 join t2 on t1.y = t2.y where exists(select * from t3 where t1.x = t3.x and t2.y = t3.y) order by t2.x;
----
2 1 2 2
2 2 2 2
2 NULL 2 2

statement ok
drop table t1;
Expand Down
13 changes: 13 additions & 0 deletions e2e_test/batch/top_n/group_top_n.slt
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,19 @@ where rank <= 3 AND rank > 1;
3 2
3 3

query II rowsort
select * except (rank) from (
select *, ROW_NUMBER() OVER (PARTITION BY x ORDER BY y) as rank from t
)
where rank <= 3 AND rank > 1;
----
1 2
1 3
2 2
2 3
3 2
3 3

query II rowsort
select x, y from (
select *, RANK() OVER (ORDER BY y) as rank from t
Expand Down
11 changes: 11 additions & 0 deletions src/frontend/planner_test/tests/testdata/input/basic_query.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,17 @@
expected_outputs:
- stream_plan
- batch_plan
- sql: |
create table t (v1 int, v2 int, v3 int);
select * except (v1, v2) from t;
expected_outputs:
- stream_plan
- batch_plan
- sql: |
create table t (v1 int, v2 int, v3 int);
select * except (v1, v2), v3 from t;
expected_outputs:
- batch_plan
- name: test boolean expression common factor extraction
sql: |
create table t (v1 Boolean, v2 Boolean, v3 Boolean);
Expand Down
16 changes: 16 additions & 0 deletions src/frontend/planner_test/tests/testdata/output/basic_query.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,22 @@
StreamMaterialize { columns: [v1, t._row_id(hidden)], stream_key: [t._row_id], pk_columns: [t._row_id], pk_conflict: NoCheck }
└─StreamFilter { predicate: (t.v1 < 1:Int32) }
└─StreamTableScan { table: t, columns: [t.v1, t._row_id], pk: [t._row_id], dist: UpstreamHashShard(t._row_id) }
- sql: |
create table t (v1 int, v2 int, v3 int);
select * except (v1, v2) from t;
Copy link
Contributor

@kwannoel kwannoel Jun 21, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we add more tests which explicitly select except columns?
Something like:

select v3 from (select * except (v1, v2) from t

Can we also add tests which have selected the excepted row, and re-select already included rows?

select * except (v1, v2), v1 from t; -- v1 excluded from except.
select * except (v1, v2), v3 from t; -- v3 included from except.

batch_plan: |-
BatchExchange { order: [], dist: Single }
└─BatchScan { table: t, columns: [t.v3], distribution: SomeShard }
stream_plan: |-
StreamMaterialize { columns: [v3, t._row_id(hidden)], stream_key: [t._row_id], pk_columns: [t._row_id], pk_conflict: NoCheck }
└─StreamTableScan { table: t, columns: [t.v3, t._row_id], pk: [t._row_id], dist: UpstreamHashShard(t._row_id) }
- sql: |
create table t (v1 int, v2 int, v3 int);
select * except (v1, v2), v3 from t;
batch_plan: |-
BatchExchange { order: [], dist: Single }
└─BatchProject { exprs: [t.v3, t.v3] }
└─BatchScan { table: t, columns: [t.v3], distribution: SomeShard }
- name: test boolean expression common factor extraction
sql: |
create table t (v1 Boolean, v2 Boolean, v3 Boolean);
Expand Down
3 changes: 2 additions & 1 deletion src/frontend/src/binder/expr/function.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1018,7 +1018,8 @@ impl Binder {
FunctionArgExpr::Expr(expr) => Ok(vec![self.bind_expr_inner(expr)?]),
FunctionArgExpr::QualifiedWildcard(_) => todo!(),
FunctionArgExpr::ExprQualifiedWildcard(_, _) => todo!(),
FunctionArgExpr::Wildcard => Ok(vec![]),
FunctionArgExpr::WildcardOrWithExcept(None) => Ok(vec![]),
FunctionArgExpr::WildcardOrWithExcept(Some(_)) => unreachable!(),
}
}

Expand Down
35 changes: 25 additions & 10 deletions src/frontend/src/binder/select.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::HashMap;
use std::collections::{HashMap, HashSet};
use std::fmt::Debug;

use itertools::Itertools;
Expand Down Expand Up @@ -259,20 +259,33 @@ impl Binder {
select_list.extend(exprs);
aliases.extend(names);
}
SelectItem::Wildcard => {
SelectItem::WildcardOrWithExcept(w) => {
if self.context.range_of.is_empty() {
return Err(ErrorCode::BindError(
"SELECT * with no tables specified is not valid".into(),
)
.into());
}

// Bind the column groups
// In psql, the USING and NATURAL columns come before the rest of the columns in
// a SELECT * statement
// In psql, the USING and NATURAL columns come before the rest of the
// columns in a SELECT * statement
let (exprs, names) = self.iter_column_groups();
select_list.extend(exprs);
aliases.extend(names);

let mut except_indices: HashSet<usize> = HashSet::new();
if let Some(exprs) = w {
for expr in exprs {
let bound = self.bind_expr(expr)?;
if let ExprImpl::InputRef(inner) = bound {
except_indices.insert(inner.index);
} else {
unreachable!();
}
}
}

// Bind columns that are not in groups
let (exprs, names) =
Self::iter_bound_columns(self.context.columns[..].iter().filter(|c| {
Expand All @@ -282,17 +295,19 @@ impl Binder {
.column_group_context
.mapping
.contains_key(&c.index)
&& !except_indices.contains(&c.index)
}));

select_list.extend(exprs);
aliases.extend(names);

// TODO: we will need to be able to handle wildcard expressions bound to aliases
// in the future. We'd then need a `NaturalGroupContext`
// bound to each alias to correctly disambiguate column
// TODO: we will need to be able to handle wildcard expressions bound to
// aliases in the future. We'd then need a
// `NaturalGroupContext` bound to each alias
// to correctly disambiguate column
// references
//
// We may need to refactor `NaturalGroupContext` to become span aware in that
// case.
// We may need to refactor `NaturalGroupContext` to become span aware in
// that case.
}
}
}
Expand Down
2 changes: 1 addition & 1 deletion src/frontend/src/handler/create_sink.rs
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ pub fn gen_sink_query_from_name(from_name: ObjectName) -> Result<Query> {
}];
let select = Select {
from,
projection: vec![SelectItem::Wildcard],
projection: vec![SelectItem::WildcardOrWithExcept(None)],
..Default::default()
};
let body = SetExpr::Select(Box::new(select));
Expand Down
15 changes: 13 additions & 2 deletions src/meta/src/manager/catalog/utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -244,7 +244,13 @@ impl QueryRewriter<'_> {
FunctionArgExpr::Expr(expr) | FunctionArgExpr::ExprQualifiedWildcard(expr, _) => {
self.visit_expr(expr)
}
FunctionArgExpr::QualifiedWildcard(_) | FunctionArgExpr::Wildcard => {}
FunctionArgExpr::QualifiedWildcard(_)
| FunctionArgExpr::WildcardOrWithExcept(None) => {}
FunctionArgExpr::WildcardOrWithExcept(Some(exprs)) => {
for expr in exprs {
self.visit_expr(expr);
}
}
},
}
}
Expand Down Expand Up @@ -346,7 +352,12 @@ impl QueryRewriter<'_> {
SelectItem::UnnamedExpr(expr)
| SelectItem::ExprQualifiedWildcard(expr, _)
| SelectItem::ExprWithAlias { expr, .. } => self.visit_expr(expr),
SelectItem::QualifiedWildcard(_) | SelectItem::Wildcard => {}
SelectItem::QualifiedWildcard(_) | SelectItem::WildcardOrWithExcept(None) => {}
SelectItem::WildcardOrWithExcept(Some(exprs)) => {
for expr in exprs {
self.visit_expr(expr);
}
}
}
}
}
Expand Down
18 changes: 15 additions & 3 deletions src/sqlparser/src/ast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1917,8 +1917,8 @@ pub enum FunctionArgExpr {
ExprQualifiedWildcard(Expr, Vec<Ident>),
/// Qualified wildcard, e.g. `alias.*` or `schema.table.*`.
QualifiedWildcard(ObjectName),
/// An unqualified `*`
Wildcard,
/// An unqualified `*` or `* with (columns)`
WildcardOrWithExcept(Option<Vec<Expr>>),
}

impl fmt::Display for FunctionArgExpr {
Expand All @@ -1936,7 +1936,19 @@ impl fmt::Display for FunctionArgExpr {
)
}
FunctionArgExpr::QualifiedWildcard(prefix) => write!(f, "{}.*", prefix),
FunctionArgExpr::Wildcard => f.write_str("*"),
FunctionArgExpr::WildcardOrWithExcept(w) => match w {
Some(exprs) => write!(
f,
"EXCEPT ({})",
exprs
.iter()
.map(|v| v.to_string())
.collect::<Vec<String>>()
.as_slice()
.join(", ")
),
None => f.write_str("*"),
},
}
}
}
Expand Down
18 changes: 15 additions & 3 deletions src/sqlparser/src/ast/query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -309,8 +309,8 @@ pub enum SelectItem {
ExprWithAlias { expr: Expr, alias: Ident },
/// `alias.*` or even `schema.table.*`
QualifiedWildcard(ObjectName),
/// An unqualified `*`
Wildcard,
/// An unqualified `*`, or `* except (exprs)`
WildcardOrWithExcept(Option<Vec<Expr>>),
}

impl fmt::Display for SelectItem {
Expand All @@ -327,7 +327,19 @@ impl fmt::Display for SelectItem {
.format_with("", |i, f| f(&format_args!(".{i}")))
),
SelectItem::QualifiedWildcard(prefix) => write!(f, "{}.*", prefix),
SelectItem::Wildcard => write!(f, "*"),
SelectItem::WildcardOrWithExcept(w) => match w {
Some(exprs) => write!(
f,
"* EXCEPT ({})",
exprs
.iter()
.map(|v| v.to_string())
.collect::<Vec<String>>()
.as_slice()
.join(", ")
),
None => write!(f, "*"),
},
}
}
}
Expand Down
24 changes: 17 additions & 7 deletions src/sqlparser/src/parser.rs
Original file line number Diff line number Diff line change
Expand Up @@ -82,7 +82,8 @@ pub enum WildcardOrExpr {
/// See also [`Expr::FieldIdentifier`] for behaviors of parentheses.
ExprQualifiedWildcard(Expr, Vec<Ident>),
QualifiedWildcard(ObjectName),
Wildcard,
// Either it's `*` or `* excepts (columns)`
WildcardOrWithExcept(Option<Vec<Expr>>),
}

impl From<WildcardOrExpr> for FunctionArgExpr {
Expand All @@ -93,7 +94,7 @@ impl From<WildcardOrExpr> for FunctionArgExpr {
Self::ExprQualifiedWildcard(expr, prefix)
}
WildcardOrExpr::QualifiedWildcard(prefix) => Self::QualifiedWildcard(prefix),
WildcardOrExpr::Wildcard => Self::Wildcard,
WildcardOrExpr::WildcardOrWithExcept(w) => Self::WildcardOrWithExcept(w),
}
}
}
Expand Down Expand Up @@ -313,7 +314,14 @@ impl Parser {
return self.word_concat_wildcard_expr(w.to_ident()?, wildcard_expr);
}
Token::Mul => {
return Ok(WildcardOrExpr::Wildcard);
if self.parse_keyword(Keyword::EXCEPT) && self.consume_token(&Token::LParen) {
let exprs = self.parse_comma_separated(Parser::parse_expr)?;
if self.consume_token(&Token::RParen) {
return Ok(WildcardOrExpr::WildcardOrWithExcept(Some(exprs)));
}
} else {
return Ok(WildcardOrExpr::WildcardOrWithExcept(None));
}
}
// parses wildcard field selection expression.
// Code is similar to `parse_struct_selection`
Expand Down Expand Up @@ -346,9 +354,10 @@ impl Parser {
let mut idents = vec![ident];
match simple_wildcard_expr {
WildcardOrExpr::QualifiedWildcard(ids) => idents.extend(ids.0),
WildcardOrExpr::Wildcard => {}
WildcardOrExpr::WildcardOrWithExcept(None) => {}
WildcardOrExpr::ExprQualifiedWildcard(_, _) => unreachable!(),
WildcardOrExpr::Expr(e) => return Ok(WildcardOrExpr::Expr(e)),
WildcardOrExpr::WildcardOrWithExcept(Some(_)) => unreachable!(),
}
Ok(WildcardOrExpr::QualifiedWildcard(ObjectName(idents)))
}
Expand Down Expand Up @@ -387,9 +396,10 @@ impl Parser {

match simple_wildcard_expr {
WildcardOrExpr::QualifiedWildcard(ids) => idents.extend(ids.0),
WildcardOrExpr::Wildcard => {}
WildcardOrExpr::WildcardOrWithExcept(None) => {}
WildcardOrExpr::ExprQualifiedWildcard(_, _) => unreachable!(),
WildcardOrExpr::Expr(_) => unreachable!(),
WildcardOrExpr::WildcardOrWithExcept(Some(_)) => unreachable!(),
}
Ok(WildcardOrExpr::ExprQualifiedWildcard(expr, idents))
}
Expand All @@ -408,7 +418,7 @@ impl Parser {
Token::Word(w) => id_parts.push(w.to_ident()?),
Token::Mul => {
return if id_parts.is_empty() {
Ok(WildcardOrExpr::Wildcard)
Ok(WildcardOrExpr::WildcardOrWithExcept(None))
} else {
Ok(WildcardOrExpr::QualifiedWildcard(ObjectName(id_parts)))
}
Expand Down Expand Up @@ -4253,7 +4263,7 @@ impl Parser {
WildcardOrExpr::ExprQualifiedWildcard(expr, prefix) => {
Ok(SelectItem::ExprQualifiedWildcard(expr, prefix))
}
WildcardOrExpr::Wildcard => Ok(SelectItem::Wildcard),
WildcardOrExpr::WildcardOrWithExcept(w) => Ok(SelectItem::WildcardOrWithExcept(w)),
}
}

Expand Down
Loading