feat: expose spark-compatible functions (#1564)
* feat: expose Spark-compatible functions (#1482)
Add `datafusion.functions.spark` module exposing the upstream
`datafusion-spark` crate's UDF/UDAF library (~87 functions across string,
math, datetime, hash, array, aggregate, bitwise, bitmap, conditional,
collection, conversion, json, map, url categories).
For DataFrame use, import the typed Python wrappers from
`datafusion.functions.spark`. For SQL use, call
`SessionContext.enable_spark_functions()` to register the Spark UDFs by
name (overriding DataFusion built-ins of the same name with their Spark
semantics — NULL-propagating `concat`, 1-indexed `substring`, HALF_UP
`round`, etc.).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* chore: drop unused borrow_deref_ref allows
Seven `#[allow(clippy::borrow_deref_ref)]` attributes on module
declarations in `crates/core/src/lib.rs` had become stale — the only
remaining lint hit was a redundant `&*x.as_str()` pattern in
`parse_file_compression_type`. Rewriting that call to
`&x.unwrap_or_default()` lets every allow come off, removing noise that
new modules were copying without need.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* refactor: tighten spark_functions macros via expr_fn
Switch most spark wrappers from UDF-direct path (which forced
`spark_udf_fixed!(name, fn_category::name, args...)` repetition) to a
`spark_expr_fn!` macro that mirrors the existing `expr_fn!` macro in
`functions.rs`, so calls collapse to `spark_expr_fn!(sha2, arg1
bit_length);`.
UDF-direct retained for genuinely variadic functions whose upstream
`expr_fn` wrappers were generated with a single-`Expr` arm by
`export_functions!` (concat, array, xxhash64, parse_url family, etc.) so
that the Python side keeps its `*args` ergonomics.
Aggregates collapse the same way via `spark_aggregate!` mirroring
`aggregate_function!`. Net 173 lines removed.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* docs: clarify spark functions cover DataFrame API too
The intro wording implied "SQL functions" only; the same wrappers are the
primary entry point for the DataFrame API as well.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* docs: rewrite spark DataFrame intro for users
Replace API-speak ("Import the submodule", "Returned values are Expr
instances that compose") with a concrete description of where users can
actually drop these calls.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* docs: defer spark function list to API reference
Hand-maintained category list would drift from the actual module as
upstream `datafusion-spark` adds/removes functions. Replace with a
pointer to the AutoAPI-generated reference, which renders from the
module itself.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* test: replace spark doctest skips with verified examples
38 wrappers carried `# doctest: +SKIP` because outputs weren't verified at
authoring time. Run each with concrete inputs, capture actual outputs, and
inline the values so the doctests execute and stay correct.
Covers datetime (20), URL (5), bitmap (3), map (3), and remaining hash,
JSON, math, string, conversion, and format_string cases. Net new doctest
coverage: 65 examples now run that were skipped before; total skipped
across the suite drops from 53 to 12.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* refactor(spark): rename function params to match pyspark
Align positional parameter names in `functions.spark` with pyspark.sql.functions:
- aggregate first positional → `col` (avg, try_sum, collect_list, collect_set)
- unary `arg` → `col` across math/string/byte/datetime helpers
- multi-arg renames: array_contains (col, value), array (*cols), shuffle (col),
array_repeat (col, count), slice (x, start, length), shiftleft/right/rightunsigned
(col, numBits), add_months (start, months), date_add/sub (start, days),
date_diff (end, start), date_trunc (format, timestamp), time_trunc (unit, time),
trunc (date, format), next_day (date, dayOfWeek), from/to_utc_timestamp
(timestamp, tz), sha2 (col, numBits), xxhash64 (*cols), map_from_arrays
(col1, col2), width_bucket (v, min, max, numBucket), substring (str, pos, len),
concat (*cols), elt (*inputs), is_valid_utf8/make_valid_utf8 (str)
Bodies updated to reference the new names; positional callers unaffected.
This finishes Category 1 / Category 4 (spark-side BOTH-bucket) renames from
PYSPARK_ALIGNMENT_PLAN.md PR 1.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* feat(spark): make pyspark-optional params optional
Match pyspark's optional-parameter surface in the spark namespace:
- make_dt_interval, make_interval: all parts default to zero (int32 0 / lit 0.0)
- str_to_map: pair_delim defaults to ',', key_value_delim defaults to ':'
- round: scale defaults to 0 (HALF_UP rounding to nearest integer)
- shuffle: accepts `seed` kwarg for pyspark parity; raises NotImplementedError
for non-None values until the Rust binding supports it
- like, ilike: accept `escapeChar` for pyspark parity; same NotImplementedError
guard; first positional renamed `string` → `str` to match pyspark
ceil/floor `scale=` deferred — the underlying Rust expr_fn is single-arg.
Added a module-level `_ZERO_I32` literal to avoid rebuilding the pyarrow
int32 zero scalar on every call.
Tests: positional-compat coverage for aggregates (`spark.avg(col)` etc.),
defaults-omitted cases for the optional-arg functions, and
NotImplementedError cases for `shuffle(seed=)` and `like/ilike(escapeChar=)`.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* refactor(spark): reshape varargs to match pyspark signatures
Replace generic ``*args`` with explicit pyspark-style signatures:
- json_tuple(col, *fields) — first positional is the JSON expr
- format_string(format, *cols) — `format` is the printf template; a plain
``str`` is auto-promoted to a literal
- parse_url(url, partToExtract, key=None) — `key` is optional and only
meaningful with ``partToExtract='QUERY'``
- try_parse_url(url, partToExtract, key=None) — same shape
- url_decode(str), try_url_decode(str), url_encode(str) — single-argument
forms (multi-arg calls were always semantically wrong)
Tests cover the three-arg parse_url path and the plain-str format_string
auto-promotion.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* docs(skills): cover the new spark function namespace
`functions.spark` mirrors `pyspark.sql.functions` and now ships on this
branch. Update every skill that references the function surface:
- skills/datafusion_python/SKILL.md (user-facing): add an import
reference, a Core Abstractions row, and a "Spark-Compatible Functions"
subsection listing coverage by category, the SQL-vs-DataFrame usage
(`enable_spark_functions`), and the divergent-semantics table
(concat NULL, round HALF_UP, trunc) so callers know which namespace
to pick.
- .ai/skills/check-upstream/SKILL.md: new area for the `datafusion-spark`
crate with the coverage policy (parity with pyspark, extras allowed
when positional pyspark calls still work). Hygiene check also now
spans `functions/spark.py`'s `__all__`.
- .ai/skills/audit-skill-md/SKILL.md: add `functions.spark` to the
surface table and a `spark-functions` scope so this audit also
validates the new subsection and divergent-semantics table.
- .ai/skills/make-pythonic/SKILL.md: explicit scope note that the
spark namespace is a deliberate pyspark mirror — generic native-type
coercion does not apply there. Path references updated to the new
`functions/__init__.py` module layout.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* docs(skills): drop references to PYSPARK_ALIGNMENT_PLAN.md
The plan file is a working document, not a committed artifact, so skills
must not point at it. Inline the one substantive reference (the
"deferred to follow-up PRs" callout in make-pythonic) and drop the
cross-cutting pointer from check-upstream.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* docs(skills): make-pythonic also targets functions.spark
Previous guidance said to skip the spark namespace entirely. That was
wrong: the spark namespace should also feel pythonic — it just carries
the extra constraint that every signature must remain compatible with
pyspark.sql.functions (parameter names, positional order, accepted input
types). Pythonic widenings like `Expr → Expr | int` are on-brand there
because pyspark itself accepts the int form.
Rewrite the scope section to spell out the compatibility rules (keep
parameter names/order; widen input types, never narrow; extra kwargs
default to None) and extend "How to Identify Candidates" to include
`functions/spark.py`.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* docs(skill): point at spark __all__ instead of enumerating
Enumerating spark functions in the user-facing skill duplicates the
__all__ list in python/datafusion/functions/spark.py and will drift the
moment a new function lands or is renamed. Replace the per-function
listing with a category summary and a discovery snippet that queries
the actual __all__ at runtime, which is the authoritative source.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* docs(spark): use isoformat in spark_cast doctest
pyarrow tzinfo repr differs across versions (<UTC> vs
zoneinfo.ZoneInfo(key='UTC')), breaking the doctest on some platforms.
isoformat is stable across versions.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* docs(spark): fix map_from_entries doctest to use the right function
The example called map_from_arrays, so it never exercised
map_from_entries. Build an array-of-struct input and call the
documented function.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* docs(spark): list Spark-compatible aggregates in aggregations guide
Add avg, try_sum, collect_list, and collect_set under a dedicated
Spark-Compatible Functions entry, with a note that the
datafusion.functions.spark namespace mirrors Spark semantics and may
differ from the like-named built-ins. Adds a (spark-functions) anchor
for the cross-reference.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* docs(spark): align if_ doctest with the single-value accessor style
Use the same lit-based single-row select and [0].as_py() accessor as
the other wrappers instead of the lone to_pylist() call.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* feat(spark): accept native Python literals for literal-friendly args
Audit the functions.spark namespace against pyspark.sql.functions and
widen arguments that pyspark types as a non-column literal so callers can
pass bare int/float/str instead of wrapping in lit():
- int args: array_repeat count, slice start/length, shiftleft/shiftright/
shiftrightunsigned numBits, sha2 numBits, round scale, substring pos/len,
width_bucket numBucket
- int32-coerced args (binding requires int32): add_months months,
date_add/date_sub days, space n, make_dt_interval/make_interval parts
- float args: modulus/pmod operands; make_*_interval secs
- str args: next_day dayOfWeek, date_trunc/trunc format, date_part field,
from_utc_timestamp/to_utc_timestamp tz, spark_cast type_str,
json_tuple *fields
- Any: array_contains value, if_ if_true/if_false
Arguments that pyspark types as ColumnOrName (str means column name, not a
literal) are left as Expr to avoid diverging from pyspark semantics:
ilike/like pattern, parse_url partToExtract/key, str_to_map delimiters,
bit_get pos, time_trunc unit.
Also rename str_to_map's delimiter params to pairDelim/keyValueDelim to
match pyspark exactly (they were pair_delim/key_value_delim).
Add a coercion test matrix and update docstring examples to show the
native-literal calling convention.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* feat(spark): accept column-name str for ColumnOrName args
For arguments that pyspark types as ColumnOrName, a bare str means a
column name (not a literal). Widen these to Expr | str and resolve a str
to a column reference via _to_raw_expr, matching pyspark semantics:
- ilike/like pattern
- parse_url/try_parse_url partToExtract and key
- str_to_map pairDelim/keyValueDelim
- bit_get pos
- time_trunc unit
Document the column-name behavior in each docstring and add a test
confirming a bare str resolves to a per-row column value.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* fix: correct grammar in file_compression_type error message
"must one of" → "must be one of".
Co-Authored-By: Claude Opus 4.8 (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