feat: ergonomic native DataFrame/Expr API (substrait.api)#204
Conversation
…ter_functions Extend the builder literal() to construct a literal for every Substrait type (decimal, uuid, precision time/timestamp[_tz], all interval kinds, struct, list, map with empty-list/empty-map handling, and typed nulls via value=None) through a new recursive _make_literal helper. Existing kinds remain byte-identical. Add the missing precision_time case to type_inference.infer_literal_type so every kind round-trips, and add ExtensionRegistry.iter_functions() to enumerate every registered (urn, name, function_type).
…arwhals The module is the Narwhals integration layer, not a general DataFrame; rename it to reflect that role and free the "DataFrame" name for the native frame. BREAKING CHANGE: import substrait.narwhals instead of substrait.dataframe. The module was a minimal, experimental Narwhals wrapper, so impact is expected low.
… type coverage Add substrait.api, a shallow front door over the existing builders: - expr.Expr: operator overloading (comparison/arithmetic/boolean), literal auto-wrap with peer-type coercion, and .cast()/.alias()/.is_null() - frame.DataFrame: chainable filter/select/with_columns/sort/limit/join/ group_by().agg(), carrying an ExtensionRegistry so it is not threaded through every call - functions.f: every scalar/aggregate/window function, generated lazily from the registry; multi-extension names resolved by argument type; functions_for() and DataFrame.f expose custom-registry functions - dtypes: nullability-aware type shortcuts (sub.i64 / sub.i64.non_null) covering every concrete Substrait type The facade is faithful: it emits byte-identical protobuf to the equivalent builder calls. Adds tests/api covering expressions, frame verbs, function coverage, type coverage and literal construction.
Add examples/api_example.py demonstrating the native substrait.api. Update narwhals_example.py to the renamed substrait.narwhals and label it as the Narwhals integration example. Remove dataframe_example.py, whose direct wrapper usage is superseded by api_example.py (native) and narwhals_example.py.
|
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. |
|
Also +1 for the general approach. I think there was a discussion on a Dask issue about DataFrames/Substrait... will need to find it |
Add an optional post_join_filter to plan.join (applied to the join output) and let plan.fetch accept count=None, emitting a FetchRel with count_expr left unset -- "all remaining rows" after the offset. Both changes are backward compatible: existing callers pass the same arguments and get identical protobuf.
Expand the Expr surface with expressions the operators could not express, each delegating to an existing extended_expression builder: - when().then()...otherwise() CASE (if_then) and Expr.switch (switch) - Expr.is_in (singular_or_list) - % and ** operators (modulus, power) plus reflected forms - ^ (boolean xor) - between, coalesce, is_nan, is_distinct_from, is_not_distinct_from Re-export when() and coalesce() from substrait.api. Each addition is covered by a serialize-equality test against the raw builder path.
…ame/drop Surface more Substrait relations on the native fluent frame, delegating to the plan builders (byte-identical protobuf, verified per verb): - set ops: union / intersect / except_ with SQL-accurate distinct/all defaults, plus n-ary union - cross_join and the write_named_table sink (error/append/replace/ignore) - all 13 JoinRel join types (was 6): right_semi/anti, left/right_single, left/right_mark; plus a post_filter predicate on join - per-column sort direction and null placement (bool or per-column list), covering all four asc/desc x nulls-first/last SortDirections - head / offset conveniences over FetchRel - rename / drop projections (schema-aware; unknown columns raise) Overlapping column names after a join are kept as-is and disambiguated explicitly via rename/drop -- auto-suffixing would break the join condition's by-name references.
…d ordered aggregates Extend the aggregate machinery, all backward compatible: - plan.aggregate gains grouping_sets (a list of index lists into grouping_expressions, one Grouping each -- GROUPING SETS / ROLLUP / CUBE) and filters (a per-measure FILTER (WHERE ...) list). Omitting both reproduces the previous single-grouping, unfiltered output. - aggregate_function gains invocation (ALL vs DISTINCT) and sorts (order-sensitive aggregates), merging any extensions the sort keys introduce.
…der_by Surface the aggregation features on the native frame and Expr: - DataFrame.rollup / cube and group_by(..., grouping_sets=[...]) (sets given by key name or position); GroupBy and one-shot aggregate() route through the same path - Expr.filter(pred) -> a Measure (agg(x) FILTER (WHERE ...)), Polars-style; agg(...) accepts Expr or Measure and threads per-measure filters - Expr.distinct() (COUNT(DISTINCT x)) and Expr.order_by(*keys, ...) (string_agg(x ORDER BY ...)), post-processing the resolved measure; Measure delegates alias/distinct/order_by so any chain order works Each verb is covered by a serialize-equality test against the builder path.
Only read_named_table existed; add builders for the other ReadRel read types, each with a golden-proto test: - virtual_table: inline VALUES rows as Expression.Nested.Struct - local_files: a ReadRel over pre-built FileOrFiles items - extension_table: a custom source carrying a google.protobuf.Any detail Factor the base-schema nullability check and the read Plan wrapper into _require_schema / _read_plan helpers.
Add DataFrame entry points beyond read_named_table, re-exported from substrait.api: - from_records(data, schema): inline rows from dicts or positional tuples, typed per schema (None -> typed null) - read_parquet / read_csv / read_orc / read_arrow: local-file reads (read_csv takes delimiter / header_lines_to_skip) - read_extension_table(schema, detail): a custom source Each is covered by a serialize-equality test against the builder path.
…Y-ALL) Add builders that embed an inner query's Rel into an Expression.Subquery, merging the inner plan's extension declarations into the outer expression: - scalar_subquery: one-row/one-column Scalar subquery - set_predicate: EXISTS / UNIQUE (SetPredicate) - in_predicate: needles IN (subquery) (InPredicate) - set_comparison: left <op> ANY/ALL (subquery) (SetComparison) Each accepts a Plan or an UnboundPlan (registry -> Plan) as the query, so a DataFrame's underlying plan can be passed directly.
Nested access, post-processing the FieldReference segment chain so it works by index/offset/key without nested-schema introspection: - Expr.struct_field(i), Expr[i] / Expr.list_element(i), Expr.map_key(k), chainable Subquery expressions over an inner DataFrame (uncorrelated): - scalar_subquery(df), exists(df), unique(df) as sub.* constructors - Expr.in_subquery(df) (InPredicate) - col > any_(df) / all_(df): comparison operators detect an ANY/ALL reduction sentinel and build a SetComparison Comparison operators refactored through a shared _compare helper. New constructors re-exported from substrait.api.
- write_named_table gains op (WriteOp) and output_mode, defaulting to CTAS so existing callers are unchanged - ddl: a DdlRel builder for CREATE / CREATE_OR_REPLACE / DROP / DROP_IF_EXIST of a TABLE or VIEW; CREATE VIEW embeds the query Rel and infers the view schema when none is given - update: an UpdateRel builder with per-column TransformExpressions (by column index) and an optional condition Golden-proto tests for ddl and update added alongside the write test.
- DataFrame.write_named_table gains op ("ctas" default / "insert" / ...)
alongside the existing create mode
- create_table / create_view / drop_table / drop_view DDL constructors
(create_view embeds a query DataFrame and infers the view schema)
- update_table(name, schema, {col: expr}, where=): columns by name or
index, resolved against the given schema
Re-exported from substrait.api; covered by structural and
extension-merge tests plus a write-insert serialize-equality test.
Add Expr.over(...) which turns a window function (f.row_number/rank/ lead/lag/...) into a windowed expression, post-processing the resolved WindowFunction: - partition_by / order_by: a column name/expression or a list of them; descending / nulls_last control the ordering - rows=(start, end) / range=(start, end): a frame where each endpoint is an int offset (negative preceding, 0 current row, positive following) or None (unbounded), mapping to BOUNDS_TYPE_ROWS/RANGE + bounds Extensions introduced by partition/order keys are merged in. Both the window_function expression and window relation already resolve through the existing type inference, so no builder or core changes are needed.
First deliberate type_inference extension to unlock a new relation:
- infer_rel_schema gains an "expand" case: one output column per expand
field (from the consistent_field expression or the first switching_field
duplicate) plus the trailing i32 duplicate-index column ExpandRel appends;
the generic emit handling still applies
- plan.expand builds an ExpandRel from ("switching", [exprs]) /
("consistent", expr) field specs
Covered by a golden-proto test and an infer_plan_schema check.
Add a Polars-style unpivot(on, index=, variable_name=, value_name=): id columns become consistent expand fields, a "variable" switching field holds the unpivoted column names and a "value" switching field holds their values. The auto-appended i32 duplicate-index column is dropped via a trailing select for a clean [index..., variable, value] output. Covered by structure tests and a post-unpivot filter test that exercises the new expand schema inference.
… contextvar Add a `reference_subtrees` ContextVar and a "reference" case to infer_rel_schema: a ReferenceRel's schema is the schema of the subtree its subtree_ordinal indexes into. The frame sets the contextvar to the active build's shared subtrees while materializing a plan, so a cached subplan referenced elsewhere still resolves its schema. Additive; the contextvar defaults to None and only the new case reads it.
cache() marks a frame as a reusable common subplan. Within a to_plan() build, its resolver registers the subplan once (by identity) and returns a single-relation plan rooted at a ReferenceRel, so existing builders inline the tiny reference unchanged. _materialize() runs the build under a context, then prepends the collected subtrees as PlanRel(rel=...) with the ReferenceRel ordinals indexing into them; to_plan/to_substrait route through it. Subtree extensions propagate via the normal merge path. No builder changes were needed. Covered by tests for shared-subtree emission, uncached inlining, schema inference through a reference, and extension propagation.
…ence - infer_rel_schema gains "nested_loop_join" and "exchange" cases; a _join_output_struct helper computes join output columns by join-type NAME (the physical joins' JoinType enums have differing integer values) - plan.nested_loop_join joins over the Cartesian product with a predicate - plan.exchange redistributes rows (round-robin or broadcast), schema unchanged Golden + schema-inference tests, including left_semi -> left-only output.
- nested_loop_join(other, on, how): a physical nested-loop join accepting the same how keys as join(), mapped to NestedLoopJoinRel's own enum - repartition(n) / broadcast(): ExchangeRel distribution verbs Covered by chaining tests (including a left_semi schema check) that exercise the new nested-loop-join and exchange schema inference.
…legacy types Step 1 of the spec bump (bisected: the ANTLR grammar refactor lands at 0.95, so 0.94 is a clean checkpoint that only requires absorbing removed types). Removed spec features absorbed (gone by 0.90): - AggregateRel.Grouping.grouping_expressions -> emit only expression_references - legacy timestamp / time / timestamp_tz types: drop their branches in derivation_expression (they broke the scalar-type isinstance chain for every later type once the ANTLR contexts were removed) and the dead literal cases in type_inference; drop the obsolete coverage tests Engine round-trip tests (pyarrow/DuckDB/DataFusion) now segfault natively because their bundled Substrait consumers lag the pinned spec; per the best-effort policy they are skipped unless SUBSTRAIT_ENGINE_TESTS=1. Non-engine suite: 453 passed, 30 skipped.
…te type-derivation grammar Step 2 of the spec bump. substrait-antlr 0.95 refactored the type-derivation grammar: the single BinaryExpr rule was split into per-operator rules (MulDiv / AddSub / Comparison / Equality / And / Or) plus IfExpr / NotExpr, alongside the existing Ternary. Rewrite derivation_expression._evaluate accordingly: a shared operator table drives the arithmetic/comparison/equality contexts, And/Or/Not evaluate boolean logic, and IfExpr folds into the existing ternary handling. The typeDef/scalarType/parameterizedType structure is unchanged. Now on the latest spec (0.96). Non-engine suite: 453 passed, 30 skipped.
Two relations/expressions newly available on the bumped spec: - plan.top_n builds a TopNRel (fused sort + fetch) with expression-valued count/offset and FETCH_MODE_ROWS_ONLY / WITH_TIES; infer_rel_schema gets a "top_n" pass-through case - execution_context_variable builds a leaf ExecutionContextVariable expression (current_timestamp / current_date / current_timezone), whose oneof variant also carries the value's type; infer_expression_type derives that type
- DataFrame.top_n(n, by, descending=, nulls_last=, offset=, with_ties=): a fused ORDER BY ... LIMIT over TopNRel - current_timestamp(precision=) / current_date() / current_timezone() execution-context-variable expressions, re-exported from substrait.api Covered by structure + chaining tests and a schema-inference check per context variable.
hint(row_count=, record_size=, alias=, output_names=) attaches advisory optimizer statistics / an alias / output-name annotations to the current relation's RelCommon.Hint, without affecting results or schema. Works on any relation (the top rel's common is set generically).
Arithmetic operators (+ - * / % **) fall through to functions_arithmetic_decimal for decimal operands, mirroring the f.* namespace's multi-URN resolution. Comparison and arithmetic against int/Decimal literals coerce the literal to decimal: peer-exact for comparisons (any1,any1 overloads require identical operand types), natural type for arithmetic (the return-type derivation computes the result). Peer-mode comparison coercion raises a clear error when a literal is not exactly representable in the column's decimal type (finer scale or precision overflow) rather than silently rounding or truncating the comparison; _decimal_type rejects precision > 38. Extract the shared multi-URN resolver (_resolve_over_urns), now used by both the operator path and the f.* namespace so their resolution and error text cannot drift. Fix _encode_decimal to scale with exact integer arithmetic instead of Decimal multiplication under the ambient decimal context, so literals with more than the context precision (28) significant digits are no longer silently rounded during encoding.
|
I've found one 'gotcha' ... this API is creating the FetchRel (https://github.com/substrait-io/substrait/blob/858397de179721218bb5ca4a329121ee053e9cab/proto/substrait/algebra.proto#L362) using Sounds like a defect in Substrait-Java I think. |
Yes, |
…ce builder Move the raw ReferenceRel/Plan construction out of DataFrame.cache() into a new substrait.builders.plan.reference builder, so shared-subtree references are built the same way as every other relation. Also document why cache()'s contextvar accumulator is load-bearing: the UnboundPlan = Callable[[ExtensionRegistry], Plan] contract has no slot to thread it, and schema inference resolves ReferenceRels from inside the builder layer (e.g. plan.filter -> infer_plan_schema on a cached input), so the subtree list is reachable there only via a contextvar. Addresses review feedback from @tokoko on substrait-io#204.
…t-io#211) Remove the cache() CTE feature and its ReferenceRel plumbing to keep this PR focused. Per review, shared subplans should be implemented at both the builder and dataframe layers with subtrees carried in-band in the Plan (like extensions) rather than via a contextvar; tracked in substrait-io#211. - frame.py: remove DataFrame.cache(), _CteContext, the _cte_context contextvar, and _materialize's subtree prepend (to_plan/to_substrait resolve directly); drop the now-unreachable hint() guard for common-less relations. - type_inference.py: remove the reference_subtrees contextvar and the ReferenceRel schema-inference case. - builders/plan.py: remove the reference() builder added in fb1c49b. - tests: remove the cache/reference and hint-after-cache tests.
a28b235 to
dee6b8b
Compare
_built was a one-line passthrough to _make_literal left over from an earlier refactor; call _make_literal directly instead. Addresses review feedback from @tokoko on substrait-io#204.
tokoko
left a comment
There was a problem hiding this comment.
looks great, thanks. let's continue with follow-up PRs
Summary
substrait-python can already build any plan, but the day-to-day building experience lags behind sibling libraries — notably substrait-java, whose
coremodule 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.builderslayer that gives Python users the idioms they expect from pandas / Polars / PySpark / Ibis: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/joinand has since grown to cover essentially the whole Substrait query-building surface, and moves the project onto the latest spec (0.96).What's here
Entry point —
import substrait.dataframe as sub. The higher-level API lives under a dedicatedsubstrait.dataframesubpackage rather than at thesubstraitnamespace root:substraitis a PEP 420 namespace package shared withsubstrait-protobuf(so the root can't hold an__init__.py), and grouping the modules under one owned subpackage — a sibling tobuilders/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 inbuilders/protobehaviour 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. Pluscast,alias,is_null/is_not_null/is_nan,between,is_distinct_from,is_in,coalesce,when().then()…otherwise()andswitch(CASE), nested access (struct_field/[]/map_key), window functions via.over(partition_by=, order_by=, rows=/range=), higher-order list functions (list_transform/list_filterwith Python-lambda callbacks), subqueries (scalar_subquery/exists/unique/in_subquery/any_/all_, including correlated viaouter()), dynamic parameters (parameter), and execution-context variables (current_timestamp/current_date/current_timezone).DataFrame verbs
read_named_table,from_records(VALUES),read_parquet/read_csv/read_orc/read_arrow,read_extension_table,extension_leaffilter,select,with_columns,rename,drop,sort(per-column direction / nulls),limit/head/offset,top_njoin(all 13 join types +post_filter),cross_join,nested_loop_join,hash_join,merge_joinunion/intersect/except_group_by().agg(),rollup,cube, explicit grouping sets, per-measureFILTER,DISTINCT, ordered aggregatesunpivot,repartition,broadcast,extension/extension_multihintwrite_named_table(CTAS / insert),create_table/create_view,drop_table/drop_view,update_tableFunction 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 viaand_/or_/not_; function options as kwargs (f.add(x, y, overflow="ERROR")).functions_for(registry)/DataFrame.fsurface custom-extension functions.Types — nullability-aware shortcuts (
sub.i64nullable,sub.i64.non_nullrequired) covering every concrete Substrait type; parametrized builders;user_definedUDTs.User-defined extension relations —
ExtensionLeafDetail/ExtensionSingleDetail/ExtensionMultiDetailABCs (mirroring substrait-java'sExtension.*RelDetail) withto_any/from_any/derive_schema, registered viaExtensionRegistry.register_extension_relationso schema inference follows a custom relation like a built-in one.Spec bump: 0.86 → 0.96
Moves the
substrait-protobuf/substrait-extensions/substrait-antlrpins to the latest spec (needed forTopNReland execution-context variables, and to track upstream). Done as a bisected, two-step migration:AggregateRel.Grouping.grouping_expressionsand the deprecatedtimestamp/time/timestamp_tztypes (superseded byexpression_references/ theprecision_*types), acrosstype_inference,derivation_expression, and tests;BinaryExprrule intoMulDiv/AddSub/Comparison/Equality/And/Or+If/Not) migrated inderivation_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_inferencegained 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;func<…>matching so higher-order functions resolve and bind type parameters;Breaking changes
timestamp/time/timestamp_tztypes andAggregateRel.Grouping.grouping_expressionsare gone (superseded by theprecision_*types /expression_references).substrait.dataframereclaimed — the old (experimental, minimal) Narwhals wrapper that lived atsubstrait.dataframemoves tosubstrait.narwhals, andsubstrait.dataframenow is the native DataFrame/Expr API (the entry point earlier drafted assubstrait.api). Consequence: existingimport substrait.dataframecode silently binds to the native API instead of the Narwhals wrapper (a behaviour change, not anImportError) — migrate those imports tosubstrait.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 unlessSUBSTRAIT_ENGINE_TESTS=1; the non-engine suite is the gate and is green.Follow-ups (not in this PR)
ReferenceRel) at both the builder and dataframe layers (DataFrame.cache()), carrying subtrees in-band in thePlanrather 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).substrait.narwhalsout into a full Narwhals compliant backend (CompliantLazyFrame / Expr / Namespace) on top ofsubstrait.dataframe.frame. The compliant protocol is experimental, andcollect()implies execution (a non-goal), so it would raise or delegate to a consumer.FunctionArgument.enum); type variations on built-in types; Iceberg read specifics (ReadRel.IcebergTable).+,>, …) over every extension the runtime registry defines for the function name — via the shared_resolve_over_urnsloop, replacing the hardcoded standard +_decimalURN 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 fromf.*, barely applies to UDTs). Broadens the overloads operators consider (e.g.date + intervalviafunctions_datetime), so it needs its own tests and an intent decision.🤖 Generated with AI