chore: extend existing window frames

kszucs · Oct 12, 2023 · 32ea7d0 · 32ea7d0
1 parent 1f95cf0
commit 32ea7d0
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 13 deletions.
diff --git a/ibis/common/patterns.py b/ibis/common/patterns.py
@@ -294,11 +294,9 @@ def __rmatmul__(self, name: str) -> Capture:
         return Capture(name, self)
 
 
-class Is(Slotted, Pattern):
-    """Pattern that matches a value against a reference value.
+class Builder(Hashable):
+    """A builder is a function that takes a context and returns a new object.
 
-    Parameters
-    ----------
     value
         The reference value to match against.
     """

diff --git a/ibis/expr/analysis.py b/ibis/expr/analysis.py
@@ -338,13 +338,26 @@ def propagate_down_window(func: ops.Value, frame: ops.WindowFrame):
     return type(func)(*clean_args)
 
 
-def windowize_function(expr, frame):
-    from ibis.common.patterns import _
+def windowize_function(expr, default_frame):
+    from ibis.common.deferred import _, var
+
+    func = var("func")
+    frame = var("frame")
 
     node = expr.op()
     res = node.replace(
-        (p.Analytic | p.Reduction) >> c.WindowFunction(_, frame),
-        stop=(ops.WindowFunction, ops.Relation),
+        (
+            p.WindowFunction(func, frame)
+            >> c.WindowFunction(
+                func,
+                frame.copy(
+                    order_by=frame.order_by + default_frame.order_by,
+                    group_by=frame.group_by + default_frame.group_by,
+                ),
+            )
+        )
+        | ((p.Analytic | p.Reduction) >> c.WindowFunction(_, default_frame)),
+        stop=ops.Relation,
     )
 
     return res.to_expr()
@@ -426,8 +439,7 @@ def __init__(self, parent, proj_exprs):
 
         default_frame = ops.RowsWindowFrame(table=parent)
         self.clean_exprs = [
-            windowize_function(expr, frame=default_frame)
-            for expr in self.resolved_exprs
+            windowize_function(expr, default_frame) for expr in self.resolved_exprs
         ]
 
     def get_result(self):

diff --git a/ibis/expr/types/groupby.py b/ibis/expr/types/groupby.py
@@ -230,7 +230,6 @@ def mutate(
             A table expression with window functions applied
         """
         exprs = self._selectables(*exprs, **kwexprs)
-        print(exprs)
         return self.table.mutate(exprs)
 
     def select(self, *exprs, **kwexprs) -> ir.Table:
@@ -253,12 +252,12 @@ def _selectables(self, *exprs, **kwexprs):
         table = self.table
         default_frame = self._get_window()
         return [
-            an.windowize_function(e2, frame=default_frame)
+            an.windowize_function(e2, default_frame)
             for expr in exprs
             for e1 in util.promote_list(expr)
             for e2 in util.promote_list(table._ensure_expr(e1))
         ] + [
-            an.windowize_function(e, frame=default_frame).name(k)
+            an.windowize_function(e, default_frame).name(k)
             for k, expr in kwexprs.items()
             for e in util.promote_list(table._ensure_expr(expr))
         ]