From cc661ea146af79b24e2d98e7672b2a6592ff83c4 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Sat, 20 Jul 2024 21:52:44 +0200 Subject: [PATCH 01/56] chore: poc --- datafusion/sql/src/select.rs | 17 +- datafusion/sql/src/utils.rs | 181 +++++++++++++++++- datafusion/sqllogictest/test_files/debug.slt | 37 ++++ datafusion/sqllogictest/test_files/unnest.slt | 34 +++- 4 files changed, 261 insertions(+), 8 deletions(-) create mode 100644 datafusion/sqllogictest/test_files/debug.slt diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs index 236403e83d74..96477b625a72 100644 --- a/datafusion/sql/src/select.rs +++ b/datafusion/sql/src/select.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; use crate::planner::{ @@ -23,7 +23,7 @@ use crate::planner::{ }; use crate::utils::{ check_columns_satisfy_exprs, extract_aliases, rebase_expr, resolve_aliases_to_exprs, - resolve_columns, resolve_positions_to_exprs, transform_bottom_unnest, + resolve_columns, resolve_positions_to_exprs, transform_bottom_unnest_v2, }; use datafusion_common::{not_impl_err, plan_err, DataFusionError, Result}; @@ -302,6 +302,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { // Each expr in select_exprs can contains multiple unnest stage // The transformation happen bottom up, one at a time for each iteration // Ony exaust the loop if no more unnest transformation is found + let mut memo = HashMap::new(); for i in 0.. { let mut unnest_columns = vec![]; // from which column used for projection, before the unnest happen @@ -316,10 +317,11 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { let outer_projection_exprs: Vec = intermediate_select_exprs .iter() .map(|expr| { - transform_bottom_unnest( + transform_bottom_unnest_v2( &intermediate_plan, &mut unnest_columns, &mut inner_projection_exprs, + &mut memo, expr, ) }) @@ -341,10 +343,19 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { let columns = unnest_columns.into_iter().map(|col| col.into()).collect(); // Set preserve_nulls to false to ensure compatibility with DuckDB and PostgreSQL let unnest_options = UnnestOptions::new().with_preserve_nulls(false); + // deduplicate expr in inner_projection_exprs + inner_projection_exprs.dedup_by(|a, b| -> bool { + a.display_name().unwrap() == b.display_name().unwrap() + }); + let plan = LogicalPlanBuilder::from(intermediate_plan) .project(inner_projection_exprs)? .unnest_columns_with_options(columns, unnest_options)? .build()?; + println!( + "intermediate plan {:?}\n, selected exprs {:?}", + plan, outer_projection_exprs + ); intermediate_plan = plan; intermediate_select_exprs = outer_projection_exprs; } diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index 2eacbd174fc2..0de4517b9443 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -17,6 +17,7 @@ //! SQL Utility Functions +use std::cell::RefCell; use std::collections::HashMap; use arrow_schema::{ @@ -263,6 +264,117 @@ pub(crate) fn normalize_ident(id: Ident) -> String { } } +pub(crate) fn transform_bottom_unnest_v2( + input: &LogicalPlan, + unnest_placeholder_columns: &mut Vec, + inner_projection_exprs: &mut Vec, + memo: &mut HashMap>, + original_expr: &Expr, +) -> Result> { + // TODO: depth first search to rewrite the bottom most unnest expressions + // until all the leave branches having at least one transformation + let mut transform = + |unnest_expr: &Expr, expr_in_unnest: &Expr| -> Result> { + if let Some(previou_transformed) = memo.get(unnest_expr) { + return Ok(previou_transformed.clone()); + } + // Full context, we are trying to plan the execution as InnerProjection->Unnest->OuterProjection + // inside unnest execution, each column inside the inner projection + // will be transformed into new columns. Thus we need to keep track of these placeholding column names + let placeholder_name = unnest_expr.display_name()?; + + unnest_placeholder_columns.push(placeholder_name.clone()); + // Add alias for the argument expression, to avoid naming conflicts + // with other expressions in the select list. For example: `select unnest(col1), col1 from t`. + // this extra projection is used to unnest transforming + inner_projection_exprs + .push(expr_in_unnest.clone().alias(placeholder_name.clone())); + let schema = input.schema(); + + let (data_type, _) = expr_in_unnest.data_type_and_nullable(schema)?; + + let outer_projection_columns = + get_unnested_columns(&placeholder_name, &data_type)?; + let expr = outer_projection_columns + .iter() + .map(|col| Expr::Column(col.0.clone())) + .collect::>(); + + memo.insert(unnest_expr.clone(), expr.clone()); + Ok(expr) + }; + // let mut stack = vec![]; + // let mut unnest_visitted = HashSet::new(); + let latest_visited = RefCell::new(None); + // we need to mark only the latest unnest expr that was visitted during the down traversal + let transform_down = |expr: Expr| -> Result> { + if let Expr::Unnest(Unnest { .. }) = expr { + *latest_visited.borrow_mut() = Some(expr.clone()); + Ok(Transformed::no(expr)) + } else { + Ok(Transformed::no(expr)) + } + }; + let transform_up = |expr: Expr| -> Result> { + if let Expr::Unnest(Unnest { expr: ref arg }) = expr { + // only transform the bottom most unnest expr + if let Some(ref mut last_visitted_expr) = *latest_visited.borrow_mut() { + if last_visitted_expr != &expr { + return Ok(Transformed::no(expr)); + } + // this is (one of) the bottom most unnest expr + let (data_type, _) = arg.data_type_and_nullable(input.schema())?; + if &expr != original_expr { + if let DataType::Struct(_) = data_type { + return internal_err!("unnest on struct can ony be applied at the root level of select expression"); + } + } + + let mut transformed_exprs = transform(&expr, arg)?; + // root_expr.push(transformed_exprs[0].clone()); + return Ok(Transformed::new( + transformed_exprs.swap_remove(0), + true, + TreeNodeRecursion::Continue, + )); + } + } + Ok(Transformed::no(expr)) + }; + + // This transformation is only done for list unnest + // struct unnest is done at the root level, and at the later stage + // because the syntax of TreeNode only support transform into 1 Expr, while + // Unnest struct will be transformed into multiple Exprs + // TODO: This can be resolved after this issue is resolved: https://github.com/apache/datafusion/issues/10102 + // + // The transformation looks like: + // - unnest(array_col) will be transformed into unnest(array_col) + // - unnest(array_col) + 1 will be transformed into unnest(array_col) + 1 + let Transformed { + data: transformed_expr, + transformed, + tnr: _, + } = original_expr + .clone() + .transform_down_up(transform_down, transform_up)?; + + if !transformed { + if matches!(&transformed_expr, Expr::Column(_)) { + inner_projection_exprs.push(transformed_expr.clone()); + Ok(vec![transformed_expr]) + } else { + // We need to evaluate the expr in the inner projection, + // outer projection just select its name + let column_name = transformed_expr.display_name()?; + inner_projection_exprs.push(transformed_expr); + Ok(vec![Expr::Column(Column::from_name(column_name))]) + } + } else { + Ok(vec![transformed_expr]) + } +} + /// The context is we want to rewrite unnest() into InnerProjection->Unnest->OuterProjection /// Given an expression which contains unnest expr as one of its children, /// Try transform depends on unnest type @@ -367,7 +479,7 @@ pub(crate) fn transform_bottom_unnest( // write test for recursive_transform_unnest #[cfg(test)] mod tests { - use std::{ops::Add, sync::Arc}; + use std::{collections::HashMap, ops::Add, sync::Arc}; use arrow::datatypes::{DataType as ArrowDataType, Field, Schema}; use arrow_schema::Fields; @@ -376,7 +488,72 @@ mod tests { use datafusion_functions::core::expr_ext::FieldAccessor; use datafusion_functions_aggregate::expr_fn::count; - use crate::utils::{resolve_positions_to_exprs, transform_bottom_unnest}; + use crate::utils::{ + resolve_positions_to_exprs, transform_bottom_unnest, transform_bottom_unnest_v2, + }; + + #[test] + fn test_transform_bottom_unnest_v2() -> Result<()> { + let schema = Schema::new(vec![ + Field::new( + "struct_col", + ArrowDataType::Struct(Fields::from(vec![ + Field::new("field1", ArrowDataType::Int32, false), + Field::new("field2", ArrowDataType::Int32, false), + ])), + false, + ), + Field::new( + "3d_col", + ArrowDataType::List(Arc::new(Field::new( + "2d_col", + ArrowDataType::List(Arc::new(Field::new( + "elements", + ArrowDataType::Int64, + true, + ))), + true, + ))), + true, + ), + Field::new("int_col", ArrowDataType::Int32, false), + ]); + + let dfschema = DFSchema::try_from(schema)?; + + let input = LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: Arc::new(dfschema), + }); + + let mut unnest_placeholder_columns = vec![]; + let mut inner_projection_exprs = vec![]; + + // unnest(struct_col) + let original_expr = + unnest(unnest(col("3d_col"))).add(unnest(unnest(col("3d_col")))); + let mut memo = HashMap::new(); + let transformed_exprs = transform_bottom_unnest_v2( + &input, + &mut unnest_placeholder_columns, + &mut inner_projection_exprs, + &mut memo, + &original_expr, + )?; + assert_eq!( + transformed_exprs, + vec![unnest(col("unnest(3d_col)")).add(unnest(col("unnest(3d_col)")))] + ); + assert_eq!(unnest_placeholder_columns, vec!["unnest(array_col)"]); + // still reference struct_col in original schema but with alias, + // to avoid colliding with the projection on the column itself if any + assert_eq!( + inner_projection_exprs, + vec![col("array_col").alias("unnest(array_col)"),] + ); + + Ok(()) + } #[test] fn test_transform_bottom_unnest() -> Result<()> { diff --git a/datafusion/sqllogictest/test_files/debug.slt b/datafusion/sqllogictest/test_files/debug.slt new file mode 100644 index 000000000000..a2728b76976d --- /dev/null +++ b/datafusion/sqllogictest/test_files/debug.slt @@ -0,0 +1,37 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +############################ +# Unnest Expressions Tests # +############################ + +statement ok +CREATE TABLE unnest_table +AS VALUES + ([[1,2,3]], [struct([1,2,3],[4,5,6])], [13, 14], struct(1,2)), + -- null array to verify the `preserve_nulls` option + (null, null, [17, 18], null) +; + +query II +select unnest(unnest(column1)) + unnest(unnest(column1)), unnest(unnest(column2)['c0']) + unnest(unnest(column2)['c1']) from unnest_table; +---- +2 [1, 2, 3] +4 [1, 2, 3] +6 [1, 2, 3] + + diff --git a/datafusion/sqllogictest/test_files/unnest.slt b/datafusion/sqllogictest/test_files/unnest.slt index 06733f7b1e40..ccdfb139727c 100644 --- a/datafusion/sqllogictest/test_files/unnest.slt +++ b/datafusion/sqllogictest/test_files/unnest.slt @@ -165,6 +165,9 @@ select unnest(column1), column1 from unnest_table; 6 [6] 12 [12] +query I? +select unnest(column1) + unnest(column1) from unnest_table; + ## unnest as children of other expr query I? select unnest(column1) + 1 , column1 from unnest_table; @@ -500,6 +503,30 @@ select unnest(column1) from (select * from (values([1,2,3]), ([4,5,6])) limit 1 query error DataFusion error: Error during planning: Projections require unique expression names but the expression "UNNEST\(Column\(Column \{ relation: Some\(Bare \{ table: "unnest_table" \}\), name: "column1" \}\)\)" at position 0 and "UNNEST\(Column\(Column \{ relation: Some\(Bare \{ table: "unnest_table" \}\), name: "column1" \}\)\)" at position 1 have the same name. Consider aliasing \("AS"\) one of them. select unnest(column1), unnest(column1) from unnest_table; + +## the same unnest expr is referened multiple times +query ??II +select unnest(column2), unnest(unnest(column2)), unnest(unnest(unnest(column2))), unnest(unnest(unnest(column2))) + 1 from recursive_unnest_table; +---- +[[1], [2]] [1] 1 2 +[[1], [2]] [2] 2 3 +[[1, 1]] [1, 1] 1 2 +[[1, 1]] [1, 1] 1 2 +[[3, 4], [5]] [3, 4] 3 4 +[[3, 4], [5]] [3, 4] 4 5 +[[3, 4], [5]] [5] 5 6 +[[, 6], , [7, 8]] [, 6] NULL NULL +[[, 6], , [7, 8]] [, 6] 6 7 +[[, 6], , [7, 8]] [7, 8] 7 8 +[[, 6], , [7, 8]] [7, 8] 8 9 + + +query ??II +select unnest(column3), unnest(column3)['c0'], unnest(unnest(column3)['c0']), unnest(unnest(column3)['c0']) + unnest(unnest(column3)['c0']) from recursive_unnest_table; + +## (struct([1], 'a'), [[[1],[2]],[[1,1]]], [struct([1],[[1,2]])]), + + statement ok drop table unnest_table; @@ -537,12 +564,13 @@ NULL [[[3, 4], [5]], [[, 6], , [7, 8]]] + query TT -explain select unnest(unnest(unnest(column3)['c1'])), column3 from recursive_unnest_table; +explain select unnest(unnest(unnest(column3)['c1'])), unnest(unnest(column3)['c1']), column3 from recursive_unnest_table; ---- logical_plan 01)Unnest: lists[unnest(unnest(unnest(recursive_unnest_table.column3)[c1]))] structs[] -02)--Projection: unnest(unnest(recursive_unnest_table.column3)[c1]) AS unnest(unnest(unnest(recursive_unnest_table.column3)[c1])), recursive_unnest_table.column3 +02)--Projection: unnest(unnest(recursive_unnest_table.column3)[c1]) AS unnest(unnest(unnest(recursive_unnest_table.column3)[c1])), unnest(unnest(recursive_unnest_table.column3)[c1]), recursive_unnest_table.column3 03)----Unnest: lists[unnest(unnest(recursive_unnest_table.column3)[c1])] structs[] 04)------Projection: get_field(unnest(recursive_unnest_table.column3), Utf8("c1")) AS unnest(unnest(recursive_unnest_table.column3)[c1]), recursive_unnest_table.column3 05)--------Unnest: lists[unnest(recursive_unnest_table.column3)] structs[] @@ -550,7 +578,7 @@ logical_plan 07)------------TableScan: recursive_unnest_table projection=[column3] physical_plan 01)UnnestExec -02)--ProjectionExec: expr=[unnest(unnest(recursive_unnest_table.column3)[c1])@0 as unnest(unnest(unnest(recursive_unnest_table.column3)[c1])), column3@1 as column3] +02)--ProjectionExec: expr=[unnest(unnest(recursive_unnest_table.column3)[c1])@0 as unnest(unnest(unnest(recursive_unnest_table.column3)[c1])), unnest(unnest(recursive_unnest_table.column3)[c1])@0 as unnest(unnest(recursive_unnest_table.column3)[c1]), column3@1 as column3] 03)----UnnestExec 04)------ProjectionExec: expr=[get_field(unnest(recursive_unnest_table.column3)@0, c1) as unnest(unnest(recursive_unnest_table.column3)[c1]), column3@1 as column3] 05)--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 From d7d45b131772bce5456d09eb604fab24d5c3b845 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Sun, 21 Jul 2024 11:35:34 +0200 Subject: [PATCH 02/56] fix unnest struct --- datafusion/sql/src/select.rs | 9 ++--- datafusion/sql/src/utils.rs | 32 ++++++++++++----- datafusion/sqllogictest/test_files/debug.slt | 37 -------------------- 3 files changed, 25 insertions(+), 53 deletions(-) delete mode 100644 datafusion/sqllogictest/test_files/debug.slt diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs index 96477b625a72..7607e4c30b47 100644 --- a/datafusion/sql/src/select.rs +++ b/datafusion/sql/src/select.rs @@ -299,9 +299,8 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { ) -> Result { let mut intermediate_plan = input; let mut intermediate_select_exprs = select_exprs; - // Each expr in select_exprs can contains multiple unnest stage - // The transformation happen bottom up, one at a time for each iteration - // Ony exaust the loop if no more unnest transformation is found + + // impl memoization to store all previous unnest transformation let mut memo = HashMap::new(); for i in 0.. { let mut unnest_columns = vec![]; @@ -352,10 +351,6 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { .project(inner_projection_exprs)? .unnest_columns_with_options(columns, unnest_options)? .build()?; - println!( - "intermediate plan {:?}\n, selected exprs {:?}", - plan, outer_projection_exprs - ); intermediate_plan = plan; intermediate_select_exprs = outer_projection_exprs; } diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index 0de4517b9443..c9766266a7ba 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -264,6 +264,14 @@ pub(crate) fn normalize_ident(id: Ident) -> String { } } +/// The context is we want to rewrite unnest() into InnerProjection->Unnest->OuterProjection +/// Given an expression which contains unnest expr as one of its children, +/// Try transform depends on unnest type +/// - For list column: unnest(col) with type list -> unnest(col) with type list::item +/// - For struct column: unnest(struct(field1, field2)) -> unnest(struct).field1, unnest(struct).field2 +/// The transformed exprs will be used in the outer projection +/// If along the path from root to bottom, there are multiple unnest expressions, the transformation +/// is done only for the bottom expression pub(crate) fn transform_bottom_unnest_v2( input: &LogicalPlan, unnest_placeholder_columns: &mut Vec, @@ -271,8 +279,6 @@ pub(crate) fn transform_bottom_unnest_v2( memo: &mut HashMap>, original_expr: &Expr, ) -> Result> { - // TODO: depth first search to rewrite the bottom most unnest expressions - // until all the leave branches having at least one transformation let mut transform = |unnest_expr: &Expr, expr_in_unnest: &Expr| -> Result> { if let Some(previou_transformed) = memo.get(unnest_expr) { @@ -303,8 +309,6 @@ pub(crate) fn transform_bottom_unnest_v2( memo.insert(unnest_expr.clone(), expr.clone()); Ok(expr) }; - // let mut stack = vec![]; - // let mut unnest_visitted = HashSet::new(); let latest_visited = RefCell::new(None); // we need to mark only the latest unnest expr that was visitted during the down traversal let transform_down = |expr: Expr| -> Result> { @@ -317,17 +321,20 @@ pub(crate) fn transform_bottom_unnest_v2( }; let transform_up = |expr: Expr| -> Result> { if let Expr::Unnest(Unnest { expr: ref arg }) = expr { - // only transform the bottom most unnest expr + // only transform the first unnest expr(s) from the bottom up + // if the expr tree contains mulitple unnest exprs, as long as neither of them + // is the direct ancestor of one another, we do all the transformation if let Some(ref mut last_visitted_expr) = *latest_visited.borrow_mut() { if last_visitted_expr != &expr { return Ok(Transformed::no(expr)); } // this is (one of) the bottom most unnest expr let (data_type, _) = arg.data_type_and_nullable(input.schema())?; - if &expr != original_expr { - if let DataType::Struct(_) = data_type { - return internal_err!("unnest on struct can ony be applied at the root level of select expression"); - } + if &expr == original_expr { + return Ok(Transformed::no(expr)); + } + if let DataType::Struct(_) = data_type { + return internal_err!("unnest on struct can ony be applied at the root level of select expression"); } let mut transformed_exprs = transform(&expr, arg)?; @@ -360,6 +367,13 @@ pub(crate) fn transform_bottom_unnest_v2( .transform_down_up(transform_down, transform_up)?; if !transformed { + // Because root expr need to transform separately + // unnest struct is only possible here + // The transformation looks like + // - unnest(struct_col) will be transformed into unnest(struct_col).field1, unnest(struct_col).field2 + if let Expr::Unnest(Unnest { expr: ref arg }) = transformed_expr { + return transform(&transformed_expr, arg); + } if matches!(&transformed_expr, Expr::Column(_)) { inner_projection_exprs.push(transformed_expr.clone()); Ok(vec![transformed_expr]) diff --git a/datafusion/sqllogictest/test_files/debug.slt b/datafusion/sqllogictest/test_files/debug.slt deleted file mode 100644 index a2728b76976d..000000000000 --- a/datafusion/sqllogictest/test_files/debug.slt +++ /dev/null @@ -1,37 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at - -# http://www.apache.org/licenses/LICENSE-2.0 - -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. - -############################ -# Unnest Expressions Tests # -############################ - -statement ok -CREATE TABLE unnest_table -AS VALUES - ([[1,2,3]], [struct([1,2,3],[4,5,6])], [13, 14], struct(1,2)), - -- null array to verify the `preserve_nulls` option - (null, null, [17, 18], null) -; - -query II -select unnest(unnest(column1)) + unnest(unnest(column1)), unnest(unnest(column2)['c0']) + unnest(unnest(column2)['c1']) from unnest_table; ----- -2 [1, 2, 3] -4 [1, 2, 3] -6 [1, 2, 3] - - From f8aa23a7e31c689c6878af7eca9b6ae2d2092ef8 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Sun, 21 Jul 2024 13:04:15 +0200 Subject: [PATCH 03/56] UT for memoization --- datafusion/sql/src/select.rs | 8 +- datafusion/sql/src/utils.rs | 233 ++++++++---------- datafusion/sqllogictest/test_files/unnest.slt | 18 +- 3 files changed, 118 insertions(+), 141 deletions(-) diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs index 7607e4c30b47..ea372fafa5a0 100644 --- a/datafusion/sql/src/select.rs +++ b/datafusion/sql/src/select.rs @@ -23,7 +23,7 @@ use crate::planner::{ }; use crate::utils::{ check_columns_satisfy_exprs, extract_aliases, rebase_expr, resolve_aliases_to_exprs, - resolve_columns, resolve_positions_to_exprs, transform_bottom_unnest_v2, + resolve_columns, resolve_positions_to_exprs, transform_bottom_unnest, }; use datafusion_common::{not_impl_err, plan_err, DataFusionError, Result}; @@ -316,7 +316,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { let outer_projection_exprs: Vec = intermediate_select_exprs .iter() .map(|expr| { - transform_bottom_unnest_v2( + transform_bottom_unnest( &intermediate_plan, &mut unnest_columns, &mut inner_projection_exprs, @@ -343,9 +343,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { // Set preserve_nulls to false to ensure compatibility with DuckDB and PostgreSQL let unnest_options = UnnestOptions::new().with_preserve_nulls(false); // deduplicate expr in inner_projection_exprs - inner_projection_exprs.dedup_by(|a, b| -> bool { - a.display_name().unwrap() == b.display_name().unwrap() - }); + inner_projection_exprs.dedup_by(|a, b| -> bool { a == b }); let plan = LogicalPlanBuilder::from(intermediate_plan) .project(inner_projection_exprs)? diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index c9766266a7ba..f2da2710556e 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -272,7 +272,7 @@ pub(crate) fn normalize_ident(id: Ident) -> String { /// The transformed exprs will be used in the outer projection /// If along the path from root to bottom, there are multiple unnest expressions, the transformation /// is done only for the bottom expression -pub(crate) fn transform_bottom_unnest_v2( +pub(crate) fn transform_bottom_unnest( input: &LogicalPlan, unnest_placeholder_columns: &mut Vec, inner_projection_exprs: &mut Vec, @@ -389,107 +389,6 @@ pub(crate) fn transform_bottom_unnest_v2( } } -/// The context is we want to rewrite unnest() into InnerProjection->Unnest->OuterProjection -/// Given an expression which contains unnest expr as one of its children, -/// Try transform depends on unnest type -/// - For list column: unnest(col) with type list -> unnest(col) with type list::item -/// - For struct column: unnest(struct(field1, field2)) -> unnest(struct).field1, unnest(struct).field2 -/// The transformed exprs will be used in the outer projection -/// If along the path from root to bottom, there are multiple unnest expressions, the transformation -/// is done only for the bottom expression -pub(crate) fn transform_bottom_unnest( - input: &LogicalPlan, - unnest_placeholder_columns: &mut Vec, - inner_projection_exprs: &mut Vec, - original_expr: &Expr, -) -> Result> { - let mut transform = - |unnest_expr: &Expr, expr_in_unnest: &Expr| -> Result> { - // Full context, we are trying to plan the execution as InnerProjection->Unnest->OuterProjection - // inside unnest execution, each column inside the inner projection - // will be transformed into new columns. Thus we need to keep track of these placeholding column names - let placeholder_name = unnest_expr.display_name()?; - - unnest_placeholder_columns.push(placeholder_name.clone()); - // Add alias for the argument expression, to avoid naming conflicts - // with other expressions in the select list. For example: `select unnest(col1), col1 from t`. - // this extra projection is used to unnest transforming - inner_projection_exprs - .push(expr_in_unnest.clone().alias(placeholder_name.clone())); - let schema = input.schema(); - - let (data_type, _) = expr_in_unnest.data_type_and_nullable(schema)?; - - let outer_projection_columns = - get_unnested_columns(&placeholder_name, &data_type)?; - let expr = outer_projection_columns - .iter() - .map(|col| Expr::Column(col.0.clone())) - .collect::>(); - Ok(expr) - }; - // This transformation is only done for list unnest - // struct unnest is done at the root level, and at the later stage - // because the syntax of TreeNode only support transform into 1 Expr, while - // Unnest struct will be transformed into multiple Exprs - // TODO: This can be resolved after this issue is resolved: https://github.com/apache/datafusion/issues/10102 - // - // The transformation looks like: - // - unnest(array_col) will be transformed into unnest(array_col) - // - unnest(array_col) + 1 will be transformed into unnest(array_col) + 1 - let Transformed { - data: transformed_expr, - transformed, - tnr: _, - } = original_expr.clone().transform_up(|expr: Expr| { - let is_root_expr = &expr == original_expr; - // Root expr is transformed separately - if is_root_expr { - return Ok(Transformed::no(expr)); - } - if let Expr::Unnest(Unnest { expr: ref arg }) = expr { - let (data_type, _) = arg.data_type_and_nullable(input.schema())?; - - if let DataType::Struct(_) = data_type { - return internal_err!("unnest on struct can ony be applied at the root level of select expression"); - } - - let mut transformed_exprs = transform(&expr, arg)?; - // root_expr.push(transformed_exprs[0].clone()); - Ok(Transformed::new( - transformed_exprs.swap_remove(0), - true, - TreeNodeRecursion::Stop, - )) - } else { - Ok(Transformed::no(expr)) - } - })?; - - if !transformed { - // Because root expr need to transform separately - // unnest struct is only possible here - // The transformation looks like - // - unnest(struct_col) will be transformed into unnest(struct_col).field1, unnest(struct_col).field2 - if let Expr::Unnest(Unnest { expr: ref arg }) = transformed_expr { - return transform(&transformed_expr, arg); - } - - if matches!(&transformed_expr, Expr::Column(_)) { - inner_projection_exprs.push(transformed_expr.clone()); - Ok(vec![transformed_expr]) - } else { - // We need to evaluate the expr in the inner projection, - // outer projection just select its name - let column_name = transformed_expr.display_name()?; - inner_projection_exprs.push(transformed_expr); - Ok(vec![Expr::Column(Column::from_name(column_name))]) - } - } else { - Ok(vec![transformed_expr]) - } -} - // write test for recursive_transform_unnest #[cfg(test)] mod tests { @@ -497,41 +396,30 @@ mod tests { use arrow::datatypes::{DataType as ArrowDataType, Field, Schema}; use arrow_schema::Fields; - use datafusion_common::{DFSchema, Result}; - use datafusion_expr::{col, lit, unnest, EmptyRelation, LogicalPlan}; + use datafusion_common::{DFSchema, Result, UnnestOptions}; + use datafusion_expr::{ + col, lit, unnest, EmptyRelation, LogicalPlan, LogicalPlanBuilder, + }; use datafusion_functions::core::expr_ext::FieldAccessor; use datafusion_functions_aggregate::expr_fn::count; - use crate::utils::{ - resolve_positions_to_exprs, transform_bottom_unnest, transform_bottom_unnest_v2, - }; + use crate::utils::{resolve_positions_to_exprs, transform_bottom_unnest}; #[test] - fn test_transform_bottom_unnest_v2() -> Result<()> { - let schema = Schema::new(vec![ - Field::new( - "struct_col", - ArrowDataType::Struct(Fields::from(vec![ - Field::new("field1", ArrowDataType::Int32, false), - Field::new("field2", ArrowDataType::Int32, false), - ])), - false, - ), - Field::new( - "3d_col", + fn test_transform_bottom_unnest_recursive_memoization() -> Result<()> { + let schema = Schema::new(vec![Field::new( + "3d_col", + ArrowDataType::List(Arc::new(Field::new( + "2d_col", ArrowDataType::List(Arc::new(Field::new( - "2d_col", - ArrowDataType::List(Arc::new(Field::new( - "elements", - ArrowDataType::Int64, - true, - ))), + "elements", + ArrowDataType::Int64, true, ))), true, - ), - Field::new("int_col", ArrowDataType::Int32, false), - ]); + ))), + true, + )]); let dfschema = DFSchema::try_from(schema)?; @@ -543,27 +431,102 @@ mod tests { let mut unnest_placeholder_columns = vec![]; let mut inner_projection_exprs = vec![]; - // unnest(struct_col) + // unnest(unnest(3d_col)) + unnest(unnest(3d_col)) let original_expr = unnest(unnest(col("3d_col"))).add(unnest(unnest(col("3d_col")))); let mut memo = HashMap::new(); - let transformed_exprs = transform_bottom_unnest_v2( + let transformed_exprs = transform_bottom_unnest( &input, &mut unnest_placeholder_columns, &mut inner_projection_exprs, &mut memo, &original_expr, )?; + // only the bottom most unnest exprs are transformed assert_eq!( transformed_exprs, vec![unnest(col("unnest(3d_col)")).add(unnest(col("unnest(3d_col)")))] ); - assert_eq!(unnest_placeholder_columns, vec!["unnest(array_col)"]); + // memoization only contains 1 transformation + assert_eq!(memo.len(), 1); + assert_eq!( + memo.get(&unnest(col("3d_col"))), + Some(&vec![col("unnest(3d_col)")]) + ); + assert_eq!(unnest_placeholder_columns, vec!["unnest(3d_col)"]); + // still reference struct_col in original schema but with alias, + // to avoid colliding with the projection on the column itself if any + assert_eq!( + inner_projection_exprs, + vec![col("3d_col").alias("unnest(3d_col)"),] + ); + + // unnest(3d_col) as 2d_col + let original_expr_2 = unnest(col("3d_col")).alias("2d_col"); + let transformed_exprs = transform_bottom_unnest( + &input, + &mut unnest_placeholder_columns, + &mut inner_projection_exprs, + &mut memo, + &original_expr_2, + )?; + + assert_eq!( + transformed_exprs, + vec![col("unnest(3d_col)").alias("2d_col")] + ); + // memoization still contains 1 transformation + // and the previous transformation is reused + assert_eq!(memo.len(), 1); + assert_eq!( + memo.get(&unnest(col("3d_col"))), + Some(&vec![col("unnest(3d_col)")]) + ); + assert_eq!(unnest_placeholder_columns, vec!["unnest(3d_col)"]); // still reference struct_col in original schema but with alias, // to avoid colliding with the projection on the column itself if any assert_eq!( inner_projection_exprs, - vec![col("array_col").alias("unnest(array_col)"),] + vec![col("3d_col").alias("unnest(3d_col)")] + ); + + // Start a new cycle, to run unnest again on previous transformation + let intermediate_columns = unnest_placeholder_columns + .into_iter() + .map(|col| col.into()) + .collect(); + let intermediate_input = LogicalPlanBuilder::from(input) + .project(inner_projection_exprs)? + .unnest_columns_with_options(intermediate_columns, UnnestOptions::default())? + .build()?; + + let mut new_unnest_placeholder_columns = vec![]; + let mut new_inner_projection_exprs = vec![]; + // Run unnest again on previously transformed expr + let transformed_exprs = transform_bottom_unnest( + &intermediate_input, + &mut new_unnest_placeholder_columns, + &mut new_inner_projection_exprs, + &mut memo, + &unnest(col("unnest(3d_col)")).add(unnest(col("unnest(3d_col)"))), + )?; + assert_eq!( + transformed_exprs, + vec![col("unnest(unnest(3d_col))").add(col("unnest(unnest(3d_col))"))], + ); + // memoization having extra transformation cached + assert_eq!(memo.len(), 2); + assert_eq!( + memo.get(&unnest(col("unnest(3d_col)"))), + Some(&vec![col("unnest(unnest(3d_col))")]) + ); + assert_eq!( + new_unnest_placeholder_columns, + vec!["unnest(unnest(3d_col))"] + ); + assert_eq!( + new_inner_projection_exprs, + vec![col("unnest(3d_col)").alias("unnest(unnest(3d_col))")] ); Ok(()) @@ -602,12 +565,14 @@ mod tests { let mut unnest_placeholder_columns = vec![]; let mut inner_projection_exprs = vec![]; + let mut memo = HashMap::new(); // unnest(struct_col) let original_expr = unnest(col("struct_col")); let transformed_exprs = transform_bottom_unnest( &input, &mut unnest_placeholder_columns, &mut inner_projection_exprs, + &mut memo, &original_expr, )?; assert_eq!( @@ -625,12 +590,14 @@ mod tests { vec![col("struct_col").alias("unnest(struct_col)"),] ); + memo.clear(); // unnest(array_col) + 1 let original_expr = unnest(col("array_col")).add(lit(1i64)); let transformed_exprs = transform_bottom_unnest( &input, &mut unnest_placeholder_columns, &mut inner_projection_exprs, + &mut memo, &original_expr, )?; assert_eq!( @@ -685,6 +652,7 @@ mod tests { let mut unnest_placeholder_columns = vec![]; let mut inner_projection_exprs = vec![]; + memo.clear(); // An expr with multiple unnest let original_expr = unnest(unnest(col("struct_col").field("matrix"))); @@ -692,6 +660,7 @@ mod tests { &input, &mut unnest_placeholder_columns, &mut inner_projection_exprs, + &mut memo, &original_expr, )?; // Only the inner most/ bottom most unnest is transformed diff --git a/datafusion/sqllogictest/test_files/unnest.slt b/datafusion/sqllogictest/test_files/unnest.slt index ccdfb139727c..4145cb02a84e 100644 --- a/datafusion/sqllogictest/test_files/unnest.slt +++ b/datafusion/sqllogictest/test_files/unnest.slt @@ -165,8 +165,16 @@ select unnest(column1), column1 from unnest_table; 6 [6] 12 [12] -query I? +query I select unnest(column1) + unnest(column1) from unnest_table; +---- +2 +4 +6 +8 +10 +12 +24 ## unnest as children of other expr query I? @@ -504,7 +512,7 @@ query error DataFusion error: Error during planning: Projections require unique select unnest(column1), unnest(column1) from unnest_table; -## the same unnest expr is referened multiple times +## the same unnest expr is referened multiple times (unnest is the bottom-most expr) query ??II select unnest(column2), unnest(unnest(column2)), unnest(unnest(unnest(column2))), unnest(unnest(unnest(column2))) + 1 from recursive_unnest_table; ---- @@ -521,10 +529,12 @@ select unnest(column2), unnest(unnest(column2)), unnest(unnest(unnest(column2))) [[, 6], , [7, 8]] [7, 8] 8 9 +## the same unnest expr is referened multiple times (unnest is not the bottom-most expr) query ??II select unnest(column3), unnest(column3)['c0'], unnest(unnest(column3)['c0']), unnest(unnest(column3)['c0']) + unnest(unnest(column3)['c0']) from recursive_unnest_table; - -## (struct([1], 'a'), [[[1],[2]],[[1,1]]], [struct([1],[[1,2]])]), +---- +{c0: [1], c1: [[1, 2]]} [1] 1 2 +{c0: [2], c1: [[3], [4]]} [2] 2 4 statement ok From cf202a8d43ad1058a56c42c1dec474f593f7348e Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Sun, 21 Jul 2024 19:57:21 +0200 Subject: [PATCH 04/56] remove unnessary projection --- datafusion/sql/src/select.rs | 37 +++++-- datafusion/sql/src/utils.rs | 97 ++++++++++++------- datafusion/sqllogictest/test_files/debug.slt | 45 +++++++++ datafusion/sqllogictest/test_files/unnest.slt | 30 ++++-- 4 files changed, 159 insertions(+), 50 deletions(-) create mode 100644 datafusion/sqllogictest/test_files/debug.slt diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs index ea372fafa5a0..3297928c243c 100644 --- a/datafusion/sql/src/select.rs +++ b/datafusion/sql/src/select.rs @@ -333,9 +333,10 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { if unnest_columns.is_empty() { // The original expr does not contain any unnest if i == 0 { - return LogicalPlanBuilder::from(intermediate_plan) - .project(inner_projection_exprs)? - .build(); + return Ok(intermediate_plan); + // return LogicalPlanBuilder::from(intermediate_plan) + // .project(inner_projection_exprs)? + // .build(); } break; } else { @@ -343,19 +344,41 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { // Set preserve_nulls to false to ensure compatibility with DuckDB and PostgreSQL let unnest_options = UnnestOptions::new().with_preserve_nulls(false); // deduplicate expr in inner_projection_exprs - inner_projection_exprs.dedup_by(|a, b| -> bool { a == b }); + // BIGTODO: retain projection order instead of sorting + let mut checklist = HashSet::new(); + inner_projection_exprs.retain(|expr| -> bool { + if checklist.get(&expr.display_name().unwrap()).is_some() { + false + } else { + checklist.insert(expr.display_name().unwrap()); + true + } + }); let plan = LogicalPlanBuilder::from(intermediate_plan) - .project(inner_projection_exprs)? + .project(inner_projection_exprs.clone())? .unnest_columns_with_options(columns, unnest_options)? .build()?; intermediate_plan = plan; intermediate_select_exprs = outer_projection_exprs; } } - LogicalPlanBuilder::from(intermediate_plan) + + let mut checklist = HashSet::new(); + intermediate_select_exprs.retain(|expr| -> bool { + if checklist.get(&expr.display_name().unwrap()).is_some() { + false + } else { + checklist.insert(expr.display_name().unwrap()); + true + } + }); + + let ret = LogicalPlanBuilder::from(intermediate_plan) .project(intermediate_select_exprs)? - .build() + .build()?; + + Ok(ret) } fn plan_selection( diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index f2da2710556e..099f0522e504 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -18,7 +18,7 @@ //! SQL Utility Functions use std::cell::RefCell; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use arrow_schema::{ DataType, DECIMAL128_MAX_PRECISION, DECIMAL256_MAX_PRECISION, DECIMAL_DEFAULT_SCALE, @@ -264,6 +264,7 @@ pub(crate) fn normalize_ident(id: Ident) -> String { } } +/// TODO: explain me /// The context is we want to rewrite unnest() into InnerProjection->Unnest->OuterProjection /// Given an expression which contains unnest expr as one of its children, /// Try transform depends on unnest type @@ -279,40 +280,53 @@ pub(crate) fn transform_bottom_unnest( memo: &mut HashMap>, original_expr: &Expr, ) -> Result> { - let mut transform = - |unnest_expr: &Expr, expr_in_unnest: &Expr| -> Result> { - if let Some(previou_transformed) = memo.get(unnest_expr) { - return Ok(previou_transformed.clone()); - } - // Full context, we are trying to plan the execution as InnerProjection->Unnest->OuterProjection - // inside unnest execution, each column inside the inner projection - // will be transformed into new columns. Thus we need to keep track of these placeholding column names - let placeholder_name = unnest_expr.display_name()?; - - unnest_placeholder_columns.push(placeholder_name.clone()); - // Add alias for the argument expression, to avoid naming conflicts - // with other expressions in the select list. For example: `select unnest(col1), col1 from t`. - // this extra projection is used to unnest transforming - inner_projection_exprs - .push(expr_in_unnest.clone().alias(placeholder_name.clone())); - let schema = input.schema(); - - let (data_type, _) = expr_in_unnest.data_type_and_nullable(schema)?; - - let outer_projection_columns = - get_unnested_columns(&placeholder_name, &data_type)?; - let expr = outer_projection_columns - .iter() - .map(|col| Expr::Column(col.0.clone())) - .collect::>(); - - memo.insert(unnest_expr.clone(), expr.clone()); - Ok(expr) - }; + let mut transform = |unnest_expr: &Expr, + expr_in_unnest: &Expr, + inner_projection_exprs: &mut Vec| + -> Result> { + if let Some(previous_transformed) = memo.get(unnest_expr) { + return Ok(previous_transformed.clone()); + } + // Full context, we are trying to plan the execution as InnerProjection->Unnest->OuterProjection + // inside unnest execution, each column inside the inner projection + // will be transformed into new columns. Thus we need to keep track of these placeholding column names + let placeholder_name = unnest_expr.display_name()?; + + unnest_placeholder_columns.push(placeholder_name.clone()); + // Add alias for the argument expression, to avoid naming conflicts + // with other expressions in the select list. For example: `select unnest(col1), col1 from t`. + // this extra projection is used to unnest transforming + inner_projection_exprs + .push(expr_in_unnest.clone().alias(placeholder_name.clone())); + let schema = input.schema(); + + let (data_type, _) = expr_in_unnest.data_type_and_nullable(schema)?; + + let outer_projection_columns = + get_unnested_columns(&placeholder_name, &data_type)?; + let expr = outer_projection_columns + .iter() + .map(|col| Expr::Column(col.0.clone())) + .collect::>(); + + memo.insert(unnest_expr.clone(), expr.clone()); + Ok(expr) + }; let latest_visited = RefCell::new(None); + let column_under_unnest = RefCell::new(HashSet::new()); + let down_unnest = RefCell::new(None); // we need to mark only the latest unnest expr that was visitted during the down traversal let transform_down = |expr: Expr| -> Result> { - if let Expr::Unnest(Unnest { .. }) = expr { + if let Expr::Unnest(Unnest { + expr: ref inner_expr, + }) = expr + { + let mut down_unnest_mut = down_unnest.borrow_mut(); + if down_unnest_mut.is_none() { + *down_unnest_mut = Some(expr.clone()); + } + + column_under_unnest.borrow_mut().insert(inner_expr.clone()); *latest_visited.borrow_mut() = Some(expr.clone()); Ok(Transformed::no(expr)) } else { @@ -337,7 +351,8 @@ pub(crate) fn transform_bottom_unnest( return internal_err!("unnest on struct can ony be applied at the root level of select expression"); } - let mut transformed_exprs = transform(&expr, arg)?; + let mut transformed_exprs = + transform(&expr, arg, inner_projection_exprs)?; // root_expr.push(transformed_exprs[0].clone()); return Ok(Transformed::new( transformed_exprs.swap_remove(0), @@ -346,6 +361,20 @@ pub(crate) fn transform_bottom_unnest( )); } } + // For column exprs that are not descendants of any unnest node + // retain their projection + // e.g given expr tree unnest(col_a) + col_b, we have to retain projection of col_b + // down_unnest is non means current upward traversal is not descendant of any unnest + if matches!(&expr, Expr::Column(_)) && down_unnest.borrow().is_none() { + inner_projection_exprs.push(expr.clone()); + } + let mut down_unnest_mut = down_unnest.borrow_mut(); + // upward traversal has reached the top unnest expr again + // reset it to None + if *down_unnest_mut == Some(expr.clone()) { + down_unnest_mut.take(); + } + Ok(Transformed::no(expr)) }; @@ -372,7 +401,7 @@ pub(crate) fn transform_bottom_unnest( // The transformation looks like // - unnest(struct_col) will be transformed into unnest(struct_col).field1, unnest(struct_col).field2 if let Expr::Unnest(Unnest { expr: ref arg }) = transformed_expr { - return transform(&transformed_expr, arg); + return transform(&transformed_expr, arg, inner_projection_exprs); } if matches!(&transformed_expr, Expr::Column(_)) { inner_projection_exprs.push(transformed_expr.clone()); diff --git a/datafusion/sqllogictest/test_files/debug.slt b/datafusion/sqllogictest/test_files/debug.slt new file mode 100644 index 000000000000..f310f90b9ce7 --- /dev/null +++ b/datafusion/sqllogictest/test_files/debug.slt @@ -0,0 +1,45 @@ +##query TT +##explain select unnest(unnest([[1,2,3]])) + unnest([4,5]); +##---- +##logical_plan +##01)Projection: unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))) + unnest(make_array(Int64(4),Int64(5))) +##02)--Unnest: lists[unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))))] structs[] +##03)----Projection: unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))), unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))) AS unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))), unnest(make_array(Int64(4),Int64(5))) +##04)------Unnest: lists[unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))), unnest(make_array(Int64(4),Int64(5)))] structs[] +##05)--------Projection: List([[1, 2, 3]]) AS unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))), List([4, 5]) AS unnest(make_array(Int64(4),Int64(5))) +##06)----------EmptyRelation +##physical_plan +##01)ProjectionExec: expr=[unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))))@1 + unnest(make_array(Int64(4),Int64(5)))@2 as unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))) + unnest(make_array(Int64(4),Int64(5)))] +##02)--UnnestExec +##03)----RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +##04)------ProjectionExec: expr=[unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))@0 as unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))), unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))@0 as unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))), unnest(make_array(Int64(4),Int64(5)))@1 as unnest(make_array(Int64(4),Int64(5)))] +##05)--------UnnestExec +##06)----------ProjectionExec: expr=[[[1, 2, 3]] as unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))), [4, 5] as unnest(make_array(Int64(4),Int64(5)))] +##07)------------PlaceholderRowExec + +query I +select unnest([4,5]) + 1; +---- +5 +6 + +## FIXME: +query II +select unnest(unnest([[1,2,3]])) + unnest([4,5]), arrow_cast(unnest([4,5]),'Int64'); +---- +5 4 +6 4 +7 4 + + + +query I +select unnest([2,1,9]) + unnest(unnest([[1,1,3]])) ; +---- +3 + +query I +select unnest(unnest([[1,1]])) + unnest([2,1,9]),unnest([2,1,9]) + unnest(unnest([[1,1]])) ; +---- +3 + diff --git a/datafusion/sqllogictest/test_files/unnest.slt b/datafusion/sqllogictest/test_files/unnest.slt index 4145cb02a84e..f91cbb99b508 100644 --- a/datafusion/sqllogictest/test_files/unnest.slt +++ b/datafusion/sqllogictest/test_files/unnest.slt @@ -165,16 +165,28 @@ select unnest(column1), column1 from unnest_table; 6 [6] 12 [12] -query I -select unnest(column1) + unnest(column1) from unnest_table; + +query II +select unnest([1,2,3]) + unnest([1,2,3]), unnest([1,2,3]) + unnest([4,5]); ---- -2 -4 -6 -8 -10 -12 -24 +2 5 +4 7 +6 NULL + +##FIXME: +## Expectation: +## 2 5 4 +## 4 7 5 +## 6 NULL NULL +query III +select unnest(unnest([[1,2,3]])) + unnest(unnest([[1,2,3]])), unnest(unnest([[1,2,3]])) + unnest([4,5]), unnest([4,5]); +---- +2 5 4 +4 6 4 +6 7 4 + + + ## unnest as children of other expr query I? From 1e74176667d5951794f49fe7012b1be2e389a60e Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Mon, 22 Jul 2024 19:45:53 +0200 Subject: [PATCH 05/56] chore: temp test case --- datafusion/sqllogictest/test_files/unnest.slt | 51 ++++++++++++++++++- 1 file changed, 50 insertions(+), 1 deletion(-) diff --git a/datafusion/sqllogictest/test_files/unnest.slt b/datafusion/sqllogictest/test_files/unnest.slt index f91cbb99b508..04ce5a52f11b 100644 --- a/datafusion/sqllogictest/test_files/unnest.slt +++ b/datafusion/sqllogictest/test_files/unnest.slt @@ -606,4 +606,53 @@ physical_plan 05)--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 06)----------UnnestExec 07)------------ProjectionExec: expr=[column3@0 as unnest(recursive_unnest_table.column3), column3@0 as column3] -08)--------------MemoryExec: partitions=1, partition_sizes=[1] \ No newline at end of file +08)--------------MemoryExec: partitions=1, partition_sizes=[1] + +statement ok +CREATE TABLE temp2 +AS VALUES + ([1,2,3],4), + (null,5), + ([1],6) +; + +query II +select unnest(column1), column2 from temp2 +---- +1 4 +2 5 +3 NULL + +statement ok +CREATE TABLE temp +AS VALUES + ([[1,2,3]],[4,5]) +; + +query II +select unnest(unnest(column1)), unnest(column2) from temp +---- +1 4 +2 4 +3 4 + +##[[1,3,4]] + [4,5] +##[1,2,3] 4 + + +query TT +explain select unnest(unnest(column1)), unnest(column2) from temp +---- +logical_plan +01)Unnest: lists[unnest(unnest(temp.column1))] structs[] +02)--Projection: unnest(temp.column1) AS unnest(unnest(temp.column1)), unnest(temp.column2) +03)----Unnest: lists[unnest(temp.column1), unnest(temp.column2)] structs[] +04)------Projection: temp.column1 AS unnest(temp.column1), temp.column2 AS unnest(temp.column2) +05)--------TableScan: temp projection=[column1, column2] +physical_plan +01)UnnestExec +02)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +03)----ProjectionExec: expr=[unnest(temp.column1)@0 as unnest(unnest(temp.column1)), unnest(temp.column2)@1 as unnest(temp.column2)] +04)------UnnestExec +05)--------ProjectionExec: expr=[column1@0 as unnest(temp.column1), column2@1 as unnest(temp.column2)] +06)----------MemoryExec: partitions=1, partition_sizes=[1] From d557ff3f7e5a6304906b9e4fd26df6e8f0757dba Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Sat, 27 Jul 2024 14:33:40 +0200 Subject: [PATCH 06/56] multi depth unnest supported --- datafusion/sql/src/select.rs | 6 +- datafusion/sql/src/utils.rs | 396 +++++++++++++++++++++++++---------- 2 files changed, 295 insertions(+), 107 deletions(-) diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs index 7666579cdc81..812b0f0bef06 100644 --- a/datafusion/sql/src/select.rs +++ b/datafusion/sql/src/select.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use std::collections::{HashMap, HashSet}; +use std::collections::HashSet; use std::sync::Arc; use crate::planner::{ @@ -306,7 +306,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { let mut intermediate_select_exprs = select_exprs; // impl memoization to store all previous unnest transformation - let mut memo = HashMap::new(); + let mut memo = HashSet::new(); // Each expr in select_exprs can contains multiple unnest stage // The transformation happen bottom up, one at a time for each iteration // Only exaust the loop if no more unnest transformation is found @@ -442,6 +442,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { // ``` let mut intermediate_plan = unwrap_arc(input); let mut intermediate_select_exprs = group_expr; + let mut memo = HashSet::new(); loop { let mut unnest_columns = vec![]; @@ -454,6 +455,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { &intermediate_plan, &mut unnest_columns, &mut inner_projection_exprs, + &mut memo, expr, ) }) diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index 099f0522e504..6c30f47adad0 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -277,103 +277,160 @@ pub(crate) fn transform_bottom_unnest( input: &LogicalPlan, unnest_placeholder_columns: &mut Vec, inner_projection_exprs: &mut Vec, - memo: &mut HashMap>, + memo: &mut HashSet, original_expr: &Expr, ) -> Result> { - let mut transform = |unnest_expr: &Expr, + let mut transform = |level: usize, expr_in_unnest: &Expr, + struct_allowed: bool, inner_projection_exprs: &mut Vec| -> Result> { - if let Some(previous_transformed) = memo.get(unnest_expr) { - return Ok(previous_transformed.clone()); - } + let already_projected = memo.get(expr_in_unnest).is_some(); // Full context, we are trying to plan the execution as InnerProjection->Unnest->OuterProjection // inside unnest execution, each column inside the inner projection // will be transformed into new columns. Thus we need to keep track of these placeholding column names - let placeholder_name = unnest_expr.display_name()?; + // let placeholder_name = unnest_expr.display_name()?; + let placeholder_name = format!( + "unnest_placeholder({})", + expr_in_unnest.display_name().unwrap(), + ); + let post_unnest_name = format!( + "unnest_placeholder({},depth={})", + expr_in_unnest.display_name().unwrap(), + level + ); - unnest_placeholder_columns.push(placeholder_name.clone()); // Add alias for the argument expression, to avoid naming conflicts // with other expressions in the select list. For example: `select unnest(col1), col1 from t`. // this extra projection is used to unnest transforming - inner_projection_exprs - .push(expr_in_unnest.clone().alias(placeholder_name.clone())); + if !already_projected { + inner_projection_exprs + .push(expr_in_unnest.clone().alias(placeholder_name.clone())); + + unnest_placeholder_columns.push(placeholder_name.clone()); + } + let schema = input.schema(); let (data_type, _) = expr_in_unnest.data_type_and_nullable(schema)?; + if !struct_allowed { + if let DataType::Struct(_) = data_type { + return internal_err!("unnest on struct can only be applied at the root level of select expression"); + } + } let outer_projection_columns = - get_unnested_columns(&placeholder_name, &data_type)?; + get_unnested_columns(&post_unnest_name, &data_type)?; let expr = outer_projection_columns .iter() .map(|col| Expr::Column(col.0.clone())) .collect::>(); - memo.insert(unnest_expr.clone(), expr.clone()); + memo.insert(expr_in_unnest.clone()); Ok(expr) }; - let latest_visited = RefCell::new(None); - let column_under_unnest = RefCell::new(HashSet::new()); - let down_unnest = RefCell::new(None); + let latest_visited_unnest = RefCell::new(None); + let exprs_under_unnest = RefCell::new(HashSet::new()); + let ancestor_unnest = RefCell::new(None); + + let consecutive_unnest = RefCell::new(Vec::>::new()); // we need to mark only the latest unnest expr that was visitted during the down traversal let transform_down = |expr: Expr| -> Result> { if let Expr::Unnest(Unnest { expr: ref inner_expr, }) = expr { - let mut down_unnest_mut = down_unnest.borrow_mut(); - if down_unnest_mut.is_none() { - *down_unnest_mut = Some(expr.clone()); + let mut consecutive_unnest_mut = consecutive_unnest.borrow_mut(); + consecutive_unnest_mut.push(Some(expr.clone())); + + let mut maybe_ancestor = ancestor_unnest.borrow_mut(); + if maybe_ancestor.is_none() { + *maybe_ancestor = Some(expr.clone()); } - column_under_unnest.borrow_mut().insert(inner_expr.clone()); - *latest_visited.borrow_mut() = Some(expr.clone()); + exprs_under_unnest.borrow_mut().insert(inner_expr.clone()); + *latest_visited_unnest.borrow_mut() = Some(expr.clone()); Ok(Transformed::no(expr)) } else { + consecutive_unnest.borrow_mut().push(None); Ok(Transformed::no(expr)) } }; let transform_up = |expr: Expr| -> Result> { - if let Expr::Unnest(Unnest { expr: ref arg }) = expr { - // only transform the first unnest expr(s) from the bottom up - // if the expr tree contains mulitple unnest exprs, as long as neither of them - // is the direct ancestor of one another, we do all the transformation - if let Some(ref mut last_visitted_expr) = *latest_visited.borrow_mut() { - if last_visitted_expr != &expr { - return Ok(Transformed::no(expr)); - } - // this is (one of) the bottom most unnest expr - let (data_type, _) = arg.data_type_and_nullable(input.schema())?; - if &expr == original_expr { - return Ok(Transformed::no(expr)); + // From the bottom up, we know the latest consecutive unnest sequence + // we only do the transformation at the top unnest node + // For example given this complex expr + // - unnest(array_concat(unnest([[1,2,3]]),unnest([[4,5,6]]))) + unnest(unnest([[7,8,9])) + // traversal will be like this: + // down[binary_add] + // ->down[unnest(...)]->down[array_concat]->down/up[unnest([[1,2,3]])]->down/up[unnest([[4,5,6]])] + // ->up[array_concat]->up[unnest(...)]->down[unnest(unnest(...))]->down[unnest([[7,8,9]])] + // ->up[unnest([[7,8,9]])]->up[unnest(unnest(...))]->up[binary_add] + // the transformation only happens for unnest([[1,2,3]]), unnest([[4,5,6]]) and unnest(unnest([[7,8,9]])) + // and the complex expr will be rewritten into: + // unnest(array_concat(place_holder_col_1, place_holder_col_2)) + place_holder_col_3 + if let Expr::Unnest(Unnest { .. }) = expr { + let mut down_unnest_mut = ancestor_unnest.borrow_mut(); + // upward traversal has reached the top unnest expr again + // reset it to None + if *down_unnest_mut == Some(expr.clone()) { + down_unnest_mut.take(); + } + // find inside consecutive_unnest, the sequence of continous unnest exprs + let mut found_first_unnest = false; + let mut unnest_stack = vec![]; + for item in consecutive_unnest.borrow().iter().rev() { + if let Some(expr) = item { + found_first_unnest = true; + unnest_stack.push(expr.clone()); + } else { + if !found_first_unnest { + continue; + } + break; } - if let DataType::Struct(_) = data_type { - return internal_err!("unnest on struct can ony be applied at the root level of select expression"); + } + + // this is the top most unnest expr inside the consecutive unnest exprs + // e.g unnest(unnest(some_col)) + if expr == *unnest_stack.last().unwrap() { + let most_inner = unnest_stack.first().unwrap(); + if let Expr::Unnest(Unnest { expr: ref arg }) = most_inner { + // this is (one of) the bottom most unnest expr + let (data_type, _) = arg.data_type_and_nullable(input.schema())?; + if &expr == original_expr { + return Ok(Transformed::no(expr)); + } + if let DataType::Struct(_) = data_type { + return internal_err!("unnest on struct can only be applied at the root level of select expression"); + } + let depth = unnest_stack.len(); + let struct_allowed = (&expr == original_expr) && depth == 1; + + let mut transformed_exprs = + transform(depth, arg, struct_allowed, inner_projection_exprs)?; + return Ok(Transformed::new( + transformed_exprs.swap_remove(0), + true, + TreeNodeRecursion::Continue, + )); + } else { + return internal_err!("not reached"); } - let mut transformed_exprs = - transform(&expr, arg, inner_projection_exprs)?; - // root_expr.push(transformed_exprs[0].clone()); - return Ok(Transformed::new( - transformed_exprs.swap_remove(0), - true, - TreeNodeRecursion::Continue, - )); + // } } + } else { + consecutive_unnest.borrow_mut().push(None); } + // For column exprs that are not descendants of any unnest node // retain their projection // e.g given expr tree unnest(col_a) + col_b, we have to retain projection of col_b // down_unnest is non means current upward traversal is not descendant of any unnest - if matches!(&expr, Expr::Column(_)) && down_unnest.borrow().is_none() { + if matches!(&expr, Expr::Column(_)) && ancestor_unnest.borrow().is_none() { inner_projection_exprs.push(expr.clone()); } - let mut down_unnest_mut = down_unnest.borrow_mut(); - // upward traversal has reached the top unnest expr again - // reset it to None - if *down_unnest_mut == Some(expr.clone()) { - down_unnest_mut.take(); - } Ok(Transformed::no(expr)) }; @@ -396,13 +453,6 @@ pub(crate) fn transform_bottom_unnest( .transform_down_up(transform_down, transform_up)?; if !transformed { - // Because root expr need to transform separately - // unnest struct is only possible here - // The transformation looks like - // - unnest(struct_col) will be transformed into unnest(struct_col).field1, unnest(struct_col).field2 - if let Expr::Unnest(Unnest { expr: ref arg }) = transformed_expr { - return transform(&transformed_expr, arg, inner_projection_exprs); - } if matches!(&transformed_expr, Expr::Column(_)) { inner_projection_exprs.push(transformed_expr.clone()); Ok(vec![transformed_expr]) @@ -421,7 +471,7 @@ pub(crate) fn transform_bottom_unnest( // write test for recursive_transform_unnest #[cfg(test)] mod tests { - use std::{collections::HashMap, ops::Add, sync::Arc}; + use std::{collections::HashSet, ops::Add, sync::Arc}; use arrow::datatypes::{DataType as ArrowDataType, Field, Schema}; use arrow_schema::Fields; @@ -433,22 +483,47 @@ mod tests { use datafusion_functions_aggregate::expr_fn::count; use crate::utils::{resolve_positions_to_exprs, transform_bottom_unnest}; - #[test] - fn test_transform_bottom_unnest_recursive_memoization() -> Result<()> { - let schema = Schema::new(vec![Field::new( - "3d_col", + fn test_transform_bottom_unnest_recursive_memoization_struct() -> Result<()> { + let three_d_dtype = ArrowDataType::List(Arc::new(Field::new( + "2d_col", ArrowDataType::List(Arc::new(Field::new( - "2d_col", + "elements", + ArrowDataType::Int64, + true, + ))), + true, + ))); + let schema = Schema::new(vec![ + // list[struct(3d_data)] [([[1,2,3]])] + Field::new( + "struct_arr_col", ArrowDataType::List(Arc::new(Field::new( - "elements", - ArrowDataType::Int64, + "struct", + ArrowDataType::Struct(Fields::from(vec![Field::new( + "field1", + three_d_dtype, + true, + )])), true, ))), true, - ))), - true, - )]); + ), + Field::new( + "3d_col", + ArrowDataType::List(Arc::new(Field::new( + "2d_col", + ArrowDataType::List(Arc::new(Field::new( + "elements", + ArrowDataType::Int64, + true, + ))), + true, + ))), + true, + ), + Field::new("i64_col", ArrowDataType::Int64, true), + ]); let dfschema = DFSchema::try_from(schema)?; @@ -461,9 +536,10 @@ mod tests { let mut inner_projection_exprs = vec![]; // unnest(unnest(3d_col)) + unnest(unnest(3d_col)) - let original_expr = - unnest(unnest(col("3d_col"))).add(unnest(unnest(col("3d_col")))); - let mut memo = HashMap::new(); + let original_expr = unnest(unnest(col("3d_col"))) + .add(unnest(unnest(col("3d_col")))) + .add(col("i64_col")); + let mut memo = HashSet::new(); let transformed_exprs = transform_bottom_unnest( &input, &mut unnest_placeholder_columns, @@ -474,20 +550,25 @@ mod tests { // only the bottom most unnest exprs are transformed assert_eq!( transformed_exprs, - vec![unnest(col("unnest(3d_col)")).add(unnest(col("unnest(3d_col)")))] + vec![col("unnest_placeholder(3d_col,depth=2)") + .add(col("unnest_placeholder(3d_col,depth=2)")) + .add(col("i64_col"))] ); // memoization only contains 1 transformation assert_eq!(memo.len(), 1); + assert!(memo.get(&col("3d_col")).is_some()); assert_eq!( - memo.get(&unnest(col("3d_col"))), - Some(&vec![col("unnest(3d_col)")]) + unnest_placeholder_columns, + vec!["unnest_placeholder(3d_col)"] ); - assert_eq!(unnest_placeholder_columns, vec!["unnest(3d_col)"]); // still reference struct_col in original schema but with alias, // to avoid colliding with the projection on the column itself if any assert_eq!( inner_projection_exprs, - vec![col("3d_col").alias("unnest(3d_col)"),] + vec![ + col("3d_col").alias("unnest_placeholder(3d_col)"), + col("i64_col") + ] ); // unnest(3d_col) as 2d_col @@ -502,60 +583,165 @@ mod tests { assert_eq!( transformed_exprs, - vec![col("unnest(3d_col)").alias("2d_col")] + vec![col("unnest_placeholder(3d_col,depth=1)").alias("2d_col")] ); // memoization still contains 1 transformation // and the previous transformation is reused assert_eq!(memo.len(), 1); + assert!(memo.get(&col("3d_col")).is_some()); assert_eq!( - memo.get(&unnest(col("3d_col"))), - Some(&vec![col("unnest(3d_col)")]) + unnest_placeholder_columns, + vec!["unnest_placeholder(3d_col)"] ); - assert_eq!(unnest_placeholder_columns, vec!["unnest(3d_col)"]); // still reference struct_col in original schema but with alias, // to avoid colliding with the projection on the column itself if any assert_eq!( inner_projection_exprs, - vec![col("3d_col").alias("unnest(3d_col)")] + vec![ + col("3d_col").alias("unnest_placeholder(3d_col)"), + col("i64_col") + ] ); - // Start a new cycle, to run unnest again on previous transformation - let intermediate_columns = unnest_placeholder_columns - .into_iter() - .map(|col| col.into()) - .collect(); - let intermediate_input = LogicalPlanBuilder::from(input) - .project(inner_projection_exprs)? - .unnest_columns_with_options(intermediate_columns, UnnestOptions::default())? - .build()?; - - let mut new_unnest_placeholder_columns = vec![]; - let mut new_inner_projection_exprs = vec![]; - // Run unnest again on previously transformed expr + // unnest(unnset(unnest(struct_arr_col)['field1'])) as fully_unnested_struct_arr + let original_expr_3 = + unnest(unnest(unnest(col("struct_arr_col")).field("field1"))) + .alias("fully_unnested_struct_arr"); let transformed_exprs = transform_bottom_unnest( - &intermediate_input, - &mut new_unnest_placeholder_columns, - &mut new_inner_projection_exprs, + &input, + &mut unnest_placeholder_columns, + &mut inner_projection_exprs, &mut memo, - &unnest(col("unnest(3d_col)")).add(unnest(col("unnest(3d_col)"))), + &original_expr_3, )?; + assert_eq!( transformed_exprs, - vec![col("unnest(unnest(3d_col))").add(col("unnest(unnest(3d_col))"))], + vec![unnest(unnest( + col("unnest_placeholder(struct_arr_col,depth=1)").field("field1") + )) + .alias("fully_unnested_struct_arr")] ); - // memoization having extra transformation cached + // memoization still contains 1 transformation + // and the previous transformation is reused assert_eq!(memo.len(), 2); + assert!(memo.get(&col("struct_arr_col")).is_some()); + assert_eq!( + unnest_placeholder_columns, + vec![ + "unnest_placeholder(3d_col)", + "unnest_placeholder(struct_arr_col)" + ] + ); + // still reference struct_col in original schema but with alias, + // to avoid colliding with the projection on the column itself if any + assert_eq!( + inner_projection_exprs, + vec![ + col("3d_col").alias("unnest_placeholder(3d_col)"), + col("i64_col"), + col("struct_arr_col").alias("unnest_placeholder(struct_arr_col)") + ] + ); + + Ok(()) + } + + #[test] + fn test_transform_bottom_unnest_recursive_memoization() -> Result<()> { + let schema = Schema::new(vec![ + Field::new( + "3d_col", + ArrowDataType::List(Arc::new(Field::new( + "2d_col", + ArrowDataType::List(Arc::new(Field::new( + "elements", + ArrowDataType::Int64, + true, + ))), + true, + ))), + true, + ), + Field::new("i64_col", ArrowDataType::Int64, true), + ]); + + let dfschema = DFSchema::try_from(schema)?; + + let input = LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: Arc::new(dfschema), + }); + + let mut unnest_placeholder_columns = vec![]; + let mut inner_projection_exprs = vec![]; + + // unnest(unnest(3d_col)) + unnest(unnest(3d_col)) + let original_expr = unnest(unnest(col("3d_col"))) + .add(unnest(unnest(col("3d_col")))) + .add(col("i64_col")); + let mut memo = HashSet::new(); + let transformed_exprs = transform_bottom_unnest( + &input, + &mut unnest_placeholder_columns, + &mut inner_projection_exprs, + &mut memo, + &original_expr, + )?; + // only the bottom most unnest exprs are transformed assert_eq!( - memo.get(&unnest(col("unnest(3d_col)"))), - Some(&vec![col("unnest(unnest(3d_col))")]) + transformed_exprs, + vec![col("unnest_placeholder(3d_col,depth=2)") + .add(col("unnest_placeholder(3d_col,depth=2)")) + .add(col("i64_col"))] + ); + // memoization only contains 1 transformation + assert_eq!(memo.len(), 1); + assert!(memo.get(&col("3d_col")).is_some()); + assert_eq!( + unnest_placeholder_columns, + vec!["unnest_placeholder(3d_col)"] + ); + // still reference struct_col in original schema but with alias, + // to avoid colliding with the projection on the column itself if any + assert_eq!( + inner_projection_exprs, + vec![ + col("3d_col").alias("unnest_placeholder(3d_col)"), + col("i64_col") + ] + ); + + // unnest(3d_col) as 2d_col + let original_expr_2 = unnest(col("3d_col")).alias("2d_col"); + let transformed_exprs = transform_bottom_unnest( + &input, + &mut unnest_placeholder_columns, + &mut inner_projection_exprs, + &mut memo, + &original_expr_2, + )?; + + assert_eq!( + transformed_exprs, + vec![col("unnest_placeholder(3d_col,depth=1)").alias("2d_col")] ); + // memoization still contains 1 transformation + // and the previous transformation is reused + assert_eq!(memo.len(), 1); + assert!(memo.get(&col("3d_col")).is_some()); assert_eq!( - new_unnest_placeholder_columns, - vec!["unnest(unnest(3d_col))"] + unnest_placeholder_columns, + vec!["unnest_placeholder(3d_col)"] ); + // still reference struct_col in original schema but with alias, + // to avoid colliding with the projection on the column itself if any assert_eq!( - new_inner_projection_exprs, - vec![col("unnest(3d_col)").alias("unnest(unnest(3d_col))")] + inner_projection_exprs, + vec![ + col("3d_col").alias("unnest_placeholder(3d_col)"), + col("i64_col") + ] ); Ok(()) @@ -594,7 +780,7 @@ mod tests { let mut unnest_placeholder_columns = vec![]; let mut inner_projection_exprs = vec![]; - let mut memo = HashMap::new(); + let mut memo = HashSet::new(); // unnest(struct_col) let original_expr = unnest(col("struct_col")); let transformed_exprs = transform_bottom_unnest( From cd497fed694a8fa199feaf5f9ae489972ef01c5b Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Sun, 28 Jul 2024 08:01:33 +0200 Subject: [PATCH 07/56] chore: add map of original column and transformed col --- datafusion/expr/src/logical_plan/builder.rs | 10 ++-- datafusion/sql/src/select.rs | 27 ++--------- datafusion/sql/src/utils.rs | 48 ++++++++++++-------- datafusion/sqllogictest/test_files/debug.slt | 36 +++++++-------- 4 files changed, 56 insertions(+), 65 deletions(-) diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index 98e262f0b187..6a0fc4ba9d39 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -1136,7 +1136,7 @@ impl LogicalPlanBuilder { /// Unnest the given columns with the given [`UnnestOptions`] pub fn unnest_columns_with_options( self, - columns: Vec, + columns: HashMap>, options: UnnestOptions, ) -> Result { Ok(Self::from(unnest_with_options( @@ -1646,14 +1646,16 @@ pub fn get_unnested_columns( /// Create a [`LogicalPlan::Unnest`] plan with options pub fn unnest_with_options( input: LogicalPlan, - columns: Vec, + columns: HashMap>, options: UnnestOptions, ) -> Result { let mut list_columns = Vec::with_capacity(columns.len()); let mut struct_columns = Vec::with_capacity(columns.len()); let column_by_original_index = columns .iter() - .map(|c| Ok((input.schema().index_of_column(c)?, c))) + .map(|(inner_col, outer_cols)| { + Ok((input.schema().index_of_column(inner_col)?, inner_col)) + }) .collect::>>()?; let input_schema = input.schema(); @@ -1668,6 +1670,8 @@ pub fn unnest_with_options( .map(|(index, (original_qualifier, original_field))| { match column_by_original_index.get(&index) { Some(&column_to_unnest) => { + let transformed_cols = columns.get(&column_to_unnest).unwrap(); + println!("columnto unnest name {}", column_to_unnest.name); let flatten_columns = get_unnested_columns( &column_to_unnest.name, original_field.data_type(), diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs index 812b0f0bef06..ab5ff7dbd499 100644 --- a/datafusion/sql/src/select.rs +++ b/datafusion/sql/src/select.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; use crate::planner::{ @@ -305,12 +305,12 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { let mut intermediate_plan = input; let mut intermediate_select_exprs = select_exprs; - // impl memoization to store all previous unnest transformation - let mut memo = HashSet::new(); // Each expr in select_exprs can contains multiple unnest stage // The transformation happen bottom up, one at a time for each iteration // Only exaust the loop if no more unnest transformation is found for i in 0.. { + // impl memoization to store all previous unnest transformation + let mut memo = HashMap::new(); let mut unnest_columns = vec![]; // from which column used for projection, before the unnest happen // including non unnest column and unnest column @@ -351,17 +351,6 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { let columns = unnest_columns.into_iter().map(|col| col.into()).collect(); // Set preserve_nulls to false to ensure compatibility with DuckDB and PostgreSQL let unnest_options = UnnestOptions::new().with_preserve_nulls(false); - // deduplicate expr in inner_projection_exprs - // BIGTODO: retain projection order instead of sorting - let mut checklist = HashSet::new(); - inner_projection_exprs.retain(|expr| -> bool { - if checklist.get(&expr.display_name().unwrap()).is_some() { - false - } else { - checklist.insert(expr.display_name().unwrap()); - true - } - }); let plan = LogicalPlanBuilder::from(intermediate_plan) .project(inner_projection_exprs.clone())? @@ -372,16 +361,6 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { } } - let mut checklist = HashSet::new(); - intermediate_select_exprs.retain(|expr| -> bool { - if checklist.get(&expr.display_name().unwrap()).is_some() { - false - } else { - checklist.insert(expr.display_name().unwrap()); - true - } - }); - let ret = LogicalPlanBuilder::from(intermediate_plan) .project(intermediate_select_exprs)? .build()?; diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index 6c30f47adad0..ad7a7fc6304d 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -26,6 +26,7 @@ use arrow_schema::{ use datafusion_common::tree_node::{ Transformed, TransformedResult, TreeNode, TreeNodeRecursion, }; +use datafusion_common::utils::proxy::VecAllocExt; use datafusion_common::{ exec_err, internal_err, plan_err, Column, DataFusionError, Result, ScalarValue, }; @@ -277,7 +278,7 @@ pub(crate) fn transform_bottom_unnest( input: &LogicalPlan, unnest_placeholder_columns: &mut Vec, inner_projection_exprs: &mut Vec, - memo: &mut HashSet, + memo: &mut HashMap>, original_expr: &Expr, ) -> Result> { let mut transform = |level: usize, @@ -285,20 +286,27 @@ pub(crate) fn transform_bottom_unnest( struct_allowed: bool, inner_projection_exprs: &mut Vec| -> Result> { - let already_projected = memo.get(expr_in_unnest).is_some(); + let col = match expr_in_unnest { + Expr::Column(col) => col, + _ => { + return internal_err!("unnesting on non-column expr is not supported"); + } + }; + let already_projected = memo.get(col); + let (already_projected, transformed_cols) = match memo.get_mut(col) { + Some(vec) => (true, vec), + _ => { + memo.insert(col.clone(), vec![]); + (false, memo.get_mut(col).unwrap()) + } + }; // Full context, we are trying to plan the execution as InnerProjection->Unnest->OuterProjection // inside unnest execution, each column inside the inner projection // will be transformed into new columns. Thus we need to keep track of these placeholding column names // let placeholder_name = unnest_expr.display_name()?; - let placeholder_name = format!( - "unnest_placeholder({})", - expr_in_unnest.display_name().unwrap(), - ); - let post_unnest_name = format!( - "unnest_placeholder({},depth={})", - expr_in_unnest.display_name().unwrap(), - level - ); + let placeholder_name = format!("unnest_placeholder({})", col.name()); + let post_unnest_name = + format!("unnest_placeholder({},depth={})", col.name(), level); // Add alias for the argument expression, to avoid naming conflicts // with other expressions in the select list. For example: `select unnest(col1), col1 from t`. @@ -323,10 +331,14 @@ pub(crate) fn transform_bottom_unnest( get_unnested_columns(&post_unnest_name, &data_type)?; let expr = outer_projection_columns .iter() - .map(|col| Expr::Column(col.0.clone())) + .map(|col| { + if !transformed_cols.contains(&col.0) { + transformed_cols.push(col.0.clone()); + } + Expr::Column(col.0.clone()) + }) .collect::>(); - memo.insert(expr_in_unnest.clone()); Ok(expr) }; let latest_visited_unnest = RefCell::new(None); @@ -401,9 +413,7 @@ pub(crate) fn transform_bottom_unnest( if &expr == original_expr { return Ok(Transformed::no(expr)); } - if let DataType::Struct(_) = data_type { - return internal_err!("unnest on struct can only be applied at the root level of select expression"); - } + let depth = unnest_stack.len(); let struct_allowed = (&expr == original_expr) && depth == 1; @@ -475,10 +485,8 @@ mod tests { use arrow::datatypes::{DataType as ArrowDataType, Field, Schema}; use arrow_schema::Fields; - use datafusion_common::{DFSchema, Result, UnnestOptions}; - use datafusion_expr::{ - col, lit, unnest, EmptyRelation, LogicalPlan, LogicalPlanBuilder, - }; + use datafusion_common::{DFSchema, Result}; + use datafusion_expr::{col, lit, unnest, EmptyRelation, LogicalPlan}; use datafusion_functions::core::expr_ext::FieldAccessor; use datafusion_functions_aggregate::expr_fn::count; diff --git a/datafusion/sqllogictest/test_files/debug.slt b/datafusion/sqllogictest/test_files/debug.slt index f310f90b9ce7..76fb93a7db49 100644 --- a/datafusion/sqllogictest/test_files/debug.slt +++ b/datafusion/sqllogictest/test_files/debug.slt @@ -1,21 +1,21 @@ -##query TT -##explain select unnest(unnest([[1,2,3]])) + unnest([4,5]); -##---- -##logical_plan -##01)Projection: unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))) + unnest(make_array(Int64(4),Int64(5))) -##02)--Unnest: lists[unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))))] structs[] -##03)----Projection: unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))), unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))) AS unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))), unnest(make_array(Int64(4),Int64(5))) -##04)------Unnest: lists[unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))), unnest(make_array(Int64(4),Int64(5)))] structs[] -##05)--------Projection: List([[1, 2, 3]]) AS unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))), List([4, 5]) AS unnest(make_array(Int64(4),Int64(5))) -##06)----------EmptyRelation -##physical_plan -##01)ProjectionExec: expr=[unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))))@1 + unnest(make_array(Int64(4),Int64(5)))@2 as unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))) + unnest(make_array(Int64(4),Int64(5)))] -##02)--UnnestExec -##03)----RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -##04)------ProjectionExec: expr=[unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))@0 as unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))), unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))@0 as unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))), unnest(make_array(Int64(4),Int64(5)))@1 as unnest(make_array(Int64(4),Int64(5)))] -##05)--------UnnestExec -##06)----------ProjectionExec: expr=[[[1, 2, 3]] as unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))), [4, 5] as unnest(make_array(Int64(4),Int64(5)))] -##07)------------PlaceholderRowExec +query TT +explain select unnest(unnest([[1,2,3]])) + unnest([4,5]); +---- +logical_plan +01)Projection: unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))) + unnest(make_array(Int64(4),Int64(5))) +02)--Unnest: lists[unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))))] structs[] +03)----Projection: unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))), unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))) AS unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))), unnest(make_array(Int64(4),Int64(5))) +04)------Unnest: lists[unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))), unnest(make_array(Int64(4),Int64(5)))] structs[] +05)--------Projection: List([[1, 2, 3]]) AS unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))), List([4, 5]) AS unnest(make_array(Int64(4),Int64(5))) +06)----------EmptyRelation +physical_plan +01)ProjectionExec: expr=[unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))))@1 + unnest(make_array(Int64(4),Int64(5)))@2 as unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))) + unnest(make_array(Int64(4),Int64(5)))] +02)--UnnestExec +03)----RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +04)------ProjectionExec: expr=[unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))@0 as unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))), unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))@0 as unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))), unnest(make_array(Int64(4),Int64(5)))@1 as unnest(make_array(Int64(4),Int64(5)))] +05)--------UnnestExec +06)----------ProjectionExec: expr=[[[1, 2, 3]] as unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))), [4, 5] as unnest(make_array(Int64(4),Int64(5)))] +07)------------PlaceholderRowExec query I select unnest([4,5]) + 1; From d2cc80cf0e9325a214f2dbb3d056abced340f87c Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Sun, 28 Jul 2024 10:11:33 +0200 Subject: [PATCH 08/56] transformation map to physical layer --- datafusion/physical-plan/src/unnest.rs | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index bdd56f4b5aa4..df757a1d378c 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -62,6 +62,8 @@ pub struct UnnestExec { input: Arc, /// The schema once the unnest is applied schema: SchemaRef, + /// which original columns are transformed into which columns + transformed_col: HashMap>, /// indices of the list-typed columns in the input schema list_column_indices: Vec, /// indices of the struct-typed columns in the input schema @@ -80,6 +82,7 @@ impl UnnestExec { input: Arc, list_column_indices: Vec, struct_column_indices: Vec, + transformed_col: HashMap>, schema: SchemaRef, options: UnnestOptions, ) -> Self { @@ -88,6 +91,7 @@ impl UnnestExec { UnnestExec { input, schema, + transformed_col, list_column_indices, struct_column_indices, options, @@ -150,6 +154,7 @@ impl ExecutionPlan for UnnestExec { Arc::clone(&children[0]), self.list_column_indices.clone(), self.struct_column_indices.clone(), + self.transformed_col.clone(), Arc::clone(&self.schema), self.options.clone(), ))) @@ -169,6 +174,7 @@ impl ExecutionPlan for UnnestExec { Ok(Box::pin(UnnestStream { input, + transformed_col: self.transformed_col, schema: Arc::clone(&self.schema), list_type_columns: self.list_column_indices.clone(), struct_column_indices: self.struct_column_indices.iter().copied().collect(), @@ -226,6 +232,8 @@ struct UnnestStream { input: SendableRecordBatchStream, /// Unnested schema schema: Arc, + /// Which original columns are transformed into which columns + transformed_col: HashMap>, /// The unnest columns list_type_columns: Vec, struct_column_indices: HashSet, From 7efee194c40b432010a5a5c99bd0c11d68a94bbb Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Sun, 28 Jul 2024 14:59:59 +0200 Subject: [PATCH 09/56] prototype for recursive array length --- datafusion/physical-plan/src/unnest.rs | 67 ++++++++++++++++++++++++-- 1 file changed, 64 insertions(+), 3 deletions(-) diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index df757a1d378c..6ea2580db0cd 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -34,7 +34,7 @@ use arrow::array::{ use arrow::compute::kernels::length::length; use arrow::compute::kernels::zip::zip; use arrow::compute::{cast, is_not_null, kernels, sum}; -use arrow::datatypes::{DataType, Int64Type, Schema, SchemaRef}; +use arrow::datatypes::{DataType, Int32Type, Int64Type, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use arrow_array::{Int64Array, Scalar, StructArray}; use arrow_ord::cmp::lt; @@ -407,6 +407,31 @@ fn build_batch( transformed } +/// TODO: only prototype +/// This function does not handle NULL yet +pub fn length_recursive(array: &dyn Array, depth: usize) -> Result { + let list = array.as_list::(); + if depth == 1 { + return Ok(length(array)?); + } + let a: Vec = list + .iter() + .map(|x| { + if x.is_none() { + return 0; + } + let ret = length_recursive(&x.unwrap(), depth - 1).unwrap(); + let a = ret.as_primitive::(); + if a.is_empty() { + 0 + } else { + sum(a).unwrap() + } + }) + .collect(); + Ok(Arc::new(PrimitiveArray::::from(a))) +} + /// Find the longest list length among the given list arrays for each row. /// /// For example if we have the following two list arrays: @@ -439,11 +464,18 @@ fn find_longest_length( } else { Scalar::new(Int64Array::from_value(0, 1)) }; + // col1: [[1,2]]|[[2,3]] + // unnest(col1): [1,2]|[2,3] => length = 2 + // unnest(unnest(col1)): 1|2|2|3 => length = 4 + // col2: [3]|[4] => length =2 + // unnest(col2): 3|4 => length = 2 + // unnest(col1), unnest(unnest(col1)), unnest(col2) let list_lengths: Vec = list_arrays .iter() .map(|list_array| { - let mut length_array = length(list_array)?; + let mut length_array = length_recursive(list_array, 1)?; // Make sure length arrays have the same type. Int64 is the most general one. + // Respect the depth of unnest( current func only get the length of 1 level of unnest) length_array = cast(&length_array, &DataType::Int64)?; length_array = zip(&is_not_null(&length_array)?, &length_array, &null_length)?; @@ -667,7 +699,7 @@ fn flatten_list_cols_from_indices( #[cfg(test)] mod tests { use super::*; - use arrow::datatypes::Field; + use arrow::datatypes::{Field, Int32Type}; use arrow_array::{GenericListArray, OffsetSizeTrait, StringArray}; use arrow_buffer::{BooleanBufferBuilder, NullBuffer, OffsetBuffer}; @@ -753,6 +785,35 @@ mod tests { Ok(()) } + #[test] + fn test_length_recursive() -> datafusion_common::Result<()> { + // [[1,2,3],null,[4,5]] + let list_arr = ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(2), Some(3)]), + None, + Some(vec![Some(4), Some(5)]), + Some(vec![Some(7), Some(8), Some(9), Some(10)]), + None, + Some(vec![Some(11), Some(12), Some(13)]), + ]); + let list_arr_ref = Arc::new(list_arr) as ArrayRef; + let offsets = OffsetBuffer::from_lengths([3, 3]); + let nested_list_arr = ListArray::new( + Arc::new(Field::new_list_field( + list_arr_ref.data_type().to_owned(), + true, + )), + offsets, + list_arr_ref, + None, + ); + let b = length_recursive(&(Arc::new(nested_list_arr) as ArrayRef), 2)?; + // [3, 0, 2] + // if depth = 2, length should be 5 + println!("{:?}", b); + Ok(()) + } + #[test] fn test_unnest_list_array() -> datafusion_common::Result<()> { // [A, B, C], [], NULL, [D], NULL, [NULL, F] From a48950262778fca67013ce56109cd220255a4166 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Tue, 30 Jul 2024 23:07:44 +0200 Subject: [PATCH 10/56] chore: some compile err --- datafusion/expr/src/logical_plan/builder.rs | 121 ++++++++++++++++---- datafusion/expr/src/logical_plan/plan.rs | 2 +- datafusion/physical-plan/src/unnest.rs | 53 +++++---- datafusion/sql/src/utils.rs | 4 +- 4 files changed, 132 insertions(+), 48 deletions(-) diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index 6a0fc4ba9d39..fef6cbec2364 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -1126,9 +1126,11 @@ impl LogicalPlanBuilder { column: impl Into, options: UnnestOptions, ) -> Result { + let default_mapping = + HashMap::from_iter(vec![(column.into(), vec![column.into()])]); Ok(Self::from(unnest_with_options( self.plan, - vec![column.into()], + default_mapping, options, )?)) } @@ -1594,9 +1596,41 @@ impl TableSource for LogicalTableSource { /// Create a [`LogicalPlan::Unnest`] plan pub fn unnest(input: LogicalPlan, columns: Vec) -> Result { - unnest_with_options(input, columns, UnnestOptions::default()) + let default_map = HashMap::from_iter(columns.iter().map(|c| { + ( + *c, + vec![UnnestType::List(UnnestList { + output_column: c.clone(), + depth: 1, + })], + ) + })); + unnest_with_options(input, default_map, UnnestOptions::default()) } +pub fn get_unnested_list_column_recursive( + data_type: &DataType, + depth: usize, +) -> Result { + match data_type { + DataType::List(field) + | DataType::FixedSizeList(field, _) + | DataType::LargeList(field) => { + if depth == 1 { + return Ok(field.data_type().clone()); + } + return get_unnested_list_column_recursive(field.data_type(), depth - 1); + } + _ => { + return internal_err!( + "trying to unnest on invalid data type {:?}", + data_type + ); + } + }; +} + +// TODO: make me recursive // Based on data type, either struct or a variant of list // return a set of columns as the result of unnesting // the input columns. @@ -1606,6 +1640,7 @@ pub fn unnest(input: LogicalPlan, columns: Vec) -> Result { pub fn get_unnested_columns( col_name: &String, data_type: &DataType, + depth: usize, ) -> Result)>> { let mut qualified_columns = Vec::with_capacity(1); @@ -1613,9 +1648,10 @@ pub fn get_unnested_columns( DataType::List(field) | DataType::FixedSizeList(field, _) | DataType::LargeList(field) => { + let data_type = get_unnested_list_column_recursive(data_type, depth)?; let new_field = Arc::new(Field::new( col_name.clone(), - field.data_type().clone(), + data_type, // Unnesting may produce NULLs even if the list is not null. // For example: unnset([1], []) -> 1, null true, @@ -1643,10 +1679,27 @@ pub fn get_unnested_columns( Ok(qualified_columns) } +pub enum UnnestType { + List(UnnestList), + Struct, +} + +pub struct UnnestList { + output_column: Column, + depth: usize, +} + /// Create a [`LogicalPlan::Unnest`] plan with options +/// This function receive a map of columns to be unnested +/// because multiple unnest can be performed on the same column (e.g unnest with different depth) +/// The new schema will contains post-unnest fields replacing the original field +/// +/// input schema as: col1: int| col2: [][]int +/// Then unnest_map with { col2 -> [(col2,depth=1), (col2,depth=2)] } +/// will generate a new schema as col1: int| unnest_col2_depth1: []int| unnest_col2_depth2: int pub fn unnest_with_options( input: LogicalPlan, - columns: HashMap>, + columns: HashMap>, options: UnnestOptions, ) -> Result { let mut list_columns = Vec::with_capacity(columns.len()); @@ -1670,30 +1723,48 @@ pub fn unnest_with_options( .map(|(index, (original_qualifier, original_field))| { match column_by_original_index.get(&index) { Some(&column_to_unnest) => { - let transformed_cols = columns.get(&column_to_unnest).unwrap(); - println!("columnto unnest name {}", column_to_unnest.name); - let flatten_columns = get_unnested_columns( - &column_to_unnest.name, - original_field.data_type(), - )?; - match original_field.data_type() { - DataType::List(_) - | DataType::FixedSizeList(_, _) - | DataType::LargeList(_) => list_columns.push(index), - DataType::Struct(_) => struct_columns.push(index), - _ => { - panic!( - "not reachable, should be caught by get_unnested_columns" - ) - } - } + let unnests_on_column = columns.get(&column_to_unnest).unwrap(); + let transformed_columns: Vec<(Column, Arc)> = + unnests_on_column + .iter() + .map(|unnest_type| match unnest_type { + UnnestType::Struct => get_unnested_columns( + &column_to_unnest.name, + original_field.data_type(), + 1, + ), + UnnestType::List(UnnestList { + output_column, + depth, + }) => get_unnested_columns( + &output_column.name, + original_field.data_type(), + *depth, + ), + }) + .collect::>>()? + .into_iter() + .flatten() + .collect(); + + // match original_field.data_type() { + // DataType::List(_) + // | DataType::FixedSizeList(_, _) + // | DataType::LargeList(_) => list_columns.push(index), + // DataType::Struct(_) => struct_columns.push(index), + // _ => { + // panic!( + // "not reachable, should be caught by get_unnested_columns" + // ) + // } + // } // new columns dependent on the same original index dependency_indices - .extend(std::iter::repeat(index).take(flatten_columns.len())); - Ok(flatten_columns + .extend(std::iter::repeat(index).take(transformed_columns.len())); + Ok(transformed_columns .iter() - .map(|col: &(Column, Arc)| { - (col.0.relation.to_owned(), col.1.to_owned()) + .map(|(col, data_type)| { + (col.relation.to_owned(), data_type.to_owned()) }) .collect()) } diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index 48fa6270b202..b2f5e0e4de5b 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -2887,7 +2887,7 @@ pub struct Unnest { pub exec_columns: Vec, /// refer to the indices(in the input schema) of columns /// that have type list to run unnest on - pub list_type_columns: Vec, + pub list_type_columns: Vec<(usize, usize)>, /// refer to the indices (in the input schema) of columns /// that have type struct to run unnest on pub struct_type_columns: Vec, diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index 6ea2580db0cd..63e280b60334 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -234,8 +234,10 @@ struct UnnestStream { schema: Arc, /// Which original columns are transformed into which columns transformed_col: HashMap>, - /// The unnest columns - list_type_columns: Vec, + /// represents all unnest operations to be applied to the input (input index, depth) + /// e.g unnest(col1),unnest(unnest(col1)) where col1 has index 1 in original input schema + /// then list_type_columns = [(1,1),(1,2)] + list_type_columns: Vec<(usize, usize)>, struct_column_indices: HashSet, /// Options options: UnnestOptions, @@ -355,18 +357,21 @@ fn flatten_struct_cols( fn build_batch( batch: &RecordBatch, schema: &SchemaRef, - list_type_columns: &[usize], + list_type_columns: &[(usize, usize)], struct_column_indices: &HashSet, options: &UnnestOptions, ) -> Result { let transformed = match list_type_columns.len() { 0 => flatten_struct_cols(batch.columns(), schema, struct_column_indices), _ => { - let list_arrays: Vec = list_type_columns + let list_arrays: Vec<(ArrayRef, usize)> = list_type_columns .iter() - .map(|index| { - ColumnarValue::Array(Arc::clone(batch.column(*index))) - .into_array(batch.num_rows()) + .map(|(index, depth)| { + Ok(( + ColumnarValue::Array(Arc::clone(batch.column(*index))) + .into_array(batch.num_rows())?, + *depth, + )) }) .collect::>()?; @@ -389,7 +394,7 @@ fn build_batch( let unnested_array_map: HashMap<_, _> = unnested_arrays .into_iter() .zip(list_type_columns.iter()) - .map(|(array, column)| (*column, array)) + .map(|(array, (column, depth))| (*column, (array, depth))) .collect(); // Create the take indices array for other columns @@ -455,7 +460,7 @@ pub fn length_recursive(array: &dyn Array, depth: usize) -> Result { /// ``` /// fn find_longest_length( - list_arrays: &[ArrayRef], + list_arrays: &[(ArrayRef, usize)], options: &UnnestOptions, ) -> Result { // The length of a NULL list @@ -472,8 +477,8 @@ fn find_longest_length( // unnest(col1), unnest(unnest(col1)), unnest(col2) let list_lengths: Vec = list_arrays .iter() - .map(|list_array| { - let mut length_array = length_recursive(list_array, 1)?; + .map(|(list_array, depth)| { + let mut length_array = length_recursive(list_array, *depth)?; // Make sure length arrays have the same type. Int64 is the most general one. // Respect the depth of unnest( current func only get the length of 1 level of unnest) length_array = cast(&length_array, &DataType::Int64)?; @@ -537,19 +542,21 @@ impl ListArrayType for FixedSizeListArray { /// Unnest multiple list arrays according to the length array. fn unnest_list_arrays( - list_arrays: &[ArrayRef], + list_arrays: &[(ArrayRef, usize)], length_array: &PrimitiveArray, capacity: usize, ) -> Result> { let typed_arrays = list_arrays .iter() - .map(|list_array| match list_array.data_type() { - DataType::List(_) => Ok(list_array.as_list::() as &dyn ListArrayType), + .map(|(list_array, depth)| match list_array.data_type() { + DataType::List(_) => { + Ok((list_array.as_list::() as &dyn ListArrayType, depth)) + } DataType::LargeList(_) => { - Ok(list_array.as_list::() as &dyn ListArrayType) + Ok((list_array.as_list::() as &dyn ListArrayType, depth)) } DataType::FixedSizeList(_, _) => { - Ok(list_array.as_fixed_size_list() as &dyn ListArrayType) + Ok((list_array.as_fixed_size_list() as &dyn ListArrayType, depth)) } other => exec_err!("Invalid unnest datatype {other }"), }) @@ -557,7 +564,9 @@ fn unnest_list_arrays( typed_arrays .iter() - .map(|list_array| unnest_list_array(*list_array, length_array, capacity)) + .map(|(list_array, depth)| { + unnest_list_array(*list_array, length_array, capacity, **depth) + }) .collect::>() } @@ -586,7 +595,9 @@ fn unnest_list_array( list_array: &dyn ListArrayType, length_array: &PrimitiveArray, capacity: usize, + depth: usize, ) -> Result { + let st = list_array.as_list(); let values = list_array.values(); let mut take_indicies_builder = PrimitiveArray::::builder(capacity); for row in 0..list_array.len() { @@ -681,7 +692,7 @@ fn create_take_indicies( /// fn flatten_list_cols_from_indices( batch: &RecordBatch, - unnested_list_arrays: &HashMap, + unnested_list_arrays: &HashMap, indices: &PrimitiveArray, ) -> Result>> { let arrays = batch @@ -689,7 +700,7 @@ fn flatten_list_cols_from_indices( .iter() .enumerate() .map(|(col_idx, arr)| match unnested_list_arrays.get(&col_idx) { - Some(unnested_array) => Ok(Arc::clone(unnested_array)), + Some((unnested_array, depth)) => Ok(Arc::clone(unnested_array)), None => Ok(kernels::take::take(arr, indices, None)?), }) .collect::>>()?; @@ -779,7 +790,7 @@ mod tests { expected: Vec>, ) -> datafusion_common::Result<()> { let length_array = Int64Array::from(lengths); - let unnested_array = unnest_list_array(list_array, &length_array, 3 * 6)?; + let unnested_array = unnest_list_array(list_array, &length_array, 3 * 6, 1)?; let strs = unnested_array.as_string::().iter().collect::>(); assert_eq!(strs, expected); Ok(()) @@ -860,7 +871,7 @@ mod tests { } fn verify_longest_length( - list_arrays: &[ArrayRef], + list_arrays: &[(ArrayRef, usize)], preserve_nulls: bool, expected: Vec, ) -> datafusion_common::Result<()> { diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index ad7a7fc6304d..eea25054f014 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -327,8 +327,10 @@ pub(crate) fn transform_bottom_unnest( } } + // TODO: refactor get_unnested_columns or use a different function + // depth = 1 has no meaning here, just to provide enough argument let outer_projection_columns = - get_unnested_columns(&post_unnest_name, &data_type)?; + get_unnested_columns(&post_unnest_name, &data_type, 1)?; let expr = outer_projection_columns .iter() .map(|col| { From 80c1cf3baf365d0b49e7618753bef3a1186a8b4b Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Wed, 31 Jul 2024 22:48:57 +0200 Subject: [PATCH 11/56] finalize input type in physical layer --- datafusion/physical-plan/src/unnest.rs | 67 +++++++++++++++++++------- 1 file changed, 49 insertions(+), 18 deletions(-) diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index 63e280b60334..422d2b0f8396 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -48,6 +48,7 @@ use datafusion_physical_expr::EquivalenceProperties; use async_trait::async_trait; use futures::{Stream, StreamExt}; use hashbrown::HashSet; +use itertools::Itertools; use log::trace; /// Unnest the given columns (either with type struct or list) @@ -349,6 +350,11 @@ fn flatten_struct_cols( Ok(RecordBatch::try_new(Arc::clone(schema), columns_expanded)?) } +struct ListUnnest { + index_in_input_schema: usize, + depth: usize, +} + /// For each row in a `RecordBatch`, some list/struct columns need to be unnested. /// - For list columns: We will expand the values in each list into multiple rows, /// taking the longest length among these lists, and shorter lists are padded with NULLs. @@ -357,7 +363,7 @@ fn flatten_struct_cols( fn build_batch( batch: &RecordBatch, schema: &SchemaRef, - list_type_columns: &[(usize, usize)], + list_type_columns: &[ListUnnest], struct_column_indices: &HashSet, options: &UnnestOptions, ) -> Result { @@ -366,13 +372,20 @@ fn build_batch( _ => { let list_arrays: Vec<(ArrayRef, usize)> = list_type_columns .iter() - .map(|(index, depth)| { - Ok(( - ColumnarValue::Array(Arc::clone(batch.column(*index))) + .map( + |ListUnnest { + depth, + index_in_input_schema, + }| { + Ok(( + ColumnarValue::Array(Arc::clone( + batch.column(*index_in_input_schema), + )) .into_array(batch.num_rows())?, - *depth, - )) - }) + *depth, + )) + }, + ) .collect::>()?; let longest_length = find_longest_length(&list_arrays, options)?; @@ -391,11 +404,25 @@ fn build_batch( // Unnest all the list arrays let unnested_arrays = unnest_list_arrays(&list_arrays, unnested_length, total_length)?; - let unnested_array_map: HashMap<_, _> = unnested_arrays + let unnested_array_map: HashMap>> = unnested_arrays .into_iter() .zip(list_type_columns.iter()) - .map(|(array, (column, depth))| (*column, (array, depth))) - .collect(); + .fold( + HashMap::new(), + |mut acc, + ( + flattened_array, + ListUnnest { + index_in_input_schema, + depth, + }, + )| { + acc.entry(index_in_input_schema) + .or_insert_with(vec![]) + .push(flattened_array); + acc + }, + ); // Create the take indices array for other columns let take_indicies = create_take_indicies(unnested_length, total_length); @@ -403,7 +430,7 @@ fn build_batch( // vertical expansion because of list unnest let ret = flatten_list_cols_from_indices( batch, - &unnested_array_map, + unnested_array_map, &take_indicies, )?; flatten_struct_cols(&ret, schema, struct_column_indices) @@ -692,18 +719,22 @@ fn create_take_indicies( /// fn flatten_list_cols_from_indices( batch: &RecordBatch, - unnested_list_arrays: &HashMap, + unnested_list_arrays: HashMap>, indices: &PrimitiveArray, ) -> Result>> { let arrays = batch .columns() - .iter() + .into_iter() .enumerate() - .map(|(col_idx, arr)| match unnested_list_arrays.get(&col_idx) { - Some((unnested_array, depth)) => Ok(Arc::clone(unnested_array)), - None => Ok(kernels::take::take(arr, indices, None)?), - }) - .collect::>>()?; + .map( + |(col_idx, arr)| match unnested_list_arrays.remove(&col_idx) { + Some(unnested_arrays) => Ok(unnested_arrays), + None => Ok(vec![kernels::take::take(arr, indices, None)?]), + }, + ) + .collect::>>()? + .flatten() + .collect::>(); Ok(arrays) } From 7acb0561874eff42d81e6a7f9a93e59865164ed3 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Thu, 1 Aug 2024 22:29:38 +0200 Subject: [PATCH 12/56] chore: refactor unnest builder --- datafusion/expr/src/logical_plan/builder.rs | 81 ++++++++++----------- datafusion/expr/src/logical_plan/plan.rs | 17 ++++- datafusion/physical-plan/src/unnest.rs | 14 +--- 3 files changed, 56 insertions(+), 56 deletions(-) diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index fef6cbec2364..75bab7e8df39 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -57,6 +57,8 @@ use datafusion_common::{ TableReference, ToDFSchema, UnnestOptions, }; +use super::plan::{ColumnUnnestList, ColumnUnnestType}; + /// Default table name for unnamed table pub const UNNAMED_TABLE: &str = "?table?"; @@ -1599,7 +1601,7 @@ pub fn unnest(input: LogicalPlan, columns: Vec) -> Result { let default_map = HashMap::from_iter(columns.iter().map(|c| { ( *c, - vec![UnnestType::List(UnnestList { + vec![UnnestType::List(ColumnUnnestList { output_column: c.clone(), depth: 1, })], @@ -1679,15 +1681,9 @@ pub fn get_unnested_columns( Ok(qualified_columns) } -pub enum UnnestType { - List(UnnestList), - Struct, -} - -pub struct UnnestList { - output_column: Column, - depth: usize, -} +// a list type column can be performed differently at the same time +// e.g select unnest(col), unnest(unnest(col)) +// while unnest struct can only be performed once at a time /// Create a [`LogicalPlan::Unnest`] plan with options /// This function receive a map of columns to be unnested @@ -1699,11 +1695,11 @@ pub struct UnnestList { /// will generate a new schema as col1: int| unnest_col2_depth1: []int| unnest_col2_depth2: int pub fn unnest_with_options( input: LogicalPlan, - columns: HashMap>, + columns: HashMap, options: UnnestOptions, ) -> Result { - let mut list_columns = Vec::with_capacity(columns.len()); - let mut struct_columns = Vec::with_capacity(columns.len()); + let mut list_columns = HashMap::default(); + let mut struct_columns = HashSet::default(); let column_by_original_index = columns .iter() .map(|(inner_col, outer_cols)| { @@ -1725,39 +1721,38 @@ pub fn unnest_with_options( Some(&column_to_unnest) => { let unnests_on_column = columns.get(&column_to_unnest).unwrap(); let transformed_columns: Vec<(Column, Arc)> = - unnests_on_column - .iter() - .map(|unnest_type| match unnest_type { - UnnestType::Struct => get_unnested_columns( + match unnests_on_column { + ColumnUnnestType::Struct => { + struct_columns.insert(index); + get_unnested_columns( &column_to_unnest.name, original_field.data_type(), 1, - ), - UnnestType::List(UnnestList { - output_column, - depth, - }) => get_unnested_columns( - &output_column.name, - original_field.data_type(), - *depth, - ), - }) - .collect::>>()? - .into_iter() - .flatten() - .collect(); - - // match original_field.data_type() { - // DataType::List(_) - // | DataType::FixedSizeList(_, _) - // | DataType::LargeList(_) => list_columns.push(index), - // DataType::Struct(_) => struct_columns.push(index), - // _ => { - // panic!( - // "not reachable, should be caught by get_unnested_columns" - // ) - // } - // } + )? + } + ColumnUnnestType::List(unnest_lists) => { + list_columns.insert(index, unnest_lists.clone()); + unnest_lists + .iter() + .map( + |ColumnUnnestList { + output_column, + depth, + }| { + get_unnested_columns( + &output_column.name, + original_field.data_type(), + *depth, + ) + }, + ) + // TODO: i'm messy + .collect::)>>>>()? + .into_iter() + .flatten() + .collect::>() + } + }; // new columns dependent on the same original index dependency_indices .extend(std::iter::repeat(index).take(transformed_columns.len())); diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index b2f5e0e4de5b..c04ace2509a5 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -22,6 +22,7 @@ use std::fmt::{self, Debug, Display, Formatter}; use std::hash::{Hash, Hasher}; use std::sync::Arc; +use super::builder::UnnestList; use super::dml::CopyTo; use super::DdlStatement; use crate::builder::{change_redundant_column, unnest_with_options}; @@ -2877,6 +2878,16 @@ pub enum Partitioning { DistributeBy(Vec), } +pub enum ColumnUnnestType { + List(Vec), + Struct, +} + +pub struct ColumnUnnestList { + pub output_column: Column, + pub depth: usize, +} + /// Unnest a column that contains a nested list type. See /// [`UnnestOptions`] for more details. #[derive(Debug, Clone, PartialEq, Eq, Hash)] @@ -2884,13 +2895,13 @@ pub struct Unnest { /// The incoming logical plan pub input: Arc, /// Columns to run unnest on, can be a list of (List/Struct) columns - pub exec_columns: Vec, + pub exec_columns: HashMap, /// refer to the indices(in the input schema) of columns /// that have type list to run unnest on - pub list_type_columns: Vec<(usize, usize)>, + pub list_type_columns: HashSet>, /// refer to the indices (in the input schema) of columns /// that have type struct to run unnest on - pub struct_type_columns: Vec, + pub struct_type_columns: HashSet, /// Having items aligned with the output columns /// representing which column in the input schema each output column depends on pub dependency_indices: Vec, diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index 422d2b0f8396..af2e9b537430 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -42,6 +42,7 @@ use datafusion_common::{ exec_datafusion_err, exec_err, internal_err, Result, UnnestOptions, }; use datafusion_execution::TaskContext; +use datafusion_expr::builder::UnnestList; use datafusion_expr::ColumnarValue; use datafusion_physical_expr::EquivalenceProperties; @@ -63,10 +64,8 @@ pub struct UnnestExec { input: Arc, /// The schema once the unnest is applied schema: SchemaRef, - /// which original columns are transformed into which columns - transformed_col: HashMap>, /// indices of the list-typed columns in the input schema - list_column_indices: Vec, + list_column_indices: Vec, /// indices of the struct-typed columns in the input schema struct_column_indices: Vec, /// Options @@ -83,7 +82,6 @@ impl UnnestExec { input: Arc, list_column_indices: Vec, struct_column_indices: Vec, - transformed_col: HashMap>, schema: SchemaRef, options: UnnestOptions, ) -> Self { @@ -92,7 +90,6 @@ impl UnnestExec { UnnestExec { input, schema, - transformed_col, list_column_indices, struct_column_indices, options, @@ -175,7 +172,6 @@ impl ExecutionPlan for UnnestExec { Ok(Box::pin(UnnestStream { input, - transformed_col: self.transformed_col, schema: Arc::clone(&self.schema), list_type_columns: self.list_column_indices.clone(), struct_column_indices: self.struct_column_indices.iter().copied().collect(), @@ -233,12 +229,10 @@ struct UnnestStream { input: SendableRecordBatchStream, /// Unnested schema schema: Arc, - /// Which original columns are transformed into which columns - transformed_col: HashMap>, /// represents all unnest operations to be applied to the input (input index, depth) /// e.g unnest(col1),unnest(unnest(col1)) where col1 has index 1 in original input schema /// then list_type_columns = [(1,1),(1,2)] - list_type_columns: Vec<(usize, usize)>, + list_type_columns: Vec, struct_column_indices: HashSet, /// Options options: UnnestOptions, @@ -350,7 +344,7 @@ fn flatten_struct_cols( Ok(RecordBatch::try_new(Arc::clone(schema), columns_expanded)?) } -struct ListUnnest { +pub struct ListUnnest { index_in_input_schema: usize, depth: usize, } From 4d321875f4aa06a9b1e2b0a1ed2d923d98778191 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Sat, 3 Aug 2024 14:26:03 +0200 Subject: [PATCH 13/56] add unnesting type inferred --- datafusion/expr/src/logical_plan/builder.rs | 96 +++++++++++++------ datafusion/expr/src/logical_plan/display.rs | 3 +- datafusion/expr/src/logical_plan/plan.rs | 26 +++-- datafusion/expr/src/logical_plan/tree_node.rs | 2 +- datafusion/physical-plan/src/unnest.rs | 3 +- 5 files changed, 89 insertions(+), 41 deletions(-) diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index 75bab7e8df39..253dda3c7d95 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -1128,11 +1128,9 @@ impl LogicalPlanBuilder { column: impl Into, options: UnnestOptions, ) -> Result { - let default_mapping = - HashMap::from_iter(vec![(column.into(), vec![column.into()])]); Ok(Self::from(unnest_with_options( self.plan, - default_mapping, + vec![(column.into(), ColumnUnnestType::Inferred)], options, )?)) } @@ -1140,7 +1138,7 @@ impl LogicalPlanBuilder { /// Unnest the given columns with the given [`UnnestOptions`] pub fn unnest_columns_with_options( self, - columns: HashMap>, + columns: Vec<(Column, ColumnUnnestType)>, options: UnnestOptions, ) -> Result { Ok(Self::from(unnest_with_options( @@ -1598,16 +1596,11 @@ impl TableSource for LogicalTableSource { /// Create a [`LogicalPlan::Unnest`] plan pub fn unnest(input: LogicalPlan, columns: Vec) -> Result { - let default_map = HashMap::from_iter(columns.iter().map(|c| { - ( - *c, - vec![UnnestType::List(ColumnUnnestList { - output_column: c.clone(), - depth: 1, - })], - ) - })); - unnest_with_options(input, default_map, UnnestOptions::default()) + let unnestings = columns + .into_iter() + .map(|c| (c, ColumnUnnestType::Inferred)) + .collect(); + unnest_with_options(input, unnestings, UnnestOptions::default()) } pub fn get_unnested_list_column_recursive( @@ -1632,6 +1625,31 @@ pub fn get_unnested_list_column_recursive( }; } +fn get_unnested_columns_inferred( + col_name: &String, + data_type: &DataType, +) -> Result { + match data_type { + DataType::List(field) + | DataType::FixedSizeList(field, _) + | DataType::LargeList(field) => { + return Ok(ColumnUnnestType::List(vec![ColumnUnnestList { + output_column: Column::from_name(col_name), + depth: 1, + }])); + } + DataType::Struct(fields) => { + return Ok(ColumnUnnestType::Struct); + } + _ => { + return internal_err!( + "trying to unnest on invalid data type {:?}", + data_type + ); + } + }; +} + // TODO: make me recursive // Based on data type, either struct or a variant of list // return a set of columns as the result of unnesting @@ -1695,17 +1713,20 @@ pub fn get_unnested_columns( /// will generate a new schema as col1: int| unnest_col2_depth1: []int| unnest_col2_depth2: int pub fn unnest_with_options( input: LogicalPlan, - columns: HashMap, + columns: Vec<(Column, ColumnUnnestType)>, options: UnnestOptions, ) -> Result { - let mut list_columns = HashMap::default(); - let mut struct_columns = HashSet::default(); + let mut list_columns: Vec<(usize, ColumnUnnestList)> = vec![]; + let mut struct_columns = vec![]; let column_by_original_index = columns .iter() - .map(|(inner_col, outer_cols)| { - Ok((input.schema().index_of_column(inner_col)?, inner_col)) + .map(|col_unnesting| { + Ok(( + input.schema().index_of_column(&col_unnesting.0)?, + col_unnesting, + )) }) - .collect::>>()?; + .collect::>>()?; let input_schema = input.schema(); @@ -1718,12 +1739,18 @@ pub fn unnest_with_options( .enumerate() .map(|(index, (original_qualifier, original_field))| { match column_by_original_index.get(&index) { - Some(&column_to_unnest) => { - let unnests_on_column = columns.get(&column_to_unnest).unwrap(); + Some((column_to_unnest, unnest_type)) => { + let mut inferred_unnest_type = *unnest_type; + if let ColumnUnnestType::Inferred = unnest_type { + inferred_unnest_type = get_unnested_columns_inferred( + &column_to_unnest.name, + original_field.data_type(), + )?; + } let transformed_columns: Vec<(Column, Arc)> = - match unnests_on_column { + match inferred_unnest_type { ColumnUnnestType::Struct => { - struct_columns.insert(index); + struct_columns.push(index); get_unnested_columns( &column_to_unnest.name, original_field.data_type(), @@ -1731,7 +1758,11 @@ pub fn unnest_with_options( )? } ColumnUnnestType::List(unnest_lists) => { - list_columns.insert(index, unnest_lists.clone()); + list_columns.extend( + unnest_lists + .iter() + .map(|ul| (index, ul.to_owned().clone())), + ); unnest_lists .iter() .map( @@ -1752,6 +1783,7 @@ pub fn unnest_with_options( .flatten() .collect::>() } + _ => internal_err!("Invalid unnest type"), }; // new columns dependent on the same original index dependency_indices @@ -2227,14 +2259,24 @@ mod tests { // Unnesting multiple fields at the same time let cols = vec!["strings", "structs", "struct_singular"] .into_iter() - .map(|c| c.into()) + .map(|c| match c { + "struct_singular" => (Column::from(c), ColumnUnnestType::Struct), + _ => ( + Column::from(c), + ColumnUnnestType::List(vec![ColumnUnnestList { + output_column: Column::from(c), + // TODO: try planning with depth > 1 + depth: 1, + }]), + ), + }) .collect(); let plan = nested_table_scan("test_table")? .unnest_columns_with_options(cols, UnnestOptions::default())? .build()?; let expected = "\ - Unnest: lists[test_table.strings, test_table.structs] structs[test_table.struct_singular]\ + Unnest: lists[] structs[test_table.struct_singular]\ \n TableScan: test_table"; assert_eq!(expected, format!("{plan:?}")); diff --git a/datafusion/expr/src/logical_plan/display.rs b/datafusion/expr/src/logical_plan/display.rs index 343eda056ffe..aa183f661eb3 100644 --- a/datafusion/expr/src/logical_plan/display.rs +++ b/datafusion/expr/src/logical_plan/display.rs @@ -647,7 +647,8 @@ impl<'a, 'b> PgJsonVisitor<'a, 'b> { let input_columns = plan.schema().columns(); let list_type_columns = list_col_indices .iter() - .map(|i| &input_columns[*i]) + // TODO: fix me + .map(|(i, _)| &input_columns[*i]) .collect::>(); let struct_type_columns = struct_col_indices .iter() diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index c04ace2509a5..708b6199a9c2 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -17,12 +17,6 @@ //! Logical plan types -use std::collections::{HashMap, HashSet}; -use std::fmt::{self, Debug, Display, Formatter}; -use std::hash::{Hash, Hasher}; -use std::sync::Arc; - -use super::builder::UnnestList; use super::dml::CopyTo; use super::DdlStatement; use crate::builder::{change_redundant_column, unnest_with_options}; @@ -40,6 +34,10 @@ use crate::{ CreateMemoryTable, CreateView, Expr, ExprSchemable, LogicalPlanBuilder, Operator, TableProviderFilterPushDown, TableSource, WindowFunctionDefinition, }; +use std::collections::{HashMap, HashSet}; +use std::fmt::{self, Debug, Display, Formatter}; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; @@ -1865,7 +1863,8 @@ impl LogicalPlan { let input_columns = plan.schema().columns(); let list_type_columns = list_col_indices .iter() - .map(|i| &input_columns[*i]) + //TODO: fixme: add depth + .map(|(i,_)| &input_columns[*i]) .collect::>(); let struct_type_columns = struct_col_indices .iter() @@ -2878,11 +2877,18 @@ pub enum Partitioning { DistributeBy(Vec), } +#[derive(Debug, Clone, PartialEq, Eq, Hash)] pub enum ColumnUnnestType { List(Vec), + // for struct, there can only be one unnest performed on one column at a time Struct, + // Infer the unnest type based on column schema + // If column is a list column, the unnest depth will be 1 + // This value is to support sugar syntax of old api (unnest(columns1,columns2)) + Inferred, } +#[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct ColumnUnnestList { pub output_column: Column, pub depth: usize, @@ -2895,13 +2901,13 @@ pub struct Unnest { /// The incoming logical plan pub input: Arc, /// Columns to run unnest on, can be a list of (List/Struct) columns - pub exec_columns: HashMap, + pub exec_columns: Vec<(Column, ColumnUnnestType)>, /// refer to the indices(in the input schema) of columns /// that have type list to run unnest on - pub list_type_columns: HashSet>, + pub list_type_columns: Vec<(usize, ColumnUnnestList)>, /// refer to the indices (in the input schema) of columns /// that have type struct to run unnest on - pub struct_type_columns: HashSet, + pub struct_type_columns: Vec, /// Having items aligned with the output columns /// representing which column in the input schema each output column depends on pub dependency_indices: Vec, diff --git a/datafusion/expr/src/logical_plan/tree_node.rs b/datafusion/expr/src/logical_plan/tree_node.rs index a47906f20322..86e7740696c2 100644 --- a/datafusion/expr/src/logical_plan/tree_node.rs +++ b/datafusion/expr/src/logical_plan/tree_node.rs @@ -494,7 +494,7 @@ impl LogicalPlan { let exprs = columns .iter() - .map(|c| Expr::Column(c.clone())) + .map(|(c, unnest_type)| Expr::Column(c.clone())) .collect::>(); exprs.iter().apply_until_stop(f) } diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index af2e9b537430..ac872f351a50 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -65,7 +65,7 @@ pub struct UnnestExec { /// The schema once the unnest is applied schema: SchemaRef, /// indices of the list-typed columns in the input schema - list_column_indices: Vec, + list_column_indices: Vec, /// indices of the struct-typed columns in the input schema struct_column_indices: Vec, /// Options @@ -152,7 +152,6 @@ impl ExecutionPlan for UnnestExec { Arc::clone(&children[0]), self.list_column_indices.clone(), self.struct_column_indices.clone(), - self.transformed_col.clone(), Arc::clone(&self.schema), self.options.clone(), ))) From e41923f65a3bcac99de5155d487ba5e69f8e932f Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Sat, 3 Aug 2024 20:01:19 +0200 Subject: [PATCH 14/56] fix compile err --- datafusion/expr/src/logical_plan/builder.rs | 32 +++++++-------------- datafusion/physical-plan/src/unnest.rs | 15 +++++----- datafusion/sql/src/utils.rs | 26 ++++++++++------- 3 files changed, 34 insertions(+), 39 deletions(-) diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index 253dda3c7d95..7052120f99de 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -1616,13 +1616,10 @@ pub fn get_unnested_list_column_recursive( } return get_unnested_list_column_recursive(field.data_type(), depth - 1); } - _ => { - return internal_err!( - "trying to unnest on invalid data type {:?}", - data_type - ); - } + _ => {} }; + + internal_err!("trying to unnest on invalid data type {:?}", data_type) } fn get_unnested_columns_inferred( @@ -1630,24 +1627,17 @@ fn get_unnested_columns_inferred( data_type: &DataType, ) -> Result { match data_type { - DataType::List(field) - | DataType::FixedSizeList(field, _) - | DataType::LargeList(field) => { - return Ok(ColumnUnnestType::List(vec![ColumnUnnestList { + DataType::List(_) | DataType::FixedSizeList(_, _) | DataType::LargeList(_) => { + Ok(ColumnUnnestType::List(vec![ColumnUnnestList { output_column: Column::from_name(col_name), depth: 1, - }])); - } - DataType::Struct(fields) => { - return Ok(ColumnUnnestType::Struct); + }])) } + DataType::Struct(_) => Ok(ColumnUnnestType::Struct), _ => { - return internal_err!( - "trying to unnest on invalid data type {:?}", - data_type - ); + internal_err!("trying to unnest on invalid data type {:?}", data_type) } - }; + } } // TODO: make me recursive @@ -1740,7 +1730,7 @@ pub fn unnest_with_options( .map(|(index, (original_qualifier, original_field))| { match column_by_original_index.get(&index) { Some((column_to_unnest, unnest_type)) => { - let mut inferred_unnest_type = *unnest_type; + let mut inferred_unnest_type = unnest_type.clone(); if let ColumnUnnestType::Inferred = unnest_type { inferred_unnest_type = get_unnested_columns_inferred( &column_to_unnest.name, @@ -1783,7 +1773,7 @@ pub fn unnest_with_options( .flatten() .collect::>() } - _ => internal_err!("Invalid unnest type"), + _ => return internal_err!("Invalid unnest type"), }; // new columns dependent on the same original index dependency_indices diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index ac872f351a50..a574d6b3224a 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -42,14 +42,12 @@ use datafusion_common::{ exec_datafusion_err, exec_err, internal_err, Result, UnnestOptions, }; use datafusion_execution::TaskContext; -use datafusion_expr::builder::UnnestList; use datafusion_expr::ColumnarValue; use datafusion_physical_expr::EquivalenceProperties; use async_trait::async_trait; use futures::{Stream, StreamExt}; use hashbrown::HashSet; -use itertools::Itertools; use log::trace; /// Unnest the given columns (either with type struct or list) @@ -80,7 +78,7 @@ impl UnnestExec { /// Create a new [UnnestExec]. pub fn new( input: Arc, - list_column_indices: Vec, + list_column_indices: Vec, struct_column_indices: Vec, schema: SchemaRef, options: UnnestOptions, @@ -343,6 +341,7 @@ fn flatten_struct_cols( Ok(RecordBatch::try_new(Arc::clone(schema), columns_expanded)?) } +#[derive(Debug, Clone)] pub struct ListUnnest { index_in_input_schema: usize, depth: usize, @@ -410,8 +409,8 @@ fn build_batch( depth, }, )| { - acc.entry(index_in_input_schema) - .or_insert_with(vec![]) + acc.entry(*index_in_input_schema) + .or_insert(vec![]) .push(flattened_array); acc }, @@ -617,7 +616,6 @@ fn unnest_list_array( capacity: usize, depth: usize, ) -> Result { - let st = list_array.as_list(); let values = list_array.values(); let mut take_indicies_builder = PrimitiveArray::::builder(capacity); for row in 0..list_array.len() { @@ -712,12 +710,12 @@ fn create_take_indicies( /// fn flatten_list_cols_from_indices( batch: &RecordBatch, - unnested_list_arrays: HashMap>, + mut unnested_list_arrays: HashMap>, indices: &PrimitiveArray, ) -> Result>> { let arrays = batch .columns() - .into_iter() + .iter() .enumerate() .map( |(col_idx, arr)| match unnested_list_arrays.remove(&col_idx) { @@ -726,6 +724,7 @@ fn flatten_list_cols_from_indices( }, ) .collect::>>()? + .into_iter() .flatten() .collect::>(); Ok(arrays) diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index eea25054f014..158e20ca0e56 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -483,11 +483,16 @@ pub(crate) fn transform_bottom_unnest( // write test for recursive_transform_unnest #[cfg(test)] mod tests { - use std::{collections::HashSet, ops::Add, sync::Arc}; + use std::{ + collections::{HashMap, HashSet}, + ops::Add, + str::FromStr, + sync::Arc, + }; use arrow::datatypes::{DataType as ArrowDataType, Field, Schema}; use arrow_schema::Fields; - use datafusion_common::{DFSchema, Result}; + use datafusion_common::{Column, DFSchema, Result}; use datafusion_expr::{col, lit, unnest, EmptyRelation, LogicalPlan}; use datafusion_functions::core::expr_ext::FieldAccessor; use datafusion_functions_aggregate::expr_fn::count; @@ -549,7 +554,7 @@ mod tests { let original_expr = unnest(unnest(col("3d_col"))) .add(unnest(unnest(col("3d_col")))) .add(col("i64_col")); - let mut memo = HashSet::new(); + let mut memo = HashMap::new(); let transformed_exprs = transform_bottom_unnest( &input, &mut unnest_placeholder_columns, @@ -566,7 +571,7 @@ mod tests { ); // memoization only contains 1 transformation assert_eq!(memo.len(), 1); - assert!(memo.get(&col("3d_col")).is_some()); + assert!(memo.get(&Column::from_name("3d_col")).is_some()); assert_eq!( unnest_placeholder_columns, vec!["unnest_placeholder(3d_col)"] @@ -598,7 +603,7 @@ mod tests { // memoization still contains 1 transformation // and the previous transformation is reused assert_eq!(memo.len(), 1); - assert!(memo.get(&col("3d_col")).is_some()); + assert!(memo.get(&Column::from_name("3d_col")).is_some()); assert_eq!( unnest_placeholder_columns, vec!["unnest_placeholder(3d_col)"] @@ -635,7 +640,8 @@ mod tests { // memoization still contains 1 transformation // and the previous transformation is reused assert_eq!(memo.len(), 2); - assert!(memo.get(&col("struct_arr_col")).is_some()); + + assert!(memo.get(&Column::from_name("struct_arr_col")).is_some()); assert_eq!( unnest_placeholder_columns, vec![ @@ -690,7 +696,7 @@ mod tests { let original_expr = unnest(unnest(col("3d_col"))) .add(unnest(unnest(col("3d_col")))) .add(col("i64_col")); - let mut memo = HashSet::new(); + let mut memo = HashMap::new(); let transformed_exprs = transform_bottom_unnest( &input, &mut unnest_placeholder_columns, @@ -707,7 +713,7 @@ mod tests { ); // memoization only contains 1 transformation assert_eq!(memo.len(), 1); - assert!(memo.get(&col("3d_col")).is_some()); + assert!(memo.get(&Column::from_name("3d_col")).is_some()); assert_eq!( unnest_placeholder_columns, vec!["unnest_placeholder(3d_col)"] @@ -739,7 +745,7 @@ mod tests { // memoization still contains 1 transformation // and the previous transformation is reused assert_eq!(memo.len(), 1); - assert!(memo.get(&col("3d_col")).is_some()); + assert!(memo.get(&Column::from_name("3d_col")).is_some()); assert_eq!( unnest_placeholder_columns, vec!["unnest_placeholder(3d_col)"] @@ -790,7 +796,7 @@ mod tests { let mut unnest_placeholder_columns = vec![]; let mut inner_projection_exprs = vec![]; - let mut memo = HashSet::new(); + let mut memo = HashMap::new(); // unnest(struct_col) let original_expr = unnest(col("struct_col")); let transformed_exprs = transform_bottom_unnest( From eca58f5631a132a68ef5d8ec30b4b62c03990484 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Sat, 3 Aug 2024 20:32:08 +0200 Subject: [PATCH 15/56] fail test in builder --- datafusion/expr/src/logical_plan/builder.rs | 23 ++++++-------------- datafusion/expr/src/logical_plan/display.rs | 10 +++++++-- datafusion/expr/src/logical_plan/plan.rs | 7 +++--- datafusion/optimizer/src/push_down_filter.rs | 2 +- 4 files changed, 20 insertions(+), 22 deletions(-) diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index 7052120f99de..ec376f1fc36a 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -2201,7 +2201,7 @@ mod tests { .build()?; let expected = "\ - Unnest: lists[test_table.strings] structs[]\ + Unnest: lists[test_table.strings|depth=1] structs[]\ \n TableScan: test_table"; assert_eq!(expected, format!("{plan:?}")); @@ -2237,8 +2237,8 @@ mod tests { let expected = "\ Unnest: lists[] structs[test_table.struct_singular]\ - \n Unnest: lists[test_table.structs] structs[]\ - \n Unnest: lists[test_table.strings] structs[]\ + \n Unnest: lists[test_table.structs|depth=1] structs[]\ + \n Unnest: lists[test_table.strings|depth=1] structs[]\ \n TableScan: test_table"; assert_eq!(expected, format!("{plan:?}")); @@ -2246,27 +2246,18 @@ mod tests { let field = plan.schema().field_with_name(None, "structs").unwrap(); assert!(matches!(field.data_type(), DataType::Struct(_))); - // Unnesting multiple fields at the same time + // Unnesting multiple fields at the same time, using infer syntax let cols = vec!["strings", "structs", "struct_singular"] .into_iter() - .map(|c| match c { - "struct_singular" => (Column::from(c), ColumnUnnestType::Struct), - _ => ( - Column::from(c), - ColumnUnnestType::List(vec![ColumnUnnestList { - output_column: Column::from(c), - // TODO: try planning with depth > 1 - depth: 1, - }]), - ), - }) + .map(|c| (Column::from(c), ColumnUnnestType::Inferred)) .collect(); + let plan = nested_table_scan("test_table")? .unnest_columns_with_options(cols, UnnestOptions::default())? .build()?; let expected = "\ - Unnest: lists[] structs[test_table.struct_singular]\ + Unnest: lists[test_table.strings|depth=1, test_table.structs|depth=1] structs[test_table.struct_singular]\ \n TableScan: test_table"; assert_eq!(expected, format!("{plan:?}")); diff --git a/datafusion/expr/src/logical_plan/display.rs b/datafusion/expr/src/logical_plan/display.rs index aa183f661eb3..732ff88cc42b 100644 --- a/datafusion/expr/src/logical_plan/display.rs +++ b/datafusion/expr/src/logical_plan/display.rs @@ -648,8 +648,14 @@ impl<'a, 'b> PgJsonVisitor<'a, 'b> { let list_type_columns = list_col_indices .iter() // TODO: fix me - .map(|(i, _)| &input_columns[*i]) - .collect::>(); + .map(|(i, unnest_info)| { + format!( + "{}|depth={:?}", + &input_columns[*i].to_string(), + unnest_info.depth + ) + }) + .collect::>(); let struct_type_columns = struct_col_indices .iter() .map(|i| &input_columns[*i]) diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index 708b6199a9c2..42dd7c0618df 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -1863,9 +1863,10 @@ impl LogicalPlan { let input_columns = plan.schema().columns(); let list_type_columns = list_col_indices .iter() - //TODO: fixme: add depth - .map(|(i,_)| &input_columns[*i]) - .collect::>(); + .map(|(i,unnest_info)| + format!("{}|depth={}", &input_columns[*i].to_string(), + unnest_info.depth)) + .collect::>(); let struct_type_columns = struct_col_indices .iter() .map(|i| &input_columns[*i]) diff --git a/datafusion/optimizer/src/push_down_filter.rs b/datafusion/optimizer/src/push_down_filter.rs index a22f2e83e211..f51648572586 100644 --- a/datafusion/optimizer/src/push_down_filter.rs +++ b/datafusion/optimizer/src/push_down_filter.rs @@ -745,7 +745,7 @@ impl OptimizerRule for PushDownFilter { let mut accum: HashSet = HashSet::new(); expr_to_columns(&predicate, &mut accum)?; - if unnest.exec_columns.iter().any(|c| accum.contains(c)) { + if unnest.exec_columns.iter().any(|(c, _)| accum.contains(c)) { unnest_predicates.push(predicate); } else { non_unnest_predicates.push(predicate); From e2fae71e52e68d5c750d24b568fcc97a359463b5 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Sun, 4 Aug 2024 12:03:45 +0200 Subject: [PATCH 16/56] Compile err --- datafusion/core/src/dataframe/mod.rs | 7 +- datafusion/core/src/physical_planner.rs | 10 +- datafusion/expr/src/logical_plan/builder.rs | 10 ++ datafusion/expr/src/logical_plan/mod.rs | 10 +- datafusion/physical-plan/src/unnest.rs | 49 ++++++-- datafusion/sql/src/select.rs | 2 +- datafusion/sql/src/utils.rs | 118 +++++++++++------- datafusion/sqllogictest/test_files/unnest.slt | 7 ++ 8 files changed, 147 insertions(+), 66 deletions(-) diff --git a/datafusion/core/src/dataframe/mod.rs b/datafusion/core/src/dataframe/mod.rs index fb28b5c1ab47..84795a3ebcb4 100644 --- a/datafusion/core/src/dataframe/mod.rs +++ b/datafusion/core/src/dataframe/mod.rs @@ -51,7 +51,7 @@ use datafusion_common::config::{CsvOptions, JsonOptions}; use datafusion_common::{ plan_err, Column, DFSchema, DataFusionError, ParamValues, SchemaError, UnnestOptions, }; -use datafusion_expr::{case, is_null, lit}; +use datafusion_expr::{case, is_null, lit, ColumnUnnestType}; use datafusion_expr::{ max, min, utils::COUNT_STAR_EXPANSION, TableProviderFilterPushDown, UNNAMED_TABLE, }; @@ -361,7 +361,10 @@ impl DataFrame { columns: &[&str], options: UnnestOptions, ) -> Result { - let columns = columns.iter().map(|c| Column::from(*c)).collect(); + let columns = columns + .iter() + .map(|c| (Column::from(*c), ColumnUnnestType::Inferred)) + .collect(); let plan = LogicalPlanBuilder::from(self.plan) .unnest_columns_with_options(columns, options)? .build()?; diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 97533cd5276a..80ff292caad4 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -87,6 +87,7 @@ use datafusion_expr::{ use datafusion_physical_expr::expressions::Literal; use datafusion_physical_expr::LexOrdering; use datafusion_physical_plan::placeholder_row::PlaceholderRowExec; +use datafusion_physical_plan::unnest::ListUnnest; use datafusion_sql::utils::window_expr_common_partition_keys; use async_trait::async_trait; @@ -1106,9 +1107,16 @@ impl DefaultPhysicalPlanner { }) => { let input = children.one()?; let schema = SchemaRef::new(schema.as_ref().to_owned().into()); + let list_column_indices = list_type_columns + .iter() + .map(|(index, unnesting)| ListUnnest { + index_in_input_schema: *index, + depth: unnesting.depth, + }) + .collect(); Arc::new(UnnestExec::new( input, - list_type_columns.clone(), + list_column_indices, struct_type_columns.clone(), schema, options.clone(), diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index ec376f1fc36a..8fab8c7a7dc8 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -1640,6 +1640,16 @@ fn get_unnested_columns_inferred( } } +pub fn get_struct_unnested_columns( + col_name: &String, + inner_fields: &Fields, +) -> Vec { + inner_fields + .iter() + .map(|f| Column::from_name(format!("{}.{}", col_name, f.name()))) + .collect() +} + // TODO: make me recursive // Based on data type, either struct or a variant of list // return a set of columns as the result of unnesting diff --git a/datafusion/expr/src/logical_plan/mod.rs b/datafusion/expr/src/logical_plan/mod.rs index 8928f70cd5d2..ead1f4088c50 100644 --- a/datafusion/expr/src/logical_plan/mod.rs +++ b/datafusion/expr/src/logical_plan/mod.rs @@ -35,11 +35,11 @@ pub use ddl::{ }; pub use dml::{DmlStatement, WriteOp}; pub use plan::{ - projection_schema, Aggregate, Analyze, CrossJoin, DescribeTable, Distinct, - DistinctOn, EmptyRelation, Explain, Extension, Filter, Join, JoinConstraint, - JoinType, Limit, LogicalPlan, Partitioning, PlanType, Prepare, Projection, - RecursiveQuery, Repartition, Sort, StringifiedPlan, Subquery, SubqueryAlias, - TableScan, ToStringifiedPlan, Union, Unnest, Values, Window, + projection_schema, Aggregate, Analyze, ColumnUnnestList, ColumnUnnestType, CrossJoin, + DescribeTable, Distinct, DistinctOn, EmptyRelation, Explain, Extension, Filter, Join, + JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, PlanType, Prepare, + Projection, RecursiveQuery, Repartition, Sort, StringifiedPlan, Subquery, + SubqueryAlias, TableScan, ToStringifiedPlan, Union, Unnest, Values, Window, }; pub use statement::{ SetVariable, Statement, TransactionAccessMode, TransactionConclusion, TransactionEnd, diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index a574d6b3224a..7b00156756c9 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -343,8 +343,8 @@ fn flatten_struct_cols( #[derive(Debug, Clone)] pub struct ListUnnest { - index_in_input_schema: usize, - depth: usize, + pub index_in_input_schema: usize, + pub depth: usize, } /// For each row in a `RecordBatch`, some list/struct columns need to be unnested. @@ -841,10 +841,13 @@ mod tests { list_arr_ref, None, ); - let b = length_recursive(&(Arc::new(nested_list_arr) as ArrayRef), 2)?; // [3, 0, 2] // if depth = 2, length should be 5 - println!("{:?}", b); + verify_longest_length( + &[(Arc::new(nested_list_arr) as ArrayRef, 2)], + true, + vec![5, 7], + ); Ok(()) } @@ -916,27 +919,51 @@ mod tests { // Test with single ListArray // [A, B, C], [], NULL, [D], NULL, [NULL, F] let list_array = Arc::new(make_generic_array::()) as ArrayRef; - verify_longest_length(&[Arc::clone(&list_array)], false, vec![3, 0, 0, 1, 0, 2])?; - verify_longest_length(&[Arc::clone(&list_array)], true, vec![3, 0, 1, 1, 1, 2])?; + verify_longest_length( + &[(Arc::clone(&list_array), 1)], + false, + vec![3, 0, 0, 1, 0, 2], + )?; + verify_longest_length( + &[(Arc::clone(&list_array), 1)], + true, + vec![3, 0, 1, 1, 1, 2], + )?; // Test with single LargeListArray // [A, B, C], [], NULL, [D], NULL, [NULL, F] let list_array = Arc::new(make_generic_array::()) as ArrayRef; - verify_longest_length(&[Arc::clone(&list_array)], false, vec![3, 0, 0, 1, 0, 2])?; - verify_longest_length(&[Arc::clone(&list_array)], true, vec![3, 0, 1, 1, 1, 2])?; + verify_longest_length( + &[(Arc::clone(&list_array), 1)], + false, + vec![3, 0, 0, 1, 0, 2], + )?; + verify_longest_length( + &[(Arc::clone(&list_array), 1)], + true, + vec![3, 0, 1, 1, 1, 2], + )?; // Test with single FixedSizeListArray // [A, B], NULL, [C, D], NULL, [NULL, F], [NULL, NULL] let list_array = Arc::new(make_fixed_list()) as ArrayRef; - verify_longest_length(&[Arc::clone(&list_array)], false, vec![2, 0, 2, 0, 2, 2])?; - verify_longest_length(&[Arc::clone(&list_array)], true, vec![2, 1, 2, 1, 2, 2])?; + verify_longest_length( + &[(Arc::clone(&list_array), 1)], + false, + vec![2, 0, 2, 0, 2, 2], + )?; + verify_longest_length( + &[(Arc::clone(&list_array), 1)], + true, + vec![2, 1, 2, 1, 2, 2], + )?; // Test with multiple list arrays // [A, B, C], [], NULL, [D], NULL, [NULL, F] // [A, B], NULL, [C, D], NULL, [NULL, F], [NULL, NULL] let list1 = Arc::new(make_generic_array::()) as ArrayRef; let list2 = Arc::new(make_fixed_list()) as ArrayRef; - let list_arrays = vec![Arc::clone(&list1), Arc::clone(&list2)]; + let list_arrays = vec![(Arc::clone(&list1), 1), (Arc::clone(&list2), 1)]; verify_longest_length(&list_arrays, false, vec![3, 0, 2, 1, 2, 2])?; verify_longest_length(&list_arrays, true, vec![3, 1, 2, 1, 2, 2])?; diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs index ab5ff7dbd499..5ddebd573cac 100644 --- a/datafusion/sql/src/select.rs +++ b/datafusion/sql/src/select.rs @@ -421,7 +421,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { // ``` let mut intermediate_plan = unwrap_arc(input); let mut intermediate_select_exprs = group_expr; - let mut memo = HashSet::new(); + let mut memo = HashMap::new(); loop { let mut unnest_columns = vec![]; diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index 158e20ca0e56..f63347165d5c 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -19,6 +19,7 @@ use std::cell::RefCell; use std::collections::{HashMap, HashSet}; +use std::vec; use arrow_schema::{ DataType, DECIMAL128_MAX_PRECISION, DECIMAL256_MAX_PRECISION, DECIMAL_DEFAULT_SCALE, @@ -26,14 +27,14 @@ use arrow_schema::{ use datafusion_common::tree_node::{ Transformed, TransformedResult, TreeNode, TreeNodeRecursion, }; -use datafusion_common::utils::proxy::VecAllocExt; use datafusion_common::{ exec_err, internal_err, plan_err, Column, DataFusionError, Result, ScalarValue, }; -use datafusion_expr::builder::get_unnested_columns; +use datafusion_expr::builder::get_struct_unnested_columns; use datafusion_expr::expr::{Alias, GroupingSet, Unnest, WindowFunction}; use datafusion_expr::utils::{expr_as_column_expr, find_column_exprs}; use datafusion_expr::{expr_vec_fmt, Expr, ExprSchemable, LogicalPlan}; +use datafusion_expr::{ColumnUnnestList, ColumnUnnestType}; use sqlparser::ast::Ident; /// Make a best-effort attempt at resolving all columns in the expression tree @@ -276,7 +277,7 @@ pub(crate) fn normalize_ident(id: Ident) -> String { /// is done only for the bottom expression pub(crate) fn transform_bottom_unnest( input: &LogicalPlan, - unnest_placeholder_columns: &mut Vec, + unnest_placeholder_columns: &mut Vec<(Column, ColumnUnnestType)>, inner_projection_exprs: &mut Vec, memo: &mut HashMap>, original_expr: &Expr, @@ -292,14 +293,7 @@ pub(crate) fn transform_bottom_unnest( return internal_err!("unnesting on non-column expr is not supported"); } }; - let already_projected = memo.get(col); - let (already_projected, transformed_cols) = match memo.get_mut(col) { - Some(vec) => (true, vec), - _ => { - memo.insert(col.clone(), vec![]); - (false, memo.get_mut(col).unwrap()) - } - }; + // Full context, we are trying to plan the execution as InnerProjection->Unnest->OuterProjection // inside unnest execution, each column inside the inner projection // will be transformed into new columns. Thus we need to keep track of these placeholding column names @@ -307,41 +301,78 @@ pub(crate) fn transform_bottom_unnest( let placeholder_name = format!("unnest_placeholder({})", col.name()); let post_unnest_name = format!("unnest_placeholder({},depth={})", col.name(), level); - - // Add alias for the argument expression, to avoid naming conflicts - // with other expressions in the select list. For example: `select unnest(col1), col1 from t`. - // this extra projection is used to unnest transforming - if !already_projected { - inner_projection_exprs - .push(expr_in_unnest.clone().alias(placeholder_name.clone())); - - unnest_placeholder_columns.push(placeholder_name.clone()); - } - let schema = input.schema(); let (data_type, _) = expr_in_unnest.data_type_and_nullable(schema)?; - if !struct_allowed { - if let DataType::Struct(_) = data_type { - return internal_err!("unnest on struct can only be applied at the root level of select expression"); - } - } - // TODO: refactor get_unnested_columns or use a different function - // depth = 1 has no meaning here, just to provide enough argument - let outer_projection_columns = - get_unnested_columns(&post_unnest_name, &data_type, 1)?; - let expr = outer_projection_columns - .iter() - .map(|col| { - if !transformed_cols.contains(&col.0) { - transformed_cols.push(col.0.clone()); + match data_type { + DataType::Struct(inner_fields) => { + if !struct_allowed { + return internal_err!("unnest on struct can only be applied at the root level of select expression"); + } + inner_projection_exprs + .push(expr_in_unnest.clone().alias(placeholder_name.clone())); + unnest_placeholder_columns.push(( + Column::from_name(placeholder_name.clone()), + ColumnUnnestType::Struct, + )); + return Ok( + get_struct_unnested_columns(&placeholder_name, &inner_fields) + .into_iter() + .map(|c| Expr::Column(c)) + .collect(), + ); + } + DataType::List(field) + | DataType::FixedSizeList(field, _) + | DataType::LargeList(field) => { + // TODO: this memo only needs to be a hashset + let (already_projected, transformed_cols) = match memo.get_mut(col) { + Some(vec) => (true, vec), + _ => { + memo.insert(col.clone(), vec![]); + (false, memo.get_mut(col).unwrap()) + } + }; + if !already_projected { + inner_projection_exprs + .push(expr_in_unnest.clone().alias(placeholder_name.clone())); } - Expr::Column(col.0.clone()) - }) - .collect::>(); - Ok(expr) + let post_unnest_column = Column::from_name(post_unnest_name); + match unnest_placeholder_columns + .iter_mut() + .find(|(inner_col, _)| inner_col == col) + { + None => { + unnest_placeholder_columns.push(( + col.clone(), + ColumnUnnestType::List(vec![ColumnUnnestList { + output_column: post_unnest_column.clone(), + depth: level, + }]), + )); + } + Some((col, unnesting)) => match unnesting { + ColumnUnnestType::List(list) => { + list.push(ColumnUnnestList { + output_column: post_unnest_column.clone(), + depth: level, + }); + } + _ => { + return internal_err!("expr_in_unnest is a list type, while previous unnesting on this column is not a list type"); + } + }, + } + return Ok(vec![Expr::Column(post_unnest_column)]); + } + _ => { + return internal_err!( + "unnest on non-list or struct type is not supported" + ); + } + } }; let latest_visited_unnest = RefCell::new(None); let exprs_under_unnest = RefCell::new(HashSet::new()); @@ -483,12 +514,7 @@ pub(crate) fn transform_bottom_unnest( // write test for recursive_transform_unnest #[cfg(test)] mod tests { - use std::{ - collections::{HashMap, HashSet}, - ops::Add, - str::FromStr, - sync::Arc, - }; + use std::{collections::HashMap, sync::Arc}; use arrow::datatypes::{DataType as ArrowDataType, Field, Schema}; use arrow_schema::Fields; diff --git a/datafusion/sqllogictest/test_files/unnest.slt b/datafusion/sqllogictest/test_files/unnest.slt index db657ba43f69..17f97fed6c0a 100644 --- a/datafusion/sqllogictest/test_files/unnest.slt +++ b/datafusion/sqllogictest/test_files/unnest.slt @@ -44,6 +44,13 @@ AS VALUES (struct([2], 'b'), [[[3,4],[5]],[[null,6],null,[7,8]]], [struct([2],[[3],[4]])]) ; +query TT +explain select unnest([1,2,3]); +---- +1 +2 +3 + ## Basic unnest expression in select list query I select unnest([1,2,3]); From 825e270fad622b05db9a390de425c3906b78981f Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Sun, 4 Aug 2024 13:20:10 +0200 Subject: [PATCH 17/56] chore: detect some bugs --- datafusion/expr/src/logical_plan/plan.rs | 19 +++++ datafusion/sql/src/utils.rs | 92 ++++++++++++++---------- 2 files changed, 74 insertions(+), 37 deletions(-) diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index 42dd7c0618df..d2753f6cfb76 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -2889,12 +2889,31 @@ pub enum ColumnUnnestType { Inferred, } +impl fmt::Display for ColumnUnnestType { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + ColumnUnnestType::List(lists) => { + let list_strs: Vec = lists.iter().map(|list| list.to_string()).collect(); + write!(f, "List([{}])", list_strs.join(", ")) + } + ColumnUnnestType::Struct => write!(f, "Struct"), + ColumnUnnestType::Inferred => write!(f, "Inferred"), + } + } +} + #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct ColumnUnnestList { pub output_column: Column, pub depth: usize, } +impl fmt::Display for ColumnUnnestList { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}|depth={}", self.output_column, self.depth) + } +} + /// Unnest a column that contains a nested list type. See /// [`UnnestOptions`] for more details. #[derive(Debug, Clone, PartialEq, Eq, Hash)] diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index f63347165d5c..fda32981fd07 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -288,8 +288,9 @@ pub(crate) fn transform_bottom_unnest( inner_projection_exprs: &mut Vec| -> Result> { let col = match expr_in_unnest { - Expr::Column(col) => col, + Expr::Column(col) => col.name(), _ => { + // TODO: this failed return internal_err!("unnesting on non-column expr is not supported"); } }; @@ -301,6 +302,7 @@ pub(crate) fn transform_bottom_unnest( let placeholder_name = format!("unnest_placeholder({})", col.name()); let post_unnest_name = format!("unnest_placeholder({},depth={})", col.name(), level); + let placeholder_column = Column::from_name(placeholder_name.clone()); let schema = input.schema(); let (data_type, _) = expr_in_unnest.data_type_and_nullable(schema)?; @@ -342,11 +344,11 @@ pub(crate) fn transform_bottom_unnest( let post_unnest_column = Column::from_name(post_unnest_name); match unnest_placeholder_columns .iter_mut() - .find(|(inner_col, _)| inner_col == col) + .find(|(inner_col, _)| inner_col == &placeholder_column) { None => { unnest_placeholder_columns.push(( - col.clone(), + placeholder_column.clone(), ColumnUnnestType::List(vec![ColumnUnnestList { output_column: post_unnest_column.clone(), depth: level, @@ -355,10 +357,13 @@ pub(crate) fn transform_bottom_unnest( } Some((col, unnesting)) => match unnesting { ColumnUnnestType::List(list) => { - list.push(ColumnUnnestList { + let unnesting = ColumnUnnestList { output_column: post_unnest_column.clone(), depth: level, - }); + }; + if !list.contains(&unnesting) { + list.push(unnesting); + } } _ => { return internal_err!("expr_in_unnest is a list type, while previous unnesting on this column is not a list type"); @@ -441,17 +446,12 @@ pub(crate) fn transform_bottom_unnest( if expr == *unnest_stack.last().unwrap() { let most_inner = unnest_stack.first().unwrap(); if let Expr::Unnest(Unnest { expr: ref arg }) = most_inner { - // this is (one of) the bottom most unnest expr - let (data_type, _) = arg.data_type_and_nullable(input.schema())?; - if &expr == original_expr { - return Ok(Transformed::no(expr)); - } - let depth = unnest_stack.len(); let struct_allowed = (&expr == original_expr) && depth == 1; let mut transformed_exprs = transform(depth, arg, struct_allowed, inner_projection_exprs)?; + // TODO: if transformed_exprs has > 1 expr, handle it properly return Ok(Transformed::new( transformed_exprs.swap_remove(0), true, @@ -514,16 +514,24 @@ pub(crate) fn transform_bottom_unnest( // write test for recursive_transform_unnest #[cfg(test)] mod tests { - use std::{collections::HashMap, sync::Arc}; + use std::{collections::HashMap, ops::Add, sync::Arc}; use arrow::datatypes::{DataType as ArrowDataType, Field, Schema}; use arrow_schema::Fields; use datafusion_common::{Column, DFSchema, Result}; - use datafusion_expr::{col, lit, unnest, EmptyRelation, LogicalPlan}; + use datafusion_expr::{ + col, lit, unnest, ColumnUnnestType, EmptyRelation, LogicalPlan, + }; use datafusion_functions::core::expr_ext::FieldAccessor; use datafusion_functions_aggregate::expr_fn::count; use crate::utils::{resolve_positions_to_exprs, transform_bottom_unnest}; + fn column_unnests_eq(l: Vec<&str>, r: &[(Column, ColumnUnnestType)]) { + let formatted: Vec = + r.iter().map(|i| format!("{}|{}", i.0, i.1)).collect(); + assert_eq!(l, formatted) + } + #[test] fn test_transform_bottom_unnest_recursive_memoization_struct() -> Result<()> { let three_d_dtype = ArrowDataType::List(Arc::new(Field::new( @@ -598,10 +606,11 @@ mod tests { // memoization only contains 1 transformation assert_eq!(memo.len(), 1); assert!(memo.get(&Column::from_name("3d_col")).is_some()); - assert_eq!( - unnest_placeholder_columns, - vec!["unnest_placeholder(3d_col)"] + column_unnests_eq( + vec!["unnest_placeholder(3d_col)"], + &unnest_placeholder_columns, ); + // still reference struct_col in original schema but with alias, // to avoid colliding with the projection on the column itself if any assert_eq!( @@ -630,9 +639,9 @@ mod tests { // and the previous transformation is reused assert_eq!(memo.len(), 1); assert!(memo.get(&Column::from_name("3d_col")).is_some()); - assert_eq!( - unnest_placeholder_columns, - vec!["unnest_placeholder(3d_col)"] + column_unnests_eq( + vec!["unnest_placeholder(3d_col)"], + &mut unnest_placeholder_columns, ); // still reference struct_col in original schema but with alias, // to avoid colliding with the projection on the column itself if any @@ -668,12 +677,12 @@ mod tests { assert_eq!(memo.len(), 2); assert!(memo.get(&Column::from_name("struct_arr_col")).is_some()); - assert_eq!( - unnest_placeholder_columns, + column_unnests_eq( vec![ "unnest_placeholder(3d_col)", - "unnest_placeholder(struct_arr_col)" - ] + "unnest_placeholder(struct_arr_col)", + ], + &mut unnest_placeholder_columns, ); // still reference struct_col in original schema but with alias, // to avoid colliding with the projection on the column itself if any @@ -740,10 +749,11 @@ mod tests { // memoization only contains 1 transformation assert_eq!(memo.len(), 1); assert!(memo.get(&Column::from_name("3d_col")).is_some()); - assert_eq!( - unnest_placeholder_columns, - vec!["unnest_placeholder(3d_col)"] + column_unnests_eq( + vec!["unnest_placeholder(3d_col)|List([unnest_placeholder(3d_col,depth=2)|depth=2])"], + &unnest_placeholder_columns, ); + // still reference struct_col in original schema but with alias, // to avoid colliding with the projection on the column itself if any assert_eq!( @@ -769,12 +779,12 @@ mod tests { vec![col("unnest_placeholder(3d_col,depth=1)").alias("2d_col")] ); // memoization still contains 1 transformation - // and the previous transformation is reused + // and the for the same column, depth = 1 needs to be performed aside from depth = 2 assert_eq!(memo.len(), 1); assert!(memo.get(&Column::from_name("3d_col")).is_some()); - assert_eq!( - unnest_placeholder_columns, - vec!["unnest_placeholder(3d_col)"] + column_unnests_eq( + vec!["unnest_placeholder(3d_col)|List([unnest_placeholder(3d_col,depth=2)|depth=2, unnest_placeholder(3d_col,depth=1)|depth=1])"], + &unnest_placeholder_columns, ); // still reference struct_col in original schema but with alias, // to avoid colliding with the projection on the column itself if any @@ -839,7 +849,10 @@ mod tests { col("unnest(struct_col).field2"), ] ); - assert_eq!(unnest_placeholder_columns, vec!["unnest(struct_col)"]); + column_unnests_eq( + vec!["unnest_placeholder(struct_col)"], + &mut unnest_placeholder_columns, + ); // still reference struct_col in original schema but with alias, // to avoid colliding with the projection on the column itself if any assert_eq!( @@ -857,9 +870,12 @@ mod tests { &mut memo, &original_expr, )?; - assert_eq!( - unnest_placeholder_columns, - vec!["unnest(struct_col)", "unnest(array_col)"] + column_unnests_eq( + vec![ + "unnest_placeholder(struct_col)", + "unnest_placeholder(array_col)", + ], + &mut unnest_placeholder_columns, ); // only transform the unnest children assert_eq!( @@ -925,10 +941,12 @@ mod tests { transformed_exprs, vec![unnest(col("unnest(struct_col[matrix])"))] ); - assert_eq!( - unnest_placeholder_columns, - vec!["unnest(struct_col[matrix])"] + + column_unnests_eq( + vec!["unnest_placeholder(struct_col[matrix])"], + &mut unnest_placeholder_columns, ); + assert_eq!( inner_projection_exprs, vec![col("struct_col") From 721c92fc94d07a15ed3cdbf7e21d1f8ee0381e53 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Mon, 12 Aug 2024 21:47:22 +0200 Subject: [PATCH 18/56] some work --- datafusion/expr/src/logical_plan/builder.rs | 15 +- datafusion/physical-plan/src/unnest.rs | 536 ++++++++++++++++++- datafusion/sql/src/select.rs | 16 +- datafusion/sql/src/utils.rs | 36 +- datafusion/sqllogictest/test_files/debug.slt | 95 ++-- 5 files changed, 615 insertions(+), 83 deletions(-) diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index 8fab8c7a7dc8..656d726173f7 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -1718,7 +1718,7 @@ pub fn unnest_with_options( ) -> Result { let mut list_columns: Vec<(usize, ColumnUnnestList)> = vec![]; let mut struct_columns = vec![]; - let column_by_original_index = columns + let indices_to_unnest = columns .iter() .map(|col_unnesting| { Ok(( @@ -1732,13 +1732,20 @@ pub fn unnest_with_options( let mut dependency_indices = vec![]; // Transform input schema into new schema - // e.g int, unnest([]int), unnest(struct(varchar,varchar)) - // becomes int, int, varchar, varchar + // Given this comprehensive example + // + // input schema: + // col1_unnest_placeholder: list[list[int]], col1: list[list[int]], col2 list[int] + // with unnest on unnest(col1,depth=2), unnest(col1,depth=1) and unnest(col2,depth=1) + // output schema: + // unnest_col1_depth_2: int, unnest_col1_depth1: list[int], col1: list[list[int]], unnest_col2_depth_1: int + // Meaning the placeholder column will be replaced by its unnested variation(s), note + // the plural. let fields = input_schema .iter() .enumerate() .map(|(index, (original_qualifier, original_field))| { - match column_by_original_index.get(&index) { + match indices_to_unnest.get(&index) { Some((column_to_unnest, unnest_type)) => { let mut inferred_unnest_type = unnest_type.clone(); if let ColumnUnnestType::Inferred = unnest_type { diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index 7b00156756c9..afd24b4f005a 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -17,6 +17,7 @@ //! Define a plan for unnesting values in columns that contain a list type. +use std::cmp; use std::collections::HashMap; use std::{any::Any, sync::Arc}; @@ -36,18 +37,20 @@ use arrow::compute::kernels::zip::zip; use arrow::compute::{cast, is_not_null, kernels, sum}; use arrow::datatypes::{DataType, Int32Type, Int64Type, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; -use arrow_array::{Int64Array, Scalar, StructArray}; +use arrow_array::{Int32Array, Int64Array, Scalar, StructArray}; use arrow_ord::cmp::lt; use datafusion_common::{ exec_datafusion_err, exec_err, internal_err, Result, UnnestOptions, }; use datafusion_execution::TaskContext; +use datafusion_expr::builder::unnest; use datafusion_expr::ColumnarValue; use datafusion_physical_expr::EquivalenceProperties; use async_trait::async_trait; use futures::{Stream, StreamExt}; use hashbrown::HashSet; +use itertools::{Either, Itertools}; use log::trace; /// Unnest the given columns (either with type struct or list) @@ -347,6 +350,220 @@ pub struct ListUnnest { pub depth: usize, } +// a map of original index (track which column in the batch to be interested in) +// input schema: +// col1_unnest_placeholder: list[list[int]], col1: list[list[int]], col2 list[int] +// with unnest on unnest(col1,depth=2), unnest(col1,depth=1) and unnest(col2,depth=1) +// output schema: +// unnest_col1_depth_2: int, unnest_col1_depth1: list[int], col1: list[list[int]], unnest_col2_depth_1: int +// Meaning the placeholder column will be replaced by its unnested variation(s), note +// the plural. + +/* + Note: unnest has a big difference in behavior between Postgres and DuckDB + Take this example + 1.Postgres + ```ignored + create table temp ( + i integer[][][], j integer[] + ) + insert into temp values ('{{{1,2},{3,4}},{{5,6},{7,8}}}', '{1,2}'); + select unnest(i), unnest(j) from temp; + ``` + + Result + 1 1 + 2 2 + 3 + 4 + 5 + 6 + 7 + 8 + 2. DuckDB + ```ignore + create table temp (i integer[][][], j integer[]); + insert into temp values ([[[1,2],[3,4]],[[5,6],[7,8]]], [1,2]); + select unnest(i,recursive:=true), unnest(j,recursive:=true) from temp; + ``` + Result: + ┌────────────────────────────────────────────────┬────────────────────────────────────────────────┐ + │ unnest(i, "recursive" := CAST('t' AS BOOLEAN)) │ unnest(j, "recursive" := CAST('t' AS BOOLEAN)) │ + │ int32 │ int32 │ + ├────────────────────────────────────────────────┼────────────────────────────────────────────────┤ + │ 1 │ 1 │ + │ 2 │ 2 │ + │ 3 │ 1 │ + │ 4 │ 2 │ + │ 5 │ 1 │ + │ 6 │ 2 │ + │ 7 │ 1 │ + │ 8 │ 2 │ + └────────────────────────────────────────────────┴────────────────────────────────────────────────┘ + The following implementation refer to Postgres's implementation + For DuckDB's result to be similar, the above query can be written as + + + +*/ + +fn unnest_at_level( + batch: &[ArrayRef], + schema: &SchemaRef, + list_type_columns: &[ListUnnest], + temp_batch: &mut HashMap<(usize, usize), ArrayRef>, + level_to_unnest: usize, + options: &UnnestOptions, +) -> Result> { + // unnested columns at this depth level + // now do some kind of projection-like + // This query: + // select unnest(col1, max_depth:=3), unnest(col1,max_depth:=2), unnest(col1, max_depth:=1) from temp; + // is equivalent to + // + // unnest(depth3) , unnest(depth2), unnest(depth1) + // select (unnest) + // batch comes in as [a,b] + // in this example list_type_columns as value + // [(a,1), (a,2)] and [(b,1)] + // 1.microwork(level=2) + // new batch as [unnest_a,b] + // new list as [(a,1)] and [(b,1)] + // temp_batch_for_projection: [(a,depth=1,unnested(a))] + // 2.microwork(level=1) + // new batch as [unnest(unnest_a),b] + // new list as [] + // temp_batch: [(a,depth=1,unnested(a)),(a,depth=2,unnested(unnested(a))),(b,depth=1,unnested(b))] + // this is final, now combine the mainbatch and temp_batch + + let temp_unnest_cols = list_type_columns + .iter() + .filter_map( + |ListUnnest { + depth, + index_in_input_schema, + }| { + if *depth == level_to_unnest { + return Some(Arc::clone(&batch[*index_in_input_schema])); + } + if *depth > level_to_unnest { + return Some( + temp_batch + .get(&(*index_in_input_schema, *depth)) + .unwrap() + .clone(), + ); + } + return None; + }, + ) + .collect::>(); + + // filter out so that list_arrays only contain column with the highest depth + // at the same time, during iteration remove this depth so next time we don't have to unnest them again + let longest_length = find_longest_length(&temp_unnest_cols, options)?; + let unnested_length = longest_length.as_primitive::(); + let total_length = if unnested_length.is_empty() { + 0 + } else { + sum(unnested_length).ok_or_else(|| { + exec_datafusion_err!("Failed to calculate the total unnested length") + })? as usize + }; + if total_length == 0 { + return Ok(vec![]); + } + + // Unnest all the list arrays + let unnested_temp_arrays = + unnest_list_arrays(&temp_unnest_cols, unnested_length, total_length)?; + + unnested_temp_arrays + .into_iter() + .zip(list_type_columns.iter()) + .for_each( + |( + flatten_arr, + ListUnnest { + index_in_input_schema, + depth, + }, + )| { + temp_batch.insert((*index_in_input_schema, *depth), flatten_arr); + }, + ); + + // Create the take indices array for other columns + let take_indices = create_take_indicies(unnested_length, total_length); + + // vertical expansion because of list unnest + let ret = flatten_batch_from_indices(batch, &take_indices)?; + return Ok(ret); +} + +fn build_batch_v2( + batch: &RecordBatch, + schema: &SchemaRef, + list_type_columns: &[ListUnnest], + struct_column_indices: &HashSet, + options: &UnnestOptions, +) -> Result { + let transformed = match list_type_columns.len() { + 0 => flatten_struct_cols(batch.columns(), schema, struct_column_indices), + _ => { + let mut temp_batch = HashMap::new(); + let highest_depth = list_type_columns + .iter() + .fold(0, |highest_depth, ListUnnest { depth, .. }| { + cmp::max(highest_depth, *depth) + }); + let mut temp_unnest_result = batch.columns(); + for depth in (1..=highest_depth).rev() { + temp_unnest_result = &unnest_at_level( + batch.columns(), + schema, + list_type_columns, + &mut temp_batch, + highest_depth, + options, + )?; + } + // TODO: combine temp with the batch + let unnested_array_map: HashMap>> = temp_batch + .into_iter() + .enumerate() + // .into_iter() + // .zip(list_type_columns.iter()) + .fold( + HashMap::new(), + |mut acc, + ( + flattened_array, + ListUnnest { + index_in_input_schema, + depth, + }, + )| { + acc.entry(*index_in_input_schema) + .or_insert(vec![]) + .push(flattened_array); + acc + }, + ); + + // temp_unnest_result.iter().map(f) + // vertical expansion because of list unnest + let ret = flatten_list_cols_from_indices( + batch, + unnested_array_map, + &take_indicies, + )?; + flatten_struct_cols(&ret, schema, struct_column_indices) + } + }; + transformed +} + /// For each row in a `RecordBatch`, some list/struct columns need to be unnested. /// - For list columns: We will expand the values in each list into multiple rows, /// taking the longest length among these lists, and shorter lists are padded with NULLs. @@ -362,6 +579,16 @@ fn build_batch( let transformed = match list_type_columns.len() { 0 => flatten_struct_cols(batch.columns(), schema, struct_column_indices), _ => { + // maintain a map of temp result + // > + // given a map of unnesting + // col1: + // col2: + // 1.run for each to unnest all the top level + // col1: => to temp + // 2.run unnest on remaining level + // col1: => to temp + // col2: => to temp let list_arrays: Vec<(ArrayRef, usize)> = list_type_columns .iter() .map( @@ -433,21 +660,53 @@ fn build_batch( /// TODO: only prototype /// This function does not handle NULL yet -pub fn length_recursive(array: &dyn Array, depth: usize) -> Result { +pub fn length_recursive( + array: &dyn Array, + depth: usize, + preserve_nulls: bool, + is_top_level: bool, // null values at top level recursion is ignored, regardless of preserve_null +) -> Result { let list = array.as_list::(); + let null_length = if preserve_nulls { + Scalar::new(Int32Array::from_value(1, 1)) + } else { + Scalar::new(Int32Array::from_value(0, 1)) + }; + // preserve null only concern nullability of the element in array, not the value of the array itself + // e.g unnest(null) => 0 rows + // while unnest([null]) => 1 row with null value + let null_length_for_value = match preserve_nulls && !is_top_level { + true => 1, + false => 0, + }; if depth == 1 { - return Ok(length(array)?); + if array.is_empty() { + return Ok(Arc::new(PrimitiveArray::::from(vec![ + null_length_for_value, + ]))); + } + + // null elements will have null length + // respect preserve_nulls here + let mut ret = length(array)?; + ret = cast(&ret, &DataType::Int32)?; + return Ok(zip(&is_not_null(&ret)?, &ret, &null_length)?); } let a: Vec = list .iter() .map(|x| { if x.is_none() { - return 0; + return null_length_for_value; } - let ret = length_recursive(&x.unwrap(), depth - 1).unwrap(); + // [[1,2,3],null,[4,5]] + // [[7,8,9,10],] + // length([1,2]) + length(null) + length([4,5]) + let ret = + length_recursive(&x.unwrap(), depth - 1, preserve_nulls, false).unwrap(); let a = ret.as_primitive::(); if a.is_empty() { - 0 + // TODO respect is_preserve_nulls + 1 } else { sum(a).unwrap() } @@ -479,7 +738,7 @@ pub fn length_recursive(array: &dyn Array, depth: usize) -> Result { /// ``` /// fn find_longest_length( - list_arrays: &[(ArrayRef, usize)], + list_arrays: &[ArrayRef], options: &UnnestOptions, ) -> Result { // The length of a NULL list @@ -496,8 +755,11 @@ fn find_longest_length( // unnest(col1), unnest(unnest(col1)), unnest(col2) let list_lengths: Vec = list_arrays .iter() - .map(|(list_array, depth)| { - let mut length_array = length_recursive(list_array, *depth)?; + .map(|list_array| { + // Perhaps we don't need recursive here + // let mut length_array = + // length_recursive(list_array, *depth, options.preserve_nulls, true)?; + let mut length_array = length(list_array)?; // Make sure length arrays have the same type. Int64 is the most general one. // Respect the depth of unnest( current func only get the length of 1 level of unnest) length_array = cast(&length_array, &DataType::Int64)?; @@ -561,21 +823,19 @@ impl ListArrayType for FixedSizeListArray { /// Unnest multiple list arrays according to the length array. fn unnest_list_arrays( - list_arrays: &[(ArrayRef, usize)], + list_arrays: &[ArrayRef], length_array: &PrimitiveArray, capacity: usize, ) -> Result> { let typed_arrays = list_arrays .iter() - .map(|(list_array, depth)| match list_array.data_type() { - DataType::List(_) => { - Ok((list_array.as_list::() as &dyn ListArrayType, depth)) - } + .map(|list_array| match list_array.data_type() { + DataType::List(_) => Ok(list_array.as_list::() as &dyn ListArrayType), DataType::LargeList(_) => { - Ok((list_array.as_list::() as &dyn ListArrayType, depth)) + Ok(list_array.as_list::() as &dyn ListArrayType) } DataType::FixedSizeList(_, _) => { - Ok((list_array.as_fixed_size_list() as &dyn ListArrayType, depth)) + Ok(list_array.as_fixed_size_list() as &dyn ListArrayType) } other => exec_err!("Invalid unnest datatype {other }"), }) @@ -583,9 +843,7 @@ fn unnest_list_arrays( typed_arrays .iter() - .map(|(list_array, depth)| { - unnest_list_array(*list_array, length_array, capacity, **depth) - }) + .map(|list_array| unnest_list_array(*list_array, length_array, capacity)) .collect::>() } @@ -614,9 +872,9 @@ fn unnest_list_array( list_array: &dyn ListArrayType, length_array: &PrimitiveArray, capacity: usize, - depth: usize, ) -> Result { let values = list_array.values(); + // TODO: handle me recursively let mut take_indicies_builder = PrimitiveArray::::builder(capacity); for row in 0..list_array.len() { let mut value_length = 0; @@ -677,6 +935,35 @@ fn create_take_indicies( builder.finish() } +fn flatten_batch_from_indices( + batch: &[ArrayRef], + // temp: &[(usize, Arc)], + // mut unnested_list_arrays: HashMap>, + indices: &PrimitiveArray, +) -> Result>> { + return batch + .into_iter() + .map(|arr| Ok(kernels::take::take(arr, indices, None)?)) + .collect::>(); + // return Ok(new_arrays); + // Ok(new_arrays) + // let arrays = batch + // .columns() + // .iter() + // .enumerate() + // .map( + // |(col_idx, arr)| match unnested_list_arrays.remove(&col_idx) { + // Some(unnested_arrays) => Ok(unnested_arrays), + // None => Ok(vec![kernels::take::take(arr, indices, None)?]), + // }, + // ) + // .collect::>>()? + // .into_iter() + // .flatten() + // .collect::>(); + // Ok(arrays) +} + /// Create the final batch given the unnested column arrays and a `indices` array /// that is used by the take kernel to copy values. /// @@ -734,7 +1021,7 @@ fn flatten_list_cols_from_indices( mod tests { use super::*; use arrow::datatypes::{Field, Int32Type}; - use arrow_array::{GenericListArray, OffsetSizeTrait, StringArray}; + use arrow_array::{GenericListArray, Int32Array, OffsetSizeTrait, StringArray}; use arrow_buffer::{BooleanBufferBuilder, NullBuffer, OffsetBuffer}; // Create a GenericListArray with the following list values: @@ -818,10 +1105,179 @@ mod tests { assert_eq!(strs, expected); Ok(()) } + #[test] + fn test_todo() -> datafusion_common::Result<()> { + // original: [[1,2,3],null,[4,5]] [[7,8,9,10],null,[11,12,13]] null + // unnest depth=1:[1,2,3], null, [4,5], [7,8,9,10], null, [11,12,13] , null + // unnest depth=2: 1,2,3,null,4,5,7,8,9,10,null,11,12,13,null + let list_arr = ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(2), Some(3)]), + None, + Some(vec![Some(4), Some(5)]), + Some(vec![Some(7), Some(8), Some(9), Some(10)]), + None, + Some(vec![Some(11), Some(12), Some(13)]), + ]); + + let list_arr_ref = Arc::new(list_arr.clone()) as ArrayRef; + // last item is null + let offsets = OffsetBuffer::from_lengths([3, 3, 0]); + + let mut nulls = BooleanBufferBuilder::new(2); + nulls.append(true); + nulls.append(true); + nulls.append(false); + let nested_list_arr = ListArray::new( + Arc::new(Field::new_list_field( + list_arr_ref.data_type().to_owned(), + true, + )), + offsets, + list_arr_ref, + Some(NullBuffer::new(nulls.finish())), + ); + // TODO: for each row we have to check if it is not null, then iterate recursively to find the offset + // until the depth = 1 + // Seed some None values at the top level, second top level to check how to null works + for row in 0..list_arr.len() { + // TODO: verify null preseving behavior + if !list_arr.is_null(row) { + let (start, end) = (&list_arr as &dyn ListArrayType).value_offsets(row); + for i in start..end {} + } + } + nested_list_arr + .values() + .as_list::() + .offsets() + .iter() + .for_each(|i| { + println!("{}", i); + }); + let nested_list_arr_ref = Arc::new(nested_list_arr.clone()) as ArrayRef; + let longest_length = find_longest_length( + &vec![(nested_list_arr_ref, 2)], + &UnnestOptions { + preserve_nulls: true, + }, + )?; + let longest_length = longest_length.as_primitive::(); + let unnested = + unnest_list_array_prototype(&nested_list_arr, longest_length, 14, 2)?; + let expected_array = Int32Array::from(Vec::::from([1, 2, 3])); + assert_eq!( + unnested.as_any().downcast_ref::().unwrap(), + &expected_array + ); + + Ok(()) + // let unnested_length = longest_length.as_primitive::(); + // unnest_list_arrays(list_arrays, length_array, capacity) + // unnest_list_array(list_array, length_array, capacity, depth) + } + + fn unnest_list_array_prototype_recursive( + index_builder: &mut PrimitiveArray, + list_array: &GenericListArray, + // length_array: &PrimitiveArray, + // capacity: usize, + depth: usize, + ) -> Result { + Ok(0) + // if depth == 1 { + // for row in 0..list_array.len() { + // let mut value_length = 0; + // if !list_array.is_null(row) { + // let (start, end) = + // (list_array as &dyn ListArrayType).value_offsets(row); + // value_length = end - start; + // for i in start..end { + // if depth == 1 { + // take_indicies_builder.append_value(i) + // } else { + // } + // } + // } + // } + // } else { + // for row in 0..list_array.len() { + // let mut value_length = 0; + // if !list_array.is_null(row) { + // let (start, end) = + // (list_array as &dyn ListArrayType).value_offsets(row); + // value_length = end - start; + // for i in start..end { + // if depth == 1 { + // take_indicies_builder.append_value(i) + // } else { + // } + // } + // } + // } + // let inner_list_array = list_array.values().as_list::(); + // unnest_list_array_prototype_recursive( + // index_builder, + // inner_list_array, + // depth - 1, + // ) + // } + } + // [[1,2,3],null,[4,5]] 1 + // [[7,8,9,10],null,[11,12,13]] 1 + // null 1 + // result into + // 1,2,3,null,4,5,7,8,9,10,null,11,12,13 + // + fn unnest_list_array_prototype( + list_array: &GenericListArray, + length_array: &PrimitiveArray, + capacity: usize, + depth: usize, + ) -> Result { + let mut cur_depth = depth; + let mut backed_array = list_array.values(); + loop { + if cur_depth == 1 { + break; + } + backed_array = backed_array.as_list::().values(); + cur_depth -= 1; + } + + // TODO: handle me recursively + let mut take_indicies_builder = PrimitiveArray::::builder(capacity); + for row in 0..list_array.len() { + let mut value_length = 0; + if !list_array.is_null(row) { + let (start, end) = (list_array as &dyn ListArrayType).value_offsets(row); + value_length = end - start; + for i in start..end { + if depth == 1 { + take_indicies_builder.append_value(i) + } else { + } + } + } + let target_length = length_array.value(row); + debug_assert!( + value_length <= target_length, + "value length is beyond the longest length" + ); + // Pad with NULL values + for _ in value_length..target_length { + take_indicies_builder.append_null(); + } + } + Ok(kernels::take::take( + &backed_array, + &take_indicies_builder.finish(), + None, + )?) + } #[test] fn test_length_recursive() -> datafusion_common::Result<()> { - // [[1,2,3],null,[4,5]] + // [[1,2,3],null,[4,5]],[[7,8,9,10],null,[11,12,13]] null let list_arr = ListArray::from_iter_primitive::(vec![ Some(vec![Some(1), Some(2), Some(3)]), None, @@ -831,7 +1287,11 @@ mod tests { Some(vec![Some(11), Some(12), Some(13)]), ]); let list_arr_ref = Arc::new(list_arr) as ArrayRef; - let offsets = OffsetBuffer::from_lengths([3, 3]); + let offsets = OffsetBuffer::from_lengths([3, 3, 0]); + let mut nulls = BooleanBufferBuilder::new(3); + nulls.append(true); + nulls.append(true); + nulls.append(false); let nested_list_arr = ListArray::new( Arc::new(Field::new_list_field( list_arr_ref.data_type().to_owned(), @@ -839,15 +1299,37 @@ mod tests { )), offsets, list_arr_ref, + Some(NullBuffer::new(nulls.finish())), + ); + let list_arr_2 = ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(2), Some(3)]), + None, + Some(vec![Some(4), Some(5)]), + Some(vec![Some(7), Some(8), Some(9), Some(10)]), None, + Some(vec![Some(11), Some(12), Some(13)]), + ]); + let list_arr_ref = Arc::new(list_arr) as ArrayRef; + let offsets = OffsetBuffer::from_lengths([3, 3, 0]); + let mut nulls = BooleanBufferBuilder::new(3); + nulls.append(true); + nulls.append(true); + nulls.append(false); + let nested_list_arr = ListArray::new( + Arc::new(Field::new_list_field( + list_arr_ref.data_type().to_owned(), + true, + )), + offsets, + list_arr_ref, + Some(NullBuffer::new(nulls.finish())), ); - // [3, 0, 2] - // if depth = 2, length should be 5 + // [6,8,0] verify_longest_length( &[(Arc::new(nested_list_arr) as ArrayRef, 2)], true, - vec![5, 7], - ); + vec![6, 8, 0], + )?; Ok(()) } diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs index 5ddebd573cac..7a54bf563128 100644 --- a/datafusion/sql/src/select.rs +++ b/datafusion/sql/src/select.rs @@ -351,9 +351,23 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { let columns = unnest_columns.into_iter().map(|col| col.into()).collect(); // Set preserve_nulls to false to ensure compatibility with DuckDB and PostgreSQL let unnest_options = UnnestOptions::new().with_preserve_nulls(false); + let mut check_list: HashSet = inner_projection_exprs + .iter() + .map(|expr| expr.clone()) + .collect(); + let deduplicated: Vec = inner_projection_exprs + .into_iter() + .filter(|expr| -> bool { + if check_list.remove(expr) { + true + } else { + false + } + }) + .collect(); let plan = LogicalPlanBuilder::from(intermediate_plan) - .project(inner_projection_exprs.clone())? + .project(deduplicated.clone())? .unnest_columns_with_options(columns, unnest_options)? .build()?; intermediate_plan = plan; diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index fda32981fd07..657524e14064 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -279,7 +279,7 @@ pub(crate) fn transform_bottom_unnest( input: &LogicalPlan, unnest_placeholder_columns: &mut Vec<(Column, ColumnUnnestType)>, inner_projection_exprs: &mut Vec, - memo: &mut HashMap>, + memo: &mut HashMap>, original_expr: &Expr, ) -> Result> { let mut transform = |level: usize, @@ -287,21 +287,22 @@ pub(crate) fn transform_bottom_unnest( struct_allowed: bool, inner_projection_exprs: &mut Vec| -> Result> { - let col = match expr_in_unnest { - Expr::Column(col) => col.name(), - _ => { - // TODO: this failed - return internal_err!("unnesting on non-column expr is not supported"); - } - }; + let inner_expr_name = expr_in_unnest.display_name()?; + // let col = match expr_in_unnest { + // Expr::Column(col) => col, + // _ => { + // // TODO: this failed + // return internal_err!("unnesting on non-column expr is not supported"); + // } + // }; // Full context, we are trying to plan the execution as InnerProjection->Unnest->OuterProjection // inside unnest execution, each column inside the inner projection // will be transformed into new columns. Thus we need to keep track of these placeholding column names // let placeholder_name = unnest_expr.display_name()?; - let placeholder_name = format!("unnest_placeholder({})", col.name()); + let placeholder_name = format!("unnest_placeholder({})", inner_expr_name); let post_unnest_name = - format!("unnest_placeholder({},depth={})", col.name(), level); + format!("unnest_placeholder({},depth={})", inner_expr_name, level); let placeholder_column = Column::from_name(placeholder_name.clone()); let schema = input.schema(); @@ -329,13 +330,14 @@ pub(crate) fn transform_bottom_unnest( | DataType::FixedSizeList(field, _) | DataType::LargeList(field) => { // TODO: this memo only needs to be a hashset - let (already_projected, transformed_cols) = match memo.get_mut(col) { - Some(vec) => (true, vec), - _ => { - memo.insert(col.clone(), vec![]); - (false, memo.get_mut(col).unwrap()) - } - }; + let (already_projected, transformed_cols) = + match memo.get_mut(&inner_expr_name) { + Some(vec) => (true, vec), + _ => { + memo.insert(inner_expr_name.clone(), vec![]); + (false, memo.get_mut(&inner_expr_name).unwrap()) + } + }; if !already_projected { inner_projection_exprs .push(expr_in_unnest.clone().alias(placeholder_name.clone())); diff --git a/datafusion/sqllogictest/test_files/debug.slt b/datafusion/sqllogictest/test_files/debug.slt index 76fb93a7db49..51c267f04186 100644 --- a/datafusion/sqllogictest/test_files/debug.slt +++ b/datafusion/sqllogictest/test_files/debug.slt @@ -1,45 +1,72 @@ +statement ok +CREATE TABLE recursive_unnest_table +AS VALUES + (struct([1], 'a'), [[[1],[2]],[[1,1]]], [struct([1],[[1,2]])]), + (struct([2], 'b'), [[[3,4],[5]],[[null,6],null,[7,8]]], [struct([2],[[3],[4]])]) +; + query TT -explain select unnest(unnest([[1,2,3]])) + unnest([4,5]); +select unnest(unnest(unnest(column2))) from recursive_unnest_table; ---- logical_plan -01)Projection: unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))) + unnest(make_array(Int64(4),Int64(5))) -02)--Unnest: lists[unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))))] structs[] -03)----Projection: unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))), unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))) AS unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))), unnest(make_array(Int64(4),Int64(5))) -04)------Unnest: lists[unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))), unnest(make_array(Int64(4),Int64(5)))] structs[] -05)--------Projection: List([[1, 2, 3]]) AS unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))), List([4, 5]) AS unnest(make_array(Int64(4),Int64(5))) -06)----------EmptyRelation +01)Unnest: lists[unnest_placeholder(recursive_unnest_table.column2)|depth=2] structs[] +02)--Projection: recursive_unnest_table.column2 AS unnest_placeholder(recursive_unnest_table.column2) +03)----TableScan: recursive_unnest_table projection=[column2] physical_plan -01)ProjectionExec: expr=[unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))))@1 + unnest(make_array(Int64(4),Int64(5)))@2 as unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))) + unnest(make_array(Int64(4),Int64(5)))] -02)--UnnestExec -03)----RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -04)------ProjectionExec: expr=[unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))@0 as unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))), unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))@0 as unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))), unnest(make_array(Int64(4),Int64(5)))@1 as unnest(make_array(Int64(4),Int64(5)))] -05)--------UnnestExec -06)----------ProjectionExec: expr=[[[1, 2, 3]] as unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))), [4, 5] as unnest(make_array(Int64(4),Int64(5)))] -07)------------PlaceholderRowExec +01)UnnestExec +02)--ProjectionExec: expr=[column2@0 as unnest_placeholder(recursive_unnest_table.column2)] +03)----MemoryExec: partitions=1, partition_sizes=[1] -query I -select unnest([4,5]) + 1; +query TT +select unnest(unnest(column2)) from recursive_unnest_table; ---- -5 -6 -## FIXME: -query II -select unnest(unnest([[1,2,3]])) + unnest([4,5]), arrow_cast(unnest([4,5]),'Int64'); +query TT +select unnest(unnest(unnest(column3)['c1'])), unnest(unnest(column3)['c1']), column3,column1, column1['c0'] from recursive_unnest_table; ---- -5 4 -6 4 -7 4 - +## query TT +## explain select unnest(unnest([[1,2,3]])) + unnest([4,5]); +## ---- +## logical_plan +## 01)Projection: unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))) + unnest(make_array(Int64(4),Int64(5))) +## 02)--Unnest: lists[unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))))] structs[] +## 03)----Projection: unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))), unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))) AS unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))), unnest(make_array(Int64(4),Int64(5))) +## 04)------Unnest: lists[unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))), unnest(make_array(Int64(4),Int64(5)))] structs[] +## 05)--------Projection: List([[1, 2, 3]]) AS unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))), List([4, 5]) AS unnest(make_array(Int64(4),Int64(5))) +## 06)----------EmptyRelation +## physical_plan +## 01)ProjectionExec: expr=[unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))))@1 + unnest(make_array(Int64(4),Int64(5)))@2 as unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))) + unnest(make_array(Int64(4),Int64(5)))] +## 02)--UnnestExec +## 03)----RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +## 04)------ProjectionExec: expr=[unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))@0 as unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))), unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))@0 as unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))), unnest(make_array(Int64(4),Int64(5)))@1 as unnest(make_array(Int64(4),Int64(5)))] +## 05)--------UnnestExec +## 06)----------ProjectionExec: expr=[[[1, 2, 3]] as unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))), [4, 5] as unnest(make_array(Int64(4),Int64(5)))] +## 07)------------PlaceholderRowExec -query I -select unnest([2,1,9]) + unnest(unnest([[1,1,3]])) ; ----- -3 - -query I -select unnest(unnest([[1,1]])) + unnest([2,1,9]),unnest([2,1,9]) + unnest(unnest([[1,1]])) ; ----- -3 +## query I +## select unnest([4,5]) + 1; +## ---- +## 5 +## 6 +## +## ## FIXME: +## query II +## select unnest(unnest([[1,2,3]])) + unnest([4,5]), arrow_cast(unnest([4,5]),'Int64'); +## ---- +## 5 4 +## 6 4 +## 7 4 +## +## +## +## query I +## select unnest([2,1,9]) + unnest(unnest([[1,1,3]])) ; +## ---- +## 3 +## +## query I +## select unnest(unnest([[1,1]])) + unnest([2,1,9]),unnest([2,1,9]) + unnest(unnest([[1,1]])) ; +## ---- +## 3 From d11ed20d0f109cd3b0a953a25e37f0ab342c8ef7 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Tue, 13 Aug 2024 22:21:59 +0200 Subject: [PATCH 19/56] support recursive unnest in physical layer --- datafusion/expr/src/logical_plan/builder.rs | 9 +- datafusion/physical-plan/src/unnest.rs | 546 ++++++-------------- 2 files changed, 171 insertions(+), 384 deletions(-) diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index 656d726173f7..852009df3b63 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -1735,10 +1735,15 @@ pub fn unnest_with_options( // Given this comprehensive example // // input schema: - // col1_unnest_placeholder: list[list[int]], col1: list[list[int]], col2 list[int] + // 1.col1_unnest_placeholder: list[list[int]], + // 2.col1: list[list[int]] + // 3.col2: list[int] // with unnest on unnest(col1,depth=2), unnest(col1,depth=1) and unnest(col2,depth=1) // output schema: - // unnest_col1_depth_2: int, unnest_col1_depth1: list[int], col1: list[list[int]], unnest_col2_depth_1: int + // 1.unnest_col1_depth_2: int + // 2.unnest_col1_depth_1: list[int] + // 3.col1: list[list[int]] + // 4.unnest_col2_depth_1: int // Meaning the placeholder column will be replaced by its unnested variation(s), note // the plural. let fields = input_schema diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index afd24b4f005a..444d2c689af9 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -17,7 +17,7 @@ //! Define a plan for unnesting values in columns that contain a list type. -use std::cmp; +use std::cmp::{self, Ordering}; use std::collections::HashMap; use std::{any::Any, sync::Arc}; @@ -43,14 +43,12 @@ use datafusion_common::{ exec_datafusion_err, exec_err, internal_err, Result, UnnestOptions, }; use datafusion_execution::TaskContext; -use datafusion_expr::builder::unnest; use datafusion_expr::ColumnarValue; use datafusion_physical_expr::EquivalenceProperties; use async_trait::async_trait; use futures::{Stream, StreamExt}; use hashbrown::HashSet; -use itertools::{Either, Itertools}; use log::trace; /// Unnest the given columns (either with type struct or list) @@ -270,7 +268,7 @@ impl UnnestStream { .map(|maybe_batch| match maybe_batch { Some(Ok(batch)) => { let timer = self.metrics.elapsed_compute.timer(); - let result = build_batch( + let result = build_batch_v2( &batch, &self.schema, &self.list_type_columns, @@ -409,7 +407,6 @@ pub struct ListUnnest { fn unnest_at_level( batch: &[ArrayRef], - schema: &SchemaRef, list_type_columns: &[ListUnnest], temp_batch: &mut HashMap<(usize, usize), ArrayRef>, level_to_unnest: usize, @@ -517,47 +514,83 @@ fn build_batch_v2( .fold(0, |highest_depth, ListUnnest { depth, .. }| { cmp::max(highest_depth, *depth) }); - let mut temp_unnest_result = batch.columns(); + let mut unnested_original_columns = vec![]; for depth in (1..=highest_depth).rev() { - temp_unnest_result = &unnest_at_level( - batch.columns(), - schema, + let input = match depth { + highest_depth => batch.columns(), + _ => &unnested_original_columns, + }; + let temp = unnest_at_level( + input, list_type_columns, &mut temp_batch, highest_depth, options, )?; + unnested_original_columns = temp; } // TODO: combine temp with the batch - let unnested_array_map: HashMap>> = temp_batch - .into_iter() - .enumerate() - // .into_iter() - // .zip(list_type_columns.iter()) - .fold( + let unnested_array_map: HashMap, usize)>> = + temp_batch.into_iter().fold( HashMap::new(), - |mut acc, - ( - flattened_array, - ListUnnest { - index_in_input_schema, - depth, - }, - )| { - acc.entry(*index_in_input_schema) + |mut acc, ((index_in_input_schema, depth), flattened_array)| { + acc.entry(index_in_input_schema) .or_insert(vec![]) - .push(flattened_array); + .push((flattened_array, depth)); acc }, ); + let output_order: HashMap<(usize, usize), usize> = list_type_columns + .iter() + .enumerate() + .map(|(order, unnest_def)| { + let ListUnnest { + depth, + index_in_input_schema, + } = unnest_def; + ((*index_in_input_schema, *depth), order) + }) + .collect(); + + let mut ordered_unnested_array_map = unnested_array_map + .into_iter() + .map( + // each item in unnested_columns is the result of unnesting the same input column + // we need to sort them to conform with the unnest definition + // e.g unnest(unnest(col)) must goes before unnest(col) + |(original_index, mut unnested_columns)| { + unnested_columns.sort_by(|(_, depth), (_, depth2)| -> Ordering { + output_order.get(&(original_index, *depth)).unwrap().cmp( + output_order.get(&(original_index, *depth2)).unwrap(), + ) + }); + ( + original_index, + unnested_columns + .into_iter() + .map(|(arr, _)| arr) + .collect::>(), + ) + }, + ) + .collect::>(); + // let highest_depth = list_type_columns + // .iter() + // .fold(0, |highest_depth, ListUnnest { depth, .. }| { + // cmp::max(highest_depth, *depth) + // }); + let ret = unnested_original_columns + .iter() + .enumerate() + .map( + |(col_idx, arr)| match ordered_unnested_array_map.remove(&col_idx) { + Some(unnested_arrays) => unnested_arrays, + None => vec![arr.clone()], + }, + ) + .flatten() + .collect::>(); - // temp_unnest_result.iter().map(f) - // vertical expansion because of list unnest - let ret = flatten_list_cols_from_indices( - batch, - unnested_array_map, - &take_indicies, - )?; flatten_struct_cols(&ret, schema, struct_column_indices) } }; @@ -569,94 +602,94 @@ fn build_batch_v2( /// taking the longest length among these lists, and shorter lists are padded with NULLs. /// - For struct columns: We will expand the struct columns into multiple subfield columns. /// For columns that don't need to be unnested, repeat their values until reaching the longest length. -fn build_batch( - batch: &RecordBatch, - schema: &SchemaRef, - list_type_columns: &[ListUnnest], - struct_column_indices: &HashSet, - options: &UnnestOptions, -) -> Result { - let transformed = match list_type_columns.len() { - 0 => flatten_struct_cols(batch.columns(), schema, struct_column_indices), - _ => { - // maintain a map of temp result - // > - // given a map of unnesting - // col1: - // col2: - // 1.run for each to unnest all the top level - // col1: => to temp - // 2.run unnest on remaining level - // col1: => to temp - // col2: => to temp - let list_arrays: Vec<(ArrayRef, usize)> = list_type_columns - .iter() - .map( - |ListUnnest { - depth, - index_in_input_schema, - }| { - Ok(( - ColumnarValue::Array(Arc::clone( - batch.column(*index_in_input_schema), - )) - .into_array(batch.num_rows())?, - *depth, - )) - }, - ) - .collect::>()?; - - let longest_length = find_longest_length(&list_arrays, options)?; - let unnested_length = longest_length.as_primitive::(); - let total_length = if unnested_length.is_empty() { - 0 - } else { - sum(unnested_length).ok_or_else(|| { - exec_datafusion_err!("Failed to calculate the total unnested length") - })? as usize - }; - if total_length == 0 { - return Ok(RecordBatch::new_empty(Arc::clone(schema))); - } - - // Unnest all the list arrays - let unnested_arrays = - unnest_list_arrays(&list_arrays, unnested_length, total_length)?; - let unnested_array_map: HashMap>> = unnested_arrays - .into_iter() - .zip(list_type_columns.iter()) - .fold( - HashMap::new(), - |mut acc, - ( - flattened_array, - ListUnnest { - index_in_input_schema, - depth, - }, - )| { - acc.entry(*index_in_input_schema) - .or_insert(vec![]) - .push(flattened_array); - acc - }, - ); - - // Create the take indices array for other columns - let take_indicies = create_take_indicies(unnested_length, total_length); - - // vertical expansion because of list unnest - let ret = flatten_list_cols_from_indices( - batch, - unnested_array_map, - &take_indicies, - )?; - flatten_struct_cols(&ret, schema, struct_column_indices) - } - }; - transformed -} +// fn build_batch( +// batch: &RecordBatch, +// schema: &SchemaRef, +// list_type_columns: &[ListUnnest], +// struct_column_indices: &HashSet, +// options: &UnnestOptions, +// ) -> Result { +// let transformed = match list_type_columns.len() { +// 0 => flatten_struct_cols(batch.columns(), schema, struct_column_indices), +// _ => { +// // maintain a map of temp result +// // > +// // given a map of unnesting +// // col1: +// // col2: +// // 1.run for each to unnest all the top level +// // col1: => to temp +// // 2.run unnest on remaining level +// // col1: => to temp +// // col2: => to temp +// let list_arrays: Vec<(ArrayRef, usize)> = list_type_columns +// .iter() +// .map( +// |ListUnnest { +// depth, +// index_in_input_schema, +// }| { +// Ok(( +// ColumnarValue::Array(Arc::clone( +// batch.column(*index_in_input_schema), +// )) +// .into_array(batch.num_rows())?, +// *depth, +// )) +// }, +// ) +// .collect::>()?; + +// let longest_length = find_longest_length(&list_arrays, options)?; +// let unnested_length = longest_length.as_primitive::(); +// let total_length = if unnested_length.is_empty() { +// 0 +// } else { +// sum(unnested_length).ok_or_else(|| { +// exec_datafusion_err!("Failed to calculate the total unnested length") +// })? as usize +// }; +// if total_length == 0 { +// return Ok(RecordBatch::new_empty(Arc::clone(schema))); +// } + +// // Unnest all the list arrays +// let unnested_arrays = +// unnest_list_arrays(&list_arrays, unnested_length, total_length)?; +// let unnested_array_map: HashMap>> = unnested_arrays +// .into_iter() +// .zip(list_type_columns.iter()) +// .fold( +// HashMap::new(), +// |mut acc, +// ( +// flattened_array, +// ListUnnest { +// index_in_input_schema, +// depth, +// }, +// )| { +// acc.entry(*index_in_input_schema) +// .or_insert(vec![]) +// .push(flattened_array); +// acc +// }, +// ); + +// // Create the take indices array for other columns +// let take_indicies = create_take_indicies(unnested_length, total_length); + +// // vertical expansion because of list unnest +// let ret = flatten_list_cols_from_indices( +// batch, +// unnested_array_map, +// &take_indicies, +// )?; +// flatten_struct_cols(&ret, schema, struct_column_indices) +// } +// }; +// transformed +// } /// TODO: only prototype /// This function does not handle NULL yet @@ -1020,8 +1053,8 @@ fn flatten_list_cols_from_indices( #[cfg(test)] mod tests { use super::*; - use arrow::datatypes::{Field, Int32Type}; - use arrow_array::{GenericListArray, Int32Array, OffsetSizeTrait, StringArray}; + use arrow::datatypes::Field; + use arrow_array::{GenericListArray, OffsetSizeTrait, StringArray}; use arrow_buffer::{BooleanBufferBuilder, NullBuffer, OffsetBuffer}; // Create a GenericListArray with the following list values: @@ -1100,238 +1133,11 @@ mod tests { expected: Vec>, ) -> datafusion_common::Result<()> { let length_array = Int64Array::from(lengths); - let unnested_array = unnest_list_array(list_array, &length_array, 3 * 6, 1)?; + let unnested_array = unnest_list_array(list_array, &length_array, 3 * 6)?; let strs = unnested_array.as_string::().iter().collect::>(); assert_eq!(strs, expected); Ok(()) } - #[test] - fn test_todo() -> datafusion_common::Result<()> { - // original: [[1,2,3],null,[4,5]] [[7,8,9,10],null,[11,12,13]] null - // unnest depth=1:[1,2,3], null, [4,5], [7,8,9,10], null, [11,12,13] , null - // unnest depth=2: 1,2,3,null,4,5,7,8,9,10,null,11,12,13,null - let list_arr = ListArray::from_iter_primitive::(vec![ - Some(vec![Some(1), Some(2), Some(3)]), - None, - Some(vec![Some(4), Some(5)]), - Some(vec![Some(7), Some(8), Some(9), Some(10)]), - None, - Some(vec![Some(11), Some(12), Some(13)]), - ]); - - let list_arr_ref = Arc::new(list_arr.clone()) as ArrayRef; - // last item is null - let offsets = OffsetBuffer::from_lengths([3, 3, 0]); - - let mut nulls = BooleanBufferBuilder::new(2); - nulls.append(true); - nulls.append(true); - nulls.append(false); - let nested_list_arr = ListArray::new( - Arc::new(Field::new_list_field( - list_arr_ref.data_type().to_owned(), - true, - )), - offsets, - list_arr_ref, - Some(NullBuffer::new(nulls.finish())), - ); - // TODO: for each row we have to check if it is not null, then iterate recursively to find the offset - // until the depth = 1 - // Seed some None values at the top level, second top level to check how to null works - for row in 0..list_arr.len() { - // TODO: verify null preseving behavior - if !list_arr.is_null(row) { - let (start, end) = (&list_arr as &dyn ListArrayType).value_offsets(row); - for i in start..end {} - } - } - nested_list_arr - .values() - .as_list::() - .offsets() - .iter() - .for_each(|i| { - println!("{}", i); - }); - let nested_list_arr_ref = Arc::new(nested_list_arr.clone()) as ArrayRef; - let longest_length = find_longest_length( - &vec![(nested_list_arr_ref, 2)], - &UnnestOptions { - preserve_nulls: true, - }, - )?; - let longest_length = longest_length.as_primitive::(); - let unnested = - unnest_list_array_prototype(&nested_list_arr, longest_length, 14, 2)?; - let expected_array = Int32Array::from(Vec::::from([1, 2, 3])); - assert_eq!( - unnested.as_any().downcast_ref::().unwrap(), - &expected_array - ); - - Ok(()) - // let unnested_length = longest_length.as_primitive::(); - // unnest_list_arrays(list_arrays, length_array, capacity) - // unnest_list_array(list_array, length_array, capacity, depth) - } - - fn unnest_list_array_prototype_recursive( - index_builder: &mut PrimitiveArray, - list_array: &GenericListArray, - // length_array: &PrimitiveArray, - // capacity: usize, - depth: usize, - ) -> Result { - Ok(0) - // if depth == 1 { - // for row in 0..list_array.len() { - // let mut value_length = 0; - // if !list_array.is_null(row) { - // let (start, end) = - // (list_array as &dyn ListArrayType).value_offsets(row); - // value_length = end - start; - // for i in start..end { - // if depth == 1 { - // take_indicies_builder.append_value(i) - // } else { - // } - // } - // } - // } - // } else { - // for row in 0..list_array.len() { - // let mut value_length = 0; - // if !list_array.is_null(row) { - // let (start, end) = - // (list_array as &dyn ListArrayType).value_offsets(row); - // value_length = end - start; - // for i in start..end { - // if depth == 1 { - // take_indicies_builder.append_value(i) - // } else { - // } - // } - // } - // } - // let inner_list_array = list_array.values().as_list::(); - // unnest_list_array_prototype_recursive( - // index_builder, - // inner_list_array, - // depth - 1, - // ) - // } - } - // [[1,2,3],null,[4,5]] 1 - // [[7,8,9,10],null,[11,12,13]] 1 - // null 1 - // result into - // 1,2,3,null,4,5,7,8,9,10,null,11,12,13 - // - fn unnest_list_array_prototype( - list_array: &GenericListArray, - length_array: &PrimitiveArray, - capacity: usize, - depth: usize, - ) -> Result { - let mut cur_depth = depth; - let mut backed_array = list_array.values(); - loop { - if cur_depth == 1 { - break; - } - backed_array = backed_array.as_list::().values(); - cur_depth -= 1; - } - - // TODO: handle me recursively - let mut take_indicies_builder = PrimitiveArray::::builder(capacity); - for row in 0..list_array.len() { - let mut value_length = 0; - if !list_array.is_null(row) { - let (start, end) = (list_array as &dyn ListArrayType).value_offsets(row); - value_length = end - start; - for i in start..end { - if depth == 1 { - take_indicies_builder.append_value(i) - } else { - } - } - } - let target_length = length_array.value(row); - debug_assert!( - value_length <= target_length, - "value length is beyond the longest length" - ); - // Pad with NULL values - for _ in value_length..target_length { - take_indicies_builder.append_null(); - } - } - Ok(kernels::take::take( - &backed_array, - &take_indicies_builder.finish(), - None, - )?) - } - - #[test] - fn test_length_recursive() -> datafusion_common::Result<()> { - // [[1,2,3],null,[4,5]],[[7,8,9,10],null,[11,12,13]] null - let list_arr = ListArray::from_iter_primitive::(vec![ - Some(vec![Some(1), Some(2), Some(3)]), - None, - Some(vec![Some(4), Some(5)]), - Some(vec![Some(7), Some(8), Some(9), Some(10)]), - None, - Some(vec![Some(11), Some(12), Some(13)]), - ]); - let list_arr_ref = Arc::new(list_arr) as ArrayRef; - let offsets = OffsetBuffer::from_lengths([3, 3, 0]); - let mut nulls = BooleanBufferBuilder::new(3); - nulls.append(true); - nulls.append(true); - nulls.append(false); - let nested_list_arr = ListArray::new( - Arc::new(Field::new_list_field( - list_arr_ref.data_type().to_owned(), - true, - )), - offsets, - list_arr_ref, - Some(NullBuffer::new(nulls.finish())), - ); - let list_arr_2 = ListArray::from_iter_primitive::(vec![ - Some(vec![Some(1), Some(2), Some(3)]), - None, - Some(vec![Some(4), Some(5)]), - Some(vec![Some(7), Some(8), Some(9), Some(10)]), - None, - Some(vec![Some(11), Some(12), Some(13)]), - ]); - let list_arr_ref = Arc::new(list_arr) as ArrayRef; - let offsets = OffsetBuffer::from_lengths([3, 3, 0]); - let mut nulls = BooleanBufferBuilder::new(3); - nulls.append(true); - nulls.append(true); - nulls.append(false); - let nested_list_arr = ListArray::new( - Arc::new(Field::new_list_field( - list_arr_ref.data_type().to_owned(), - true, - )), - offsets, - list_arr_ref, - Some(NullBuffer::new(nulls.finish())), - ); - // [6,8,0] - verify_longest_length( - &[(Arc::new(nested_list_arr) as ArrayRef, 2)], - true, - vec![6, 8, 0], - )?; - Ok(()) - } #[test] fn test_unnest_list_array() -> datafusion_common::Result<()> { @@ -1379,7 +1185,7 @@ mod tests { } fn verify_longest_length( - list_arrays: &[(ArrayRef, usize)], + list_arrays: &[ArrayRef], preserve_nulls: bool, expected: Vec, ) -> datafusion_common::Result<()> { @@ -1401,51 +1207,27 @@ mod tests { // Test with single ListArray // [A, B, C], [], NULL, [D], NULL, [NULL, F] let list_array = Arc::new(make_generic_array::()) as ArrayRef; - verify_longest_length( - &[(Arc::clone(&list_array), 1)], - false, - vec![3, 0, 0, 1, 0, 2], - )?; - verify_longest_length( - &[(Arc::clone(&list_array), 1)], - true, - vec![3, 0, 1, 1, 1, 2], - )?; + verify_longest_length(&[Arc::clone(&list_array)], false, vec![3, 0, 0, 1, 0, 2])?; + verify_longest_length(&[Arc::clone(&list_array)], true, vec![3, 0, 1, 1, 1, 2])?; // Test with single LargeListArray // [A, B, C], [], NULL, [D], NULL, [NULL, F] let list_array = Arc::new(make_generic_array::()) as ArrayRef; - verify_longest_length( - &[(Arc::clone(&list_array), 1)], - false, - vec![3, 0, 0, 1, 0, 2], - )?; - verify_longest_length( - &[(Arc::clone(&list_array), 1)], - true, - vec![3, 0, 1, 1, 1, 2], - )?; + verify_longest_length(&[Arc::clone(&list_array)], false, vec![3, 0, 0, 1, 0, 2])?; + verify_longest_length(&[Arc::clone(&list_array)], true, vec![3, 0, 1, 1, 1, 2])?; // Test with single FixedSizeListArray // [A, B], NULL, [C, D], NULL, [NULL, F], [NULL, NULL] let list_array = Arc::new(make_fixed_list()) as ArrayRef; - verify_longest_length( - &[(Arc::clone(&list_array), 1)], - false, - vec![2, 0, 2, 0, 2, 2], - )?; - verify_longest_length( - &[(Arc::clone(&list_array), 1)], - true, - vec![2, 1, 2, 1, 2, 2], - )?; + verify_longest_length(&[Arc::clone(&list_array)], false, vec![2, 0, 2, 0, 2, 2])?; + verify_longest_length(&[Arc::clone(&list_array)], true, vec![2, 1, 2, 1, 2, 2])?; // Test with multiple list arrays // [A, B, C], [], NULL, [D], NULL, [NULL, F] // [A, B], NULL, [C, D], NULL, [NULL, F], [NULL, NULL] let list1 = Arc::new(make_generic_array::()) as ArrayRef; let list2 = Arc::new(make_fixed_list()) as ArrayRef; - let list_arrays = vec![(Arc::clone(&list1), 1), (Arc::clone(&list2), 1)]; + let list_arrays = vec![Arc::clone(&list1), Arc::clone(&list2)]; verify_longest_length(&list_arrays, false, vec![3, 0, 2, 1, 2, 2])?; verify_longest_length(&list_arrays, true, vec![3, 1, 2, 1, 2, 2])?; From 5b9ce5ef11e8b2e37e7cb93b7fcc7b6436fe46ce Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Thu, 15 Aug 2024 22:13:38 +0200 Subject: [PATCH 20/56] UT for new build batch function --- datafusion/physical-plan/src/unnest.rs | 242 +++++++++++++++++++------ datafusion/sql/src/utils.rs | 10 +- 2 files changed, 194 insertions(+), 58 deletions(-) diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index 444d2c689af9..f0d58a8753d6 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -43,7 +43,6 @@ use datafusion_common::{ exec_datafusion_err, exec_err, internal_err, Result, UnnestOptions, }; use datafusion_execution::TaskContext; -use datafusion_expr::ColumnarValue; use datafusion_physical_expr::EquivalenceProperties; use async_trait::async_trait; @@ -433,28 +432,36 @@ fn unnest_at_level( // temp_batch: [(a,depth=1,unnested(a)),(a,depth=2,unnested(unnested(a))),(b,depth=1,unnested(b))] // this is final, now combine the mainbatch and temp_batch - let temp_unnest_cols = list_type_columns - .iter() - .filter_map( - |ListUnnest { - depth, - index_in_input_schema, - }| { - if *depth == level_to_unnest { - return Some(Arc::clone(&batch[*index_in_input_schema])); - } - if *depth > level_to_unnest { - return Some( - temp_batch - .get(&(*index_in_input_schema, *depth)) - .unwrap() - .clone(), - ); - } - return None; - }, - ) - .collect::>(); + let (temp_unnest_cols, unnested_locations): (Vec>, Vec<_>) = + list_type_columns + .iter() + .filter_map( + |ListUnnest { + depth, + index_in_input_schema, + }| { + if *depth == level_to_unnest { + return Some(( + Arc::clone(&batch[*index_in_input_schema]), + (*index_in_input_schema, *depth), + )); + } + // this means the depth to unnest is still on going, keep on unnesting + // at current level + if *depth > level_to_unnest { + return Some(( + Arc::clone( + temp_batch + .get(&(*index_in_input_schema, *depth)) + .unwrap(), + ), + (*index_in_input_schema, *depth), + )); + } + None + }, + ) + .unzip(); // filter out so that list_arrays only contain column with the highest depth // at the same time, during iteration remove this depth so next time we don't have to unnest them again @@ -473,29 +480,31 @@ fn unnest_at_level( // Unnest all the list arrays let unnested_temp_arrays = - unnest_list_arrays(&temp_unnest_cols, unnested_length, total_length)?; - - unnested_temp_arrays - .into_iter() - .zip(list_type_columns.iter()) - .for_each( - |( - flatten_arr, - ListUnnest { - index_in_input_schema, - depth, - }, - )| { - temp_batch.insert((*index_in_input_schema, *depth), flatten_arr); - }, - ); + unnest_list_arrays(temp_unnest_cols.as_ref(), unnested_length, total_length)?; + unnested_temp_arrays.iter().for_each(|arr| { + println!("ret {}", arr.len()); + }); // Create the take indices array for other columns let take_indices = create_take_indicies(unnested_length, total_length); + println!("take indices {}", take_indices.len()); // vertical expansion because of list unnest let ret = flatten_batch_from_indices(batch, &take_indices)?; - return Ok(ret); + ret.iter().for_each(|arr| { + println!("ret {}", arr.len()); + }); + unnested_temp_arrays + .into_iter() + .zip(unnested_locations.iter()) + .for_each(|(flatten_arr, (index_in_input_schema, depth))| { + println!( + "inserting into temp value {} {}", + index_in_input_schema, depth + ); + temp_batch.insert((*index_in_input_schema, *depth), flatten_arr); + }); + Ok(ret) } fn build_batch_v2( @@ -516,26 +525,28 @@ fn build_batch_v2( }); let mut unnested_original_columns = vec![]; for depth in (1..=highest_depth).rev() { - let input = match depth { - highest_depth => batch.columns(), - _ => &unnested_original_columns, + let input = match depth == highest_depth { + true => batch.columns(), + false => &unnested_original_columns, }; + println!("unnesting at level {}", depth); let temp = unnest_at_level( input, list_type_columns, &mut temp_batch, - highest_depth, + depth, options, )?; unnested_original_columns = temp; } + println!("{} temp batch", temp_batch.len()); // TODO: combine temp with the batch let unnested_array_map: HashMap, usize)>> = temp_batch.into_iter().fold( HashMap::new(), |mut acc, ((index_in_input_schema, depth), flattened_array)| { acc.entry(index_in_input_schema) - .or_insert(vec![]) + .or_default() .push((flattened_array, depth)); acc }, @@ -552,6 +563,7 @@ fn build_batch_v2( }) .collect(); + println!("unnested array map {}", unnested_array_map.len()); let mut ordered_unnested_array_map = unnested_array_map .into_iter() .map( @@ -579,16 +591,16 @@ fn build_batch_v2( // .fold(0, |highest_depth, ListUnnest { depth, .. }| { // cmp::max(highest_depth, *depth) // }); + println!("len of map {}", ordered_unnested_array_map.len()); let ret = unnested_original_columns - .iter() + .into_iter() .enumerate() - .map( - |(col_idx, arr)| match ordered_unnested_array_map.remove(&col_idx) { + .flat_map(|(col_idx, arr)| { + match ordered_unnested_array_map.remove(&col_idx) { Some(unnested_arrays) => unnested_arrays, - None => vec![arr.clone()], - }, - ) - .flatten() + None => vec![arr], + } + }) .collect::>(); flatten_struct_cols(&ret, schema, struct_column_indices) @@ -1053,7 +1065,7 @@ fn flatten_list_cols_from_indices( #[cfg(test)] mod tests { use super::*; - use arrow::datatypes::Field; + use arrow::{datatypes::Field, util::pretty::pretty_format_batches}; use arrow_array::{GenericListArray, OffsetSizeTrait, StringArray}; use arrow_buffer::{BooleanBufferBuilder, NullBuffer, OffsetBuffer}; @@ -1138,6 +1150,130 @@ mod tests { assert_eq!(strs, expected); Ok(()) } + #[test] + fn test_build_batch() -> datafusion_common::Result<()> { + // col1 | col2 + // [[1,2,3],null,[4,5]] | ['a','b'] + // [[7,8,9,10], null, [11,12,13]] | ['c','d'] + // null | ['e'] + let list_arr1 = ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(2), Some(3)]), + None, + Some(vec![Some(4), Some(5)]), + Some(vec![Some(7), Some(8), Some(9), Some(10)]), + None, + Some(vec![Some(11), Some(12), Some(13)]), + ]); + + let list_arr1_ref = Arc::new(list_arr1) as ArrayRef; + let offsets = OffsetBuffer::from_lengths([3, 3, 0]); + let mut nulls = BooleanBufferBuilder::new(3); + nulls.append(true); + nulls.append(true); + nulls.append(false); + // list> + let col1_field = Field::new_list_field( + DataType::List(Arc::new(Field::new_list_field( + list_arr1_ref.data_type().to_owned(), + true, + ))), + true, + ); + let col1 = ListArray::new( + Arc::new(Field::new_list_field( + list_arr1_ref.data_type().to_owned(), + true, + )), + offsets, + list_arr1_ref, + Some(NullBuffer::new(nulls.finish())), + ); + + let list_arr2 = StringArray::from(vec![ + Some("a"), + Some("b"), + Some("c"), + Some("d"), + Some("e"), + ]); + + let offsets = OffsetBuffer::from_lengths([2, 2, 1]); + let mut nulls = BooleanBufferBuilder::new(3); + nulls.append_n(3, true); + let col2_field = Field::new( + "col2", + DataType::List(Arc::new(Field::new_list_field(DataType::Utf8, true))), + true, + ); + let col2 = GenericListArray::::new( + Arc::new(Field::new_list_field(DataType::Utf8, true)), + OffsetBuffer::new(offsets.into()), + Arc::new(list_arr2), + Some(NullBuffer::new(nulls.finish())), + ); + // convert col1 and col2 to a record batch + let schema = Arc::new(Schema::new(vec![col1_field, col2_field])); + let out_schema = Arc::new(Schema::new(vec![ + Field::new( + "col1_unnest_placeholder_depth_1", + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), + true, + ), + Field::new("col1_unnest_placeholder_depth_2", DataType::Int32, true), + Field::new("col2_unnest_placeholder_depth_1", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(col1) as ArrayRef, Arc::new(col2) as ArrayRef], + ) + .unwrap(); + let list_type_columns = vec![ + ListUnnest { + index_in_input_schema: 0, + depth: 1, + }, + ListUnnest { + index_in_input_schema: 0, + depth: 2, + }, + ListUnnest { + index_in_input_schema: 1, + depth: 1, + }, + ]; + let ret = build_batch_v2( + &batch, + &out_schema, + list_type_columns.as_ref(), + &HashSet::default(), + &UnnestOptions { + preserve_nulls: true, + }, + )?; + let actual = + format!("{}", pretty_format_batches(vec![ret].as_ref())?).to_lowercase(); + let expected_asc = r#" ++----------------------------------+--------------------------+ +| trace_id | max(traces.timestamp_ms) | ++----------------------------------+--------------------------+ +| 5868861a23ed31355efc5200eb80fe74 | 16909009999999 | +| 4040e64656804c3d77320d7a0e7eb1f0 | 16909009999998 | +| 02801bbe533190a9f8713d75222f445d | 16909009999997 | +| 9e31b3b5a620de32b68fefa5aeea57f1 | 16909009999996 | +| 2d88a860e9bd1cfaa632d8e7caeaa934 | 16909009999995 | +| a47edcef8364ab6f191dd9103e51c171 | 16909009999994 | +| 36a3fa2ccfbf8e00337f0b1254384db6 | 16909009999993 | +| 0756be84f57369012e10de18b57d8a2f | 16909009999992 | +| d4d6bf9845fa5897710e3a8db81d5907 | 16909009999991 | +| 3c2cc1abe728a66b61e14880b53482a0 | 16909009999990 | ++----------------------------------+--------------------------+ + "# + .trim(); + + println!("{}", actual); + // format pretty record batch + Ok(()) + } #[test] fn test_unnest_list_array() -> datafusion_common::Result<()> { diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index 657524e14064..0644c1877c78 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -607,7 +607,7 @@ mod tests { ); // memoization only contains 1 transformation assert_eq!(memo.len(), 1); - assert!(memo.get(&Column::from_name("3d_col")).is_some()); + assert!(memo.get("3d_col").is_some()); column_unnests_eq( vec!["unnest_placeholder(3d_col)"], &unnest_placeholder_columns, @@ -640,7 +640,7 @@ mod tests { // memoization still contains 1 transformation // and the previous transformation is reused assert_eq!(memo.len(), 1); - assert!(memo.get(&Column::from_name("3d_col")).is_some()); + assert!(memo.get("3d_col").is_some()); column_unnests_eq( vec!["unnest_placeholder(3d_col)"], &mut unnest_placeholder_columns, @@ -678,7 +678,7 @@ mod tests { // and the previous transformation is reused assert_eq!(memo.len(), 2); - assert!(memo.get(&Column::from_name("struct_arr_col")).is_some()); + assert!(memo.get("struct_arr_col").is_some()); column_unnests_eq( vec![ "unnest_placeholder(3d_col)", @@ -750,7 +750,7 @@ mod tests { ); // memoization only contains 1 transformation assert_eq!(memo.len(), 1); - assert!(memo.get(&Column::from_name("3d_col")).is_some()); + assert!(memo.get("3d_col").is_some()); column_unnests_eq( vec!["unnest_placeholder(3d_col)|List([unnest_placeholder(3d_col,depth=2)|depth=2])"], &unnest_placeholder_columns, @@ -783,7 +783,7 @@ mod tests { // memoization still contains 1 transformation // and the for the same column, depth = 1 needs to be performed aside from depth = 2 assert_eq!(memo.len(), 1); - assert!(memo.get(&Column::from_name("3d_col")).is_some()); + assert!(memo.get("3d_col").is_some()); column_unnests_eq( vec!["unnest_placeholder(3d_col)|List([unnest_placeholder(3d_col,depth=2)|depth=2, unnest_placeholder(3d_col,depth=1)|depth=1])"], &unnest_placeholder_columns, From 366d8aefab0e1c70b40718e5120b66d0af7bae6d Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Thu, 15 Aug 2024 22:25:47 +0200 Subject: [PATCH 21/56] compile err --- datafusion/physical-plan/src/unnest.rs | 241 ++++--------------- datafusion/sqllogictest/test_files/debug.slt | 77 +++--- 2 files changed, 72 insertions(+), 246 deletions(-) diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index f0d58a8753d6..856f8e88d680 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -35,9 +35,9 @@ use arrow::array::{ use arrow::compute::kernels::length::length; use arrow::compute::kernels::zip::zip; use arrow::compute::{cast, is_not_null, kernels, sum}; -use arrow::datatypes::{DataType, Int32Type, Int64Type, Schema, SchemaRef}; +use arrow::datatypes::{DataType, Int64Type, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; -use arrow_array::{Int32Array, Int64Array, Scalar, StructArray}; +use arrow_array::{Int64Array, Scalar, StructArray}; use arrow_ord::cmp::lt; use datafusion_common::{ exec_datafusion_err, exec_err, internal_err, Result, UnnestOptions, @@ -267,7 +267,7 @@ impl UnnestStream { .map(|maybe_batch| match maybe_batch { Some(Ok(batch)) => { let timer = self.metrics.elapsed_compute.timer(); - let result = build_batch_v2( + let result = build_batch( &batch, &self.schema, &self.list_type_columns, @@ -481,33 +481,27 @@ fn unnest_at_level( // Unnest all the list arrays let unnested_temp_arrays = unnest_list_arrays(temp_unnest_cols.as_ref(), unnested_length, total_length)?; - unnested_temp_arrays.iter().for_each(|arr| { - println!("ret {}", arr.len()); - }); // Create the take indices array for other columns let take_indices = create_take_indicies(unnested_length, total_length); - println!("take indices {}", take_indices.len()); // vertical expansion because of list unnest let ret = flatten_batch_from_indices(batch, &take_indices)?; - ret.iter().for_each(|arr| { - println!("ret {}", arr.len()); - }); unnested_temp_arrays .into_iter() .zip(unnested_locations.iter()) .for_each(|(flatten_arr, (index_in_input_schema, depth))| { - println!( - "inserting into temp value {} {}", - index_in_input_schema, depth - ); temp_batch.insert((*index_in_input_schema, *depth), flatten_arr); }); Ok(ret) } -fn build_batch_v2( +/// For each row in a `RecordBatch`, some list/struct columns need to be unnested. +/// - For list columns: We will expand the values in each list into multiple rows, +/// taking the longest length among these lists, and shorter lists are padded with NULLs. +/// - For struct columns: We will expand the struct columns into multiple subfield columns. +/// For columns that don't need to be unnested, repeat their values until reaching the longest length. +fn build_batch( batch: &RecordBatch, schema: &SchemaRef, list_type_columns: &[ListUnnest], @@ -529,7 +523,6 @@ fn build_batch_v2( true => batch.columns(), false => &unnested_original_columns, }; - println!("unnesting at level {}", depth); let temp = unnest_at_level( input, list_type_columns, @@ -539,8 +532,6 @@ fn build_batch_v2( )?; unnested_original_columns = temp; } - println!("{} temp batch", temp_batch.len()); - // TODO: combine temp with the batch let unnested_array_map: HashMap, usize)>> = temp_batch.into_iter().fold( HashMap::new(), @@ -563,7 +554,6 @@ fn build_batch_v2( }) .collect(); - println!("unnested array map {}", unnested_array_map.len()); let mut ordered_unnested_array_map = unnested_array_map .into_iter() .map( @@ -586,12 +576,7 @@ fn build_batch_v2( }, ) .collect::>(); - // let highest_depth = list_type_columns - // .iter() - // .fold(0, |highest_depth, ListUnnest { depth, .. }| { - // cmp::max(highest_depth, *depth) - // }); - println!("len of map {}", ordered_unnested_array_map.len()); + let ret = unnested_original_columns .into_iter() .enumerate() @@ -609,157 +594,6 @@ fn build_batch_v2( transformed } -/// For each row in a `RecordBatch`, some list/struct columns need to be unnested. -/// - For list columns: We will expand the values in each list into multiple rows, -/// taking the longest length among these lists, and shorter lists are padded with NULLs. -/// - For struct columns: We will expand the struct columns into multiple subfield columns. -/// For columns that don't need to be unnested, repeat their values until reaching the longest length. -// fn build_batch( -// batch: &RecordBatch, -// schema: &SchemaRef, -// list_type_columns: &[ListUnnest], -// struct_column_indices: &HashSet, -// options: &UnnestOptions, -// ) -> Result { -// let transformed = match list_type_columns.len() { -// 0 => flatten_struct_cols(batch.columns(), schema, struct_column_indices), -// _ => { -// // maintain a map of temp result -// // > -// // given a map of unnesting -// // col1: -// // col2: -// // 1.run for each to unnest all the top level -// // col1: => to temp -// // 2.run unnest on remaining level -// // col1: => to temp -// // col2: => to temp -// let list_arrays: Vec<(ArrayRef, usize)> = list_type_columns -// .iter() -// .map( -// |ListUnnest { -// depth, -// index_in_input_schema, -// }| { -// Ok(( -// ColumnarValue::Array(Arc::clone( -// batch.column(*index_in_input_schema), -// )) -// .into_array(batch.num_rows())?, -// *depth, -// )) -// }, -// ) -// .collect::>()?; - -// let longest_length = find_longest_length(&list_arrays, options)?; -// let unnested_length = longest_length.as_primitive::(); -// let total_length = if unnested_length.is_empty() { -// 0 -// } else { -// sum(unnested_length).ok_or_else(|| { -// exec_datafusion_err!("Failed to calculate the total unnested length") -// })? as usize -// }; -// if total_length == 0 { -// return Ok(RecordBatch::new_empty(Arc::clone(schema))); -// } - -// // Unnest all the list arrays -// let unnested_arrays = -// unnest_list_arrays(&list_arrays, unnested_length, total_length)?; -// let unnested_array_map: HashMap>> = unnested_arrays -// .into_iter() -// .zip(list_type_columns.iter()) -// .fold( -// HashMap::new(), -// |mut acc, -// ( -// flattened_array, -// ListUnnest { -// index_in_input_schema, -// depth, -// }, -// )| { -// acc.entry(*index_in_input_schema) -// .or_insert(vec![]) -// .push(flattened_array); -// acc -// }, -// ); - -// // Create the take indices array for other columns -// let take_indicies = create_take_indicies(unnested_length, total_length); - -// // vertical expansion because of list unnest -// let ret = flatten_list_cols_from_indices( -// batch, -// unnested_array_map, -// &take_indicies, -// )?; -// flatten_struct_cols(&ret, schema, struct_column_indices) -// } -// }; -// transformed -// } - -/// TODO: only prototype -/// This function does not handle NULL yet -pub fn length_recursive( - array: &dyn Array, - depth: usize, - preserve_nulls: bool, - is_top_level: bool, // null values at top level recursion is ignored, regardless of preserve_null -) -> Result { - let list = array.as_list::(); - let null_length = if preserve_nulls { - Scalar::new(Int32Array::from_value(1, 1)) - } else { - Scalar::new(Int32Array::from_value(0, 1)) - }; - // preserve null only concern nullability of the element in array, not the value of the array itself - // e.g unnest(null) => 0 rows - // while unnest([null]) => 1 row with null value - let null_length_for_value = match preserve_nulls && !is_top_level { - true => 1, - false => 0, - }; - if depth == 1 { - if array.is_empty() { - return Ok(Arc::new(PrimitiveArray::::from(vec![ - null_length_for_value, - ]))); - } - - // null elements will have null length - // respect preserve_nulls here - let mut ret = length(array)?; - ret = cast(&ret, &DataType::Int32)?; - return Ok(zip(&is_not_null(&ret)?, &ret, &null_length)?); - } - let a: Vec = list - .iter() - .map(|x| { - if x.is_none() { - return null_length_for_value; - } - // [[1,2,3],null,[4,5]] - // [[7,8,9,10],] - // length([1,2]) + length(null) + length([4,5]) - let ret = - length_recursive(&x.unwrap(), depth - 1, preserve_nulls, false).unwrap(); - let a = ret.as_primitive::(); - if a.is_empty() { - // TODO respect is_preserve_nulls - 1 - } else { - sum(a).unwrap() - } - }) - .collect(); - Ok(Arc::new(PrimitiveArray::::from(a))) -} - /// Find the longest list length among the given list arrays for each row. /// /// For example if we have the following two list arrays: @@ -801,9 +635,6 @@ fn find_longest_length( let list_lengths: Vec = list_arrays .iter() .map(|list_array| { - // Perhaps we don't need recursive here - // let mut length_array = - // length_recursive(list_array, *depth, options.preserve_nulls, true)?; let mut length_array = length(list_array)?; // Make sure length arrays have the same type. Int64 is the most general one. // Respect the depth of unnest( current func only get the length of 1 level of unnest) @@ -1065,7 +896,10 @@ fn flatten_list_cols_from_indices( #[cfg(test)] mod tests { use super::*; - use arrow::{datatypes::Field, util::pretty::pretty_format_batches}; + use arrow::{ + datatypes::{Field, Int32Type}, + util::pretty::pretty_format_batches, + }; use arrow_array::{GenericListArray, OffsetSizeTrait, StringArray}; use arrow_buffer::{BooleanBufferBuilder, NullBuffer, OffsetBuffer}; @@ -1150,6 +984,7 @@ mod tests { assert_eq!(strs, expected); Ok(()) } + #[test] fn test_build_batch() -> datafusion_common::Result<()> { // col1 | col2 @@ -1241,7 +1076,7 @@ mod tests { depth: 1, }, ]; - let ret = build_batch_v2( + let ret = build_batch( &batch, &out_schema, list_type_columns.as_ref(), @@ -1252,26 +1087,34 @@ mod tests { )?; let actual = format!("{}", pretty_format_batches(vec![ret].as_ref())?).to_lowercase(); - let expected_asc = r#" -+----------------------------------+--------------------------+ -| trace_id | max(traces.timestamp_ms) | -+----------------------------------+--------------------------+ -| 5868861a23ed31355efc5200eb80fe74 | 16909009999999 | -| 4040e64656804c3d77320d7a0e7eb1f0 | 16909009999998 | -| 02801bbe533190a9f8713d75222f445d | 16909009999997 | -| 9e31b3b5a620de32b68fefa5aeea57f1 | 16909009999996 | -| 2d88a860e9bd1cfaa632d8e7caeaa934 | 16909009999995 | -| a47edcef8364ab6f191dd9103e51c171 | 16909009999994 | -| 36a3fa2ccfbf8e00337f0b1254384db6 | 16909009999993 | -| 0756be84f57369012e10de18b57d8a2f | 16909009999992 | -| d4d6bf9845fa5897710e3a8db81d5907 | 16909009999991 | -| 3c2cc1abe728a66b61e14880b53482a0 | 16909009999990 | -+----------------------------------+--------------------------+ + let expected = r#" ++---------------------------------+---------------------------------+---------------------------------+ +| col1_unnest_placeholder_depth_1 | col1_unnest_placeholder_depth_2 | col2_unnest_placeholder_depth_1 | ++---------------------------------+---------------------------------+---------------------------------+ +| [1, 2, 3] | 1 | a | +| | 2 | b | +| [4, 5] | 3 | | +| [1, 2, 3] | | a | +| | | b | +| [4, 5] | | | +| [1, 2, 3] | 4 | a | +| | 5 | b | +| [4, 5] | | | +| [7, 8, 9, 10] | 7 | c | +| | 8 | d | +| [11, 12, 13] | 9 | | +| | 10 | | +| [7, 8, 9, 10] | | c | +| | | d | +| [11, 12, 13] | | | +| [7, 8, 9, 10] | 11 | c | +| | 12 | d | +| [11, 12, 13] | 13 | | +| | | e | ++---------------------------------+---------------------------------+---------------------------------+ "# .trim(); - - println!("{}", actual); - // format pretty record batch + assert_eq!(actual, expected); Ok(()) } diff --git a/datafusion/sqllogictest/test_files/debug.slt b/datafusion/sqllogictest/test_files/debug.slt index 51c267f04186..18f908e3e55b 100644 --- a/datafusion/sqllogictest/test_files/debug.slt +++ b/datafusion/sqllogictest/test_files/debug.slt @@ -6,10 +6,10 @@ AS VALUES ; query TT -select unnest(unnest(unnest(column2))) from recursive_unnest_table; +explain select unnest(unnest(unnest(column2))) from recursive_unnest_table; ---- logical_plan -01)Unnest: lists[unnest_placeholder(recursive_unnest_table.column2)|depth=2] structs[] +01)Unnest: lists[unnest_placeholder(recursive_unnest_table.column2)|depth=3] structs[] 02)--Projection: recursive_unnest_table.column2 AS unnest_placeholder(recursive_unnest_table.column2) 03)----TableScan: recursive_unnest_table projection=[column2] physical_plan @@ -18,55 +18,38 @@ physical_plan 03)----MemoryExec: partitions=1, partition_sizes=[1] query TT -select unnest(unnest(column2)) from recursive_unnest_table; +explain select unnest(unnest(column2)) from recursive_unnest_table; ---- +logical_plan +01)Unnest: lists[unnest_placeholder(recursive_unnest_table.column2)|depth=2] structs[] +02)--Projection: recursive_unnest_table.column2 AS unnest_placeholder(recursive_unnest_table.column2) +03)----TableScan: recursive_unnest_table projection=[column2] +physical_plan +01)UnnestExec +02)--ProjectionExec: expr=[column2@0 as unnest_placeholder(recursive_unnest_table.column2)] +03)----MemoryExec: partitions=1, partition_sizes=[1] -query TT +query I???? select unnest(unnest(unnest(column3)['c1'])), unnest(unnest(column3)['c1']), column3,column1, column1['c0'] from recursive_unnest_table; ---- +1 [1, 2] [{c0: [1], c1: [[1, 2]]}] {c0: [1], c1: a} [1] +2 NULL [{c0: [1], c1: [[1, 2]]}] {c0: [1], c1: a} [1] +3 [3] [{c0: [2], c1: [[3], [4]]}] {c0: [2], c1: b} [2] +NULL [4] [{c0: [2], c1: [[3], [4]]}] {c0: [2], c1: b} [2] +4 [3] [{c0: [2], c1: [[3], [4]]}] {c0: [2], c1: b} [2] +NULL [4] [{c0: [2], c1: [[3], [4]]}] {c0: [2], c1: b} [2] -## query TT -## explain select unnest(unnest([[1,2,3]])) + unnest([4,5]); -## ---- -## logical_plan -## 01)Projection: unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))) + unnest(make_array(Int64(4),Int64(5))) -## 02)--Unnest: lists[unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))))] structs[] -## 03)----Projection: unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))), unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))) AS unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))), unnest(make_array(Int64(4),Int64(5))) -## 04)------Unnest: lists[unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))), unnest(make_array(Int64(4),Int64(5)))] structs[] -## 05)--------Projection: List([[1, 2, 3]]) AS unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))), List([4, 5]) AS unnest(make_array(Int64(4),Int64(5))) -## 06)----------EmptyRelation -## physical_plan -## 01)ProjectionExec: expr=[unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))))@1 + unnest(make_array(Int64(4),Int64(5)))@2 as unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))) + unnest(make_array(Int64(4),Int64(5)))] -## 02)--UnnestExec -## 03)----RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -## 04)------ProjectionExec: expr=[unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))@0 as unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))), unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))@0 as unnest(unnest(make_array(make_array(Int64(1),Int64(2),Int64(3))))), unnest(make_array(Int64(4),Int64(5)))@1 as unnest(make_array(Int64(4),Int64(5)))] -## 05)--------UnnestExec -## 06)----------ProjectionExec: expr=[[[1, 2, 3]] as unnest(make_array(make_array(Int64(1),Int64(2),Int64(3)))), [4, 5] as unnest(make_array(Int64(4),Int64(5)))] -## 07)------------PlaceholderRowExec +query I +select unnest(unnest([[1,2,3]])) + unnest([4,5]); +---- +5 +7 +NULL -## query I -## select unnest([4,5]) + 1; -## ---- -## 5 -## 6 -## -## ## FIXME: -## query II -## select unnest(unnest([[1,2,3]])) + unnest([4,5]), arrow_cast(unnest([4,5]),'Int64'); -## ---- -## 5 4 -## 6 4 -## 7 4 -## -## -## -## query I -## select unnest([2,1,9]) + unnest(unnest([[1,1,3]])) ; -## ---- -## 3 -## -## query I -## select unnest(unnest([[1,1]])) + unnest([2,1,9]),unnest([2,1,9]) + unnest(unnest([[1,1]])) ; -## ---- -## 3 +query I +select unnest(unnest([[1,2,3]])) + unnest(unnest([[1,2,3]])); +---- +2 +4 +6 \ No newline at end of file From fd01450694933a3bb49181afadbd5fb2e839c876 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Sat, 17 Aug 2024 10:59:40 +0200 Subject: [PATCH 22/56] fix unnesting into empty arrays --- datafusion/physical-plan/src/unnest.rs | 13 ++++-- datafusion/sql/src/utils.rs | 9 +++- datafusion/sqllogictest/test_files/unnest.slt | 46 ++++++++++--------- 3 files changed, 40 insertions(+), 28 deletions(-) diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index 36a06b8e21b6..5f2b95975521 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -410,7 +410,7 @@ fn unnest_at_level( temp_batch: &mut HashMap<(usize, usize), ArrayRef>, level_to_unnest: usize, options: &UnnestOptions, -) -> Result> { +) -> Result<(Vec, usize)> { // unnested columns at this depth level // now do some kind of projection-like // This query: @@ -475,7 +475,7 @@ fn unnest_at_level( })? as usize }; if total_length == 0 { - return Ok(vec![]); + return Ok((vec![], 0)); } // Unnest all the list arrays @@ -493,7 +493,7 @@ fn unnest_at_level( .for_each(|(flatten_arr, (index_in_input_schema, depth))| { temp_batch.insert((*index_in_input_schema, *depth), flatten_arr); }); - Ok(ret) + Ok((ret, total_length)) } /// For each row in a `RecordBatch`, some list/struct columns need to be unnested. @@ -524,13 +524,16 @@ fn build_batch( true => batch.columns(), false => &unnested_original_columns, }; - let temp = unnest_at_level( + let (temp, num_rows) = unnest_at_level( input, list_type_columns, &mut temp_batch, depth, options, )?; + if num_rows == 0 { + return Ok(RecordBatch::new_empty(Arc::clone(schema))); + } unnested_original_columns = temp; } let unnested_array_map: HashMap, usize)>> = @@ -987,7 +990,7 @@ mod tests { } #[test] - fn test_build_batch() -> datafusion_common::Result<()> { + fn test_build_batch_list_arr() -> datafusion_common::Result<()> { // col1 | col2 // [[1,2,3],null,[4,5]] | ['a','b'] // [[7,8,9,10], null, [11,12,13]] | ['c','d'] diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index 49da159fc433..c6349341e7bf 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -458,6 +458,7 @@ pub(crate) fn transform_bottom_unnest( Ok(Transformed::no(expr)) } }; + let mut transformed_root_exprs = None; let transform_up = |expr: Expr| -> Result> { // From the bottom up, we know the latest consecutive unnest sequence // we only do the transformation at the top unnest node @@ -503,7 +504,9 @@ pub(crate) fn transform_bottom_unnest( let mut transformed_exprs = transform(depth, arg, struct_allowed, inner_projection_exprs)?; - // TODO: if transformed_exprs has > 1 expr, handle it properly + if struct_allowed { + transformed_root_exprs = Some(transformed_exprs.clone()); + } return Ok(Transformed::new( transformed_exprs.swap_remove(0), true, @@ -559,11 +562,13 @@ pub(crate) fn transform_bottom_unnest( Ok(vec![Expr::Column(Column::from_name(column_name))]) } } else { + if let Some(transformed_root_exprs) = transformed_root_exprs { + return Ok(transformed_root_exprs); + } Ok(vec![transformed_expr]) } } -// write test for recursive_transform_unnest #[cfg(test)] mod tests { use std::{collections::HashMap, ops::Add, sync::Arc}; diff --git a/datafusion/sqllogictest/test_files/unnest.slt b/datafusion/sqllogictest/test_files/unnest.slt index 33acf8f3deca..29117833a770 100644 --- a/datafusion/sqllogictest/test_files/unnest.slt +++ b/datafusion/sqllogictest/test_files/unnest.slt @@ -60,10 +60,10 @@ select unnest([1,2,3]); 3 ## Basic unnest expression in select struct -## query III -## select unnest(struct(1,2,3)); -## ---- -## 1 2 3 +query III +select unnest(struct(1,2,3)); +---- +1 2 3 ## Basic unnest list expression in from clause query I @@ -180,17 +180,13 @@ select unnest([1,2,3]) + unnest([1,2,3]), unnest([1,2,3]) + unnest([4,5]); 4 7 6 NULL -##FIXME: -## Expectation: -## 2 5 4 -## 4 7 5 -## 6 NULL NULL + query III select unnest(unnest([[1,2,3]])) + unnest(unnest([[1,2,3]])), unnest(unnest([[1,2,3]])) + unnest([4,5]), unnest([4,5]); ---- 2 5 4 -4 6 4 -6 7 4 +4 7 5 +6 NULL NULL @@ -536,17 +532,25 @@ query ??II select unnest(column2), unnest(unnest(column2)), unnest(unnest(unnest(column2))), unnest(unnest(unnest(column2))) + 1 from recursive_unnest_table; ---- [[1], [2]] [1] 1 2 -[[1], [2]] [2] 2 3 -[[1, 1]] [1, 1] 1 2 -[[1, 1]] [1, 1] 1 2 +[[1, 1]] [2] NULL NULL +[[1], [2]] [1, 1] 2 3 +[[1, 1]] NULL NULL NULL +[[1], [2]] [1] 1 2 +[[1, 1]] [2] 1 2 +[[1], [2]] [1, 1] NULL NULL +[[1, 1]] NULL NULL NULL [[3, 4], [5]] [3, 4] 3 4 -[[3, 4], [5]] [3, 4] 4 5 -[[3, 4], [5]] [5] 5 6 -[[, 6], , [7, 8]] [, 6] NULL NULL -[[, 6], , [7, 8]] [, 6] 6 7 -[[, 6], , [7, 8]] [7, 8] 7 8 -[[, 6], , [7, 8]] [7, 8] 8 9 - +[[, 6], , [7, 8]] [5] 4 5 +[[3, 4], [5]] [, 6] 5 6 +[[, 6], , [7, 8]] NULL NULL NULL +NULL [7, 8] NULL NULL +[[3, 4], [5]] [3, 4] NULL NULL +[[, 6], , [7, 8]] [5] 6 7 +[[3, 4], [5]] [, 6] NULL NULL +[[, 6], , [7, 8]] NULL NULL NULL +NULL [7, 8] NULL NULL +[[3, 4], [5]] NULL 7 8 +[[, 6], , [7, 8]] NULL 8 9 ## the same unnest expr is referened multiple times (unnest is not the bottom-most expr) query ??II From 2b6e70fc94b2969c34677254957a5710113005d8 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Sat, 17 Aug 2024 14:02:42 +0200 Subject: [PATCH 23/56] some comment --- datafusion/expr/src/expr.rs | 2 +- datafusion/sql/src/utils.rs | 52 ++++++++++++++------ datafusion/sqllogictest/test_files/debug.slt | 8 +++ 3 files changed, 45 insertions(+), 17 deletions(-) diff --git a/datafusion/expr/src/expr.rs b/datafusion/expr/src/expr.rs index b4d489cc7c1e..0699f39f0341 100644 --- a/datafusion/expr/src/expr.rs +++ b/datafusion/expr/src/expr.rs @@ -2361,7 +2361,7 @@ impl fmt::Display for Expr { Expr::Placeholder(Placeholder { id, .. }) => write!(f, "{id}"), Expr::Unnest(Unnest { expr }) => { // TODO: use Display instead of Debug, there is non-unique expression name in projection issue. - write!(f, "UNNEST({expr:?})") + write!(f, "UNNEST({expr})") } } } diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index c6349341e7bf..279da7a4c1e4 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -338,13 +338,6 @@ pub(crate) fn transform_bottom_unnest( inner_projection_exprs: &mut Vec| -> Result> { let inner_expr_name = expr_in_unnest.schema_name().to_string(); - // let col = match expr_in_unnest { - // Expr::Column(col) => col, - // _ => { - // // TODO: this failed - // return internal_err!("unnesting on non-column expr is not supported"); - // } - // }; // Full context, we are trying to plan the execution as InnerProjection->Unnest->OuterProjection // inside unnest execution, each column inside the inner projection @@ -376,9 +369,9 @@ pub(crate) fn transform_bottom_unnest( .collect(), ); } - DataType::List(field) - | DataType::FixedSizeList(field, _) - | DataType::LargeList(field) => { + DataType::List(_) + | DataType::FixedSizeList(_, _) + | DataType::LargeList(_) => { // TODO: this memo only needs to be a hashset let (already_projected, transformed_cols) = match memo.get_mut(&inner_expr_name) { @@ -435,6 +428,8 @@ pub(crate) fn transform_bottom_unnest( let exprs_under_unnest = RefCell::new(HashSet::new()); let ancestor_unnest = RefCell::new(None); + // two conseqcutive unnest is something look like: unnest(unnest(some_expr)) + // the last items represent the inner most exprs let consecutive_unnest = RefCell::new(Vec::>::new()); // we need to mark only the latest unnest expr that was visitted during the down traversal let transform_down = |expr: Expr| -> Result> { @@ -461,7 +456,7 @@ pub(crate) fn transform_bottom_unnest( let mut transformed_root_exprs = None; let transform_up = |expr: Expr| -> Result> { // From the bottom up, we know the latest consecutive unnest sequence - // we only do the transformation at the top unnest node + // only do the transformation at the top unnest node // For example given this complex expr // - unnest(array_concat(unnest([[1,2,3]]),unnest([[4,5,6]]))) + unnest(unnest([[7,8,9])) // traversal will be like this: @@ -473,15 +468,22 @@ pub(crate) fn transform_bottom_unnest( // and the complex expr will be rewritten into: // unnest(array_concat(place_holder_col_1, place_holder_col_2)) + place_holder_col_3 if let Expr::Unnest(Unnest { .. }) = expr { - let mut down_unnest_mut = ancestor_unnest.borrow_mut(); - // upward traversal has reached the top unnest expr again + let mut ancestor_unnest_ref = ancestor_unnest.borrow_mut(); + // upward traversal has reached the top most unnest expr again // reset it to None - if *down_unnest_mut == Some(expr.clone()) { - down_unnest_mut.take(); + if *ancestor_unnest_ref == Some(expr.clone()) { + ancestor_unnest_ref.take(); } // find inside consecutive_unnest, the sequence of continous unnest exprs let mut found_first_unnest = false; let mut unnest_stack = vec![]; + + // TODO: this is still a hack and sub-optimal + // it's trying to get the latest consecutive unnest exprs + // and check if current upward traversal is the returning to the root expr + // (e.g) unnest(unnest(col)) then the traversal happens like: + // - down(unnest) -> down(unnest) -> down(col) -> up(col) -> up(unnest) -> up(unnest) + // the result of such traversal is unnest(col,depth:=2) for item in consecutive_unnest.borrow().iter().rev() { if let Some(expr) = item { found_first_unnest = true; @@ -497,11 +499,29 @@ pub(crate) fn transform_bottom_unnest( // this is the top most unnest expr inside the consecutive unnest exprs // e.g unnest(unnest(some_col)) if expr == *unnest_stack.last().unwrap() { + // TODO: detect if the top level is an unnest on struct + // and if the unnest stack contain > 1 exprs + // then do not transform the top level unnest yet, instead do the transform + // for the inner unnest first + // e.g: unnest(unnest(unnest([[struct()]]))) + // will needs to be transformed into + // unnest(unnest([[struct()]],depth=2)) let most_inner = unnest_stack.first().unwrap(); if let Expr::Unnest(Unnest { expr: ref arg }) = most_inner { + let (data_type, _) = arg.data_type_and_nullable(input.schema())?; + // unnest(unnest(struct_arr_col)) is not allow to be done recursively + // it needs to be splitted into multiple unnest logical plan + // unnest(struct_arr) + // unnest(struct_arr_col) as struct_arr + // instead of unnest(struct_arr_col, depth = 2) + match data_type { + DataType::Struct(_) => {} + _ => {} + } let depth = unnest_stack.len(); let struct_allowed = (&expr == original_expr) && depth == 1; + // TODO: arg should be inner most let mut transformed_exprs = transform(depth, arg, struct_allowed, inner_projection_exprs)?; if struct_allowed { @@ -525,7 +545,7 @@ pub(crate) fn transform_bottom_unnest( // For column exprs that are not descendants of any unnest node // retain their projection // e.g given expr tree unnest(col_a) + col_b, we have to retain projection of col_b - // down_unnest is non means current upward traversal is not descendant of any unnest + // this condition can be checked by maintaining an Option if matches!(&expr, Expr::Column(_)) && ancestor_unnest.borrow().is_none() { inner_projection_exprs.push(expr.clone()); } diff --git a/datafusion/sqllogictest/test_files/debug.slt b/datafusion/sqllogictest/test_files/debug.slt index 18f908e3e55b..453b0e7f31d1 100644 --- a/datafusion/sqllogictest/test_files/debug.slt +++ b/datafusion/sqllogictest/test_files/debug.slt @@ -1,3 +1,11 @@ +query I +select unnest(unnest(unnest([[[1]]]))) +---- + +query I +select unnest(unnest([struct(1,2,3)])) +---- + statement ok CREATE TABLE recursive_unnest_table AS VALUES From 2b49fa01b7b8522bfdd4c72cc5333b6e5a41d97a Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Sat, 17 Aug 2024 22:35:34 +0200 Subject: [PATCH 24/56] fix unnest struct --- datafusion/sql/src/utils.rs | 19 +++-- datafusion/sqllogictest/test_files/debug.slt | 3 + datafusion/sqllogictest/test_files/unnest.slt | 72 ++++--------------- 3 files changed, 29 insertions(+), 65 deletions(-) diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index 279da7a4c1e4..69754197edb7 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -437,8 +437,21 @@ pub(crate) fn transform_bottom_unnest( expr: ref inner_expr, }) = expr { + let (data_type, _) = inner_expr.data_type_and_nullable(input.schema())?; let mut consecutive_unnest_mut = consecutive_unnest.borrow_mut(); consecutive_unnest_mut.push(Some(expr.clone())); + // if expr inside unnest is a struct, do not consider + // the next unnest as consecutive unnest (if any) + // meaning unnest(unnest(struct_arr_col)) can't + // be interpreted as unest(struct_arr_col, depth:=2) + // but has to be split into multiple unnest logical plan instead + // a.k.a: + // - unnest(struct_col) + // unnest(struct_arr_col) as struct_col + + if let DataType::Struct(_) = data_type { + consecutive_unnest_mut.push(None); + } let mut maybe_ancestor = ancestor_unnest.borrow_mut(); if maybe_ancestor.is_none() { @@ -508,16 +521,12 @@ pub(crate) fn transform_bottom_unnest( // unnest(unnest([[struct()]],depth=2)) let most_inner = unnest_stack.first().unwrap(); if let Expr::Unnest(Unnest { expr: ref arg }) = most_inner { - let (data_type, _) = arg.data_type_and_nullable(input.schema())?; // unnest(unnest(struct_arr_col)) is not allow to be done recursively // it needs to be splitted into multiple unnest logical plan // unnest(struct_arr) // unnest(struct_arr_col) as struct_arr // instead of unnest(struct_arr_col, depth = 2) - match data_type { - DataType::Struct(_) => {} - _ => {} - } + let depth = unnest_stack.len(); let struct_allowed = (&expr == original_expr) && depth == 1; diff --git a/datafusion/sqllogictest/test_files/debug.slt b/datafusion/sqllogictest/test_files/debug.slt index 453b0e7f31d1..c8074c5e1bfe 100644 --- a/datafusion/sqllogictest/test_files/debug.slt +++ b/datafusion/sqllogictest/test_files/debug.slt @@ -1,3 +1,6 @@ + + + query I select unnest(unnest(unnest([[[1]]]))) ---- diff --git a/datafusion/sqllogictest/test_files/unnest.slt b/datafusion/sqllogictest/test_files/unnest.slt index 29117833a770..bc0308569fa8 100644 --- a/datafusion/sqllogictest/test_files/unnest.slt +++ b/datafusion/sqllogictest/test_files/unnest.slt @@ -523,7 +523,7 @@ select unnest(column1) from (select * from (values([1,2,3]), ([4,5,6])) limit 1 6 ## FIXME: https://github.com/apache/datafusion/issues/11198 -query error DataFusion error: Error during planning: Projections require unique expression names but the expression "UNNEST\(Column\(Column \{ relation: Some\(Bare \{ table: "unnest_table" \}\), name: "column1" \}\)\)" at position 0 and "UNNEST\(Column\(Column \{ relation: Some\(Bare \{ table: "unnest_table" \}\), name: "column1" \}\)\)" at position 1 have the same name. Consider aliasing \("AS"\) one of them. +query error DataFusion error: Error during planning: Projections require unique expression names but the expression "UNNEST\(unnest_table.column1\)" at position 0 and "UNNEST\(unnest_table.column1\)" at position 1 have the same name. Consider aliasing \("AS"\) one of them. select unnest(column1), unnest(column1) from unnest_table; @@ -560,8 +560,7 @@ select unnest(column3), unnest(column3)['c0'], unnest(unnest(column3)['c0']), un {c0: [2], c1: [[3], [4]]} [2] 2 4 -statement ok -drop table unnest_table; + ## unnest list followed by unnest struct query ??? @@ -602,71 +601,24 @@ query TT explain select unnest(unnest(unnest(column3)['c1'])), unnest(unnest(column3)['c1']), column3 from recursive_unnest_table; ---- logical_plan -01)Unnest: lists[UNNEST(UNNEST(UNNEST(recursive_unnest_table.column3)[c1]))] structs[] -02)--Projection: UNNEST(UNNEST(recursive_unnest_table.column3)[c1]) AS UNNEST(UNNEST(UNNEST(recursive_unnest_table.column3)[c1])), recursive_unnest_table.column3 -03)----Unnest: lists[UNNEST(UNNEST(recursive_unnest_table.column3)[c1])] structs[] -04)------Projection: get_field(UNNEST(recursive_unnest_table.column3), Utf8("c1")) AS UNNEST(UNNEST(recursive_unnest_table.column3)[c1]), recursive_unnest_table.column3 -05)--------Unnest: lists[UNNEST(recursive_unnest_table.column3)] structs[] -06)----------Projection: recursive_unnest_table.column3 AS UNNEST(recursive_unnest_table.column3), recursive_unnest_table.column3 -07)------------TableScan: recursive_unnest_table projection=[column3] +01)Unnest: lists[unnest_placeholder(unnest_placeholder(recursive_unnest_table.column3,depth=1)[c1])|depth=2, unnest_placeholder(unnest_placeholder(recursive_unnest_table.column3,depth=1)[c1])|depth=1] structs[] +02)--Projection: get_field(unnest_placeholder(recursive_unnest_table.column3,depth=1), Utf8("c1")) AS unnest_placeholder(unnest_placeholder(recursive_unnest_table.column3,depth=1)[c1]), recursive_unnest_table.column3 +03)----Unnest: lists[unnest_placeholder(recursive_unnest_table.column3)|depth=1] structs[] +04)------Projection: recursive_unnest_table.column3 AS unnest_placeholder(recursive_unnest_table.column3), recursive_unnest_table.column3 +05)--------TableScan: recursive_unnest_table projection=[column3] physical_plan 01)UnnestExec -02)--ProjectionExec: expr=[UNNEST(UNNEST(recursive_unnest_table.column3)[c1])@0 as UNNEST(UNNEST(UNNEST(recursive_unnest_table.column3)[c1])), column3@1 as column3] -03)----UnnestExec -04)------ProjectionExec: expr=[get_field(UNNEST(recursive_unnest_table.column3)@0, c1) as UNNEST(UNNEST(recursive_unnest_table.column3)[c1]), column3@1 as column3] -05)--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -06)----------UnnestExec -07)------------ProjectionExec: expr=[column3@0 as UNNEST(recursive_unnest_table.column3), column3@0 as column3] -08)--------------MemoryExec: partitions=1, partition_sizes=[1] - -statement ok -CREATE TABLE temp2 -AS VALUES - ([1,2,3],4), - (null,5), - ([1],6) -; +02)--ProjectionExec: expr=[get_field(unnest_placeholder(recursive_unnest_table.column3,depth=1)@0, c1) as unnest_placeholder(unnest_placeholder(recursive_unnest_table.column3,depth=1)[c1]), column3@1 as column3] +03)----RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +04)------UnnestExec +05)--------ProjectionExec: expr=[column3@0 as unnest_placeholder(recursive_unnest_table.column3), column3@0 as column3] +06)----------MemoryExec: partitions=1, partition_sizes=[1] -query II -select unnest(column1), column2 from temp2 ----- -1 4 -2 5 -3 NULL -statement ok -CREATE TABLE temp -AS VALUES - ([[1,2,3]],[4,5]) -; -query II -select unnest(unnest(column1)), unnest(column2) from temp ----- -1 4 -2 4 -3 4 -##[[1,3,4]] + [4,5] -##[1,2,3] 4 -query TT -explain select unnest(unnest(column1)), unnest(column2) from temp ----- -logical_plan -01)Unnest: lists[unnest(unnest(temp.column1))] structs[] -02)--Projection: unnest(temp.column1) AS unnest(unnest(temp.column1)), unnest(temp.column2) -03)----Unnest: lists[unnest(temp.column1), unnest(temp.column2)] structs[] -04)------Projection: temp.column1 AS unnest(temp.column1), temp.column2 AS unnest(temp.column2) -05)--------TableScan: temp projection=[column1, column2] -physical_plan -01)UnnestExec -02)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -03)----ProjectionExec: expr=[unnest(temp.column1)@0 as unnest(unnest(temp.column1)), unnest(temp.column2)@1 as unnest(temp.column2)] -04)------UnnestExec -05)--------ProjectionExec: expr=[column1@0 as unnest(temp.column1), column2@1 as unnest(temp.column2)] -06)----------MemoryExec: partitions=1, partition_sizes=[1] ## group by unnest ### without agg exprs From 222087e74da7fe6baf29d0f67206a63fd46f1598 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Tue, 20 Aug 2024 20:44:57 +0200 Subject: [PATCH 25/56] some note --- datafusion/expr/src/logical_plan/builder.rs | 4 +- datafusion/expr/src/logical_plan/tree_node.rs | 2 +- datafusion/sql/src/select.rs | 97 +++++++++++++++---- datafusion/sql/src/utils.rs | 45 ++++----- datafusion/sqllogictest/test_files/debug.slt | 69 +++---------- 5 files changed, 114 insertions(+), 103 deletions(-) diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index b2ff26b5cb25..2f06abdeb4a6 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -1576,9 +1576,7 @@ pub fn get_unnested_columns( let mut qualified_columns = Vec::with_capacity(1); match data_type { - DataType::List(field) - | DataType::FixedSizeList(field, _) - | DataType::LargeList(field) => { + DataType::List(_) | DataType::FixedSizeList(_, _) | DataType::LargeList(_) => { let data_type = get_unnested_list_column_recursive(data_type, depth)?; let new_field = Arc::new(Field::new( col_name.clone(), diff --git a/datafusion/expr/src/logical_plan/tree_node.rs b/datafusion/expr/src/logical_plan/tree_node.rs index 1045d688dd2f..409a9f4190cd 100644 --- a/datafusion/expr/src/logical_plan/tree_node.rs +++ b/datafusion/expr/src/logical_plan/tree_node.rs @@ -495,7 +495,7 @@ impl LogicalPlan { let exprs = columns .iter() - .map(|(c, unnest_type)| Expr::Column(c.clone())) + .map(|(c, _)| Expr::Column(c.clone())) .collect::>(); exprs.iter().apply_until_stop(f) } diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs index 9539f4ee0cd7..4547fcaa2c11 100644 --- a/datafusion/sql/src/select.rs +++ b/datafusion/sql/src/select.rs @@ -22,14 +22,16 @@ use crate::planner::{ idents_to_table_reference, ContextProvider, PlannerContext, SqlToRel, }; use crate::utils::{ - check_columns_satisfy_exprs, extract_aliases, rebase_expr, resolve_aliases_to_exprs, - resolve_columns, resolve_positions_to_exprs, transform_bottom_unnests, + check_columns_satisfy_exprs, extract_aliases, group_bottom_most_consecutive_unnests, + rebase_expr, resolve_aliases_to_exprs, resolve_columns, resolve_positions_to_exprs, }; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion_common::UnnestOptions; use datafusion_common::{not_impl_err, plan_err, DataFusionError, Result}; -use datafusion_expr::expr::{Alias, PlannedReplaceSelectItem, WildcardOptions}; +use datafusion_expr::expr::{ + schema_name_from_exprs, Alias, PlannedReplaceSelectItem, WildcardOptions, +}; use datafusion_expr::expr_rewriter::{ normalize_col, normalize_col_with_schemas_and_ambiguity_check, normalize_cols, }; @@ -160,6 +162,8 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { &combined_schema, planner_context, )?; + + println!("group by expr {:?} initial", group_by_expr); // aliases from the projection can conflict with same-named expressions in the input let mut alias_map = alias_map.clone(); for f in base_plan.schema().fields() { @@ -174,6 +178,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { base_plan.schema(), &[group_by_expr.clone()], )?; + println!("group by expr {:?} final", group_by_expr); Ok(group_by_expr) }) .collect::>>()? @@ -300,7 +305,14 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { select_exprs: Vec, ) -> Result { // Try process group by unnest - let input = self.try_process_aggregate_unnest(input)?; + let temp = input.clone(); + let (input, select_exprs) = + self.try_process_aggregate_unnest(input, select_exprs)?; + // TODO: some column was renamed after try_process_aggregate_unnest + // thus select_exprs also need to be transformed + if input == temp { + return Ok(input); + } let mut intermediate_plan = input; let mut intermediate_select_exprs = select_exprs; @@ -321,7 +333,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { // - unnest(struct_col) will be transformed into unnest(struct_col).field1, unnest(struct_col).field2 // - unnest(array_col) will be transformed into unnest(array_col).element // - unnest(array_col) + 1 will be transformed into unnest(array_col).element +1 - let outer_projection_exprs = transform_bottom_unnests( + let outer_projection_exprs = group_bottom_most_consecutive_unnests( &intermediate_plan, &mut unnest_columns, &mut inner_projection_exprs, @@ -374,32 +386,48 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { Ok(ret) } - fn try_process_aggregate_unnest(&self, input: LogicalPlan) -> Result { + fn try_process_aggregate_unnest( + &self, + input: LogicalPlan, + select_exprs: Vec, + ) -> Result<(LogicalPlan, Vec)> { match input { LogicalPlan::Aggregate(agg) => { let agg_expr = agg.aggr_expr.clone(); - let (new_input, new_group_by_exprs) = - self.try_process_group_by_unnest(agg)?; - LogicalPlanBuilder::from(new_input) - .aggregate(new_group_by_exprs, agg_expr)? - .build() + // we somehow need to un_rebase column exprs appeared in select/group by exprs + // e.g unnest(col("somecole")) expr may be rewritten as col("unnest(some_col)") + let (new_input, new_select_exprs, new_group_by_exprs) = + self.try_process_group_by_unnest(agg, select_exprs)?; + Ok(( + LogicalPlanBuilder::from(new_input) + .aggregate(new_group_by_exprs, agg_expr)? + .build()?, + new_select_exprs, + )) } LogicalPlan::Filter(mut filter) => { - filter.input = Arc::new( - self.try_process_aggregate_unnest(unwrap_arc(filter.input))?, - ); - Ok(LogicalPlan::Filter(filter)) + let (new_filter_input, new_select_expr) = self + .try_process_aggregate_unnest( + unwrap_arc(filter.input), + select_exprs, + )?; + filter.input = Arc::new(new_filter_input); + Ok((LogicalPlan::Filter(filter), new_select_expr)) } - _ => Ok(input), + _ => Ok((input, select_exprs)), } } /// 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 instance: + /// - select unnest(arr) as c1, unnest(arr) + unnest(arr) as c2 group by c1 + /// We need both for this funciton argument to check how select exprs has been transformed fn try_process_group_by_unnest( &self, agg: Aggregate, - ) -> Result<(LogicalPlan, Vec)> { + select_exprs: Vec, + ) -> Result<(LogicalPlan, Vec, Vec)> { let mut aggr_expr_using_columns: Option> = None; let Aggregate { @@ -426,20 +454,42 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { // TableScan: tab // ``` let mut intermediate_plan = unwrap_arc(input); - let mut intermediate_select_exprs = group_expr; + let mut intermediate_select_exprs = select_exprs; + let mut intermediate_group_by_exprs = group_expr; let mut memo = HashMap::new(); loop { let mut unnest_columns = vec![]; let mut inner_projection_exprs = vec![]; - let outer_projection_exprs = transform_bottom_unnests( + let outer_projection_exprs = group_bottom_most_consecutive_unnests( &intermediate_plan, &mut unnest_columns, &mut inner_projection_exprs, &mut memo, &intermediate_select_exprs, )?; + println!("inner_projection_exprs {:?}", inner_projection_exprs); + println!( + "inner_projection_exprs {}", + schema_name_from_exprs(&inner_projection_exprs)?, + ); + println!( + "select exprs {}", + schema_name_from_exprs(&intermediate_select_exprs)?, + ); + let temp_new_group_exprs = group_bottom_most_consecutive_unnests( + &intermediate_plan, + &mut unnest_columns, + &mut inner_projection_exprs, + &mut memo, + &intermediate_group_by_exprs, + )?; + println!("inner_projection_exprs {:?}", inner_projection_exprs); + println!( + "inner_projection_exprs {}", + schema_name_from_exprs(&inner_projection_exprs)?, + ); if unnest_columns.is_empty() { break; @@ -466,6 +516,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { } }; projection_exprs.extend(inner_projection_exprs); + println!("projection_exprs {:?}", projection_exprs); intermediate_plan = LogicalPlanBuilder::from(intermediate_plan) .project(projection_exprs)? @@ -473,10 +524,15 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { .build()?; intermediate_select_exprs = outer_projection_exprs; + intermediate_group_by_exprs = temp_new_group_exprs; } } - Ok((intermediate_plan, intermediate_select_exprs)) + Ok(( + intermediate_plan, + intermediate_select_exprs, + intermediate_group_by_exprs, + )) } fn plan_selection( @@ -760,6 +816,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { .iter() .map(|expr| resolve_columns(expr, input)) .collect::>>()?; + println!("aggr_projection_exprs {:?}", aggr_projection_exprs); // next we replace any expressions that are not a column with a column referencing // an output column from the aggregate schema diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index 69754197edb7..4727b987637d 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -291,7 +291,7 @@ pub(crate) fn value_to_string(value: &Value) -> Option { } } -pub(crate) fn transform_bottom_unnests( +pub(crate) fn group_bottom_most_consecutive_unnests( input: &LogicalPlan, unnest_placeholder_columns: &mut Vec<(Column, ColumnUnnestType)>, inner_projection_exprs: &mut Vec, @@ -301,7 +301,7 @@ pub(crate) fn transform_bottom_unnests( Ok(original_exprs .iter() .map(|expr| { - transform_bottom_unnest( + group_bottom_most_consecutive_unnest( input, unnest_placeholder_columns, inner_projection_exprs, @@ -315,7 +315,7 @@ pub(crate) fn transform_bottom_unnests( .collect::>()) } -/// Explain me +/// TODO: maybe renamed into transform_bottom_most_consecutive_unnests /// The context is we want to rewrite unnest() into InnerProjection->Unnest->OuterProjection /// Given an expression which contains unnest expr as one of its children, /// Try transform depends on unnest type @@ -325,7 +325,7 @@ pub(crate) fn transform_bottom_unnests( /// The transformed exprs will be used in the outer projection /// If along the path from root to bottom, there are multiple unnest expressions, the transformation /// is done only for the bottom expression -pub(crate) fn transform_bottom_unnest( +pub(crate) fn group_bottom_most_consecutive_unnest( input: &LogicalPlan, unnest_placeholder_columns: &mut Vec<(Column, ColumnUnnestType)>, inner_projection_exprs: &mut Vec, @@ -373,14 +373,13 @@ pub(crate) fn transform_bottom_unnest( | DataType::FixedSizeList(_, _) | DataType::LargeList(_) => { // TODO: this memo only needs to be a hashset - let (already_projected, transformed_cols) = - match memo.get_mut(&inner_expr_name) { - Some(vec) => (true, vec), - _ => { - memo.insert(inner_expr_name.clone(), vec![]); - (false, memo.get_mut(&inner_expr_name).unwrap()) - } - }; + let (already_projected, _) = match memo.get_mut(&inner_expr_name) { + Some(vec) => (true, vec), + _ => { + memo.insert(inner_expr_name.clone(), vec![]); + (false, memo.get_mut(&inner_expr_name).unwrap()) + } + }; if !already_projected { inner_projection_exprs .push(expr_in_unnest.clone().alias(placeholder_name.clone())); @@ -400,7 +399,7 @@ pub(crate) fn transform_bottom_unnest( }]), )); } - Some((col, unnesting)) => match unnesting { + Some((_, unnesting)) => match unnesting { ColumnUnnestType::List(list) => { let unnesting = ColumnUnnestList { output_column: post_unnest_column.clone(), @@ -611,7 +610,9 @@ mod tests { use datafusion_functions::core::expr_ext::FieldAccessor; use datafusion_functions_aggregate::expr_fn::count; - use crate::utils::{resolve_positions_to_exprs, transform_bottom_unnest}; + use crate::utils::{ + group_bottom_most_consecutive_unnest, resolve_positions_to_exprs, + }; fn column_unnests_eq(l: Vec<&str>, r: &[(Column, ColumnUnnestType)]) { let formatted: Vec = r.iter().map(|i| format!("{}|{}", i.0, i.1)).collect(); @@ -675,7 +676,7 @@ mod tests { .add(unnest(unnest(col("3d_col")))) .add(col("i64_col")); let mut memo = HashMap::new(); - let transformed_exprs = transform_bottom_unnest( + let transformed_exprs = group_bottom_most_consecutive_unnest( &input, &mut unnest_placeholder_columns, &mut inner_projection_exprs, @@ -709,7 +710,7 @@ mod tests { // unnest(3d_col) as 2d_col let original_expr_2 = unnest(col("3d_col")).alias("2d_col"); - let transformed_exprs = transform_bottom_unnest( + let transformed_exprs = group_bottom_most_consecutive_unnest( &input, &mut unnest_placeholder_columns, &mut inner_projection_exprs, @@ -743,7 +744,7 @@ mod tests { let original_expr_3 = unnest(unnest(unnest(col("struct_arr_col")).field("field1"))) .alias("fully_unnested_struct_arr"); - let transformed_exprs = transform_bottom_unnest( + let transformed_exprs = group_bottom_most_consecutive_unnest( &input, &mut unnest_placeholder_columns, &mut inner_projection_exprs, @@ -818,7 +819,7 @@ mod tests { .add(unnest(unnest(col("3d_col")))) .add(col("i64_col")); let mut memo = HashMap::new(); - let transformed_exprs = transform_bottom_unnest( + let transformed_exprs = group_bottom_most_consecutive_unnest( &input, &mut unnest_placeholder_columns, &mut inner_projection_exprs, @@ -852,7 +853,7 @@ mod tests { // unnest(3d_col) as 2d_col let original_expr_2 = unnest(col("3d_col")).alias("2d_col"); - let transformed_exprs = transform_bottom_unnest( + let transformed_exprs = group_bottom_most_consecutive_unnest( &input, &mut unnest_placeholder_columns, &mut inner_projection_exprs, @@ -921,7 +922,7 @@ mod tests { let mut memo = HashMap::new(); // unnest(struct_col) let original_expr = unnest(col("struct_col")); - let transformed_exprs = transform_bottom_unnest( + let transformed_exprs = group_bottom_most_consecutive_unnest( &input, &mut unnest_placeholder_columns, &mut inner_projection_exprs, @@ -949,7 +950,7 @@ mod tests { memo.clear(); // unnest(array_col) + 1 let original_expr = unnest(col("array_col")).add(lit(1i64)); - let transformed_exprs = transform_bottom_unnest( + let transformed_exprs = group_bottom_most_consecutive_unnest( &input, &mut unnest_placeholder_columns, &mut inner_projection_exprs, @@ -1015,7 +1016,7 @@ mod tests { // An expr with multiple unnest let original_expr = unnest(unnest(col("struct_col").field("matrix"))); - let transformed_exprs = transform_bottom_unnest( + let transformed_exprs = group_bottom_most_consecutive_unnest( &input, &mut unnest_placeholder_columns, &mut inner_projection_exprs, diff --git a/datafusion/sqllogictest/test_files/debug.slt b/datafusion/sqllogictest/test_files/debug.slt index c8074c5e1bfe..cf03e14f75a5 100644 --- a/datafusion/sqllogictest/test_files/debug.slt +++ b/datafusion/sqllogictest/test_files/debug.slt @@ -1,66 +1,21 @@ - - - -query I -select unnest(unnest(unnest([[[1]]]))) ----- - -query I -select unnest(unnest([struct(1,2,3)])) ----- - statement ok -CREATE TABLE recursive_unnest_table -AS VALUES - (struct([1], 'a'), [[[1],[2]],[[1,1]]], [struct([1],[[1,2]])]), - (struct([2], 'b'), [[[3,4],[5]],[[null,6],null,[7,8]]], [struct([2],[[3],[4]])]) -; +CREATE TABLE temp +AS VALUES + ([1,2,3],1,[9,10]), + ([1,2,3],2,[11,12]), + ([4,5,6],2,[13,14]) query TT -explain select unnest(unnest(unnest(column2))) from recursive_unnest_table; +explain select unnest(column1) c1 from temp group by unnest(column1); ---- -logical_plan -01)Unnest: lists[unnest_placeholder(recursive_unnest_table.column2)|depth=3] structs[] -02)--Projection: recursive_unnest_table.column2 AS unnest_placeholder(recursive_unnest_table.column2) -03)----TableScan: recursive_unnest_table projection=[column2] -physical_plan -01)UnnestExec -02)--ProjectionExec: expr=[column2@0 as unnest_placeholder(recursive_unnest_table.column2)] -03)----MemoryExec: partitions=1, partition_sizes=[1] query TT -explain select unnest(unnest(column2)) from recursive_unnest_table; +explain select unnest(column1) c1 from temp group by unnest(column1); ---- -logical_plan -01)Unnest: lists[unnest_placeholder(recursive_unnest_table.column2)|depth=2] structs[] -02)--Projection: recursive_unnest_table.column2 AS unnest_placeholder(recursive_unnest_table.column2) -03)----TableScan: recursive_unnest_table projection=[column2] -physical_plan -01)UnnestExec -02)--ProjectionExec: expr=[column2@0 as unnest_placeholder(recursive_unnest_table.column2)] -03)----MemoryExec: partitions=1, partition_sizes=[1] -query I???? -select unnest(unnest(unnest(column3)['c1'])), unnest(unnest(column3)['c1']), column3,column1, column1['c0'] from recursive_unnest_table; ----- -1 [1, 2] [{c0: [1], c1: [[1, 2]]}] {c0: [1], c1: a} [1] -2 NULL [{c0: [1], c1: [[1, 2]]}] {c0: [1], c1: a} [1] -3 [3] [{c0: [2], c1: [[3], [4]]}] {c0: [2], c1: b} [2] -NULL [4] [{c0: [2], c1: [[3], [4]]}] {c0: [2], c1: b} [2] -4 [3] [{c0: [2], c1: [[3], [4]]}] {c0: [2], c1: b} [2] -NULL [4] [{c0: [2], c1: [[3], [4]]}] {c0: [2], c1: b} [2] - -query I -select unnest(unnest([[1,2,3]])) + unnest([4,5]); ----- -5 -7 -NULL - - -query I -select unnest(unnest([[1,2,3]])) + unnest(unnest([[1,2,3]])); +query TT +explain select unnest(column1) c1 from temp group by c1; ---- -2 -4 -6 \ No newline at end of file +1 1 +1 2 +1 3 From f269e2fe25a1b156e803471052eff083dbf2b419 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Sun, 1 Sep 2024 12:45:24 +0200 Subject: [PATCH 26/56] chore: fix all test failure --- datafusion/sql/src/select.rs | 84 +++++-------------- datafusion/sql/src/utils.rs | 25 ++++-- datafusion/sqllogictest/test_files/debug.slt | 34 ++++++-- datafusion/sqllogictest/test_files/unnest.slt | 24 +++--- 4 files changed, 82 insertions(+), 85 deletions(-) diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs index 4547fcaa2c11..e15eb220bbf3 100644 --- a/datafusion/sql/src/select.rs +++ b/datafusion/sql/src/select.rs @@ -29,9 +29,7 @@ use crate::utils::{ use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion_common::UnnestOptions; use datafusion_common::{not_impl_err, plan_err, DataFusionError, Result}; -use datafusion_expr::expr::{ - schema_name_from_exprs, Alias, PlannedReplaceSelectItem, WildcardOptions, -}; +use datafusion_expr::expr::{Alias, PlannedReplaceSelectItem, WildcardOptions}; use datafusion_expr::expr_rewriter::{ normalize_col, normalize_col_with_schemas_and_ambiguity_check, normalize_cols, }; @@ -163,7 +161,6 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { planner_context, )?; - println!("group by expr {:?} initial", group_by_expr); // aliases from the projection can conflict with same-named expressions in the input let mut alias_map = alias_map.clone(); for f in base_plan.schema().fields() { @@ -178,7 +175,6 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { base_plan.schema(), &[group_by_expr.clone()], )?; - println!("group by expr {:?} final", group_by_expr); Ok(group_by_expr) }) .collect::>>()? @@ -306,13 +302,9 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { ) -> Result { // Try process group by unnest let temp = input.clone(); - let (input, select_exprs) = - self.try_process_aggregate_unnest(input, select_exprs)?; + let input = self.try_process_aggregate_unnest(input)?; // TODO: some column was renamed after try_process_aggregate_unnest // thus select_exprs also need to be transformed - if input == temp { - return Ok(input); - } let mut intermediate_plan = input; let mut intermediate_select_exprs = select_exprs; @@ -345,10 +337,10 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { if unnest_columns.is_empty() { // The original expr does not contain any unnest if i == 0 { - return Ok(intermediate_plan); - // return LogicalPlanBuilder::from(intermediate_plan) - // .project(inner_projection_exprs)? - // .build(); + // return Ok(intermediate_plan); + return LogicalPlanBuilder::from(intermediate_plan) + .project(intermediate_select_exprs)? + .build(); } break; } else { @@ -386,35 +378,25 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { Ok(ret) } - fn try_process_aggregate_unnest( - &self, - input: LogicalPlan, - select_exprs: Vec, - ) -> Result<(LogicalPlan, Vec)> { + fn try_process_aggregate_unnest(&self, input: LogicalPlan) -> Result { match input { LogicalPlan::Aggregate(agg) => { let agg_expr = agg.aggr_expr.clone(); // we somehow need to un_rebase column exprs appeared in select/group by exprs // e.g unnest(col("somecole")) expr may be rewritten as col("unnest(some_col)") - let (new_input, new_select_exprs, new_group_by_exprs) = - self.try_process_group_by_unnest(agg, select_exprs)?; - Ok(( - LogicalPlanBuilder::from(new_input) - .aggregate(new_group_by_exprs, agg_expr)? - .build()?, - new_select_exprs, - )) + let (new_input, new_group_by_exprs) = + self.try_process_group_by_unnest(agg)?; + Ok(LogicalPlanBuilder::from(new_input) + .aggregate(new_group_by_exprs, agg_expr)? + .build()?) } LogicalPlan::Filter(mut filter) => { - let (new_filter_input, new_select_expr) = self - .try_process_aggregate_unnest( - unwrap_arc(filter.input), - select_exprs, - )?; + let new_filter_input = + self.try_process_aggregate_unnest(unwrap_arc(filter.input))?; filter.input = Arc::new(new_filter_input); - Ok((LogicalPlan::Filter(filter), new_select_expr)) + Ok(LogicalPlan::Filter(filter)) } - _ => Ok((input, select_exprs)), + _ => Ok(input), } } @@ -426,8 +408,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { fn try_process_group_by_unnest( &self, agg: Aggregate, - select_exprs: Vec, - ) -> Result<(LogicalPlan, Vec, Vec)> { + ) -> Result<(LogicalPlan, Vec)> { let mut aggr_expr_using_columns: Option> = None; let Aggregate { @@ -454,8 +435,8 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { // TableScan: tab // ``` let mut intermediate_plan = unwrap_arc(input); - let mut intermediate_select_exprs = select_exprs; - let mut intermediate_group_by_exprs = group_expr; + let mut intermediate_select_exprs = group_expr; + // let mut intermediate_group_by_exprs = group_expr; let mut memo = HashMap::new(); loop { @@ -469,27 +450,6 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { &mut memo, &intermediate_select_exprs, )?; - println!("inner_projection_exprs {:?}", inner_projection_exprs); - println!( - "inner_projection_exprs {}", - schema_name_from_exprs(&inner_projection_exprs)?, - ); - println!( - "select exprs {}", - schema_name_from_exprs(&intermediate_select_exprs)?, - ); - let temp_new_group_exprs = group_bottom_most_consecutive_unnests( - &intermediate_plan, - &mut unnest_columns, - &mut inner_projection_exprs, - &mut memo, - &intermediate_group_by_exprs, - )?; - println!("inner_projection_exprs {:?}", inner_projection_exprs); - println!( - "inner_projection_exprs {}", - schema_name_from_exprs(&inner_projection_exprs)?, - ); if unnest_columns.is_empty() { break; @@ -516,7 +476,6 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { } }; projection_exprs.extend(inner_projection_exprs); - println!("projection_exprs {:?}", projection_exprs); intermediate_plan = LogicalPlanBuilder::from(intermediate_plan) .project(projection_exprs)? @@ -524,14 +483,14 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { .build()?; intermediate_select_exprs = outer_projection_exprs; - intermediate_group_by_exprs = temp_new_group_exprs; + // intermediate_group_by_exprs = temp_new_group_exprs; } } Ok(( intermediate_plan, intermediate_select_exprs, - intermediate_group_by_exprs, + // intermediate_group_by_exprs, )) } @@ -816,7 +775,6 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { .iter() .map(|expr| resolve_columns(expr, input)) .collect::>>()?; - println!("aggr_projection_exprs {:?}", aggr_projection_exprs); // next we replace any expressions that are not a column with a column referencing // an output column from the aggregate schema diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index 4727b987637d..6a62e4e009b6 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -33,7 +33,7 @@ use datafusion_common::{ use datafusion_expr::builder::get_struct_unnested_columns; use datafusion_expr::expr::{Alias, GroupingSet, Unnest, WindowFunction}; use datafusion_expr::utils::{expr_as_column_expr, find_column_exprs}; -use datafusion_expr::{expr_vec_fmt, Expr, ExprSchemable, LogicalPlan}; +use datafusion_expr::{col, expr_vec_fmt, Expr, ExprSchemable, LogicalPlan}; use datafusion_expr::{ColumnUnnestList, ColumnUnnestType}; use sqlparser::ast::Ident; use sqlparser::ast::Value; @@ -315,6 +315,14 @@ pub(crate) fn group_bottom_most_consecutive_unnests( .collect::>()) } +fn print_unnest(expr: &str, level: usize) -> String { + let mut result = String::from(expr); + for _ in 0..level { + result = format!("UNNEST({})", result); + } + format!("{}", result) +} + /// TODO: maybe renamed into transform_bottom_most_consecutive_unnests /// The context is we want to rewrite unnest() into InnerProjection->Unnest->OuterProjection /// Given an expression which contains unnest expr as one of its children, @@ -346,6 +354,9 @@ pub(crate) fn group_bottom_most_consecutive_unnest( let placeholder_name = format!("unnest_placeholder({})", inner_expr_name); let post_unnest_name = format!("unnest_placeholder({},depth={})", inner_expr_name, level); + // This is due to the fact that unnest transformation should keep the original + // column name as is, to comply with group by and order by + let post_unnest_alias = print_unnest(&inner_expr_name, level); let placeholder_column = Column::from_name(placeholder_name.clone()); let schema = input.schema(); @@ -385,16 +396,18 @@ pub(crate) fn group_bottom_most_consecutive_unnest( .push(expr_in_unnest.clone().alias(placeholder_name.clone())); } - let post_unnest_column = Column::from_name(post_unnest_name); + // let post_unnest_column = Column::from_name(post_unnest_name); + let post_unnest_expr = + col(post_unnest_name.clone()).alias(post_unnest_alias); match unnest_placeholder_columns .iter_mut() .find(|(inner_col, _)| inner_col == &placeholder_column) { None => { unnest_placeholder_columns.push(( - placeholder_column.clone(), + Column::from_name(placeholder_name.clone()), ColumnUnnestType::List(vec![ColumnUnnestList { - output_column: post_unnest_column.clone(), + output_column: Column::from_name(post_unnest_name), depth: level, }]), )); @@ -402,7 +415,7 @@ pub(crate) fn group_bottom_most_consecutive_unnest( Some((_, unnesting)) => match unnesting { ColumnUnnestType::List(list) => { let unnesting = ColumnUnnestList { - output_column: post_unnest_column.clone(), + output_column: Column::from_name(post_unnest_name), depth: level, }; if !list.contains(&unnesting) { @@ -414,7 +427,7 @@ pub(crate) fn group_bottom_most_consecutive_unnest( } }, } - return Ok(vec![Expr::Column(post_unnest_column)]); + return Ok(vec![post_unnest_expr]); } _ => { return internal_err!( diff --git a/datafusion/sqllogictest/test_files/debug.slt b/datafusion/sqllogictest/test_files/debug.slt index cf03e14f75a5..53089785b602 100644 --- a/datafusion/sqllogictest/test_files/debug.slt +++ b/datafusion/sqllogictest/test_files/debug.slt @@ -1,9 +1,33 @@ -statement ok -CREATE TABLE temp + +statement ok +CREATE TABLE unnest_table AS VALUES - ([1,2,3],1,[9,10]), - ([1,2,3],2,[11,12]), - ([4,5,6],2,[13,14]) + ([1,2,3], [7], 1, [13, 14], struct(1,2)) +; + +statement ok +CREATE TABLE recursive_unnest_table +AS VALUES + (struct([1], 'a'), [[[1],[2]],[[1,1]]], [struct([1],[[1,2]])]), + (struct([2], 'b'), [[[3,4],[5]],[[null,6],null,[7,8]]], [struct([2],[[3],[4]])]) +; + +query TT +explain select unnest(column2), unnest(unnest(column2)), unnest(unnest(unnest(column2))), unnest(unnest(unnest(column2))) + 1 from recursive_unnest_table; +---- + + +# query TT +# explain select column1 as c1 from unnest_table; +# ---- + + + + + +query TT +explain select unnest(column1) c1 from unnest_table group by c1 order by c1; +---- query TT explain select unnest(column1) c1 from temp group by unnest(column1); diff --git a/datafusion/sqllogictest/test_files/unnest.slt b/datafusion/sqllogictest/test_files/unnest.slt index bc0308569fa8..b4ec3fe100db 100644 --- a/datafusion/sqllogictest/test_files/unnest.slt +++ b/datafusion/sqllogictest/test_files/unnest.slt @@ -601,18 +601,20 @@ query TT explain select unnest(unnest(unnest(column3)['c1'])), unnest(unnest(column3)['c1']), column3 from recursive_unnest_table; ---- logical_plan -01)Unnest: lists[unnest_placeholder(unnest_placeholder(recursive_unnest_table.column3,depth=1)[c1])|depth=2, unnest_placeholder(unnest_placeholder(recursive_unnest_table.column3,depth=1)[c1])|depth=1] structs[] -02)--Projection: get_field(unnest_placeholder(recursive_unnest_table.column3,depth=1), Utf8("c1")) AS unnest_placeholder(unnest_placeholder(recursive_unnest_table.column3,depth=1)[c1]), recursive_unnest_table.column3 -03)----Unnest: lists[unnest_placeholder(recursive_unnest_table.column3)|depth=1] structs[] -04)------Projection: recursive_unnest_table.column3 AS unnest_placeholder(recursive_unnest_table.column3), recursive_unnest_table.column3 -05)--------TableScan: recursive_unnest_table projection=[column3] +01)Projection: unnest_placeholder(UNNEST(recursive_unnest_table.column3)[c1],depth=2) AS UNNEST(UNNEST(UNNEST(recursive_unnest_table.column3)[c1])), unnest_placeholder(UNNEST(recursive_unnest_table.column3)[c1],depth=1) AS UNNEST(UNNEST(recursive_unnest_table.column3)[c1]), recursive_unnest_table.column3 +02)--Unnest: lists[unnest_placeholder(UNNEST(recursive_unnest_table.column3)[c1])|depth=2, unnest_placeholder(UNNEST(recursive_unnest_table.column3)[c1])|depth=1] structs[] +03)----Projection: get_field(unnest_placeholder(recursive_unnest_table.column3,depth=1) AS UNNEST(recursive_unnest_table.column3), Utf8("c1")) AS unnest_placeholder(UNNEST(recursive_unnest_table.column3)[c1]), recursive_unnest_table.column3 +04)------Unnest: lists[unnest_placeholder(recursive_unnest_table.column3)|depth=1] structs[] +05)--------Projection: recursive_unnest_table.column3 AS unnest_placeholder(recursive_unnest_table.column3), recursive_unnest_table.column3 +06)----------TableScan: recursive_unnest_table projection=[column3] physical_plan -01)UnnestExec -02)--ProjectionExec: expr=[get_field(unnest_placeholder(recursive_unnest_table.column3,depth=1)@0, c1) as unnest_placeholder(unnest_placeholder(recursive_unnest_table.column3,depth=1)[c1]), column3@1 as column3] -03)----RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -04)------UnnestExec -05)--------ProjectionExec: expr=[column3@0 as unnest_placeholder(recursive_unnest_table.column3), column3@0 as column3] -06)----------MemoryExec: partitions=1, partition_sizes=[1] +01)ProjectionExec: expr=[unnest_placeholder(UNNEST(recursive_unnest_table.column3)[c1],depth=2)@0 as UNNEST(UNNEST(UNNEST(recursive_unnest_table.column3)[c1])), unnest_placeholder(UNNEST(recursive_unnest_table.column3)[c1],depth=1)@1 as UNNEST(UNNEST(recursive_unnest_table.column3)[c1]), column3@2 as column3] +02)--UnnestExec +03)----ProjectionExec: expr=[get_field(unnest_placeholder(recursive_unnest_table.column3,depth=1)@0, c1) as unnest_placeholder(UNNEST(recursive_unnest_table.column3)[c1]), column3@1 as column3] +04)------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +05)--------UnnestExec +06)----------ProjectionExec: expr=[column3@0 as unnest_placeholder(recursive_unnest_table.column3), column3@0 as column3] +07)------------MemoryExec: partitions=1, partition_sizes=[1] From 39fab44666fc72ba1fb996c3c0c882c315e52e8e Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Sun, 1 Sep 2024 20:59:14 +0200 Subject: [PATCH 27/56] fix projection pushdown --- datafusion/optimizer/src/push_down_filter.rs | 20 ++++++++-- datafusion/sqllogictest/test_files/debug.slt | 39 ------------------ .../test_files/push_down_filter.slt | 40 +++++++++---------- 3 files changed, 37 insertions(+), 62 deletions(-) diff --git a/datafusion/optimizer/src/push_down_filter.rs b/datafusion/optimizer/src/push_down_filter.rs index 9e75e9c2d4fd..450b88aad729 100644 --- a/datafusion/optimizer/src/push_down_filter.rs +++ b/datafusion/optimizer/src/push_down_filter.rs @@ -36,8 +36,8 @@ use datafusion_expr::utils::{ conjunction, expr_to_columns, split_conjunction, split_conjunction_owned, }; use datafusion_expr::{ - and, build_join_schema, or, BinaryExpr, Expr, Filter, LogicalPlanBuilder, Operator, - Projection, TableProviderFilterPushDown, + and, build_join_schema, or, BinaryExpr, ColumnUnnestType, Expr, Filter, + LogicalPlanBuilder, Operator, Projection, TableProviderFilterPushDown, }; use crate::optimizer::ApplyOrder; @@ -745,7 +745,21 @@ impl OptimizerRule for PushDownFilter { let mut accum: HashSet = HashSet::new(); expr_to_columns(&predicate, &mut accum)?; - if unnest.exec_columns.iter().any(|(c, _)| accum.contains(c)) { + if unnest.exec_columns.iter().any(|(c, unnest_detail)| { + match unnest_detail { + ColumnUnnestType::List(vec) => { + return vec + .iter() + .any(|c| accum.contains(&c.output_column)); + } + ColumnUnnestType::Struct => { + return false; + } + ColumnUnnestType::Inferred => { + return accum.contains(c); + } + } + }) { unnest_predicates.push(predicate); } else { non_unnest_predicates.push(predicate); diff --git a/datafusion/sqllogictest/test_files/debug.slt b/datafusion/sqllogictest/test_files/debug.slt index 53089785b602..951c9132cb82 100644 --- a/datafusion/sqllogictest/test_files/debug.slt +++ b/datafusion/sqllogictest/test_files/debug.slt @@ -4,42 +4,3 @@ CREATE TABLE unnest_table AS VALUES ([1,2,3], [7], 1, [13, 14], struct(1,2)) ; - -statement ok -CREATE TABLE recursive_unnest_table -AS VALUES - (struct([1], 'a'), [[[1],[2]],[[1,1]]], [struct([1],[[1,2]])]), - (struct([2], 'b'), [[[3,4],[5]],[[null,6],null,[7,8]]], [struct([2],[[3],[4]])]) -; - -query TT -explain select unnest(column2), unnest(unnest(column2)), unnest(unnest(unnest(column2))), unnest(unnest(unnest(column2))) + 1 from recursive_unnest_table; ----- - - -# query TT -# explain select column1 as c1 from unnest_table; -# ---- - - - - - -query TT -explain select unnest(column1) c1 from unnest_table group by c1 order by c1; ----- - -query TT -explain select unnest(column1) c1 from temp group by unnest(column1); ----- - -query TT -explain select unnest(column1) c1 from temp group by unnest(column1); ----- - -query TT -explain select unnest(column1) c1 from temp group by c1; ----- -1 1 -1 2 -1 3 diff --git a/datafusion/sqllogictest/test_files/push_down_filter.slt b/datafusion/sqllogictest/test_files/push_down_filter.slt index 2d74c1fc6994..86aa07b04ce1 100644 --- a/datafusion/sqllogictest/test_files/push_down_filter.slt +++ b/datafusion/sqllogictest/test_files/push_down_filter.slt @@ -36,9 +36,9 @@ query TT explain select uc2 from (select unnest(column2) as uc2, column1 from v) where column1 = 2; ---- logical_plan -01)Projection: UNNEST(v.column2) AS uc2 -02)--Unnest: lists[UNNEST(v.column2)] structs[] -03)----Projection: v.column2 AS UNNEST(v.column2), v.column1 +01)Projection: unnest_placeholder(v.column2,depth=1) AS uc2 +02)--Unnest: lists[unnest_placeholder(v.column2)|depth=1] structs[] +03)----Projection: v.column2 AS unnest_placeholder(v.column2), v.column1 04)------Filter: v.column1 = Int64(2) 05)--------TableScan: v projection=[column1, column2] @@ -53,11 +53,11 @@ query TT explain select uc2 from (select unnest(column2) as uc2, column1 from v) where uc2 > 3; ---- logical_plan -01)Projection: UNNEST(v.column2) AS uc2 -02)--Filter: UNNEST(v.column2) > Int64(3) -03)----Projection: UNNEST(v.column2) -04)------Unnest: lists[UNNEST(v.column2)] structs[] -05)--------Projection: v.column2 AS UNNEST(v.column2), v.column1 +01)Projection: unnest_placeholder(v.column2,depth=1) AS uc2 +02)--Filter: unnest_placeholder(v.column2,depth=1) > Int64(3) +03)----Projection: unnest_placeholder(v.column2,depth=1) +04)------Unnest: lists[unnest_placeholder(v.column2)|depth=1] structs[] +05)--------Projection: v.column2 AS unnest_placeholder(v.column2), v.column1 06)----------TableScan: v projection=[column1, column2] query II @@ -71,10 +71,10 @@ query TT explain select uc2, column1 from (select unnest(column2) as uc2, column1 from v) where uc2 > 3 AND column1 = 2; ---- logical_plan -01)Projection: UNNEST(v.column2) AS uc2, v.column1 -02)--Filter: UNNEST(v.column2) > Int64(3) -03)----Unnest: lists[UNNEST(v.column2)] structs[] -04)------Projection: v.column2 AS UNNEST(v.column2), v.column1 +01)Projection: unnest_placeholder(v.column2,depth=1) AS uc2, v.column1 +02)--Filter: unnest_placeholder(v.column2,depth=1) > Int64(3) +03)----Unnest: lists[unnest_placeholder(v.column2)|depth=1] structs[] +04)------Projection: v.column2 AS unnest_placeholder(v.column2), v.column1 05)--------Filter: v.column1 = Int64(2) 06)----------TableScan: v projection=[column1, column2] @@ -90,10 +90,10 @@ query TT explain select uc2, column1 from (select unnest(column2) as uc2, column1 from v) where uc2 > 3 OR column1 = 2; ---- logical_plan -01)Projection: UNNEST(v.column2) AS uc2, v.column1 -02)--Filter: UNNEST(v.column2) > Int64(3) OR v.column1 = Int64(2) -03)----Unnest: lists[UNNEST(v.column2)] structs[] -04)------Projection: v.column2 AS UNNEST(v.column2), v.column1 +01)Projection: unnest_placeholder(v.column2,depth=1) AS uc2, v.column1 +02)--Filter: unnest_placeholder(v.column2,depth=1) > Int64(3) OR v.column1 = Int64(2) +03)----Unnest: lists[unnest_placeholder(v.column2)|depth=1] structs[] +04)------Projection: v.column2 AS unnest_placeholder(v.column2), v.column1 05)--------TableScan: v projection=[column1, column2] statement ok @@ -112,10 +112,10 @@ query TT explain select * from (select column1, unnest(column2) as o from d) where o['a'] = 1; ---- logical_plan -01)Projection: d.column1, UNNEST(d.column2) AS o -02)--Filter: get_field(UNNEST(d.column2), Utf8("a")) = Int64(1) -03)----Unnest: lists[UNNEST(d.column2)] structs[] -04)------Projection: d.column1, d.column2 AS UNNEST(d.column2) +01)Projection: d.column1, unnest_placeholder(d.column2,depth=1) AS o +02)--Filter: get_field(unnest_placeholder(d.column2,depth=1), Utf8("a")) = Int64(1) +03)----Unnest: lists[unnest_placeholder(d.column2)|depth=1] structs[] +04)------Projection: d.column1, d.column2 AS unnest_placeholder(d.column2) 05)--------TableScan: d projection=[column1, column2] From 444d741ab6d16474bf85eb56d9a065f7c76a4a1b Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Sun, 1 Sep 2024 22:17:04 +0200 Subject: [PATCH 28/56] custom rewriter for recursive unnest --- datafusion/optimizer/src/push_down_filter.rs | 17 +- datafusion/sql/src/select.rs | 5 +- datafusion/sql/src/utils.rs | 185 ++++++++++--------- 3 files changed, 104 insertions(+), 103 deletions(-) diff --git a/datafusion/optimizer/src/push_down_filter.rs b/datafusion/optimizer/src/push_down_filter.rs index 450b88aad729..7f16dcb6d9bf 100644 --- a/datafusion/optimizer/src/push_down_filter.rs +++ b/datafusion/optimizer/src/push_down_filter.rs @@ -745,19 +745,14 @@ impl OptimizerRule for PushDownFilter { let mut accum: HashSet = HashSet::new(); expr_to_columns(&predicate, &mut accum)?; - if unnest.exec_columns.iter().any(|(c, unnest_detail)| { - match unnest_detail { + if unnest.exec_columns.iter().any(|(unnest_col, unnest_type)| { + match unnest_type { ColumnUnnestType::List(vec) => { - return vec - .iter() - .any(|c| accum.contains(&c.output_column)); - } - ColumnUnnestType::Struct => { - return false; - } - ColumnUnnestType::Inferred => { - return accum.contains(c); + vec.iter().any(|c| accum.contains(&c.output_column)) } + ColumnUnnestType::Struct => false, + // for inferred unnest, output column will be the same with input column + ColumnUnnestType::Inferred => accum.contains(unnest_col), } }) { unnest_predicates.push(predicate); diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs index e15eb220bbf3..343d59a091d5 100644 --- a/datafusion/sql/src/select.rs +++ b/datafusion/sql/src/select.rs @@ -301,10 +301,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { select_exprs: Vec, ) -> Result { // Try process group by unnest - let temp = input.clone(); let input = self.try_process_aggregate_unnest(input)?; - // TODO: some column was renamed after try_process_aggregate_unnest - // thus select_exprs also need to be transformed let mut intermediate_plan = input; let mut intermediate_select_exprs = select_exprs; @@ -344,7 +341,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { } break; } else { - let columns = unnest_columns.into_iter().map(|col| col.into()).collect(); + let columns = unnest_columns.into_iter().map(|col| col).collect(); // Set preserve_nulls to false to ensure compatibility with DuckDB and PostgreSQL let unnest_options = UnnestOptions::new().with_preserve_nulls(false); let mut check_list: HashSet = inner_projection_exprs diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index 6a62e4e009b6..d8d933460905 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -17,18 +17,18 @@ //! SQL Utility Functions -use std::cell::RefCell; -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use std::vec; use arrow_schema::{ DataType, DECIMAL128_MAX_PRECISION, DECIMAL256_MAX_PRECISION, DECIMAL_DEFAULT_SCALE, }; use datafusion_common::tree_node::{ - Transformed, TransformedResult, TreeNode, TreeNodeRecursion, + Transformed, TransformedResult, TreeNode, TreeNodeRecursion, TreeNodeRewriter, }; use datafusion_common::{ - exec_err, internal_err, plan_err, Column, DataFusionError, Result, ScalarValue, + exec_err, internal_err, plan_err, Column, DFSchemaRef, DataFusionError, Result, + ScalarValue, }; use datafusion_expr::builder::get_struct_unnested_columns; use datafusion_expr::expr::{Alias, GroupingSet, Unnest, WindowFunction}; @@ -320,31 +320,31 @@ fn print_unnest(expr: &str, level: usize) -> String { for _ in 0..level { result = format!("UNNEST({})", result); } - format!("{}", result) + result } -/// TODO: maybe renamed into transform_bottom_most_consecutive_unnests -/// The context is we want to rewrite unnest() into InnerProjection->Unnest->OuterProjection -/// Given an expression which contains unnest expr as one of its children, -/// Try transform depends on unnest type -/// - For list column: unnest(col) with type list -> unnest(col) with type list::item -/// - For struct column: unnest(struct(field1, field2)) -> unnest(struct).field1, unnest(struct).field2 -/// -/// The transformed exprs will be used in the outer projection -/// If along the path from root to bottom, there are multiple unnest expressions, the transformation -/// is done only for the bottom expression -pub(crate) fn group_bottom_most_consecutive_unnest( - input: &LogicalPlan, - unnest_placeholder_columns: &mut Vec<(Column, ColumnUnnestType)>, - inner_projection_exprs: &mut Vec, - memo: &mut HashMap>, - original_expr: &Expr, -) -> Result> { - let mut transform = |level: usize, - expr_in_unnest: &Expr, - struct_allowed: bool, - inner_projection_exprs: &mut Vec| - -> Result> { +/* +This is only usedful when used with transform down up +A full example of how the transformation works: + */ +struct RecursiveUnnestRewriter<'a> { + memo: &'a mut HashMap>, + input_schema: &'a DFSchemaRef, + original_expr: &'a Expr, + latest_visited_unnest: Option, + ancestor_unnest: Option, + consecutive_unnest: Vec>, + inner_projection_exprs: &'a mut Vec, + unnest_placeholder_columns: &'a mut Vec<(Column, ColumnUnnestType)>, + transformed_root_exprs: Option>, +} +impl<'a> RecursiveUnnestRewriter<'a> { + fn transform( + &mut self, + level: usize, + expr_in_unnest: &Expr, + struct_allowed: bool, + ) -> Result> { let inner_expr_name = expr_in_unnest.schema_name().to_string(); // Full context, we are trying to plan the execution as InnerProjection->Unnest->OuterProjection @@ -358,18 +358,18 @@ pub(crate) fn group_bottom_most_consecutive_unnest( // column name as is, to comply with group by and order by let post_unnest_alias = print_unnest(&inner_expr_name, level); let placeholder_column = Column::from_name(placeholder_name.clone()); - let schema = input.schema(); + // let schema = input.schema(); - let (data_type, _) = expr_in_unnest.data_type_and_nullable(schema)?; + let (data_type, _) = expr_in_unnest.data_type_and_nullable(self.input_schema)?; match data_type { DataType::Struct(inner_fields) => { if !struct_allowed { return internal_err!("unnest on struct can only be applied at the root level of select expression"); } - inner_projection_exprs + self.inner_projection_exprs .push(expr_in_unnest.clone().alias(placeholder_name.clone())); - unnest_placeholder_columns.push(( + self.unnest_placeholder_columns.push(( Column::from_name(placeholder_name.clone()), ColumnUnnestType::Struct, )); @@ -384,27 +384,28 @@ pub(crate) fn group_bottom_most_consecutive_unnest( | DataType::FixedSizeList(_, _) | DataType::LargeList(_) => { // TODO: this memo only needs to be a hashset - let (already_projected, _) = match memo.get_mut(&inner_expr_name) { + let (already_projected, _) = match self.memo.get_mut(&inner_expr_name) { Some(vec) => (true, vec), _ => { - memo.insert(inner_expr_name.clone(), vec![]); - (false, memo.get_mut(&inner_expr_name).unwrap()) + self.memo.insert(inner_expr_name.clone(), vec![]); + (false, self.memo.get_mut(&inner_expr_name).unwrap()) } }; if !already_projected { - inner_projection_exprs + self.inner_projection_exprs .push(expr_in_unnest.clone().alias(placeholder_name.clone())); } // let post_unnest_column = Column::from_name(post_unnest_name); let post_unnest_expr = col(post_unnest_name.clone()).alias(post_unnest_alias); - match unnest_placeholder_columns + match self + .unnest_placeholder_columns .iter_mut() .find(|(inner_col, _)| inner_col == &placeholder_column) { None => { - unnest_placeholder_columns.push(( + self.unnest_placeholder_columns.push(( Column::from_name(placeholder_name.clone()), ColumnUnnestType::List(vec![ColumnUnnestList { output_column: Column::from_name(post_unnest_name), @@ -435,23 +436,19 @@ pub(crate) fn group_bottom_most_consecutive_unnest( ); } } - }; - let latest_visited_unnest = RefCell::new(None); - let exprs_under_unnest = RefCell::new(HashSet::new()); - let ancestor_unnest = RefCell::new(None); - - // two conseqcutive unnest is something look like: unnest(unnest(some_expr)) - // the last items represent the inner most exprs - let consecutive_unnest = RefCell::new(Vec::>::new()); - // we need to mark only the latest unnest expr that was visitted during the down traversal - let transform_down = |expr: Expr| -> Result> { + } +} + +impl<'a> TreeNodeRewriter for RecursiveUnnestRewriter<'a> { + type Node = Expr; + + fn f_down(&mut self, expr: Expr) -> Result> { if let Expr::Unnest(Unnest { expr: ref inner_expr, }) = expr { - let (data_type, _) = inner_expr.data_type_and_nullable(input.schema())?; - let mut consecutive_unnest_mut = consecutive_unnest.borrow_mut(); - consecutive_unnest_mut.push(Some(expr.clone())); + let (data_type, _) = inner_expr.data_type_and_nullable(self.input_schema)?; + self.consecutive_unnest.push(Some(expr.clone())); // if expr inside unnest is a struct, do not consider // the next unnest as consecutive unnest (if any) // meaning unnest(unnest(struct_arr_col)) can't @@ -462,42 +459,28 @@ pub(crate) fn group_bottom_most_consecutive_unnest( // unnest(struct_arr_col) as struct_col if let DataType::Struct(_) = data_type { - consecutive_unnest_mut.push(None); + self.consecutive_unnest.push(None); } - let mut maybe_ancestor = ancestor_unnest.borrow_mut(); - if maybe_ancestor.is_none() { - *maybe_ancestor = Some(expr.clone()); + if self.ancestor_unnest.is_none() { + self.ancestor_unnest = Some(expr.clone()); } - exprs_under_unnest.borrow_mut().insert(inner_expr.clone()); - *latest_visited_unnest.borrow_mut() = Some(expr.clone()); + self.latest_visited_unnest = Some(expr.clone()); Ok(Transformed::no(expr)) } else { - consecutive_unnest.borrow_mut().push(None); + self.consecutive_unnest.push(None); Ok(Transformed::no(expr)) } - }; - let mut transformed_root_exprs = None; - let transform_up = |expr: Expr| -> Result> { - // From the bottom up, we know the latest consecutive unnest sequence - // only do the transformation at the top unnest node - // For example given this complex expr - // - unnest(array_concat(unnest([[1,2,3]]),unnest([[4,5,6]]))) + unnest(unnest([[7,8,9])) - // traversal will be like this: - // down[binary_add] - // ->down[unnest(...)]->down[array_concat]->down/up[unnest([[1,2,3]])]->down/up[unnest([[4,5,6]])] - // ->up[array_concat]->up[unnest(...)]->down[unnest(unnest(...))]->down[unnest([[7,8,9]])] - // ->up[unnest([[7,8,9]])]->up[unnest(unnest(...))]->up[binary_add] - // the transformation only happens for unnest([[1,2,3]]), unnest([[4,5,6]]) and unnest(unnest([[7,8,9]])) - // and the complex expr will be rewritten into: - // unnest(array_concat(place_holder_col_1, place_holder_col_2)) + place_holder_col_3 + } + + fn f_up(&mut self, expr: Expr) -> Result> { if let Expr::Unnest(Unnest { .. }) = expr { - let mut ancestor_unnest_ref = ancestor_unnest.borrow_mut(); + // let mut ancestor_unnest_ref = ancestor_unnest.borrow_mut(); // upward traversal has reached the top most unnest expr again // reset it to None - if *ancestor_unnest_ref == Some(expr.clone()) { - ancestor_unnest_ref.take(); + if self.ancestor_unnest == Some(expr.clone()) { + self.ancestor_unnest.take(); } // find inside consecutive_unnest, the sequence of continous unnest exprs let mut found_first_unnest = false; @@ -509,7 +492,7 @@ pub(crate) fn group_bottom_most_consecutive_unnest( // (e.g) unnest(unnest(col)) then the traversal happens like: // - down(unnest) -> down(unnest) -> down(col) -> up(col) -> up(unnest) -> up(unnest) // the result of such traversal is unnest(col,depth:=2) - for item in consecutive_unnest.borrow().iter().rev() { + for item in self.consecutive_unnest.iter().rev() { if let Some(expr) = item { found_first_unnest = true; unnest_stack.push(expr.clone()); @@ -540,13 +523,13 @@ pub(crate) fn group_bottom_most_consecutive_unnest( // instead of unnest(struct_arr_col, depth = 2) let depth = unnest_stack.len(); - let struct_allowed = (&expr == original_expr) && depth == 1; + let struct_allowed = (&expr == self.original_expr) && depth == 1; // TODO: arg should be inner most let mut transformed_exprs = - transform(depth, arg, struct_allowed, inner_projection_exprs)?; + self.transform(depth, arg, struct_allowed)?; if struct_allowed { - transformed_root_exprs = Some(transformed_exprs.clone()); + self.transformed_root_exprs = Some(transformed_exprs.clone()); } return Ok(Transformed::new( transformed_exprs.swap_remove(0), @@ -556,22 +539,50 @@ pub(crate) fn group_bottom_most_consecutive_unnest( } else { return internal_err!("not reached"); } - - // } } } else { - consecutive_unnest.borrow_mut().push(None); + self.consecutive_unnest.push(None); } // For column exprs that are not descendants of any unnest node // retain their projection // e.g given expr tree unnest(col_a) + col_b, we have to retain projection of col_b // this condition can be checked by maintaining an Option - if matches!(&expr, Expr::Column(_)) && ancestor_unnest.borrow().is_none() { - inner_projection_exprs.push(expr.clone()); + if matches!(&expr, Expr::Column(_)) && self.ancestor_unnest.is_none() { + self.inner_projection_exprs.push(expr.clone()); } Ok(Transformed::no(expr)) + } +} + +/// TODO: maybe renamed into transform_bottom_most_consecutive_unnests +/// The context is we want to rewrite unnest() into InnerProjection->Unnest->OuterProjection +/// Given an expression which contains unnest expr as one of its children, +/// Try transform depends on unnest type +/// - For list column: unnest(col) with type list -> unnest(col) with type list::item +/// - For struct column: unnest(struct(field1, field2)) -> unnest(struct).field1, unnest(struct).field2 +/// +/// The transformed exprs will be used in the outer projection +/// If along the path from root to bottom, there are multiple unnest expressions, the transformation +/// is done only for the bottom expression +pub(crate) fn group_bottom_most_consecutive_unnest( + input: &LogicalPlan, + unnest_placeholder_columns: &mut Vec<(Column, ColumnUnnestType)>, + inner_projection_exprs: &mut Vec, + memo: &mut HashMap>, + original_expr: &Expr, +) -> Result> { + let mut rewriter = RecursiveUnnestRewriter { + memo, + input_schema: input.schema(), + original_expr, + latest_visited_unnest: None, + ancestor_unnest: None, + consecutive_unnest: vec![], + inner_projection_exprs, + unnest_placeholder_columns, + transformed_root_exprs: None, }; // This transformation is only done for list unnest @@ -587,9 +598,7 @@ pub(crate) fn group_bottom_most_consecutive_unnest( data: transformed_expr, transformed, tnr: _, - } = original_expr - .clone() - .transform_down_up(transform_down, transform_up)?; + } = original_expr.clone().rewrite(&mut rewriter)?; if !transformed { if matches!(&transformed_expr, Expr::Column(_)) { @@ -603,7 +612,7 @@ pub(crate) fn group_bottom_most_consecutive_unnest( Ok(vec![Expr::Column(Column::from_name(column_name))]) } } else { - if let Some(transformed_root_exprs) = transformed_root_exprs { + if let Some(transformed_root_exprs) = rewriter.transformed_root_exprs { return Ok(transformed_root_exprs); } Ok(vec![transformed_expr]) From bba148a2427a1bb7512870cde81797378628fef7 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Mon, 2 Sep 2024 22:18:57 +0200 Subject: [PATCH 29/56] simplify --- datafusion/sql/src/select.rs | 39 ++++----- datafusion/sql/src/utils.rs | 154 +++++++++++++---------------------- 2 files changed, 75 insertions(+), 118 deletions(-) diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs index 343d59a091d5..b9b332b6bacb 100644 --- a/datafusion/sql/src/select.rs +++ b/datafusion/sql/src/select.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use std::collections::{HashMap, HashSet}; +use std::collections::HashSet; use std::sync::Arc; use crate::planner::{ @@ -310,8 +310,6 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { // The transformation happen bottom up, one at a time for each iteration // Only exaust the loop if no more unnest transformation is found for i in 0.. { - // impl memoization to store all previous unnest transformation - let mut memo = HashMap::new(); let mut unnest_columns = vec![]; // from which column used for projection, before the unnest happen // including non unnest column and unnest column @@ -326,7 +324,6 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { &intermediate_plan, &mut unnest_columns, &mut inner_projection_exprs, - &mut memo, &intermediate_select_exprs, )?; @@ -344,23 +341,24 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { let columns = unnest_columns.into_iter().map(|col| col).collect(); // Set preserve_nulls to false to ensure compatibility with DuckDB and PostgreSQL let unnest_options = UnnestOptions::new().with_preserve_nulls(false); - let mut check_list: HashSet = inner_projection_exprs - .iter() - .map(|expr| expr.clone()) - .collect(); - let deduplicated: Vec = inner_projection_exprs - .into_iter() - .filter(|expr| -> bool { - if check_list.remove(expr) { - true - } else { - false - } - }) - .collect(); + // let mut check_list: HashSet = inner_projection_exprs + // .iter() + // .map(|expr| expr.clone()) + // .collect(); + // let deduplicated: Vec = inner_projection_exprs + // .into_iter() + // .filter(|expr| -> bool { + // return true; + // if check_list.remove(expr) { + // true + // } else { + // false + // } + // }) + // .collect(); let plan = LogicalPlanBuilder::from(intermediate_plan) - .project(deduplicated.clone())? + .project(inner_projection_exprs)? .unnest_columns_with_options(columns, unnest_options)? .build()?; intermediate_plan = plan; @@ -433,8 +431,6 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { // ``` let mut intermediate_plan = unwrap_arc(input); let mut intermediate_select_exprs = group_expr; - // let mut intermediate_group_by_exprs = group_expr; - let mut memo = HashMap::new(); loop { let mut unnest_columns = vec![]; @@ -444,7 +440,6 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { &intermediate_plan, &mut unnest_columns, &mut inner_projection_exprs, - &mut memo, &intermediate_select_exprs, )?; diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index d8d933460905..8550e1adf94a 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -295,7 +295,6 @@ pub(crate) fn group_bottom_most_consecutive_unnests( input: &LogicalPlan, unnest_placeholder_columns: &mut Vec<(Column, ColumnUnnestType)>, inner_projection_exprs: &mut Vec, - memo: &mut HashMap>, original_exprs: &[Expr], ) -> Result> { Ok(original_exprs @@ -305,7 +304,6 @@ pub(crate) fn group_bottom_most_consecutive_unnests( input, unnest_placeholder_columns, inner_projection_exprs, - memo, expr, ) }) @@ -328,14 +326,13 @@ This is only usedful when used with transform down up A full example of how the transformation works: */ struct RecursiveUnnestRewriter<'a> { - memo: &'a mut HashMap>, input_schema: &'a DFSchemaRef, original_expr: &'a Expr, latest_visited_unnest: Option, ancestor_unnest: Option, - consecutive_unnest: Vec>, + consecutive_unnest: Vec>, inner_projection_exprs: &'a mut Vec, - unnest_placeholder_columns: &'a mut Vec<(Column, ColumnUnnestType)>, + columns_unnestings: &'a mut Vec<(Column, ColumnUnnestType)>, transformed_root_exprs: Option>, } impl<'a> RecursiveUnnestRewriter<'a> { @@ -358,7 +355,6 @@ impl<'a> RecursiveUnnestRewriter<'a> { // column name as is, to comply with group by and order by let post_unnest_alias = print_unnest(&inner_expr_name, level); let placeholder_column = Column::from_name(placeholder_name.clone()); - // let schema = input.schema(); let (data_type, _) = expr_in_unnest.data_type_and_nullable(self.input_schema)?; @@ -369,43 +365,36 @@ impl<'a> RecursiveUnnestRewriter<'a> { } self.inner_projection_exprs .push(expr_in_unnest.clone().alias(placeholder_name.clone())); - self.unnest_placeholder_columns.push(( + self.columns_unnestings.push(( Column::from_name(placeholder_name.clone()), ColumnUnnestType::Struct, )); - return Ok( + Ok( get_struct_unnested_columns(&placeholder_name, &inner_fields) .into_iter() - .map(|c| Expr::Column(c)) + .map(Expr::Column) .collect(), - ); + ) } DataType::List(_) | DataType::FixedSizeList(_, _) | DataType::LargeList(_) => { - // TODO: this memo only needs to be a hashset - let (already_projected, _) = match self.memo.get_mut(&inner_expr_name) { - Some(vec) => (true, vec), - _ => { - self.memo.insert(inner_expr_name.clone(), vec![]); - (false, self.memo.get_mut(&inner_expr_name).unwrap()) - } - }; - if !already_projected { - self.inner_projection_exprs - .push(expr_in_unnest.clone().alias(placeholder_name.clone())); - } + push_projection_dedupl( + self.inner_projection_exprs, + expr_in_unnest.clone().alias(placeholder_name.clone()), + ); // let post_unnest_column = Column::from_name(post_unnest_name); let post_unnest_expr = col(post_unnest_name.clone()).alias(post_unnest_alias); match self - .unnest_placeholder_columns + .columns_unnestings .iter_mut() .find(|(inner_col, _)| inner_col == &placeholder_column) { + // there is not unnesting done on this column yet None => { - self.unnest_placeholder_columns.push(( + self.columns_unnestings.push(( Column::from_name(placeholder_name.clone()), ColumnUnnestType::List(vec![ColumnUnnestList { output_column: Column::from_name(post_unnest_name), @@ -413,6 +402,8 @@ impl<'a> RecursiveUnnestRewriter<'a> { }]), )); } + // some unnesting(at some level) has been done on this column + // e.g select unnest(column3), unnest(unnest(column3)) Some((_, unnesting)) => match unnesting { ColumnUnnestType::List(list) => { let unnesting = ColumnUnnestList { @@ -424,7 +415,7 @@ impl<'a> RecursiveUnnestRewriter<'a> { } } _ => { - return internal_err!("expr_in_unnest is a list type, while previous unnesting on this column is not a list type"); + return internal_err!(""); } }, } @@ -443,12 +434,10 @@ impl<'a> TreeNodeRewriter for RecursiveUnnestRewriter<'a> { type Node = Expr; fn f_down(&mut self, expr: Expr) -> Result> { - if let Expr::Unnest(Unnest { - expr: ref inner_expr, - }) = expr - { - let (data_type, _) = inner_expr.data_type_and_nullable(self.input_schema)?; - self.consecutive_unnest.push(Some(expr.clone())); + if let Expr::Unnest(ref unnest_expr) = expr { + let (data_type, _) = + unnest_expr.expr.data_type_and_nullable(self.input_schema)?; + self.consecutive_unnest.push(Some(unnest_expr.clone())); // if expr inside unnest is a struct, do not consider // the next unnest as consecutive unnest (if any) // meaning unnest(unnest(struct_arr_col)) can't @@ -475,8 +464,7 @@ impl<'a> TreeNodeRewriter for RecursiveUnnestRewriter<'a> { } fn f_up(&mut self, expr: Expr) -> Result> { - if let Expr::Unnest(Unnest { .. }) = expr { - // let mut ancestor_unnest_ref = ancestor_unnest.borrow_mut(); + if let Expr::Unnest(ref traversing_unnest) = expr { // upward traversal has reached the top most unnest expr again // reset it to None if self.ancestor_unnest == Some(expr.clone()) { @@ -504,9 +492,11 @@ impl<'a> TreeNodeRewriter for RecursiveUnnestRewriter<'a> { } } - // this is the top most unnest expr inside the consecutive unnest exprs - // e.g unnest(unnest(some_col)) - if expr == *unnest_stack.last().unwrap() { + // There exist a sequence of unnest exprs + // and this traversal has reached the top most + // e.g Unnest(top) -> Unnest(2nd) -> Column(bottom) + // -> Unnest(2nd) -> Unnest(top) a.k.a here + if traversing_unnest == unnest_stack.last().unwrap() { // TODO: detect if the top level is an unnest on struct // and if the unnest stack contain > 1 exprs // then do not transform the top level unnest yet, instead do the transform @@ -515,30 +505,26 @@ impl<'a> TreeNodeRewriter for RecursiveUnnestRewriter<'a> { // will needs to be transformed into // unnest(unnest([[struct()]],depth=2)) let most_inner = unnest_stack.first().unwrap(); - if let Expr::Unnest(Unnest { expr: ref arg }) = most_inner { - // unnest(unnest(struct_arr_col)) is not allow to be done recursively - // it needs to be splitted into multiple unnest logical plan - // unnest(struct_arr) - // unnest(struct_arr_col) as struct_arr - // instead of unnest(struct_arr_col, depth = 2) - - let depth = unnest_stack.len(); - let struct_allowed = (&expr == self.original_expr) && depth == 1; - - // TODO: arg should be inner most - let mut transformed_exprs = - self.transform(depth, arg, struct_allowed)?; - if struct_allowed { - self.transformed_root_exprs = Some(transformed_exprs.clone()); - } - return Ok(Transformed::new( - transformed_exprs.swap_remove(0), - true, - TreeNodeRecursion::Continue, - )); - } else { - return internal_err!("not reached"); + let arg = most_inner.expr.as_ref(); + // unnest(unnest(struct_arr_col)) is not allow to be done recursively + // it needs to be splitted into multiple unnest logical plan + // unnest(struct_arr) + // unnest(struct_arr_col) as struct_arr + // instead of unnest(struct_arr_col, depth = 2) + + let depth = unnest_stack.len(); + let struct_allowed = (&expr == self.original_expr) && depth == 1; + + // TODO: arg should be inner most + let mut transformed_exprs = self.transform(depth, arg, struct_allowed)?; + if struct_allowed { + self.transformed_root_exprs = Some(transformed_exprs.clone()); } + return Ok(Transformed::new( + transformed_exprs.swap_remove(0), + true, + TreeNodeRecursion::Continue, + )); } } else { self.consecutive_unnest.push(None); @@ -549,13 +535,22 @@ impl<'a> TreeNodeRewriter for RecursiveUnnestRewriter<'a> { // e.g given expr tree unnest(col_a) + col_b, we have to retain projection of col_b // this condition can be checked by maintaining an Option if matches!(&expr, Expr::Column(_)) && self.ancestor_unnest.is_none() { - self.inner_projection_exprs.push(expr.clone()); + push_projection_dedupl(self.inner_projection_exprs, expr.clone()); } Ok(Transformed::no(expr)) } } +fn push_projection_dedupl(projection: &mut Vec, expr: Expr) { + let schema_name = expr.schema_name().to_string(); + if !projection + .iter() + .any(|e| e.schema_name().to_string() == schema_name) + { + projection.push(expr); + } +} /// TODO: maybe renamed into transform_bottom_most_consecutive_unnests /// The context is we want to rewrite unnest() into InnerProjection->Unnest->OuterProjection /// Given an expression which contains unnest expr as one of its children, @@ -570,18 +565,16 @@ pub(crate) fn group_bottom_most_consecutive_unnest( input: &LogicalPlan, unnest_placeholder_columns: &mut Vec<(Column, ColumnUnnestType)>, inner_projection_exprs: &mut Vec, - memo: &mut HashMap>, original_expr: &Expr, ) -> Result> { let mut rewriter = RecursiveUnnestRewriter { - memo, input_schema: input.schema(), original_expr, latest_visited_unnest: None, ancestor_unnest: None, consecutive_unnest: vec![], inner_projection_exprs, - unnest_placeholder_columns, + columns_unnestings: unnest_placeholder_columns, transformed_root_exprs: None, }; @@ -602,13 +595,13 @@ pub(crate) fn group_bottom_most_consecutive_unnest( if !transformed { if matches!(&transformed_expr, Expr::Column(_)) { - inner_projection_exprs.push(transformed_expr.clone()); + push_projection_dedupl(inner_projection_exprs, transformed_expr.clone()); Ok(vec![transformed_expr]) } else { // We need to evaluate the expr in the inner projection, // outer projection just select its name let column_name = transformed_expr.schema_name().to_string(); - inner_projection_exprs.push(transformed_expr); + push_projection_dedupl(inner_projection_exprs, transformed_expr); Ok(vec![Expr::Column(Column::from_name(column_name))]) } } else { @@ -697,12 +690,10 @@ mod tests { let original_expr = unnest(unnest(col("3d_col"))) .add(unnest(unnest(col("3d_col")))) .add(col("i64_col")); - let mut memo = HashMap::new(); let transformed_exprs = group_bottom_most_consecutive_unnest( &input, &mut unnest_placeholder_columns, &mut inner_projection_exprs, - &mut memo, &original_expr, )?; // only the bottom most unnest exprs are transformed @@ -712,9 +703,6 @@ mod tests { .add(col("unnest_placeholder(3d_col,depth=2)")) .add(col("i64_col"))] ); - // memoization only contains 1 transformation - assert_eq!(memo.len(), 1); - assert!(memo.get("3d_col").is_some()); column_unnests_eq( vec!["unnest_placeholder(3d_col)"], &unnest_placeholder_columns, @@ -736,7 +724,6 @@ mod tests { &input, &mut unnest_placeholder_columns, &mut inner_projection_exprs, - &mut memo, &original_expr_2, )?; @@ -744,10 +731,6 @@ mod tests { transformed_exprs, vec![col("unnest_placeholder(3d_col,depth=1)").alias("2d_col")] ); - // memoization still contains 1 transformation - // and the previous transformation is reused - assert_eq!(memo.len(), 1); - assert!(memo.get("3d_col").is_some()); column_unnests_eq( vec!["unnest_placeholder(3d_col)"], &mut unnest_placeholder_columns, @@ -770,7 +753,6 @@ mod tests { &input, &mut unnest_placeholder_columns, &mut inner_projection_exprs, - &mut memo, &original_expr_3, )?; @@ -781,11 +763,7 @@ mod tests { )) .alias("fully_unnested_struct_arr")] ); - // memoization still contains 1 transformation - // and the previous transformation is reused - assert_eq!(memo.len(), 2); - assert!(memo.get("struct_arr_col").is_some()); column_unnests_eq( vec![ "unnest_placeholder(3d_col)", @@ -840,12 +818,10 @@ mod tests { let original_expr = unnest(unnest(col("3d_col"))) .add(unnest(unnest(col("3d_col")))) .add(col("i64_col")); - let mut memo = HashMap::new(); let transformed_exprs = group_bottom_most_consecutive_unnest( &input, &mut unnest_placeholder_columns, &mut inner_projection_exprs, - &mut memo, &original_expr, )?; // only the bottom most unnest exprs are transformed @@ -855,9 +831,6 @@ mod tests { .add(col("unnest_placeholder(3d_col,depth=2)")) .add(col("i64_col"))] ); - // memoization only contains 1 transformation - assert_eq!(memo.len(), 1); - assert!(memo.get("3d_col").is_some()); column_unnests_eq( vec!["unnest_placeholder(3d_col)|List([unnest_placeholder(3d_col,depth=2)|depth=2])"], &unnest_placeholder_columns, @@ -879,7 +852,6 @@ mod tests { &input, &mut unnest_placeholder_columns, &mut inner_projection_exprs, - &mut memo, &original_expr_2, )?; @@ -887,10 +859,6 @@ mod tests { transformed_exprs, vec![col("unnest_placeholder(3d_col,depth=1)").alias("2d_col")] ); - // memoization still contains 1 transformation - // and the for the same column, depth = 1 needs to be performed aside from depth = 2 - assert_eq!(memo.len(), 1); - assert!(memo.get("3d_col").is_some()); column_unnests_eq( vec!["unnest_placeholder(3d_col)|List([unnest_placeholder(3d_col,depth=2)|depth=2, unnest_placeholder(3d_col,depth=1)|depth=1])"], &unnest_placeholder_columns, @@ -941,14 +909,12 @@ mod tests { let mut unnest_placeholder_columns = vec![]; let mut inner_projection_exprs = vec![]; - let mut memo = HashMap::new(); // unnest(struct_col) let original_expr = unnest(col("struct_col")); let transformed_exprs = group_bottom_most_consecutive_unnest( &input, &mut unnest_placeholder_columns, &mut inner_projection_exprs, - &mut memo, &original_expr, )?; assert_eq!( @@ -969,14 +935,12 @@ mod tests { vec![col("struct_col").alias("UNNEST(struct_col)"),] ); - memo.clear(); // unnest(array_col) + 1 let original_expr = unnest(col("array_col")).add(lit(1i64)); let transformed_exprs = group_bottom_most_consecutive_unnest( &input, &mut unnest_placeholder_columns, &mut inner_projection_exprs, - &mut memo, &original_expr, )?; column_unnests_eq( @@ -1034,7 +998,6 @@ mod tests { let mut unnest_placeholder_columns = vec![]; let mut inner_projection_exprs = vec![]; - memo.clear(); // An expr with multiple unnest let original_expr = unnest(unnest(col("struct_col").field("matrix"))); @@ -1042,7 +1005,6 @@ mod tests { &input, &mut unnest_placeholder_columns, &mut inner_projection_exprs, - &mut memo, &original_expr, )?; // Only the inner most/ bottom most unnest is transformed From c10812d9e8cd18250be3dd90830680051ca41423 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Tue, 3 Sep 2024 22:31:48 +0200 Subject: [PATCH 30/56] rm unnecessary projection --- datafusion/sql/src/utils.rs | 101 ++++++++---------- datafusion/sqllogictest/test_files/unnest.slt | 15 ++- 2 files changed, 58 insertions(+), 58 deletions(-) diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index 8550e1adf94a..1b72380127a0 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -327,15 +327,30 @@ A full example of how the transformation works: */ struct RecursiveUnnestRewriter<'a> { input_schema: &'a DFSchemaRef, - original_expr: &'a Expr, - latest_visited_unnest: Option, - ancestor_unnest: Option, + root_expr: &'a Expr, + // useful to detect which child expr is a part of/ not a part of unnest operation + top_most_unnest: Option, consecutive_unnest: Vec>, inner_projection_exprs: &'a mut Vec, columns_unnestings: &'a mut Vec<(Column, ColumnUnnestType)>, transformed_root_exprs: Option>, } impl<'a> RecursiveUnnestRewriter<'a> { + // given a sequence of [None,Unnest,Unnest,None,None] + // returns [Unnest,Unnest] + // The first items is the inner most unnest + fn get_latest_consecutive_unnest(&self) -> Vec { + self.consecutive_unnest + .iter() + .rev() + .skip_while(|item| item.is_none()) + .take_while(|item| item.is_some()) + .to_owned() + .cloned() + .map(|item| item.unwrap()) + .collect() + } + fn transform( &mut self, level: usize, @@ -363,8 +378,10 @@ impl<'a> RecursiveUnnestRewriter<'a> { if !struct_allowed { return internal_err!("unnest on struct can only be applied at the root level of select expression"); } - self.inner_projection_exprs - .push(expr_in_unnest.clone().alias(placeholder_name.clone())); + push_projection_dedupl( + self.inner_projection_exprs, + expr_in_unnest.clone().alias(placeholder_name.clone()), + ); self.columns_unnestings.push(( Column::from_name(placeholder_name.clone()), ColumnUnnestType::Struct, @@ -415,16 +432,14 @@ impl<'a> RecursiveUnnestRewriter<'a> { } } _ => { - return internal_err!(""); + return internal_err!("not reached"); } }, } - return Ok(vec![post_unnest_expr]); + Ok(vec![post_unnest_expr]) } _ => { - return internal_err!( - "unnest on non-list or struct type is not supported" - ); + internal_err!("unnest on non-list or struct type is not supported") } } } @@ -450,12 +465,10 @@ impl<'a> TreeNodeRewriter for RecursiveUnnestRewriter<'a> { if let DataType::Struct(_) = data_type { self.consecutive_unnest.push(None); } - - if self.ancestor_unnest.is_none() { - self.ancestor_unnest = Some(expr.clone()); + if self.top_most_unnest.is_none() { + self.top_most_unnest = Some(unnest_expr.clone()); } - self.latest_visited_unnest = Some(expr.clone()); Ok(Transformed::no(expr)) } else { self.consecutive_unnest.push(None); @@ -465,58 +478,37 @@ impl<'a> TreeNodeRewriter for RecursiveUnnestRewriter<'a> { fn f_up(&mut self, expr: Expr) -> Result> { if let Expr::Unnest(ref traversing_unnest) = expr { - // upward traversal has reached the top most unnest expr again - // reset it to None - if self.ancestor_unnest == Some(expr.clone()) { - self.ancestor_unnest.take(); + if traversing_unnest == self.top_most_unnest.as_ref().unwrap() { + self.top_most_unnest = None; } // find inside consecutive_unnest, the sequence of continous unnest exprs - let mut found_first_unnest = false; - let mut unnest_stack = vec![]; - // TODO: this is still a hack and sub-optimal - // it's trying to get the latest consecutive unnest exprs + // Get the latest consecutive unnest exprs // and check if current upward traversal is the returning to the root expr - // (e.g) unnest(unnest(col)) then the traversal happens like: - // - down(unnest) -> down(unnest) -> down(col) -> up(col) -> up(unnest) -> up(unnest) - // the result of such traversal is unnest(col,depth:=2) - for item in self.consecutive_unnest.iter().rev() { - if let Some(expr) = item { - found_first_unnest = true; - unnest_stack.push(expr.clone()); - } else { - if !found_first_unnest { - continue; - } - break; - } - } + // for example given a expr `unnest(unnest(col))` then the traversal happens like: + // down(unnest) -> down(unnest) -> down(col) -> up(col) -> up(unnest) -> up(unnest) + // the result of such traversal is unnest(col, depth:=2) + let unnest_stack = self.get_latest_consecutive_unnest(); - // There exist a sequence of unnest exprs - // and this traversal has reached the top most + // This traversal has reached the top most unnest again // e.g Unnest(top) -> Unnest(2nd) -> Column(bottom) // -> Unnest(2nd) -> Unnest(top) a.k.a here + // Thus + // Unnest(Unnest(some_col)) is rewritten into Unnest(some_col, depth:=2) if traversing_unnest == unnest_stack.last().unwrap() { - // TODO: detect if the top level is an unnest on struct - // and if the unnest stack contain > 1 exprs - // then do not transform the top level unnest yet, instead do the transform - // for the inner unnest first - // e.g: unnest(unnest(unnest([[struct()]]))) - // will needs to be transformed into - // unnest(unnest([[struct()]],depth=2)) let most_inner = unnest_stack.first().unwrap(); - let arg = most_inner.expr.as_ref(); + let inner_expr = most_inner.expr.as_ref(); // unnest(unnest(struct_arr_col)) is not allow to be done recursively // it needs to be splitted into multiple unnest logical plan // unnest(struct_arr) // unnest(struct_arr_col) as struct_arr // instead of unnest(struct_arr_col, depth = 2) - let depth = unnest_stack.len(); - let struct_allowed = (&expr == self.original_expr) && depth == 1; + let unnest_recursion = unnest_stack.len(); + let struct_allowed = (&expr == self.root_expr) && unnest_recursion == 1; - // TODO: arg should be inner most - let mut transformed_exprs = self.transform(depth, arg, struct_allowed)?; + let mut transformed_exprs = + self.transform(unnest_recursion, inner_expr, struct_allowed)?; if struct_allowed { self.transformed_root_exprs = Some(transformed_exprs.clone()); } @@ -534,7 +526,7 @@ impl<'a> TreeNodeRewriter for RecursiveUnnestRewriter<'a> { // retain their projection // e.g given expr tree unnest(col_a) + col_b, we have to retain projection of col_b // this condition can be checked by maintaining an Option - if matches!(&expr, Expr::Column(_)) && self.ancestor_unnest.is_none() { + if matches!(&expr, Expr::Column(_)) && self.top_most_unnest.is_none() { push_projection_dedupl(self.inner_projection_exprs, expr.clone()); } @@ -569,9 +561,8 @@ pub(crate) fn group_bottom_most_consecutive_unnest( ) -> Result> { let mut rewriter = RecursiveUnnestRewriter { input_schema: input.schema(), - original_expr, - latest_visited_unnest: None, - ancestor_unnest: None, + root_expr: original_expr, + top_most_unnest: None, consecutive_unnest: vec![], inner_projection_exprs, columns_unnestings: unnest_placeholder_columns, @@ -614,7 +605,7 @@ pub(crate) fn group_bottom_most_consecutive_unnest( #[cfg(test)] mod tests { - use std::{collections::HashMap, ops::Add, sync::Arc}; + use std::{ops::Add, sync::Arc}; use arrow::datatypes::{DataType as ArrowDataType, Field, Schema}; use arrow_schema::Fields; diff --git a/datafusion/sqllogictest/test_files/unnest.slt b/datafusion/sqllogictest/test_files/unnest.slt index b4ec3fe100db..34231f9f92db 100644 --- a/datafusion/sqllogictest/test_files/unnest.slt +++ b/datafusion/sqllogictest/test_files/unnest.slt @@ -595,8 +595,18 @@ NULL [[[3, 4], [5]], [[, 6], , [7, 8]]] 8 [[[3, 4], [5]], [[, 6], , [7, 8]]] - - +query I?? +select unnest(unnest(unnest(column3)['c1'])), unnest(unnest(column3)['c1']), column3 from recursive_unnest_table; +---- +1 [1, 2] [{c0: [1], c1: [[1, 2]]}] +2 NULL [{c0: [1], c1: [[1, 2]]}] +3 [3] [{c0: [2], c1: [[3], [4]]}] +NULL [4] [{c0: [2], c1: [[3], [4]]}] +4 [3] [{c0: [2], c1: [[3], [4]]}] +NULL [4] [{c0: [2], c1: [[3], [4]]}] + +## demonstrate where multiple unnest plan is needed +## e.g unnest -> field_access -> unnest query TT explain select unnest(unnest(unnest(column3)['c1'])), unnest(unnest(column3)['c1']), column3 from recursive_unnest_table; ---- @@ -620,7 +630,6 @@ physical_plan - ## group by unnest ### without agg exprs From 577805879a4b225d61338706615cc30b715ca1be Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Thu, 5 Sep 2024 21:47:02 +0200 Subject: [PATCH 31/56] chore: better comments --- datafusion/physical-plan/src/unnest.rs | 236 +++++++++++-------------- 1 file changed, 100 insertions(+), 136 deletions(-) diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index 5f2b95975521..c345ae9666f8 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -341,127 +341,103 @@ fn flatten_struct_cols( Ok(RecordBatch::try_new(Arc::clone(schema), columns_expanded)?) } -#[derive(Debug, Clone)] +#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)] pub struct ListUnnest { pub index_in_input_schema: usize, pub depth: usize, } -// a map of original index (track which column in the batch to be interested in) -// input schema: -// col1_unnest_placeholder: list[list[int]], col1: list[list[int]], col2 list[int] -// with unnest on unnest(col1,depth=2), unnest(col1,depth=1) and unnest(col2,depth=1) -// output schema: -// unnest_col1_depth_2: int, unnest_col1_depth1: list[int], col1: list[list[int]], unnest_col2_depth_1: int -// Meaning the placeholder column will be replaced by its unnested variation(s), note -// the plural. - -/* - Note: unnest has a big difference in behavior between Postgres and DuckDB - Take this example - 1.Postgres - ```ignored - create table temp ( - i integer[][][], j integer[] - ) - insert into temp values ('{{{1,2},{3,4}},{{5,6},{7,8}}}', '{1,2}'); - select unnest(i), unnest(j) from temp; - ``` - - Result - 1 1 - 2 2 - 3 - 4 - 5 - 6 - 7 - 8 - 2. DuckDB - ```ignore - create table temp (i integer[][][], j integer[]); - insert into temp values ([[[1,2],[3,4]],[[5,6],[7,8]]], [1,2]); - select unnest(i,recursive:=true), unnest(j,recursive:=true) from temp; - ``` - Result: - ┌────────────────────────────────────────────────┬────────────────────────────────────────────────┐ - │ unnest(i, "recursive" := CAST('t' AS BOOLEAN)) │ unnest(j, "recursive" := CAST('t' AS BOOLEAN)) │ - │ int32 │ int32 │ - ├────────────────────────────────────────────────┼────────────────────────────────────────────────┤ - │ 1 │ 1 │ - │ 2 │ 2 │ - │ 3 │ 1 │ - │ 4 │ 2 │ - │ 5 │ 1 │ - │ 6 │ 2 │ - │ 7 │ 1 │ - │ 8 │ 2 │ - └────────────────────────────────────────────────┴────────────────────────────────────────────────┘ - The following implementation refer to Postgres's implementation - For DuckDB's result to be similar, the above query can be written as - - - -*/ - +/// Note: unnest has a big difference in behavior between Postgres and DuckDB +/// Take this example +/// 1.Postgres +/// ```ignored +/// create table temp ( +/// i integer[][][], j integer[] +/// ) +/// insert into temp values ('{{{1,2},{3,4}},{{5,6},{7,8}}}', '{1,2}'); +/// select unnest(i), unnest(j) from temp; +/// ``` +/// +/// Result +/// 1 1 +/// 2 2 +/// 3 +/// 4 +/// 5 +/// 6 +/// 7 +/// 8 +/// 2. DuckDB +/// ```ignore +/// create table temp (i integer[][][], j integer[]); +/// insert into temp values ([[[1,2],[3,4]],[[5,6],[7,8]]], [1,2]); +/// select unnest(i,recursive:=true), unnest(j,recursive:=true) from temp; +/// ``` +/// Result: +/// ┌────────────────────────────────────────────────┬────────────────────────────────────────────────┐ +/// │ unnest(i, "recursive" := CAST('t' AS BOOLEAN)) │ unnest(j, "recursive" := CAST('t' AS BOOLEAN)) │ +/// │ int32 │ int32 │ +/// ├────────────────────────────────────────────────┼────────────────────────────────────────────────┤ +/// │ 1 │ 1 │ +/// │ 2 │ 2 │ +/// │ 3 │ 1 │ +/// │ 4 │ 2 │ +/// │ 5 │ 1 │ +/// │ 6 │ 2 │ +/// │ 7 │ 1 │ +/// │ 8 │ 2 │ +/// └────────────────────────────────────────────────┴────────────────────────────────────────────────┘ +/// The following implementation refer to DuckDB's implementation +/// +/// Recursion happens for the highest level first +/// Demonstatring with examples: +/// +/// +/// Set "A" as a 3-dimension columns and "B" as an array (1-dimension) +/// Query: select unnest(A, max_depth:=3), unnest(A,max_depth:=2), unnest(B, max_depth:=1) from temp; +/// Let's given these projection names P1,P2,P3 respectively +/// +/// Each combination of (column,depth) result in an entry in temp_batch +/// This is needed, even if the same column is being unnested for different recursion levels +/// +/// This function is called with the descending order of recursion +/// Depth = 3 +/// - P1(3-dimension) unnest into temp column temp_P1(2_dimension) +/// - A(3-dimension) having indice repeated by the unnesting above +/// Depth = 2 +/// - temp_P1(2-dimension) unnest into temp column temp_P1(1-dimension) +/// - A(3-dimension) unnest into temp column temp_P2(2-dimension) +/// Depth = 1 +/// - temp_P1(1-dimension) unnest into P1 +/// - temp_P2(2-dimension) unnest into P2 +/// - B(1-dimension) unnest into P3 fn unnest_at_level( batch: &[ArrayRef], - list_type_columns: &[ListUnnest], - temp_batch: &mut HashMap<(usize, usize), ArrayRef>, + list_type_unnests: &[ListUnnest], + temp_batch: &mut HashMap, level_to_unnest: usize, options: &UnnestOptions, ) -> Result<(Vec, usize)> { - // unnested columns at this depth level - // now do some kind of projection-like - // This query: - // select unnest(col1, max_depth:=3), unnest(col1,max_depth:=2), unnest(col1, max_depth:=1) from temp; - // is equivalent to - // - // unnest(depth3) , unnest(depth2), unnest(depth1) - // select (unnest) - // batch comes in as [a,b] - // in this example list_type_columns as value - // [(a,1), (a,2)] and [(b,1)] - // 1.microwork(level=2) - // new batch as [unnest_a,b] - // new list as [(a,1)] and [(b,1)] - // temp_batch_for_projection: [(a,depth=1,unnested(a))] - // 2.microwork(level=1) - // new batch as [unnest(unnest_a),b] - // new list as [] - // temp_batch: [(a,depth=1,unnested(a)),(a,depth=2,unnested(unnested(a))),(b,depth=1,unnested(b))] - // this is final, now combine the mainbatch and temp_batch - - let (temp_unnest_cols, unnested_locations): (Vec>, Vec<_>) = - list_type_columns - .iter() - .filter_map( - |ListUnnest { - depth, - index_in_input_schema, - }| { - if *depth == level_to_unnest { - return Some(( - Arc::clone(&batch[*index_in_input_schema]), - (*index_in_input_schema, *depth), - )); - } - // this means the depth to unnest is still on going, keep on unnesting - // at current level - if *depth > level_to_unnest { - return Some(( - Arc::clone( - temp_batch - .get(&(*index_in_input_schema, *depth)) - .unwrap(), - ), - (*index_in_input_schema, *depth), - )); - } - None - }, - ) - .unzip(); + let (temp_unnest_cols, unnestings): (Vec>, Vec<_>) = list_type_unnests + .iter() + .filter_map(|unnesting| { + if level_to_unnest == unnesting.depth { + return Some(( + Arc::clone(&batch[unnesting.index_in_input_schema]), + *unnesting, + )); + } + // this means the unnesting on this item has started at higher level + // and need to continue until depth reaches 1 + if level_to_unnest < unnesting.depth { + return Some(( + Arc::clone(temp_batch.get(unnesting).unwrap()), + *unnesting, + )); + } + None + }) + .unzip(); // filter out so that list_arrays only contain column with the highest depth // at the same time, during iteration remove this depth so next time we don't have to unnest them again @@ -489,9 +465,9 @@ fn unnest_at_level( let ret = flatten_batch_from_indices(batch, &take_indices)?; unnested_temp_arrays .into_iter() - .zip(unnested_locations.iter()) - .for_each(|(flatten_arr, (index_in_input_schema, depth))| { - temp_batch.insert((*index_in_input_schema, *depth), flatten_arr); + .zip(unnestings.iter()) + .for_each(|(flatten_arr, unnesting)| { + temp_batch.insert(*unnesting, flatten_arr); }); Ok((ret, total_length)) } @@ -539,7 +515,14 @@ fn build_batch( let unnested_array_map: HashMap, usize)>> = temp_batch.into_iter().fold( HashMap::new(), - |mut acc, ((index_in_input_schema, depth), flattened_array)| { + |mut acc, + ( + ListUnnest { + index_in_input_schema, + depth, + }, + flattened_array, + )| { acc.entry(index_in_input_schema) .or_default() .push((flattened_array, depth)); @@ -817,31 +800,12 @@ fn create_take_indicies( fn flatten_batch_from_indices( batch: &[ArrayRef], - // temp: &[(usize, Arc)], - // mut unnested_list_arrays: HashMap>, indices: &PrimitiveArray, ) -> Result>> { - return batch + batch .into_iter() .map(|arr| Ok(kernels::take::take(arr, indices, None)?)) - .collect::>(); - // return Ok(new_arrays); - // Ok(new_arrays) - // let arrays = batch - // .columns() - // .iter() - // .enumerate() - // .map( - // |(col_idx, arr)| match unnested_list_arrays.remove(&col_idx) { - // Some(unnested_arrays) => Ok(unnested_arrays), - // None => Ok(vec![kernels::take::take(arr, indices, None)?]), - // }, - // ) - // .collect::>>()? - // .into_iter() - // .flatten() - // .collect::>(); - // Ok(arrays) + .collect::>() } /// Create the final batch given the unnested column arrays and a `indices` array From 70bccdb001e84abb732e5c20bf1f79398bd52f34 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Thu, 5 Sep 2024 22:22:24 +0200 Subject: [PATCH 32/56] more comments --- datafusion/physical-plan/src/unnest.rs | 68 +++++++++++++++++--------- 1 file changed, 45 insertions(+), 23 deletions(-) diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index c345ae9666f8..ac2ec22278a3 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -43,6 +43,7 @@ use datafusion_common::{ exec_datafusion_err, exec_err, internal_err, Result, UnnestOptions, }; use datafusion_execution::TaskContext; +use datafusion_expr::Unnest; use datafusion_physical_expr::EquivalenceProperties; use async_trait::async_trait; @@ -471,6 +472,10 @@ fn unnest_at_level( }); Ok((ret, total_length)) } +struct UnnestingResult { + arr: ArrayRef, + depth: usize, +} /// For each row in a `RecordBatch`, some list/struct columns need to be unnested. /// - For list columns: We will expand the values in each list into multiple rows, @@ -489,14 +494,17 @@ fn build_batch( 0 => flatten_struct_cols(batch.columns(), schema, struct_column_indices), _ => { let mut temp_batch = HashMap::new(); - let highest_depth = list_type_columns + let max_recursion = list_type_columns .iter() .fold(0, |highest_depth, ListUnnest { depth, .. }| { cmp::max(highest_depth, *depth) }); let mut unnested_original_columns = vec![]; - for depth in (1..=highest_depth).rev() { - let input = match depth == highest_depth { + + // original batch has the same columns + // all unnesting results are written to temp_batch + for depth in (1..=max_recursion).rev() { + let input = match depth == max_recursion { true => batch.columns(), false => &unnested_original_columns, }; @@ -512,7 +520,7 @@ fn build_batch( } unnested_original_columns = temp; } - let unnested_array_map: HashMap, usize)>> = + let unnested_array_map: HashMap> = temp_batch.into_iter().fold( HashMap::new(), |mut acc, @@ -523,41 +531,54 @@ fn build_batch( }, flattened_array, )| { - acc.entry(index_in_input_schema) - .or_default() - .push((flattened_array, depth)); + acc.entry(index_in_input_schema).or_default().push( + UnnestingResult { + arr: flattened_array, + depth, + }, + ); acc }, ); - let output_order: HashMap<(usize, usize), usize> = list_type_columns + let output_order: HashMap = list_type_columns .iter() .enumerate() - .map(|(order, unnest_def)| { - let ListUnnest { - depth, - index_in_input_schema, - } = unnest_def; - ((*index_in_input_schema, *depth), order) - }) + .map(|(order, unnest_def)| (*unnest_def, order)) .collect(); - let mut ordered_unnested_array_map = unnested_array_map + // one original column may be unnested multiple times into separate columns + let mut multi_unnested_per_original_index = unnested_array_map .into_iter() .map( // each item in unnested_columns is the result of unnesting the same input column // we need to sort them to conform with the unnest definition // e.g unnest(unnest(col)) must goes before unnest(col) |(original_index, mut unnested_columns)| { - unnested_columns.sort_by(|(_, depth), (_, depth2)| -> Ordering { - output_order.get(&(original_index, *depth)).unwrap().cmp( - output_order.get(&(original_index, *depth2)).unwrap(), - ) - }); + unnested_columns.sort_by( + |UnnestingResult { depth: depth1, .. }, + UnnestingResult { depth: depth2, .. }| + -> Ordering { + output_order + .get(&ListUnnest { + depth: *depth1, + index_in_input_schema: original_index, + }) + .unwrap() + .cmp( + output_order + .get(&ListUnnest { + depth: *depth2, + index_in_input_schema: original_index, + }) + .unwrap(), + ) + }, + ); ( original_index, unnested_columns .into_iter() - .map(|(arr, _)| arr) + .map(|result| result.arr) .collect::>(), ) }, @@ -568,7 +589,8 @@ fn build_batch( .into_iter() .enumerate() .flat_map(|(col_idx, arr)| { - match ordered_unnested_array_map.remove(&col_idx) { + // this column has some unnested output(s) + match multi_unnested_per_original_index.remove(&col_idx) { Some(unnested_arrays) => unnested_arrays, None => vec![arr], } From cc8169aec38b6449a8a6c5651e0ae74f6cd77640 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Sat, 7 Sep 2024 12:26:38 +0200 Subject: [PATCH 33/56] chore: better comments --- datafusion/physical-plan/src/unnest.rs | 114 +++++++++++-------------- 1 file changed, 50 insertions(+), 64 deletions(-) diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index ac2ec22278a3..ffd00283681f 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -43,7 +43,6 @@ use datafusion_common::{ exec_datafusion_err, exec_err, internal_err, Result, UnnestOptions, }; use datafusion_execution::TaskContext; -use datafusion_expr::Unnest; use datafusion_physical_expr::EquivalenceProperties; use async_trait::async_trait; @@ -412,37 +411,42 @@ pub struct ListUnnest { /// - temp_P1(1-dimension) unnest into P1 /// - temp_P2(2-dimension) unnest into P2 /// - B(1-dimension) unnest into P3 +/// +/// The returned array will has the same size as the input batch +/// and only contains original columns that are not being unnested fn unnest_at_level( batch: &[ArrayRef], list_type_unnests: &[ListUnnest], - temp_batch: &mut HashMap, + temp_unnested_arrs: &mut HashMap, level_to_unnest: usize, options: &UnnestOptions, ) -> Result<(Vec, usize)> { - let (temp_unnest_cols, unnestings): (Vec>, Vec<_>) = list_type_unnests - .iter() - .filter_map(|unnesting| { - if level_to_unnest == unnesting.depth { - return Some(( - Arc::clone(&batch[unnesting.index_in_input_schema]), - *unnesting, - )); - } - // this means the unnesting on this item has started at higher level - // and need to continue until depth reaches 1 - if level_to_unnest < unnesting.depth { - return Some(( - Arc::clone(temp_batch.get(unnesting).unwrap()), - *unnesting, - )); - } - None - }) - .unzip(); + // extract unnestable columns at this level + let (arrs_to_unnest, list_unnest_specs): (Vec>, Vec<_>) = + list_type_unnests + .iter() + .filter_map(|unnesting| { + if level_to_unnest == unnesting.depth { + return Some(( + Arc::clone(&batch[unnesting.index_in_input_schema]), + *unnesting, + )); + } + // this means the unnesting on this item has started at higher level + // and need to continue until depth reaches 1 + if level_to_unnest < unnesting.depth { + return Some(( + Arc::clone(temp_unnested_arrs.get(unnesting).unwrap()), + *unnesting, + )); + } + None + }) + .unzip(); // filter out so that list_arrays only contain column with the highest depth // at the same time, during iteration remove this depth so next time we don't have to unnest them again - let longest_length = find_longest_length(&temp_unnest_cols, options)?; + let longest_length = find_longest_length(&arrs_to_unnest, options)?; let unnested_length = longest_length.as_primitive::(); let total_length = if unnested_length.is_empty() { 0 @@ -457,18 +461,18 @@ fn unnest_at_level( // Unnest all the list arrays let unnested_temp_arrays = - unnest_list_arrays(temp_unnest_cols.as_ref(), unnested_length, total_length)?; + unnest_list_arrays(arrs_to_unnest.as_ref(), unnested_length, total_length)?; // Create the take indices array for other columns let take_indices = create_take_indicies(unnested_length, total_length); // vertical expansion because of list unnest - let ret = flatten_batch_from_indices(batch, &take_indices)?; + let ret = flatten_arrs_from_indices(batch, &take_indices)?; unnested_temp_arrays .into_iter() - .zip(unnestings.iter()) + .zip(list_unnest_specs.iter()) .for_each(|(flatten_arr, unnesting)| { - temp_batch.insert(*unnesting, flatten_arr); + temp_unnested_arrs.insert(*unnesting, flatten_arr); }); Ok((ret, total_length)) } @@ -493,35 +497,37 @@ fn build_batch( let transformed = match list_type_columns.len() { 0 => flatten_struct_cols(batch.columns(), schema, struct_column_indices), _ => { - let mut temp_batch = HashMap::new(); + let mut temp_unnested_result = HashMap::new(); let max_recursion = list_type_columns .iter() .fold(0, |highest_depth, ListUnnest { depth, .. }| { cmp::max(highest_depth, *depth) }); - let mut unnested_original_columns = vec![]; + + // This arr always has the same column count with the input batch + let mut flatten_arrs = vec![]; // original batch has the same columns // all unnesting results are written to temp_batch for depth in (1..=max_recursion).rev() { let input = match depth == max_recursion { true => batch.columns(), - false => &unnested_original_columns, + false => &flatten_arrs, }; - let (temp, num_rows) = unnest_at_level( + let (temp_result, num_rows) = unnest_at_level( input, list_type_columns, - &mut temp_batch, + &mut temp_unnested_result, depth, options, )?; if num_rows == 0 { return Ok(RecordBatch::new_empty(Arc::clone(schema))); } - unnested_original_columns = temp; + flatten_arrs = temp_result; } let unnested_array_map: HashMap> = - temp_batch.into_iter().fold( + temp_unnested_result.into_iter().fold( HashMap::new(), |mut acc, ( @@ -585,11 +591,13 @@ fn build_batch( ) .collect::>(); - let ret = unnested_original_columns + let ret = flatten_arrs .into_iter() .enumerate() .flat_map(|(col_idx, arr)| { - // this column has some unnested output(s) + // convert original column into its unnested version(s) + // Plural because one column can be unnested with different recursion level + // and into separate output columns match multi_unnested_per_original_index.remove(&col_idx) { Some(unnested_arrays) => unnested_arrays, None => vec![arr], @@ -820,17 +828,7 @@ fn create_take_indicies( builder.finish() } -fn flatten_batch_from_indices( - batch: &[ArrayRef], - indices: &PrimitiveArray, -) -> Result>> { - batch - .into_iter() - .map(|arr| Ok(kernels::take::take(arr, indices, None)?)) - .collect::>() -} - -/// Create the final batch given the unnested column arrays and a `indices` array +/// Create the batch given the unnested column arrays and a `indices` array /// that is used by the take kernel to copy values. /// /// For example if we have the following `RecordBatch`: @@ -861,26 +859,14 @@ fn flatten_batch_from_indices( /// c2: 'a', 'b', 'c', 'c', 'c', null, 'd', 'd' /// ``` /// -fn flatten_list_cols_from_indices( - batch: &RecordBatch, - mut unnested_list_arrays: HashMap>, +fn flatten_arrs_from_indices( + batch: &[ArrayRef], indices: &PrimitiveArray, ) -> Result>> { - let arrays = batch - .columns() + batch .iter() - .enumerate() - .map( - |(col_idx, arr)| match unnested_list_arrays.remove(&col_idx) { - Some(unnested_arrays) => Ok(unnested_arrays), - None => Ok(vec![kernels::take::take(arr, indices, None)?]), - }, - ) - .collect::>>()? - .into_iter() - .flatten() - .collect::>(); - Ok(arrays) + .map(|arr| Ok(kernels::take::take(arr, indices, None)?)) + .collect::>() } #[cfg(test)] From 89e454708c78439602d73eeca201997c23179304 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Sun, 8 Sep 2024 10:26:02 +0200 Subject: [PATCH 34/56] remove breaking api --- datafusion/core/src/dataframe/mod.rs | 6 +- datafusion/expr/src/logical_plan/builder.rs | 75 +++++++++++++++++---- datafusion/sql/src/select.rs | 20 +----- 3 files changed, 66 insertions(+), 35 deletions(-) diff --git a/datafusion/core/src/dataframe/mod.rs b/datafusion/core/src/dataframe/mod.rs index 2336724c3c86..2c132ca93cce 100644 --- a/datafusion/core/src/dataframe/mod.rs +++ b/datafusion/core/src/dataframe/mod.rs @@ -365,10 +365,8 @@ impl DataFrame { columns: &[&str], options: UnnestOptions, ) -> Result { - let columns = columns - .iter() - .map(|c| (Column::from(*c), ColumnUnnestType::Inferred)) - .collect(); + let columns = columns.iter().map(|c| Column::from(*c)).collect(); + let plan = LogicalPlanBuilder::from(self.plan) .unnest_columns_with_options(columns, options)? .build()?; diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index 2f06abdeb4a6..66251905ec80 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -1135,6 +1135,22 @@ impl LogicalPlanBuilder { /// Unnest the given columns with the given [`UnnestOptions`] pub fn unnest_columns_with_options( + self, + // columns: Vec<(Column, ColumnUnnestType)>, + columns: Vec, + options: UnnestOptions, + ) -> Result { + Ok(Self::from(unnest_with_options( + self.plan, + columns + .into_iter() + .map(|c| (c, ColumnUnnestType::Inferred)).collect(), + options, + )?)) + } + + /// TODO: rename + pub fn unnest_columns_with_options_v2( self, columns: Vec<(Column, ColumnUnnestType)>, options: UnnestOptions, @@ -1561,13 +1577,14 @@ pub fn get_struct_unnested_columns( .collect() } -// TODO: make me recursive // Based on data type, either struct or a variant of list // return a set of columns as the result of unnesting // the input columns. // For example, given a column with name "a", // - List(Element) returns ["a"] with data type Element // - Struct(field1, field2) returns ["a.field1","a.field2"] +// For list data type, an argument depth is used to specify +// the recursion level pub fn get_unnested_columns( col_name: &String, data_type: &DataType, @@ -1608,26 +1625,39 @@ pub fn get_unnested_columns( Ok(qualified_columns) } -// a list type column can be performed differently at the same time -// e.g select unnest(col), unnest(unnest(col)) -// while unnest struct can only be performed once at a time - /// Create a [`LogicalPlan::Unnest`] plan with options -/// This function receive a map of columns to be unnested +/// This function receive a list of columns to be unnested /// because multiple unnest can be performed on the same column (e.g unnest with different depth) /// The new schema will contains post-unnest fields replacing the original field /// -/// input schema as: col1: int| col2: [][]int -/// Then unnest_map with { col2 -> [(col2,depth=1), (col2,depth=2)] } -/// will generate a new schema as col1: int| unnest_col2_depth1: []int| unnest_col2_depth2: int +/// For example: +/// Input schema as +/// +/// +---------------------+-----------+ +/// | col1 | col2 | +/// +---------------------+-----------+ +/// | Struct(INT64,INT32) | [[Int64]] | +/// +---------------------+-----------+ +/// +/// Then unnesting columns with: +/// - [col1,Struct] +/// - [col2,List([depth=1,depth=2])] +/// +/// will generate a new schema as +/// +/// +---------+---------+---------------------+---------------------+ +/// | col1.c0 | col1.c1 | unnest_col2_depth_1 | unnest_col2_depth_2 | +/// +---------+---------+---------------------+---------------------+ +/// | Int64 | Int32 | [Int64] | Int64 | +/// +---------+---------+---------------------+---------------------+ pub fn unnest_with_options( input: LogicalPlan, - columns: Vec<(Column, ColumnUnnestType)>, + columns_to_unnest: Vec<(Column, ColumnUnnestType)>, options: UnnestOptions, ) -> Result { let mut list_columns: Vec<(usize, ColumnUnnestList)> = vec![]; let mut struct_columns = vec![]; - let indices_to_unnest = columns + let indices_to_unnest = columns_to_unnest .iter() .map(|col_unnesting| { Ok(( @@ -1738,7 +1768,7 @@ pub fn unnest_with_options( Ok(LogicalPlan::Unnest(Unnest { input: Arc::new(input), - exec_columns: columns, + exec_columns: columns_to_unnest, list_type_columns: list_columns, struct_type_columns: struct_columns, dependency_indices, @@ -1749,6 +1779,8 @@ pub fn unnest_with_options( #[cfg(test)] mod tests { + use std::string; + use super::*; use crate::logical_plan::StringifiedPlan; use crate::{col, expr, expr_fn::exists, in_subquery, lit, scalar_subquery}; @@ -2143,7 +2175,7 @@ mod tests { // Unnesting multiple fields at the same time, using infer syntax let cols = vec!["strings", "structs", "struct_singular"] .into_iter() - .map(|c| (Column::from(c), ColumnUnnestType::Inferred)) + .map(|c| c.into()) .collect(); let plan = nested_table_scan("test_table")? @@ -2159,6 +2191,21 @@ mod tests { let plan = nested_table_scan("test_table")?.unnest_column("missing"); assert!(plan.is_err()); + // Recursive unnest at different recursion levels + let plan = nested_table_scan("test_table")? + .unnest_column("strings")? + .unnest_column_with_options(column) + .build()?; + + let expected = "\ + Unnest: lists[test_table.strings|depth=1] structs[]\ + \n TableScan: test_table"; + assert_eq!(expected, format!("{plan}")); + + // Check unnested field is a scalar + let field = plan.schema().field_with_name(None, "strings").unwrap(); + assert_eq!(&DataType::Utf8, field.data_type()); + Ok(()) } @@ -2174,6 +2221,7 @@ mod tests { false, ); let string_field = Field::new("item", DataType::Utf8, false); + let strings_field = Field::new_list("item", string, false); let schema = Schema::new(vec![ Field::new("scalar", DataType::UInt32, false), Field::new_list("strings", string_field, false), @@ -2186,6 +2234,7 @@ mod tests { ])), false, ), + Field::new("stringss", strings_field, false), ]); table_scan(Some(table_name), &schema, None) diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs index b9b332b6bacb..2e497cb0b290 100644 --- a/datafusion/sql/src/select.rs +++ b/datafusion/sql/src/select.rs @@ -338,28 +338,12 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { } break; } else { - let columns = unnest_columns.into_iter().map(|col| col).collect(); // Set preserve_nulls to false to ensure compatibility with DuckDB and PostgreSQL let unnest_options = UnnestOptions::new().with_preserve_nulls(false); - // let mut check_list: HashSet = inner_projection_exprs - // .iter() - // .map(|expr| expr.clone()) - // .collect(); - // let deduplicated: Vec = inner_projection_exprs - // .into_iter() - // .filter(|expr| -> bool { - // return true; - // if check_list.remove(expr) { - // true - // } else { - // false - // } - // }) - // .collect(); let plan = LogicalPlanBuilder::from(intermediate_plan) .project(inner_projection_exprs)? - .unnest_columns_with_options(columns, unnest_options)? + .unnest_columns_with_options_v2(unnest_columns, unnest_options)? .build()?; intermediate_plan = plan; intermediate_select_exprs = outer_projection_exprs; @@ -471,7 +455,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { intermediate_plan = LogicalPlanBuilder::from(intermediate_plan) .project(projection_exprs)? - .unnest_columns_with_options(columns, unnest_options)? + .unnest_columns_with_options_v2(columns, unnest_options)? .build()?; intermediate_select_exprs = outer_projection_exprs; From eb215b12aebab69223ad71dc2a188fe4ca094d79 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Sun, 8 Sep 2024 10:32:52 +0200 Subject: [PATCH 35/56] rename --- datafusion/core/src/dataframe/mod.rs | 2 +- datafusion/expr/src/logical_plan/builder.rs | 13 ++++++++----- datafusion/sql/src/select.rs | 7 +++++-- 3 files changed, 14 insertions(+), 8 deletions(-) diff --git a/datafusion/core/src/dataframe/mod.rs b/datafusion/core/src/dataframe/mod.rs index 2c132ca93cce..8a5771454d7f 100644 --- a/datafusion/core/src/dataframe/mod.rs +++ b/datafusion/core/src/dataframe/mod.rs @@ -51,7 +51,7 @@ use datafusion_common::config::{CsvOptions, JsonOptions}; use datafusion_common::{ plan_err, Column, DFSchema, DataFusionError, ParamValues, SchemaError, UnnestOptions, }; -use datafusion_expr::{case, is_null, lit, ColumnUnnestType}; +use datafusion_expr::{case, is_null, lit}; use datafusion_expr::{ utils::COUNT_STAR_EXPANSION, TableProviderFilterPushDown, UNNAMED_TABLE, }; diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index 66251905ec80..d929d2dba452 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -1136,21 +1136,24 @@ impl LogicalPlanBuilder { /// Unnest the given columns with the given [`UnnestOptions`] pub fn unnest_columns_with_options( self, - // columns: Vec<(Column, ColumnUnnestType)>, columns: Vec, options: UnnestOptions, ) -> Result { Ok(Self::from(unnest_with_options( self.plan, columns - .into_iter() - .map(|c| (c, ColumnUnnestType::Inferred)).collect(), + .into_iter() + .map(|c| (c, ColumnUnnestType::Inferred)) + .collect(), options, )?)) } - /// TODO: rename - pub fn unnest_columns_with_options_v2( + /// Unnest the given columns with the given [`UnnestOptions`] + /// if one column is an list type, it can be recursively and simultaneously + /// unnested into the desired recursion levels + /// e.g select unnest(list_col,depth=1), unnest(list_col,depth=2) + pub fn unnest_columns_recursive_with_options( self, columns: Vec<(Column, ColumnUnnestType)>, options: UnnestOptions, diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs index 2e497cb0b290..b00eba431e57 100644 --- a/datafusion/sql/src/select.rs +++ b/datafusion/sql/src/select.rs @@ -343,7 +343,10 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { let plan = LogicalPlanBuilder::from(intermediate_plan) .project(inner_projection_exprs)? - .unnest_columns_with_options_v2(unnest_columns, unnest_options)? + .unnest_columns_recursive_with_options( + unnest_columns, + unnest_options, + )? .build()?; intermediate_plan = plan; intermediate_select_exprs = outer_projection_exprs; @@ -455,7 +458,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { intermediate_plan = LogicalPlanBuilder::from(intermediate_plan) .project(projection_exprs)? - .unnest_columns_with_options_v2(columns, unnest_options)? + .unnest_columns_recursive_with_options(columns, unnest_options)? .build()?; intermediate_select_exprs = outer_projection_exprs; From 1deb566d697053d09eda83c7d2b0300260cd4104 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Sun, 8 Sep 2024 11:51:33 +0200 Subject: [PATCH 36/56] more unit test --- datafusion/expr/src/expr.rs | 1 - datafusion/expr/src/logical_plan/builder.rs | 56 ++++++++++++++++---- datafusion/expr/src/logical_plan/plan.rs | 9 ++-- datafusion/optimizer/src/push_down_filter.rs | 15 ++---- 4 files changed, 54 insertions(+), 27 deletions(-) diff --git a/datafusion/expr/src/expr.rs b/datafusion/expr/src/expr.rs index 0699f39f0341..8a6af4dc53cd 100644 --- a/datafusion/expr/src/expr.rs +++ b/datafusion/expr/src/expr.rs @@ -2360,7 +2360,6 @@ impl fmt::Display for Expr { }, Expr::Placeholder(Placeholder { id, .. }) => write!(f, "{id}"), Expr::Unnest(Unnest { expr }) => { - // TODO: use Display instead of Debug, there is non-unique expression name in projection issue. write!(f, "UNNEST({expr})") } } diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index d929d2dba452..35e2f2c5ec4e 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -1150,7 +1150,7 @@ impl LogicalPlanBuilder { } /// Unnest the given columns with the given [`UnnestOptions`] - /// if one column is an list type, it can be recursively and simultaneously + /// if one column is a list type, it can be recursively and simultaneously /// unnested into the desired recursion levels /// e.g select unnest(list_col,depth=1), unnest(list_col,depth=2) pub fn unnest_columns_recursive_with_options( @@ -1731,7 +1731,6 @@ pub fn unnest_with_options( ) }, ) - // TODO: i'm messy .collect::)>>>>()? .into_iter() .flatten() @@ -1782,7 +1781,6 @@ pub fn unnest_with_options( #[cfg(test)] mod tests { - use std::string; use super::*; use crate::logical_plan::StringifiedPlan; @@ -2194,20 +2192,56 @@ mod tests { let plan = nested_table_scan("test_table")?.unnest_column("missing"); assert!(plan.is_err()); - // Recursive unnest at different recursion levels + // Simultaneously unnesting a list (with different depth) and a struct column let plan = nested_table_scan("test_table")? - .unnest_column("strings")? - .unnest_column_with_options(column) + .unnest_columns_recursive_with_options( + vec![ + ( + "stringss".into(), + ColumnUnnestType::List(vec![ + ColumnUnnestList { + output_column: Column::from_name("stringss_depth_1"), + depth: 1, + }, + ColumnUnnestList { + output_column: Column::from_name("stringss_depth_2"), + depth: 2, + }, + ]), + ), + ("struct_singular".into(), ColumnUnnestType::Inferred), + ], + UnnestOptions::default(), + )? .build()?; let expected = "\ - Unnest: lists[test_table.strings|depth=1] structs[]\ + Unnest: lists[test_table.stringss|depth=1, test_table.stringss|depth=2] structs[test_table.struct_singular]\ \n TableScan: test_table"; assert_eq!(expected, format!("{plan}")); - // Check unnested field is a scalar - let field = plan.schema().field_with_name(None, "strings").unwrap(); + // Check output columns has correct type + let field = plan + .schema() + .field_with_name(None, "stringss_depth_1") + .unwrap(); + assert_eq!( + &DataType::new_list(DataType::Utf8, false), + field.data_type() + ); + let field = plan + .schema() + .field_with_name(None, "stringss_depth_2") + .unwrap(); assert_eq!(&DataType::Utf8, field.data_type()); + // unnesting struct is still correct + for field_name in &["a", "b"] { + let field = plan + .schema() + .field_with_name(None, &format!("struct_singular.{}", field_name)) + .unwrap(); + assert_eq!(&DataType::UInt32, field.data_type()); + } Ok(()) } @@ -2224,7 +2258,7 @@ mod tests { false, ); let string_field = Field::new("item", DataType::Utf8, false); - let strings_field = Field::new_list("item", string, false); + let strings_field = Field::new_list("item", string_field.clone(), false); let schema = Schema::new(vec![ Field::new("scalar", DataType::UInt32, false), Field::new_list("strings", string_field, false), @@ -2237,7 +2271,7 @@ mod tests { ])), false, ), - Field::new("stringss", strings_field, false), + Field::new_list("stringss", strings_field, false), ]); table_scan(Some(table_name), &schema, None) diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index 5a7283e25d3c..f5e000dd3e5d 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -17,6 +17,11 @@ //! Logical plan types +use std::collections::{HashMap, HashSet}; +use std::fmt::{self, Debug, Display, Formatter}; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; + use super::dml::CopyTo; use super::DdlStatement; use crate::builder::{change_redundant_column, unnest_with_options}; @@ -35,10 +40,6 @@ use crate::{ CreateMemoryTable, CreateView, Expr, ExprSchemable, LogicalPlanBuilder, Operator, TableProviderFilterPushDown, TableSource, WindowFunctionDefinition, }; -use std::collections::{HashMap, HashSet}; -use std::fmt::{self, Debug, Display, Formatter}; -use std::hash::{Hash, Hasher}; -use std::sync::Arc; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; diff --git a/datafusion/optimizer/src/push_down_filter.rs b/datafusion/optimizer/src/push_down_filter.rs index 7f16dcb6d9bf..e2b19ab33f73 100644 --- a/datafusion/optimizer/src/push_down_filter.rs +++ b/datafusion/optimizer/src/push_down_filter.rs @@ -36,8 +36,8 @@ use datafusion_expr::utils::{ conjunction, expr_to_columns, split_conjunction, split_conjunction_owned, }; use datafusion_expr::{ - and, build_join_schema, or, BinaryExpr, ColumnUnnestType, Expr, Filter, - LogicalPlanBuilder, Operator, Projection, TableProviderFilterPushDown, + and, build_join_schema, or, BinaryExpr, Expr, Filter, LogicalPlanBuilder, Operator, + Projection, TableProviderFilterPushDown, }; use crate::optimizer::ApplyOrder; @@ -745,15 +745,8 @@ impl OptimizerRule for PushDownFilter { let mut accum: HashSet = HashSet::new(); expr_to_columns(&predicate, &mut accum)?; - if unnest.exec_columns.iter().any(|(unnest_col, unnest_type)| { - match unnest_type { - ColumnUnnestType::List(vec) => { - vec.iter().any(|c| accum.contains(&c.output_column)) - } - ColumnUnnestType::Struct => false, - // for inferred unnest, output column will be the same with input column - ColumnUnnestType::Inferred => accum.contains(unnest_col), - } + if unnest.list_type_columns.iter().any(|(_, unnest_list)| { + accum.contains(&unnest_list.output_column) }) { unnest_predicates.push(predicate); } else { From 2fb984280fd56e131887bf73fab42eb309b82128 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Sun, 8 Sep 2024 11:53:08 +0200 Subject: [PATCH 37/56] remove debug --- datafusion/sqllogictest/test_files/debug.slt | 6 ------ 1 file changed, 6 deletions(-) delete mode 100644 datafusion/sqllogictest/test_files/debug.slt diff --git a/datafusion/sqllogictest/test_files/debug.slt b/datafusion/sqllogictest/test_files/debug.slt deleted file mode 100644 index 951c9132cb82..000000000000 --- a/datafusion/sqllogictest/test_files/debug.slt +++ /dev/null @@ -1,6 +0,0 @@ - -statement ok -CREATE TABLE unnest_table -AS VALUES - ([1,2,3], [7], 1, [13, 14], struct(1,2)) -; From 88b8edd3d24c34b3f0b510b9eb1ab811568eb12a Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Sun, 8 Sep 2024 12:27:11 +0200 Subject: [PATCH 38/56] clean up --- datafusion/sql/src/select.rs | 16 ++++------------ datafusion/sql/src/utils.rs | 10 +--------- datafusion/sqllogictest/test_files/unnest.slt | 14 +++++--------- 3 files changed, 10 insertions(+), 30 deletions(-) diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs index 39c8340ea881..0d487684db7b 100644 --- a/datafusion/sql/src/select.rs +++ b/datafusion/sql/src/select.rs @@ -329,7 +329,6 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { if unnest_columns.is_empty() { // The original expr does not contain any unnest if i == 0 { - // return Ok(intermediate_plan); return LogicalPlanBuilder::from(intermediate_plan) .project(intermediate_select_exprs)? .build(); @@ -362,16 +361,14 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { match input { LogicalPlan::Aggregate(agg) => { let agg_expr = agg.aggr_expr.clone(); - // we somehow need to un_rebase column exprs appeared in select/group by exprs - // e.g unnest(col("somecole")) expr may be rewritten as col("unnest(some_col)") let (new_input, new_group_by_exprs) = self.try_process_group_by_unnest(agg)?; - Ok(LogicalPlanBuilder::from(new_input) + LogicalPlanBuilder::from(new_input) .aggregate(new_group_by_exprs, agg_expr)? - .build()?) + .build() } LogicalPlan::Filter(mut filter) => { - filter.input = + filter.input = Arc::new(self.try_process_aggregate_unnest(Arc::unwrap_or_clone( filter.input, ))?); @@ -461,15 +458,10 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { .build()?; intermediate_select_exprs = outer_projection_exprs; - // intermediate_group_by_exprs = temp_new_group_exprs; } } - Ok(( - intermediate_plan, - intermediate_select_exprs, - // intermediate_group_by_exprs, - )) + Ok((intermediate_plan, intermediate_select_exprs)) } fn plan_selection( diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index 88b8ee4ca718..e5199059f2dc 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -33,16 +33,8 @@ use datafusion_common::{ use datafusion_expr::builder::get_struct_unnested_columns; use datafusion_expr::expr::{Alias, GroupingSet, Unnest, WindowFunction}; use datafusion_expr::utils::{expr_as_column_expr, find_column_exprs}; -<<<<<<< HEAD -use datafusion_expr::{col, expr_vec_fmt, Expr, ExprSchemable, LogicalPlan}; -use datafusion_expr::{ColumnUnnestList, ColumnUnnestType}; -use sqlparser::ast::Ident; -use sqlparser::ast::Value; -======= -use datafusion_expr::{expr_vec_fmt, Expr, ExprSchemable, LogicalPlan}; +use datafusion_expr::{col,expr_vec_fmt, Expr, ExprSchemable, LogicalPlan,ColumnUnnestList, ColumnUnnestType}; use sqlparser::ast::{Ident, Value}; -use std::collections::HashMap; ->>>>>>> origin/main /// Make a best-effort attempt at resolving all columns in the expression tree pub(crate) fn resolve_columns(expr: &Expr, plan: &LogicalPlan) -> Result { diff --git a/datafusion/sqllogictest/test_files/unnest.slt b/datafusion/sqllogictest/test_files/unnest.slt index 34231f9f92db..76b75fdee831 100644 --- a/datafusion/sqllogictest/test_files/unnest.slt +++ b/datafusion/sqllogictest/test_files/unnest.slt @@ -44,13 +44,6 @@ AS VALUES (struct([2], 'b'), [[[3,4],[5]],[[null,6],null,[7,8]]], [struct([2],[[3],[4]])]) ; -query I -select unnest([1,2,3]); ----- -1 -2 -3 - ## Basic unnest expression in select list query I select unnest([1,2,3]); @@ -173,6 +166,7 @@ select unnest(column1), column1 from unnest_table; 12 [12] +# binary expr linking different unnest exprs query II select unnest([1,2,3]) + unnest([1,2,3]), unnest([1,2,3]) + unnest([4,5]); ---- @@ -181,6 +175,7 @@ select unnest([1,2,3]) + unnest([1,2,3]), unnest([1,2,3]) + unnest([4,5]); 6 NULL +# binary expr linking different recursive unnest exprs query III select unnest(unnest([[1,2,3]])) + unnest(unnest([[1,2,3]])), unnest(unnest([[1,2,3]])) + unnest([4,5]), unnest([4,5]); ---- @@ -552,7 +547,7 @@ NULL [7, 8] NULL NULL [[3, 4], [5]] NULL 7 8 [[, 6], , [7, 8]] NULL 8 9 -## the same unnest expr is referened multiple times (unnest is not the bottom-most expr) +## the same composite expr (unnest(field_access(unnest(col)))) which containing unnest is referened multiple times query ??II select unnest(column3), unnest(column3)['c0'], unnest(unnest(column3)['c0']), unnest(unnest(column3)['c0']) + unnest(unnest(column3)['c0']) from recursive_unnest_table; ---- @@ -605,7 +600,8 @@ NULL [4] [{c0: [2], c1: [[3], [4]]}] 4 [3] [{c0: [2], c1: [[3], [4]]}] NULL [4] [{c0: [2], c1: [[3], [4]]}] -## demonstrate where multiple unnest plan is needed +## demonstrate where recursive unnest is impossible +## and need multiple unnesting logical plans ## e.g unnest -> field_access -> unnest query TT explain select unnest(unnest(unnest(column3)['c1'])), unnest(unnest(column3)['c1']), column3 from recursive_unnest_table; From 00dab2f47db3606f6929857772a0663461345881 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Tue, 10 Sep 2024 20:59:59 +0200 Subject: [PATCH 39/56] fix proto --- datafusion/proto/proto/datafusion.proto | 33 +- datafusion/proto/src/generated/pbjson.rs | 579 +++++++++++++++++- datafusion/proto/src/generated/prost.rs | 61 +- datafusion/proto/src/logical_plan/mod.rs | 93 ++- datafusion/proto/src/physical_plan/mod.rs | 20 +- .../tests/cases/roundtrip_physical_plan.rs | 23 +- 6 files changed, 778 insertions(+), 31 deletions(-) diff --git a/datafusion/proto/proto/datafusion.proto b/datafusion/proto/proto/datafusion.proto index 645df14a0337..1204c843fdb1 100644 --- a/datafusion/proto/proto/datafusion.proto +++ b/datafusion/proto/proto/datafusion.proto @@ -262,13 +262,35 @@ message CopyToNode { message UnnestNode { LogicalPlanNode input = 1; - repeated datafusion_common.Column exec_columns = 2; - repeated uint64 list_type_columns = 3; + repeated ColumnUnnestExec exec_columns = 2; + repeated ColumnUnnestListItem list_type_columns = 3; repeated uint64 struct_type_columns = 4; repeated uint64 dependency_indices = 5; datafusion_common.DfSchema schema = 6; UnnestOptions options = 7; } +message ColumnUnnestListItem { + uint32 input_index = 1; + ColumnUnnestListRecursion recursion = 2; +} + +message ColumnUnnestListRecursions { + repeated ColumnUnnestListRecursion recursions = 2; +} + +message ColumnUnnestListRecursion { + datafusion_common.Column output_column = 1; + uint32 depth = 2; +} + +message ColumnUnnestExec { + datafusion_common.Column column = 1; + oneof UnnestType { + ColumnUnnestListRecursions list = 2; + datafusion_common.EmptyMessage struct = 3; + datafusion_common.EmptyMessage inferred = 4; + } +} message UnnestOptions { bool preserve_nulls = 1; @@ -758,11 +780,16 @@ message ParquetSinkExecNode { message UnnestExecNode { PhysicalPlanNode input = 1; datafusion_common.Schema schema = 2; - repeated uint64 list_type_columns = 3; + repeated ListUnnest list_type_columns = 3; repeated uint64 struct_type_columns = 4; UnnestOptions options = 5; } +message ListUnnest { + uint32 index_in_input_schema = 1; + uint32 depth = 2; +} + message PhysicalExtensionNode { bytes node = 1; repeated PhysicalPlanNode inputs = 2; diff --git a/datafusion/proto/src/generated/pbjson.rs b/datafusion/proto/src/generated/pbjson.rs index e493f761b51f..0614e33b7a4b 100644 --- a/datafusion/proto/src/generated/pbjson.rs +++ b/datafusion/proto/src/generated/pbjson.rs @@ -2321,6 +2321,458 @@ impl<'de> serde::Deserialize<'de> for ColumnIndex { deserializer.deserialize_struct("datafusion.ColumnIndex", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for ColumnUnnestExec { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.column.is_some() { + len += 1; + } + if self.unnest_type.is_some() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.ColumnUnnestExec", len)?; + if let Some(v) = self.column.as_ref() { + struct_ser.serialize_field("column", v)?; + } + if let Some(v) = self.unnest_type.as_ref() { + match v { + column_unnest_exec::UnnestType::List(v) => { + struct_ser.serialize_field("list", v)?; + } + column_unnest_exec::UnnestType::Struct(v) => { + struct_ser.serialize_field("struct", v)?; + } + column_unnest_exec::UnnestType::Inferred(v) => { + struct_ser.serialize_field("inferred", v)?; + } + } + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for ColumnUnnestExec { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "column", + "list", + "struct", + "inferred", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Column, + List, + Struct, + Inferred, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "column" => Ok(GeneratedField::Column), + "list" => Ok(GeneratedField::List), + "struct" => Ok(GeneratedField::Struct), + "inferred" => Ok(GeneratedField::Inferred), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = ColumnUnnestExec; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.ColumnUnnestExec") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut column__ = None; + let mut unnest_type__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Column => { + if column__.is_some() { + return Err(serde::de::Error::duplicate_field("column")); + } + column__ = map_.next_value()?; + } + GeneratedField::List => { + if unnest_type__.is_some() { + return Err(serde::de::Error::duplicate_field("list")); + } + unnest_type__ = map_.next_value::<::std::option::Option<_>>()?.map(column_unnest_exec::UnnestType::List) +; + } + GeneratedField::Struct => { + if unnest_type__.is_some() { + return Err(serde::de::Error::duplicate_field("struct")); + } + unnest_type__ = map_.next_value::<::std::option::Option<_>>()?.map(column_unnest_exec::UnnestType::Struct) +; + } + GeneratedField::Inferred => { + if unnest_type__.is_some() { + return Err(serde::de::Error::duplicate_field("inferred")); + } + unnest_type__ = map_.next_value::<::std::option::Option<_>>()?.map(column_unnest_exec::UnnestType::Inferred) +; + } + } + } + Ok(ColumnUnnestExec { + column: column__, + unnest_type: unnest_type__, + }) + } + } + deserializer.deserialize_struct("datafusion.ColumnUnnestExec", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for ColumnUnnestListItem { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.input_index != 0 { + len += 1; + } + if self.recursion.is_some() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.ColumnUnnestListItem", len)?; + if self.input_index != 0 { + struct_ser.serialize_field("inputIndex", &self.input_index)?; + } + if let Some(v) = self.recursion.as_ref() { + struct_ser.serialize_field("recursion", v)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for ColumnUnnestListItem { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "input_index", + "inputIndex", + "recursion", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + InputIndex, + Recursion, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "inputIndex" | "input_index" => Ok(GeneratedField::InputIndex), + "recursion" => Ok(GeneratedField::Recursion), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = ColumnUnnestListItem; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.ColumnUnnestListItem") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut input_index__ = None; + let mut recursion__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::InputIndex => { + if input_index__.is_some() { + return Err(serde::de::Error::duplicate_field("inputIndex")); + } + input_index__ = + Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) + ; + } + GeneratedField::Recursion => { + if recursion__.is_some() { + return Err(serde::de::Error::duplicate_field("recursion")); + } + recursion__ = map_.next_value()?; + } + } + } + Ok(ColumnUnnestListItem { + input_index: input_index__.unwrap_or_default(), + recursion: recursion__, + }) + } + } + deserializer.deserialize_struct("datafusion.ColumnUnnestListItem", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for ColumnUnnestListRecursion { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.output_column.is_some() { + len += 1; + } + if self.depth != 0 { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.ColumnUnnestListRecursion", len)?; + if let Some(v) = self.output_column.as_ref() { + struct_ser.serialize_field("outputColumn", v)?; + } + if self.depth != 0 { + struct_ser.serialize_field("depth", &self.depth)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for ColumnUnnestListRecursion { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "output_column", + "outputColumn", + "depth", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + OutputColumn, + Depth, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "outputColumn" | "output_column" => Ok(GeneratedField::OutputColumn), + "depth" => Ok(GeneratedField::Depth), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = ColumnUnnestListRecursion; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.ColumnUnnestListRecursion") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut output_column__ = None; + let mut depth__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::OutputColumn => { + if output_column__.is_some() { + return Err(serde::de::Error::duplicate_field("outputColumn")); + } + output_column__ = map_.next_value()?; + } + GeneratedField::Depth => { + if depth__.is_some() { + return Err(serde::de::Error::duplicate_field("depth")); + } + depth__ = + Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) + ; + } + } + } + Ok(ColumnUnnestListRecursion { + output_column: output_column__, + depth: depth__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.ColumnUnnestListRecursion", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for ColumnUnnestListRecursions { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if !self.recursions.is_empty() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.ColumnUnnestListRecursions", len)?; + if !self.recursions.is_empty() { + struct_ser.serialize_field("recursions", &self.recursions)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for ColumnUnnestListRecursions { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "recursions", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Recursions, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "recursions" => Ok(GeneratedField::Recursions), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = ColumnUnnestListRecursions; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.ColumnUnnestListRecursions") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut recursions__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Recursions => { + if recursions__.is_some() { + return Err(serde::de::Error::duplicate_field("recursions")); + } + recursions__ = Some(map_.next_value()?); + } + } + } + Ok(ColumnUnnestListRecursions { + recursions: recursions__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.ColumnUnnestListRecursions", FIELDS, GeneratedVisitor) + } +} impl serde::Serialize for CopyToNode { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result @@ -8763,6 +9215,119 @@ impl<'de> serde::Deserialize<'de> for ListRange { deserializer.deserialize_struct("datafusion.ListRange", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for ListUnnest { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.index_in_input_schema != 0 { + len += 1; + } + if self.depth != 0 { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.ListUnnest", len)?; + if self.index_in_input_schema != 0 { + struct_ser.serialize_field("indexInInputSchema", &self.index_in_input_schema)?; + } + if self.depth != 0 { + struct_ser.serialize_field("depth", &self.depth)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for ListUnnest { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "index_in_input_schema", + "indexInInputSchema", + "depth", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + IndexInInputSchema, + Depth, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "indexInInputSchema" | "index_in_input_schema" => Ok(GeneratedField::IndexInInputSchema), + "depth" => Ok(GeneratedField::Depth), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = ListUnnest; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.ListUnnest") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut index_in_input_schema__ = None; + let mut depth__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::IndexInInputSchema => { + if index_in_input_schema__.is_some() { + return Err(serde::de::Error::duplicate_field("indexInInputSchema")); + } + index_in_input_schema__ = + Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) + ; + } + GeneratedField::Depth => { + if depth__.is_some() { + return Err(serde::de::Error::duplicate_field("depth")); + } + depth__ = + Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) + ; + } + } + } + Ok(ListUnnest { + index_in_input_schema: index_in_input_schema__.unwrap_or_default(), + depth: depth__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.ListUnnest", FIELDS, GeneratedVisitor) + } +} impl serde::Serialize for ListingTableScanNode { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result @@ -19403,7 +19968,7 @@ impl serde::Serialize for UnnestExecNode { struct_ser.serialize_field("schema", v)?; } if !self.list_type_columns.is_empty() { - struct_ser.serialize_field("listTypeColumns", &self.list_type_columns.iter().map(ToString::to_string).collect::>())?; + struct_ser.serialize_field("listTypeColumns", &self.list_type_columns)?; } if !self.struct_type_columns.is_empty() { struct_ser.serialize_field("structTypeColumns", &self.struct_type_columns.iter().map(ToString::to_string).collect::>())?; @@ -19505,10 +20070,7 @@ impl<'de> serde::Deserialize<'de> for UnnestExecNode { if list_type_columns__.is_some() { return Err(serde::de::Error::duplicate_field("listTypeColumns")); } - list_type_columns__ = - Some(map_.next_value::>>()? - .into_iter().map(|x| x.0).collect()) - ; + list_type_columns__ = Some(map_.next_value()?); } GeneratedField::StructTypeColumns => { if struct_type_columns__.is_some() { @@ -19576,7 +20138,7 @@ impl serde::Serialize for UnnestNode { struct_ser.serialize_field("execColumns", &self.exec_columns)?; } if !self.list_type_columns.is_empty() { - struct_ser.serialize_field("listTypeColumns", &self.list_type_columns.iter().map(ToString::to_string).collect::>())?; + struct_ser.serialize_field("listTypeColumns", &self.list_type_columns)?; } if !self.struct_type_columns.is_empty() { struct_ser.serialize_field("structTypeColumns", &self.struct_type_columns.iter().map(ToString::to_string).collect::>())?; @@ -19694,10 +20256,7 @@ impl<'de> serde::Deserialize<'de> for UnnestNode { if list_type_columns__.is_some() { return Err(serde::de::Error::duplicate_field("listTypeColumns")); } - list_type_columns__ = - Some(map_.next_value::>>()? - .into_iter().map(|x| x.0).collect()) - ; + list_type_columns__ = Some(map_.next_value()?); } GeneratedField::StructTypeColumns => { if struct_type_columns__.is_some() { diff --git a/datafusion/proto/src/generated/prost.rs b/datafusion/proto/src/generated/prost.rs index 1d086a610ce4..bc732ed764ab 100644 --- a/datafusion/proto/src/generated/prost.rs +++ b/datafusion/proto/src/generated/prost.rs @@ -430,9 +430,9 @@ pub struct UnnestNode { #[prost(message, optional, boxed, tag = "1")] pub input: ::core::option::Option<::prost::alloc::boxed::Box>, #[prost(message, repeated, tag = "2")] - pub exec_columns: ::prost::alloc::vec::Vec, - #[prost(uint64, repeated, tag = "3")] - pub list_type_columns: ::prost::alloc::vec::Vec, + pub exec_columns: ::prost::alloc::vec::Vec, + #[prost(message, repeated, tag = "3")] + pub list_type_columns: ::prost::alloc::vec::Vec, #[prost(uint64, repeated, tag = "4")] pub struct_type_columns: ::prost::alloc::vec::Vec, #[prost(uint64, repeated, tag = "5")] @@ -443,6 +443,49 @@ pub struct UnnestNode { pub options: ::core::option::Option, } #[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct ColumnUnnestListItem { + #[prost(uint32, tag = "1")] + pub input_index: u32, + #[prost(message, optional, tag = "2")] + pub recursion: ::core::option::Option, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct ColumnUnnestListRecursions { + #[prost(message, repeated, tag = "2")] + pub recursions: ::prost::alloc::vec::Vec, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct ColumnUnnestListRecursion { + #[prost(message, optional, tag = "1")] + pub output_column: ::core::option::Option, + #[prost(uint32, tag = "2")] + pub depth: u32, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct ColumnUnnestExec { + #[prost(message, optional, tag = "1")] + pub column: ::core::option::Option, + #[prost(oneof = "column_unnest_exec::UnnestType", tags = "2, 3, 4")] + pub unnest_type: ::core::option::Option, +} +/// Nested message and enum types in `ColumnUnnestExec`. +pub mod column_unnest_exec { + #[allow(clippy::derive_partial_eq_without_eq)] + #[derive(Clone, PartialEq, ::prost::Oneof)] + pub enum UnnestType { + #[prost(message, tag = "2")] + List(super::ColumnUnnestListRecursions), + #[prost(message, tag = "3")] + Struct(super::super::datafusion_common::EmptyMessage), + #[prost(message, tag = "4")] + Inferred(super::super::datafusion_common::EmptyMessage), + } +} +#[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, Copy, PartialEq, ::prost::Message)] pub struct UnnestOptions { #[prost(bool, tag = "1")] @@ -1204,14 +1247,22 @@ pub struct UnnestExecNode { pub input: ::core::option::Option<::prost::alloc::boxed::Box>, #[prost(message, optional, tag = "2")] pub schema: ::core::option::Option, - #[prost(uint64, repeated, tag = "3")] - pub list_type_columns: ::prost::alloc::vec::Vec, + #[prost(message, repeated, tag = "3")] + pub list_type_columns: ::prost::alloc::vec::Vec, #[prost(uint64, repeated, tag = "4")] pub struct_type_columns: ::prost::alloc::vec::Vec, #[prost(message, optional, tag = "5")] pub options: ::core::option::Option, } #[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, Copy, PartialEq, ::prost::Message)] +pub struct ListUnnest { + #[prost(uint32, tag = "1")] + pub index_in_input_schema: u32, + #[prost(uint32, tag = "2")] + pub depth: u32, +} +#[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] pub struct PhysicalExtensionNode { #[prost(bytes = "vec", tag = "1")] diff --git a/datafusion/proto/src/logical_plan/mod.rs b/datafusion/proto/src/logical_plan/mod.rs index bf5394ec01de..5cd4267f426a 100644 --- a/datafusion/proto/src/logical_plan/mod.rs +++ b/datafusion/proto/src/logical_plan/mod.rs @@ -19,8 +19,12 @@ use std::collections::HashMap; use std::fmt::Debug; use std::sync::Arc; +use crate::protobuf::column_unnest_exec::UnnestType; use crate::protobuf::logical_plan_node::LogicalPlanType::CustomScan; -use crate::protobuf::{CustomTableScanNode, SortExprNodeCollection}; +use crate::protobuf::{ + ColumnUnnestExec, ColumnUnnestListItem, ColumnUnnestListRecursion, + ColumnUnnestListRecursions, CustomTableScanNode, SortExprNodeCollection, +}; use crate::{ convert_required, into_required, protobuf::{ @@ -65,7 +69,8 @@ use datafusion_expr::{ DistinctOn, DropView, Expr, LogicalPlan, LogicalPlanBuilder, ScalarUDF, SortExpr, WindowUDF, }; -use datafusion_expr::{AggregateUDF, Unnest}; +use datafusion_expr::{AggregateUDF, ColumnUnnestList, ColumnUnnestType, Unnest}; +use datafusion_proto_common::EmptyMessage; use self::to_proto::{serialize_expr, serialize_exprs}; use crate::logical_plan::to_proto::serialize_sorts; @@ -860,11 +865,50 @@ impl AsLogicalPlan for LogicalPlanNode { into_logical_plan!(unnest.input, ctx, extension_codec)?; Ok(datafusion_expr::LogicalPlan::Unnest(Unnest { input: Arc::new(input), - exec_columns: unnest.exec_columns.iter().map(|c| c.into()).collect(), + exec_columns: unnest + .exec_columns + .iter() + .map(|c| { + ( + c.column.as_ref().unwrap().to_owned().into(), + match c.unnest_type.as_ref().unwrap() { + UnnestType::Inferred(_) => ColumnUnnestType::Inferred, + UnnestType::Struct(_) => ColumnUnnestType::Struct, + UnnestType::List(l) => ColumnUnnestType::List( + l.recursions + .iter() + .map(|ul| ColumnUnnestList { + output_column: ul + .output_column + .as_ref() + .unwrap() + .to_owned() + .into(), + depth: ul.depth as usize, + }) + .collect(), + ), + }, + ) + }) + .collect(), list_type_columns: unnest .list_type_columns .iter() - .map(|c| *c as usize) + .map(|c| { + let recursion_item = c.recursion.as_ref().unwrap(); + ( + c.input_index as _, + ColumnUnnestList { + output_column: recursion_item + .output_column + .as_ref() + .unwrap() + .into(), + depth: recursion_item.depth as _, + }, + ) + }) .collect(), struct_type_columns: unnest .struct_type_columns @@ -1536,15 +1580,50 @@ impl AsLogicalPlan for LogicalPlanNode { input, extension_codec, )?; + let todo = list_type_columns + .iter() + .map(|(index, ul)| ColumnUnnestListItem { + input_index: *index as _, + recursion: Some(ColumnUnnestListRecursion { + output_column: Some(ul.output_column.to_owned().into()), + depth: ul.depth as _, + }), + }) + .collect(); Ok(protobuf::LogicalPlanNode { logical_plan_type: Some(LogicalPlanType::Unnest(Box::new( protobuf::UnnestNode { input: Some(Box::new(input)), - exec_columns: exec_columns.iter().map(|c| c.into()).collect(), - list_type_columns: list_type_columns + exec_columns: exec_columns .iter() - .map(|c| *c as u64) + .map(|(col, unnesting)| ColumnUnnestExec { + column: Some(col.into()), + unnest_type: Some(match unnesting { + ColumnUnnestType::Inferred => { + UnnestType::Inferred(EmptyMessage {}) + } + ColumnUnnestType::Struct => { + UnnestType::Struct(EmptyMessage {}) + } + ColumnUnnestType::List(list) => { + UnnestType::List(ColumnUnnestListRecursions { + recursions: list + .iter() + .map(|ul| ColumnUnnestListRecursion { + output_column: Some( + ul.output_column + .to_owned() + .into(), + ), + depth: ul.depth as _, + }) + .collect(), + }) + } + }), + }) .collect(), + list_type_columns: todo, struct_type_columns: struct_type_columns .iter() .map(|c| *c as u64) diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 74b6073a415e..6e485c4f0cd0 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -58,7 +58,7 @@ use datafusion::physical_plan::repartition::RepartitionExec; use datafusion::physical_plan::sorts::sort::SortExec; use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion::physical_plan::union::{InterleaveExec, UnionExec}; -use datafusion::physical_plan::unnest::UnnestExec; +use datafusion::physical_plan::unnest::{ListUnnest, UnnestExec}; use datafusion::physical_plan::windows::{BoundedWindowAggExec, WindowAggExec}; use datafusion::physical_plan::{ ExecutionPlan, InputOrderMode, PhysicalExpr, WindowExpr, @@ -78,7 +78,9 @@ use crate::physical_plan::to_proto::{ use crate::protobuf::physical_aggregate_expr_node::AggregateFunction; use crate::protobuf::physical_expr_node::ExprType; use crate::protobuf::physical_plan_node::PhysicalPlanType; -use crate::protobuf::{self, proto_error, window_agg_exec_node}; +use crate::protobuf::{ + self, proto_error, window_agg_exec_node, ListUnnest as ProtoListUnnest, +}; use crate::{convert_required, into_required}; use self::from_proto::parse_protobuf_partitioning; @@ -1108,7 +1110,14 @@ impl AsExecutionPlan for protobuf::PhysicalPlanNode { Ok(Arc::new(UnnestExec::new( input, - unnest.list_type_columns.iter().map(|c| *c as _).collect(), + unnest + .list_type_columns + .iter() + .map(|c| ListUnnest { + index_in_input_schema: c.index_in_input_schema as _, + depth: c.depth as _, + }) + .collect(), unnest.struct_type_columns.iter().map(|c| *c as _).collect(), Arc::new(convert_required!(unnest.schema)?), into_required!(unnest.options)?, @@ -1984,7 +1993,10 @@ impl AsExecutionPlan for protobuf::PhysicalPlanNode { list_type_columns: exec .list_column_indices() .iter() - .map(|c| *c as _) + .map(|c| ProtoListUnnest { + index_in_input_schema: c.index_in_input_schema as _, + depth: c.depth as _, + }) .collect(), struct_type_columns: exec .struct_column_indices() diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index 97d65ba87197..514ba21fd4da 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -72,7 +72,7 @@ use datafusion::physical_plan::repartition::RepartitionExec; use datafusion::physical_plan::sorts::sort::SortExec; use datafusion::physical_plan::udaf::AggregateFunctionExpr; use datafusion::physical_plan::union::{InterleaveExec, UnionExec}; -use datafusion::physical_plan::unnest::UnnestExec; +use datafusion::physical_plan::unnest::{ListUnnest, UnnestExec}; use datafusion::physical_plan::windows::{ BuiltInWindowExpr, PlainAggregateWindowExpr, WindowAggExec, }; @@ -1373,6 +1373,25 @@ fn roundtrip_unnest() -> Result<()> { Arc::new(Schema::new(vec![fa, fb0, fc1, fc2, fd0, fe1, fe2, fe3])); let input = Arc::new(EmptyExec::new(input_schema)); let options = UnnestOptions::default(); - let unnest = UnnestExec::new(input, vec![1, 3], vec![2, 4], output_schema, options); + let unnest = UnnestExec::new( + input, + vec![ + ListUnnest { + index_in_input_schema: 1, + depth: 1, + }, + ListUnnest { + index_in_input_schema: 1, + depth: 2, + }, + ListUnnest { + index_in_input_schema: 3, + depth: 2, + }, + ], + vec![2, 4], + output_schema, + options, + ); roundtrip_test(Arc::new(unnest)) } From 784bf3ef47eea8e7675bd39175699cb48457a614 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Tue, 10 Sep 2024 21:09:04 +0200 Subject: [PATCH 40/56] fix dataframe --- datafusion/core/tests/dataframe/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs index 19ce9294cfad..21a3229a9617 100644 --- a/datafusion/core/tests/dataframe/mod.rs +++ b/datafusion/core/tests/dataframe/mod.rs @@ -1388,7 +1388,7 @@ async fn unnest_with_redundant_columns() -> Result<()> { let optimized_plan = df.clone().into_optimized_plan()?; let expected = vec![ "Projection: shapes.shape_id [shape_id:UInt32]", - " Unnest: lists[shape_id2] structs[] [shape_id:UInt32, shape_id2:UInt32;N]", + " Unnest: lists[shape_id2|depth=1] structs[] [shape_id:UInt32, shape_id2:UInt32;N]", " Aggregate: groupBy=[[shapes.shape_id]], aggr=[[array_agg(shapes.shape_id) AS shape_id2]] [shape_id:UInt32, shape_id2:List(Field { name: \"item\", data_type: UInt32, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} });N]", " TableScan: shapes projection=[shape_id] [shape_id:UInt32]", ]; From 4d3f5081542f7d24426b9512e6114062276e2baa Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Tue, 10 Sep 2024 21:20:11 +0200 Subject: [PATCH 41/56] fix clippy --- datafusion/expr/src/logical_plan/builder.rs | 16 ++++++++-------- datafusion/sql/src/select.rs | 12 ++++++------ 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index 40f8664da5e9..9d241f71e34b 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -1692,22 +1692,22 @@ pub fn get_unnested_columns( /// For example: /// Input schema as /// -/// +---------------------+-----------+ -/// | col1 | col2 | -/// +---------------------+-----------+ -/// | Struct(INT64,INT32) | [[Int64]] | -/// +---------------------+-----------+ +/// +---------------------+-------------------+ +/// | col1 | col2 | +/// +---------------------+-------------------+ +/// | Struct(INT64,INT32) | List(List(Int64)) | +/// +---------------------+-------------------+ /// /// Then unnesting columns with: -/// - [col1,Struct] -/// - [col2,List([depth=1,depth=2])] +/// - (col1,Struct) +/// - (col2,List(\[depth=1,depth=2\])) /// /// will generate a new schema as /// /// +---------+---------+---------------------+---------------------+ /// | col1.c0 | col1.c1 | unnest_col2_depth_1 | unnest_col2_depth_2 | /// +---------+---------+---------------------+---------------------+ -/// | Int64 | Int32 | [Int64] | Int64 | +/// | Int64 | Int32 | List(Int64) | Int64 | /// +---------+---------+---------------------+---------------------+ pub fn unnest_with_options( input: LogicalPlan, diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs index 0d487684db7b..b0ca7bf5d110 100644 --- a/datafusion/sql/src/select.rs +++ b/datafusion/sql/src/select.rs @@ -378,11 +378,9 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { } } - /// Try converting Unnest(Expr) of group by to Unnest/Projection + /// 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 instance: - /// - select unnest(arr) as c1, unnest(arr) + unnest(arr) as c2 group by c1 - /// We need both for this funciton argument to check how select exprs has been transformed + /// Select exprs can be different from agg exprs, for example: fn try_process_group_by_unnest( &self, agg: Aggregate, @@ -429,7 +427,6 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { if unnest_columns.is_empty() { break; } else { - let columns = unnest_columns.into_iter().map(|col| col.into()).collect(); let unnest_options = UnnestOptions::new().with_preserve_nulls(false); let mut projection_exprs = match &aggr_expr_using_columns { @@ -454,7 +451,10 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { intermediate_plan = LogicalPlanBuilder::from(intermediate_plan) .project(projection_exprs)? - .unnest_columns_recursive_with_options(columns, unnest_options)? + .unnest_columns_recursive_with_options( + unnest_columns, + unnest_options, + )? .build()?; intermediate_select_exprs = outer_projection_exprs; From 926efd53f08f60ee672cac20777605ab0acd5aa9 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Tue, 10 Sep 2024 21:25:22 +0200 Subject: [PATCH 42/56] cargo fmt --- datafusion/expr/src/logical_plan/plan.rs | 5 +++-- datafusion/sql/src/utils.rs | 5 ++++- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index 0f1e28227196..eea12b84a8a0 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -1973,7 +1973,7 @@ impl LogicalPlan { let input_columns = plan.schema().columns(); let list_type_columns = list_col_indices .iter() - .map(|(i,unnest_info)| + .map(|(i,unnest_info)| format!("{}|depth={}", &input_columns[*i].to_string(), unnest_info.depth)) .collect::>(); @@ -3066,7 +3066,8 @@ impl fmt::Display for ColumnUnnestType { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { ColumnUnnestType::List(lists) => { - let list_strs: Vec = lists.iter().map(|list| list.to_string()).collect(); + let list_strs: Vec = + lists.iter().map(|list| list.to_string()).collect(); write!(f, "List([{}])", list_strs.join(", ")) } ColumnUnnestType::Struct => write!(f, "Struct"), diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index e5199059f2dc..c7e218f66592 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -33,7 +33,10 @@ use datafusion_common::{ use datafusion_expr::builder::get_struct_unnested_columns; use datafusion_expr::expr::{Alias, GroupingSet, Unnest, WindowFunction}; use datafusion_expr::utils::{expr_as_column_expr, find_column_exprs}; -use datafusion_expr::{col,expr_vec_fmt, Expr, ExprSchemable, LogicalPlan,ColumnUnnestList, ColumnUnnestType}; +use datafusion_expr::{ + col, expr_vec_fmt, ColumnUnnestList, ColumnUnnestType, Expr, ExprSchemable, + LogicalPlan, +}; use sqlparser::ast::{Ident, Value}; /// Make a best-effort attempt at resolving all columns in the expression tree From bf4068812d994f82c3b00119afa8df155a211b3d Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Tue, 10 Sep 2024 22:05:24 +0200 Subject: [PATCH 43/56] fix some test --- datafusion/sql/src/utils.rs | 212 +++++++----------------------------- 1 file changed, 42 insertions(+), 170 deletions(-) diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index c7e218f66592..dba972ef6d1e 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -621,38 +621,17 @@ mod tests { use crate::utils::{ group_bottom_most_consecutive_unnest, resolve_positions_to_exprs, }; - fn column_unnests_eq(l: Vec<&str>, r: &[(Column, ColumnUnnestType)]) { - let formatted: Vec = + fn column_unnests_eq(l: Vec<(&str, &str)>, r: &[(Column, ColumnUnnestType)]) { + let r_formatted: Vec = r.iter().map(|i| format!("{}|{}", i.0, i.1)).collect(); - assert_eq!(l, formatted) + let l_formatted: Vec = + l.iter().map(|i| format!("{}|{}", i.0, i.1)).collect(); + assert_eq!(l_formatted, r_formatted); } #[test] - fn test_transform_bottom_unnest_recursive_memoization_struct() -> Result<()> { - let three_d_dtype = ArrowDataType::List(Arc::new(Field::new( - "2d_col", - ArrowDataType::List(Arc::new(Field::new( - "elements", - ArrowDataType::Int64, - true, - ))), - true, - ))); + fn test_transform_bottom_unnest_recursive() -> Result<()> { let schema = Schema::new(vec![ - // list[struct(3d_data)] [([[1,2,3]])] - Field::new( - "struct_arr_col", - ArrowDataType::List(Arc::new(Field::new( - "struct", - ArrowDataType::Struct(Fields::from(vec![Field::new( - "field1", - three_d_dtype, - true, - )])), - true, - ))), - true, - ), Field::new( "3d_col", ArrowDataType::List(Arc::new(Field::new( @@ -693,11 +672,18 @@ mod tests { assert_eq!( transformed_exprs, vec![col("unnest_placeholder(3d_col,depth=2)") - .add(col("unnest_placeholder(3d_col,depth=2)")) + .alias("UNNEST(UNNEST(3d_col))") + .add( + col("unnest_placeholder(3d_col,depth=2)") + .alias("UNNEST(UNNEST(3d_col))") + ) .add(col("i64_col"))] ); column_unnests_eq( - vec!["unnest_placeholder(3d_col)"], + vec![( + "unnest_placeholder(3d_col)", + "List([unnest_placeholder(3d_col,depth=2)|depth=2])", + )], &unnest_placeholder_columns, ); @@ -722,138 +708,15 @@ mod tests { assert_eq!( transformed_exprs, - vec![col("unnest_placeholder(3d_col,depth=1)").alias("2d_col")] - ); - column_unnests_eq( - vec!["unnest_placeholder(3d_col)"], - &mut unnest_placeholder_columns, - ); - // still reference struct_col in original schema but with alias, - // to avoid colliding with the projection on the column itself if any - assert_eq!( - inner_projection_exprs, - vec![ - col("3d_col").alias("unnest_placeholder(3d_col)"), - col("i64_col") - ] - ); - - // unnest(unnset(unnest(struct_arr_col)['field1'])) as fully_unnested_struct_arr - let original_expr_3 = - unnest(unnest(unnest(col("struct_arr_col")).field("field1"))) - .alias("fully_unnested_struct_arr"); - let transformed_exprs = group_bottom_most_consecutive_unnest( - &input, - &mut unnest_placeholder_columns, - &mut inner_projection_exprs, - &original_expr_3, - )?; - - assert_eq!( - transformed_exprs, - vec![unnest(unnest( - col("unnest_placeholder(struct_arr_col,depth=1)").field("field1") - )) - .alias("fully_unnested_struct_arr")] - ); - - column_unnests_eq( - vec![ - "unnest_placeholder(3d_col)", - "unnest_placeholder(struct_arr_col)", - ], - &mut unnest_placeholder_columns, - ); - // still reference struct_col in original schema but with alias, - // to avoid colliding with the projection on the column itself if any - assert_eq!( - inner_projection_exprs, - vec![ - col("3d_col").alias("unnest_placeholder(3d_col)"), - col("i64_col"), - col("struct_arr_col").alias("unnest_placeholder(struct_arr_col)") - ] - ); - - Ok(()) - } - - #[test] - fn test_transform_bottom_unnest_recursive_memoization() -> Result<()> { - let schema = Schema::new(vec![ - Field::new( - "3d_col", - ArrowDataType::List(Arc::new(Field::new( - "2d_col", - ArrowDataType::List(Arc::new(Field::new( - "elements", - ArrowDataType::Int64, - true, - ))), - true, - ))), - true, - ), - Field::new("i64_col", ArrowDataType::Int64, true), - ]); - - let dfschema = DFSchema::try_from(schema)?; - - let input = LogicalPlan::EmptyRelation(EmptyRelation { - produce_one_row: false, - schema: Arc::new(dfschema), - }); - - let mut unnest_placeholder_columns = vec![]; - let mut inner_projection_exprs = vec![]; - - // unnest(unnest(3d_col)) + unnest(unnest(3d_col)) - let original_expr = unnest(unnest(col("3d_col"))) - .add(unnest(unnest(col("3d_col")))) - .add(col("i64_col")); - let transformed_exprs = group_bottom_most_consecutive_unnest( - &input, - &mut unnest_placeholder_columns, - &mut inner_projection_exprs, - &original_expr, - )?; - // only the bottom most unnest exprs are transformed - assert_eq!( - transformed_exprs, - vec![col("unnest_placeholder(3d_col,depth=2)") - .add(col("unnest_placeholder(3d_col,depth=2)")) - .add(col("i64_col"))] - ); - column_unnests_eq( - vec!["unnest_placeholder(3d_col)|List([unnest_placeholder(3d_col,depth=2)|depth=2])"], - &unnest_placeholder_columns, - ); - - // still reference struct_col in original schema but with alias, - // to avoid colliding with the projection on the column itself if any - assert_eq!( - inner_projection_exprs, vec![ - col("3d_col").alias("unnest_placeholder(3d_col)"), - col("i64_col") + (col("unnest_placeholder(3d_col,depth=1)").alias("UNNEST(3d_col)")) + .alias("2d_col") ] ); - - // unnest(3d_col) as 2d_col - let original_expr_2 = unnest(col("3d_col")).alias("2d_col"); - let transformed_exprs = group_bottom_most_consecutive_unnest( - &input, - &mut unnest_placeholder_columns, - &mut inner_projection_exprs, - &original_expr_2, - )?; - - assert_eq!( - transformed_exprs, - vec![col("unnest_placeholder(3d_col,depth=1)").alias("2d_col")] - ); column_unnests_eq( - vec!["unnest_placeholder(3d_col)|List([unnest_placeholder(3d_col,depth=2)|depth=2, unnest_placeholder(3d_col,depth=1)|depth=1])"], + vec![("unnest_placeholder(3d_col)", + "List([unnest_placeholder(3d_col,depth=2)|depth=2, unnest_placeholder(3d_col,depth=1)|depth=1])"), + ], &unnest_placeholder_columns, ); // still reference struct_col in original schema but with alias, @@ -913,19 +776,19 @@ mod tests { assert_eq!( transformed_exprs, vec![ - col("UNNEST(struct_col).field1"), - col("UNNEST(struct_col).field2"), + col("unnest_placeholder(struct_col).field1"), + col("unnest_placeholder(struct_col).field2"), ] ); column_unnests_eq( - vec!["unnest_placeholder(struct_col)"], - &mut unnest_placeholder_columns, + vec![("unnest_placeholder(struct_col)", "Struct")], + &unnest_placeholder_columns, ); // still reference struct_col in original schema but with alias, // to avoid colliding with the projection on the column itself if any assert_eq!( inner_projection_exprs, - vec![col("struct_col").alias("UNNEST(struct_col)"),] + vec![col("struct_col").alias("unnest_placeholder(struct_col)"),] ); // unnest(array_col) + 1 @@ -938,15 +801,20 @@ mod tests { )?; column_unnests_eq( vec![ - "unnest_placeholder(struct_col)", - "unnest_placeholder(array_col)", + ("unnest_placeholder(struct_col)", "Struct"), + ( + "unnest_placeholder(array_col)", + "List([unnest_placeholder(array_col,depth=1)|depth=1])", + ), ], &mut unnest_placeholder_columns, ); // only transform the unnest children assert_eq!( transformed_exprs, - vec![col("UNNEST(array_col)").add(lit(1i64))] + vec![col("unnest_placeholder(array_col,depth=1)") + .alias("UNNEST(array_col)") + .add(lit(1i64))] ); // keep appending to the current vector @@ -955,8 +823,8 @@ mod tests { assert_eq!( inner_projection_exprs, vec![ - col("struct_col").alias("UNNEST(struct_col)"), - col("array_col").alias("UNNEST(array_col)") + col("struct_col").alias("unnest_placeholder(struct_col)"), + col("array_col").alias("unnest_placeholder(array_col)") ] ); @@ -1003,11 +871,15 @@ mod tests { // Only the inner most/ bottom most unnest is transformed assert_eq!( transformed_exprs, - vec![unnest(col("UNNEST(struct_col[matrix])"))] + vec![unnest(col( + "unnest_placeholder(struct_col[matrix],depth=2)" + ))] ); + // TODO: add a test case where + // unnest -> field access -> unnest column_unnests_eq( - vec!["unnest_placeholder(struct_col[matrix])"], + vec![("unnest_placeholder(struct_col[matrix])", "Struct")], &mut unnest_placeholder_columns, ); @@ -1015,7 +887,7 @@ mod tests { inner_projection_exprs, vec![col("struct_col") .field("matrix") - .alias("UNNEST(struct_col[matrix])"),] + .alias("unnest_placeholder(struct_col[matrix])"),] ); Ok(()) From 1e5442228ad5ec450aca70181b79048b89ea2dee Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Tue, 10 Sep 2024 22:06:58 +0200 Subject: [PATCH 44/56] fix all test --- datafusion/sql/src/utils.rs | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index dba972ef6d1e..af06c6a91f17 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -871,15 +871,17 @@ mod tests { // Only the inner most/ bottom most unnest is transformed assert_eq!( transformed_exprs, - vec![unnest(col( - "unnest_placeholder(struct_col[matrix],depth=2)" - ))] + vec![col("unnest_placeholder(struct_col[matrix],depth=2)") + .alias("UNNEST(UNNEST(struct_col[matrix]))")] ); // TODO: add a test case where // unnest -> field access -> unnest column_unnests_eq( - vec![("unnest_placeholder(struct_col[matrix])", "Struct")], + vec![( + "unnest_placeholder(struct_col[matrix])", + "List([unnest_placeholder(struct_col[matrix],depth=2)|depth=2])", + )], &mut unnest_placeholder_columns, ); From abfd4fc89e275d8602989a52f1ea28107c09f414 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Tue, 10 Sep 2024 22:07:42 +0200 Subject: [PATCH 45/56] fix unnest in join --- datafusion/sqllogictest/test_files/joins.slt | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/datafusion/sqllogictest/test_files/joins.slt b/datafusion/sqllogictest/test_files/joins.slt index 7d0262952b31..b6ee3a291532 100644 --- a/datafusion/sqllogictest/test_files/joins.slt +++ b/datafusion/sqllogictest/test_files/joins.slt @@ -4055,9 +4055,9 @@ logical_plan 03)----TableScan: join_t1 projection=[t1_id, t1_name] 04)--SubqueryAlias: series 05)----Subquery: -06)------Projection: UNNEST(generate_series(Int64(1),outer_ref(t1.t1_int))) AS i -07)--------Unnest: lists[UNNEST(generate_series(Int64(1),outer_ref(t1.t1_int)))] structs[] -08)----------Projection: generate_series(Int64(1), CAST(outer_ref(t1.t1_int) AS Int64)) AS UNNEST(generate_series(Int64(1),outer_ref(t1.t1_int))) +06)------Projection: unnest_placeholder(generate_series(Int64(1),outer_ref(t1.t1_int)),depth=1) AS i +07)--------Unnest: lists[unnest_placeholder(generate_series(Int64(1),outer_ref(t1.t1_int)))|depth=1] structs[] +08)----------Projection: generate_series(Int64(1), CAST(outer_ref(t1.t1_int) AS Int64)) AS unnest_placeholder(generate_series(Int64(1),outer_ref(t1.t1_int))) 09)------------EmptyRelation From eb4cb446303d4094b372fba3185639bea59f5969 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Wed, 11 Sep 2024 20:38:19 +0200 Subject: [PATCH 46/56] fix doc and tests --- datafusion/physical-plan/src/unnest.rs | 4 +++- datafusion/sql/tests/cases/plan_to_sql.rs | 10 +++++----- datafusion/sqllogictest/test_files/joins.slt | 6 +++--- 3 files changed, 11 insertions(+), 9 deletions(-) diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index c53c93b92512..ae1870a9345c 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -411,7 +411,6 @@ pub struct ListUnnest { /// Recursion happens for the highest level first /// Demonstatring with examples: /// -/// /// Set "A" as a 3-dimension columns and "B" as an array (1-dimension) /// Query: select unnest(A, max_depth:=3), unnest(A,max_depth:=2), unnest(B, max_depth:=1) from temp; /// Let's given these projection names P1,P2,P3 respectively @@ -420,12 +419,15 @@ pub struct ListUnnest { /// This is needed, even if the same column is being unnested for different recursion levels /// /// This function is called with the descending order of recursion +/// /// Depth = 3 /// - P1(3-dimension) unnest into temp column temp_P1(2_dimension) /// - A(3-dimension) having indice repeated by the unnesting above +/// /// Depth = 2 /// - temp_P1(2-dimension) unnest into temp column temp_P1(1-dimension) /// - A(3-dimension) unnest into temp column temp_P2(2-dimension) +/// /// Depth = 1 /// - temp_P1(1-dimension) unnest into P1 /// - temp_P2(2-dimension) unnest into P2 diff --git a/datafusion/sql/tests/cases/plan_to_sql.rs b/datafusion/sql/tests/cases/plan_to_sql.rs index fa95d05c3275..89f17f627b06 100644 --- a/datafusion/sql/tests/cases/plan_to_sql.rs +++ b/datafusion/sql/tests/cases/plan_to_sql.rs @@ -466,11 +466,11 @@ fn test_unnest_logical_plan() -> Result<()> { }; let sql_to_rel = SqlToRel::new(&context); let plan = sql_to_rel.sql_statement_to_plan(statement).unwrap(); - - let expected = "Projection: UNNEST(unnest_table.struct_col).field1, UNNEST(unnest_table.struct_col).field2, UNNEST(unnest_table.array_col), unnest_table.struct_col, unnest_table.array_col\ - \n Unnest: lists[UNNEST(unnest_table.array_col)] structs[UNNEST(unnest_table.struct_col)]\ - \n Projection: unnest_table.struct_col AS UNNEST(unnest_table.struct_col), unnest_table.array_col AS UNNEST(unnest_table.array_col), unnest_table.struct_col, unnest_table.array_col\ - \n TableScan: unnest_table"; + let expected = r#" +Projection: unnest_placeholder(unnest_table.struct_col).field1, unnest_placeholder(unnest_table.struct_col).field2, unnest_placeholder(unnest_table.array_col,depth=1) AS UNNEST(unnest_table.array_col), unnest_table.struct_col, unnest_table.array_col + Unnest: lists[unnest_placeholder(unnest_table.array_col)|depth=1] structs[unnest_placeholder(unnest_table.struct_col)] + Projection: unnest_table.struct_col AS unnest_placeholder(unnest_table.struct_col), unnest_table.array_col AS unnest_placeholder(unnest_table.array_col), unnest_table.struct_col, unnest_table.array_col + TableScan: unnest_table"#.trim_start(); assert_eq!(format!("{plan}"), expected); diff --git a/datafusion/sqllogictest/test_files/joins.slt b/datafusion/sqllogictest/test_files/joins.slt index b6ee3a291532..6196087dd530 100644 --- a/datafusion/sqllogictest/test_files/joins.slt +++ b/datafusion/sqllogictest/test_files/joins.slt @@ -4077,9 +4077,9 @@ logical_plan 03)----TableScan: join_t1 projection=[t1_id, t1_name] 04)--SubqueryAlias: series 05)----Subquery: -06)------Projection: UNNEST(generate_series(Int64(1),outer_ref(t2.t1_int))) AS i -07)--------Unnest: lists[UNNEST(generate_series(Int64(1),outer_ref(t2.t1_int)))] structs[] -08)----------Projection: generate_series(Int64(1), CAST(outer_ref(t2.t1_int) AS Int64)) AS UNNEST(generate_series(Int64(1),outer_ref(t2.t1_int))) +06)------Projection: unnest_placeholder(generate_series(Int64(1),outer_ref(t2.t1_int)),depth=1) AS i +07)--------Unnest: lists[unnest_placeholder(generate_series(Int64(1),outer_ref(t2.t1_int)))|depth=1] structs[] +08)----------Projection: generate_series(Int64(1), CAST(outer_ref(t2.t1_int) AS Int64)) AS unnest_placeholder(generate_series(Int64(1),outer_ref(t2.t1_int))) 09)------------EmptyRelation From 7aee552286c4adce26d5695262b3afdb58e40731 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Wed, 11 Sep 2024 21:48:50 +0200 Subject: [PATCH 47/56] chore: better doc --- datafusion/sql/src/utils.rs | 47 +++++++++++++++++++++++++++++++++---- 1 file changed, 42 insertions(+), 5 deletions(-) diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index af06c6a91f17..55eb867a0c53 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -338,9 +338,12 @@ struct RecursiveUnnestRewriter<'a> { transformed_root_exprs: Option>, } impl<'a> RecursiveUnnestRewriter<'a> { - // given a sequence of [None,Unnest,Unnest,None,None] - // returns [Unnest,Unnest] - // The first items is the inner most unnest + /// This struct stores the history of expr + /// during its tree-traversal with a notation of + /// \[None,**Unnest(exprA)**,**Unnest(exprB)**,None,None\] + /// then this function will returns \[**Unnest(exprA)**,**Unnest(exprB)**\] + /// + /// The first item will be the inner most expr fn get_latest_consecutive_unnest(&self) -> Vec { self.consecutive_unnest .iter() @@ -450,6 +453,10 @@ impl<'a> RecursiveUnnestRewriter<'a> { impl<'a> TreeNodeRewriter for RecursiveUnnestRewriter<'a> { type Node = Expr; + /// This downward traversal needs to keep track of: + /// - Whether or not some unnest expr has been visited from the top util the current node + /// - If some unnest expr has been visited, maintain a stack of such information, this + /// is used to detect if some recursive unnest expr exists (e.g **unnest(unnest(unnest(3d column))))** fn f_down(&mut self, expr: Expr) -> Result> { if let Expr::Unnest(ref unnest_expr) = expr { let (data_type, _) = @@ -478,6 +485,36 @@ impl<'a> TreeNodeRewriter for RecursiveUnnestRewriter<'a> { } } + /// The rewriting only happens when the traversal has reached the top-most unnest expr + /// within a sequence of consecutive unnest exprs. + /// node, for example given a stack of expr + /// + /// For example an expr of **unnest(unnest(column1)) + unnest(unnest(unnest(column2)))** + /// ```text + /// ┌──────────────────┐ + /// │ binaryexpr │ + /// │ │ + /// └──────────────────┘ + /// f_down / / │ │ + /// / / f_up │ │ + /// / / f_down│ │f_up + /// unnest │ │ + /// │ │ + /// f_down / / f_up(rewriting) │ │ + /// / / + /// / / unnest + /// unnest + /// f_down / / f_up(rewriting) + /// f_down / /f_up / / + /// / / / / + /// / / unnest + /// column1 + /// f_down / /f_up + /// / / + /// / / + /// column2 + /// ``` + /// fn f_up(&mut self, expr: Expr) -> Result> { if let Expr::Unnest(ref traversing_unnest) = expr { if traversing_unnest == self.top_most_unnest.as_ref().unwrap() { @@ -807,7 +844,7 @@ mod tests { "List([unnest_placeholder(array_col,depth=1)|depth=1])", ), ], - &mut unnest_placeholder_columns, + &unnest_placeholder_columns, ); // only transform the unnest children assert_eq!( @@ -882,7 +919,7 @@ mod tests { "unnest_placeholder(struct_col[matrix])", "List([unnest_placeholder(struct_col[matrix],depth=2)|depth=2])", )], - &mut unnest_placeholder_columns, + &unnest_placeholder_columns, ); assert_eq!( From df4fee3f780410d29f079b733b5accd98b7bc604 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Wed, 11 Sep 2024 22:30:11 +0200 Subject: [PATCH 48/56] better doc --- datafusion/physical-plan/src/unnest.rs | 131 ++++++++++++++----------- datafusion/sql/src/utils.rs | 2 +- 2 files changed, 77 insertions(+), 56 deletions(-) diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index ae1870a9345c..92c0e498d000 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -366,76 +366,47 @@ pub struct ListUnnest { pub depth: usize, } -/// Note: unnest has a big difference in behavior between Postgres and DuckDB -/// Take this example -/// 1.Postgres -/// ```ignored -/// create table temp ( -/// i integer[][][], j integer[] -/// ) -/// insert into temp values ('{{{1,2},{3,4}},{{5,6},{7,8}}}', '{1,2}'); -/// select unnest(i), unnest(j) from temp; -/// ``` +/// This function is used to execute the unnesting on multiple columns all at once, but +/// one level at a time, and is called n times, where n is the highest recursion level among +/// the unnest exprs in the query. /// -/// Result -/// 1 1 -/// 2 2 -/// 3 -/// 4 -/// 5 -/// 6 -/// 7 -/// 8 -/// 2. DuckDB -/// ```ignore -/// create table temp (i integer[][][], j integer[]); -/// insert into temp values ([[[1,2],[3,4]],[[5,6],[7,8]]], [1,2]); -/// select unnest(i,recursive:=true), unnest(j,recursive:=true) from temp; +/// For example giving the following query: +/// ```sql +/// select unnest(colA, max_depth:=3) as P1, unnest(colA,max_depth:=2) as P2, unnest(colB, max_depth:=1) as P3 from temp; /// ``` -/// Result: -/// ┌────────────────────────────────────────────────┬────────────────────────────────────────────────┐ -/// │ unnest(i, "recursive" := CAST('t' AS BOOLEAN)) │ unnest(j, "recursive" := CAST('t' AS BOOLEAN)) │ -/// │ int32 │ int32 │ -/// ├────────────────────────────────────────────────┼────────────────────────────────────────────────┤ -/// │ 1 │ 1 │ -/// │ 2 │ 2 │ -/// │ 3 │ 1 │ -/// │ 4 │ 2 │ -/// │ 5 │ 1 │ -/// │ 6 │ 2 │ -/// │ 7 │ 1 │ -/// │ 8 │ 2 │ -/// └────────────────────────────────────────────────┴────────────────────────────────────────────────┘ -/// The following implementation refer to DuckDB's implementation -/// -/// Recursion happens for the highest level first -/// Demonstatring with examples: +/// Then the total times this function being called is 3 /// -/// Set "A" as a 3-dimension columns and "B" as an array (1-dimension) -/// Query: select unnest(A, max_depth:=3), unnest(A,max_depth:=2), unnest(B, max_depth:=1) from temp; -/// Let's given these projection names P1,P2,P3 respectively +/// It needs to be aware of which level the current unnesting is, because if there exists +/// multiple unnesting on the same column, but with different recursion levels, say +/// **unnest(colA, max_depth:=3)** and **unnest(colA, max_depth:=2)**, then the unnesting +/// of expr **unnest(colA, max_depth:=3)** will start at level 3, while unnesting for expr +/// **unnest(colA, max_depth:=2)** has to start at level 2 /// -/// Each combination of (column,depth) result in an entry in temp_batch -/// This is needed, even if the same column is being unnested for different recursion levels -/// -/// This function is called with the descending order of recursion +/// Set *colA* as a 3-dimension columns and *colB* as an array (1-dimension). As stated, +/// this function is called with the descending order of recursion depth /// /// Depth = 3 -/// - P1(3-dimension) unnest into temp column temp_P1(2_dimension) -/// - A(3-dimension) having indice repeated by the unnesting above +/// - colA(3-dimension) unnest into temp column temp_P1(2_dimension) (unnesting of P1 starts +/// from this level) +/// - colA(3-dimension) having indices repeated by the unnesting operation above +/// - colB(1-dimension) having indices repeated by the unnesting operation above /// /// Depth = 2 /// - temp_P1(2-dimension) unnest into temp column temp_P1(1-dimension) -/// - A(3-dimension) unnest into temp column temp_P2(2-dimension) +/// - colA(3-dimension) unnest into temp column temp_P2(2-dimension) (unnesting of P2 starts +/// from this level) +/// - colB(1-dimension) having indices repeated by the unnesting operation above /// /// Depth = 1 /// - temp_P1(1-dimension) unnest into P1 /// - temp_P2(2-dimension) unnest into P2 -/// - B(1-dimension) unnest into P3 +/// - colB(1-dimension) unnest into P3 (unnesting of P3 starts from this level) /// /// The returned array will has the same size as the input batch /// and only contains original columns that are not being unnested -fn unnest_at_level( +/// If there are multiple unnest operations with different recursion level, +/// the total time of unnesting will be max(recursion) +fn list_unnest_at_level( batch: &[ArrayRef], list_type_unnests: &[ListUnnest], temp_unnested_arrs: &mut HashMap, @@ -508,6 +479,56 @@ struct UnnestingResult { /// - For struct columns: We will expand the struct columns into multiple subfield columns. /// /// For columns that don't need to be unnested, repeat their values until reaching the longest length. +/// +/// Note: unnest has a big difference in behavior between Postgres and DuckDB +/// +/// Take this example +/// +/// 1. Postgres +/// ```ignored +/// create table temp ( +/// i integer[][][], j integer[] +/// ) +/// insert into temp values ('{{{1,2},{3,4}},{{5,6},{7,8}}}', '{1,2}'); +/// select unnest(i), unnest(j) from temp; +/// ``` +/// +/// Result +/// ```text +/// 1 1 +/// 2 2 +/// 3 +/// 4 +/// 5 +/// 6 +/// 7 +/// 8 +/// ``` +/// 2. DuckDB +/// ```ignore +/// create table temp (i integer[][][], j integer[]); +/// insert into temp values ([[[1,2],[3,4]],[[5,6],[7,8]]], [1,2]); +/// select unnest(i,recursive:=true), unnest(j,recursive:=true) from temp; +/// ``` +/// Result: +/// ```text +/// +/// ┌────────────────────────────────────────────────┬────────────────────────────────────────────────┐ +/// │ unnest(i, "recursive" := CAST('t' AS BOOLEAN)) │ unnest(j, "recursive" := CAST('t' AS BOOLEAN)) │ +/// │ int32 │ int32 │ +/// ├────────────────────────────────────────────────┼────────────────────────────────────────────────┤ +/// │ 1 │ 1 │ +/// │ 2 │ 2 │ +/// │ 3 │ 1 │ +/// │ 4 │ 2 │ +/// │ 5 │ 1 │ +/// │ 6 │ 2 │ +/// │ 7 │ 1 │ +/// │ 8 │ 2 │ +/// └────────────────────────────────────────────────┴────────────────────────────────────────────────┘ +/// ``` +/// +/// The following implementation refer to DuckDB's implementation fn build_batch( batch: &RecordBatch, schema: &SchemaRef, @@ -535,7 +556,7 @@ fn build_batch( true => batch.columns(), false => &flatten_arrs, }; - let (temp_result, num_rows) = unnest_at_level( + let (temp_result, num_rows) = list_unnest_at_level( input, list_type_columns, &mut temp_unnested_result, diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index 55eb867a0c53..6ea3808fa2c8 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -456,7 +456,7 @@ impl<'a> TreeNodeRewriter for RecursiveUnnestRewriter<'a> { /// This downward traversal needs to keep track of: /// - Whether or not some unnest expr has been visited from the top util the current node /// - If some unnest expr has been visited, maintain a stack of such information, this - /// is used to detect if some recursive unnest expr exists (e.g **unnest(unnest(unnest(3d column))))** + /// is used to detect if some recursive unnest expr exists (e.g **unnest(unnest(unnest(3d column))))** fn f_down(&mut self, expr: Expr) -> Result> { if let Expr::Unnest(ref unnest_expr) = expr { let (data_type, _) = From f96f6e9fbae3c05d71c9d11ac0daf93fbcef997b Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Wed, 11 Sep 2024 22:33:15 +0200 Subject: [PATCH 49/56] tune comment --- datafusion/physical-plan/src/unnest.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index 92c0e498d000..d2c69063de7d 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -403,9 +403,7 @@ pub struct ListUnnest { /// - colB(1-dimension) unnest into P3 (unnesting of P3 starts from this level) /// /// The returned array will has the same size as the input batch -/// and only contains original columns that are not being unnested -/// If there are multiple unnest operations with different recursion level, -/// the total time of unnesting will be max(recursion) +/// and only contains original columns that are not being unnested. fn list_unnest_at_level( batch: &[ArrayRef], list_type_unnests: &[ListUnnest], From c695b5ec53bb15841ff7abc8c59a369b5a9f17c7 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Wed, 11 Sep 2024 23:15:40 +0200 Subject: [PATCH 50/56] rm todo --- datafusion/core/src/dataframe/mod.rs | 1 - datafusion/expr/src/logical_plan/display.rs | 1 - datafusion/physical-plan/src/unnest.rs | 1 - datafusion/proto/src/logical_plan/mod.rs | 4 ++-- datafusion/sql/src/select.rs | 14 ++++++-------- datafusion/sql/src/utils.rs | 21 +++++++++------------ 6 files changed, 17 insertions(+), 25 deletions(-) diff --git a/datafusion/core/src/dataframe/mod.rs b/datafusion/core/src/dataframe/mod.rs index 72ca1eedf898..2138bd1294b4 100644 --- a/datafusion/core/src/dataframe/mod.rs +++ b/datafusion/core/src/dataframe/mod.rs @@ -367,7 +367,6 @@ impl DataFrame { options: UnnestOptions, ) -> Result { let columns = columns.iter().map(|c| Column::from(*c)).collect(); - let plan = LogicalPlanBuilder::from(self.plan) .unnest_columns_with_options(columns, options)? .build()?; diff --git a/datafusion/expr/src/logical_plan/display.rs b/datafusion/expr/src/logical_plan/display.rs index 388071f3bae8..26d54803d403 100644 --- a/datafusion/expr/src/logical_plan/display.rs +++ b/datafusion/expr/src/logical_plan/display.rs @@ -643,7 +643,6 @@ impl<'a, 'b> PgJsonVisitor<'a, 'b> { let input_columns = plan.schema().columns(); let list_type_columns = list_col_indices .iter() - // TODO: fix me .map(|(i, unnest_info)| { format!( "{}|depth={:?}", diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index d2c69063de7d..8c1b63a7fa63 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -807,7 +807,6 @@ fn unnest_list_array( capacity: usize, ) -> Result { let values = list_array.values(); - // TODO: handle me recursively let mut take_indicies_builder = PrimitiveArray::::builder(capacity); for row in 0..list_array.len() { let mut value_length = 0; diff --git a/datafusion/proto/src/logical_plan/mod.rs b/datafusion/proto/src/logical_plan/mod.rs index 5cd4267f426a..1d2b631682cb 100644 --- a/datafusion/proto/src/logical_plan/mod.rs +++ b/datafusion/proto/src/logical_plan/mod.rs @@ -1580,7 +1580,7 @@ impl AsLogicalPlan for LogicalPlanNode { input, extension_codec, )?; - let todo = list_type_columns + let proto_unnest_list_items = list_type_columns .iter() .map(|(index, ul)| ColumnUnnestListItem { input_index: *index as _, @@ -1623,7 +1623,7 @@ impl AsLogicalPlan for LogicalPlanNode { }), }) .collect(), - list_type_columns: todo, + list_type_columns: proto_unnest_list_items, struct_type_columns: struct_type_columns .iter() .map(|c| *c as u64) diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs index b0ca7bf5d110..eb79f67e4c88 100644 --- a/datafusion/sql/src/select.rs +++ b/datafusion/sql/src/select.rs @@ -22,8 +22,8 @@ use crate::planner::{ idents_to_table_reference, ContextProvider, PlannerContext, SqlToRel, }; use crate::utils::{ - check_columns_satisfy_exprs, extract_aliases, group_bottom_most_consecutive_unnests, - rebase_expr, resolve_aliases_to_exprs, resolve_columns, resolve_positions_to_exprs, + check_columns_satisfy_exprs, extract_aliases, rebase_expr, resolve_aliases_to_exprs, + resolve_columns, resolve_positions_to_exprs, rewrite_recursive_unnests_bottom_up, }; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; @@ -318,7 +318,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { // - unnest(struct_col) will be transformed into unnest(struct_col).field1, unnest(struct_col).field2 // - unnest(array_col) will be transformed into unnest(array_col).element // - unnest(array_col) + 1 will be transformed into unnest(array_col).element +1 - let outer_projection_exprs = group_bottom_most_consecutive_unnests( + let outer_projection_exprs = rewrite_recursive_unnests_bottom_up( &intermediate_plan, &mut unnest_columns, &mut inner_projection_exprs, @@ -350,11 +350,9 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { } } - let ret = LogicalPlanBuilder::from(intermediate_plan) + LogicalPlanBuilder::from(intermediate_plan) .project(intermediate_select_exprs)? - .build()?; - - Ok(ret) + .build() } fn try_process_aggregate_unnest(&self, input: LogicalPlan) -> Result { @@ -417,7 +415,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { let mut unnest_columns = vec![]; let mut inner_projection_exprs = vec![]; - let outer_projection_exprs = group_bottom_most_consecutive_unnests( + let outer_projection_exprs = rewrite_recursive_unnests_bottom_up( &intermediate_plan, &mut unnest_columns, &mut inner_projection_exprs, diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index 6ea3808fa2c8..f12441ec955e 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -293,7 +293,7 @@ pub(crate) fn value_to_string(value: &Value) -> Option { } } -pub(crate) fn group_bottom_most_consecutive_unnests( +pub(crate) fn rewrite_recursive_unnests_bottom_up( input: &LogicalPlan, unnest_placeholder_columns: &mut Vec<(Column, ColumnUnnestType)>, inner_projection_exprs: &mut Vec, @@ -302,7 +302,7 @@ pub(crate) fn group_bottom_most_consecutive_unnests( Ok(original_exprs .iter() .map(|expr| { - group_bottom_most_consecutive_unnest( + rewrite_recursive_unnest_bottom_up( input, unnest_placeholder_columns, inner_projection_exprs, @@ -582,7 +582,6 @@ fn push_projection_dedupl(projection: &mut Vec, expr: Expr) { projection.push(expr); } } -/// TODO: maybe renamed into transform_bottom_most_consecutive_unnests /// The context is we want to rewrite unnest() into InnerProjection->Unnest->OuterProjection /// Given an expression which contains unnest expr as one of its children, /// Try transform depends on unnest type @@ -592,7 +591,7 @@ fn push_projection_dedupl(projection: &mut Vec, expr: Expr) { /// The transformed exprs will be used in the outer projection /// If along the path from root to bottom, there are multiple unnest expressions, the transformation /// is done only for the bottom expression -pub(crate) fn group_bottom_most_consecutive_unnest( +pub(crate) fn rewrite_recursive_unnest_bottom_up( input: &LogicalPlan, unnest_placeholder_columns: &mut Vec<(Column, ColumnUnnestType)>, inner_projection_exprs: &mut Vec, @@ -655,9 +654,7 @@ mod tests { use datafusion_functions::core::expr_ext::FieldAccessor; use datafusion_functions_aggregate::expr_fn::count; - use crate::utils::{ - group_bottom_most_consecutive_unnest, resolve_positions_to_exprs, - }; + use crate::utils::{resolve_positions_to_exprs, rewrite_recursive_unnest_bottom_up}; fn column_unnests_eq(l: Vec<(&str, &str)>, r: &[(Column, ColumnUnnestType)]) { let r_formatted: Vec = r.iter().map(|i| format!("{}|{}", i.0, i.1)).collect(); @@ -699,7 +696,7 @@ mod tests { let original_expr = unnest(unnest(col("3d_col"))) .add(unnest(unnest(col("3d_col")))) .add(col("i64_col")); - let transformed_exprs = group_bottom_most_consecutive_unnest( + let transformed_exprs = rewrite_recursive_unnest_bottom_up( &input, &mut unnest_placeholder_columns, &mut inner_projection_exprs, @@ -736,7 +733,7 @@ mod tests { // unnest(3d_col) as 2d_col let original_expr_2 = unnest(col("3d_col")).alias("2d_col"); - let transformed_exprs = group_bottom_most_consecutive_unnest( + let transformed_exprs = rewrite_recursive_unnest_bottom_up( &input, &mut unnest_placeholder_columns, &mut inner_projection_exprs, @@ -804,7 +801,7 @@ mod tests { // unnest(struct_col) let original_expr = unnest(col("struct_col")); - let transformed_exprs = group_bottom_most_consecutive_unnest( + let transformed_exprs = rewrite_recursive_unnest_bottom_up( &input, &mut unnest_placeholder_columns, &mut inner_projection_exprs, @@ -830,7 +827,7 @@ mod tests { // unnest(array_col) + 1 let original_expr = unnest(col("array_col")).add(lit(1i64)); - let transformed_exprs = group_bottom_most_consecutive_unnest( + let transformed_exprs = rewrite_recursive_unnest_bottom_up( &input, &mut unnest_placeholder_columns, &mut inner_projection_exprs, @@ -899,7 +896,7 @@ mod tests { // An expr with multiple unnest let original_expr = unnest(unnest(col("struct_col").field("matrix"))); - let transformed_exprs = group_bottom_most_consecutive_unnest( + let transformed_exprs = rewrite_recursive_unnest_bottom_up( &input, &mut unnest_placeholder_columns, &mut inner_projection_exprs, From f73921fd4f8e1fa177607d0c74e2cb098e8e1943 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Thu, 12 Sep 2024 23:15:22 +0200 Subject: [PATCH 51/56] refactor --- datafusion/expr/src/logical_plan/builder.rs | 8 +++++-- datafusion/physical-plan/src/unnest.rs | 24 ++++++++------------- datafusion/sql/src/utils.rs | 22 ++++++++----------- 3 files changed, 24 insertions(+), 30 deletions(-) diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index 9d241f71e34b..c469ab3a4dee 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -1691,24 +1691,28 @@ pub fn get_unnested_columns( /// /// For example: /// Input schema as -/// +/// ```text /// +---------------------+-------------------+ /// | col1 | col2 | /// +---------------------+-------------------+ /// | Struct(INT64,INT32) | List(List(Int64)) | /// +---------------------+-------------------+ +/// ``` +/// +/// /// /// Then unnesting columns with: /// - (col1,Struct) /// - (col2,List(\[depth=1,depth=2\])) /// /// will generate a new schema as -/// +/// ```text /// +---------+---------+---------------------+---------------------+ /// | col1.c0 | col1.c1 | unnest_col2_depth_1 | unnest_col2_depth_2 | /// +---------+---------+---------------------+---------------------+ /// | Int64 | Int32 | List(Int64) | Int64 | /// +---------+---------+---------------------+---------------------+ +/// ``` pub fn unnest_with_options( input: LogicalPlan, columns_to_unnest: Vec<(Column, ColumnUnnestType)>, diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index 8c1b63a7fa63..50af6b4960a5 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -247,7 +247,7 @@ struct UnnestStream { schema: Arc, /// represents all unnest operations to be applied to the input (input index, depth) /// e.g unnest(col1),unnest(unnest(col1)) where col1 has index 1 in original input schema - /// then list_type_columns = [(1,1),(1,2)] + /// then list_type_columns = [ListUnnest{1,1},ListUnnest{1,2}] list_type_columns: Vec, struct_column_indices: HashSet, /// Options @@ -456,8 +456,9 @@ fn list_unnest_at_level( // Create the take indices array for other columns let take_indices = create_take_indicies(unnested_length, total_length); - // vertical expansion because of list unnest - let ret = flatten_arrs_from_indices(batch, &take_indices)?; + // dimension of arrays in batch is untouch, but the values are repeated + // as the side effect of unnesting + let ret = repeat_arrs_from_indices(batch, &take_indices)?; unnested_temp_arrays .into_iter() .zip(list_unnest_specs.iter()) @@ -597,7 +598,7 @@ fn build_batch( .into_iter() .map( // each item in unnested_columns is the result of unnesting the same input column - // we need to sort them to conform with the unnest definition + // we need to sort them to conform with the original expression order // e.g unnest(unnest(col)) must goes before unnest(col) |(original_index, mut unnested_columns)| { unnested_columns.sort_by( @@ -683,18 +684,11 @@ fn find_longest_length( } else { Scalar::new(Int64Array::from_value(0, 1)) }; - // col1: [[1,2]]|[[2,3]] - // unnest(col1): [1,2]|[2,3] => length = 2 - // unnest(unnest(col1)): 1|2|2|3 => length = 4 - // col2: [3]|[4] => length =2 - // unnest(col2): 3|4 => length = 2 - // unnest(col1), unnest(unnest(col1)), unnest(col2) let list_lengths: Vec = list_arrays .iter() .map(|list_array| { let mut length_array = length(list_array)?; // Make sure length arrays have the same type. Int64 is the most general one. - // Respect the depth of unnest( current func only get the length of 1 level of unnest) length_array = cast(&length_array, &DataType::Int64)?; length_array = zip(&is_not_null(&length_array)?, &length_array, &null_length)?; @@ -867,10 +861,10 @@ fn create_take_indicies( builder.finish() } -/// Create the batch given the unnested column arrays and a `indices` array +/// Create the batch given an arrays and a `indices` array /// that is used by the take kernel to copy values. /// -/// For example if we have the following `RecordBatch`: +/// For example if we have the following batch: /// /// ```ignore /// c1: [1], null, [2, 3, 4], null, [5, 6] @@ -898,7 +892,7 @@ fn create_take_indicies( /// c2: 'a', 'b', 'c', 'c', 'c', null, 'd', 'd' /// ``` /// -fn flatten_arrs_from_indices( +fn repeat_arrs_from_indices( batch: &[ArrayRef], indices: &PrimitiveArray, ) -> Result>> { @@ -1001,7 +995,7 @@ mod tests { } #[test] - fn test_build_batch_list_arr() -> datafusion_common::Result<()> { + fn test_build_batch_list_arr_recursive() -> datafusion_common::Result<()> { // col1 | col2 // [[1,2,3],null,[4,5]] | ['a','b'] // [[7,8,9,10], null, [11,12,13]] | ['c','d'] diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index f12441ec955e..656e4b851aa8 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -315,14 +315,6 @@ pub(crate) fn rewrite_recursive_unnests_bottom_up( .collect::>()) } -fn print_unnest(expr: &str, level: usize) -> String { - let mut result = String::from(expr); - for _ in 0..level { - result = format!("UNNEST({})", result); - } - result -} - /* This is only usedful when used with transform down up A full example of how the transformation works: @@ -359,6 +351,7 @@ impl<'a> RecursiveUnnestRewriter<'a> { fn transform( &mut self, level: usize, + alias_name: String, expr_in_unnest: &Expr, struct_allowed: bool, ) -> Result> { @@ -373,7 +366,7 @@ impl<'a> RecursiveUnnestRewriter<'a> { format!("unnest_placeholder({},depth={})", inner_expr_name, level); // This is due to the fact that unnest transformation should keep the original // column name as is, to comply with group by and order by - let post_unnest_alias = print_unnest(&inner_expr_name, level); + // let post_unnest_alias = print_unnest(&inner_expr_name, level); let placeholder_column = Column::from_name(placeholder_name.clone()); let (data_type, _) = expr_in_unnest.data_type_and_nullable(self.input_schema)?; @@ -407,8 +400,7 @@ impl<'a> RecursiveUnnestRewriter<'a> { ); // let post_unnest_column = Column::from_name(post_unnest_name); - let post_unnest_expr = - col(post_unnest_name.clone()).alias(post_unnest_alias); + let post_unnest_expr = col(post_unnest_name.clone()).alias(alias_name); match self .columns_unnestings .iter_mut() @@ -546,8 +538,12 @@ impl<'a> TreeNodeRewriter for RecursiveUnnestRewriter<'a> { let unnest_recursion = unnest_stack.len(); let struct_allowed = (&expr == self.root_expr) && unnest_recursion == 1; - let mut transformed_exprs = - self.transform(unnest_recursion, inner_expr, struct_allowed)?; + let mut transformed_exprs = self.transform( + unnest_recursion, + expr.schema_name().to_string(), + inner_expr, + struct_allowed, + )?; if struct_allowed { self.transformed_root_exprs = Some(transformed_exprs.clone()); } From c3b68c6837e7e51b63fc991c9464f15cb132468a Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Sat, 14 Sep 2024 10:14:43 +0200 Subject: [PATCH 52/56] chore: reserve test --- datafusion/expr/src/logical_plan/builder.rs | 13 +++++---- datafusion/sqllogictest/test_files/unnest.slt | 28 +++++++++++++++---- 2 files changed, 30 insertions(+), 11 deletions(-) diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index c469ab3a4dee..c38ce2d02fd2 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -1590,7 +1590,7 @@ pub fn unnest(input: LogicalPlan, columns: Vec) -> Result { unnest_with_options(input, unnestings, UnnestOptions::default()) } -pub fn get_unnested_list_column_recursive( +pub fn get_unnested_list_datatype_recursive( data_type: &DataType, depth: usize, ) -> Result { @@ -1601,7 +1601,7 @@ pub fn get_unnested_list_column_recursive( if depth == 1 { return Ok(field.data_type().clone()); } - return get_unnested_list_column_recursive(field.data_type(), depth - 1); + return get_unnested_list_datatype_recursive(field.data_type(), depth - 1); } _ => {} }; @@ -1609,7 +1609,10 @@ pub fn get_unnested_list_column_recursive( internal_err!("trying to unnest on invalid data type {:?}", data_type) } -fn get_unnested_columns_inferred( +/// Infer the unnest type based on the data type: +/// - list type: infer to unnest(list(col, depth=1)) +/// - struct type: infer to unnest(struct) +fn infer_unnest_type( col_name: &String, data_type: &DataType, ) -> Result { @@ -1654,7 +1657,7 @@ pub fn get_unnested_columns( match data_type { DataType::List(_) | DataType::FixedSizeList(_, _) | DataType::LargeList(_) => { - let data_type = get_unnested_list_column_recursive(data_type, depth)?; + let data_type = get_unnested_list_datatype_recursive(data_type, depth)?; let new_field = Arc::new(Field::new( col_name, data_type, // Unnesting may produce NULLs even if the list is not null. @@ -1756,7 +1759,7 @@ pub fn unnest_with_options( Some((column_to_unnest, unnest_type)) => { let mut inferred_unnest_type = unnest_type.clone(); if let ColumnUnnestType::Inferred = unnest_type { - inferred_unnest_type = get_unnested_columns_inferred( + inferred_unnest_type = infer_unnest_type( &column_to_unnest.name, original_field.data_type(), )?; diff --git a/datafusion/sqllogictest/test_files/unnest.slt b/datafusion/sqllogictest/test_files/unnest.slt index 76b75fdee831..4086f465788a 100644 --- a/datafusion/sqllogictest/test_files/unnest.slt +++ b/datafusion/sqllogictest/test_files/unnest.slt @@ -564,6 +564,24 @@ select unnest(unnest(column3)), column3 from recursive_unnest_table; [1] [[1, 2]] [{c0: [1], c1: [[1, 2]]}] [2] [[3], [4]] [{c0: [2], c1: [[3], [4]]}] + +query TT +explain select unnest(unnest(column3)), column3 from recursive_unnest_table; +---- +logical_plan +01)Unnest: lists[] structs[unnest_placeholder(UNNEST(recursive_unnest_table.column3))] +02)--Projection: unnest_placeholder(recursive_unnest_table.column3,depth=1) AS UNNEST(recursive_unnest_table.column3) AS unnest_placeholder(UNNEST(recursive_unnest_table.column3)), recursive_unnest_table.column3 +03)----Unnest: lists[unnest_placeholder(recursive_unnest_table.column3)|depth=1] structs[] +04)------Projection: recursive_unnest_table.column3 AS unnest_placeholder(recursive_unnest_table.column3), recursive_unnest_table.column3 +05)--------TableScan: recursive_unnest_table projection=[column3] +physical_plan +01)UnnestExec +02)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +03)----ProjectionExec: expr=[unnest_placeholder(recursive_unnest_table.column3,depth=1)@0 as unnest_placeholder(UNNEST(recursive_unnest_table.column3)), column3@1 as column3] +04)------UnnestExec +05)--------ProjectionExec: expr=[column3@0 as unnest_placeholder(recursive_unnest_table.column3), column3@0 as column3] +06)----------MemoryExec: partitions=1, partition_sizes=[1] + ## unnest->field_access->unnest->unnest query I? select unnest(unnest(unnest(column3)['c1'])), column3 from recursive_unnest_table; @@ -604,17 +622,17 @@ NULL [4] [{c0: [2], c1: [[3], [4]]}] ## and need multiple unnesting logical plans ## e.g unnest -> field_access -> unnest query TT -explain select unnest(unnest(unnest(column3)['c1'])), unnest(unnest(column3)['c1']), column3 from recursive_unnest_table; +explain select unnest(unnest(unnest(column3)['c1'])), column3 from recursive_unnest_table; ---- logical_plan -01)Projection: unnest_placeholder(UNNEST(recursive_unnest_table.column3)[c1],depth=2) AS UNNEST(UNNEST(UNNEST(recursive_unnest_table.column3)[c1])), unnest_placeholder(UNNEST(recursive_unnest_table.column3)[c1],depth=1) AS UNNEST(UNNEST(recursive_unnest_table.column3)[c1]), recursive_unnest_table.column3 -02)--Unnest: lists[unnest_placeholder(UNNEST(recursive_unnest_table.column3)[c1])|depth=2, unnest_placeholder(UNNEST(recursive_unnest_table.column3)[c1])|depth=1] structs[] +01)Projection: unnest_placeholder(UNNEST(recursive_unnest_table.column3)[c1],depth=2) AS UNNEST(UNNEST(UNNEST(recursive_unnest_table.column3)[c1])), recursive_unnest_table.column3 +02)--Unnest: lists[unnest_placeholder(UNNEST(recursive_unnest_table.column3)[c1])|depth=2] structs[] 03)----Projection: get_field(unnest_placeholder(recursive_unnest_table.column3,depth=1) AS UNNEST(recursive_unnest_table.column3), Utf8("c1")) AS unnest_placeholder(UNNEST(recursive_unnest_table.column3)[c1]), recursive_unnest_table.column3 04)------Unnest: lists[unnest_placeholder(recursive_unnest_table.column3)|depth=1] structs[] 05)--------Projection: recursive_unnest_table.column3 AS unnest_placeholder(recursive_unnest_table.column3), recursive_unnest_table.column3 06)----------TableScan: recursive_unnest_table projection=[column3] physical_plan -01)ProjectionExec: expr=[unnest_placeholder(UNNEST(recursive_unnest_table.column3)[c1],depth=2)@0 as UNNEST(UNNEST(UNNEST(recursive_unnest_table.column3)[c1])), unnest_placeholder(UNNEST(recursive_unnest_table.column3)[c1],depth=1)@1 as UNNEST(UNNEST(recursive_unnest_table.column3)[c1]), column3@2 as column3] +01)ProjectionExec: expr=[unnest_placeholder(UNNEST(recursive_unnest_table.column3)[c1],depth=2)@0 as UNNEST(UNNEST(UNNEST(recursive_unnest_table.column3)[c1])), column3@1 as column3] 02)--UnnestExec 03)----ProjectionExec: expr=[get_field(unnest_placeholder(recursive_unnest_table.column3,depth=1)@0, c1) as unnest_placeholder(UNNEST(recursive_unnest_table.column3)[c1]), column3@1 as column3] 04)------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 @@ -624,8 +642,6 @@ physical_plan - - ## group by unnest ### without agg exprs From 43e7a04fba03a14e42bb816aead756db93d29d5e Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Sat, 14 Sep 2024 10:18:44 +0200 Subject: [PATCH 53/56] add a basic test --- datafusion/sqllogictest/test_files/unnest.slt | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/datafusion/sqllogictest/test_files/unnest.slt b/datafusion/sqllogictest/test_files/unnest.slt index 4086f465788a..63ca74e9714c 100644 --- a/datafusion/sqllogictest/test_files/unnest.slt +++ b/datafusion/sqllogictest/test_files/unnest.slt @@ -165,6 +165,13 @@ select unnest(column1), column1 from unnest_table; 6 [6] 12 [12] +# unnest at different level at the same time +query II +select unnest([1,2,3]), unnest(unnest([[1,2,3]])); +---- +1 1 +2 2 +3 3 # binary expr linking different unnest exprs query II From 6dae034bcd714e6d3a7167266cc3715f7b9b124e Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Sun, 22 Sep 2024 11:09:45 +0200 Subject: [PATCH 54/56] chore: more document --- datafusion/expr/src/logical_plan/plan.rs | 31 +++++++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index eea12b84a8a0..c90aedc989d2 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -3051,6 +3051,12 @@ pub enum Partitioning { DistributeBy(Vec), } +/// Represents the unnesting operation on a column based on the context (a known struct +/// column, a list column, or let the planner infer the unnesting type). +/// +/// The inferred unnesting type works for both struct and list column, but the unnesting +/// will only be done once (depth = 1). In case recursion is needed on a multi-dimensional +/// list type, use [`ColumnUnnestList`] #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub enum ColumnUnnestType { List(Vec), @@ -3058,7 +3064,7 @@ pub enum ColumnUnnestType { Struct, // Infer the unnest type based on column schema // If column is a list column, the unnest depth will be 1 - // This value is to support sugar syntax of old api (unnest(columns1,columns2)) + // This value is to support sugar syntax of old api in Dataframe (unnest(either_list_or_struct_column)) Inferred, } @@ -3076,6 +3082,29 @@ impl fmt::Display for ColumnUnnestType { } } +/// Represent the unnesting operation on a list column, such as the recursion depth and +/// the output column name after unnesting +/// +/// Example: given `ColumnUnnestList { output_column: "output_name", depth: 2 }` +/// and input as a two dimentional array column values +/// ```text +/// input +/// --- +/// [[1,2]] +/// [[3]] +/// [[4],[5]] +/// ``` +/// +/// This operation will result into a column with values +/// ```text +/// output_name +/// --- +/// 1 +/// 2 +/// 3 +/// 4 +/// 5 +/// ``` #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct ColumnUnnestList { pub output_column: Column, From e72bb61532d37aaf7892cf5244f930f01318fc67 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Sun, 22 Sep 2024 11:25:33 +0200 Subject: [PATCH 55/56] doc on ColumnUnnestType List --- datafusion/expr/src/logical_plan/plan.rs | 32 +++++++++++------------- 1 file changed, 15 insertions(+), 17 deletions(-) diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index c90aedc989d2..b91f88dc0d25 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -3059,6 +3059,8 @@ pub enum Partitioning { /// list type, use [`ColumnUnnestList`] #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub enum ColumnUnnestType { + // Unnesting a list column, a vector of ColumnUnnestList is used because + // a column can be unnested at different levels, resulting different output columns List(Vec), // for struct, there can only be one unnest performed on one column at a time Struct, @@ -3086,24 +3088,20 @@ impl fmt::Display for ColumnUnnestType { /// the output column name after unnesting /// /// Example: given `ColumnUnnestList { output_column: "output_name", depth: 2 }` -/// and input as a two dimentional array column values -/// ```text -/// input -/// --- -/// [[1,2]] -/// [[3]] -/// [[4],[5]] -/// ``` -/// -/// This operation will result into a column with values +/// /// ```text -/// output_name -/// --- -/// 1 -/// 2 -/// 3 -/// 4 -/// 5 +/// input output_name +/// ┌─────────┐ ┌─────────┐ +/// │{{1,2}} │ │ 1 │ +/// ├─────────┼─────►├─────────┤ +/// │{{3}} │ │ 2 │ +/// ├─────────┤ ├─────────┤ +/// │{{4},{5}}│ │ 3 │ +/// └─────────┘ ├─────────┤ +/// │ 4 │ +/// ├─────────┤ +/// │ 5 │ +/// └─────────┘ /// ``` #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct ColumnUnnestList { From bc3f4bf98451a10292756f690b9f5ca0d4395342 Mon Sep 17 00:00:00 2001 From: Duong Cong Toai Date: Sun, 22 Sep 2024 13:19:54 +0200 Subject: [PATCH 56/56] chore: add partialord to new types --- datafusion/expr/src/logical_plan/plan.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index 3ad200b68f46..443d23804adb 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -3306,7 +3306,7 @@ pub enum Partitioning { /// The inferred unnesting type works for both struct and list column, but the unnesting /// will only be done once (depth = 1). In case recursion is needed on a multi-dimensional /// list type, use [`ColumnUnnestList`] -#[derive(Debug, Clone, PartialEq, Eq, Hash)] +#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd)] pub enum ColumnUnnestType { // Unnesting a list column, a vector of ColumnUnnestList is used because // a column can be unnested at different levels, resulting different output columns @@ -3352,7 +3352,7 @@ impl fmt::Display for ColumnUnnestType { /// │ 5 │ /// └─────────┘ /// ``` -#[derive(Debug, Clone, PartialEq, Eq, Hash)] +#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd)] pub struct ColumnUnnestList { pub output_column: Column, pub depth: usize, @@ -3395,10 +3395,10 @@ impl PartialOrd for Unnest { /// The incoming logical plan pub input: &'a Arc, /// Columns to run unnest on, can be a list of (List/Struct) columns - pub exec_columns: &'a Vec, + pub exec_columns: &'a Vec<(Column, ColumnUnnestType)>, /// refer to the indices(in the input schema) of columns /// that have type list to run unnest on - pub list_type_columns: &'a Vec, + pub list_type_columns: &'a Vec<(usize, ColumnUnnestList)>, /// refer to the indices (in the input schema) of columns /// that have type struct to run unnest on pub struct_type_columns: &'a Vec,