implement unnest

apache · Mar 29, 2024 · da634a2 · da634a2
1 parent a17bbb8
commit da634a2
Showing 1 changed file with 22 additions and 4 deletions.
diff --git 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,