API to register behavior for Extension Types

[alamb]

Is your feature request related to a problem or challenge?

This is is part of implementing LogicalTypes / Extension Types in DataFusion, as described by @findepi

• Support Extension Types / User Defined Types in DataFusion #12644

ExtensionTypes are defined using the metadata on an arrow Field (not the DataType) and stored physically as one of the existing arrow types. This system has the nice benefit that extension types can be processed (passed through) by systems that don't support them as their underlying arrow type, and then additional semantics added by systems that do.

As people continue using DataFusion to implement more sophisticated extension types such as geometry and geography (@paleolimbot) and Variant @friendlymatthew ), they are finding is important to customize certain operations that are currently hardcoded in DataFusion based on physical type.

Some example of operations where special semantics are sometimes needed for extension types

1. printing / displaying values (e.g. printing Variant values in a JSON like manner rather than their raw bytes)
2. casting values to/from other types
3. Comparing values (e.g. it is not correct to compare two variant values byte by byte)

There are a few challenges challenges now:

1. Extension type information is carried on Field (rather than DataType), and the Field is not yet available everywhere (though @paleolimbot and others are working on this)
2. Even once we have Field available everywhere, the callsites for many cast/print and binary operations call directly into the arrow kernels which have no way to customize behavior for extension types.

Describe the solution you'd like

I think we need some sort of DataFusion API for users of extension types to specify and customize their behavior.

Describe alternatives you've considered

One possibility is to add a DFExtensionType trait, that extends the exiting ExtensionType trait, similar to DFSchema

Maybe something like

/// DataFusion Extension Type support
pub trait DFExtensionType: ExtensionTrait {
  /// Cast a column of this extension type to the target
  fn cast(&self, input: ColumarValue, output_type: &Field) -> Result<ColumnarValue>;
  // .. other functions ...
}

We would also need some way to register these types dynamically with the SessionContext as well as pass along the registry into the places they are needed.

let ctx = SessionContext::new();
ctx.register_extension(Arc::new(DFVariantExtension));
...

I am not quite sure if this is the right API, we would need to try it out probably

Additional context

No response

[alamb]

@tobixdev has a great suggestion on #12644 (comment)

Having something Arc is nice because it limits the points at which a "registry" is necessary.

That's a valid point. If we further believe that all the metadata handling should be attached to Field, maybe the field should not directly have a physical DataType. More like something akin to:

enum FieldDataType {
Native(DataType),
Extension(Arc)
}
This could i) allow arrow kernels to keep on focusing on DataType to interpret the underlying physical layout (which I think was the primary reason for rejecting the DataType::Extension(...) solution and ii) prevent accidentally using the wrong implementation of, for example, the + operator for an extension type as the type information stems from the Field. Again, a huge breaking change but I think this cannot be prevented. We could offer a utility functions that allows ignoring any logical type to ease the migration.

Any thoughts on that?

This would be a similar approach to TypeSignature(https://github.com/apache/datafusion/blob/main/datafusion/common/src/types/logical.rs) .

[tobixdev]

I am currently working on a prototype for this feature. However, the last and next few days are rather full. I hope to have a draft in the near future.

However, if someone has more time and wants to expedite the process, please feel free to do so! The faster we set the stage for additional functionality based on extension types, the better. 🚀

[paleolimbot]

Thank you for writing this up, and thank you @tobixdev for #15106 and the reviews along the way (that PR I think hits the mark in a number of ways!)

I wonder if a low impact way and reasonably satisfying way to start would be:

// Maybe the lifetimes will be too annoying here, but the idea is that this can cheaply
// represent any of the combinations of ArrayRef/FieldMetadata/FieldRef/ScalarValue/DataType
// without cloning anything. This is basically nanoarrow's ArrowSchemaView :)
pub struct SerializedTypeView<'a, 'b, 'c> {
    arrow_type: &'a DataType,
    extension_name: Option<&'b str>,
    extension_metadata: Option<&'c str>,
}

pub trait TypeExtensions {
    // None means just use the storage type implementation. Maybe Box<> or &'static could work here
    fn pretty_print_extension(&self, extension_name: &str) -> Option<Arc<dyn PrettyPrintExtension>> { None }
    // ...a future version could also return an Option<Arc<dyn CustomOrdering>> from #18124 
    // ...and in some magical future maybe we can just do
    fn logical_type(&self, type_view: &SerializedTypeView) -> Result<Arc<dyn ThePerfectLogicalTypeTrait>>;
    // ...where the LogicalType knows how to create LogicalArray/Scalars that can do this stuff on their own
    // with convenient APIs so people actually use them.
}

pub trait PrettyPrintExtension {
    // Probably there's a more established pretty print API.
    fn pretty_print_serialized(&self, type_view: &SerializedTypeView, storage: &ArrayRef, options_of_some_kind) -> Result<ArrayRef>; 
    // ...a separate trait allows this API to evolve without affecting TypeExtensions
}

It might be particularly satisfying to implement that for variant so that queries against the test files in the CLI show nice pretty JSON. It would be less satisfying, but possibly easier, to implement that for UUID. It also makes it so we don't have to come up with ThePerfectLogicalTypeTrait because that is hard and right now we represent type information in a lot of different ways.

[tobixdev]

One possibility is to add a DFExtensionType trait, that extends the exiting ExtensionType trait, similar to DFSchema

We have one problem here that ExtensionType is not dyn-compatible due to the use of an associated constant and an associated type. I believe we need dyn-compatibility as our registry will look something like Arc<dyn ExtensionTypeThingey>.

I've experimented a bit more with this and ran into the problem that we have been discussing earlier with custom printing. If I want to define a custom string representation for a type, I need to have access to that in the printing logic. Currently, this is happening in the respective Debug/Display implementations and they do not have access to a registry. Therefore, from my perspective, there are only two sane options:

1. Use some kind of pretty-printing visitor that has access to the registry
2. Directly store the extension in the Field or a possible DFField.

I think 2. would be the better approach but I may be mistaken. Here, a problem is that the "use Field instead of DataType" approach would likely not be enough if Field cannot provide access to a dyn ExtensionTypeThingey. This would require us to make another round of "replace Field with PowerfulField" which I think we want to avoid if possible.

So, can we use Field as a carrier for our DataFusion extension type trait? Not in its current form but I've created an example of how it might could look:

https://github.com/apache/arrow-rs/compare/main...tobixdev:arrow-rs:crazy-field-experiment?expand=1

Code that uses arrow-rs can then provide their own enriched extension type traits without managing this themselves. It has of course multiple drawbacks:

• Downcasting can create problems that a DFDataType enum would prevent.
• I got an error with unwind safety in the tests due to the dynamic dispatch (see diff).
• Less ergonomic and complex as another trait exists

But this would probably be the approach of least resistance (from the API breakage perspective). Any thoughts on that @alamb @paleolimbot ?

I think for Option 2. the other way would be to have a DFField and a DFDataType. I think this could also be fine if we use that in DFSchema.

@paleolimbot Thanks for your input 👍 ! I am currently exploring directly integrating the dyn ExtensionTypeThingey in the expressions and plans. If we decide to go with another option (pass registry or TypeExtensions) this could be an interessting approach!

[paleolimbot]

Cool!

If I want to define a custom string representation for a type, I need to have access to that in the printing logic.

Luckily I think this one might be a fairly isolated problem...error messages and explain plans are mostly where these show up and we can do a solidly OK job at rendering extension types without a registry (arrow.fixed_size_tensor({stuff}) < storage >). Not ideal but gets the point across. Type equality is a bit more important I think but we can also do a solidly OK job with this (see PartialEq and Display for the SerializedTypeView in #18277). The perfect solution would obviously dispatch dynamically for both of those.

We have one problem here that ExtensionType is not dyn-compatible due to the use of an associated constant and an associated type.

I also ran into this in apache/arrow-rs#7398 ...it's possibly useful for helping to fill in some of the boilerplate but I didn't find it all that helpful for customizing the behaviour I was interested in customizing.

I got an error with unwind safety in the tests due to the dynamic dispatch (see diff).

I think I solved that one with pub trait DynExtensionType: Debug + RefUnwindSafe { ... } but I'm not an expert here and maybe that's a separate issue.

But this would probably be the approach of least resistance (from the API breakage perspective). Any thoughts on that ?

I have historically had bad luck guessing what will and won't meet resistance in arrow-rs and have found it more enjoyable to focus on solutions elsewhere. For example, much of what we're trying to do here can be done with logical plan transformations or a separate UDF framework (which is what we do in SedonaDB).

I am currently exploring directly integrating the dyn ExtensionTypeThingey in the expressions and plans. If we decide to go with another option (pass registry or TypeExtensions) this could be an interessting approach!

This sounds hard but I look forward to seeing what you come up with!

[tobixdev]

I have historically had bad luck guessing what will and won't meet resistance in arrow-rs and have found it more enjoyable to focus on solutions elsewhere. For example, much of what we're trying to do here can be done with logical plan transformations or a separate UDF framework (which is what we do in SedonaDB).

apache/arrow-rs#8730

I've tried to write together an issue. Feel free to provide additional context.

[tobixdev]

A reply to apache/arrow-rs#8730 (comment) as some of the issues seem to be more fitting to this issue.

Firstly, thank you very much for your input!

I think the core of the issue is that almost everybody who types DataType or ArrayRef intends for this concept to capture both the logical and physical data type, which is an easy mistake to make given that the enum includes Timestamp, Duration, and Interval and the meaning of the DataType structures in Spark/Arrow C++/go/pyarrow/Java.

For SedonaDB, the appearance of DataType or ArrayRef in arrow-rs or DataFusion is usually a bug, a feature I can't use, or something I have to rewrite (some recent rewrites include the table formatter and the Python UDF framework, which both build on non-extensible arrow-rs APIs built on the DataType and ArrayRef). The places where arrow kernels focus only on the physical data layout are places where it is very easy for DataFusion to introduce bugs for us: when +, - or cast to operate only on Physical + Datetime portion of the type these are correctness errors for our engine.

The recent replacement of FieldRef for DataType is problematic in the other direction...rather than store too little information, it stores too much. The name, nullability, and non-extension metadata sometimes matter and sometimes don't. This usually shows up as schema mismatch errors and/or confusion over how exactly how to propagate the information (e.g., there are two PRs implementing a logical Cast in the DataFusion logical plan that do very different things with all three of those!). These end up affecting everybody, not just users of extension types.

That's a very interesting point! Would using the existing LogicalType instead of Field resolve this problem? There we don't distinguish between, for example, Utf8 and Utf8View.

[...] nor am I sure whether the propagation of the extensions member from Field instance to Field instance will happen in a sufficient number of places to be useful [...]

I think this is a fair critique point here. IMO, at some utopian future, having validations in-place that check whether the stored logical type matches the type metadata could prevent many issues. Nevertheless, this will make problems until the usage of logical types is applied consistently. Maybe some workarounds such as an initial "optimizer pass" that sets the correct logical types can help to ease the migration.

[...] (probably an extension collection in our SessionContext or ConfigOptions would have a better chance of being applied consistently to all fields).

I think the extension collection in SessionContext will be central to any design of extension types. What I am exploring is that you can, in addition to the collection in SessionContext, pass around the reference to the logical type within Exprs and Fields etc. I guess this is similar to how we have a central UDF registry but also carry the reference to the appropriate UDF for each Expr::ScalarFunction which allows us to work with the scalar function call (e.g., printing) without referencing the registry again.

I wonder if implementing specific extendable functionality may be a way to start. The problem of supporting (combinations of) extension types in the pretty printing of batches (including dictionaries, lists, and structs that contain them) highlights a few of these issues, as does Schema::try_merge() (variant in particular will be an issue for that one). I think those have broad applicability beyond DataFusion although DataFusion would certainly beneft.

In my head, the idea is that the logical type reference will provide access to these extensions. For example, the PrettyPrintExtension from above can be retrieved from a particular LogicalType and then be used to print Array values and ScalarValues. Then there would be two ways of retrieving the extension: Use a "regular" Field (or similar), obtain the logical type from the registry and call the printer (this would always be an "escape hatch" option) or if the pointer to the LogicalType is stored inline, skip the call to the registry.

To be concrete there is some design sketch below. Please point out any flaws in that design if you catch anything. I am really unsure whether this can cater to all the needs.

Extend the LogicalType:

pub trait LogicalType: {
    // ....

    /// Returns a reference to the pretty printer for this type.
    fn pretty_printer(&self) -> &dyn ValuePrettyPrinter;

   // Other extensions will follow (e.g., built-in operators, comparisons, sorting, etc.). These may not directly implement the operations but feed the necessary information to DataFusion (similar to the `PlanningInformation` of the custom sorting PR).
}

The printer extensions (the function signatures are subject to change):

/// Implements pretty printing for a set of types.
///
/// For example, the default pretty-printer for a byte array might not be adequate for a UUID type,
/// which is physically stored as a fixed-length byte array. This extension allows the user to
/// override the default pretty-printer for a given type.
pub trait ValuePrettyPrinter {
    /// Pretty print a scalar value.
    ///
    /// # Error
    ///
    /// Will return an error if the given `value` is not supported by this pretty printer.
    fn pretty_print_scalar(&self, value: &ScalarValue) -> Result<String>;

    /// Pretty print a specific value of a given array.
    ///
    /// # Error
    ///
    /// Will return an error if the `array` value at `index` is not supported by this pretty printer.
    fn pretty_print_array(&self, array: &dyn Array, index: usize) -> Result<String> {
        let value = ScalarValue::try_from_array(array, index)?;
        self.pretty_print_scalar(&value)
    }
}

And then in impl Display for Expr:

Expr::Literal(v, metadata) => {
    let logical_type = metadata.logical_type(); // Assumes that this is now stored somewhere in-line. 
    let pretty_printer = logical_type.pretty_printer(); // Will return a default implementation for `NativeType`.
    let v = pretty_printer.pretty_print_scalar(v);

    match metadata.as_ref().map(|m| m.is_empty()).unwrap_or(true) {
         false => write!(f, "{v} {:?}", metadata.as_ref().unwrap()),
        true => write!(f, "{v}"),
    }
}

[paleolimbot]

I like the pretty print extension!

That's a very interesting point! Would using the existing LogicalType instead of Field resolve this problem? There we don't distinguish between, for example, Utf8 and Utf8View.

It's close! The LogicalType implements a few things (like Utf8/Utf8View as the same type) that if swapped into the Expr today might cause some disruption. None of the below should keep you from pursuing this more ambitious/eutopian future...I want that future but I also want PRs that can merge and getting anybody to care about/review this stuff is hard (thank you for caring and reviewing!).

In my head, the idea is that the logical type reference will provide access to these extensions.

Agreed! Inlining that reference is great. Probably the immediate challenge is ensuring the Column references are able to resolve their dyn type (maybe you've solved this already)

What is specified in the Expr/LogicalPlan enum vairants today is basically a Field whose name we ignore:

struct ExprType {
   // on merge mismatch is an error, casts + printing based on this
   storage_type: DataType, 
   // Alias merges metadata keys, Values error on mismatch, casts and printing ignore
   metadata: FieldMetadata, 
   // Mostly propagated everywhere but it's not quite clear how this should work with
   // a Cast or Alias
   nullability: bool
}

The SerializedTypeView garbage I was playing with was trying to formalize a few of the ambiguities with respect to how the information is propagated. Basically I would like to separate FieldMetadata from extension_name/metadata to make it harder to accidentally overwrite or drop the type information.

struct ExprType {
   // merge mismatch, cast, and printing always consider these three pieces of information
   // the 'safe' but slightly too strict way to do this is to always treat them with exact equality
   storage_type: DataType, 
   extension_name: String,
   extension_metadata: String,
   // Not sure whether merge or drop on mismatch is better but at least nothing
   // overwrite or drops the extension type
   metadata: FieldMetadata, 
   // Allows Alias and Cast to propagate the parent expression's nullability (which they
   // currently do)
   nullability: Option<bool>
}

...maybe swapping in Arc<dyn TypeExtension> for extension_name and extension_metadata in that example would work? Or maybe an ExprType is DataType + Arc<dyn TypeExtension> and ExprField adds the FieldMetadata and nullability? ScalarAndMetadata could maybe then store the ExprType as well? Or something better?