Tracking Issue: Extension Type Derives #11115

This is a tracking issue for letting extension types opt into storage behavior like equality and ordering in a concise manner, similar to Rust's #[derive(...)]. The claims live on ExtVTable and DType in vortex-array/src/dtype/.

• Subissue of Epic: Harden Type System #10730 and Epic: Extension Types #7683.
• Builds on Tracking Issue: Fix Extension Types #11114, which makes extension types opaque by default.

Motivation

Once extension types are opaque (#11114), every semantic (not structural, see the next section) operation on them is rejected unless the type opts in. But most types often want common behaviors from their storage. For example, timestamps should compare and sort like their i64 storage, and UUIDs should compare equality by their bytes. Without a concise opt-in, every extension type would need its own handwritten rule for every operation that relies on equality or ordering, which would likely just be boilerplate, pushing stuff down into the storage.

Today, the extension compare kernel compares storage for every extension type (which is wrong for JSON), while is_sorted and min_max only run on storage for the temporal types and Uuid, through hand-written registrations (#10566, #10616).

Design

Semantic vs structural operations

Generally speaking, there are 2 kinds of operations in Vortex. Structural operations move or select values without reading them, like slice, take, filter, and zip. Semantic operations read values, so their result depends on what the values mean, like =, <, min, hashing, and casts.

Structural operations are the only ones that can "always work" on extension types. We can always filter an extension array, but we can't run some function add_1 on it, because we don't know anything about the semantics of its values (even if the storage type is an integer).

So loosely speaking, we can call operations that depend on the semantics of the values "semantic" operations, and things that operate over the structure of arrays / columns "structural" operations. Generally speaking, we can push down structural operations but not semantic ones.

Note that there are a few functions, like replace_time_zone or the spatial functions, that implement the semantics of specific extension types themselves, so we are not concerned with those.

Extension "claims"

There are cases where we can forward semantic operations to the storage, but only if we know something about the relationship between the extension type and its storage type.

For most extension types, 2 values are equal if and only if their storage values are equal. This isn't true when a type considers 2 different storage values to be equal (think JSON, or integers modulo some number). Similarly, for many types the order of the values is exactly the order of their storage values (which implies equality). The datetime types have both properties, so we can push equality and ordering down into their integer storage.

So we want extension types to be able to "claim" these properties, similar to Rust's #[derive(...)]. I propose something like this (which can evolve in the future):

// In the future, this might become an open set instead of hardcoded.
#[non_exhaustive]
pub enum Claim { Eq, Ord }

pub trait ExtVTable {
    fn claims(dtype: &ExtDType<Self>) -> &'static [Claim] { &[] }
}

impl DType {
    pub fn holds(&self, claim: Claim) -> bool;
}

Note that a claim only covers values of 1 dtype, so Timestamp(ms) and Timestamp(ns) never compare even though both claim Ord. Claims are listed by hand rather than derived from which kernels exist, and ExtDType::try_new checks that the declared storage dtype actually holds them. Claims beyond Eq and Ord (for example, for text functions on Json) and claims defined outside of Vortex are future work.

Using claims

Each scalar function and aggregate also declares how it uses extension operands. #11114 adds this declaration without Storage, and this issue adds Storage:

pub enum ExtensionUse {
    None,
    Structural,
    Specific,
    Storage { claims: Vec<Option<Claim>>, result: ResultMapping },
}

A Storage function names the claim that each argument needs and how its result gets the extension dtype back. When the claims hold, Vortex rewrites the call to run on storage (ts_a < ts_b becomes ext_storage(ts_a) < ext_storage(ts_b)), so the existing kernels, statistics, and pruning apply for free. Otherwise, the call is rejected.

Storage carries the result mapping, so every function that runs on storage has to give one:

pub enum ResultMapping {
    Rewrap(ExtDTypeRef),
    Keep,
    Fields(Vec<ResultMapping>),
    Custom(Arc<dyn CustomMapping>),
}

• Rewrap wraps the storage result in the given extension dtype. For example, min needs to answer Rewrap with its input dtype, since the minimum of timestamps is also a timestamp. The target is a dtype and not an argument index, so a result that is part of an input works too: a hypothetical list_max(List<Timestamp>) answers Rewrap(Timestamp). Note that I think there is some work to be done with nested paths here, but I think this idea still works.
• Keep leaves the result as it is. Comparisons and is_sorted answer Keep, since they return Bool no matter what.
• Fields gives one mapping for each field of a struct result. min_max answers Fields with 2 Rewraps.
• Custom lets the function map the result itself, for results that the other 3 can't describe. For example, min over List<Timestamp> returns a List<Timestamp>, which isn't an extension dtype.

Steps

• Add Claim, DType::holds, and ExtVTable::claims, checked when a dtype is built (needs the storage declaration from Tracking Issue: Fix Extension Types #11114)
• Give every extension type in the repository its claims
• Conformance checks and reference test types for claims
• Run a function on storage when its claims hold, with ExtensionUse::Storage, ResultMapping, and the generic storage rewrite (needs the typing check from Tracking Issue: Fix Extension Types #11114)
• Replace the hardcoded opt-ins of Tracking Issue: Fix Extension Types #11114 with claims for compare, the aggregates, list_contains, the row encoding, statistics, pruning, and engine pushdown
• Document claims for extension type authors

Until then, #11114 hardcodes the opt-ins for the extension types in the repository. Its lists match the claims of these types, so each switch to claims allows and rejects the same calls.

Unresolved questions

Implementation history