Skip to content

Commit

Permalink
fix(snowflake): bring back where filter support in group_concat; …
Browse files Browse the repository at this point in the history
…fix `array_agg` ordering (#9758)
  • Loading branch information
cpcloud authored Aug 3, 2024
1 parent a18cb5d commit 6e7e4de
Showing 1 changed file with 22 additions and 4 deletions.
26 changes: 22 additions & 4 deletions ibis/backends/sql/compilers/snowflake.py
Original file line number Diff line number Diff line change
Expand Up @@ -361,12 +361,27 @@ def visit_TimestampFromUNIX(self, op, *, arg, unit):
timestamp_units_to_scale = {"s": 0, "ms": 3, "us": 6, "ns": 9}
return self.f.to_timestamp(arg, timestamp_units_to_scale[unit.short])

def _array_collect(self, *, arg, where, order_by):
if where is not None:
arg = self.if_(where, arg, NULL)

out = self.f.array_agg(arg)

if order_by:
out = sge.WithinGroup(this=out, expression=sge.Order(expressions=order_by))

return out

def visit_ArrayCollect(self, op, *, arg, where, order_by):
return self._array_collect(arg=arg, where=where, order_by=order_by)

def visit_First(self, op, *, arg, where, order_by):
return self.f.get(self.agg.array_agg(arg, where=where, order_by=order_by), 0)
out = self._array_collect(arg=arg, where=where, order_by=order_by)
return self.f.get(out, 0)

def visit_Last(self, op, *, arg, where, order_by):
expr = self.agg.array_agg(arg, where=where, order_by=order_by)
return self.f.get(expr, self.f.array_size(expr) - 1)
out = self._array_collect(arg=arg, where=where, order_by=order_by)
return self.f.get(out, self.f.array_size(out) - 1)

def visit_GroupConcat(self, op, *, arg, where, sep, order_by):
if where is not None:
Expand All @@ -377,7 +392,10 @@ def visit_GroupConcat(self, op, *, arg, where, sep, order_by):
if order_by:
out = sge.WithinGroup(this=out, expression=sge.Order(expressions=order_by))

return out
if where is None:
return out

return self.if_(self.f.count_if(where) > 0, out, NULL)

def visit_TimestampBucket(self, op, *, arg, interval, offset):
if offset is not None:
Expand Down

0 comments on commit 6e7e4de

Please sign in to comment.