GitHub - apache/datafusion-python: Apache DataFusion Python Bindings
This is a Python library that binds to
in-memory query engine
.
DataFusion's Python bindings can be used as a foundation for building new data systems in Python. Here are some examples:
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.
is a distributed SQL query engine that extends DataFusion's Python bindings for distributed use cases.
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
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:
https://www.nyc.gov/site/tlc/about/tlc-trip-record-data.page
fromdatafusionimportSessionContext# Create a DataFusion contextctx=SessionContext() # Register table with contextctx.register_parquet('taxi', 'yellow_tripdata_2021-01.parquet') # Execute SQLdf=ctx.sql("select passenger_count, count(*) ""from taxi ""where passenger_count is not null ""group by passenger_count ""order by passenger_count") # convert to Pandaspandas_df=df.to_pandas() # create a chartfig=pandas_df.plot(kind="bar", title="Trip Count by Number of Passengers").get_figure() fig.savefig('chart.png')This produces the following 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.
fromdatafusionimportSessionContext, col, literal# Create a DataFusion contextctx=SessionContext() # Create sample datadata= {"a": [1, 2, 3, 4, 5], "b": [10, 20, 30, 40, 50]} # Create a DataFrame from the dictionarydf=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 contextctx.register_view("view1", df_filtered) # Now run a SQL query against the registered viewdf_view=ctx.sql("SELECT * FROM view1") # Collect the resultsresults=df_view.collect() # Convert results to a list of dictionaries for displayresult_dicts= [batch.to_pydict() forbatchinresults] 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
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
.
More Examples
See
for more information.
Executing Queries with DataFusion
Query a Parquet file using SQL
Query a Parquet file using the DataFrame API
Run a SQL query and store the results in a Pandas DataFrame
Run a SQL query with a Python user-defined function (UDF)
Run a SQL query with a Python user-defined aggregation function (UDAF)
Running User-Defined Python Code
Register a Python UDF with DataFusion
Register a Python UDAF with DataFusion
Substrait Support
Serialize query plans using Substrait
How to install
uv
uv add datafusionPip
pip install datafusion # or python -m pip install datafusionConda
conda install -c conda-forge datafusionYou can verify the installation by running:
>>>importdatafusion>>>datafusion.__version__'0.6.0'Using DataFusion with AI coding assistants
This project ships a
that teaches AI coding assistants how to write idiomatic DataFusion Python. It follows the
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
and
. 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
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 [email protected]:apache/datafusion-python.git # cd to the repo rootcd datafusion-python/ # create the virtual environment uv sync --dev --no-install-package datafusion # activate the environmentsource .venv/bin/activateBootstrap (pip):
# fetch this repo git clone [email protected]:apache/datafusion-python.git # cd to the repo rootcd datafusion-python/ # prepare development environment (used to build wheel / install in development) python3 -m venv .venv # activate the venvsource .venv/bin/activate # update pip itself if necessary python -m pip install -U pip # install dependencies python -m pip install -r pyproject.tomlThe tests rely on test data in git submodules.
git submodule update --initWhenever 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 pytestAlternatively 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 pytestTo 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_*pyRunning & Installing pre-commit hooks
datafusion-python takes advantage of
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.shChecking Upstream DataFusion Coverage
This project includes an
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
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