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!

[plotor]

Of course, is there anything I can do?

[universalmind303]

I think the starting point would be to come up with an architectural proposal for how we plan on supporting rust based extensions for sources, sinks, functions and io protocols. We could likely narrow the scope to just sources initially.

[universalmind303]

One idea I had -- now that we're going to be on arrow-rs instead of arrow2, we could try to bring in datafusion and use that as the ffi layer. They have a crate for defining ffi sources. If daft were to implement a datafusion-ffi compatible ffi api, then we'd get access to a lot of the datafusion ecosystem without necessarily needing to include datafusion as a dependency.

[universalmind303]

So I did some experimentation on what it would look like to have ffi based extensions. cc @plotor

For IO protocols,
I think this is easily solved via opendal. See #6177

---

For Functions, I'd like to propose a datafusion ffi based approach, Here's how I think it'd look from a library author experience.

Extension Author Experience (target state)

Everything the author writes is Rust + a pyproject.toml. No Python code required.

Cargo.toml

[lib]
crate-type = ["cdylib"]

[dependencies]
arrow-array = "57"
arrow-schema = "57"
datafusion-expr = "52"
datafusion-common = "52"
datafusion-ffi = "52"
pyo3 = { version = "0.23", features = ["extension-module"] }

Note: No dependency on any daft Rust crate. The only coupling is:

• datafusion-expr + datafusion-ffi — for the ScalarUDFImpl trait and FFI_ScalarUDF
• pyo3 — for PyCapsule creation and #[pymodule]
• daft Python package — runtime dependency for daft.register_udf()

lib.rs

use std::sync::Arc;
use datafusion_expr::*;
use datafusion_ffi::udf::FFI_ScalarUDF;
use pyo3::prelude::*;
use pyo3::types::PyCapsule;

// 1. Implement the function using DataFusion's trait
#[derive(Debug, PartialEq, Eq, Hash)]
struct CosineDistance { signature: Signature }

impl ScalarUDFImpl for CosineDistance {
    fn name(&self) -> &str { "cosine_distance" }
    fn return_type(&self, _: &[DataType]) -> Result<DataType> { Ok(DataType::Float64) }
    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
        // Pure arrow-rs computation — no Daft dependency
    }
    // ...
}

// 2. Register via PyCapsule on import
#[pymodule]
fn daft_distance(py: Python, _m: &Bound<PyModule>) -> PyResult<()> {
    let udf = Arc::new(ScalarUDF::from(CosineDistance::new()));
    let ffi_udf = FFI_ScalarUDF::from(udf);
    let capsule = PyCapsule::new(py, ffi_udf, Some(c"daft.udf"))?;

    let daft = py.import("daft")?;
    daft.call_method1("register_udf", (capsule,))?;
    Ok(())
}

pyproject.toml

[build-system]
requires = ["maturin>=1.0"]
build-backend = "maturin"

[project]
name = "daft-distance"
requires-python = ">=3.9"
dependencies = ["daft>=0.3"]  # needed for daft.register_udf

User experience

import daft_distance  # triggers #[pymodule] → PyCapsule → register — that's it

import daft
from daft import col, udf

df = daft.from_pydict({"a": [...], "b": [...]})
df = df.with_column("sim", 
    daft_distance.cosine_distance(
        col("a"), 
        col("b")
    )
)

---

For table sources, I think we could do something very similar to functions. Datafusion also provides a FFI_TableProvider. This is a bit more work as we'd have to translate daft exprs into datafusion expr.

Extension Author Experience

Nearly identical to scalar UDFs — implement a DataFusion trait, wrap in FFI, register via PyCapsule.

Example: generate_series(n) — returns a table with numbers 0..n

#[derive(Debug)]
struct GenerateSeries;

impl TableFunctionImpl for GenerateSeries {
    fn call(&self, args: &[Expr]) -> Result<Arc<dyn TableProvider>> {
        // Extract the count from the first argument (a literal u64)
        let n = match &args[0] {
            Expr::Literal(ScalarValue::UInt64(Some(n)), _) => *n as usize,
            _ => return exec_err!("generate_series expects a u64 literal"),
        };

        let schema = Arc::new(Schema::new(vec![
            Field::new("i", DataType::UInt64, false),
        ]));

        let array = UInt64Array::from_iter_values(0..n as u64);
        let batch = RecordBatch::try_new(
            Arc::clone(&schema),
            vec![Arc::new(array)],
        )?;

        Ok(Arc::new(MemTable::try_new(schema, vec![vec![batch]])?))
    }
}

// Register via PyCapsule on import
#[pymodule]
fn daft_generate(py: Python, _m: &Bound<PyModule>) -> PyResult<()> {
    let udtf = Arc::new(GenerateSeries) as Arc<dyn TableFunctionImpl>;
    let ffi_udtf = FFI_TableFunction::new(udtf, None, /* task_ctx */, None);
    let capsule = PyCapsule::new(py, ffi_udtf, Some(c"daft.udtf"))?;

    let daft = py.import("daft")?;
    daft.call_method1("register_udtf", ("generate_series", capsule))?;
    Ok(())
}

User experience

import daft_generate  # triggers registration

import daft
df = daft_generate.generate_series(n=1000)
df = df.where(col("i") % 2 == 0)  # even numbers only
df.show()
# ╭───────╮
# │ i     │
# │ ---   │
# │ UInt64│
# ╞═══════╡
# │ 0     │
# │ 2     │
# │ 4     │
# │ ...   │
# ╰───────╯

Sinks get pretty tricky as theres nothing in datafusion_ffi to handle these.

[caican00]

Currently, the implementation of AI Functions is rather complex. If the extension api can be supported, allowing other developers to customize AI Functions without modifying the source code of daft, it will bring about a significant improvement in efficiency.

Now for the same requirements, different users may need to repeatedly implement the infrerence UDF. If the extension api is supported, users can conveniently build a shareable AI Functions lib based on this.

[rchowell]

Currently, the implementation of AI Functions is rather complex. If the extension api can be supported, allowing other developers to customize AI Functions without modifying the source code of daft, it will bring about a significant improvement in efficiency.

The AI functions are designed to be extensible from python, I left more details at #5765 but agreed that this could be better documented to avoid confusion in the future, #5920