feat: consume and produce Substrait type extensions #11510

Blizzara · 2024-07-17T11:04:44Z

Which issue does this PR close?

Closes #.

Rationale for this change

#10646 introduced Substrait UDTs for interval types. However as that PR didn't yet write proper Substrait extensions for those types, it's hard to produce matching UDTs from other Substrait producers. Until substrait-io/substrait#665 gets merged and passed through to Java & DataFusion, I'd need a UDT to represent Spark's CalendarInterval / DF's IntervalMonthDayNano.

What changes are included in this PR?

Produce and consume SimpleExtensionDeclarations for types in addition to functions. This makes user-defined types interoperable across different consumers/producers, assuming they utilize matching names for the same UDT.

Also adds most of the support for doing it for type variations, but doesn't use that yet.

Similarly as for the existing support for functions, this doesn't read any actual extension files, so the extension URL references are bogus (max int). However this makes it easier to add that support in future.

Are these changes tested?

Existing unit tests, especially the round trip tests for IntervalMonthDayNano type and literal

Are there any user-facing changes?

Blizzara · 2024-07-17T11:06:03Z

datafusion/substrait/src/variation_const.rs

 /// For [`DataType::Interval`] with [`IntervalUnit::MonthDayNano`].
 ///
 /// [`DataType::Interval`]: datafusion::arrow::datatypes::DataType::Interval
 /// [`IntervalUnit::MonthDayNano`]: datafusion::arrow::datatypes::IntervalUnit::MonthDayNano
-pub const INTERVAL_MONTH_DAY_NANO_TYPE_URL: &str = "interval-month-day-nano";
+pub const INTERVAL_MONTH_DAY_NANO_TYPE_NAME: &str = "interval-month-day-nano";


substrait uses "url" for referring to the yaml file containing the declaration (which we don't have/use in DF yet), "name" is better for referring to the name of the type

Sounds good -- a url to another file sounds a lot like XML .xsd files from back in the day 😱

Blizzara · 2024-07-17T11:17:27Z

fyi @alamb @waynexia

alamb

Thank you for this PR @Blizzara -- I don't entirely follow the uri vs Identifier thing but the tests are passing and I think this is a good step forward

One thing I didn't see in this PR was any tests for the new functionaly, though I see you say "Existing unit tests, especially the round trip tests for IntervalMonthDayNano type and literal"

I think we need to add some tests prior to merging this PR to demonstrate the new functionality to register extension types, unless there is some reason we can't (or maybe they are coming in a follow on PR or something)

I also think it would be good to add some comments to Extensions

alamb · 2024-07-18T20:02:44Z

datafusion/substrait/src/extensions.rs

+use substrait::proto::extensions::SimpleExtensionDeclaration;
+
+#[derive(Default)]
+pub struct Extensions {


Could we please document this struct (I think you have a lot of context on the PR description you could just copy here)

Specifically I think it would help to document what the u32 was / represented (is it a substrait id or something?

Likewise I think we should document the various functions below

added docs in 21654ec

alamb · 2024-07-18T20:03:46Z

datafusion/substrait/src/extensions.rs

+
+        for (t_anchor, t_name) in val.types {
+            let type_extension = ExtensionType {
+                extension_uri_reference: u32::MAX, // We don't register proper extension URIs yet


Should we maybe file a ticket to track adding URIs?

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

datafusion/substrait/src/logical_plan/consumer.rs

-            None => not_impl_err!("Cannot parse empty extension"),
-        })
-        .collect::<Result<HashMap<_, _>>>()?;
+    let extensions = Extensions::try_from(&plan.extensions)?;


Very nice 👨‍🍳 👌

alamb · 2024-07-18T20:05:54Z

datafusion/substrait/src/variation_const.rs

 /// For [`DataType::Interval`] with [`IntervalUnit::MonthDayNano`].
 ///
 /// [`DataType::Interval`]: datafusion::arrow::datatypes::DataType::Interval
 /// [`IntervalUnit::MonthDayNano`]: datafusion::arrow::datatypes::IntervalUnit::MonthDayNano
-pub const INTERVAL_MONTH_DAY_NANO_TYPE_URL: &str = "interval-month-day-nano";
+pub const INTERVAL_MONTH_DAY_NANO_TYPE_NAME: &str = "interval-month-day-nano";


Sounds good -- a url to another file sounds a lot like XML .xsd files from back in the day 😱

alamb · 2024-07-18T20:07:05Z

datafusion/substrait/src/variation_const.rs

@@ -55,6 +55,7 @@ pub const DECIMAL_256_TYPE_VARIATION_REF: u32 = 1;
 /// [`DataType::Interval`]: datafusion::arrow::datatypes::DataType::Interval
 /// [`IntervalUnit::YearMonth`]: datafusion::arrow::datatypes::IntervalUnit::YearMonth
 /// [`ScalarValue::IntervalYearMonth`]: datafusion::common::ScalarValue::IntervalYearMonth
+#[deprecated]


Perhaps we can add a note here that they were deprecated in version 41 and leave a suggestion of what to do with them instead?

Something like

datafusion/datafusion/sqllogictest/test_files/aggregate.slt

Line 186 in d67b0fb

# so they are covered in `datafusion/physical-expr/src/aggregate/array_agg_distinct.rs`

Blizzara

Thanks for the review @alamb - I added docs in 21654ec and tests in 1042941, and filed two followup issues #11544 and #11545

Blizzara · 2024-07-19T11:57:29Z

datafusion/substrait/src/extensions.rs

+use substrait::proto::extensions::SimpleExtensionDeclaration;
+
+#[derive(Default)]
+pub struct Extensions {


added docs in 21654ec

alamb

LGTM -- thank you very much @Blizzara

* support reading type extensions in consumer * read extension for UDTs * support also type extensions in producer * produce extensions for MonthDayNano UDT * unify extensions between consumer and producer * fixes * add doc comments * add extension tests * fix * fix docs * fix test * fix clipppy

Blizzara added 5 commits July 17, 2024 10:57

support reading type extensions in consumer

1577d13

read extension for UDTs

660f04e

support also type extensions in producer

2ae2c8e

produce extensions for MonthDayNano UDT

4ca6da0

unify extensions between consumer and producer

54c9ec3

Blizzara marked this pull request as ready for review July 17, 2024 11:04

github-actions bot added the substrait Changes to the substrait crate label Jul 17, 2024

Blizzara commented Jul 17, 2024

View reviewed changes

fixes

4b0be1d

Blizzara force-pushed the avo/substrait-type-extensions branch from 3f6416d to 4b0be1d Compare July 18, 2024 13:33

alamb reviewed Jul 18, 2024

View reviewed changes

Blizzara mentioned this pull request Jul 19, 2024

Use SimpleExtensions for Substrait type variations #11544

Open

Blizzara added 2 commits July 19, 2024 11:48

add doc comments

21654ec

add extension tests

1042941

Blizzara commented Jul 19, 2024

View reviewed changes

Blizzara added 2 commits July 19, 2024 14:21

fix

99a9869

fix docs

b5a4c1a

Blizzara mentioned this pull request Jul 19, 2024

feat: support Map literals in Substrait consumer and producer #11547

Merged

Blizzara added 2 commits July 19, 2024 14:38

fix test

f328c4a

fix clipppy

493570d

alamb approved these changes Jul 19, 2024

View reviewed changes

alamb merged commit a4c9bb4 into apache:main Jul 19, 2024
23 checks passed

Blizzara deleted the avo/substrait-type-extensions branch July 19, 2024 21:19

feat: consume and produce Substrait type extensions #11510

feat: consume and produce Substrait type extensions #11510

Uh oh!

Conversation

Blizzara commented Jul 17, 2024

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?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Blizzara commented Jul 17, 2024

Uh oh!

alamb left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Blizzara left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

alamb left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Blizzara left a comment •

edited

Loading