feat: close upstream coverage gaps for DataFusion 55.1.0

[timsaucer]

Which issue does this PR close?

No issue filed. The gaps were found by auditing the Python API against upstream DataFusion 55.1.0 with the check-upstream skill. Gaps that already have open issues (#1571–#1575, #1577, #1668, #1669) are left out.

Rationale for this change

Several upstream functions, optional arguments, and DataFrame methods were not reachable from Python. The audit also turned up bugs where options on aggregate and window functions were silently dropped.

What changes are included in this PR?

• New functions: array math and array_first, any_value, file metadata functions, and a batch of Spark functions and pyspark aliases.
• Optional arguments upstream already supports, such as distinct= on several aggregates, characters= on the trim functions, null_treatment= on lead/lag, and new DataFrame.explain options. explain raises ValueError when show_statistics is combined with analyze, matching upstream SQL.
• DataFrame.fill_nan.
• FFI: ScalarUDF and WindowUDF accept a bare PyCapsule, and the capsule overloads now type-check correctly. Passing the wrong kind of capsule to udf, udaf, or udwf now names the expected and the actual capsule.
• More Pythonic range and gen_series: they accept plain ints, and range(stop) works like Python's built-in.
• Fixes:
  • Chaining builder methods or .over() onto an aggregate or window function no longer discards options already set on it (ordering, null treatment, window frame). Whether a window frame is re-derived depends only on the expression, so it behaves the same after copy, pickle, or parsing from SQL. A builder option that does not apply to the function's kind still raises.
  • percentile_cont, quantile_cont, approx_percentile_cont, and approx_percentile_cont_with_weight keep the direction of their sort expression.
  • mean(x, filter=...) no longer raises TypeError.
  • string_agg rejects a non-bool distinct, so an old positional order_by call fails loudly instead of being used as the filter.
  • fill_null(subset=[]) and fill_nan(subset=[]) fill no columns instead of all of them.
  • spark.printf treats a bare str as a column name, as pyspark does.
  • The keyword-only form of the udf, udaf, and udwf decorators no longer raises IndexError.
• Docs: the substr guidance in skills/datafusion_python/SKILL.md now shows the length argument as valid, and the aggregate docstrings link to the guide with real :ref: roles.

Every new function has a doctest and pytest coverage. The same options are still dropped when .over() converts an aggregate to a window function; that is tracked separately in #1764 and noted in the over() docstring. fill_nan and fill_null fail on uppercase or dotted column names because of an upstream bug, filed as apache/datafusion#25829.

Are there any user-facing changes?

Yes, new functions, arguments, and methods. The following are documented in docs/source/user-guide/upgrade-guides.md:

• distinct is inserted before filter in bit_and, bit_or, mean, percentile_cont, quantile_cont, and string_agg, matching sum and avg in 54.0.0. Pass filter (and, for string_agg, order_by) by keyword.
• Chaining keeps options already set. Results can change, some chains that ran before now raise (for example array_agg(..., distinct=True).order_by(...)), and generated names for first_value, last_value, and nth_value now include RESPECT NULLS.
• The percentile functions honor a descending sort, and their generated names include the ordering.
• fill_null(subset=[]) fills no columns.
• spark.last_day's parameter is renamed from col to date to match pyspark.

🤖 Generated with Claude Code

[Copilot]

Copilot review overview

🟡 Changes recommended

Public explain enums need consistent exports, and identified documentation and distinct-option coverage gaps remain unresolved.

Get a fresh assessment by requesting another Copilot review.

Review effort: Balanced
Findings: 1 · 2

Open (3)

• Re-export new explain option enums at package root · New
• Use Sphinx ref role for aggregation guide link · New
• Assert distinct flag is preserved in bitwise expressions · New

What changed in this PR

Expands Python coverage for DataFusion 55.1.0 APIs and fixes option preservation across expression builders.

Changes:

• Adds array, aggregate, metadata, Spark, range, and DataFrame APIs.
• Preserves aggregate/window builder options and supports bare FFI capsules.
• Adds documentation, migration guidance, and Python integration tests.

File	Description
.ai/​skills/​check-upstream/​SKILL.md	Records to_string audit guidance.
crates/​core/​src/​dataframe.rs	Binds explain options and fill_nan.
crates/​core/​src/​expr.rs	Preserves expression-builder options.
crates/​core/​src/​functions.rs	Binds new functions and arguments.
crates/​core/​src/​spark_functions.rs	Adds Spark function bindings.
crates/​core/​src/​udf.rs	Imports bare scalar UDF capsules.
crates/​core/​src/​udwf.rs	Imports bare window UDF capsules.
docs/​source/​user-guide/​common-operations/​aggregations.md	Updates aggregate catalog.
docs/​source/​user-guide/​upgrade-guides.md	Documents breaking signature changes.
examples/​datafusion-ffi-example/​python/​tests/​_test_aggregate_udf.py	Tests bare aggregate capsules.
examples/​datafusion-ffi-example/​python/​tests/​_test_scalar_udf.py	Tests bare scalar capsules.
examples/​datafusion-ffi-example/​python/​tests/​_test_window_udf.py	Tests bare window capsules.
python/​datafusion/​catalog.py	Exports TableProviderFactory.
python/​datafusion/​dataframe.py	Exposes explain options and fill_nan.
python/​datafusion/​expr.py	Tracks explicit window frames.
python/​datafusion/​functions/​__init__.py	Adds functions and optional arguments.
python/​datafusion/​functions/​spark.py	Adds Spark functions and aliases.
python/​datafusion/​user_defined.py	Accepts bare UDF capsules.
python/​tests/​test_aggregation.py	Covers aggregate additions.
python/​tests/​test_dataframe.py	Covers DataFrame and window options.
python/​tests/​test_expr.py	Covers builder option preservation.
python/​tests/​test_functions.py	Covers scalar and array additions.
python/​tests/​test_lambda.py	Covers array_first.
python/​tests/​test_spark_functions.py	Covers Spark additions and aliases.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[timsaucer]

Not taking this one. The package root has grown to hold nearly every public class, which makes it hard to predict where anything lives: ExplainFormat is at the root, and so is WindowFrame, but Window is not. 55.0.0 already started the other direction: the new extension protocols live only in datafusion.extensions, and SessionExtensionComponents stays at the root "because a bundle constructs one rather than merely naming it". ExplainAnalyzeLevel and ExplainMetricCategory are only ever named as arguments to DataFrame.explain, so they stay in datafusion.dataframe. I've opened #1771 to apply the same rule to the rest of the root, with a deprecation period, rather than extend the current pattern here.

[timsaucer]

Fixed in 7259819: docs: fix the broken guide links in aggregate docstrings

[timsaucer]

Fixed in d779502: test: check that bit_and and bit_or keep distinct in the expression

[andygrove]

Thanks for this, it closes a lot of gaps. Each inline comment has a repro, run against both this PR's head (96f990b) and its merge base (90f8709).

The biggest theme is the new keep-existing-options behavior when chaining builder methods or .over(). It makes some things work that didn't before (filter/distinct on an aggregate window now build correctly), but:

• options that don't apply to the function's kind are now silently dropped instead of raising (builder_from_expr),
• the explicit-frame flag lives only on the Python object, so it's lost through copy/pickle/from_bytes/SQL,
• the result and column-name changes aren't in the upgrade guide yet.

A few comments are about bugs that predate this PR but sit in code it touches (percentile_cont's sort, .over() on aggregates / #1764, the keyword-only udf decorator). Those are fine to split out.

One more that can't go inline because the file isn't in the diff: skills/datafusion_python/SKILL.md (lines 747-760) still says F.substr takes only two arguments, that a third raises TypeError, and labels F.substr(col("c_phone"), lit(1), lit(2)) as WRONG. With the new length argument that call works (returns '25'), so the skill now steers agents away from a valid form.

[andygrove]

Because the seeded builder always carries a function kind, upstream's per-method kind checks in ExprFunctionExt no longer apply. partition_by/window_frame on a plain aggregate, and distinct on a non-aggregate window function, now build and the option is silently dropped. On main each of these raises at build() with "ExprFunctionExt can only be used with Expr::AggregateFunction or Expr::WindowFunction":

from datafusion import SessionContext, col, functions as f
from datafusion.expr import WindowFrame

ctx = SessionContext()
df = ctx.from_pydict({"g": [1, 1, 1, 2], "v": [3, 1, 2, 4]})

e = f.sum(col("v")).partition_by(col("g")).build()
print(e)                                             # Expr(sum(v))
print(df.aggregate([], [e.alias("s")]).to_pydict())  # {'s': [10]}

print(f.sum(col("v")).window_frame(WindowFrame("rows", 1, 0)).build())  # Expr(sum(v))
print(f.lead(col("v"), order_by="v").distinct().build())
# Expr(lead(DISTINCT v, Int64(1), NULL) ORDER BY [...]) -- runs, DISTINCT ignored

The new filter/distinct support on aggregate windows (f.sum(..).over(..).distinct()) does work, so one way to keep it while restoring the errors: in partition_by/window_frame, fall back to upstream's self.expr.clone().partition_by(..) / .window_frame(..) when self.expr is an AggregateFunction, and in distinct do the same for a window function whose definition isn't WindowFunctionDefinition::AggregateUDF.

[timsaucer]

Fixed in 92ace16: fix: raise when a builder option does not apply to the function kind

[andygrove]

The upgrade guide covers the distinct insertion, but this one positional shape doesn't fail loudly like the others do. string_agg(expr, ",", None, order_col) now binds None to distinct (the Rust Option<bool> accepts it) and the ordering column to filter:

from datafusion import SessionContext, col, functions as f

ctx = SessionContext()
df = ctx.from_pydict({"s": ["c", "a", "b", "d"], "n": [3, 1, 0, 2]})
e = f.string_agg(col("s"), ",", None, col("n"))
print(df.aggregate([], [e.alias("r")]).to_pydict())
# main: {'r': ['b,a,d,c']}  (ORDER BY n)
# PR:   {'r': ['c,a,d']}    (n used as the FILTER)

The other shifted calls raise TypeError: 'Expr' object is not an instance of 'bool'. A guard such as if not isinstance(distinct, bool): raise TypeError(...) would make this one fail too.

[timsaucer]

Fixed in 0ac494e: fix: reject a non-bool distinct in string_agg

[andygrove]

Pre-existing, but since this call is being edited: add_builder_fns_to_aggregate starts from an empty builder (agg_fn.null_treatment(None)), so build() overwrites the WITHIN GROUP ordering that upstream's percentile_cont(sort, p) stored. A descending sort is silently ignored:

from datafusion import SessionContext, col, functions as f

ctx = SessionContext()
df = ctx.from_pydict({"a": [1.0, 2.0, 3.0, 4.0, 5.0]}, name="t")

e = f.percentile_cont(col("a").sort(ascending=False), 0.25)
print(e)                                             # Expr(percentile_cont(a, Float64(0.25))), no DESC
print(df.aggregate([], [e.alias("p")]).to_pydict())  # {'p': [2.0]}
print(ctx.sql(
    "SELECT percentile_cont(0.25) WITHIN GROUP (ORDER BY a DESC) AS p FROM t"
).to_pydict())                                       # {'p': [4.0]}

quantile_cont, approx_percentile_cont (1.75 vs 4.25 from SQL) and approx_percentile_cont_with_weight go through the same path. Passing the sort expression through as order_by here would keep it. Relatedly, the docstring's "this function ignores the options order_by and null_treatment" isn't right for order_by: f.percentile_cont(col("a"), 0.25).order_by(col("a").sort(ascending=False)).build() returns 4.0.

[timsaucer]

Fixed in 1b01845: fix: keep the sort direction in percentile_cont and friends

[andygrove]

Upstream DataFrame::fill_nan (via fill_columns) rebuilds every column with col(field.name()), which parses the name as an identifier (lowercasing it, splitting on .). So fill_nan fails on any DataFrame with an uppercase or dotted column name, even when that column isn't in subset:

from datafusion import SessionContext

nan = float("nan")
df = SessionContext().from_pydict({"Name": ["a", "b"], "v": [nan, 1.0]})
df.fill_nan(0.0, subset=["v"])
# Schema error: No field named name. Did you mean '..."Name"'?

fill_null has the same upstream bug, so this is probably worth an upstream issue (ident(field.name()) there would fix both). Until then the wrapper could build the projection itself from unparsed column references plus nanvl.

[timsaucer]

Confirmed in pure Rust on 55.1.0, for both Name and a.b column names, and filed upstream as apache/datafusion#25829 with a suggested fix (Expr::Column(Column::from((qualifier, field))), which also keeps the qualifier). I've left the wrapper as-is so fill_nan and fill_null keep behaving the same way, and we'll pick up the upstream fix when it lands.

[andygrove]

The PR description says chaining builder methods or .over() onto an already-configured function may now give different results, but that isn't in the upgrade guide yet. Some examples of what changes:

from datafusion import SessionContext, col, functions as f
from datafusion.expr import Window

ctx = SessionContext()
df = ctx.from_pydict(
    {"g": [1, 1, 1, 2, 2], "t": [1, 2, 3, 4, 5], "v": [10, None, 30, 40, 50]}
)
e = f.lead(col("v"), 1, partition_by=[col("g")], order_by="t").over(Window(order_by="t"))
print(df.select(col("t"), e.alias("r")).sort(col("t")).collect_column("r").to_pylist())
# main: [None, 30, 40, 50, None]
# PR:   [None, 30, None, 50, None]  (partition now kept)

df2 = ctx.from_pydict({"s": ["a", "b", "a"], "v": [3, 1, 2]})
e2 = f.array_agg(col("s"), distinct=True).order_by(col("v")).build()
print(df2.aggregate([], [e2.alias("r")]).to_pydict())
# main: {'r': [['b', 'a', 'a']]}  (distinct silently dropped)
# PR:   Execution error: In an aggregate with DISTINCT, ORDER BY expressions must appear in argument list

Auto-generated names change too. f.first_value(col("a")).order_by(col("b")).build(), the chained form shown in aggregations.md, was named first_value(a) ORDER BY [b ASC NULLS FIRST] on main and is now first_value(a) RESPECT NULLS ORDER BY [b ASC NULLS FIRST]. That matches the non-chained form, but code that selects the unaliased column by name will break. last_value and nth_value change the same way. Expr.over's docstring also still describes only aggregates.

[timsaucer]

Fixed in 6b52225: docs: document that chaining now keeps options already set

[andygrove]

show_statistics is ignored when analyze=True. Upstream's LogicalPlan::Analyze has statement-level overrides for analyze_level and analyze_categories, but none for statistics:

from datafusion import SessionContext, col, lit

df = SessionContext().from_pydict({"a": [1, 2]}).filter(col("a") > lit(1))
df.explain(show_statistics=True)                # physical plan includes statistics=[...]
df.explain(analyze=True, show_statistics=True)  # no statistics

Upstream SQL rejects the combination (EXPLAIN option COSTS cannot be combined with ANALYZE), so raising a ValueError here, or at least saying so in the docstring, would match.

[timsaucer]

Fixed in 607dde4: fix: reject show_statistics combined with analyze in explain

[andygrove]

pyspark's printf(format: ColumnOrName, *cols) treats a bare str as a column name, unlike format_string(format: str, ...). As an alias of format_string, this silently returns the column name's text for pyspark-style calls:

from datafusion import SessionContext, col
from datafusion.functions import spark

df = SessionContext().from_pydict({"a": ["aa%d%s"], "b": [123], "c": ["cc"]})
print(df.select(spark.printf("a", col("b"), col("c")).alias("v")).to_pydict())
# {'v': ['a']}; pyspark's documented sf.printf("a", "b", "c") gives 'aa123cc'
print(df.select(spark.printf(col("a"), col("b"), col("c")).alias("v")).to_pydict())
# {'v': ['aa123cc']}

Other functions in this module already treat a bare str as a column name to match pyspark (e.g. bit_get's pos).

[timsaucer]

Fixed in 4933fa2: fix: treat a bare str as a column name in spark.printf

[andygrove]

Now that a bare capsule is accepted, passing the wrong kind is easy, and pointer_checked on its own only gives CPython's fixed message:

from datafusion import SessionContext, udf

ctx = SessionContext()
udf(ctx.__datafusion_logical_extension_codec__())
# ValueError: PyCapsule_GetPointer called with incorrect name
ctx.register_table("t", ctx.__datafusion_logical_extension_codec__())
# ValueError: Expected name 'datafusion_table_provider' in PyCapsule, instead got 'datafusion_logical_extension_codec'

validate_pycapsule in crates/util/src/lib.rs exists for this, and its doc comment says to call it first at every extraction site. Adding validate_pycapsule(capsule, "datafusion_scalar_udf")?; here, plus the equivalent in window_udf_from_capsule and aggregate_udf_from_capsule, would fix it.

[timsaucer]

Fixed in 23ed533: fix: name the expected and found capsule when importing a UDF

[andygrove]

subset=[] fills every floating-point column, not none. The binding maps both None and [] to an empty Vec, which upstream treats as "all columns":

from datafusion import SessionContext

nan = float("nan")
df = SessionContext().from_pydict({"a": [1.0, nan, None], "b": [nan, 2.0, 3.0]})
print(df.fill_nan(0.0, subset=[]).to_pydict())
# {'a': [1.0, 0.0, None], 'b': [0.0, 2.0, 3.0]}

A computed subset that happens to be empty ([c for c in cols if pred(c)]) would rewrite every float column. Returning self unchanged for [], or documenting it, would avoid the surprise. fill_null has the same behavior.

[timsaucer]

Fixed in 8454a09: fix: treat an empty fill_null or fill_nan subset as no columns

[andygrove]

Pre-existing (main behaves the same, and udaf too), but since this line changed: args[0] is read before the if args and ... guard below, so the keyword-only decorator form raises IndexError:

import pyarrow as pa
from datafusion import udf, udwf

udf(input_fields=[pa.int64()], return_field=pa.int64(), volatility="immutable")
# IndexError: tuple index out of range
udwf(input_types=[pa.int64()], return_type=pa.int64(), volatility="immutable")
# IndexError: tuple index out of range

if args and (hasattr(args[0], "__datafusion_scalar_udf__") or _is_pycapsule(args[0])): fixes it, and the same at line 1101 for udwf.

[timsaucer]

Fixed in f8d7b70: fix: allow the keyword-only form of the udf, udaf, and udwf decorators

[timsaucer]

@andygrove re the SKILL.md substr note in the review summary: fixed in db5b7cc: docs(skill): show substr's length argument as valid

[timsaucer]

@andygrove Thanks for the review! I think we've addressed all of the comments, Claude and I ;)