feat(pandas): implement Contains with a group by

ibis-project · Aug 2, 2022 · c534848 · c534848
1 parent 45612ce
commit c534848
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/ibis/backends/pandas/execution/generic.py b/ibis/backends/pandas/execution/generic.py
@@ -982,6 +982,17 @@ def execute_node_contains_series_sequence(op, data, elements, **kwargs):
     return data.isin(elements)
 
 
+@execute_node.register(
+    ops.Contains,
+    SeriesGroupBy,
+    (collections.abc.Sequence, collections.abc.Set, pd.Series),
+)
+def execute_node_contains_series_group_by_sequence(
+    op, data, elements, **kwargs
+):
+    return data.obj.isin(elements).groupby(data.grouper.groupings)
+
+
 @execute_node.register(
     ops.NotContains,
     pd.Series,