add support for order by

yyy1000 · May 28, 2024 · e79b3d5 · e79b3d5
1 parent acb6171
commit e79b3d5
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 0 deletions.
diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs
@@ -282,6 +282,25 @@ impl Unparser<'_> {
                             .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())
                     }
                 };

diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs
@@ -4632,6 +4632,7 @@ fn roundtrip_statement() -> Result<()> {
             "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 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",