[C++] Supporting compute functions on ExtensionTypes

[thisisnic]

Hi folks,

I wanted to open a discussion around supporting compute kernels on ExtensionType Arrays or casting from ExtensionTypes to supported Arrow types. There was a recent request from an R user around this [1], and we’ve supplied them with a workaround, but after discussion with Alenka, we noticed this theme has come up repeatedly.

In the case of simple ExtensionTypes where it would make sense to operate on the storage type or be able to cast to it, it feels like this could be a great win to enhance interoperability, though I see from previous discussions there are questions to resolve around how this would work if a straightforward mapping didn’t exist [2].

There are multiple existing issues around this [3] (thanks @AlenkaF for collating these), and I recognise it could take significant dev time to implement, but if we had an idea of how this would look, perhaps we could break off smaller chunks, and potentially apply for funding to work on this? It feels like it would be a big win for interoperability.

My goals for this discussion:

1. To understand the constraints and unresolved questions around this
2. To establish whether others agree this would be impactful to the project or not
3. If it’s something we could take forward, to get a high-level understanding of how we might want to implement it, before having a more in-depth discussion on the mailing list

Keen to hear your thoughts!

Nic

[1] R arrow cannot handle labelled data in arrow tables · Issue #45601 · apache/arrow · GitHub
[2] [C++][Python] Support ExtensionType arrays in more kernels
[3] List of previous issues on the topic:

R

User request for extension type compute support
R arrow cannot handle labelled data in arrow tables · Issue #45601 · apache/arrow · GitHub

Python

Extension types and compute functions
[C++][Python] Support ExtensionType arrays in more kernels
[C++] Provide a default implementation for ExtensionType::ExtensionEquals
Enhancement Request: Custom Operator Support for PyArrow Extension Types in Compute Functions
[Python] For extension types, compute kernels should default to storage types?

Extension types and PyArrow constructors

[Python] Error using extension types in struct in PyArrow
Extension types not fully supported in list arrays
[Python] Instantiating arrays with type ListType[ExtensionType] is not supported
[Python] Python extension types aren't usable in struct arrays, maybe the same: [Python] pyarrow.Table.cast does not handle extension types which wrap list types
[C++][Python] SEGFAULT when casting FixedSizeTensorArray to storage type then back to FixedSizeTensorArray

[paleolimbot]

Thanks Nic for opening this and Nic/Alenka for collecting these!

I'll also add

• GH-46659: [C++] Fix export of extension arrays with binary view/string view storage #46660 (which fixes an invalid string view/binary view export of extension types)

I'm happy to poke through these as I have time (since many of them arose from Geo extension types) or review of others have funding to do the work!

[AlenkaF]

Thanks Dewey for joining in the discussion!
Would it be possible for us to dig into the topic just enough to get a high-level understanding, break the implementation idea into smaller chunks and then apply for the funding to make this happen? We would need some C++ guidance from the start already and if Dewey, you are OK with taking the lead, I am happy to help!

[thisisnic]

A few notes from points raised in the Arrow dev meeting:

• there is some interest in a "fast path" for implicitly casting extension types to their physical types which would be useful for things like Arrow's variant type
• in principle, you can remove the metadata and treat the column as its physical type (e.g. a JSON extension type just becomes a string). This can be done in-place without copying data. However, problems arise when kernels or functions check types using the metadata and don’t recognize the mapping
• in R, due to optimisation and execution order, suggestions like mutating a dataset before processing might not work as expected due to query optimisation done in R, though this could possibly be fixed at the R wrapper level if we were to build lower-level R bindings to mutate the schema metadata directly as a workaround.
• the extension type feature is incomplete across the ecosystem - not just compute - so there could be some value in consolidating known issues and prioritizing C++ fixes to benefit R and Python
• some cases are particularly complex and this has been seen in other Arrow-compatible ecosystems, i.e. DuckDB now preserves extension metadata when reading Arrow tables, but modifying data (e.g. appending strings) inevitably strips extension types and so recasting is needed after such transformations.

[paleolimbot]

in principle, you can remove the metadata and treat the column as its physical type

Just a note that implicit casting to storage only makes sense for some extension types (although it's appropriate for many, like JSON). Something that DuckDB does which is quite nice with respect to its extension types ("aliases") is the ability to register a cast implementation between two types (which includes the option of whether it is implicit or not). That said the implicit cast to storage is not the end of the world (just allows non-sensical operations to occur that might be clearer as an error, like multiplying an S2_CELL identifier stored as a uint64 by something, since the result is meaningless).

modifying data (e.g. appending strings) inevitably strips extension types

I think this is usually the desired behaviour (i.e., a substring of a JSON item is no longer necessarily JSON?)

My low-priority personal wishlist for extension type functionality in Arrow C++/pyarrow based on my experience in geoarrow-pyarrow would be:

• Ability to register a cast to string (so that my geometry ChunkedArrays and tables are printed more nicely!)
• A compute function to strip extensions (that also works on things that aren't extensions). This is sort of an opt-in version of the implicit cast to storage.
• Ability to register a type2 function (like vctrs::vec_ptype2()) and a cast function (like vctrs::vec_cast()) to support concatenating extension type arrays that don't have identical storage. Variant will probably need this to be able to handle shredded and unshredded versions in the same Dataset.

[AlenkaF]

I think this is great—thanks!
I was initially inclined to open an umbrella issue to track all the ideas, but I’ve decided to keep this discussion as the central place for now. Once we have enough bandwidth to take any of it on, we can revisit this thread and start from here.