Try converting Expr(Unnest(Expr)) to Projection/Unnest/Projection
(
&self,
input: LogicalPlan,
select_exprs: Vec<Expr>,
)
| 437 | |
| 438 | /// Try converting Expr(Unnest(Expr)) to Projection/Unnest/Projection |
| 439 | pub(super) fn try_process_unnest( |
| 440 | &self, |
| 441 | input: LogicalPlan, |
| 442 | select_exprs: Vec<Expr>, |
| 443 | ) -> Result<LogicalPlan> { |
| 444 | // Try process group by unnest |
| 445 | let input = self.try_process_aggregate_unnest(input)?; |
| 446 | |
| 447 | let mut intermediate_plan = input; |
| 448 | let mut intermediate_select_exprs = select_exprs; |
| 449 | // Fast path: If there is are no unnests in the select_exprs, wrap the plan in a projection |
| 450 | if !intermediate_select_exprs |
| 451 | .iter() |
| 452 | .any(has_unnest_expr_recursively) |
| 453 | { |
| 454 | return LogicalPlanBuilder::from(intermediate_plan) |
| 455 | .project(intermediate_select_exprs)? |
| 456 | .build(); |
| 457 | } |
| 458 | |
| 459 | // Each expr in select_exprs can contains multiple unnest stage |
| 460 | // The transformation happen bottom up, one at a time for each iteration |
| 461 | // Only exhaust the loop if no more unnest transformation is found |
| 462 | for i in 0.. { |
| 463 | let mut unnest_columns = IndexMap::new(); |
| 464 | // from which column used for projection, before the unnest happen |
| 465 | // including non unnest column and unnest column |
| 466 | let mut inner_projection_exprs = vec![]; |
| 467 | |
| 468 | // expr returned here maybe different from the originals in inner_projection_exprs |
| 469 | // for example: |
| 470 | // - unnest(struct_col) will be transformed into struct_col.field1, struct_col.field2 |
| 471 | // - unnest(array_col) will be transformed into array_col.element |
| 472 | // - unnest(array_col) + 1 will be transformed into array_col.element +1 |
| 473 | let mut outer_projection_exprs = vec![]; |
| 474 | for expr in &intermediate_select_exprs { |
| 475 | let mut rewritten_exprs = rewrite_recursive_unnest_bottom_up( |
| 476 | &intermediate_plan, |
| 477 | &mut unnest_columns, |
| 478 | &mut inner_projection_exprs, |
| 479 | expr, |
| 480 | )?; |
| 481 | |
| 482 | if let Some(columns) = |
| 483 | self.get_struct_unnest_columns(&intermediate_plan, expr)? |
| 484 | { |
| 485 | rewritten_exprs = rewritten_exprs |
| 486 | .into_iter() |
| 487 | .zip(columns) |
| 488 | .map(|(expr, column)| expr.alias(column.flat_name())) |
| 489 | .collect(); |
| 490 | } |
| 491 | |
| 492 | outer_projection_exprs.extend(rewritten_exprs); |
| 493 | } |
| 494 | |
| 495 | // No more unnest is possible |
| 496 | if unnest_columns.is_empty() { |
no test coverage detected