ARROW-14003: [C++][Python] Not providing a sort_key in the "select_k_unstable" kernel crashes

cpp/src/arrow/compute/kernels/vector_sort.cc

@@ -2337,9 +2337,12 @@ class SelectKUnstableMetaFunction : public MetaFunction {
                             const FunctionOptions* options, ExecContext* ctx) const {
     const SelectKOptions& select_k_options = static_cast<const SelectKOptions&>(*options);
     if (select_k_options.k < 0) {
-      return Status::Invalid("SelectK requires a nonnegative `k`, got ",
+      return Status::Invalid("select_k_unstable requires a nonnegative `k`, got ",
                              select_k_options.k);
     }
+    if (select_k_options.sort_keys.size() == 0) {
+      return Status::Invalid("select_k_unstable requires a non-empty `sort_keys`");
+    }
     switch (args[0].kind()) {
       case Datum::ARRAY: {
         return SelectKth(*args[0].make_array(), select_k_options, ctx);

python/pyarrow/compute.py

@@ -643,3 +643,85 @@ def fill_null(values, fill_value):
         fill_value = pa.scalar(fill_value.as_py(), type=values.type)
 
     return call_function("coalesce", [values, fill_value])
+
+
+def top_k_unstable(values, k, sort_keys=None, memory_pool=None):
+    """
+    Select the indices of the top-k ordered elements from array- or table-like
+    data.
+
+    This is a specialization for :func:`select_k_unstable`. Output is not
+    guaranteed to be stable.
+
+    Parameters
+    ----------
+    values : Array, ChunkedArray, RecordBatch, or Table
+    k :  The number of `k` elements to keep.
+    sort_keys : Column key names to order by when input is table-like data.
+
+    Returns
+    -------
+    result : Array of indices
+
+    Examples
+    --------
+    >>> import pyarrow as pa
+    >>> import pyarrow.compute as pc
+    >>> arr = pa.array(["a", "b", "c", None, "e", "f"])
+    >>> pc.top_k_unstable(arr, k=3)
+    <pyarrow.lib.UInt64Array object at 0x7fdcb19d7f30>
+    [
+      5,
+      4,
+      2
+    ]
+    """
+    if sort_keys is None:
+        sort_keys = []
+    if isinstance(values, (pa.Array, pa.ChunkedArray)):
+        sort_keys.append(("dummy", "descending"))
+    else:
+        sort_keys = map(lambda key_name: (key_name, "descending"), sort_keys)
+    options = SelectKOptions(k=k, sort_keys=sort_keys)
+    return call_function("select_k_unstable", [values], options, memory_pool)
+
+
+def bottom_k_unstable(values, k, sort_keys=None, memory_pool=None):
+    """
+    Select the indices of the bottom-k ordered elements from
+    array- or table-like data.
+
+    This is a specialization for :func:`select_k_unstable`. Output is not
+    guaranteed to be stable.
+
+    Parameters
+    ----------
+    values : Array, ChunkedArray, RecordBatch, or Table
+    k :  The number of `k` elements to keep.
+    sort_keys : Column key names to order by when input is table-like data.
+
+    Returns
+    -------
+    result : Array of indices
+
+    Examples
+    --------
+    >>> import pyarrow as pa
+    >>> import pyarrow.compute as pc
+    >>> arr = pa.array(["a", "b", "c", None, "e", "f"])
+    >>> pc.bottom_k_unstable(arr, k=3)
+    <pyarrow.lib.UInt64Array object at 0x7fdcb19d7fa0>
+    [
+      0,
+      1,
+      2
+    ]
+    """
+    if sort_keys is None:
+        sort_keys = []
+    if isinstance(values, (pa.Array, pa.ChunkedArray)):
+        sort_keys.append(("dummy", "ascending"))
+    else:
+        sort_keys = map(lambda key_name: (key_name, "ascending"), sort_keys)
+    options = SelectKOptions(k=k, sort_keys=sort_keys)
+    return call_function("select_k_unstable", [values], options, memory_pool)

python/pyarrow/tests/test_compute.py

@@ -1859,35 +1859,42 @@ def validate_select_k(select_k_indices, arr, order, stable_sort=False):
             assert actual == expected
 
     arr = pa.array([1, 2, None, 0])
-    for order in ["descending", "ascending"]:
-        for k in [0, 2, 4]:
+    for k in [0, 2, 4]:
+        for order in ["descending", "ascending"]:
             result = pc.select_k_unstable(
                 arr, k=k, sort_keys=[("dummy", order)])
             validate_select_k(result, arr, order)
 
-    result = pc.select_k_unstable(arr, options=pc.SelectKOptions(
-        k=2, sort_keys=[("dummy", "descending")]))
+        result = pc.top_k_unstable(arr, k=k)
+        validate_select_k(result, arr, "descending")
+
+        result = pc.bottom_k_unstable(arr, k=k)
+        validate_select_k(result, arr, "ascending")
+
+    result = pc.select_k_unstable(
+        arr, options=pc.SelectKOptions(
+            k=2, sort_keys=[("dummy", "descending")])
+    )
     validate_select_k(result, arr, "descending")
 
-    result = pc.select_k_unstable(arr, options=pc.SelectKOptions(
-        k=2, sort_keys=[("dummy", "ascending")]))
+    result = pc.select_k_unstable(
+        arr, options=pc.SelectKOptions(k=2, sort_keys=[("dummy", "ascending")])
+    )
     validate_select_k(result, arr, "ascending")
 
 
 def test_select_k_table():
-    table = pa.table({"a": [1, 2, 0], "b": [1, 0, 1]})
-
-    def validate_select_k(select_k_indices, table, sort_keys,
-                          stable_sort=False):
-        sorted_indices = pc.sort_indices(table, sort_keys=sort_keys)
+    def validate_select_k(select_k_indices, tbl, sort_keys, stable_sort=False):
+        sorted_indices = pc.sort_indices(tbl, sort_keys=sort_keys)
         head_k_indices = sorted_indices.slice(0, len(select_k_indices))
         if stable_sort:
             assert select_k_indices == head_k_indices
         else:
-            expected = pc.take(table, head_k_indices)
-            actual = pc.take(table, select_k_indices)
+            expected = pc.take(tbl, head_k_indices)
+            actual = pc.take(tbl, select_k_indices)
             assert actual == expected
 
+    table = pa.table({"a": [1, 2, 0], "b": [1, 0, 1]})
     for k in [0, 2, 4]:
         result = pc.select_k_unstable(
             table, k=k, sort_keys=[("a", "ascending")])
@@ -1896,13 +1903,26 @@ def validate_select_k(select_k_indices, table, sort_keys,
         result = pc.select_k_unstable(
             table, k=k, sort_keys=[("a", "ascending"), ("b", "ascending")]
         )
-        validate_select_k(result, table, sort_keys=[
-                          ("a", "ascending"), ("b", "ascending")])
+        validate_select_k(
+            result, table, sort_keys=[("a", "ascending"), ("b", "ascending")]
+        )
+
+        result = pc.top_k_unstable(table, k=k, sort_keys=["a"])
+        validate_select_k(result, table, sort_keys=[("a", "descending")])
+
+        result = pc.bottom_k_unstable(table, k=k, sort_keys=["a", "b"])
+        validate_select_k(
+            result, table, sort_keys=[("a", "ascending"), ("b", "ascending")]
+        )
 
     with pytest.raises(ValueError,
-                       match="SelectK requires a nonnegative `k`"):
+                       match="select_k_unstable requires a nonnegative `k`"):
         pc.select_k_unstable(table)
 
+    with pytest.raises(ValueError, match="select_k_unstable requires "
+                       "a non-empty `sort_keys`"):
+        pc.select_k_unstable(table, k=2)
+
     with pytest.raises(ValueError, match="not a valid order"):
         pc.select_k_unstable(table, k=k, sort_keys=[("a", "nonscending")])