Support Substrait's VirtualTables #10531

Blizzara · 2024-05-15T16:51:56Z

Which issue does this PR close?

This (well, now together with the earlier PRs for chunks of this) is my first (but hopefully not last!) contribution to DataFusion, and I'm pretty new with Rust as well - so all feedback on both style and details is very much welcome!

Rationale for this change

What changes are included in this PR?

Adds support for Substrait's VirtualTables, ie. tables with data baked-in into the Substrait plan instead of being read from a source.

Adds conversion in both ways (Substrait -> DataFusion and DataFusion -> Substrait)

Adds support for both empty VirtualTables (mapped into LogicalPlan::EmptyTable) and those containing values (mapped into LogicalPlan::Values).

Are these changes tested?

Tested through round-trip tests, similar to other Substrait support.

Are there any user-facing changes?

More support through Substrait

datafusion/substrait/src/logical_plan/consumer.rs

Blizzara · 2024-05-15T17:02:16Z

datafusion/substrait/src/logical_plan/consumer.rs

@@ -1277,6 +1407,56 @@ pub(crate) fn from_substrait_literal(lit: &Literal) -> Result<ScalarValue> {
                s,
            )
        }
+        Some(LiteralType::List(l)) => {


I don't necessarily need these for this PR, so if it's preferrable I'm happy to split them into separate PRs - I just added them while coding to confirm that lists and structs inside lists are handled correctly.

Same applies for the ones in producer.rs

If it is convenient, using separate PRs would make the review process easier. Perhaps one PR could focus on the virtual table, while the other two could address the list and struct separately.

Sure, will do!

Split those into #10615, #10622, #10640

Blizzara · 2024-05-15T17:03:47Z

datafusion/substrait/tests/cases/roundtrip_logical_plan.rs

+    let ctx = create_context().await?;
+    let row1 = vec![lit(1), lit("a")];
+    let row2 = vec![lit(2), lit("b")];
+    let plan = LogicalPlanBuilder::values(vec![row1, row2])?.build()?;


I didn't find VALUES in the DataFusion SQL - is it not implemented yet?

It is implemented, and you can use it like this.

roundtrip("VALUES (1, 'a', [[1,2,3], []], STRUCT(true, 1))").await

Ah, I'm blind, thanks!

I tried switching to that, but it produces somewhat ugly plans as it uses make_array and struct as UDFs, making the assert fail:

left: "Values: (Int64(1), Utf8(\"a\"), List([[-213.1, , 5.5, 2.0, 1.0], []]) AS make_array(make_array(Float64(-213.1),NULL,Float64(5.5),Float64(2),Float64(1)),make_array()), Struct({c0:true,c1:1}) AS struct(Boolean(true),Int64(1))), (Int64(NULL), Utf8(NULL), List(), Struct({c0:,c1:}))" right: "Values: (Int64(1), Utf8(\"a\"), List([[-213.1, , 5.5, 2.0, 1.0], []]), Struct({c0:true,c1:1})), (Int64(NULL), Utf8(NULL), List(), Struct({c0:,c1:}))"

The AS aliases it produces seem to get ignored (the column names are still column1, column2, column3, ..), so it doesn't seem to make sense to map them into Substrait (and I'm not even sure if that'd be possible). Any ideas how to go around that?

I agree. For Values plan, these aliases should be able to be ignored.
I think we could use assert_expected_plan during testing to ignore aliases, for example:

#[tokio::test] async fn roundtrip_values() -> Result<()> { assert_expected_plan( "VALUES (1, 'a', [[1,2,3], []], STRUCT(true, 1))", "Values: (Int64(1), Utf8(\"a\"), List([[1, 2, 3], []]), Struct({c0:true,c1:1}))", ) .await }

Additionally, aliases should be ignored when handling LogicalPlan::Values in to_substrait_rel.

Nice, that works - 18b7f08.

alamb · 2024-05-17T00:31:09Z

Thanks @Blizzara -- I started the CI on this PR

Blizzara · 2024-05-21T10:16:18Z

@jonahgao @alamb I think this is ready by me :)

jonahgao · 2024-05-21T14:11:23Z

datafusion/substrait/src/logical_plan/consumer.rs

-        Some(LiteralType::Null(ntype)) => from_substrait_null(ntype)?,
+        Some(LiteralType::List(l)) => {
+            let element_dt: Option<&DataType> = match dt {
+                Some(DataType::List(l)) => Ok(Some(l.data_type())),


I am wondering if dt is redundant here and can be removed. We still need to check whether it is a LargeList, just as we check type_variation_reference below.

Is it possible to only extract element_dt from elements[0]? I guess that elements are not empty because substrait has a special type called EmptyList.

pub struct List { /// A homogeneously-typed list of one or more expressions that form the /// list entries. To specify an empty list, use Literal.empty_list /// (otherwise type information would be missing). #[prost(message, repeated, tag="1")] pub values: ::prost::alloc::vec::Vec<super::super::Expression>, }

Yeah, we already do that (look at the first element) if dt is not given. I think the reason I included dt was because Substrait types don't include the field names for structs, so not using dt would mean using something like names and name_idx instead (like in from_substrait_type_with_names).

Thank you for your explanation @Blizzara . My concern is that using two types for List simultaneously is a bit verbose unless it is necessary.

I tested the following two cases: the first one panicked, and the second one also failed.

#[tokio::test] async fn roundtrip_list_iteral() -> Result<()> { roundtrip("SELECT [] from data").await } #[tokio::test] async fn roundtrip_struct_iteral() -> Result<()> { roundtrip("select struct(1 as name0, 3.14 as name1, 'e', true as name3) from data") .await }

---- cases::roundtrip_logical_plan::roundtrip_list_iteral stdout ---- thread 'cases::roundtrip_logical_plan::roundtrip_list_iteral' panicked at consumer.rs:1425:41: index out of bounds: the len is 0 but the index is 0

I think lists and structs are more complex than the virtual table, so it's better for them to have separate PRs.

This was because I wasn't handling empty lists properly (producer created Lists rather than EmptyList) - that was fixed in the Lists PR. Thanks for catching!

jonahgao · 2024-05-21T14:25:46Z

datafusion/substrait/src/logical_plan/consumer.rs

+                    .expect("expected schema to exist for virtual table");
+
+                let fields = from_substrait_named_struct(&base_schema)?;
+                let schema = DFSchemaRef::new(DFSchema::try_from(Schema::new(fields))?);


I suspect that the schema here may not be necessary, perhaps it can be restored through LogicalPlanBuilder::values.

It probably could, except for the column names - LogicalPlanBuilder::values names everything to just column1 etc. Do you see some benefit to using that (and then handling the column names somehow) instead of passing in the schema explicitly?

I once wanted to remove the schema, but we need to use it to pass the struct's names, so the current approach should be fine.

Blizzara · 2024-05-23T16:28:30Z

@jonahgao Thanks for the review - I pushed a new version but it builds on top of #10622 and #10640 so will need those merged and this rebased first before this PRs looks clean.

This is mostly what it used to be with one more major change: I removed passing the datatype into from_substrait_literal and instead do the same thing with names as I did for from_substrait_type. That maybe seems a bit cleaner (even though I dislike the name handling, but I don't see a better way to deal with it based on how it works in Substrait :/)

Adds support for Substrait's VirtualTables, ie. tables with data baked-in into the Substrait plan instead of being read from a source. Adds conversion in both ways (Substrait -> DataFusion and DataFusion -> Substrait) and a roundtrip test.

Blizzara · 2024-05-24T14:57:25Z

@jonahgao @alamb this last one is ready now too :)

jonahgao · 2024-05-25T13:25:57Z

datafusion/substrait/src/logical_plan/producer.rs

+                }
+                Ok(names)
+            }
+            DataType::List(l) => names_dfs(l.data_type()),


I think we also need to consider DataType::LargeList here because to_substrait_type supports it.

Thanks, fixed 9c83ba4

Is there a way to create a LargeList using VALUES for testing, without writing a billion values in there? 😅

arrow_cast is a solution .

VALUES(arrow_cast([1,2], 'LargeList(Int64)'));

Nice! I added that in 4b8babe, though it seems that arrow_cast doesn't yet support structs, so the test still wouldn't have caught this. Unless you can think of some workaround? I'm not familiar at all with arrow_cast and the like 😅

I think so too, let's keep it as a TODO.

jonahgao · 2024-05-25T13:32:03Z

datafusion/substrait/tests/cases/roundtrip_logical_plan.rs

+
+#[tokio::test]
+async fn roundtrip_empty_relation() -> Result<()> {
+    roundtrip("SELECT * FROM (VALUES (1)) LIMIT 0").await


I've slightly modified this test and printed out the schema on the consumer side, but it seems to be incorrect.

#[tokio::test] async fn roundtrip_empty_relation() -> Result<()> { roundtrip("SELECT * FROM (VALUES ([STRUCT(1 as a)], 2)) LIMIT 0").await }

plan2 schema: DFSchema { inner: Schema { fields: [ Field { name: "column1", data_type: List(Field { name: "item", data_type: Struct([Field { name: "c0", data_type: Int64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }]), nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }), nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }, # Its name should be "column2" Field { name: "a", data_type: Int64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }], metadata: {} }, field_qualifiers: [None, None], functional_dependencies: FunctionalDependencies { deps: [] } }

We might need a test to verify the schemas. For this test, the plan str on both sides would just be EmptyRelation.

Good catch! Fixed in 24d07c9 and added schema verification to most tests

Also adds roundtrip schema assertions for cases where possible

…al-tables

jonahgao

Looks good to me👍. Thank you for your contribution @Blizzara .

* Add support for Substrait VirtualTables Adds support for Substrait's VirtualTables, ie. tables with data baked-in into the Substrait plan instead of being read from a source. Adds conversion in both ways (Substrait -> DataFusion and DataFusion -> Substrait) and a roundtrip test. * fix clippy * Add support for empty relations * Fix consuming Structs inside Lists and Structs Also adds roundtrip schema assertions for cases where possible * Rename from_substrait_struct -> from_substrait_struct_type for clarity * Add DataType::LargeList to to_substrait_named_struct * cargo fmt --all * Add validation that names list matches schema exactly * Add a LargeList into VALUES test

github-actions bot added the substrait label May 15, 2024

Blizzara commented May 15, 2024

View reviewed changes

datafusion/substrait/src/logical_plan/consumer.rs Outdated Show resolved Hide resolved

Blizzara commented May 15, 2024

View reviewed changes

Blizzara marked this pull request as ready for review May 16, 2024 16:24

jonahgao reviewed May 21, 2024

View reviewed changes

This was referenced May 22, 2024

Add support for Substrait List/EmptyList literals #10615

Merged

Add support for Substrait Struct literals and type #10622

Merged

More properly handle nullability of types/literals in Substrait Blizzara/datafusion#1

Closed

Blizzara force-pushed the avo/substrait-virtual-tables branch from eb98045 to 7ec7d42 Compare May 23, 2024 16:22

Blizzara mentioned this pull request May 23, 2024

More properly handle nullability of types/literals in Substrait #10640

Merged

Blizzara marked this pull request as draft May 23, 2024 16:28

Blizzara force-pushed the avo/substrait-virtual-tables branch 4 times, most recently from 01ca26f to 908e533 Compare May 24, 2024 14:37

Add support for Substrait VirtualTables

02e7bc0

Adds support for Substrait's VirtualTables, ie. tables with data baked-in into the Substrait plan instead of being read from a source. Adds conversion in both ways (Substrait -> DataFusion and DataFusion -> Substrait) and a roundtrip test.

Blizzara force-pushed the avo/substrait-virtual-tables branch from 908e533 to 02e7bc0 Compare May 24, 2024 14:56

Blizzara marked this pull request as ready for review May 24, 2024 14:56

Blizzara added 2 commits May 24, 2024 18:17

fix clippy

b5b4430

Add support for empty relations

3f40ecf

jonahgao reviewed May 25, 2024

View reviewed changes

Blizzara added 2 commits May 27, 2024 11:34

Fix consuming Structs inside Lists and Structs

24d07c9

Also adds roundtrip schema assertions for cases where possible

Rename from_substrait_struct -> from_substrait_struct_type for clarity

e40f721

Blizzara added 5 commits May 27, 2024 12:01

Add DataType::LargeList to to_substrait_named_struct

9c83ba4

Merge remote-tracking branch 'upstream/main' into avo/substrait-virtu…

6b580ce

…al-tables

cargo fmt --all

0680b4f

Add validation that names list matches schema exactly

d7821bc

Add a LargeList into VALUES test

4b8babe

jonahgao approved these changes May 27, 2024

View reviewed changes

alamb merged commit 2762754 into apache:main May 28, 2024
23 checks passed

Blizzara deleted the avo/substrait-virtual-tables branch May 28, 2024 15:00

alamb mentioned this pull request May 28, 2024

DataFusion weekly project plan (Andrew Lamb) - May 27, 2024 #10699

Closed

9 tasks

Blizzara mentioned this pull request May 29, 2024

[EPIC] Tracking issue of support substrait logical plan #8149

Open

36 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support Substrait's VirtualTables #10531

Support Substrait's VirtualTables #10531

Blizzara commented May 15, 2024 •

edited

Loading

Blizzara May 15, 2024 •

edited

Loading

jonahgao May 21, 2024

Blizzara May 21, 2024

Blizzara May 23, 2024

Blizzara May 15, 2024

jonahgao May 20, 2024

Blizzara May 20, 2024

jonahgao May 21, 2024 •

edited

Loading

Blizzara May 21, 2024

alamb commented May 17, 2024

Blizzara commented May 21, 2024

jonahgao May 21, 2024

Blizzara May 21, 2024

jonahgao May 22, 2024

Blizzara May 23, 2024

jonahgao May 21, 2024

Blizzara May 21, 2024

jonahgao May 22, 2024

Blizzara commented May 23, 2024

Blizzara commented May 24, 2024

jonahgao May 25, 2024

Blizzara May 27, 2024

Blizzara May 27, 2024

jonahgao May 27, 2024

Blizzara May 27, 2024

jonahgao May 27, 2024

jonahgao May 25, 2024

jonahgao May 25, 2024

Blizzara May 27, 2024

jonahgao left a comment

Support Substrait's VirtualTables #10531

Support Substrait's VirtualTables #10531

Conversation

Blizzara commented May 15, 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?

Blizzara May 15, 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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jonahgao May 21, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb commented May 17, 2024

Blizzara commented May 21, 2024

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

Blizzara commented May 23, 2024

Blizzara commented May 24, 2024

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

jonahgao left a comment

Choose a reason for hiding this comment

Blizzara commented May 15, 2024 •

edited

Loading

Blizzara May 15, 2024 •

edited

Loading

jonahgao May 21, 2024 •

edited

Loading