Add `ExtensionType` trait and `CanonicalExtensionType` enum #5822

mbrobbel · 2024-05-31T13:09:37Z

closes Add support for Arrow Extension Types #4472?

Rationale for this change

It would be nice to better support reading and writing the Arrow canonical uuid and json extension types with the arrow and parquet crate i.e. mapping between the arrow extension type and the parquet logical uuid and json types.

What changes are included in this PR?

This adds an ExtensionType trait, some impls for canonical extension types and a CanonicalExtensionType enum for canonical extension types.

Are there any user-facing changes?

Users can now annotate their fields with extension types, and for uuid and json they are propagated via the arrow writer to map to the parquet uuid and json logical types.

kylebarron · 2024-06-01T07:49:33Z

Maybe ExtensionType could be a trait to be externally implementable and not limited to canonical extension types?

tustvold

This makes sense to me, and seems like an unobtrusive way to provide better ergonomics for extension types.

That being said I've limited exposure to them so getting some broader perspectives might be valuable, perhaps on the mailing list or something?

mbrobbel · 2024-06-18T13:33:21Z

I haven't had time to work on this, but I'm planning to pick this up later.

alamb · 2024-07-01T19:17:30Z

Thanks @mbrobbel -- marking this PR as draft as I think it still has planned but not yet completed work

aykut-bozkurt · 2024-08-10T12:26:11Z

Thanks for the PR. This seems very useful to support not yet mapped logical types. e.g. json

mbrobbel · 2024-09-26T21:27:59Z

~~I updated the PR to define a trait for extension types instead. Ready for another round of feedback.~~
Edit: just realized I need to change some trait methods to make it work for other extension types.

…onst was added in 1.81

mbrobbel · 2025-01-20T19:22:14Z

This is ready for another round of reviews.

I updated the trait and changed some field methods, added implementations for the current set of canonical extension types (behind a feature flag) and added more tests and docs. Maybe this PR is too big now, but I don't mind splitting it up if needed.

Roundtripping through parquet is not implemented in this PR.

emilk · 2025-01-22T11:44:47Z

This looks great to me! I really like this approach

jayzhan211

It looks pretty nice!

kylebarron

Oops forgot to commit changes

kylebarron · 2025-01-22T16:33:17Z

parquet/src/arrow/schema/mod.rs

+        // TODO: roundtrip
+        // let arrow_schema = parquet_to_arrow_schema(&parquet_schema, None)?;
+        // assert_eq!(
+        //     arrow_schema.field(0).try_extension_type::<Json>()?,
+        //     Json::default()
+        // );


Is this left for a future PR?

Yes, this requires a bit more work to propagate correctly.

Filed Support round tripping extension types in parquet #7063 to track

kylebarron · 2025-01-22T16:49:37Z

arrow-schema/src/extension/mod.rs

+///
+///     // We just return a reference to the Uuid version.
+///     fn metadata(&self) -> &Self::Metadata {
+///         &self.0


My first go at implementing this trait I used:

pub struct PointType(CoordType, Dimension); impl ExtensionType for PointType { const NAME: &'static str = "geoarrow.point"; type Metadata = (CoordType, Dimension); fn metadata(&self) -> &Self::Metadata { &(self.0, self.1) } }

But that won't work for this because cannot return reference to temporary value.

So essentially type Metadata needs to be the same as whatever is defined inside the extension type struct, and that needs to be a single item, so that it can return Self::Metadata without creating any new items?

Yes, the current signature requires that. We could change it to return Self::Metadata?

The pattern I used for the canonical extension types with metadata applied to your example:

pub struct PointType(PointTypeMetadata); pub struct PointTypeMetadata(CoordType, Dimension); impl ExtensionType for PointType { const NAME: &'static str = "geoarrow.point"; type Metadata = PointTypeMetadata; fn metadata(&self) -> &Self::Metadata { &self.0 } }

kylebarron · 2025-01-22T17:17:04Z

arrow-schema/src/extension/mod.rs

+    ///
+    /// [`Field`]: crate::Field
+    /// [`Field::metadata`]: crate::Field::metadata
+    fn serialize_metadata(&self) -> Option<String>;


This is interesting because it requires the type to contain all information to construct its Arrow extension metadata. This is not how I've implemented it so far in geoarrow-rs. For example, I define a Point type as parameterized only by its "coordinate type" (i.e. FixedSizeList coordinates or Struct[x: Float, y: Float] coordinates) and its dimension (XY, XYZ, and so on). Separately, I store more metadata, and each array type stores both pieces of data.

So the data type alone is not enough information to serialize the full GeoArrow metadata, which also needs the Coordinate Reference System (CRS). It's a good question whether the CRS should be a part of the type here. Perhaps it should, but then we need to support CRS equality, which is quite tricky to do, so a single logical CRS can be represented in different ways.

cc @paleolimbot , this might be a question for you.

Ah this is great, I'll take a better look at this tomorrow and figure out if the current definition works with GeoArrow.
Note that in the spec extension types are defined for Field, not for Array.

I consider the CRS a property of the type in geoarrow-c and geoarrow-python (which is also how it is implemented in the draft Parquet LogicalType, GeoParquet, and the draft Iceberg spec). My reading of this is that it's compatible with GeoArrow (and will be great!). (A Python implementation is here: https://github.com/geoarrow/geoarrow-python/blob/main/geoarrow-types/src/geoarrow/types/type_pyarrow.py ).

It's true that CRS equality is thorny (generally it's expressed as a probability)...I implement this as part of the type and implement == on the type.

It looks like all the pieces are there (in geoarrow-rs), so I guess it will be straightforward to add ExtensionType implementations that can be used with the added extension type related methods of Field.

Tpt

Just passing by and had a lot of fun reading this code (thank you)!

Please just ignore my comments if you feel so.

Tpt · 2025-01-23T18:36:20Z

arrow-schema/src/extension/canonical/bool8.rs

+impl ExtensionType for Bool8 {
+    const NAME: &'static str = "arrow.bool8";
+
+    type Metadata = &'static str;


Maybe bad idea: use () here. This way it's not possible to encode bad metadata value (any string that is not empty)

With the proposed definition of the trait, () would indicate no metadata and results in the metadata key to be in unset in the field custom metadata map. Note that there is no actual storage for the metadata in this extension type. Both serialization and deserialization make sure the metadata is always an empty string.

Ok! Missed that! Thanks!

Tpt · 2025-01-23T18:38:18Z

arrow-schema/src/extension/canonical/bool8.rs

+                "" => Ok(""),
+                _ => Err(ArrowError::InvalidArgumentError(ERR.to_owned())),
+            },
+        )


slightly shorter:

Suggested change

)

if metadata.map_or(false, str::is_empty) {

Ok("")

} else {

Err(ArrowError::InvalidArgumentError("Bool8 extension type expects an empty string as metadata".into()))

}

Tpt · 2025-01-23T18:46:42Z

arrow-schema/src/extension/canonical/json.rs

+///
+/// This type does not have any parameters.
+///
+/// Metadata is either an empty string or a JSON string with an empty


Based on implementation I guess it's not "an empty string" but "a JSON string with an empty string" (ie "\"\"" and not "")

Good catch! 29a94cb

Maybe bad idea: instead of an Option<Map<...>> what about just using a is_empty_object: bool field or something like this to make invalid states unrepresentable?

Also, you might even hack something like this to only deserialize the empty object:

#[derive(Deserialize)] struct Empty {} let is_empty_struct = serde_json::from_str<Empty>(value).is_ok()

and then use the Option<Empty> type in JsonMetadata

I did consider doing this. The spec states (https://arrow.apache.org/docs/format/CanonicalExtensions.html#json):

Description of the serialization:
Metadata is either an empty string or a JSON string with an empty object. In the future, additional fields may be added, but they are not required to interpret the array.

So I figured, I'll support the future use-case, however, if/when in the future fields are added here it would actually be better to use a strongly typed definition. I like your suggestion of using Empty here and since the metadata content currently is not pub, if fields are added in the future, we can just rename and modify the definition of the inner metadata type without introducing a breaking change.

c6f0443

Great! It might be possible to also use a similar approach for the Bool8 metadata.

There is currently no metadata storage for Bool8, and the difference is that there is no note about future additions to the metadata serialization:

Description of the serialization:
Metadata is an empty string.

So, if there are breaking changes in the spec for this type, it would have to be a breaking change here too anyway.

alamb · 2025-01-24T21:41:45Z

I think we should leave this PR open at least through the weekend to gather additional feedback. I think given the size/scope of it perhaps we can plan to merge it after I cut the release candidate for the next arrow-rs release (mid-late next week) if there are no other concerns / comments.

THank you again @mbrobbel for pushing this along

alamb · 2025-01-31T01:00:36Z

I'll merge this PR in the next day or two unless anyone else has any additional comments or wants more time to comment

alamb · 2025-02-02T11:36:04Z

Now that we hare release 54.1.0 let's do it

Release arrow-rs / parquet minor version 54.1.0 (Jan 2025) #6929

🚀 Thanks again @mbrobbel for seeing this through.

alamb · 2025-02-02T11:36:37Z

I merged up from main to run the CI once more before merging this

alamb · 2025-02-02T12:06:30Z

Epic work -- can't wait to see how this works in the real world

alamb · 2025-02-12T15:32:30Z

For some reason this PR isn't showing up in the release notes 🤔

Update version to 54.2.0 and add CHANGELOG #7110

It is in the release though

Add ExtensionType for uuid and map to parquet logical type

a7244f5

github-actions bot added parquet Changes to the parquet crate arrow Changes to the arrow crate labels May 31, 2024

Fix docs

6b2e7aa

tustvold reviewed Jun 11, 2024

View reviewed changes

alamb marked this pull request as draft July 1, 2024 19:17

mbrobbel added 4 commits September 26, 2024 21:10

Use an ExtensionType trait instead

bdeab9f

Merge branch 'master' into parquet-uuid-schema

315bd92

Fix clippy warnings

8428653

Add type annotation to fix build

7896455

Merge branch 'master' into parquet-uuid-schema

374d017

aykut-bozkurt mentioned this pull request Oct 18, 2024

Support UUID type CrunchyData/pg_parquet#46

Open

Update ExtensionType trait to support more canonical extension types

e35630a

mbrobbel force-pushed the parquet-uuid-schema branch from 6b883c5 to e35630a Compare January 17, 2025 22:49

Add Json support to parquet, schema roundtrip not working yet

0966a0f

mbrobbel changed the title ~~Add ExtensionType for uuid and map to parquet logical type~~ Add ExtensionType trait and CanonicalExtensionType enum Jan 17, 2025

mbrobbel added 7 commits January 18, 2025 00:28

Fix some clippy warnings

f5c06b1

Add explicit lifetime, resolving elided lifetime to static in assoc c…

b602412

…onst was added in 1.81

Merge branch 'main' into parquet-uuid-schema

5cdfa3f

Replace use of deprecated method, mark roundtrip as todo

81594d9

Add more tests and missing impls

bb7c86a

Add missing type annotations

38c7255

Fix doc warning

1a21e96

mbrobbel marked this pull request as ready for review January 20, 2025 18:59

This was referenced Jan 22, 2025

Jan 18, 2025: This week(s) in DataFusion apache/datafusion#14179

Closed

Extension Types apache/datafusion#12644

Open

jayzhan211 approved these changes Jan 22, 2025

View reviewed changes

mbrobbel added 2 commits January 22, 2025 16:02

Add the feature to the arrow crate and use underscores

069642f

Update feature name in parquet crate

8519344

kylebarron reviewed Jan 22, 2025

View reviewed changes

mbrobbel added 2 commits January 22, 2025 21:52

Add experimental warning to extensions module docs

5fec56d

Add a note about the associated metadata type

b78e692

alamb mentioned this pull request Jan 23, 2025

Change ReturnTypeInfo to return a Field rather than DataType apache/datafusion#14247

Open

Tpt reviewed Jan 23, 2025

View reviewed changes

mbrobbel added 4 commits January 24, 2025 12:07

Fix Json canonical extension type empty string metadata

29a94cb

Simplify Bool8::deserialize_metadata

757b041

Use Empty instead of serde_json::Map in JsonMetadata

c6f0443

Use map_or instead of is_some_and (msrv)

75f56a4

the80srobot mentioned this pull request Jan 31, 2025

Let arrow schemas specify the ENUM type from parquet schemas #7050

Open

kylebarron mentioned this pull request Jan 31, 2025

Refactor NativeType to support upstream ExtensionType information geoarrow/geoarrow-rs#987

Open

Merge remote-tracking branch 'apache/main' into parquet-uuid-schema

4c62785

alamb mentioned this pull request Feb 2, 2025

Support round tripping extension types in parquet #7063

Open

alamb merged commit 8baaa8b into apache:main Feb 2, 2025
29 checks passed

alamb added the enhancement Any new improvement worthy of a entry in the changelog label Feb 2, 2025

alamb mentioned this pull request Feb 2, 2025

Minor: add ticket reference for todo #7064

Merged

mbrobbel deleted the parquet-uuid-schema branch February 2, 2025 12:32

alamb mentioned this pull request Feb 10, 2025

Release arrow-rs / parquet minor version 54.1.0 (Jan 2025) #6929

Closed

5 tasks

-        )
+        if metadata.map_or(false, str::is_empty) {
+           Ok("")
+        } else {
+           Err(ArrowError::InvalidArgumentError("Bool8 extension type expects an empty string as metadata".into()))
+        }

Add ExtensionType trait and CanonicalExtensionType enum #5822

Add ExtensionType trait and CanonicalExtensionType enum #5822

Conversation

mbrobbel commented May 31, 2024 • edited by alamb Loading

Rationale for this change

What changes are included in this PR?

Are there any user-facing changes?

kylebarron commented Jun 1, 2024

tustvold left a comment

Choose a reason for hiding this comment

mbrobbel commented Jun 18, 2024

alamb commented Jul 1, 2024

aykut-bozkurt commented Aug 10, 2024

mbrobbel commented Sep 26, 2024 • edited Loading

mbrobbel commented Jan 20, 2025

emilk commented Jan 22, 2025

jayzhan211 left a comment

Choose a reason for hiding this comment

kylebarron left a 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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Tpt left a 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

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

alamb commented Jan 24, 2025

alamb commented Jan 31, 2025

alamb commented Feb 2, 2025

alamb commented Feb 2, 2025

alamb commented Feb 2, 2025

alamb commented Feb 12, 2025

Add `ExtensionType` trait and `CanonicalExtensionType` enum #5822

Add `ExtensionType` trait and `CanonicalExtensionType` enum #5822

mbrobbel commented May 31, 2024 •

edited by alamb

Loading

mbrobbel commented Sep 26, 2024 •

edited

Loading