Refactor to support recursive unnest in physical plan #11577

duongcongtoai · 2024-07-20T19:54:08Z

Which issue does this PR close?

Rationale for this change

We are not properly handling unnest where there are multiple unnest exprs at different level involved in the select exprs

What changes are included in this PR?

Implement `RecursiveUnnestRewriter`

It can transform a composite expr of unnest(unnest(unnest(columnA))) into a logical plan unnest(columnA, depth=3)

This implementation also fixed #11198, by the usage of push_projection_dedupl, meaning during the traversal of the complex expr, if a column is already projected, it should not be projected again (check #11198 (comment) for further context)

Support recursive unnest at physical layer by calling `list_unnest_at_level` multiple times

To fix #11689, i took a look how DuckDB and Postgres behave, where in the query there are multiple unnesting exprs, but with different recursion level.
This implementation follows DuckDB's behavior. For more context on the behavior difference, check this comment

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.

For example giving the following query:

select unnest(colA, max_depth:=3) as P1, unnest(colA,max_depth:=2) as P2, unnest(colB, max_depth:=1) as P3 from temp;

Then the total times this function being called is 3

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

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

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)
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
colB(1-dimension) unnest into P3 (unnesting of P3 starts from this level)

Are these changes tested?

Are there any user-facing changes?

…iple-projection-reference

duongcongtoai · 2024-09-14T08:19:24Z

datafusion/sqllogictest/test_files/unnest.slt

@@ -165,6 +165,34 @@ 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]]));


alamb

Thank you @duongcongtoai -- I started looking through this PR and it looks good so far -- I need to find a few more minutes to complete the review but I left some comments for now

alamb · 2024-09-21T11:16:17Z

datafusion/expr/src/logical_plan/plan.rs

+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]


Can we please add documentation to this struct about what it means (what computation it represents)? In the context of this PR it is clear, but once we merge I think the context will be lost

Thank you, comment addressed

alamb · 2024-09-21T11:16:58Z

datafusion/expr/src/logical_plan/plan.rs

@@ -3049,17 +3051,54 @@ pub enum Partitioning {
    DistributeBy(Vec<Expr>),
 }

+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub enum ColumnUnnestType {
+    List(Vec<ColumnUnnestList>),


Can you please add documentation to List about what it is (ideally with an example of when it is used and how it is different to Inferred?

comment addressed

…iple-projection-reference

alamb

Thank you very much @duongcongtoai -- this is an epic PR. Very nice 👏

I think it is worth considering the public API (aka adding fields to UnnestOptions rather than new methods), but otherwise this is looking really nice.

It is well tested and well commented

I left a few other comments that would be nice to address too -- comments and adding a few more tests.

Thanks again

cc @jayzhan211

alamb · 2024-09-23T18:58:35Z

datafusion/expr/src/logical_plan/builder.rs

+    /// 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(


An alternate API would be to add the ColumnUnnestType as a field on UnnestOptions

(for that matter, maybe we should just add columns: Vec<(Column, ColumnUnnestType)>, as a field 🤔

alamb · 2024-09-23T18:59:03Z

datafusion/expr/src/logical_plan/builder.rs

+    unnest_with_options(input, unnestings, UnnestOptions::default())
+}
+
+pub fn get_unnested_list_datatype_recursive(


If this needs to be pub i think it should have some doc strings

alamb · 2024-09-23T18:59:27Z

datafusion/expr/src/logical_plan/builder.rs

+    }
+}
+
+pub fn get_struct_unnested_columns(


likewise here

alamb · 2024-09-23T19:03:23Z

datafusion/physical-plan/src/unnest.rs

+                preserve_nulls: true,
+            },
+        )?;
+        let actual =


for the future, you an use assert_batches_eq for this kind of comparison

alamb · 2024-09-23T19:05:41Z

datafusion/sql/src/utils.rs

-        assert_eq!(
-            unnest_placeholder_columns,
-            vec!["UNNEST(struct_col[matrix])"]
+        // TODO: add a test case where


Is this still something you mean to do?

alamb · 2024-09-23T19:06:45Z

datafusion/sqllogictest/test_files/unnest.slt

@@ -165,6 +165,34 @@ select unnest(column1), column1 from unnest_table;
 6 [6]
 12 [12]



Are there any .slt tests for unnesting structures? these tests seem to be on arrays (Lists) rather than Structs

jayzhan211 · 2024-09-25T09:11:40Z

datafusion/sql/src/utils.rs

+                    expr_in_unnest.clone().alias(placeholder_name.clone()),
+                );
+
+                // let post_unnest_column = Column::from_name(post_unnest_name);


seems like leftover

jayzhan211 · 2024-09-25T09:14:28Z

datafusion/sql/src/utils.rs

+        // 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()?;


seems like leftover

alamb · 2024-09-25T10:26:59Z

I think we can handle the follow ons / comments above as a follow on PR. Let's merge this one to keep things flowing.

@duongcongtoai let me know if you plan to work on the suggestions.

Thanks again. 🚀 onwards. 🙏 @jayzhan211

toaiduongdh · 2024-09-25T11:34:16Z

Yes thank you for the reviews, i'll create a follow up Pr

* chore: poc * fix unnest struct * UT for memoization * remove unnessary projection * chore: temp test case * multi depth unnest supported * chore: add map of original column and transformed col * transformation map to physical layer * prototype for recursive array length * chore: some compile err * finalize input type in physical layer * chore: refactor unnest builder * add unnesting type inferred * fix compile err * fail test in builder * Compile err * chore: detect some bugs * some work * support recursive unnest in physical layer * UT for new build batch function * compile err * fix unnesting into empty arrays * some comment * fix unnest struct * some note * chore: fix all test failure * fix projection pushdown * custom rewriter for recursive unnest * simplify * rm unnecessary projection * chore: better comments * more comments * chore: better comments * remove breaking api * rename * more unit test * remove debug * clean up * fix proto * fix dataframe * fix clippy * cargo fmt * fix some test * fix all test * fix unnest in join * fix doc and tests * chore: better doc * better doc * tune comment * rm todo * refactor * chore: reserve test * add a basic test * chore: more document * doc on ColumnUnnestType List * chore: add partialord to new types

chore: poc

cc661ea

github-actions bot added sql SQL Planner sqllogictest SQL Logic Tests (.slt) labels Jul 20, 2024

duongcongtoai added 9 commits July 21, 2024 11:35

fix unnest struct

d7d45b1

UT for memoization

f8aa23a

remove unnessary projection

cf202a8

chore: temp test case

1e74176

Merge remote-tracking branch 'origin/main' into 11198-fix-unnest-mult…

b3ff0d7

…iple-projection-reference

multi depth unnest supported

d557ff3

chore: add map of original column and transformed col

cd497fe

transformation map to physical layer

d2cc80c

prototype for recursive array length

7efee19

github-actions bot added the logical-expr Logical plan and expressions label Jul 30, 2024

chore: some compile err

a489502

duongcongtoai mentioned this pull request Jul 31, 2024

Reduce repetition in try_process_group_by_unnest and try_process_unnest #11714

Merged

duongcongtoai added 5 commits July 31, 2024 22:48

finalize input type in physical layer

80c1cf3

chore: refactor unnest builder

7acb056

add unnesting type inferred

4d32187

fix compile err

e41923f

fail test in builder

eca58f5

github-actions bot added the optimizer Optimizer rules label Aug 3, 2024

duongcongtoai added 2 commits August 4, 2024 12:03

Compile err

e2fae71

chore: detect some bugs

825e270

github-actions bot added the core Core DataFusion crate label Aug 4, 2024

some work

721c92f

github-actions bot added the physical-expr Physical Expressions label Aug 12, 2024

duongcongtoai added 4 commits August 13, 2024 22:21

support recursive unnest in physical layer

d11ed20

UT for new build batch function

5b9ce5e

compile err

366d8ae

Merge remote-tracking branch 'origin/main' into 11198-fix-unnest-mult…

a06dbad

…iple-projection-reference

duongcongtoai added 9 commits September 10, 2024 21:20

fix clippy

4d3f508

cargo fmt

926efd5

fix some test

bf40688

fix all test

1e54422

fix unnest in join

abfd4fc

fix doc and tests

eb4cb44

chore: better doc

7aee552

better doc

df4fee3

tune comment

f96f6e9

duongcongtoai marked this pull request as ready for review September 11, 2024 20:41

duongcongtoai marked this pull request as draft September 11, 2024 21:00

duongcongtoai added 4 commits September 11, 2024 23:15

rm todo

c695b5e

refactor

f73921f

chore: reserve test

c3b68c6

add a basic test

43e7a04

duongcongtoai commented Sep 14, 2024

View reviewed changes

duongcongtoai marked this pull request as ready for review September 14, 2024 08:20

alamb mentioned this pull request Sep 16, 2024

DataFusion weekly project plan (Andrew Lamb) - Sep 16, 2024 #12494

Closed

8 tasks

alamb reviewed Sep 21, 2024

View reviewed changes

duongcongtoai added 4 commits September 22, 2024 11:09

chore: more document

6dae034

doc on ColumnUnnestType List

e72bb61

Merge remote-tracking branch 'origin/main' into 11198-fix-unnest-mult…

2e53740

…iple-projection-reference

chore: add partialord to new types

bc3f4bf

alamb approved these changes Sep 23, 2024

View reviewed changes

jayzhan211 reviewed Sep 25, 2024

View reviewed changes

alamb merged commit b35e720 into apache:main Sep 25, 2024
24 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Refactor to support recursive unnest in physical plan #11577

Refactor to support recursive unnest in physical plan #11577

duongcongtoai commented Jul 20, 2024 •

edited

Loading

duongcongtoai Sep 14, 2024

alamb left a comment

alamb Sep 21, 2024

duongcongtoai Sep 22, 2024

alamb Sep 21, 2024

duongcongtoai Sep 22, 2024

alamb left a comment

alamb Sep 23, 2024

alamb Sep 23, 2024

alamb Sep 23, 2024

alamb Sep 23, 2024

alamb Sep 23, 2024

alamb Sep 23, 2024

jayzhan211 Sep 25, 2024

jayzhan211 Sep 25, 2024

alamb commented Sep 25, 2024 •

edited

Loading

toaiduongdh commented Sep 25, 2024

		@@ -165,6 +165,34 @@ select unnest(column1), column1 from unnest_table;
		6 [6]
		12 [12]

Refactor to support recursive unnest in physical plan #11577

Refactor to support recursive unnest in physical plan #11577

Conversation

duongcongtoai commented Jul 20, 2024 • edited Loading

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Implement RecursiveUnnestRewriter

Support recursive unnest at physical layer by calling list_unnest_at_level multiple times

Are these changes tested?

Are there any user-facing changes?

Choose a reason for hiding this comment

alamb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb commented Sep 25, 2024 • edited Loading

toaiduongdh commented Sep 25, 2024

duongcongtoai commented Jul 20, 2024 •

edited

Loading

Implement `RecursiveUnnestRewriter`

Support recursive unnest at physical layer by calling `list_unnest_at_level` multiple times

alamb commented Sep 25, 2024 •

edited

Loading