RFC: Extension API for Daft

[universalmind303]

Extension API for Daft

We're exploring an extension API to allow third-party libraries to build on top of Daft. Draft PR: #5751

Motivation

External contributors want to build domain-specific functionality without modifying Daft core. An extension API allows the ecosystem to grow independently. Extensions can iterate at their own pace, maintain separate release cycles, and avoid coupling to Daft's core development timeline.

Design Approaches

Approach 1: Dynamic Registration (Polars/Pandas model)

Lightweight Python-level extensions via decorators:

# DataFrame namespace extensions
@register_dataframe_extension("geo")
class GeoExtension:
    def __init__(self, df: daft.DataFrame):
        self._df = df

    def to_wkt(self) -> daft.DataFrame:
        return self._df.with_column("wkt", ...)

df.geo.to_wkt().collect()

Pros: Simple to implement, flexible, pure Python, easy for extension authors
Cons: No type hints, limited discoverability, doesn't integrate with query planner

Approach 2: Global Registry (Spark/DataFusion model)

# Registering custom data sources
daft.register_datasource("clickhouse", ClickHouseSource)
df = daft.read("clickhouse").options(**clickhouse_options).load()

# Registering custom sinks
daft.register_sink("clickhouse", ClickHouseSink)
df.write("clickhouse").options(**clickhouse_write_options).load()

# register custom io sources
daft.register_io_source("s3a", MyS3Handler)
df = daft.read_parquet("s3a://bucket/path")

Pros: Tighter integration, can participate in optimization passes, first-class citizens in query planning
Cons: Much more complex implementation, needs careful API boundaries, harder to extend

We do have a lot of this in place for sources and sinks already, but don't provide any registry mechanism

Key Extension Points

Potential registration targets:

• Helper/Utility functions: Distributed operations (like merge_columns, compaction utilities, etc)
• Io Protocols: hdfs://, gvfs://, etc
• DataSources: Custom read_*() implementations
• Sinks: Custom write_*() implementations
• Functions/Expressions: Custom operations that participate in logical plan. (can do this via udfs for the most part)
• Catalogs: External metadata systems
• Namespace extensions: DataFrame method namespaces (draft PR approach)

Questions

1. Which extension points are most valuable to potential users?
2. Should we support both approaches - registry for core integrations, decorators for convenience methods?
3. How do registered extensions interact with optimizer?
4. For datasources/sinks - what's the minimal interface extensions must implement?
5. Does the registry model work with Daft's Rust core, or Python-only?

The draft PR implements the simpler decorator approach. Looking for feedback on whether we need the heavier registry system, and which extension points matter most.

[plotor]

I have recently been refactoring the read_lance interface based on the Lance Rust SDK. I think this kind of abstraction should not only consider the Connector implementation in the Python language but also support the Rust language.

[universalmind303]

So i agree that rust based extensions would be ideal and I would prefer doing this as well. By using FFI as the extension layer, it'd be rather easy to add a python wrapper for it.

However introducing FFI based extensions is quite an undertaking. If this is something you or other would be willing to contribute, I'd love to move forward with a rust FFI implementation!