Merge branch 'main' into project-select

Author: tokoko

.github/workflows/codegen-check.yml

@@ -20,12 +20,12 @@ jobs:
         uses: devcontainers/ci@v0.3
         with:
           runCmd: |
+            # fetch submodule tags since actions/checkout does not
+            git submodule foreach 'git fetch --unshallow || true'
             # Ensure dependencies are installed
             uv sync --extra test --extra gen_proto
             # Run all code generation steps
-            make antlr
-            ./gen_proto.sh
-            make codegen-extensions
+            make codegen
 
       - name: Check for uncommitted changes
         run: |
@@ -36,9 +36,7 @@ jobs:
             git diff src/substrait/gen/
             echo ""
             echo "To fix this, run:"
-            echo "  make antlr"
-            echo "  ./gen_proto.sh"
-            echo "  make codegen-extensions"
+            echo "  make codegen"
             echo "Then commit the changes."
             exit 1
           fi

CONTRIBUTING.md

@@ -22,6 +22,12 @@ git submodule update --init --recursive
 
 # Code generation
 
+You can run the full code generation using the following command or use the individual commands to selectively regenerate the generated code. This does not update the Substrait Git submodule.
+
+```
+make codegen
+```
+
 ## Protobuf stubs
 
 Run the upgrade script to upgrade the submodule and regenerate the protobuf stubs.
@@ -31,6 +37,12 @@ uv sync --extra gen_proto
 uv run ./update_proto.sh <version>
 ```
 
+Or run the proto codegen without updating the Substrait Git submodule:
+
+```
+make codegen-proto
+```
+
 ## Antlr grammar
 
 Substrait uses antlr grammar to derive output types of extension functions. Make sure java is installed and ANTLR_JAR environment variable is set. Take a look at .devcontainer/Dockerfile for example setup.

Makefile

@@ -1,15 +1,27 @@
+codegen: antlr codegen-proto codegen-extensions codegen-version
+
+
 antlr:
 	cd third_party/substrait/grammar \
 		&& java -jar ${ANTLR_JAR} -o ../../../src/substrait/gen/antlr -Dlanguage=Python3 SubstraitType.g4 \
 		&& rm ../../../src/substrait/gen/antlr/*.tokens \
 		&& rm ../../../src/substrait/gen/antlr/*.interp
 
+codegen-version:
+	echo -n 'substrait_version = "' > src/substrait/gen/version.py \
+		&& cd third_party/substrait && git describe --tags | tr -d 'v\n' >> ../../src/substrait/gen/version.py && cd ../.. \
+		&& echo '"' >> src/substrait/gen/version.py
+
+codegen-proto:
+	./gen_proto.sh
+
 codegen-extensions:
 	uv run --with datamodel-code-generator datamodel-codegen \
 		--input-file-type jsonschema \
 		--input third_party/substrait/text/simple_extensions_schema.yaml \
 		--output src/substrait/gen/json/simple_extensions.py \
 		--output-model-type dataclasses.dataclass \
+		--target-python-version 3.10  \
 		--disable-timestamp
 
 lint:

src/substrait/builders/extended_expression.py

@@ -1,18 +1,19 @@
-from datetime import date
 import itertools
+from datetime import date
+from typing import Any, Callable, Iterable, Union
+
 import substrait.gen.proto.algebra_pb2 as stalg
-import substrait.gen.proto.type_pb2 as stp
 import substrait.gen.proto.extended_expression_pb2 as stee
 import substrait.gen.proto.extensions.extensions_pb2 as ste
+import substrait.gen.proto.type_pb2 as stp
 from substrait.extension_registry import ExtensionRegistry
+from substrait.type_inference import infer_extended_expression_schema
 from substrait.utils import (
-    type_num_names,
-    merge_extension_urns,
-    merge_extension_uris,
     merge_extension_declarations,
+    merge_extension_uris,
+    merge_extension_urns,
+    type_num_names,
 )
-from substrait.type_inference import infer_extended_expression_schema
-from typing import Callable, Any, Union, Iterable
 
 UnboundExtendedExpression = Callable[
     [stp.NamedStruct, ExtensionRegistry], stee.ExtendedExpression
@@ -21,7 +22,7 @@
 
 
 def _alias_or_inferred(
-    alias: Union[Iterable[str], str],
+    alias: Union[Iterable[str], str, None],
     op: str,
     args: Iterable[str],
 ):
@@ -44,7 +45,7 @@ def resolve_expression(
 
 
 def literal(
-    value: Any, type: stp.Type, alias: Union[Iterable[str], str] = None
+    value: Any, type: stp.Type, alias: Union[Iterable[str], str, None] = None
 ) -> UnboundExtendedExpression:
     """Builds a resolver for ExtendedExpression containing a literal expression"""
 
@@ -154,7 +155,7 @@ def resolve(
     return resolve
 
 
-def column(field: Union[str, int], alias: Union[Iterable[str], str] = None):
+def column(field: Union[str, int], alias: Union[Iterable[str], str, None] = None):
     """Builds a resolver for ExtendedExpression containing a FieldReference expression
 
     Accepts either an index or a field name of a desired field.
@@ -208,7 +209,7 @@ def scalar_function(
     urn: str,
     function: str,
     expressions: Iterable[ExtendedExpressionOrUnbound],
-    alias: Union[Iterable[str], str] = None,
+    alias: Union[Iterable[str], str, None] = None,
 ):
     """Builds a resolver for ExtendedExpression containing a ScalarFunction expression"""
 
@@ -306,7 +307,7 @@ def aggregate_function(
     urn: str,
     function: str,
     expressions: Iterable[ExtendedExpressionOrUnbound],
-    alias: Union[Iterable[str], str] = None,
+    alias: Union[Iterable[str], str, None] = None,
 ):
     """Builds a resolver for ExtendedExpression containing a AggregateFunction measure"""
 
@@ -402,7 +403,7 @@ def window_function(
     function: str,
     expressions: Iterable[ExtendedExpressionOrUnbound],
     partitions: Iterable[ExtendedExpressionOrUnbound] = [],
-    alias: Union[Iterable[str], str] = None,
+    alias: Union[Iterable[str], str, None] = None,
 ):
     """Builds a resolver for ExtendedExpression containing a WindowFunction expression"""
 
@@ -512,7 +513,7 @@ def resolve(
 def if_then(
     ifs: Iterable[tuple[ExtendedExpressionOrUnbound, ExtendedExpressionOrUnbound]],
     _else: ExtendedExpressionOrUnbound,
-    alias: Union[Iterable[str], str] = None,
+    alias: Union[Iterable[str], str, None] = None,
 ):
     """Builds a resolver for ExtendedExpression containing an IfThen expression"""
 
@@ -767,7 +768,11 @@ def resolve(
     return resolve
 
 
-def cast(input: ExtendedExpressionOrUnbound, type: stp.Type):
+def cast(
+    input: ExtendedExpressionOrUnbound,
+    type: stp.Type,
+    alias: Union[Iterable[str], str, None] = None,
+):
     """Builds a resolver for ExtendedExpression containing a cast expression"""
 
     def resolve(
@@ -785,7 +790,9 @@ def resolve(
                             failure_behavior=stalg.Expression.Cast.FAILURE_BEHAVIOR_RETURN_NULL,
                         )
                     ),
-                    output_names=["cast"],  # TODO construct name from inputs
+                    output_names=_alias_or_inferred(
+                        alias, "cast", [bound_input.referred_expr[0].output_names[0]]
+                    ),
                 )
             ],
             base_schema=base_schema,

src/substrait/builders/plan.py

@@ -5,30 +5,46 @@
 See `examples/builder_example.py` for usage.
 """
 
-from typing import Iterable, Optional, Union, Callable
+import re
+from typing import Callable, Iterable, Optional, Union
 
 import substrait.gen.proto.algebra_pb2 as stalg
-from substrait.gen.proto.extensions.extensions_pb2 import AdvancedExtension
+import substrait.gen.proto.extended_expression_pb2 as stee
 import substrait.gen.proto.plan_pb2 as stp
 import substrait.gen.proto.type_pb2 as stt
-import substrait.gen.proto.extended_expression_pb2 as stee
-from substrait.extension_registry import ExtensionRegistry
 from substrait.builders.extended_expression import (
     ExtendedExpressionOrUnbound,
     resolve_expression,
 )
+from substrait.extension_registry import ExtensionRegistry
+from substrait.gen.proto.extensions.extensions_pb2 import AdvancedExtension
 from substrait.type_inference import infer_plan_schema
 from substrait.utils import (
     merge_extension_declarations,
-    merge_extension_urns,
     merge_extension_uris,
+    merge_extension_urns,
 )
+from substrait.gen.version import substrait_version
 
 UnboundPlan = Callable[[ExtensionRegistry], stp.Plan]
 
 PlanOrUnbound = Union[stp.Plan, UnboundPlan]
 
 
+def _create_default_version():
+    p = re.compile(r"(\d+)\.(\d+)\.(\d+)")
+    m = p.match(substrait_version)
+    global default_version
+    default_version = stp.Version(
+        major_number=int(m.group(1)),
+        minor_number=int(m.group(2)),
+        patch_number=int(m.group(3)),
+    )
+
+
+_create_default_version()
+
+
 def _merge_extensions(*objs):
     """Merge extension URIs, URNs, and declarations from multiple plan/expression objects.
 
@@ -65,9 +81,10 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
         )
 
         return stp.Plan(
+            version=default_version,
             relations=[
                 stp.PlanRel(root=stalg.RelRoot(input=rel, names=named_struct.names))
-            ]
+            ],
         )
 
     return resolve
@@ -169,6 +186,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
         )
 
         return stp.Plan(
+            version=default_version,
             relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=names))],
             **_merge_extensions(_plan, *bound_expressions),
         )
@@ -199,6 +217,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
         names = ns.names
 
         return stp.Plan(
+            version=default_version,
             relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=names))],
             **_merge_extensions(bound_plan, bound_expression),
         )
@@ -245,6 +264,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
         )
 
         return stp.Plan(
+            version=default_version,
             relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=ns.names))],
             **_merge_extensions(bound_plan, *[e[0] for e in bound_expressions]),
         )
@@ -262,6 +282,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
         )
 
         return stp.Plan(
+            version=default_version,
             relations=[
                 stp.PlanRel(
                     root=stalg.RelRoot(
@@ -300,6 +321,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
         )
 
         return stp.Plan(
+            version=default_version,
             relations=[
                 stp.PlanRel(
                     root=stalg.RelRoot(
@@ -348,6 +370,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
         )
 
         return stp.Plan(
+            version=default_version,
             relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=ns.names))],
             **_merge_extensions(bound_left, bound_right, bound_expression),
         )
@@ -383,6 +406,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
         )
 
         return stp.Plan(
+            version=default_version,
             relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=ns.names))],
             **_merge_extensions(bound_left, bound_right),
         )
@@ -434,10 +458,41 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
         ] + [e.referred_expr[0].output_names[0] for e in bound_measures]
 
         return stp.Plan(
+            version=default_version,
             relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=names))],
             **_merge_extensions(
                 bound_input, *bound_grouping_expressions, *bound_measures
             ),
         )
 
     return resolve
+
+
+def write_named_table(
+    table_names: Union[str, Iterable[str]],
+    input: PlanOrUnbound,
+    create_mode: Union[stalg.WriteRel.CreateMode.ValueType, None] = None,
+) -> UnboundPlan:
+    def resolve(registry: ExtensionRegistry) -> stp.Plan:
+        bound_input = input if isinstance(input, stp.Plan) else input(registry)
+        ns = infer_plan_schema(bound_input)
+        _table_names = [table_names] if isinstance(table_names, str) else table_names
+        _create_mode = create_mode or stalg.WriteRel.CREATE_MODE_ERROR_IF_EXISTS
+
+        write_rel = stalg.Rel(
+            write=stalg.WriteRel(
+                input=bound_input.relations[-1].root.input,
+                table_schema=ns,
+                op=stalg.WriteRel.WRITE_OP_CTAS,
+                create_mode=_create_mode,
+                named_table=stalg.NamedObjectWrite(names=_table_names),
+            )
+        )
+        return stp.Plan(
+            relations=[
+                stp.PlanRel(root=stalg.RelRoot(input=write_rel, names=ns.names))
+            ],
+            **_merge_extensions(bound_input),
+        )
+
+    return resolve

src/substrait/builders/type.py

@@ -1,4 +1,5 @@
 from typing import Iterable
+
 import substrait.gen.proto.type_pb2 as stt