diff --git a/datafusion/expr/src/logical_plan/mutate.rs b/datafusion/expr/src/logical_plan/mutate.rs index 9cb507f0ed069..5a0026282122f 100644 --- a/datafusion/expr/src/logical_plan/mutate.rs +++ b/datafusion/expr/src/logical_plan/mutate.rs @@ -19,8 +19,8 @@ use super::plan::*; use crate::expr::{Exists, InSubquery}; use crate::Expr; use datafusion_common::tree_node::Transformed; -use datafusion_common::Result; -use datafusion_common::{DFSchema, DFSchemaRef}; +use datafusion_common::{internal_err, Result}; +use datafusion_common::{Column, DFSchema, DFSchemaRef}; use std::cell::OnceCell; use std::sync::Arc; @@ -106,8 +106,26 @@ impl LogicalPlan { rewrite_expr_iter_mut(filters.iter_mut(), f) } LogicalPlan::Unnest(Unnest { column, .. }) => { - //f(&Expr::Column(column.clone())) - todo!(); + // Unnest doesn't have an Expr to visit, so need to make one + // pick a literal here that is clear it didn't come from + // user query and is easy to grep and find in the source code + // it would be really nice to avoid this and instead have unnest have Expr + + let mut swap_column = Column::new_unqualified("TEMP_unnest_column"); + std::mem::swap(column, &mut swap_column); + + let mut expr = Expr::Column(swap_column); + let result = f(&mut expr)?; + // put the column back + let Expr::Column(mut swap_column) = expr else { + return internal_err!( + "Rewrite of Unnest expr must return Column, returned {:?}", + expr + ); + }; + // put the rewritten column back + std::mem::swap(column, &mut swap_column); + Ok(result) } LogicalPlan::Distinct(Distinct::On(DistinctOn { on_expr,