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

Aberta
  • Algo errado?
#273 0 comentários 0 reações 0 responsáveis Ver no GitHub

Mantenedores costumam responder em até 3 dias

Ninguém assumiu esta issue ainda.

Avaliação

Dificuldade
5/5
Tempo estimado
Mais de uma semana
Facilidade para iniciantes
35/100
Tipo de issue
Bug
Clareza
Razoavelmente clara
Status de atividade
Ativa
Stack de tecnologia
python, rust
Domínio
backend, performance

Direção de pesquisa

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.

Escrita pelo modelo de indexação a partir do texto da issue.

Descrição

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.

Linguagem predominante
Python
Estrelas
142
Forks
19
Merge médio
3d 47min
PRs com merge (30d)
8

Preparar o ambiente

Primeiros passos

  1. Leia a issue inteira e depois o guia de contribuição do projeto.
  2. Comente na issue dizendo que vai assumir — evita que duas pessoas façam o mesmo trabalho.
  3. Faça um fork do repositório e trabalhe em uma branch.
  4. Abra um pull request que referencie o número da issue.

Mais de xqlsystems/xarray-sql

Todas as issues de xqlsystems/xarray-sql

Issues semelhantes

Mais issues de Python

Receba novas issues na sua caixa de entrada

Um resumo curto de issues do GitHub para quem está começando.