Skip to content

Commit

Permalink
fix(ir): convert analytic functions to window functions in filters
Browse files Browse the repository at this point in the history
  • Loading branch information
cpcloud authored and jcrist committed Aug 13, 2024
1 parent b63e0fd commit 31295dd
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 3 deletions.
8 changes: 5 additions & 3 deletions ibis/expr/rewrites.py
Original file line number Diff line number Diff line change
Expand Up @@ -223,7 +223,7 @@ def lower_stringslice(_, **kwargs):


@replace(p.Analytic)
def project_wrap_analytic(_, rel):
def wrap_analytic(_, **__):
# Wrap analytic functions in a window function
return ops.WindowFunction(_)

Expand All @@ -250,7 +250,7 @@ def rewrite_project_input(value, relation):
# or scalar subqueries depending on whether they are originating from the
# relation
return value.replace(
project_wrap_analytic | project_wrap_reduction,
wrap_analytic | project_wrap_reduction,
filter=p.Value & ~p.WindowFunction,
context={"rel": relation},
)
Expand All @@ -272,7 +272,9 @@ def filter_wrap_reduction(_):


def rewrite_filter_input(value):
return value.replace(filter_wrap_reduction, filter=p.Value & ~p.WindowFunction)
return value.replace(
wrap_analytic | filter_wrap_reduction, filter=p.Value & ~p.WindowFunction
)


@replace(p.Analytic | p.Reduction)
Expand Down
9 changes: 9 additions & 0 deletions ibis/expr/tests/test_newrels.py
Original file line number Diff line number Diff line change
Expand Up @@ -1829,3 +1829,12 @@ def test_expr_in_join_projection():
"lit6": "int64",
}
)


def test_analytic_dereference():
t = ibis.table({"a": "int"})
ix = ibis.row_number()
expr = t.mutate(ix=ix).filter(ix == 5)
assert expr.op().predicates == (
ops.Equals(ops.WindowFunction(ops.RowNumber()), ops.Literal(5, dtype="int8")),
)

0 comments on commit 31295dd

Please sign in to comment.