feat: Python UDFs: per-session inlining toggle and strict refusal setting (#1546)
* feat: per-session Python UDF inlining toggle + sender ctx + strict refusal
Adds a per-session toggle that turns inline Python UDF encoding on or
off, plus the supporting plumbing to make it usable through
pickle.dumps.
Codec layer:
* PythonLogicalCodec / PythonPhysicalCodec gain a python_udf_inlining
bool (default true) and a with_python_udf_inlining(enabled) builder.
Each try_encode_udf{,af,wf} short-circuits to inner when the toggle
is off; each try_decode_udf{,af,wf} that recognizes a DFPY* magic
on a strict codec returns a clean Execution error instead of
invoking cloudpickle.loads. The refusal message names the UDF and
the wire family so an operator can see at a glance whether to
re-encode the bytes or register the UDF on the receiver.
Session layer:
* PySessionContext::with_python_udf_inlining(enabled) returns a new
session whose stacked logical + physical codecs both carry the
toggle. The Arc<SessionState> is cloned (cheap), only the codec
pair is rebuilt, so registrations and config stay attached.
* SessionContext.with_python_udf_inlining(*, enabled) is the Python
wrapper. enabled is keyword-only because positional booleans at
the call site read as opaque.
Sender-side context:
* datafusion.ipc gains set_sender_ctx / get_sender_ctx /
clear_sender_ctx thread-locals. Expr.__reduce__ now consults
get_sender_ctx() to pick the codec for outbound pickles, which is
the only path through which a strict session affects pickle.dumps
(the protocol calls __reduce__ with no arguments). Without a
sender context the default codec is used.
Tests:
* test_pickle_expr.py picks up TestPythonUdfInliningToggle (covers
both directions of the toggle plus the explicit-ctx fast path),
TestWorkerCtxLifecycle (set/clear/threading), and
TestSenderCtxLifecycle.
* New test_pickle_multiprocessing.py + helpers exercise the full
driver -> worker round-trip on a multiprocessing.Pool with set_*_ctx
installed in the worker initializer.
* CI workflow gets a 30-minute timeout-minutes backstop so a hung
pickle worker can't block the matrix indefinitely.
User-guide docs and the runnable examples land in PR4 of this series.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* update uv lock
* docs: clarify Python UDF inlining docstring; drop unresolved :doc: refs
Rewrite with_python_udf_inlining docstring for readability and remove
references to /user-guide/io/distributing_work, which does not exist
yet. Keep security warning inline as a .. warning:: Security block,
matching the existing pattern in Expr.to_bytes / from_bytes /
__reduce__. The central doc will land in a follow-on PR.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* docs: add doctest examples for sender ctx + UDF inlining toggle
Per CLAUDE.md, every Python function needs a docstring example.
Adds examples to with_python_udf_inlining, set_sender_ctx,
clear_sender_ctx, and get_sender_ctx. Also clarifies that
with_python_udf_inlining returns a new SessionContext and leaves
the original unchanged, matching the with_logical_extension_codec
pattern.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* refactor: address review nits for UDF inlining toggle + sender ctx
* codec: strict refusal routes through `read_framed_payload` so
malformed inline bytes surface their own diagnostic; the
"inlining is disabled" message now fires only when the payload
would have decoded.
* codec: add summary line above `PythonPhysicalCodec::with_python_udf_inlining`
cross-link for rustdoc rendering.
* expr: hoist `get_sender_ctx` import to module top; note that
`__reduce__` also drives `copy.copy` / `copy.deepcopy`.
* context: accept `with_python_udf_inlining` positionally or as
kwarg (drop `*,`).
* tests: replace size-ratio heuristic with semantic check for the
`DFPYUDF` family prefix; switch single-batch closure test to
`pool.apply`.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* refactor: keyword-only inlining flag, skip GIL on prefix mismatch
- `SessionContext.with_python_udf_inlining` now keyword-only (`*, enabled`)
to match the documented call style and the existing doctests/tests.
- `refuse_if_inline` and the three `try_decode_python_*` decoders short-
circuit on a `starts_with(family)` check before `Python::attach`, so
plans whose UDFs are not Python-defined no longer pay a GIL acquisition
per decode call. Semantics preserved: `strip_wire_header` already
returns `Ok(None)` when the prefix does not match.
- `datafusion.ipc` module docstring wraps the `set_sender_ctx` example in
`try`/`finally` and notes that the thread-local holds a strong
reference to the installed `SessionContext` until cleared.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
* Add dev dependency
* Add testing for CI failure
* Additional debugging for mp tests in CI
* Set path for workers
* more path updates for unit tests
* test(pickle): remove multiprocessing CI debug instrumentation
Multiprocessing forkserver/spawn hang was diagnosed and fixed: workers
could not import `tests._pickle_multiprocessing_helpers` because
`pytest --import-mode=importlib` does not add the test parent dir to
`sys.path`. The fix (appending the parent dir to `sys.path` so it is
inherited by mp workers without shadowing the installed `datafusion`
wheel) is retained. This commit drops the diagnostic scaffolding that
was added to identify the hang point:
- `_diag` + per-import / per-task log writes to /tmp
- `snapshot_processes` and the `threading.Timer` that captured worker
state mid-hang
- `diag_init` Pool initializer
- "Dump multiprocessing diagnostic log" CI step
Pre-existing infrastructure is kept: per-test `@pytest.mark.timeout(120)`
(backed by `pytest-timeout` dev dep) and the job-level
`timeout-minutes: 30` backstop on the test matrix.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
* Shorten rust side docstring since it's duplicative of the exposed python docstring
* docs: clarify strict-mode refusal message and to_bytes inlining docs
Address PR review feedback:
- codec.rs: rewrite strict-refusal error to present the two real
remediations (sender re-encode by-name + receiver register; or
receiver enables inlining, accepting cloudpickle risk) instead of
bundling registration with both-side inlining.
- expr.py: qualify to_bytes docstring so Python UDF self-contained
behavior is conditional on with_python_udf_inlining being enabled.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* docs: clarify with_python_udf_inlining enabled arg is required
Reword docstring to drop misleading "(the default)" claim. The
`enabled` parameter is keyword-only and required — there is no
argument default. Note instead that fresh sessions inline UDFs
until the toggle overrides them (a session-level default, not an
argument default).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* docs: demonstrate strict-mode refusal in with_python_udf_inlining docstring
Replace placeholder isinstance check with a doctest that registers
a Python UDF, encodes an expression on the default session, then
shows the strict session refusing to decode the inline payload.
Exercises the actual behavior the toggle controls.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* docs: convert sender-ctx example to executable doctest
Replace the code-block in the ipc module docstring that demonstrated
set_sender_ctx with a doctest that actually runs. Worker-init example
remains a code-block since it documents a Pool-initializer pattern
that does not fit naturally into a doctest.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* docs: use 'thread-local sender context' as adjectival phrase
Bare 'thread-local' as a noun reads ambiguously next to the
_local.ctx attribute name. Hyphenate as adjective with explicit
'sender context' noun so the referent is unambiguous.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* docs: drop trailing clear_sender_ctx from set_sender_ctx example
The trailing cleanup call was test hygiene, not API teaching, and
risked implying callers must always pair set with clear. Adjacent
clear_sender_ctx and get_sender_ctx doctests are self-contained
(they explicitly set or clear before asserting), so removing the
cleanup line does not affect doctest outcomes.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
---------
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This is a Python library that binds to Apache Arrow in-memory query engine DataFusion.
DataFusion's Python bindings can be used as a foundation for building new data systems in Python. Here are some examples:
For tips on tuning parallelism, see Maximizing CPU Usage in the configuration guide.
The following example demonstrates running a SQL query against a Parquet file using DataFusion, storing the results in a Pandas DataFrame, and then plotting a chart.
The Parquet file used in this example can be downloaded from the following page:
from datafusion import SessionContext # Create a DataFusion context ctx = SessionContext() # Register table with context ctx.register_parquet('taxi', 'yellow_tripdata_2021-01.parquet') # Execute SQL df = ctx.sql("select passenger_count, count(*) " "from taxi " "where passenger_count is not null " "group by passenger_count " "order by passenger_count") # convert to Pandas pandas_df = df.to_pandas() # create a chart fig = pandas_df.plot(kind="bar", title="Trip Count by Number of Passengers").get_figure() fig.savefig('chart.png')
This produces the following chart:
You can use SessionContext's register_view method to convert a DataFrame into a view and register it with the context.
from datafusion import SessionContext, col, literal # Create a DataFusion context ctx = SessionContext() # Create sample data data = {"a": [1, 2, 3, 4, 5], "b": [10, 20, 30, 40, 50]} # Create a DataFrame from the dictionary df = ctx.from_pydict(data, "my_table") # Filter the DataFrame (for example, keep rows where a > 2) df_filtered = df.filter(col("a") > literal(2)) # Register the dataframe as a view with the context ctx.register_view("view1", df_filtered) # Now run a SQL query against the registered view df_view = ctx.sql("SELECT * FROM view1") # Collect the results results = df_view.collect() # Convert results to a list of dictionaries for display result_dicts = [batch.to_pydict() for batch in results] print(result_dicts)
This will output:
[{'a': [3, 4, 5], 'b': [30, 40, 50]}]
It is possible to configure runtime (memory and disk settings) and configuration settings when creating a context.
runtime = ( RuntimeEnvBuilder() .with_disk_manager_os() .with_fair_spill_pool(10000000) ) config = ( SessionConfig() .with_create_default_catalog_and_schema(True) .with_default_catalog_and_schema("foo", "bar") .with_target_partitions(8) .with_information_schema(True) .with_repartition_joins(False) .with_repartition_aggregations(False) .with_repartition_windows(False) .with_parquet_pruning(False) .set("datafusion.execution.parquet.pushdown_filters", "true") ) ctx = SessionContext(config, runtime)
Refer to the API documentation for more information.
Printing the context will show the current configuration settings.
print(ctx)
For information about how to extend DataFusion Python, please see the extensions page of the online documentation.
See examples for more information.
uv add datafusion
pip install datafusion # or python -m pip install datafusion
conda install -c conda-forge datafusion
You can verify the installation by running:
>>> import datafusion >>> datafusion.__version__ '0.6.0'
This project ships a SKILL.md that teaches AI coding assistants how to write idiomatic DataFusion Python. It follows the Agent Skills open standard.
Preferred: npx skills add apache/datafusion-python — installs the skill in Claude Code, Cursor, Windsurf, Cline, Codex, Copilot, Gemini CLI, and other supported agents.
Manual: paste this line into your project's AGENTS.md / CLAUDE.md:
For DataFusion Python code, see https://github.com/apache/datafusion-python/blob/main/skills/datafusion_python/SKILL.md
This assumes that you have rust and cargo installed. We use the workflow recommended by pyo3 and maturin. The Maturin tools used in this workflow can be installed either via uv or pip. Both approaches should offer the same experience. It is recommended to use uv since it has significant performance improvements over pip.
Currently for protobuf support either protobuf or cmake must be installed.
Bootstrap (uv):
By default uv will attempt to build the datafusion python package. For our development we prefer to build manually. This means that when creating your virtual environment using uv sync you need to pass in the additional --no-install-package datafusion and for uv run commands the additional parameter --no-project
# fetch this repo git clone git@github.com:apache/datafusion-python.git # cd to the repo root cd datafusion-python/ # create the virtual environment uv sync --dev --no-install-package datafusion # activate the environment source .venv/bin/activate
Bootstrap (pip):
# fetch this repo git clone git@github.com:apache/datafusion-python.git # cd to the repo root cd datafusion-python/ # prepare development environment (used to build wheel / install in development) python3 -m venv .venv # activate the venv source .venv/bin/activate # update pip itself if necessary python -m pip install -U pip # install dependencies python -m pip install -r pyproject.toml
The tests rely on test data in git submodules.
git submodule update --init
Whenever rust code changes (your changes or via git pull):
# make sure you activate the venv using "source venv/bin/activate" first maturin develop --uv python -m pytest
Alternatively if you are using uv you can do the following without needing to activate the virtual environment:
uv run --no-project maturin develop --uv uv run --no-project pytest
To run the FFI tests within the examples folder, after you have built datafusion-python with the previous commands:
cd examples/datafusion-ffi-example uv run --no-project maturin develop --uv uv run --no-project pytest python/tests/_test_*py
datafusion-python takes advantage of pre-commit to assist developers with code linting to help reduce the number of commits that ultimately fail in CI due to linter errors. Using the pre-commit hooks is optional for the developer but certainly helpful for keeping PRs clean and concise.
Our pre-commit hooks can be installed by running pre-commit install, which will install the configurations in your DATAFUSION_PYTHON_ROOT/.github directory and run each time you perform a commit, failing to complete the commit if an offending lint is found allowing you to make changes locally before pushing.
The pre-commit hooks can also be run adhoc without installing them by simply running pre-commit run --all-files.
NOTE: the current pre-commit hooks require docker, and cmake. See note on protobuf above.
There are scripts in ci/scripts for running Rust and Python linters.
./ci/scripts/python_lint.sh ./ci/scripts/rust_clippy.sh ./ci/scripts/rust_fmt.sh ./ci/scripts/rust_toml_fmt.sh
This project includes an AI agent skill for auditing which features from the upstream Apache DataFusion Rust library are not yet exposed in these Python bindings. This is useful when adding missing functions, auditing API coverage, or ensuring parity with upstream.
The skill accepts an optional area argument:
scalar functions aggregate functions window functions dataframe session context ffi types all
If no argument is provided, it defaults to checking all areas. The skill will fetch the upstream DataFusion documentation, compare it against the functions and methods exposed in this project, and produce a coverage report listing what is currently exposed and what is missing.
The skill definition lives in .ai/skills/check-upstream/SKILL.md and follows the Agent Skills open standard. It can be used by any AI coding agent that supports skill discovery, or followed manually.
To change test dependencies, change the pyproject.toml and run
uv sync --dev --no-install-package datafusion