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

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

closes Add an ExtensionType to DataType enum #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.

alamb · 2025-01-21T23:17:59Z

Thank you @mbrobbel -- I plan to review this carefully tomorrow

alamb

😮 -- thank you SO much @mbrobbel -- I looked at the structure of this PR and I really like it. It is really nicely done. I have some minor comments, but really I found this PR really well structured, commented and tested 🏅

This also implements the consensus reached on #4472 as I understand it

I have some API design questions, but I also think this is good enough to merge as is and I would be happy to maintain the code as is.

If there are no objections to this approach from other reviewers in theory I will find the time to review the tests a bit more closely before merging, but the ones I have looked at so far looks 👍 👍 👍

@westonpace @kylebarron @tustvold @yukkit perhap you would have some time to review / look over this PR as you were quite active on #4472

@emilk, @danking @jleibs, @crepererum and @findepi as you seem to have mentioned this PR in the past

alamb · 2025-01-22T10:59:57Z

arrow-schema/src/extension/mod.rs

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


Is it ever possible to have an error serializialzing the metadata? If so perhaps this should return Result<Option<String>>

Probably Result<String> is enough, I think None and "" in this case doesn't matter

I can't think of a practical situation where serialization of the extension type metadata would error, but I'm fine with adding a Result here.

The Option is required to support extension types without metadata (they should return None here). Note that this is different from "", because in that case the extension metadata key is added to the field custom metadata, whereas for None this key is skipped.

So with support for serialization errors this would return: Option<Result<String, ArrowError>>.

because in that case the extension metadata key is added to the field custom metadata, whereas for None this key is skipped.

I forget what version of Arrow C++ I fixed this in, but old versions or Arrow C++ (maybe older than ~a year) will segfault if ARROW:extension:metadata is missing when importing via the C Data interface, so it might be worth always serializing an empty string.

Per https://arrow.apache.org/docs/format/Columnar.html#extension-types:

Extension types may or may not use the 'ARROW:extension:metadata' field.

Which I interpret as; it should also work when this key is missing in the field metadata for extension types that don't require metadata.

It's your call

Probably a better way to put that would have been: I don't have a strong feeling about it, but it did take me quite a long time to figure out why geoarrow-c was crashing causing pyarrow to crash 🙂

Yes, I can imagine. Thanks for the suggestion.

Maybe we can add a note in the comments and suggest for maximum compatibility extension types should always return Some("")

This would be in conflict with the docs for deserialize_metadata:

An extension type that defines no metadata should expect None for the serialized metadata and return Ok(()).

This function should return an error when unexpected metadata is set (for extension types without metadata)

We could add a note as you suggested for serialize_metadata and then modify the comments in deserialize_metadata to be suggestions only?

The Uuid canonical extension type requires no metadata, so if we change the recommendation here we should also change that implementation maybe?

At least Arrow C++ and DuckDB will export ARROW:extension:metadata as an empty string if there are no parameters, so something will have to handle that (if they are interested in interacting with other Arrow implementations). In GeoArrow I personally accept "{}", "" or missing metadata.

alamb · 2025-01-22T11:00:18Z

arrow-schema/src/extension/mod.rs

+///
+/// # Example
+///
+/// The example below demonstrates how to implement this trait for a `Uuid`


this is a really nice example and illustrates how this API works

alamb · 2025-01-22T11:07:30Z

arrow-schema/src/extension/mod.rs

+    ///
+    /// If an extension type defines no metadata it should use `()` to indicate
+    /// this.
+    type Metadata;


I wonder what is the usecase for the Metadata indirection?

For instance, the example above seems like it would be simpler if it simply serialized / deserialized UuidVersion directly

I think that would mean the API would look like

trait ExtensionType { // serialize self to a string fn serialize(&self) -> Option<String>; // deserialize into self fn deserialize(metadata: Option<&str>) -> Result<Self, ArrowError>; }

In Json, serde_json::Value is the metadata instead of String, with your design I think we are not able to store metadata other than String

I do think any metadata needs to be serialized to a string (which is the same as what the PR does)

fn serialize_metadata(&self) -> Option<String>;

I am just wondering why have a separate Rust type for Metadata

For eample, if I were doing this my first inclination would be to do

impl ExtensionType for MyAwesomeType { ... type Metadata = Self; .... } So then I could just serialize/deserialize instances of my awesome type

I see. The only benefit is that we can have a separate clean struct IF the metadata becomes pretty complex. For usual case, simply se/de without type Metadata is enough.

This doesn't work for all extension types. For example the FixedShapeTensor and VariableShapeTensor canonical extension types can't be constructed from just the metadata, they also need the field datatype (which explains the ExtensionType::try_new signature).

Ah, got it. Maybe we can add some comments explaining the rationale a bit more.

Or maybe we can default the Metadata type to Self

type Metadata = Self;

To be clear we can do this as a follow on PR (or never). I don't think it is necessary for this PR

Unfortunately associated type defaults are not stabilized yet. I added a note in b78e692.

alamb · 2025-01-22T11:11:50Z

arrow-schema/src/extension/mod.rs

+
+//! Extension types.
+
+#[cfg(feature = "canonical-extension-types")]


💯 for making this behind a feature flag

It would be nice to add this to the documentation:

https://github.com/apache/arrow-rs/tree/main/arrow#feature-flags

Perhaps we could also note it is "experimental" and say we reserve the right to make breaking changes between even minor version releases -- that would give us flexibilty to iterate in the short term if necessary

069642f

Perhaps we could also note it is "experimental"

Do you want to add an experimental feature to the arrow crate (same as in the parquet crate)?

arrow-rs/parquet/Cargo.toml

Lines 112 to 113 in 1664214

# Experimental, unstable functionality primarily used for testing

experimental = []

Or just a note in the module docs?

I think adding a note in the module docs is good enough

Added a warning in 5fec56d.

alamb · 2025-01-22T11:17:20Z

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

+/// Canonical extension types.
+///
+/// <https://arrow.apache.org/docs/format/CanonicalExtensions.html#format-canonical-extensions>
+#[non_exhaustive]


👍 for non exhaustive so we can add new extension types without breaking API changes

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.

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.

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

alamb approved these changes Jan 22, 2025

View reviewed changes

alamb mentioned this pull request Jan 22, 2025

Add an ExtensionType to DataType enum #4472

Open

This was referenced Jan 22, 2025

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

Open

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


		//! Extension types.

		#[cfg(feature = "canonical-extension-types")]

	# Experimental, unstable functionality primarily used for testing
	experimental = []

Add ExtensionType trait and CanonicalExtensionType enum #5822

Are you sure you want to change the base?

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

alamb commented Jan 21, 2025

alamb 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

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

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

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