feat: add lead and lag functions (#25)

Author: tokoko

subframe/__init__.py

@@ -125,11 +125,10 @@ def literal(value: Any, type: str = None) -> Value:
     else:
         raise Exception(f"Unknown literal type - {type}")
 
-    print(literal)
-
     return Value(
         expression=stalg.Expression(literal=literal),
         data_type=infer_literal_type(literal),
+        name=str(value),
     )
 
 

subframe/utils.py

@@ -119,12 +119,18 @@ def infer_expression_type(
         return infer_literal_type(expression.literal)
     elif rex_type == "scalar_function":
         return expression.scalar_function.output_type
-    # WindowFunction window_function = 5;
-    # IfThen if_then = 6;
-    # SwitchExpression switch_expression = 7;
-    # SingularOrList singular_or_list = 8;
-    # MultiOrList multi_or_list = 9;
-    # Cast cast = 11;
+    elif rex_type == "window_function":
+        return expression.window_function.output_type
+    elif rex_type == "if_then":
+        return infer_expression_type(expression.if_then.ifs[0].then)
+    elif rex_type == "switch_expression":
+        return infer_expression_type(expression.switch_expression.ifs[0].then)
+    elif rex_type == "cast":
+        return expression.cast.type
+    elif rex_type == "singular_or_list" or rex_type == "multi_or_list":
+        return Type(
+            bool=Type.Boolean(nullability=stt.Type.Nullability.NULLABILITY_NULLABLE)
+        )
     # Subquery subquery = 12;
     # Nested nested = 13;
     else:

subframe/value.py

@@ -4,20 +4,6 @@
 # from .table import Table
 
 
-def substrait_type_from_substrait_str(data_type: str) -> stt.Type:
-    data_type = data_type.replace("?", "")  # TODO
-    if data_type == "i64":
-        return stt.Type(i64=stt.Type.I64(nullability=stt.Type.NULLABILITY_NULLABLE))
-    elif data_type == "fp64":
-        return stt.Type(fp64=stt.Type.FP64(nullability=stt.Type.NULLABILITY_NULLABLE))
-    elif data_type == "boolean":
-        return stt.Type(
-            bool=stt.Type.Boolean(nullability=stt.Type.NULLABILITY_NULLABLE)
-        )
-    else:
-        raise Exception(f"Unknown data type {data_type}")
-
-
 class Value:
     def __init__(
         self,
@@ -205,6 +191,66 @@ def count(self):
             col_name="Count",
         )
 
+    def _apply_window_function(
+        self, additional_arguments: list["Value"], url: str, func: str, col_name: str
+    ):
+        from subframe import registry
+
+        (func_entry, rtn) = registry.lookup_function(
+            url,
+            function_name=func,
+            signature=[
+                self.data_type,
+                *[a.data_type for a in additional_arguments],
+            ],
+        )
+
+        output_type = rtn
+
+        expression = stalg.Expression(
+            window_function=stalg.Expression.WindowFunction(
+                function_reference=func_entry.anchor,
+                arguments=[
+                    stalg.FunctionArgument(value=self.expression),
+                    *[
+                        stalg.FunctionArgument(value=a.expression)
+                        for a in additional_arguments
+                    ],
+                ],
+                options=[],
+                phase=stalg.AggregationPhase.AGGREGATION_PHASE_INITIAL_TO_RESULT,
+                sorts=[],
+                partitions=[],
+            )
+        )
+
+        return Value(
+            expression=expression,
+            data_type=output_type,
+            name=f"{col_name}({self._name}, {' ,'.join([a._name for a in additional_arguments])})",
+            extensions={func_entry.uri: {str(func_entry): func_entry.anchor}},
+        )
+
+    def lead(self, offset):  # TODO default
+        from subframe import literal
+
+        if type(offset) == int:
+            offset = literal(offset, type="i32")
+
+        return self._apply_window_function(
+            [offset], "functions_arithmetic.yaml", "lead", "Lead"
+        )
+
+    def lag(self, offset):  # TODO default
+        from subframe import literal
+
+        if type(offset) == int:
+            offset = literal(offset, type="i32")
+
+        return self._apply_window_function(
+            [offset], "functions_arithmetic.yaml", "lag", "Lag"
+        )
+
 
 class AggregateValue:
     def __init__(

tests/test_execution.py

@@ -569,3 +569,32 @@ def transform(module):
     sf_expr = transform(subframe)
 
     run_parity_test(request.getfixturevalue(consumer), ibis_expr, sf_expr)
+
+
+@pytest.mark.parametrize(
+    "consumer",
+    [
+        pytest.param(
+            "acero_consumer",
+            marks=[pytest.mark.xfail(Exception, reason="Unimplemented")],
+        ),
+        "datafusion_consumer",
+        pytest.param(
+            "duckdb_consumer",
+            marks=[pytest.mark.xfail(Exception, reason="Unimplemented")],
+        ),
+    ],
+)
+def test_lead_lag(consumer, request):
+
+    def transform(module):
+        t1 = _orders(module)
+
+        return t1.select(
+            t1["order_total"].lead(offset=2), t1["order_total"].lag(offset=2)
+        )
+
+    ibis_expr = transform(ibis)
+    sf_expr = transform(subframe)
+
+    run_parity_test(request.getfixturevalue(consumer), ibis_expr, sf_expr)