feat(pandas): add approx_median

ibis/backends/pandas/execution/generic.py

@@ -916,6 +916,23 @@ def mode(x):
     return aggcontext.agg(data, mode)
 
 
+@execute_node.register(ops.ApproxMedian, pd.Series, (pd.Series, type(None)))
+def execute_approx_median_series(_, data, mask, aggcontext=None, **kwargs):
+    return aggcontext.agg(
+        data[mask] if mask is not None else data, lambda x: x.median()
+    )
+
+
+@execute_node.register(ops.ApproxMedian, SeriesGroupBy, (SeriesGroupBy, type(None)))
+def execute_approx_median_series_groupby(_, data, mask, aggcontext=None, **kwargs):
+    median = pd.Series.median
+
+    if mask is not None:
+        median = functools.partial(_filtered_reduction, mask.obj, median)
+
+    return aggcontext.agg(data, median)
+
+
 @execute_node.register((ops.Not, ops.Negate), (bool, np.bool_))
 def execute_not_bool(_, data, **kwargs):
     return not data

ibis/backends/pandas/tests/execution/test_functions.py

@@ -270,3 +270,20 @@ def func(x):
     result = result.tolist()
     assert result == [value]
     assert type(result[0]) is type(value)
+
+
+@pytest.mark.parametrize(
+    ('ibis_func', 'pandas_func'),
+    [
+        (
+            lambda x: x.approx_median(),
+            lambda x: x.median(),
+        )
+    ],
+)
+@pytest.mark.parametrize('column', ['float64_with_zeros', 'int64_with_zeros'])
+def test_approx_median(t, df, ibis_func, pandas_func, column):
+    expr = ibis_func(t[column])
+    result = expr.execute()
+    expected = pandas_func(df[column])
+    assert expected == result

ibis/backends/pandas/tests/execution/test_window.py

@@ -209,6 +209,15 @@ def test_batting_quantile(players, players_df):
     tm.assert_frame_equal(result, expected)
 
 
+def test_batting_approx_median(players, players_df):
+    expr = players.mutate(hits_median=lambda t: t.H.approx_median())
+    hits_median = players_df.groupby('playerID').H.transform('median')
+    expected = players_df.assign(hits_median=hits_median)
+    cols = expected.columns.tolist()
+    result = expr.execute()[cols].sort_values(cols).reset_index(drop=True)
+    tm.assert_frame_equal(result, expected)
+
+
 @pytest.mark.parametrize('op', ['sum', 'mean', 'min', 'max'])
 def test_batting_specific_cumulative(batting, batting_df, op, sort_kind):
     ibis_method = methodcaller(f'cum{op}')

ibis/backends/tests/test_aggregation.py

@@ -1034,7 +1034,7 @@ def test_corr_cov(
     raises=com.OperationNotDefinedError,
 )
 @pytest.mark.broken(
-    ["dask", "pandas"],
+    ["dask"],
     raises=AttributeError,
     reason="'Series' object has no attribute 'approx_median'",
 )