feat(cast): add alias for cast expression (#127)

Added alias for the cast expression builder.
Includes test coverage for cast alias functionality.

Author: giospada

Makefile

@@ -10,6 +10,7 @@ codegen-extensions:
 		--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/gen/json/simple_extensions.py

@@ -1,6 +1,6 @@
 import substrait.gen.proto.algebra_pb2 as stalg
-import substrait.gen.proto.type_pb2 as stt
 import substrait.gen.proto.extended_expression_pb2 as stee
+import substrait.gen.proto.type_pb2 as stt
 from substrait.builders.extended_expression import cast, literal
 from substrait.builders.type import i8, i16
 from substrait.extension_registry import ExtensionRegistry
@@ -37,7 +37,7 @@ def test_cast():
                         failure_behavior=stalg.Expression.Cast.FAILURE_BEHAVIOR_RETURN_NULL,
                     )
                 ),
-                output_names=["cast"],
+                output_names=["cast(Literal(3))"],
             )
         ],
         base_schema=named_struct,
@@ -48,6 +48,7 @@ def test_cast():
 
 def test_cast_with_extension():
     import yaml
+
     import substrait.gen.proto.extensions.extensions_pb2 as ste
     from substrait.builders.extended_expression import scalar_function
 
@@ -134,7 +135,7 @@ def test_cast_with_extension():
                         failure_behavior=stalg.Expression.Cast.FAILURE_BEHAVIOR_RETURN_NULL,
                     )
                 ),
-                output_names=["cast"],
+                output_names=["cast(add(Literal(1),Literal(2)))"],
             )
         ],
         base_schema=named_struct,