pitrou commented on a change in pull request #11164:
URL: https://github.com/apache/arrow/pull/11164#discussion_r709901888
##########
File path: 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
Review comment:
Can you check the result actually has `k` elements?
--
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]
