Name a bundle's bare capsules after the bundle

A codec handed to `with_extensions` as a bare `PyCapsule` fell through to
`anon:<uuid4>`, an id private to the session that installed it. Plans written
through it are undecodable anywhere else, and `with_extensions` accepts no
`codec_id=` to override that — so the workaround was to wrap the capsule in an
object declaring `__datafusion_codec_id__`, which nothing documented. A
distributed engine has to decode its plans in another process, so the shape it
would naturally ship — a Rust bundle handing over capsules, as
`MyPlannerExtension` does — was the one shape that could not work.

The bundle is the stable name that was missing. It is a plain Python object, so
its `module.QualName` is library-owned and exactly as stable across processes as
an exporting codec class's, which arm 3 of `derive_codec_id` already trusts. The
capsule was unnameable only because a capsule carries no type of its own, not
because nothing stable was in reach.

Resolve a capsule's id through the contributing bundle, using `derive_codec_id`
itself so the bundle inherits the same `__datafusion_codec_id__` escape hatch
against a class rename. The fallback applies only where randomness would have:
an id declared on the handed-over object, or that object's own class, still
wins, so an extension can name a codec directly.

Two bare capsules of one kind from one bundle collide and are refused. Numbering
them by position would be exactly the id `codec.rs` rejects for `anon:` — one
another library can mint the same value from — and would break stored plans the
first time the bundle reordered what it returns.

`resolve_codec_id` gains the bundle argument, `_install_extensions` takes
(codec, bundle) pairs, and the collision message now names both routes to a
distinct identity; it previously offered only `codec_id=`, which is unreachable
from `with_extensions`.

Covered in `python/tests/test_context.py`, which reaches every arm without a
built extension library: the bundle-derived name, an extension pinning its own
id, an id on the handed-over object winning, an exporting object keeping its
own, and the two-capsule collision. The cross-FFI case is pinned in the query
planner example, where a Rust bundle's capsules must report
`datafusion_ffi_query_planner_example.MyPlannerExtension` and no id may be
`anon:`.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
6 files changed
tree: d7a75bd55a44735c75ec313612177f783fff324d
  1. .ai/
  2. .claude/
  3. .github/
  4. benchmarks/
  5. ci/
  6. crates/
  7. dev/
  8. docs/
  9. examples/
  10. python/
  11. skills/
  12. .asf.yaml
  13. .dockerignore
  14. .gitignore
  15. .gitmodules
  16. .pre-commit-config.yaml
  17. AGENTS.md
  18. Cargo.lock
  19. Cargo.toml
  20. CHANGELOG.md
  21. conftest.py
  22. LICENSE.txt
  23. pyproject.toml
  24. README.md
  25. rust-toolchain.toml
  26. rustfmt.toml
  27. uv.lock
README.md

DataFusion in Python

Python test Python Release Build

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:

  • Dask SQL uses DataFusion's Python bindings for SQL parsing, query planning, and logical plan optimizations, and then transpiles the logical plan to Dask operations for execution.
  • DataFusion Ballista is a distributed SQL query engine that extends DataFusion's Python bindings for distributed use cases.
  • DataFusion Ray is another distributed query engine that uses DataFusion's Python bindings.

Features

  • Execute queries using SQL or DataFrames against CSV, Parquet, and JSON data sources.
  • Queries are optimized using DataFusion's query optimizer.
  • Execute user-defined Python code from SQL.
  • Exchange data with Pandas and other DataFrame libraries that support PyArrow.
  • Serialize and deserialize query plans in Substrait format.
  • Experimental support for transpiling SQL queries to DataFrame calls with Polars, Pandas, and cuDF.

For tips on tuning parallelism, see Maximizing CPU Usage in the configuration guide.

Example Usage

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:

Chart

Registering a DataFrame as a View

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]}]

Configuration

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)

Extensions

For information about how to extend DataFusion Python, please see the extensions page of the online documentation.

More Examples

See examples for more information.

Executing Queries with DataFusion

Running User-Defined Python Code

Substrait Support

How to install

uv

uv add datafusion

Pip

pip install datafusion
# or
python -m pip install datafusion

Conda

conda install -c conda-forge datafusion

You can verify the installation by running:

>>> import datafusion
>>> datafusion.__version__
'0.6.0'

Using DataFusion with AI coding assistants

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

How to develop

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

Running & Installing pre-commit hooks

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.

Running linters without using pre-commit

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

Checking Upstream DataFusion Coverage

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.

How to update dependencies

To change test dependencies, change the pyproject.toml and run

uv sync --dev --no-install-package datafusion