Improve recursive `unnest` options API #12836

duongcongtoai · 2024-10-09T18:11:16Z

Which issue does this PR close?

Cleanup #11577

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

duongcongtoai · 2024-10-09T18:12:18Z

datafusion/expr/src/logical_plan/builder.rs

    unnest_with_options(input, unnestings, UnnestOptions::default())
 }

-pub fn get_unnested_list_datatype_recursive(


remove this method, recursion specification can be declared using UnnestOptions

duongcongtoai · 2024-10-09T18:13:16Z

datafusion/expr/src/logical_plan/plan.rs

-/// 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, PartialOrd)]
-pub enum ColumnUnnestType {


This struct is now a part of the builder, no need to expose to api

alamb · 2024-10-18T19:50:22Z

Merging up from main to get CI to pass

alamb

Thank you @duongcongtoai -- I think this is a significant improvement in code cleaniless

I think it would be nice to add some comments, as I have described below, but I also don't think it is strictly necessary and could be done after merge

alamb · 2024-10-18T19:52:00Z

datafusion/common/src/unnest.rs

@@ -64,13 +66,24 @@
 pub struct UnnestOptions {
    /// Should nulls in the input be preserved? Defaults to true
    pub preserve_nulls: bool,
+    /// Without explicit recursions, all


Can you please document what the Vec means in this context? Is the option for each of the columns in the schema? If so can you say it is expected to have the same number of elements?

If possible it would also be great to update the documentation on this struct to explain what is going on with recursive options

alamb · 2024-10-18T19:52:12Z

datafusion/common/src/unnest.rs

+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd)]
+pub struct RecursionUnnestOption {


Can we please add documentation to this structure as it is public?

alamb · 2024-10-18T19:53: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(


alamb · 2024-10-18T19:56:06Z

datafusion/expr/src/logical_plan/builder.rs

+                UnnestOptions::default().with_recursions(vec![
+                    RecursionUnnestOption {
+                        input_column: "stringss".into(),
+                        output_column: "stringss_depth_1".into(),
+                        depth: 1,
+                    },
+                    RecursionUnnestOption {
+                        input_column: "stringss".into(),
+                        output_column: "stringss_depth_2".into(),
+                        depth: 2,
+                    },
+                ]),


Here is another potential way for this API to look would be to take individual options rather than a single Vec

I don't feel strongly about this

Suggested change

UnnestOptions::default().with_recursions(vec![

RecursionUnnestOption {

input_column: "stringss".into(),

output_column: "stringss_depth_1".into(),

depth: 1,

},

RecursionUnnestOption {

input_column: "stringss".into(),

output_column: "stringss_depth_2".into(),

depth: 2,

},

]),

UnnestOptions::new()

.with_recursion(

RecursionUnnestOption {

input_column: "stringss".into(),

output_column: "stringss_depth_1".into(),

depth: 1,

})

.with_recusion(

RecursionUnnestOption {

input_column: "stringss".into(),

output_column: "stringss_depth_2".into(),

depth: 2,

},

),

thank you, i refactored so that the recursion is optional at column level, and since None can be replaced by an empty vector, i simplified the type of this field from Option to Vec

alamb · 2024-10-18T19:56:44Z

datafusion/physical-plan/src/unnest.rs

+"|                                 |                                 | e                               |",
+"+---------------------------------+---------------------------------+---------------------------------+",
+        ];
+        assert_batches_eq!(expected, &[ret]);


alamb · 2024-10-18T19:59:18Z

datafusion/sql/src/select.rs

@@ -301,7 +301,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
        // The transformation happen bottom up, one at a time for each iteration
        // Only exhaust the loop if no more unnest transformation is found
        for i in 0.. {
-            let mut unnest_columns = vec![];
+            let mut unnest_columns = HashMap::new();


It might be better to use an IndexMap here (to retain the same order as was specified in the query) -- a HashMap can mix up the order

alamb · 2024-10-18T20:04:10Z

datafusion/sql/src/select.rs

+                        match maybe_list_unnest {
+                            None => {}
+                            Some(list_unnest) => {
+                                list_recursions.extend(list_unnest.into_iter().map(
+                                    |unnest_list| RecursionUnnestOption {
+                                        input_column: col.clone(),
+                                        output_column: unnest_list.output_column,
+                                        depth: unnest_list.depth,
+                                    },
+                                ));
+                            }
+                        }


I think you could simplify this code (reduce recursion level) with if let like

Suggested change

match maybe_list_unnest {

None => {}

Some(list_unnest) => {

list_recursions.extend(list_unnest.into_iter().map(

|unnest_list| RecursionUnnestOption {

input_column: col.clone(),

output_column: unnest_list.output_column,

depth: unnest_list.depth,

},

));

}

}

if let Some(list_unnest) = maybe_list_unnest {

list_recursions.extend(list_unnest.into_iter().map(

|unnest_list| RecursionUnnestOption {

input_column: col.clone(),

output_column: unnest_list.output_column,

depth: unnest_list.depth,

},

));

}

}

Also I wonder if it would be possible / easier for rewrite_recursive_unnests_bottom_up to simply return a HashSet<Column, RecursionUnnestOption> directly rather than having to convert the results

true, i can refactor this function so it returns 3 result:

inner_projection_exprs

unnest_placeholder_columns

outer_projection_exprs

alamb · 2024-10-20T10:53:12Z

Looks great -- thanks again @duongcongtoai

Ref: apache/datafusion#12836

duongcongtoai added 3 commits October 5, 2024 12:55

refactor

9b1ab42

refactor unnest options

a16799b

more test

e43bc66

github-actions bot added sql SQL Planner logical-expr Logical plan and expressions common Related to common crate labels Oct 9, 2024

duongcongtoai commented Oct 9, 2024

View reviewed changes

duongcongtoai mentioned this pull request Oct 14, 2024

unnest + field access of "list of struct" broken #12921

Closed

resolve comments

7d4ae95

github-actions bot added the physical-expr Physical Expressions label Oct 17, 2024

duongcongtoai added 2 commits October 17, 2024 21:23

add back doc

8e824a9

Merge remote-tracking branch 'origin/main' into follow-up-11577

8ed2563

duongcongtoai marked this pull request as ready for review October 17, 2024 19:35

duongcongtoai changed the title ~~Cleanup TODO in recursive unnest~~ Cleanup TODO in recursive unnest Oct 17, 2024

fix proto

6ea0c93

github-actions bot added the proto Related to proto crate label Oct 17, 2024

duongcongtoai and others added 3 commits October 17, 2024 22:52

flaky test

1f92614

clippy

2d19955

Merge remote-tracking branch 'apache/main' into follow-up-11577

002c1c3

alamb added the api change Changes the API exposed to users of the crate label Oct 18, 2024

alamb approved these changes Oct 18, 2024

View reviewed changes

duongcongtoai added 4 commits October 20, 2024 10:49

use indexmap

aa9327a

chore: compile err

e6aef1f

chore: update cargo

c9eb970

chore: fmt cargotoml

52db541

alamb changed the title ~~Cleanup TODO in recursive unnest~~ Improve recursive unnest options API Oct 20, 2024

alamb merged commit c7e5d8d into apache:main Oct 20, 2024
27 checks passed

alamb mentioned this pull request Oct 21, 2024

Oct 21, 2024: This week in DataFusion #13035

Closed

4 tasks

Michael-J-Ward added a commit to Michael-J-Ward/datafusion-python that referenced this pull request Oct 28, 2024

update API for logical expr Limit

4cf0179

Ref: apache/datafusion#12836

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Improve recursive `unnest` options API #12836

Improve recursive `unnest` options API #12836

duongcongtoai commented Oct 9, 2024 •

edited

Loading

duongcongtoai Oct 9, 2024

duongcongtoai Oct 9, 2024

alamb commented Oct 18, 2024

alamb left a comment

alamb Oct 18, 2024

alamb Oct 18, 2024

alamb Oct 18, 2024

alamb Oct 18, 2024

duongcongtoai Oct 20, 2024

alamb Oct 18, 2024

alamb Oct 18, 2024

alamb Oct 18, 2024

duongcongtoai Oct 20, 2024

alamb commented Oct 20, 2024

Improve recursive unnest options API #12836

Improve recursive unnest options API #12836

Conversation

duongcongtoai commented Oct 9, 2024 • edited Loading

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb commented Oct 18, 2024

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

Choose a reason for hiding this comment

alamb commented Oct 20, 2024

Improve recursive `unnest` options API #12836

Improve recursive `unnest` options API #12836

duongcongtoai commented Oct 9, 2024 •

edited

Loading