datafusion-ffi rebuilds the session and task context, with every function, on each scan and execute (~1.5 ms per scan)

This is a draft intended for review and re-filing in apache/datafusion, so the relevant entry points are in datafusion-ffi 55.2.0: table_provider.rs (scan_fn_wrapper), session/mod.rs, and execution/task_ctx.rs (execute_fn_wrapper). Start by reproducing the timings and measuring session and task context conversion separately. Done means an upstream change reduces repeated conversion overhead while preserving correct session and task context behavior; the issue lists lazy registries and caching as options, not settled designs.

🌀 Claude (xarray-sql)

A draft to review, then re-file upstream in apache/datafusion (datafusion-ffi). Found while profiling #272.

Summary

Scanning a table provider through datafusion-ffi costs about 1.45 ms more than scanning an in-memory table of the same rows, almost all of it on the FFI boundary, not in the provider. On every scan, datafusion-ffi rebuilds the caller's session on the provider's side (ForeignSession::try_from(&FFI_SessionRef) in scan_fn_wrapper). It rebuilds the task context again on every partition's execute (From<FFI_TaskContext> for Arc<TaskContext> in execute_fn_wrapper). Each rebuild converts every registered scalar, aggregate and window function across the boundary, and the session config.

For one large scan that's noise. But a plan that scans small tables many times pays it each time. DataFusion inlines a CTE or a view at every reference, so the GPT in xarray-sql's ML demos scans its eleven small tables 2,989 times in one query, and 6.7 s of its 9.2 s forward pass is this overhead (#272).

Measurements

datafusion-python 54.0.0, which has 302 functions registered by default, with xarray-sql's FFI table provider, on an Intel i7-8700. SELECT * FROM t over 1,040 rows:

Physical planning (the provider's scan, through scan_fn_wrapper) adds 0.72 ms, and execution (execute_fn_wrapper) adds about 0.75 ms. The provider's own work is small: its Python side, which builds the record batches, takes 0.28 ms per scan, and its Rust scan() only clones partition handles.

The cost grows with the number of registered functions, though only on the FFI scan:

That's about 1.5 µs per function per scan, or about 0.45 ms of the overhead at the default 302 functions. I haven't broken down the remaining ~1 ms. SessionConfig::try_from on the FFI config, which converts every option, and the catalog list are candidates.

Reproduction

With xarray-sql (any FFI TableProvider would do):

import time
import numpy as np
import pyarrow as pa
import xarray as xr
from datafusion import udf
import xarray_sql as xql

ds = xr.Dataset({"w": (("tok", "e"), np.random.rand(65, 16))})
ctx = xql.Context()
ctx.from_dataset("ffi", ds)
ctx.register_record_batches("mem", [ctx.sql("SELECT * FROM ffi").to_arrow_table().to_batches()])

def per_scan(table, n=300):
    q = f"SELECT * FROM {table}"
    ctx.sql(q).to_arrow_table()
    start = time.perf_counter()
    for _ in range(n):
        ctx.sql(q).to_arrow_table()
    return (time.perf_counter() - start) / n * 1000

print(per_scan("ffi"), per_scan("mem"))   # ~1.7 ms, ~0.25 ms
for i in range(1000):
    ctx.register_udf(udf(lambda x: x, [pa.float64()], pa.float64(), "immutable", name=f"f{i}"))
print(per_scan("ffi"), per_scan("mem"))   # ~3.2 ms, ~0.37 ms

Where, in datafusion-ffi 55.2.0

• table_provider.rs, scan_fn_wrapper: ForeignSession::try_from(&session), unless session.as_local() applies, which it does only when both sides are the same library.
• session/mod.rs, impl TryFrom<&FFI_SessionRef> for ForeignSession: collects scalar_functions, aggregate_functions and window_functions into new maps, converts each function's FFI wrapper, and converts the config and table options.
• execution/task_ctx.rs, impl From<FFI_TaskContext> for Arc<TaskContext>: the same for every execute, once per partition.

Possible fixes upstream

• Lazy registries: a ForeignSession or foreign TaskContext could look functions up across the boundary on demand (udf(name)), instead of copying all of them up front. Most providers never call a function during scan or execute.
• Cache per session: keep the converted session or task context keyed by the session id, so it's rebuilt only when the session changes.
• Measure the config conversion separately; if it is the other ~1 ms, it could be cached the same way.

Workaround in xarray-sql

Register small in-memory Datasets as native in-memory tables, which bypass FFI (#272). Lazy or large Datasets still need the FFI provider, so they keep paying this per scan.