Extension API for third-party DataFusion integrations (UDFs, optimizer rules, planner extensions)

[james-willis]

Summary

There is no public API for plugging a third-party extension into Sail. Today, integrating something like Apache SedonaDB — which contributes scalar UDFs, aggregate UDAFs, a session-config extension, logical optimizer rules, and a physical planner extension — requires forking Sail and editing internals across at least four crates.

This issue proposes one composable extension API that covers every dimension a real-world extension contributes (SailExtension trait), plus a Python entry-points discovery layer so end users on pip install pysail can consume extensions without rebuilding pysail. One small prerequisite fix to ExtensionPhysicalPlanner is folded in.

Motivating example

Apache SedonaDB is a representative case because it exercises most of the extension dimensions at once. To make ST_Intersects joins plan as SpatialJoinExec instead of NestedLoopJoinExec, an embedder needs Sail to:

1. Resolve ST_* scalar UDFs at plan time.
2. Carry SedonaOptions on the SessionConfig so Sedona's optimizer rules can read their config.
3. Run Sedona's logical optimizer rules so spatial filters on cross-joins become SpatialJoinPlanNode extension nodes.
4. Plan those extension nodes via Sedona's physical planner extension so they become SpatialJoinExec.
5. Re-resolve those UDFs by name on remote workers when codec-deserializing physical plans for distributed execution.

Items 1–4 are all required to plan a single SQL query. Item 5 is required the moment that query runs in a multi-node cluster. There is no way to contribute any of them from outside the sail-* workspace. (Sedona also ships aggregate UDAFs, window UDFs, and table functions — those are separate dimensions covered by the inventory below, but they aren't on the critical path for the spatial-join example.)

Where each dimension is hardcoded today

This is the inventory an embedder discovers after attempting the integration. All of these are currently sail-*-internal touch points that any third-party extension would need to mutate.

Dimension	Hardcoded location
Scalar UDFs (plan time)	crates/sail-plan/src/function/mod.rs — BUILT_IN_SCALAR_FUNCTIONS lazy_static HashMap. Sedona functions are merged in via a direct call to sail_sedona::sedona_scalar_udfs().
Generator / table functions (plan time)	Same file — BUILT_IN_GENERATOR_FUNCTIONS, BUILT_IN_TABLE_FUNCTIONS.
Aggregate UDAFs (plan time)	crates/sail-plan/src/function/aggregate.rs:541 — direct call to sail_sedona::sedona_aggregate_udfs() populates BUILT_IN_AGGREGATE_FUNCTIONS.
Window UDFs (plan time)	crates/sail-plan/src/function/window.rs — BUILT_IN_WINDOW_FUNCTIONS.
UDF / UDAF re-resolution (execution time)	crates/sail-execution/src/codec.rs:1539 and :1741 — hardcoded fallback lookups against sail_sedona::get_sedona_scalar_udf / get_sedona_aggregate_udf so workers can rebuild UDFs from a serialized plan.
Session-config extensions	crates/sail-session/src/session_factory/server.rs::create_session_config — Sedona requires add_sedona_option_extension(config) here.
Logical optimizer rules	crates/sail-session/src/session_factory/server.rs::create_session_state — Sedona requires register_spatial_join_logical_optimizer(builder) here.
Physical planner extensions	crates/sail-session/src/planner.rs::ExtensionQueryPlanner::create_physical_plan — Sedona's SpatialJoinExtensionPlanner must be pushed onto the extension_planners chain (see prerequisite below for ordering constraints).

The ServerSessionMutator trait exists and is the closest existing extension hook, but it only exposes mutation of SessionConfig, SessionStateBuilder, and RuntimeEnvBuilder — it doesn't help with the plan-time function registries, the codec fallbacks, or the ExtensionQueryPlanner chain. And it requires the embedder to use Sail as a library, not as the sail CLI binary or pysail.

Proposal

Prerequisite: stop short-circuiting on unknown extension nodes

crates/sail-session/src/planner.rs ends its node match with:

} else {
    return internal_err!("unsupported logical extension node: {:?}", node);
};

DataFusion's DefaultPhysicalPlanner walks each registered ExtensionPlanner and stops at the first one that returns Ok(Some(plan)). The convention is that planners which don't recognize a node return Ok(None) and the chain continues. Sail's ExtensionPhysicalPlanner instead returns Err(...), so any planner that needs to handle a node Sail doesn't know about must be listed before ExtensionPhysicalPlanner — otherwise the chain short-circuits with a confusing internal error.

This took non-trivial debugging to diagnose for the Sedona case: the logical rule fired correctly, the new node appeared in the plan, but the physical build failed with unsupported logical extension node: SpatialJoin.

Five-line change, no public API impact, but mandatory for the rest of the proposal to compose cleanly.

Part 1: a unified SailExtension trait

Introduce a single trait that an embedder implements once, replacing the need for any per-extension bridge crate inside sail-*. The trait covers every dimension listed above:

pub trait SailExtension: Send + Sync {
    fn name(&self) -> &str;

    // Function registries — merged into both the plan-time HashMaps in
    // sail-plan and the execution-time codec fallbacks in sail-execution.
    fn scalar_udfs(&self)     -> HashMap<String, Arc<ScalarUDF>>     { Default::default() }
    fn aggregate_udfs(&self)  -> HashMap<String, Arc<AggregateUDF>>  { Default::default() }
    fn window_udfs(&self)     -> HashMap<String, Arc<WindowUDF>>     { Default::default() }
    fn generators(&self)      -> HashMap<String, Arc<ScalarUDF>>     { Default::default() }
    fn table_functions(&self) -> HashMap<String, Arc<TableFunction>> { Default::default() }

    // Session-state contributions — applied during create_session_state.
    fn configure_session(&self, config: SessionConfig) -> Result<SessionConfig> { Ok(config) }
    fn logical_optimizer_rules(&self)  -> Vec<Arc<dyn OptimizerRule + Send + Sync>>         { vec![] }
    fn physical_optimizer_rules(&self) -> Vec<Arc<dyn PhysicalOptimizerRule + Send + Sync>> { vec![] }
    fn extension_planners(&self)       -> Vec<Arc<dyn ExtensionPlanner + Send + Sync>>      { vec![] }
}

Sail merges these contributions in well-defined order:

• Function registries: extension entries override built-ins (or vice-versa — pick one and document; overriding feels right since the extension was loaded explicitly).
• Logical / physical optimizer rules: appended to Sail's defaults.
• Extension planners: prepended to Sail's chain.

Registration via the existing ServerSessionMutator boundary or a new with_extension(Arc<dyn SailExtension>) builder method on ServerSessionFactory. The embedder's main becomes:

ServerSessionFactory::new(...)
    .with_extension(Arc::new(SedonaExtension::new()))

sail-sedona becomes a regular external crate that implements SailExtension. No edits to any sail-* crate.

Part 2: Python entry-points discovery for pysail

The trait alone unblocks anyone who can build a binary, but it does nothing for pip install pysail users. To close that gap, pysail discovers extensions at startup via Python entry points:

# pyproject.toml of a hypothetical pysail-sedona wheel
[project.entry-points."pysail.extensions"]
sedona = "pysail_sedona:register"

At session-create time, pysail iterates importlib.metadata.entry_points(group="pysail.extensions") and calls each register() function, which hands back an Arc<dyn SailExtension> (via a thin pyo3 bridge) that gets passed to with_extension.

End-user UX:

pip install pysail pysail-sedona
sail spark server --port 50051     # Sedona is just there

No rebuild of pysail. This is the closest practical analog to Spark's "drop a jar in $SPARK_HOME/jars, call the registrator" story.

Coverage

What each part of the proposal enables, by use case:

	Custom Rust binary	Custom pysail wheel	pip install plugin
Prerequisite alone	partial — no UDFs	partial — no UDFs	no
+ SailExtension trait	full	full (one-line change)	no
+ entry-points discovery	full	full	yes

All three parts ship together; the entry-points layer is what makes the end-user UX comparable to Spark.

Outstanding decisions

The proposal supports loading multiple extensions composably (multiple with_extension calls, or multiple plugin wheels installed alongside each other). What it doesn't yet specify is how Sail should resolve collisions between extensions. These need to be decided before implementation:

• Function name collisions. Two extensions register a scalar UDF with the same name (e.g. Sedona's ST_Distance and a custom geo library's ST_Distance). Should this error at registration, error at function-resolution, last-registered wins, first-registered wins, or require namespacing (sedona.ST_Distance)? Same question for UDAFs, window UDFs, generators, and table functions.
• SessionConfig extension collisions. DataFusion's add_extension is keyed by Rust TypeId, so two extensions contributing different config types compose freely; two contributing the same type silently replace. Is silent replace acceptable, or should Sail detect and surface this?
• Optimizer rule ordering. Logical and physical optimizer rules from multiple extensions are appended in extension-registration order. Is registration order sufficient, or do extensions need a way to declare ordering constraints relative to each other (e.g. "must run before PushDownFilter", "must run after Sedona's MergeSpatialFilterIntoJoin")?
• Extension planner ordering. Same question for the extension_planners chain. Two extensions could both claim to handle the same logical extension node; today the first one in the chain wins by virtue of returning Ok(Some(...)) first.
• Entry-point discovery conflicts. If two installed wheels both register a pysail.extensions entry point with the same name, what should pysail do at startup? Error, warn, or load both (and then face the function-collision question above)?
• Per-session enable/disable. Are extensions always-on once registered, or should there be a way to opt in/out per session (e.g. SET sail.extensions = 'sedona,comet')?

Confidence and prior art

Roughly 60–70% confident the entry-points layer works end-to-end without unpleasant surprises. The closest demonstrated case is Polars plugins, which use the same pattern (compile-time-coupled native extension wheels discovered at runtime). It works, but with strict version coupling: every plugin must be rebuilt on every Polars upgrade because Rust has no stable ABI for trait objects.

For Sail, "version coupling" means each plugin wheel must be built against matching versions of pysail, datafusion, arrow, and pyo3 — the types crossing the pyo3 boundary (Arc<dyn SailExtension>, Arc<dyn ScalarUDF>, Arc<dyn OptimizerRule>, …) all come from those crates. The plugin's own internal dependencies (e.g. which sedona-* versions a pysail-sedona wheel pulls in) are private to the plugin; pysail never sees those types. Worth prototyping the Arc<dyn SailExtension> handoff across two independently-built pyo3 cdylibs before fully committing to this layer — that's the part with no direct precedent I'm aware of.

The trait itself (Part 1) has no such risk and can ship independently if Part 2 turns out to need more design work.

Reference: a working but forked Sedona integration

A complete Sedona integration that would have been roughly:

ServerSessionFactory::new(...).with_extension(Arc::new(SedonaExtension::new()))

under this proposal currently exists as a fork at james-willis/sail branch sedona-integration. The diff against sail-session, sail-plan, sail-execution, and the workspace Cargo.toml is the inventory in the table above made concrete.

[linhr]

Nice write-up! It would likely take some time to finalize the design but it's great to see the concrete proposal here. Let me share my thoughts.

Indeed, there are many extension points we need to have (functions, logical/physical plan nodes, and optimizers) to support the full query execution logic. (I also noticed it here in a previous discussion.) Most of these interfaces are from DataFusion, so they are relatively stable fortunately.

The session mutator is something not very stable and I'd consider it a deep implementation detail, so we need to think more around the best way to alter session config.

Using Python package for distribution and the discovery mechanism make sense to me. However, since DataFusion has an FFI, we won't use Rust trait as the API, but we'll have an extension FFI built on top of DataFusion so that we won't need to recompile the extension on every Sail version or Rust version change. This allows Sail and the extension to release under different schedule. See my comment here.

I would love to see end-to-end proof of concept around FFI.

Open questions:

1. How do we handle breaking changes to the extension API?
2. How does this fit into Spark Connect's existing plugin mechanism?
3. How to specify optimizer and planner orders? (This is mentioned in the proposal already.)

[paleolimbot]

However, since DataFusion has an FFI, we won't use Rust trait as the API, but we'll have an extension FFI built on top of DataFusion so that we won't need to recompile the extension on every Sail version or Rust version change.

This is the main trick. If I've kept up with the DataFusion FFI work, the main thing missing is that the logical plan is not part of the extension FFI (so our current spatial join optimization is tricky to generically plug in to another DataFusion based engine that doesn't already know about spatial joins and predicates). We haven't tried very hard to plug in SedonaDB to sedona-python but it would be roughly the same work.

FWIW my main complaint with the DataFusion FFI is that the major version is not part of the FFI, so you get a crash if you try to mix datafusion-python versions with whatever your extension was compiled against but you could work around that here.

Also FWIW, DuckDB handles this by providing the CI job that aligns everything so that no FFI is needed. This is super powerful (allows access to literally all of DuckDB's internal APIs) but it means extension authors have to keep up (I have three DuckDB extensions that are all out of date as a result of this 🙂 ). It took them quite a long time to get this working.

[shehabgamin]

cc @alexy

[linhr]

These are great discussions! I'm going to move this to a GitHub discussion thread and we can continue there.

[linhr]

Ok done.

[alexy]

Friends -- great to join this effort, leading the community at LakeSail. Saw the Sedona presentation at the DataFusion meetup recently. Sedona is an excellent project and a great extension to think through and add, as it exercises all of the components of Sail. I'm fairly new to Sail and Rust, but with an extensive background in similar languages, functional programming, and distributed systems (and a Spark pioneer). So I'll structure it as a learning process for those who want to understand how Sail works, and we can confirm the right touchpoints for extensions.

We wrote a book, with Codex, reviewing all parts of Sail with an eye on their participation in the extensions:

https://github.com/alexy/sail-rust-book

We also have a community Slack, and I've just started a #sail-extensions channel there. Please join and chime in, so we have an informal way to bounce off ideas before they become a part of the record here.

Since @james-willis has a full prototype, one question is: what prevents us from using it as is? I believe we need to answer @linhr's questions, correct?

[james-willis]

@alexy my prototype is a hard fork of sail that takes a sedona-db dependency. I side stepped all of the interesting questions about plugins this way.

In an ideal world I would be able to build a plugin that augments sail with sedona functionality.

It would look something like:

1. I build some sedona-sail repo/wheel file that implements a sail plugin from the rust internals of sedona-db
2. Users pip install sail and sedona sail . They'll do something like:

# This is all pseudo code - I did not refer to how you actually init a sail session
sail = SailContextGetter.get()
sedonaSail.registerSedona(sail)

sail.sql("SELECT ST_POINT(1.0, 2.0)").show()

3. Users can update their versions of sail and sedona-sail (somewhat) independently to get new function form each as they need it.

[SemyonSinchenko]

@james-willis what do you think about SparkConnect protocol? Since Sail is already implementing the protocol and google.protobuf.Any extension = 998; is already a part of sail, it looks an obvious choice for plugins system. And in this case users can switch back-end (between Spark+Sedona and Sail+SedonaDB) without changing the code. In Spark there are four types of "plugins": Relation, Expression, Command and MLBackend and for me it looks like they are covering all the possible cases.

Using spark's way to plugins a) fits natively to the SparkConnect protocol (that is adopted more and more -- Sail, Snowflake, etc.) as well in the clients code no changes are required (exactly as we did in GraphFrames: runtime dispatch under the hood and users do not need to care what kind of session/backend is it).

[james-willis]

As you know im not the spark connect expert. But my inclination is. that spark connect protocol is more for extensions that operate on top of pyspark connect interfaces. This topic is more about bringing in extensions that work on datafusion interfaces.

[james-willis]

@paleolimbot and I (mostly Dewey) are working through this plugin feature now in sedonaDB (ie plugging into sedonadb from some other datafusion based library).

Give us a few weeks and we might have some more battle-tested knowledge about this work.

[paleolimbot]

A few updates on our end...I've been adding these as we need them to avoid too much hypothetical design. I'm vaguely planning to propose them to DataFusion as well since these have come up in the context of Comet needing a concrete C API for JNI integration. DataFusion also has its own FFI work that is evolving and would also work for some parts of this as well. I like knowing the details of what is passed across the FFI boundary (and version-independence from DataFusion) but I'm not offended if one of their APIs is sufficient for this.

Some things we've done:

• Getting scalar functions across a boundary: feat(c/sedona-extension): Move FFI scalar kernel definitions to c/sedona-extension and write them in C apache/sedona-db#407 . We needed this for implementing functions in C++ for a library that only exists in C++.
• Getting tables and execution plans across an FFI boundary: feat(c/sedona-extension): Add FFI table provider, exec plan, and synchronous stream exchange apache/sedona-db#1004 . We're using this to separate some of the more specialized/obscure file formats (e.g., PointCloud, Zarr, proprietary catalogs).
• Getting expressions across a boundary (specifically for pruning): feat(c/sedona-extension): Push Expr filters into the TableProvider across the FFI boundary  apache/sedona-db#1094 . This relies heavily on DataFusion protobuf which has mostly been backward compatible for some time (particularly for logical expressions), but could use something else. We need this to implement pruning in our data sources.

The join is probably the most unique extension point that would need to be supported...either Sail would have to hard-code some of the logical planning that identifies a logical and/or KNN join and have a specific extension point for "spatial join" (which could resolve the FFI version of the executionplan), or implement some general "join extension" (where it passes the join condition to an extension which can optionally return an FFI execution plan if it applies). This would be ambitious (there is a lot of functionality loaded on to the execution plan that makes the interaction complicated).