chore: support nulls_first in pandas and dask

ibis-project · Jun 21, 2024 · be221d4 · be221d4
1 parent 2f805c8
commit be221d4
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 4 deletions.
diff --git a/ibis/backends/dask/executor.py b/ibis/backends/dask/executor.py
@@ -364,8 +364,14 @@ def visit(cls, op: ops.Sort, parent, keys):
         # 2. sort the dataframe using those columns
         # 3. drop the sort key columns
         ascending = [key.ascending for key in op.keys]
-        na_pos_dict = {True: "first", False: "last"}
-        na_position = na_pos_dict[all(key.nulls_first for key in op.keys)]
+        nulls_first = [key.nulls_first for key in op.keys]
+
+        if all(nulls_first):
+            na_position = "first"
+        elif not any(nulls_first):
+            na_position = "last"
+        else:
+            raise ValueError("dask does not support different columns ordering")
 
         newcols = {gen_name("sort_key"): col for col in keys}
         names = list(newcols.keys())

diff --git a/ibis/backends/pandas/executor.py b/ibis/backends/pandas/executor.py
@@ -593,8 +593,14 @@ def visit(cls, op: ops.Sort, parent, keys):
         # 2. sort the dataframe using those columns
         # 3. drop the sort key columns
         ascending = [key.ascending for key in op.keys]
-        na_pos_dict = {True: "first", False: "last"}
-        na_position = na_pos_dict[all(key.nulls_first for key in op.keys)]
+        nulls_first = [key.nulls_first for key in op.keys]
+
+        if all(nulls_first):
+            na_position = "first"
+        elif not any(nulls_first):
+            na_position = "last"
+        else:
+            raise ValueError("pandas does not support different columns ordering")
 
         newcols = {gen_name("sort_key"): col for col in keys}
         names = list(newcols.keys())