| 468 | } |
| 469 | |
| 470 | fn simplify( |
| 471 | &self, |
| 472 | args: Vec<Expr>, |
| 473 | _info: &datafusion_expr::simplify::SimplifyContext, |
| 474 | ) -> Result<ExprSimplifyResult> { |
| 475 | // Need at least 2 args (base + field) |
| 476 | if args.len() < 2 { |
| 477 | return Ok(ExprSimplifyResult::Original(args)); |
| 478 | } |
| 479 | |
| 480 | // Flatten all nested get_field calls in a single pass |
| 481 | // Pattern: get_field(get_field(get_field(base, a), b), c) => get_field(base, a, b, c) |
| 482 | |
| 483 | // Collect path arguments from all nested levels |
| 484 | let mut path_args_stack = Vec::new(); |
| 485 | let mut current_expr = &args[0]; |
| 486 | |
| 487 | // Push the outermost path arguments first |
| 488 | path_args_stack.push(&args[1..]); |
| 489 | |
| 490 | // Walk down the chain of nested get_field calls |
| 491 | let base_expr = loop { |
| 492 | if let Expr::ScalarFunction(ScalarFunction { |
| 493 | func, |
| 494 | args: inner_args, |
| 495 | }) = current_expr |
| 496 | && func.inner().is::<GetFieldFunc>() |
| 497 | { |
| 498 | // Store this level's path arguments (all except the first, which is base/nested call) |
| 499 | path_args_stack.push(&inner_args[1..]); |
| 500 | |
| 501 | // Move to the next level down |
| 502 | current_expr = &inner_args[0]; |
| 503 | continue; |
| 504 | } |
| 505 | // Not a get_field call, this is the base expression |
| 506 | break current_expr; |
| 507 | }; |
| 508 | |
| 509 | // If no nested get_field calls were found, return original |
| 510 | if path_args_stack.len() == args.len() - 1 { |
| 511 | return Ok(ExprSimplifyResult::Original(args)); |
| 512 | } |
| 513 | |
| 514 | // If we found any nested get_field calls, flatten them |
| 515 | // Build merged args: [base, ...all_path_args_in_correct_order] |
| 516 | let mut merged_args = vec![base_expr.clone()]; |
| 517 | |
| 518 | // Add path args in reverse order (innermost to outermost) |
| 519 | // Stack is: [outermost_paths, ..., innermost_paths] |
| 520 | // We want: [base, innermost_paths, ..., outermost_paths] |
| 521 | for path_slice in path_args_stack.iter().rev() { |
| 522 | merged_args.extend_from_slice(path_slice); |
| 523 | } |
| 524 | |
| 525 | Ok(ExprSimplifyResult::Simplified(Expr::ScalarFunction( |
| 526 | ScalarFunction::new_udf( |
| 527 | Arc::new(ScalarUDF::new_from_impl(GetFieldFunc::new())), |