feat: support `unnest` in GROUP BY clause #11469

JasonLi-cn · 2024-07-15T04:13:46Z

Which issue does this PR close?

Rationale for this change

support SQLs like:

SELECT unnest(list_type_col) c0, count(1) FROM tab GROUP BY c0;
SELECT unnest(unnest(recursive_list_type_col)) c0, count(1) FROM tab GROUP BY c0;

What changes are included in this PR?

Are these changes tested?

Yes. Has added some test cases into unnest.slt.

Are there any user-facing changes?

Simplified SQL syntax, eg:

SELECT t0.c0, count(1) FROM (SELECT unnest(list_col) c0 FROM tal) t0 GROUP BY t0.c0;

can replace by

SELECT unnest(list_col) c0, count(1) FROM tal GROUP BY c0;

jayzhan211 · 2024-07-15T06:08:47Z

I wonder why we need to specialize group by expression for unnest, does that mean we need to specialize unnest with filter by too? 🤔

JasonLi-cn · 2024-07-15T06:54:59Z

I wonder why we need to specialize group by expression for unnest, does that mean we need to specialize unnest with filter by too? 🤔

Pg supports group by unnest expression;
Our user needs this syntax support.

So, I tried to implement it in the community version.

As far as I know, pg does not currently support unnest in filter, so I'm not sure it needs to be implemented.

JasonLi-cn · 2024-07-15T07:41:48Z

Thank you for the new feature, I have some thoughts for an alternative design:

Recursive aggregate functions are generally unsupported: e.g. select sum(sum(v1)) from t1 is not supported in DataFusion/Postgres/DuckDB. For this unnest() case, a recursive call is required, DuckDB uses: https://duckdb.org/docs/sql/query_syntax/unnest.html#recursive-unnest

What do you think about implementing DuckDB syntax, I think it's better for:

Keep the SQL syntax consistent and avoid nested aggregate function

Avoid changing planner code (which might be error-prone), and keep the recursive unnest logic within the aggregate function body

Thank you for your reply.

1 If "keep the recursive unnest logic within the aggregate function body" is implemented, at least the following work needs to be done:
(1) Improve the logic of Expr::Unnest processing in TypeCoercionRewriter;
(2) Implement create_physical_expr function to support Expr::Unnest;
(3) Maybe there are other tasks...
Also, I would like to check whether you are talking about aggregate function or AggregateExec here. Because pg currently does not support aggregation functions that include unnest. (Of course, it is not appropriate to illustrate this issue only by pg here)

2 If 'Recursive Unnest' is handled like DuckDB, how do we implement the unnest(unnest(list_column)) syntax when the user needs it? By optimization rules?

2010YOUY01 · 2024-07-15T07:53:34Z

Thank you for the new feature, I have some thoughts for an alternative design:
Recursive aggregate functions are generally unsupported: e.g. select sum(sum(v1)) from t1 is not supported in DataFusion/Postgres/DuckDB. For this unnest() case, a recursive call is required, DuckDB uses: https://duckdb.org/docs/sql/query_syntax/unnest.html#recursive-unnest
What do you think about implementing DuckDB syntax, I think it's better for:

Keep the SQL syntax consistent and avoid nested aggregate function

Avoid changing planner code (which might be error-prone), and keep the recursive unnest logic within the aggregate function body

Thank you for your reply.

1 If "keep the recursive unnest logic within the aggregate function body" is implemented, at least the following work needs to be done: (1) Improve the logic of Expr::Unnest processing in TypeCoercionRewriter; (2) Implement create_physical_expr function to support Expr::Unnest; (3) Maybe there are other tasks... Also, I would like to check whether you are talking about aggregate function or AggregateExec here. Because pg currently does not support aggregation functions that include unnest. (Of course, it is not appropriate to illustrate this issue only by pg here)

2 If 'Recursive Unnest' is handled like DuckDB, how do we implement the unnest(unnest(list_column)) syntax when the user needs it? By optimization rules?

@JasonLi-cn I mistakenly thought unnest was some aggregate function... please ignore my previous reply, and also sorry for the confusion.

jayzhan211 · 2024-07-15T12:17:45Z

@JasonLi-cn Is it better to move the logic into try_process_unnest?

alamb

Thank you @JasonLi-cn -- I think this PR looks good to me. Thank you for the contribution

I think adding some comments about what try_process_group_by_unnest is doing and why would help improve this code but I don't think it is required prior to merge

c @jonahgao

datafusion/sql/src/select.rs

datafusion/sqllogictest/test_files/unnest.slt

JasonLi-cn · 2024-07-16T03:27:19Z

Thank you @JasonLi-cn -- I think this PR looks good to me. Thank you for the contribution

I think adding some comments about what try_process_group_by_unnest is doing and why would help improve this code but I don't think it is required prior to merge

c @jonahgao

Thank you @alamb for your review. I have completed the code modification according to your modification suggestions.

jonahgao

I left some comments, but I think it can be merged as is. It works well. Perhaps we can try to improve it through subsequent PRs.

jonahgao · 2024-07-16T08:32:33Z

datafusion/sql/src/select.rs

+        //     Projection: tab.array_col AS unnest(tab.array_col)
+        //       TableScan: tab
+        // ```
+        let mut intermediate_plan = input.clone();


Is it possible to avoid this clone?

I pushed 810ca8b which avoids the cloning in this function

jonahgao · 2024-07-16T08:38:20Z

datafusion/sql/src/select.rs

+
+    /// Try converting Unnest(Expr) of group by to Unnest/Projection
+    /// Return the new input and group_by_exprs of Aggregate.
+    fn try_process_group_by_unnest(


This function seems to have some code repetition with function try_process_unnest.

I wonder if there's a better way to handle this, such as planning unnest before aggregation, and then reusing the current group-by planning logic. This seems more intuitive to me. But I'm not sure about it.

Tracked in #11498

jonahgao · 2024-07-16T08:40:04Z

datafusion/sql/src/select.rs

@@ -297,6 +298,9 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
        input: LogicalPlan,
        select_exprs: Vec<Expr>,
    ) -> Result<LogicalPlan> {
+        // Try process group by unnest
+        let input = self.try_process_aggregate_unnest(input)?;


If unnest has already been processed by try_process_aggregate_unnest, does the following logic for handling unnest become redundant?

I agree with you. We need to verify this logic in the future

Tracked in #11498

alamb · 2024-07-16T21:26:02Z

I merged up to resolve a conflict as well

JasonLi-cn · 2024-07-16T23:28:06Z

I merged up to resolve a conflict as well

Thank you @alamb for your help to modify the code, I learned a lot from it.

jonahgao · 2024-07-17T01:32:48Z

Thanks @JasonLi-cn @alamb

* feat: support group by unnest * pass slt * refactor: mv process_group_by_unnest into try_process_unnest * chore: add some documentation comments and tests * Avoid cloning input * use consistent field names --------- Co-authored-by: Andrew Lamb <[email protected]>

feat: support group by unnest

f45f1b0

github-actions bot added sql SQL Planner sqllogictest SQL Logic Tests (.slt) labels Jul 15, 2024

pass slt

3012936

refactor: mv process_group_by_unnest into try_process_unnest

8d7937d

alamb approved these changes Jul 15, 2024

View reviewed changes

datafusion/sql/src/select.rs Outdated Show resolved Hide resolved

datafusion/sql/src/select.rs Outdated Show resolved Hide resolved

datafusion/sqllogictest/test_files/unnest.slt Show resolved Hide resolved

chore: add some documentation comments and tests

17a5158

jonahgao approved these changes Jul 16, 2024

View reviewed changes

alamb added 2 commits July 16, 2024 16:14

Merge remote-tracking branch 'apache/main' into feature/group_by_unnest

3998224

Avoid cloning input

810ca8b

alamb mentioned this pull request Jul 16, 2024

Reduce repetition in try_process_group_by_unnest and try_process_unnest #11498

Closed

alamb added 2 commits July 16, 2024 16:32

use consistent field names

2309e8b

Merge remote-tracking branch 'apache/main' into feature/group_by_unnest

b37a1e3

alamb changed the title ~~feat: support group by unnest~~ feat: support unnest in GROUP BY clause Jul 16, 2024

jonahgao merged commit a979f3e into apache:main Jul 17, 2024
24 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: support `unnest` in GROUP BY clause #11469

feat: support `unnest` in GROUP BY clause #11469

JasonLi-cn commented Jul 15, 2024 •

edited

Loading

jayzhan211 commented Jul 15, 2024

JasonLi-cn commented Jul 15, 2024

JasonLi-cn commented Jul 15, 2024

2010YOUY01 commented Jul 15, 2024

jayzhan211 commented Jul 15, 2024

alamb left a comment

JasonLi-cn commented Jul 16, 2024

jonahgao left a comment

jonahgao Jul 16, 2024

alamb Jul 16, 2024

jonahgao Jul 16, 2024

alamb Jul 16, 2024

jonahgao Jul 16, 2024

JasonLi-cn Jul 16, 2024

alamb Jul 16, 2024

alamb commented Jul 16, 2024

JasonLi-cn commented Jul 16, 2024

jonahgao commented Jul 17, 2024

feat: support unnest in GROUP BY clause #11469

feat: support unnest in GROUP BY clause #11469

Conversation

JasonLi-cn commented Jul 15, 2024 • edited Loading

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

jayzhan211 commented Jul 15, 2024

JasonLi-cn commented Jul 15, 2024

JasonLi-cn commented Jul 15, 2024

2010YOUY01 commented Jul 15, 2024

jayzhan211 commented Jul 15, 2024

alamb left a comment

Choose a reason for hiding this comment

JasonLi-cn commented Jul 16, 2024

jonahgao left a comment

Choose a reason for hiding this comment

jonahgao Jul 16, 2024

Choose a reason for hiding this comment

alamb Jul 16, 2024

Choose a reason for hiding this comment

jonahgao Jul 16, 2024

Choose a reason for hiding this comment

alamb Jul 16, 2024

Choose a reason for hiding this comment

jonahgao Jul 16, 2024

Choose a reason for hiding this comment

JasonLi-cn Jul 16, 2024

Choose a reason for hiding this comment

alamb Jul 16, 2024

Choose a reason for hiding this comment

alamb commented Jul 16, 2024

JasonLi-cn commented Jul 16, 2024

jonahgao commented Jul 17, 2024

feat: support `unnest` in GROUP BY clause #11469

feat: support `unnest` in GROUP BY clause #11469

JasonLi-cn commented Jul 15, 2024 •

edited

Loading