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

examples/dataframe_example.py

@@ -1,17 +1,73 @@
-import substrait.dataframe as sdf
-from substrait.builders.plan import read_named_table
-from substrait.builders.type import boolean, i64, named_struct, struct
-from substrait.extension_registry import ExtensionRegistry
+"""Example usage of the ergonomic `substrait.dataframe` facade.
 
-registry = ExtensionRegistry(load_default_extensions=True)
+Compare this with `builder_example.py`, which builds the same kinds of plans
+with the lower-level `substrait.builders.*` functions.
+"""
 
-ns = named_struct(
-    names=["id", "is_applicable"],
-    struct=struct(types=[i64(nullable=False), boolean()], nullable=False),
-)
+import substrait.dataframe as sub
+from substrait.utils.display import pretty_print_plan
 
-table = read_named_table("example_table", ns)
 
-frame = sdf.DataFrame(read_named_table("example_table", ns))
-frame = frame.select(sdf.col("id"))
-print(frame.to_substrait(registry))
+def filter_select_example():
+    """read -> filter -> with_columns -> select, with operator expressions."""
+    plan = (
+        sub.read_named_table(
+            "people", {"id": sub.i64, "age": sub.i64, "name": sub.string}
+        )
+        .filter((sub.col("age") > 25) & sub.col("name").is_not_null())
+        .with_columns(next_year=sub.col("age") + 1)
+        .select("id", "name", "next_year")
+        .to_plan()
+    )
+    pretty_print_plan(plan, use_colors=True)
+
+
+def aggregate_example():
+    """group_by().agg() with the named-function namespace `f`.
+
+    Note the explicit nullability: ``region`` is required, ``amount`` nullable.
+    ``amount > 0`` also shows literal coercion -- the int ``0`` is typed to match
+    the fp64 column so the comparison overload resolves.
+    """
+    plan = (
+        sub.read_named_table(
+            "sales", {"region": sub.string.non_null, "amount": sub.fp64}
+        )
+        .filter(sub.col("amount") > 0)
+        .group_by("region")
+        .agg(
+            sub.f.sum(sub.col("amount")).alias("total"),
+            sub.f.count(sub.col("amount")).alias("n"),
+        )
+        .to_plan()
+    )
+    pretty_print_plan(plan, use_colors=True)
+
+
+def join_example():
+    """Join two tables and project across the combined schema."""
+    customers = sub.read_named_table(
+        "customers", {"cust_id": sub.i64, "name": sub.string}
+    )
+    orders = sub.read_named_table(
+        "orders", {"order_id": sub.i64, "cust_ref": sub.i64, "amount": sub.fp64}
+    )
+    plan = (
+        customers.join(
+            orders, on=sub.col("cust_id") == sub.col("cust_ref"), how="inner"
+        )
+        .select("name", "amount")
+        .sort("amount", descending=True)
+        .limit(10)
+        .to_plan()
+    )
+    pretty_print_plan(plan, use_colors=True)
+
+
+if __name__ == "__main__":
+    print("=== filter / with_columns / select ===")
+    filter_select_example()
+    print("\n=== group_by / agg ===")
+    aggregate_example()
+    print("\n=== join / sort / limit ===")
+    join_example()

examples/narwhals_example.py

@@ -1,4 +1,10 @@
-# Install duckdb and pyarrow before running this example
+# Example of the `substrait.narwhals` integration layer: drive Substrait plan
+# construction through Narwhals (`nw.from_native`), so backend-agnostic Narwhals
+# code compiles to a Substrait plan.
+#
+# For building plans directly (without Narwhals), see `dataframe_example.py`,
+# which uses the Substrait-native DataFrame in `substrait.dataframe`.
+#
 # /// script
 # dependencies = [
 #   "narwhals==2.9.0",
@@ -9,7 +15,7 @@
 import narwhals as nw
 from narwhals.typing import FrameT
 
-import substrait.dataframe as sdf
+import substrait.narwhals as sn
 from substrait.builders.plan import read_named_table
 from substrait.builders.type import boolean, i64, named_struct, struct
 from substrait.extension_registry import ExtensionRegistry
@@ -21,15 +27,11 @@
     struct=struct(types=[i64(nullable=False), boolean()], nullable=False),
 )
 
-table = read_named_table("example_table", ns)
-
-
-lazy_frame: FrameT = nw.from_native(
-    sdf.DataFrame(read_named_table("example_table", ns))
-)
+# Wrap the Substrait Narwhals backend and drive it with the Narwhals API.
+lazy_frame: FrameT = nw.from_native(sn.DataFrame(read_named_table("example_table", ns)))
 
 lazy_frame = lazy_frame.select(nw.col("id").abs(), new_id=nw.col("id"))
 
-df: sdf.DataFrame = lazy_frame.to_native()
+df: sn.DataFrame = lazy_frame.to_native()
 
 print(df.to_substrait(registry))

pyproject.toml

@@ -7,20 +7,20 @@ readme = "README.md"
 requires-python = ">=3.10,<3.15"
 dependencies = [
     "protobuf >=5,<7",
-    "substrait-protobuf==0.86.0",
-    "substrait-extensions==0.86.0",
+    "substrait-protobuf==0.96.0",
+    "substrait-extensions==0.96.0",
 ]
 dynamic = ["version"]
 
 [tool.setuptools_scm]
 write_to = "src/substrait/_version.py"
 
 [project.optional-dependencies]
-extensions = ["substrait-antlr==0.86.0", "pyyaml"]
+extensions = ["substrait-antlr==0.96.0", "pyyaml"]
 sql = ["sqloxide", "deepdiff"]
 
 [dependency-groups]
-dev = ["pytest >= 7.0.0", "substrait-antlr==0.86.0", "pyyaml", "sqloxide", "deepdiff", "duckdb<=1.2.2; python_version < '3.14'", "datafusion"]
+dev = ["pytest >= 7.0.0", "substrait-antlr==0.96.0", "pyyaml", "sqloxide", "deepdiff", "duckdb<=1.2.2; python_version < '3.14'", "datafusion"]
 
 [tool.pytest.ini_options]
 pythonpath = "src"