feat: add select builder (#125)

The project builder in plan.py now builds a plan that outputs all inputs
field
followed by all expressions.

The select builders captures the old behavior of project and outputs
only the
expressions.

BREAKING CHANGE: the project builder in plan.py has been renamed to
select

Author: tokoko

examples/builder_example.py

@@ -1,6 +1,6 @@
 from substrait.builders.plan import (
     read_named_table,
-    project,
+    select,
     filter,
     sort,
     fetch,
@@ -34,7 +34,7 @@ def basic_example():
             expressions=[column("id"), literal(100, i64(nullable=False))],
         ),
     )
-    table = project(table, expressions=[column("id")])
+    table = select(table, expressions=[column("id")])
 
     print(table(registry))
     pretty_print_plan(table(registry), use_colors=True)
@@ -177,7 +177,7 @@ def advanced_example():
             expressions=[column("id"), literal(100, i64(nullable=False))],
         ),
     )
-    table = project(table, expressions=[column("id")])
+    table = select(table, expressions=[column("id")])
 
     print("Simple filtered table:")
     pretty_print_plan(table(registry), use_colors=True)
@@ -212,7 +212,7 @@ def advanced_example():
     )
 
     # Project with calculated fields
-    enriched_users = project(
+    enriched_users = select(
         adult_users,
         expressions=[
             column("user_id"),
@@ -322,7 +322,7 @@ def expression_only_example():
         struct=struct(types=[fp64(nullable=False)], nullable=False),
     )
     dummy_table = read_named_table("dummy", dummy_schema)
-    dummy_plan = project(dummy_table, expressions=[complex_expr])
+    dummy_plan = select(dummy_table, expressions=[complex_expr])
     pretty_print_plan(dummy_plan(registry), use_colors=True)
 
     print("\n" + "=" * 50 + "\n")

examples/duckdb_example.py

@@ -9,7 +9,7 @@
 
 
 import duckdb
-from substrait.builders.plan import read_named_table, project, filter
+from substrait.builders.plan import read_named_table, select, filter
 from substrait.builders.extended_expression import column, scalar_function, literal
 from substrait.builders.type import i32
 from substrait.extension_registry import ExtensionRegistry
@@ -46,7 +46,7 @@ def read_duckdb_named_table(name: str, conn):
         expressions=[column("c_nationkey"), literal(3, i32())],
     ),
 )
-table = project(
+table = select(
     table, expressions=[column("c_name"), column("c_address"), column("c_nationkey")]
 )
 sql = "CALL from_substrait(?)"

examples/pyarrow_example.py

@@ -9,7 +9,7 @@
 import pyarrow.compute as pc
 import pyarrow.substrait as pa_substrait
 import substrait
-from substrait.builders.plan import project, read_named_table
+from substrait.builders.plan import select, read_named_table
 
 arrow_schema = pa.schema([pa.field("x", pa.int32()), pa.field("y", pa.int32())])
 
@@ -24,5 +24,5 @@
 pysubstrait_expr = substrait.proto.ExtendedExpression.FromString(bytes(substrait_expr))
 
 table = read_named_table("example", substrait_schema)
-table = project(table, expressions=[pysubstrait_expr])(None)
+table = select(table, expressions=[pysubstrait_expr])(None)
 print(table)

src/substrait/builders/plan.py

@@ -95,6 +95,69 @@ def project(
     expressions: Iterable[ExtendedExpressionOrUnbound],
     extension: Optional[AdvancedExtension] = None,
 ) -> UnboundPlan:
+    """
+    Builds an UnboundPlan with ProjectRel as the root node. Expressions are appended to the parent relation fields to produce an output.
+    Semantically similar to a withColumn transformation.
+
+    :param plan: Parent plan
+    :type plan: PlanOrUnbound
+    :param expressions: Expressions to project
+    :type expressions: Iterable[ExtendedExpressionOrUnbound]
+    :param extension: Optional user-defined extension
+    :type extension: Optional[AdvancedExtension]
+    :return: UnboundPlan with ProjectRel as the root node
+    :rtype: UnboundPlan
+    """
+
+    def resolve(registry: ExtensionRegistry) -> stp.Plan:
+        _plan = plan if isinstance(plan, stp.Plan) else plan(registry)
+        ns = infer_plan_schema(_plan)
+        bound_expressions: Iterable[stee.ExtendedExpression] = [
+            resolve_expression(e, ns, registry) for e in expressions
+        ]
+
+        names = list(_plan.relations[-1].root.names) + [
+            e.output_names[0] for ee in bound_expressions for e in ee.referred_expr
+        ]
+
+        rel = stalg.Rel(
+            project=stalg.ProjectRel(
+                input=_plan.relations[-1].root.input,
+                expressions=[
+                    e.expression for ee in bound_expressions for e in ee.referred_expr
+                ],
+                advanced_extension=extension,
+            )
+        )
+
+        return stp.Plan(
+            version=default_version,
+            relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=names))],
+            **_merge_extensions(_plan, *bound_expressions),
+        )
+
+    return resolve
+
+
+def select(
+    plan: PlanOrUnbound,
+    expressions: Iterable[ExtendedExpressionOrUnbound],
+    extension: Optional[AdvancedExtension] = None,
+) -> UnboundPlan:
+    """
+    Builds an UnboundPlan with ProjectRel as the root node. Expressions make up the fields of an output relation.
+    Semantically similar to a select transformation.
+
+    :param plan: Parent plan
+    :type plan: PlanOrUnbound
+    :param expressions: Expressions to project
+    :type expressions: Iterable[ExtendedExpressionOrUnbound]
+    :param extension: Optional user-defined extension
+    :type extension: Optional[AdvancedExtension]
+    :return: UnboundPlan with ProjectRel as the root node
+    :rtype: UnboundPlan
+    """
+
     def resolve(registry: ExtensionRegistry) -> stp.Plan:
         _plan = plan if isinstance(plan, stp.Plan) else plan(registry)
         ns = infer_plan_schema(_plan)

src/substrait/extension_registry.py

@@ -339,12 +339,14 @@ def satisfies_signature(self, signature: tuple) -> Optional[str]:
         output_type = evaluate(self.impl.return_, parameters)
 
         if self.nullability == se.NullabilityHandling.MIRROR:
-            sig_contains_nullable = any([
-                p.__getattribute__(p.WhichOneof("kind")).nullability
-                == Type.NULLABILITY_NULLABLE
-                for p in signature
-                if isinstance(p, Type)
-            ])
+            sig_contains_nullable = any(
+                [
+                    p.__getattribute__(p.WhichOneof("kind")).nullability
+                    == Type.NULLABILITY_NULLABLE
+                    for p in signature
+                    if isinstance(p, Type)
+                ]
+            )
             output_type.__getattribute__(output_type.WhichOneof("kind")).nullability = (
                 Type.NULLABILITY_NULLABLE
                 if sig_contains_nullable

src/substrait/sql/sql_to_substrait.py

@@ -11,7 +11,7 @@
 )
 from substrait.builders.plan import (
     read_named_table,
-    project,
+    select,
     filter,
     sort,
     fetch,
@@ -309,7 +309,7 @@ def translate(ast: dict, schema_resolver: SchemaResolver, registry: ExtensionReg
         if having_predicate:
             relation = filter(relation, having_predicate)(registry)
 
-        return project(relation, expressions=projection)(registry)
+        return select(relation, expressions=projection)(registry)
     elif op == "Table":
         name = ast["name"][0]["Identifier"]["value"]
         return read_named_table(name, schema_resolver(name))

tests/builders/plan/test_project.py

@@ -2,7 +2,7 @@
 import substrait.gen.proto.plan_pb2 as stp
 import substrait.gen.proto.algebra_pb2 as stalg
 from substrait.builders.type import boolean, i64
-from substrait.builders.plan import read_named_table, project, default_version
+from substrait.builders.plan import read_named_table, select, project, default_version
 from substrait.builders.extended_expression import column
 from substrait.extension_registry import ExtensionRegistry
 
@@ -22,6 +22,41 @@ def test_project():
 
     expected = stp.Plan(
         version=default_version,
+        relations=[
+            stp.PlanRel(
+                root=stalg.RelRoot(
+                    input=stalg.Rel(
+                        project=stalg.ProjectRel(
+                            input=table(None).relations[-1].root.input,
+                            expressions=[
+                                stalg.Expression(
+                                    selection=stalg.Expression.FieldReference(
+                                        direct_reference=stalg.Expression.ReferenceSegment(
+                                            struct_field=stalg.Expression.ReferenceSegment.StructField(
+                                                field=0
+                                            )
+                                        ),
+                                        root_reference=stalg.Expression.FieldReference.RootReference(),
+                                    )
+                                )
+                            ],
+                        )
+                    ),
+                    names=["id", "is_applicable", "id"],
+                )
+            )
+        ],
+    )
+
+    assert actual == expected
+
+
+def test_select():
+    table = read_named_table("table", named_struct)
+
+    actual = select(table, [column("id")])(registry)
+
+    expected = stp.Plan(
         relations=[
             stp.PlanRel(
                 root=stalg.RelRoot(
@@ -49,6 +84,7 @@ def test_project():
                 )
             )
         ],
+        version=default_version,
     )
 
     assert actual == expected

tests/test_uri_urn_migration.py

@@ -24,7 +24,7 @@
 from substrait.builders.plan import (
     read_named_table,
     aggregate,
-    project,
+    select,
     filter,
     default_version,
 )
@@ -149,7 +149,7 @@ def test_project_outputs_both_uri_and_urn():
         alias=["add"],
     )
 
-    actual = project(table, [add_expr])(registry)
+    actual = select(table, [add_expr])(registry)
 
     ns = infer_plan_schema(table(None))