fix(dask): fix argmin/argmax implementation for dask

ibis-project · Apr 18, 2024 · 93834f1 · 93834f1
1 parent 03825b6
commit 93834f1
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/ibis/backends/dask/executor.py b/ibis/backends/dask/executor.py
@@ -52,7 +52,8 @@ def limit_df(
 
 
 def argminmax_chunk(df, keycol, valcol, method):
-    return df.iloc[getattr(df[keycol], method)()]
+    idx = getattr(df[keycol], method)()
+    return df[[keycol, valcol]].iloc[idx : idx + 1]
 
 
 def argminmax_aggregate(df, keycol, valcol, method):