feat: support `unnest` multiple arrays #10044

jonahgao · 2024-04-11T05:44:19Z

Which issue does this PR close?

Closes #212.
Closes #1608.
Closes #6555.
Closes #6796.
Closes #7087.

Rationale for this change

From the PostgreSQL's document.

Expands multiple arrays (possibly of different data types) into a set of rows. If the arrays are not all the same length then the shorter ones are padded with NULLs. This form is only allowed in a query's FROM clause

An unnest function with N parameters is equivalent to N single-parameter unnest functions, that is,
select * from unnest(arg1, arg2, arg3) is equivalent to select unnest(arg1), unnest(arg2), unnest(arg3).
This PR will unify them all into the latter to be handled by UnnestExec.

In order to support UnnestExec with multiple unnest columns:

For each row in a RecordBatch, some list columns need to be unnested. 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 columns that don't need to be unnested, repeat their values until reaching the longest length.

What changes are included in this PR?

Expand UnnestExec to support unnesting multiple arrays.
Also, enable the following two types of queries:

There are multiple unnest functions in the select list.
In the from clause, unnest can support multiple arguments.

DataFusion CLI v37.0.0
❯ select unnest([1,2]), unnest([3]);
+---------------------------------------+------------------------------+
| unnest(make_array(Int64(1),Int64(2))) | unnest(make_array(Int64(3))) |
+---------------------------------------+------------------------------+
| 1                                     | 3                            |
| 2                                     |                              |
+---------------------------------------+------------------------------+
2 row(s) fetched.
Elapsed 0.072 seconds.

❯ select * from unnest([1,2], [3]);
+---------------------------------------+------------------------------+
| unnest(make_array(Int64(1),Int64(2))) | unnest(make_array(Int64(3))) |
+---------------------------------------+------------------------------+
| 1                                     | 3                            |
| 2                                     |                              |
+---------------------------------------+------------------------------+

Are these changes tested?

Yes

Are there any user-facing changes?

Yes.
Several unnest related structs have changed.

jonahgao · 2024-04-12T14:55:06Z

datafusion/expr/src/expr.rs

@@ -186,7 +186,16 @@ pub enum Expr {

 #[derive(Clone, PartialEq, Eq, Hash, Debug)]
 pub struct Unnest {
-    pub exprs: Vec<Expr>,
+    pub expr: Box<Expr>,


In the expression position, only one argument can be accepted.

alamb · 2024-04-13T00:24:31Z

That is a pretty sweet list of things this closes. Thank you @jonahgao -- I will try and review this PR over the next day or two if no one beats me to it

jayzhan211 · 2024-04-13T09:00:39Z

datafusion/sqllogictest/test_files/unnest.slt

-query error DataFusion error: This feature is not implemented: Only support single unnest expression for now
-select unnest(column1), unnest(column2) from unnest_table;
+query ?I
+select unnest([]), unnest(NULL::int[]);


I think we should support null as well 🤔

query error DataFusion error: This feature is not implemented: unnest\(\) does not support null yet select unnest([]), unnest(NULL);

I'm not sure how to support it yet. PostgreSQL, BigQuery, and Clickhouse don't support unnesting untyped NULLs.

postgres=# select unnest(NULL); ERROR: function unnest(unknown) is not unique LINE 1: select unnest(NULL); ^ HINT: Could not choose a best candidate function. You might need to add explicit type casts.

Although DuckDB supports it, its behavior seems a bit strange to me.

D select unnest(NULL), unnest([1,2]); ┌──────────────┬───────────────────────────────┐ │ unnest(NULL) │ unnest(main.list_value(1, 2)) │ │ int32 │ int32 │ ├──────────────────────────────────────────────┤ │ 0 rows │ └──────────────────────────────────────────────┘ D select unnest(NULL::int[]), unnest([1,2]); ┌─────────────────────────────────┬───────────────────────────────┐ │ unnest(CAST(NULL AS INTEGER[])) │ unnest(main.list_value(1, 2)) │ │ int32 │ int32 │ ├─────────────────────────────────┼───────────────────────────────┤ │ │ 1 │ │ │ 2 │ └─────────────────────────────────┴───────────────────────────────┘

Unnesting an untyped NULL and a typed NULL produced different results.

I think there may be two appropriate methods:

Use type coercion to convert an untype NULL into a certain List type like cast(NULL, NULL::int[]).

Explicitly reject unnesting untyped NULLs.

I prefer 2, since PostgresSQL, and Clickhouse return errors in this case, we follow the majority.

I prefer 2, since PostgresSQL, and Clickhouse return errors in this case, we follow the majority.

Make sense to me!

I double checked and this PR does error as expected. I'll add a small test for this case here and merge this PR.

❯ select unnest(null) from t; This feature is not implemented: unnest() does not support null yet

jayzhan211 · 2024-04-13T08:58:24Z

datafusion/expr/src/logical_plan/builder.rs

+                unnest_field.name(),
+                field.data_type().clone(),
+                // Unnesting may produce NULLs even if the list is not null.
+                // For example: unnset([1], []) -> 1, null


Although the test does not cover this case, sadly I don't think there is currently a way to create a list with nullable as false 🤔 Maybe we can write a simple rust test for this?

I think so too. true should be equivalent to the previous unnest_field.is_nullable().
I planned to write some tests related to DataFrame, which could potentially implement it. But the changes in this PR are larger than I expected, so those tests have been delayed.

I planned to write some tests related to DataFrame, which could potentially implement it. But the changes in this PR are larger than I expected, so those tests have been delayed.

We can do it in another PR!

We can do it in another PR!

👌

jayzhan211

it looks pretty nice!

jayzhan211 · 2024-04-13T09:05:20Z

datafusion/expr/src/logical_plan/builder.rs

-            return Ok(input);
-        }
-    };
+    let mut unnested_fields: HashMap<usize, _> = HashMap::with_capacity(columns.len());


alamb · 2024-04-15T10:36:16Z

Closes #212.
Closes #1608.
Closes #6555.
Closes #6796.
Closes #7087.

😍

Thank you @jonahgao and thank you @jayzhan211 for the review. I merged this branch up from main and added a test for untyped nulls and plan to merge it once CI passes

alamb · 2024-04-15T11:05:05Z

🚀

jonahgao · 2024-04-15T12:53:33Z

Thank you @alamb for the new test and the review. Thank you @jayzhan211 for the review.

jonahgao added 2 commits April 11, 2024 00:36

Impl find_longest_length

d9fd370

impl unnesting multi columns

393c0fb

jonahgao marked this pull request as draft April 11, 2024 05:44

Merge branch 'main' into unnest_multi_args

b73c884

jonahgao changed the title ~~feat: support unnest with multiple arguments~~ feat: support unnest multiple columns Apr 11, 2024

Change plans

d4a8625

github-actions bot added sql SQL Planner logical-expr Logical plan and expressions core Core DataFusion crate sqllogictest SQL Logic Tests (.slt) labels Apr 11, 2024

jonahgao changed the title ~~feat: support unnest multiple columns~~ feat: support unnest multiple arrays Apr 11, 2024

jonahgao added 5 commits April 11, 2024 23:06

remove println

5fa7894

fix tests

8ced636

simplify unnested fields

992d0b2

update doc and tests

5d6c932

more tests

92a7efa

jonahgao marked this pull request as ready for review April 12, 2024 14:43

jonahgao commented Apr 12, 2024

View reviewed changes

jonahgao added 2 commits April 12, 2024 23:25

add test

6b82040

fix comment

5ad6d77

alamb mentioned this pull request Apr 13, 2024

DataFusion weekly project plan (Andrew Lamb) - April 8, 2024 #10002

Closed

9 tasks

jayzhan211 reviewed Apr 13, 2024

View reviewed changes

jayzhan211 approved these changes Apr 13, 2024

View reviewed changes

jayzhan211 reviewed Apr 13, 2024

View reviewed changes

alamb added 2 commits April 15, 2024 06:35

Add test for untyped null

8e7af12

Merge remote-tracking branch 'apache/main' into unnest_multi_args

c20546c

alamb merged commit 34eda15 into apache:main Apr 15, 2024
25 checks passed

jonahgao deleted the unnest_multi_args branch April 15, 2024 12:53

jonahgao mentioned this pull request Apr 17, 2024

feat: DataFrame supports unnesting multiple columns #10118

Merged

alamb mentioned this pull request Apr 22, 2024

DataFusion weekly project plan (Andrew Lamb) - April 22, 2024 #10172

Closed

7 tasks

ahirner mentioned this pull request Jun 2, 2024

Getting arrays of objects datafusion-contrib/datafusion-functions-json#13

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: support `unnest` multiple arrays #10044

feat: support `unnest` multiple arrays #10044

jonahgao commented Apr 11, 2024 •

edited

Loading

jonahgao Apr 12, 2024

This comment was marked as outdated.

alamb commented Apr 13, 2024

This comment was marked as resolved.

jayzhan211 Apr 13, 2024

jonahgao Apr 13, 2024 •

edited

Loading

jayzhan211 Apr 15, 2024

jonahgao Apr 15, 2024

alamb Apr 15, 2024

jayzhan211 Apr 13, 2024 •

edited

Loading

jonahgao Apr 13, 2024

jayzhan211 Apr 15, 2024

jonahgao Apr 15, 2024 •

edited

Loading

jayzhan211 left a comment

jayzhan211 Apr 13, 2024

alamb commented Apr 15, 2024

alamb commented Apr 15, 2024

jonahgao commented Apr 15, 2024

feat: support unnest multiple arrays #10044

feat: support unnest multiple arrays #10044

Conversation

jonahgao commented Apr 11, 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

This comment was marked as outdated.

alamb commented Apr 13, 2024

This comment was marked as resolved.

Choose a reason for hiding this comment

jonahgao Apr 13, 2024 • edited Loading

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

jayzhan211 Apr 13, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jonahgao Apr 15, 2024 • edited Loading

Choose a reason for hiding this comment

jayzhan211 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb commented Apr 15, 2024

alamb commented Apr 15, 2024

jonahgao commented Apr 15, 2024

feat: support `unnest` multiple arrays #10044

feat: support `unnest` multiple arrays #10044

jonahgao commented Apr 11, 2024 •

edited

Loading

jonahgao Apr 13, 2024 •

edited

Loading

jayzhan211 Apr 13, 2024 •

edited

Loading

jonahgao Apr 15, 2024 •

edited

Loading