jorisvandenbossche commented on a change in pull request #11164:
URL: https://github.com/apache/arrow/pull/11164#discussion_r709881145
##########
File path: python/pyarrow/tests/test_compute.py
##########
@@ -1896,11 +1903,19 @@ 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`"):
+ with pytest.raises(ValueError, match="SelectK requires a nonnegative `k`"):
pc.select_k_unstable(table)
Review comment:
Maybe add here the case of `pc.select_k_unstable(table, k=2)` (without
sort_keys, the original report)
##########
File path: 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=[], memory_pool=None):
Review comment:
```suggestion
def top_k_unstable(values, k, sort_keys=None, memory_pool=None):
```
and then
```
if sort_keys is None:
sort_keys = []
```
as you did in SelectKOptions
##########
File path: cpp/src/arrow/compute/kernels/vector_sort.cc
##########
@@ -2340,6 +2340,9 @@ class SelectKUnstableMetaFunction : public MetaFunction {
return Status::Invalid("SelectK requires a nonnegative `k`, got ",
select_k_options.k);
}
+ if (select_k_options.sort_keys.size() == 0) {
+ return Status::Invalid("SelectK requires `sort_keys` option paremeter,
got {}");
Review comment:
I would maybe leave out the `got {}`, as this will look a bit strange
for users of the bindings (the {} stands for an empty vector in C++, but
Python/R users won't necessarily understand that)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
