fix: UnnestExec preserves relevant equivalence properties of input

[vegarsti]

Which issue does this PR close?

• Closes unnest should preserve the input's equivalence properties for uninvolved columns #15231.

What changes are included in this PR?

• In UnnestExec's compute_properties we now construct itsEquivalenceProperties using what we can from the input plan, so that we preserve sort ordering of unrelated columns (and avoid unnecessary sorts further up in the plan).

Are these changes tested?

• Adds test cases to the sqllogictests for UnnestExec in unnest.slt

Are there any user-facing changes?

No

Explanation

Given a struct or array value col, unnest(col) takes the N elements of col and "spreads" these onto N rows, where all other columns in the statement are preserved. Said another way, when we unnest a column we are inserting a lateral cross-join against its elements, which by construction:

• Duplicates every other column once for each array/map element
• Replaces the original collection column with one (or more) “element” columns
• Expands one input row into zero (if empty) or many output rows

E.g. (from unnest.slt):

datafusion/datafusion/sqllogictest/test_files/unnest.slt

Lines 699 to 712 in 6d9b76e

	query III
	select unnest(column1) c1, unnest(column2) c2, column3 c3 from unnest_table group by c1, c2, c3 order by c1, c2, c3;
	----
	1 7 1
	2 NULL 1
	3 NULL 1
	4 8 2
	5 9 2
	6 11 3
	12 NULL NULL
	NULL 10 2
	NULL 12 3
	NULL 42 NULL
	NULL NULL NULL

The EquivalenceProperties struct has three types of properties:

1. equivalence groups (expressions with the same value)
2. ordering equivalence classes (expressions that define the same ordering)
3. table constraints - a set of columns that form a primary key or a unique key

In this PR we construct the UnnestExec node's EquivalenceProperties by using the input plan's equivalence properties for the columns that are not transformed - except for table constraints, which we discard entirely. The reasoning for discarding constraints is that because we're duplicating the other columns across rows, we are invalidating any uniqueness or primary-key constraint. We also need to some twiddling with the mapping of the projection (indices change due to the unnesting).

[vegarsti]

Tagging @alamb, maybe you can trigger CI? 🙏🏻

[vegarsti]

Grammar fix

[asubiotto]

Nice!

[vegarsti]

I discovered EquivalenceProperties.project, which seems to do what we need: We can get the unnest plan's equivalence properties by doing input_eq_properties.project(unnested_columns, schema), i.e. discarding the properties of the columns that are being unnested.

I have updated the PR doing that.

Equivalence properties are

1. equivalence groups (expressions with the same value)
2. ordering equivalence classes (expressions that define the same ordering)
3. table constraints - these can be primary key or unique

I am pretty sure that this takes care of 1 and 2, since we now have no equivalence properties for the columns. I am not yet sure about 3, though - if the original expression uses a column that is a primary key, after the unnest we will have multiple rows with the same column. Does that mean we need to remove that constraint from the eq properties? It kinda sounds like yes, but I need to see exactly what it's being used for.

[vegarsti]

I discovered EquivalenceProperties.project, which seems to do what we need: We can get the unnest plan's equivalence properties by doing input_eq_properties.project(unnested_columns, schema), i.e. discarding the properties of the columns that are being unnested.

I have updated the PR doing that.

Equivalence properties are

1. equivalence groups (expressions with the same value)
2. ordering equivalence classes (expressions that define the same ordering)
3. table constraints - these can be primary key or unique

I am pretty sure that this takes care of 1 and 2, since we now have no equivalence properties for the columns. I am not yet sure about 3, though - if the original expression uses a column that is a primary key, after the unnest we will have multiple rows with the same column. Does that mean we need to remove that constraint from the eq properties? It kinda sounds like yes, but I need to see exactly what it's being used for.

After reading some more I have now updated it so that we remove any constraint from the properties. I've updated the PR description.

I think this is semantically sound now.

FYI @alamb and @asubiotto

[asubiotto]

Nice work! This LGTM, I'll leave it to @alamb for a final review and CI kickoff.

[alamb]

Sorry to bug you again @alamb, do you have time to review this today or tomorrow? 👀 Do let me know if there's anything that I can do to make it easier to review.

Hi @vegarsti - sorry I didn't see this earlier. I will try and review it over the next day or two

Maybe @berkaysynnada or @suremarc has some time to review as well

[vegarsti]

Friendly ping @alamb 😄

[alamb]

Thank you for this contribution @vegarsti

I am sorry for the delayed review -- I am always trying to encourage others to review PRs, but indeed I often function as the reviewer of last resort. Anything you can do to help (like help review PRs yourself) would be most appreciated!

This is definitely the right direction, but when I did some testing of this PR some of the behavior didn't make sense to me

Could you look at the test I provided, as well as add additional cases:

1. That unnest a struct (the one I provided unnests a list)
2. has multiple list/structs unnested (as the code seems to handle such a case)

[vegarsti]

Thank you for this contribution @vegarsti

I am sorry for the delayed review -- I am always trying to encourage others to review PRs, but indeed I often function as the reviewer of last resort. Anything you can do to help (like help review PRs yourself) would be most appreciated!

This is definitely the right direction, but when I did some testing of this PR some of the behavior didn't make sense to me

Could you look at the test I provided, as well as add additional cases:

1. That unnest a struct (the one I provided unnests a list)

2. has multiple list/structs unnested (as the code seems to handle such a case)

Thank you so much for the detailed and gracious review. Thanks for catching the weird behavior, I will address this.

And I am happy to start reviewing PRs!

[alamb]

Marking as draft as I think this PR is no longer waiting on feedback and I am trying to make it easier to find PRs in need of review. Please mark it as ready for review when it is ready for another look

[vegarsti]

Marking as draft as I think this PR is no longer waiting on feedback and I am trying to make it easier to find PRs in need of review. Please mark it as ready for review when it is ready for another look

Indeed, thanks!

[vegarsti]

Figured out why the test @alamb added failed -- the way I was creating the projection mapping was too simplistic, causing indexes to not match. Will add the two requested test cases as well.

[vegarsti]

Figured out why the test @alamb added failed -- the way I was creating the projection mapping was too simplistic, causing indexes to not match. Will add the two requested test cases as well.

Added two similar test cases:

1. with struct
2. with nested array, array, and struct

# cargo test --test sqllogictests -- unnest
    Finished `test` profile [unoptimized + debuginfo] target(s) in 0.19s
     Running bin/sqllogictests.rs (target/debug/deps/sqllogictests-8ad6b462cb0c808e)
Completed 1 test files in 0 seconds

[vegarsti]

Since CI ran on this one, I'll leave it here without updating the branch until this gets reviewed again 👍🏻

[vegarsti]

@berkaysynnada @suremarc @alamb Gentle ping for a review!

[tobixdev]

To me, the changes and tests make sense. Thanks!

CAVEAT: I am by no means a DataFusion pro. Just trying to learn more while providing some feedback. :)

[tobixdev]

I think we have had multiple issues with quadratic planning time for a large amount of columns. I think we could get the same problem here as the contains is another linear scan, thus creating a quadratic runtime depending on the number of columns. I could also be wrong and this doesn't cause an issue.

Maybe we could build a buffer and then simply index into the buffer on whether this is unnested (not tested):

        let input_schema = input.schema();
        let mut unnested_indices = BooleanBufferBuilder::new(input.len());
        unnested_indices.append_n(input.len(), false);
        for list_unnest in list_column_indices {
            unnested_indices.set_bit(list_unnest.index_in_input_schema, true);
        }
        for list_unnest in struct_column_indices {
            unnested_indices.set_bit(*list_unnest, true)
        }
        let unnested_indices = unnested_indices.finish();

        let non_unnested_indices: Vec<usize> = (0..input_schema.fields().len())
            .filter(|idx| !unnested_indices.value(*idx))
            .collect();

Otherwise, I think changing the iterator to (0..input_schema.fields().len()) would help with readability as you don't seem to be using the actual field.

[vegarsti]

Otherwise, I think changing the iterator to (0..input_schema.fields().len()) would help with readability as you don't seem to be using the actual field.

Definitely doing this! Thank you.

Good idea to build a buffer and index into it. I'll give that a shot and see how it turns out!

[tobixdev]

Yeah I'd only change that if its easy to do with similar complexity. I think the quadratic behavior only makes a problem if we have many many columns and most of them use unnest.

[vegarsti]

This worked very well! Added in c42c8c1. I think the buffer approach you gave is more readable as well, so it's win win! Thanks a lot

[tobixdev]

Maybe one additional note: I think the resulting sort properties can be improved for unnesting structs if we know that the struct columns themselves are ordered.

If that makes sense we could also somehow expand the LexSort entry for the struct column.

But as this is already an improvement I think that tracking this in a separate issue is fine.

[vegarsti]

Thank you so much @tobixdev!

[vegarsti]

Maybe one additional note: I think the resulting sort properties can be improved for unnesting structs if we know that the struct columns themselves are ordered.

If that makes sense we could also somehow expand the LexSort entry for the struct column.

But as this is already an improvement I think that tracking this in a separate issue is fine.

Great idea!

[adriangb]

I took a look and it seems all good to me but given there's already been a lot of review on it I think the existing reviewers need to approve for it to be mergeable, so I will defer to them. Consider this my token ✅

[alamb]

Looks like there are some outstanding comments from @tobixdev -- please ping me @vegarsti when you have addressed them @vegarsti and are ready for a final review / stamp

[vegarsti]

Thanks a lot everyone! @alamb ready for the stamp now ;)

[alamb]

Thank you for this contribution @vegarsti

I think this is very close. I think it should have:

1. Some additional tests / comments cleanup (see comments)
2. Avoid unwrap / expect to minimize the severity of symptoms

[alamb]

this plan shows the data being sorted, but the comment suggests it should not be 🤔

Could you please explain in more detail what you expect this explain plan to be showing? Given there is no ORDER BY in the query (or in the OVER clause) it is not clear why this is testing ordering

[alamb]

Could you also please add two additional tests:

1. a negative test case here. order by the output of the unnest and verify that it is in fact sorted correctly
2. A case with the ordering column as the first index (e.g. tuples like (100, [3,2,1], 'a') and then order by 100