ref: refactor extension_registry better exception handling, naming and added function type

Author: giospada

src/substrait/builders/type.py

@@ -223,6 +223,16 @@ def precision_timestamp_tz(precision: int, nullable=True) -> stt.Type:
     )
 
 
+def timestamp(nullable=True) -> stt.Type:
+    return stt.Type(
+        timestamp=stt.Type.Timestamp(
+            nullability=stt.Type.NULLABILITY_NULLABLE
+            if nullable
+            else stt.Type.NULLABILITY_REQUIRED,
+        )
+    )
+
+
 def struct(types: Iterable[stt.Type], nullable=True) -> stt.Type:
     return stt.Type(
         struct=stt.Type.Struct(

src/substrait/dataframe/__init__.py

@@ -11,6 +11,7 @@ def col(name: str) -> Expression:
     """Column selection."""
     return Expression(column(name))
 
+
 # TODO handle str_as_lit argument
 def parse_into_expr(expr, str_as_lit: bool):
     return expr._to_compliant_expr(substrait.dataframe)

src/substrait/dataframe/dataframe.py

@@ -29,7 +29,7 @@ def select(
         expressions = [e.expr for e in exprs] + [
             expr.alias(alias).expr for alias, expr in named_exprs.items()
         ]
-        return DataFrame(select(self.plan, expressions=expressions))        
+        return DataFrame(select(self.plan, expressions=expressions))
 
     # TODO handle version
     def _with_version(self, version):

src/substrait/dataframe/expression.py

@@ -2,7 +2,7 @@
     UnboundExtendedExpression,
     ExtendedExpressionOrUnbound,
     resolve_expression,
-    scalar_function
+    scalar_function,
 )
 import substrait.gen.proto.type_pb2 as stp
 import substrait.gen.proto.extended_expression_pb2 as stee
@@ -30,7 +30,9 @@ def __init__(self, expr: UnboundExtendedExpression):
     def alias(self, alias: str):
         self.expr = _alias(self.expr, alias)
         return self
-    
+
     def abs(self):
-        self.expr = scalar_function("functions_arithmetic.yaml", "abs", expressions=[self.expr]) 
+        self.expr = scalar_function(
+            "functions_arithmetic.yaml", "abs", expressions=[self.expr]
+        )
         return self