Skip to content

Commit

Permalink
apacheGH-41464: [Python] Fix StructArray.sort() for by=None (apache#4…
Browse files Browse the repository at this point in the history
…1495)

### Rationale for this change
Closes issue apache#41464. Fix `StructArray.sort` method's `by` param to work in the case of `by=None` which was documented to mean sort by all fields (the default), but would raise an exception.

### What changes are included in this PR?
* Add a unit test with by=None in `test_struct_array_sort` that fails on main
* Fix the sort method

### Are these changes tested?
yes

### Are there any user-facing changes?
yes
* GitHub Issue: apache#41464

Authored-by: a-reich <[email protected]>
Signed-off-by: Joris Van den Bossche <[email protected]>
  • Loading branch information
a-reich authored May 14, 2024
1 parent fd84ec0 commit d7c2260
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 4 deletions.
7 changes: 3 additions & 4 deletions python/pyarrow/array.pxi
Original file line number Diff line number Diff line change
Expand Up @@ -3920,12 +3920,11 @@ cdef class StructArray(Array):
result : StructArray
"""
if by is not None:
tosort = self._flattened_field(by)
tosort, sort_keys = self._flattened_field(by), [("", order)]
else:
tosort = self
tosort, sort_keys = self, [(field.name, order) for field in self.type]
indices = _pc().sort_indices(
tosort,
options=_pc().SortOptions(sort_keys=[("", order)], **kwargs)
tosort, options=_pc().SortOptions(sort_keys=sort_keys, **kwargs)
)
return self.take(indices)

Expand Down
8 changes: 8 additions & 0 deletions python/pyarrow/tests/test_array.py
Original file line number Diff line number Diff line change
Expand Up @@ -3536,6 +3536,14 @@ def test_struct_array_sort():
{"a": 5, "b": "foo"},
]

sorted_arr = arr.sort()
assert sorted_arr.to_pylist() == [
{"a": 5, "b": "foo"},
{"a": 7, "b": "bar"},
{"a": 7, "b": "car"},
{"a": 35, "b": "foobar"},
]

arr_with_nulls = pa.StructArray.from_arrays([
pa.array([5, 7, 7, 35], type=pa.int64()),
pa.array(["foo", "car", "bar", "foobar"])
Expand Down

0 comments on commit d7c2260

Please sign in to comment.