raulcd commented on code in PR #13327:
URL: https://github.com/apache/arrow/pull/13327#discussion_r891410163
##########
python/pyarrow/_compute.pyx:
##########
@@ -2047,6 +2047,60 @@ class RandomOptions(_RandomOptions):
self._set_options(length, initializer)
+cdef class _RankOptions(FunctionOptions):
+
+ _tiebreaker_map = {
+ "Min": CRankOptionsTiebreaker_MIN,
+ "Max": CRankOptionsTiebreaker_MAX,
+ "First": CRankOptionsTiebreaker_FIRST,
+ "Dense": CRankOptionsTiebreaker_DENSE,
+ }
+
+ def _set_options(self, sort_keys, null_placement, tiebreaker):
+ cdef vector[CSortKey] c_sort_keys
+ for name, order in sort_keys:
+ c_sort_keys.push_back(
+ CSortKey(tobytes(name), unwrap_sort_order(order))
+ )
+ try:
+ self.wrapped.reset(
+ new CRankOptions(c_sort_keys,
unwrap_null_placement(null_placement), self._tiebreaker_map[tiebreaker])
+ )
+ except KeyError:
+ _raise_invalid_function_option(tiebreaker, "tiebreaker")
+
+
+class RankOptions(_RankOptions):
+ """
+ Options for the `rank` function.
+
+ Parameters
+ ----------
+ sort_keys : sequence of (name, order) tuples
Review Comment:
@pitrou do you mean something like this change?
```
$ git diff
diff --git a/python/pyarrow/_compute.pyx b/python/pyarrow/_compute.pyx
index 6df6637..a4830d9 100644
--- a/python/pyarrow/_compute.pyx
+++ b/python/pyarrow/_compute.pyx
@@ -2076,7 +2076,7 @@ class RankOptions(_RankOptions):
Parameters
----------
- sort_keys : sequence of (name, order) tuples
+ sort_keys : sequence of (name, order) tuples, optional
Names of field/column keys to sort the input on,
along with the order each field/column is sorted in.
Accepted values for `order` are "ascending", "descending".
@@ -2097,7 +2097,7 @@ class RankOptions(_RankOptions):
number of distinct values in the input.
"""
- def __init__(self, sort_keys, *, null_placement="at_end",
tiebreaker="First"):
+ def __init__(self, *, sort_keys=(), null_placement="at_end",
tiebreaker="First"):
self._set_options(sort_keys, null_placement, tiebreaker)
diff --git a/python/pyarrow/tests/test_compute.py
b/python/pyarrow/tests/test_compute.py
index 4b7aa3e..aaefefc 100644
--- a/python/pyarrow/tests/test_compute.py
+++ b/python/pyarrow/tests/test_compute.py
@@ -2764,8 +2764,7 @@ def test_rank_options_tiebreaker(tiebreaker,
expected_values):
def test_rank_options():
arr = pa.array([1.2, 0.0, 5.3, None, 5.3, None, 0.0])
- rank_options = pc.RankOptions(sort_keys=[("b", "ascending")])
- result = pc.rank(arr, options=rank_options)
+ result = pc.rank(arr, options=pc.RankOptions())
expected = pa.array([3, 1, 4, 6, 5, 7, 2], type=pa.uint64())
assert result.equals(expected)
```
--
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]
