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

Ouverte
  • Quelque chose ne va pas ?
#273 0 commentaires 0 réactions 0 personnes assignées Voir sur GitHub

Les mainteneurs répondent en général sous 3 jours

Personne n'a encore pris cette issue.

Évaluation

Difficulté
5/5
Temps estimé
Plus d'une semaine
Accessibilité débutants
35/100
Type d'issue
Bug
Clarté
Plutôt claire
Activité
Active
Stack technique
python, rust
Domaine
backend, performance

Piste de recherche

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.

Rédigé par le modèle d'indexation à partir du texte de l'issue.

Description

ml

🌀 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:

logical plan + optimize + physical plan collect
FFI table provider 0.085 ms 0.165 ms 0.918 ms 1.744 ms
in-memory table (register_record_batches) 0.079 ms 0.154 ms 0.201 ms 0.251 ms

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:

extra UDFs registered FFI scan in-memory scan
0 1.69 ms 0.25 ms
250 2.06 ms 0.27 ms
500 2.46 ms 0.30 ms
1,000 3.19 ms 0.37 ms

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.

Langage dominant
Python
Étoiles
142
Forks
19
Merge moyen
3 j 47 min
PR mergées (30 j)
8

Préparer son environnement

Par où commencer

  1. Lisez l'issue en entier, puis le guide de contribution du projet.
  2. Signalez en commentaire que vous la prenez — cela évite que deux personnes fassent le même travail.
  3. Forkez le dépôt et travaillez sur une branche.
  4. Ouvrez une pull request qui référence le numéro de l'issue.

Autres issues de xqlsystems/xarray-sql

Toutes les issues de xqlsystems/xarray-sql

Issues similaires

Plus d'issues Python

Recevez les nouvelles issues par e-mail

Un résumé court des issues GitHub adaptées aux débutants.