feat: ergonomic native DataFrame/Expr API (substrait.api)

[nielspardon]

Summary

substrait-python can already build any plan, but the day-to-day building experience lags behind sibling libraries — notably substrait-java, whose core module ships a fluent builder, a nullability-aware type factory, a preloaded extension collection, and named function helpers.

This PR adds a thin, additive facade over the existing substrait.builders layer that gives Python users the idioms they expect from pandas / Polars / PySpark / Ibis:

import substrait.dataframe as sub

plan = (
    sub.read_named_table("people", {"id": sub.i64, "age": sub.i64, "name": sub.string})
    .filter(sub.col("age") > 25)
    .with_columns(bonus=(sub.col("age") + 1) * 2)
    .select("id", "name", "bonus")
    .to_plan()
)

The facade is faithful: each verb / expression emits byte-identical protobuf to the equivalent raw builder call (asserted per feature). It began as filter/select/join and has since grown to cover essentially the whole Substrait query-building surface, and moves the project onto the latest spec (0.96).

Draft — opening for feedback on the API shape, the spec bump, and the Narwhals direction.

What's here

Entry point — import substrait.dataframe as sub. The higher-level API lives under a dedicated substrait.dataframe subpackage rather than at the substrait namespace root: substrait is a PEP 420 namespace package shared with substrait-protobuf (so the root can't hold an __init__.py), and grouping the modules under one owned subpackage — a sibling to builders / sql / narwhals, named for what it is rather than the generic "api" — keeps a single import surface and avoids scattering generic names (expr, frame, …) across the shared namespace. Nothing in builders / proto behaviour changes.

Expressions (Expr) — operator overloading (< <= > >= == != + - * / % ** & | ^ ~) mapping to the standard function extensions and resolving lazily, including decimal arithmetic (functions_arithmetic_decimal) for decimal operands; literal auto-wrap with peer-type coercion (col_fp64 * 2, col_i32 > 25, col_decimal > Decimal("9.99")), which raises rather than silently rounding a literal the column's decimal type can't represent exactly. Plus cast, alias, is_null / is_not_null / is_nan, between, is_distinct_from, is_in, coalesce, when().then()…otherwise() and switch (CASE), nested access (struct_field / [] / map_key), window functions via .over(partition_by=, order_by=, rows=/range=), higher-order list functions (list_transform / list_filter with Python-lambda callbacks), subqueries (scalar_subquery / exists / unique / in_subquery / any_ / all_, including correlated via outer()), dynamic parameters (parameter), and execution-context variables (current_timestamp / current_date / current_timezone).

DataFrame verbs

• reads: read_named_table, from_records (VALUES), read_parquet / read_csv / read_orc / read_arrow, read_extension_table, extension_leaf
• projection/rows: filter, select, with_columns, rename, drop, sort (per-column direction / nulls), limit / head / offset, top_n
• joins: join (all 13 join types + post_filter), cross_join, nested_loop_join, hash_join, merge_join
• set ops: union / intersect / except_
• aggregation: group_by().agg(), rollup, cube, explicit grouping sets, per-measure FILTER, DISTINCT, ordered aggregates
• reshape/physical: unpivot, repartition, broadcast, extension / extension_multi
• plan structure: hint
• write / DDL: write_named_table (CTAS / insert), create_table / create_view, drop_table / drop_view, update_table

Function namespace (f) — every scalar / aggregate / window function from the default extensions, generated lazily from the registry (dir(sub.f)-discoverable). Multi-extension names resolve to the right extension by argument type; keyword names via and_ / or_ / not_; function options as kwargs (f.add(x, y, overflow="ERROR")). functions_for(registry) / DataFrame.f surface custom-extension functions.

Types — nullability-aware shortcuts (sub.i64 nullable, sub.i64.non_null required) covering every concrete Substrait type; parametrized builders; user_defined UDTs.

User-defined extension relations — ExtensionLeafDetail / ExtensionSingleDetail / ExtensionMultiDetail ABCs (mirroring substrait-java's Extension.*RelDetail) with to_any / from_any / derive_schema, registered via ExtensionRegistry.register_extension_relation so schema inference follows a custom relation like a built-in one.

Spec bump: 0.86 → 0.96

Moves the substrait-protobuf / substrait-extensions / substrait-antlr pins to the latest spec (needed for TopNRel and execution-context variables, and to track upstream). Done as a bisected, two-step migration:

• removed spec features absorbed — AggregateRel.Grouping.grouping_expressions and the deprecated timestamp / time / timestamp_tz types (superseded by expression_references / the precision_* types), across type_inference, derivation_expression, and tests;
• ANTLR type-derivation grammar refactor (0.95 split the single BinaryExpr rule into MulDiv / AddSub / Comparison / Equality / And / Or + If / Not) migrated in derivation_expression.

Core enablers (beyond the pure facade)

Most verbs are thin wrappers, but a few features required additive changes to the lower layers — never a reimplementation; existing paths stay byte-identical:

• type_inference gained schema-inference cases for the new relations (expand, top_n, exchange, physical joins, extension relations) and expression types (lambda → func<>, dynamic_parameter, execution_context_variable), plus outer-schema resolution for correlated references;
• the registry's signature matcher gained func<…> matching so higher-order functions resolve and bind type parameters;
• correlated subqueries thread a small contextvar (an outer-schema stack) so inference resolves outer references.

Breaking changes

• Spec 0.86 → 0.96: the deprecated timestamp / time / timestamp_tz types and AggregateRel.Grouping.grouping_expressions are gone (superseded by the precision_* types / expression_references).
• substrait.dataframe reclaimed — the old (experimental, minimal) Narwhals wrapper that lived at substrait.dataframe moves to substrait.narwhals, and substrait.dataframe now is the native DataFrame/Expr API (the entry point earlier drafted as substrait.api). Consequence: existing import substrait.dataframe code silently binds to the native API instead of the Narwhals wrapper (a behaviour change, not an ImportError) — migrate those imports to substrait.narwhals.

Testing

New tests/dataframe/ plus builder golden tests, all asserting byte-identical protobuf to the raw builder path, with schema-inference and round-trip checks. Engine round-trip tests (pyarrow / DuckDB / DataFusion) lag the spec and can crash the interpreter natively on a newer-than-supported plan, so per a best-effort policy they are skipped unless SUBSTRAIT_ENGINE_TESTS=1; the non-engine suite is the gate and is green.

Follow-ups (not in this PR)

• CTEs / shared subplans (ReferenceRel) at both the builder and dataframe layers (DataFrame.cache()), carrying subtrees in-band in the Plan rather than via a contextvar — tracked in Native DataFrame + builders: CTEs / shared subplans (ReferenceRel) with in-band subtrees #211 (prototyped here, then pulled out to keep this PR focused).
• Build substrait.narwhals out into a full Narwhals compliant backend (CompliantLazyFrame / Expr / Namespace) on top of substrait.dataframe.frame. The compliant protocol is experimental, and collect() implies execution (a non-goal), so it would raise or delegate to a consumer.
• Positional enum function arguments (FunctionArgument.enum); type variations on built-in types; Iceberg read specifics (ReadRel.IcebergTable).
• Registry-derived operator resolution: resolve the arithmetic/comparison operators (+, >, …) over every extension the runtime registry defines for the function name — via the shared _resolve_over_urns loop, replacing the hardcoded standard + _decimal URN list — so operators work over user-defined types and type variations from custom extension YAMLs (the peer literal-coercion step, the operators' reason to differ from f.*, barely applies to UDTs). Broadens the overloads operators consider (e.g. date + interval via functions_datetime), so it needs its own tests and an intent decision.

🤖 Generated with AI

[tokoko]

I haven't looked through everything yet, but my +1 on the overall approach. There are some cases (lambda functions most notably) where narwhals/polars api and substrait representation genuinely diverge from one another. Having a native api and a narwhals wrapper is probably a better long-term decision.

[nielspardon]

I haven't looked through everything yet, but my +1 on the overall approach. There are some cases (lambda functions most notably) where narwhals/polars api and substrait representation genuinely diverge from one another. Having a native api and a narwhals wrapper is probably a better long-term decision.

Yeah, definitely a lot of Substrait features to be covered. Wasn't sure yet how to stage it and thought I stop after the first couple iterations to see what the community thinks.

[mbwhite]

Also +1 for the general approach.

I think there was a discussion on a Dask issue about DataFrames/Substrait... will need to find it

EDIT dask-contrib/dask-sql#1041

[tokoko]

looks great, thanks. let's continue with follow-up PRs