Skip to content

Commit

Permalink
Support LogicalPlan::Distinct in unparser (apache#10690)
Browse files Browse the repository at this point in the history
* support distinct

* cargo fmt

* better fmt

* add support for order by

* add another test
  • Loading branch information
yyy1000 authored May 28, 2024
1 parent 9fba34d commit 156a525
Show file tree
Hide file tree
Showing 2 changed files with 39 additions and 3 deletions.
39 changes: 36 additions & 3 deletions datafusion/sql/src/unparser/plan.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,9 @@
// under the License.

use datafusion_common::{internal_err, not_impl_err, plan_err, DataFusionError, Result};
use datafusion_expr::{expr::Alias, Expr, JoinConstraint, JoinType, LogicalPlan};
use datafusion_expr::{
expr::Alias, Distinct, Expr, JoinConstraint, JoinType, LogicalPlan,
};
use sqlparser::ast::{self, SetExpr};

use crate::unparser::utils::unproject_agg_exprs;
Expand Down Expand Up @@ -271,8 +273,39 @@ impl Unparser<'_> {
relation,
)
}
LogicalPlan::Distinct(_distinct) => {
not_impl_err!("Unsupported operator: {plan:?}")
LogicalPlan::Distinct(distinct) => {
let (select_distinct, input) = match distinct {
Distinct::All(input) => (ast::Distinct::Distinct, input.as_ref()),
Distinct::On(on) => {
let exprs = on
.on_expr
.iter()
.map(|e| self.expr_to_sql(e))
.collect::<Result<Vec<_>>>()?;
let items = on
.select_expr
.iter()
.map(|e| self.select_item_to_sql(e))
.collect::<Result<Vec<_>>>()?;
match &on.sort_expr {
Some(sort_expr) => {
if let Some(query_ref) = query {
query_ref
.order_by(self.sort_to_sql(sort_expr.clone())?);
} else {
return internal_err!(
"Sort operator only valid in a statement context."
);
}
}
None => {}
}
select.projection(items);
(ast::Distinct::On(exprs), on.input.as_ref())
}
};
select.distinct(Some(select_distinct));
self.select_to_sql_recursively(input, query, select, relation)
}
LogicalPlan::Join(join) => {
match join.join_constraint {
Expand Down
3 changes: 3 additions & 0 deletions datafusion/sql/tests/sql_integration.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4631,6 +4631,9 @@ fn roundtrip_statement() -> Result<()> {
"select (id-1)/2, count(*) / (sum(id/10)-1) as agg_expr from (select (id-1) as id from person) group by id",
"select CAST(id/2 as VARCHAR) NOT LIKE 'foo*' from person where NOT EXISTS (select ta.j1_id, tb.j2_string from j1 ta join j2 tb on (ta.j1_id = tb.j2_id))",
r#"select "First Name" from person_quoted_cols"#,
"select DISTINCT id FROM person",
"select DISTINCT on (id) id, first_name from person",
"select DISTINCT on (id) id, first_name from person order by id",
r#"select id, count("First Name") as cnt from (select id, "First Name" from person_quoted_cols) group by id"#,
"select id, count(*) as cnt from (select p1.id as id from person p1 inner join person p2 on p1.id=p2.id) group by id",
"select id, count(*), first_name from person group by first_name, id",
Expand Down

0 comments on commit 156a525

Please sign in to comment.