Try converting Unnest(Expr) of group by to Unnest/Projection. Return the new input and group_by_exprs of Aggregate. Select exprs can be different from agg exprs, for example:
(
&self,
agg: Aggregate,
)
| 598 | /// Return the new input and group_by_exprs of Aggregate. |
| 599 | /// Select exprs can be different from agg exprs, for example: |
| 600 | fn try_process_group_by_unnest( |
| 601 | &self, |
| 602 | agg: Aggregate, |
| 603 | ) -> Result<(LogicalPlan, Vec<Expr>)> { |
| 604 | let mut aggr_expr_using_columns: Option<HashSet<Expr>> = None; |
| 605 | |
| 606 | let Aggregate { |
| 607 | input, |
| 608 | group_expr, |
| 609 | aggr_expr, |
| 610 | .. |
| 611 | } = agg; |
| 612 | |
| 613 | // Process unnest of group_by_exprs, and input of agg will be rewritten |
| 614 | // for example: |
| 615 | // |
| 616 | // ``` |
| 617 | // Aggregate: groupBy=[[UNNEST(Column(Column { relation: Some(Bare { table: "tab" }), name: "array_col" }))]], aggr=[[]] |
| 618 | // TableScan: tab |
| 619 | // ``` |
| 620 | // |
| 621 | // will be transformed into |
| 622 | // |
| 623 | // ``` |
| 624 | // Aggregate: groupBy=[[unnest(tab.array_col)]], aggr=[[]] |
| 625 | // Unnest: lists[unnest(tab.array_col)] structs[] |
| 626 | // Projection: tab.array_col AS unnest(tab.array_col) |
| 627 | // TableScan: tab |
| 628 | // ``` |
| 629 | let mut intermediate_plan = Arc::unwrap_or_clone(input); |
| 630 | let mut intermediate_select_exprs = group_expr; |
| 631 | |
| 632 | loop { |
| 633 | let mut unnest_columns = IndexMap::new(); |
| 634 | let mut inner_projection_exprs = vec![]; |
| 635 | |
| 636 | let outer_projection_exprs = rewrite_recursive_unnests_bottom_up( |
| 637 | &intermediate_plan, |
| 638 | &mut unnest_columns, |
| 639 | &mut inner_projection_exprs, |
| 640 | &intermediate_select_exprs, |
| 641 | )?; |
| 642 | |
| 643 | if unnest_columns.is_empty() { |
| 644 | break; |
| 645 | } else { |
| 646 | let mut unnest_options = UnnestOptions::new().with_preserve_nulls(false); |
| 647 | |
| 648 | #[allow(clippy::allow_attributes, clippy::mutable_key_type)] |
| 649 | // Expr contains Arc with interior mutability but is intentionally used as hash key |
| 650 | let mut projection_exprs = match &aggr_expr_using_columns { |
| 651 | Some(exprs) => (*exprs).clone(), |
| 652 | None => { |
| 653 | #[allow(clippy::allow_attributes, clippy::mutable_key_type)] |
| 654 | let mut columns = HashSet::new(); |
| 655 | for expr in &aggr_expr { |
| 656 | expr.apply(|expr| { |
| 657 | if let Expr::Column(c) = expr { |
no test coverage detected