unnest should preserve the input's equivalence properties for uninvolved columns

[vlad-arista]

Continuing the discussion from last comments in #14991

It looks like unnest's implementation does preserve the output ordering of the columns that aren't being unnested, but it doesn't preserve their equivalence properties. This leads to situation where ordering information is lost and results in unnecessary unordered aggregation (which also breaks streaming queries that unnest and re-aggregate columns).

[alamb]

Here is the example / test case from @asubiotto in #14991 (comment)

In the following plan, both Aggregate should have ordering_mode=Sorted

So instead of

Aggregate: groupBy=[[unnested.generated_id]], aggr=[[array_agg(unnested.ar)]]                                                                                                                                                                                                                                                                                                                                                                      

It should be

AggregateExec: mode=FinalPartitioned, gby=[generated_id@0 as generated_id], aggr=[sum(unnested.ar)], ordering_mode=Sorted                                                                                                                                                                                                                                                                                                                         

Here is the full query

EXPLAIN WITH unnested AS (SELECT
    ROW_NUMBER() OVER () AS generated_id,
    unnest(array[value]) as ar
  FROM range(1,5)) SELECT array_agg(ar) FROM unnested group by generated_id;
+---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| plan_type     | plan                                                                                                                                                                                                                                                                                                                                                                                                                                                  |
+---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| logical_plan  | Projection: array_agg(unnested.ar)                                                                                                                                                                                                                                                                                                                                                                                                                    |
|               |   Aggregate: groupBy=[[unnested.generated_id]], aggr=[[array_agg(unnested.ar)]]                                                                                                                                                                                                                                                                                                                                                                       |
|               |     SubqueryAlias: unnested                                                                                                                                                                                                                                                                                                                                                                                                                           |
|               |       Projection: generated_id, __unnest_placeholder(make_array(tmp_table.value),depth=1) AS UNNEST(make_array(tmp_table.value)) AS ar                                                                                                                                                                                                                                                                                                                |
|               |         Unnest: lists[__unnest_placeholder(make_array(tmp_table.value))|depth=1] structs[]                                                                                                                                                                                                                                                                                                                                                            |
|               |           Projection: row_number() ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING, row_number() ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING AS generated_id, make_array(tmp_table.value) AS __unnest_placeholder(make_array(tmp_table.value))                                                                                                                                                                                  |
|               |             WindowAggr: windowExpr=[[row_number() ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING]]                                                                                                                                                                                                                                                                                                                                          |
|               |               TableScan: tmp_table projection=[value]                                                                                                                                                                                                                                                                                                                                                                                                 |
| physical_plan | ProjectionExec: expr=[array_agg(unnested.ar)@1 as array_agg(unnested.ar)]                                                                                                                                                                                                                                                                                                                                                                             |
|               |   AggregateExec: mode=FinalPartitioned, gby=[generated_id@0 as generated_id], aggr=[array_agg(unnested.ar)]                                                                                                                                                                                                                                                                                                                                           |
|               |     CoalesceBatchesExec: target_batch_size=8192                                                                                                                                                                                                                                                                                                                                                                                                       |
|               |       RepartitionExec: partitioning=Hash([generated_id@0], 12), input_partitions=12                                                                                                                                                                                                                                                                                                                                                                   |
|               |         AggregateExec: mode=Partial, gby=[generated_id@0 as generated_id], aggr=[array_agg(unnested.ar)]                                                                                                                                                                                                                                                                                                                                              |
|               |           ProjectionExec: expr=[generated_id@1 as generated_id, __unnest_placeholder(make_array(tmp_table.value),depth=1)@2 as ar]                                                                                                                                                                                                                                                                                                                    |
|               |             UnnestExec                                                                                                                                                                                                                                                                                                                                                                                                                                |
|               |               ProjectionExec: expr=[row_number() ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@1 as row_number() ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING, row_number() ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@1 as generated_id, make_array(value@0) as __unnest_placeholder(make_array(tmp_table.value))]                                                                                              |
|               |                 RepartitionExec: partitioning=RoundRobinBatch(12), input_partitions=1                                                                                                                                                                                                                                                                                                                                                                 |
|               |                   BoundedWindowAggExec: wdw=[row_number() ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "row_number() ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING", data_type: UInt64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }], mode=[Sorted] |
|               |                     LazyMemoryExec: partitions=1, batch_generators=[range: start=1, end=5, batch_size=8192]                                                                                                                                                                                                                                                                                                                                           |
|               |                                                                                                                                                                                                                                                                                                                                                                                                                                                       |
+---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
2 row(s) fetched.
Elapsed 0.010 seconds.

[vegarsti]

Can I try to take this?

[alamb]

Sounds good -- thanks @vegarsti

FYI @asubiotto

[vegarsti]

Just doing some initial exploration looking at plan output. The style of the output seems to have changed a bit since these were posted, so posting it here. I am running cargo run --bin datafusion-cli (please let me know if you all are using a different entrypoint).

To confirm, what we want is for both the inner and the outer AggregateExec to know that they're sorted?

> EXPLAIN WITH unnested AS (SELECT
    ROW_NUMBER() OVER () AS generated_id,
    unnest(array[value]) as ar
  FROM range(1,5)) SELECT array_agg(ar) FROM unnested group by generated_id;
+---------------+-------------------------------+
| plan_type     | plan                          |
+---------------+-------------------------------+
| physical_plan | ┌───────────────────────────┐ |
|               | │       ProjectionExec      │ |
|               | │    --------------------   │ |
|               | │  array_agg(unnested.ar):  │ |
|               | │   array_agg(unnested.ar)  │ |
|               | └─────────────┬─────────────┘ |
|               | ┌─────────────┴─────────────┐ |
|               | │       AggregateExec       │ |
|               | │    --------------------   │ |
|               | │           aggr:           │ |
|               | │   array_agg(unnested.ar)  │ |
|               | │                           │ |
|               | │         group_by:         │ |
|               | │        generated_id       │ |
|               | │                           │ |
|               | │           mode:           │ |
|               | │      FinalPartitioned     │ |
|               | └─────────────┬─────────────┘ |
|               | ┌─────────────┴─────────────┐ |
|               | │    CoalesceBatchesExec    │ |
|               | │    --------------------   │ |
|               | │     target_batch_size:    │ |
|               | │            8192           │ |
|               | └─────────────┬─────────────┘ |
|               | ┌─────────────┴─────────────┐ |
|               | │      RepartitionExec      │ |
|               | │    --------------------   │ |
|               | │ partition_count(in->out): │ |
|               | │          11 -> 11         │ |
|               | │                           │ |
|               | │    partitioning_scheme:   │ |
|               | │ Hash([generated_id@0], 11)│ |
|               | └─────────────┬─────────────┘ |
|               | ┌─────────────┴─────────────┐ |
|               | │       AggregateExec       │ |
|               | │    --------------------   │ |
|               | │           aggr:           │ |
|               | │   array_agg(unnested.ar)  │ |
|               | │                           │ |
|               | │         group_by:         │ |
|               | │        generated_id       │ |
|               | │                           │ |
|               | │       mode: Partial       │ |
|               | └─────────────┬─────────────┘ |
|               | ┌─────────────┴─────────────┐ |
|               | │       ProjectionExec      │ |
|               | │    --------------------   │ |
|               | │            ar:            │ |
|               | │    __unnest_placeholder   │ |
|               | │    (make_array(range()    │ |
|               | │      .value),depth=1)     │ |
|               | │                           │ |
|               | │       generated_id:       │ |
|               | │        generated_id       │ |
|               | └─────────────┬─────────────┘ |
|               | ┌─────────────┴─────────────┐ |
|               | │         UnnestExec        │ |
|               | └─────────────┬─────────────┘ |
|               | ┌─────────────┴─────────────┐ |
|               | │       ProjectionExec      │ |
|               | │    --------------------   │ |
|               | │    __unnest_placeholder   │ |
|               | │    (make_array(range()    │ |
|               | │         .value)):         │ |
|               | │     make_array(value)     │ |
|               | │                           │ |
|               | │       generated_id:       │ |
|               | │ row_number() ROWS BETWEEN │ |
|               | │   UNBOUNDED PRECEDING AND │ |
|               | │     UNBOUNDED FOLLOWING   │ |
|               | └─────────────┬─────────────┘ |
|               | ┌─────────────┴─────────────┐ |
|               | │      RepartitionExec      │ |
|               | │    --------------------   │ |
|               | │ partition_count(in->out): │ |
|               | │          1 -> 11          │ |
|               | │                           │ |
|               | │    partitioning_scheme:   │ |
|               | │    RoundRobinBatch(11)    │ |
|               | └─────────────┬─────────────┘ |
|               | ┌─────────────┴─────────────┐ |
|               | │    BoundedWindowAggExec   │ |
|               | │    --------------------   │ |
|               | │        mode: Sorted       │ |
|               | │                           │ |
|               | │        select_list:       │ |
|               | │ row_number() ROWS BETWEEN │ |
|               | │   UNBOUNDED PRECEDING AND │ |
|               | │     UNBOUNDED FOLLOWING   │ |
|               | └─────────────┬─────────────┘ |
|               | ┌─────────────┴─────────────┐ |
|               | │       LazyMemoryExec      │ |
|               | │    --------------------   │ |
|               | │     batch_generators:     │ |
|               | │   range: start=1, end=5,  │ |
|               | │       batch_size=8192     │ |
|               | └───────────────────────────┘ |
|               |                               |
+---------------+-------------------------------+
1 row(s) fetched.
Elapsed 0.025 seconds.

FWIW the other example from @asubiotto (without unnest) which did preserve ordering, no longer includes ordering_mode=Sorted in the plan output:

> EXPLAIN WITH unnested AS (SELECT
    ROW_NUMBER() OVER () AS generated_id,
    value as ar
  FROM range(1,5)) SELECT generated_id, sum(ar) FROM unnested group by generated_id;
+---------------+-------------------------------+
| plan_type     | plan                          |
+---------------+-------------------------------+
| physical_plan | ┌───────────────────────────┐ |
|               | │       AggregateExec       │ |
|               | │    --------------------   │ |
|               | │           aggr:           │ |
|               | │      sum(unnested.ar)     │ |
|               | │                           │ |
|               | │         group_by:         │ |
|               | │        generated_id       │ |
|               | │                           │ |
|               | │           mode:           │ |
|               | │      FinalPartitioned     │ |
|               | └─────────────┬─────────────┘ |
|               | ┌─────────────┴─────────────┐ |
|               | │          SortExec         │ |
|               | │    --------------------   │ |
|               | │  generated_id@0 ASC NULLS │ |
|               | │            LAST           │ |
|               | └─────────────┬─────────────┘ |
|               | ┌─────────────┴─────────────┐ |
|               | │    CoalesceBatchesExec    │ |
|               | │    --------------------   │ |
|               | │     target_batch_size:    │ |
|               | │            8192           │ |
|               | └─────────────┬─────────────┘ |
|               | ┌─────────────┴─────────────┐ |
|               | │      RepartitionExec      │ |
|               | │    --------------------   │ |
|               | │ partition_count(in->out): │ |
|               | │          11 -> 11         │ |
|               | │                           │ |
|               | │    partitioning_scheme:   │ |
|               | │ Hash([generated_id@0], 11)│ |
|               | └─────────────┬─────────────┘ |
|               | ┌─────────────┴─────────────┐ |
|               | │       AggregateExec       │ |
|               | │    --------------------   │ |
|               | │           aggr:           │ |
|               | │      sum(unnested.ar)     │ |
|               | │                           │ |
|               | │         group_by:         │ |
|               | │        generated_id       │ |
|               | │                           │ |
|               | │       mode: Partial       │ |
|               | └─────────────┬─────────────┘ |
|               | ┌─────────────┴─────────────┐ |
|               | │      RepartitionExec      │ |
|               | │    --------------------   │ |
|               | │ partition_count(in->out): │ |
|               | │          1 -> 11          │ |
|               | │                           │ |
|               | │    partitioning_scheme:   │ |
|               | │    RoundRobinBatch(11)    │ |
|               | └─────────────┬─────────────┘ |
|               | ┌─────────────┴─────────────┐ |
|               | │       ProjectionExec      │ |
|               | │    --------------------   │ |
|               | │         ar: value         │ |
|               | │                           │ |
|               | │       generated_id:       │ |
|               | │ row_number() ROWS BETWEEN │ |
|               | │   UNBOUNDED PRECEDING AND │ |
|               | │     UNBOUNDED FOLLOWING   │ |
|               | └─────────────┬─────────────┘ |
|               | ┌─────────────┴─────────────┐ |
|               | │    BoundedWindowAggExec   │ |
|               | │    --------------------   │ |
|               | │        mode: Sorted       │ |
|               | │                           │ |
|               | │        select_list:       │ |
|               | │ row_number() ROWS BETWEEN │ |
|               | │   UNBOUNDED PRECEDING AND │ |
|               | │     UNBOUNDED FOLLOWING   │ |
|               | └─────────────┬─────────────┘ |
|               | ┌─────────────┴─────────────┐ |
|               | │       LazyMemoryExec      │ |
|               | │    --------------------   │ |
|               | │     batch_generators:     │ |
|               | │   range: start=1, end=5,  │ |
|               | │       batch_size=8192     │ |
|               | └───────────────────────────┘ |
|               |                               |
+---------------+-------------------------------+
1 row(s) fetched.
Elapsed 0.022 seconds.

Similarly, even the plan for select min(value)from range(1,5) doesn't show that it's sorted 🤔

> explain select min(value)from range(1,5);
+---------------+-------------------------------+
| plan_type     | plan                          |
+---------------+-------------------------------+
| physical_plan | ┌───────────────────────────┐ |
|               | │       AggregateExec       │ |
|               | │    --------------------   │ |
|               | │           aggr:           │ |
|               | │     min(range().value)    │ |
|               | │                           │ |
|               | │        mode: Final        │ |
|               | └─────────────┬─────────────┘ |
|               | ┌─────────────┴─────────────┐ |
|               | │   CoalescePartitionsExec  │ |
|               | └─────────────┬─────────────┘ |
|               | ┌─────────────┴─────────────┐ |
|               | │       AggregateExec       │ |
|               | │    --------------------   │ |
|               | │           aggr:           │ |
|               | │     min(range().value)    │ |
|               | │                           │ |
|               | │       mode: Partial       │ |
|               | └─────────────┬─────────────┘ |
|               | ┌─────────────┴─────────────┐ |
|               | │      RepartitionExec      │ |
|               | │    --------------------   │ |
|               | │ partition_count(in->out): │ |
|               | │          1 -> 11          │ |
|               | │                           │ |
|               | │    partitioning_scheme:   │ |
|               | │    RoundRobinBatch(11)    │ |
|               | └─────────────┬─────────────┘ |
|               | ┌─────────────┴─────────────┐ |
|               | │       LazyMemoryExec      │ |
|               | │    --------------------   │ |
|               | │     batch_generators:     │ |
|               | │   range: start=1, end=5,  │ |
|               | │       batch_size=8192     │ |
|               | └───────────────────────────┘ |
|               |                               |
+---------------+-------------------------------+
1 row(s) fetched.
Elapsed 0.013 seconds.

[vegarsti]

Okay, my bad: I can reproduce that output by putting it in a sqllogictest.

[vegarsti]

Take

[vegarsti]

I have something that works for the query above, which is copying over the input plan's equivalence properties. I need to understand EquivalenceProperties better to understand if this is semantically valid for all cases. I think I need to come up with a query where there is an ordering on the column that is unnested, to verify that in this case we don't have an ordering anymore.

In any case, I opened a PR here: #16985

[vegarsti]

The PR is ready for review!

[alamb]

Here is a test case I wrote that should change with this code

# Data is ordered by column2 (100, 200, 300, 400)
statement error DataFusion error: SQL error: ParserError\("Expected: an SQL statement, found: COPT at Line: 1, Column: 1"\)
COPY (
  SELECT * FROM VALUES
    ([1,2,3], 100),
    ([3],     200),
    ([],      300),
    ([3,1],   400)
  ORDER BY column2
 ) TO 'test_files/scratch/unnest/ordered.parquet';


statement ok
CREATE EXTERNAL TABLE t
STORED AS PARQUET
LOCATION 'test_files/scratch/unnest/ordered.parquet'
WITH ORDER (column2)

query ?I
SELECT * FROM t;
----
[1, 2, 3] 100
[3] 200
[] 300
[3, 1] 400

# data is sorted on column2 already, so no need to sort again
query II
SELECT UNNEST(column1), column2 FROM t ORDER BY column2;
----
1 100
2 100
3 100
3 200
3 400
1 400

# Explain should not have a SortExec
query TT
EXPLAIN SELECT UNNEST(column1), column2 FROM t ORDER BY column2;
----
logical_plan
01)Sort: t.column2 ASC NULLS LAST
02)--Projection: __unnest_placeholder(t.column1,depth=1) AS UNNEST(t.column1), t.column2
03)----Unnest: lists[__unnest_placeholder(t.column1)|depth=1] structs[]
04)------Projection: t.column1 AS __unnest_placeholder(t.column1), t.column2
05)--------TableScan: t projection=[column1, column2]
physical_plan
01)SortExec: expr=[column2@1 ASC NULLS LAST], preserve_partitioning=[false]
02)--ProjectionExec: expr=[__unnest_placeholder(t.column1,depth=1)@0 as UNNEST(t.column1), column2@1 as column2]
03)----UnnestExec
04)------ProjectionExec: expr=[column1@0 as __unnest_placeholder(t.column1), column2@1 as column2]
05)--------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/unnest/ordered.parquet]]}, projection=[column1, column2], output_ordering=[column2@1 ASC NULLS LAST], file_type=parquet

# cleanup
statement ok
drop table t;

[vegarsti]

The PR now includes and passes this test case along with some new ones