Skip to content

Commit 8a34732

Browse files
apacheGH-14778: [Python] Add (Chunked)Array sort() method (apache#14781)
* Closes: apache#14778 Lead-authored-by: Alessandro Molina <amol@turbogears.org> Co-authored-by: Joris Van den Bossche <jorisvandenbossche@gmail.com> Signed-off-by: Joris Van den Bossche <jorisvandenbossche@gmail.com>
1 parent 5c1044f commit 8a34732

5 files changed

Lines changed: 285 additions & 19 deletions

File tree

‎python/pyarrow/array.pxi‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1399,6 +1399,29 @@ cdef class Array(_PandasConvertible):
13991399
"""
14001400
return _pc().index(self, value, start, end, memory_pool=memory_pool)
14011401

1402+
def sort(self, order="ascending", **kwargs):
1403+
"""
1404+
Sort the Array
1405+
1406+
Parameters
1407+
----------
1408+
order : str, default "ascending"
1409+
Which order to sort values in.
1410+
Accepted values are "ascending", "descending".
1411+
**kwargs : dict, optional
1412+
Additional sorting options.
1413+
As allowed by :class:`SortOptions`
1414+
1415+
Returns
1416+
-------
1417+
result : Array
1418+
"""
1419+
indices = _pc().sort_indices(
1420+
self,
1421+
options=_pc().SortOptions(sort_keys=[("", order)], **kwargs)
1422+
)
1423+
return self.take(indices)
1424+
14021425
def _to_pandas(self, options, types_mapper=None, **kwargs):
14031426
return _array_like_to_pandas(self, options, types_mapper=types_mapper)
14041427

@@ -2641,6 +2664,39 @@ cdef class StructArray(Array):
26412664

26422665
return pyarrow_wrap_array(child)
26432666

2667+
def _flattened_field(self, index, MemoryPool memory_pool=None):
2668+
"""
2669+
Retrieves the child array belonging to field,
2670+
accounting for the parent array null bitmap.
2671+
2672+
Parameters
2673+
----------
2674+
index : Union[int, str]
2675+
Index / position or name of the field.
2676+
memory_pool : MemoryPool, default None
2677+
For memory allocations, if required, otherwise use default pool.
2678+
2679+
Returns
2680+
-------
2681+
result : Array
2682+
"""
2683+
cdef:
2684+
CStructArray* arr = <CStructArray*> self.ap
2685+
shared_ptr[CArray] child
2686+
CMemoryPool* pool = maybe_unbox_memory_pool(memory_pool)
2687+
2688+
if isinstance(index, (bytes, str)):
2689+
int_index = self.type.get_field_index(index)
2690+
if int_index < 0:
2691+
raise KeyError(index)
2692+
elif isinstance(index, int):
2693+
int_index = _normalize_index(index, self.ap.num_fields())
2694+
else:
2695+
raise TypeError('Expected integer or string index')
2696+
2697+
child = GetResultValue(arr.GetFlattenedField(int_index, pool))
2698+
return pyarrow_wrap_array(child)
2699+
26442700
def flatten(self, MemoryPool memory_pool=None):
26452701
"""
26462702
Return one individual array for each field in the struct.
@@ -2743,6 +2799,36 @@ cdef class StructArray(Array):
27432799
result.validate()
27442800
return result
27452801

2802+
def sort(self, order="ascending", by=None, **kwargs):
2803+
"""
2804+
Sort the StructArray
2805+
2806+
Parameters
2807+
----------
2808+
order : str, default "ascending"
2809+
Which order to sort values in.
2810+
Accepted values are "ascending", "descending".
2811+
by : str or None, default None
2812+
If to sort the array by one of its fields
2813+
or by the whole array.
2814+
**kwargs : dict, optional
2815+
Additional sorting options.
2816+
As allowed by :class:`SortOptions`
2817+
2818+
Returns
2819+
-------
2820+
result : StructArray
2821+
"""
2822+
if by is not None:
2823+
tosort = self._flattened_field(by)
2824+
else:
2825+
tosort = self
2826+
indices = _pc().sort_indices(
2827+
tosort,
2828+
options=_pc().SortOptions(sort_keys=[("", order)], **kwargs)
2829+
)
2830+
return self.take(indices)
2831+
27462832

27472833
cdef class ExtensionArray(Array):
27482834
"""

‎python/pyarrow/includes/libarrow.pxd‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -766,6 +766,7 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
766766

767767
shared_ptr[CArray] field(int pos)
768768
shared_ptr[CArray] GetFieldByName(const c_string& name) const
769+
CResult[shared_ptr[CArray]] GetFlattenedField(int index, CMemoryPool* pool) const
769770

770771
CResult[vector[shared_ptr[CArray]]] Flatten(CMemoryPool* pool)
771772

‎python/pyarrow/table.pxi‎

Lines changed: 57 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1039,6 +1039,29 @@ cdef class ChunkedArray(_PandasConvertible):
10391039
"""
10401040
return _pc().drop_null(self)
10411041

1042+
def sort(self, order="ascending", **kwargs):
1043+
"""
1044+
Sort the ChunkedArray
1045+
1046+
Parameters
1047+
----------
1048+
order : str, default "ascending"
1049+
Which order to sort values in.
1050+
Accepted values are "ascending", "descending".
1051+
**kwargs : dict, optional
1052+
Additional sorting options.
1053+
As allowed by :class:`SortOptions`
1054+
1055+
Returns
1056+
-------
1057+
result : ChunkedArray
1058+
"""
1059+
indices = _pc().sort_indices(
1060+
self,
1061+
options=_pc().SortOptions(sort_keys=[("", order)], **kwargs)
1062+
)
1063+
return self.take(indices)
1064+
10421065
def unify_dictionaries(self, MemoryPool memory_pool=None):
10431066
"""
10441067
Unify dictionaries across all chunks.
@@ -2231,6 +2254,35 @@ cdef class RecordBatch(_PandasConvertible):
22312254
"""
22322255
return _pc().drop_null(self)
22332256

2257+
def sort_by(self, sorting, **kwargs):
2258+
"""
2259+
Sort the RecordBatch by one or multiple columns.
2260+
2261+
Parameters
2262+
----------
2263+
sorting : str or list[tuple(name, order)]
2264+
Name of the column to use to sort (ascending), or
2265+
a list of multiple sorting conditions where
2266+
each entry is a tuple with column name
2267+
and sorting order ("ascending" or "descending")
2268+
**kwargs : dict, optional
2269+
Additional sorting options.
2270+
As allowed by :class:`SortOptions`
2271+
2272+
Returns
2273+
-------
2274+
RecordBatch
2275+
A new record batch sorted according to the sort keys.
2276+
"""
2277+
if isinstance(sorting, str):
2278+
sorting = [(sorting, "ascending")]
2279+
2280+
indices = _pc().sort_indices(
2281+
self,
2282+
options=_pc().SortOptions(sort_keys=sorting, **kwargs)
2283+
)
2284+
return self.take(indices)
2285+
22342286
def to_pydict(self):
22352287
"""
22362288
Convert the RecordBatch to a dict or OrderedDict.
@@ -4664,7 +4716,7 @@ cdef class Table(_PandasConvertible):
46644716
"""
46654717
return TableGroupBy(self, keys)
46664718

4667-
def sort_by(self, sorting):
4719+
def sort_by(self, sorting, **kwargs):
46684720
"""
46694721
Sort the table by one or multiple columns.
46704722
@@ -4675,6 +4727,9 @@ cdef class Table(_PandasConvertible):
46754727
a list of multiple sorting conditions where
46764728
each entry is a tuple with column name
46774729
and sorting order ("ascending" or "descending")
4730+
**kwargs : dict, optional
4731+
Additional sorting options.
4732+
As allowed by :class:`SortOptions`
46784733
46794734
Returns
46804735
-------
@@ -4705,7 +4760,7 @@ cdef class Table(_PandasConvertible):
47054760

47064761
indices = _pc().sort_indices(
47074762
self,
4708-
sort_keys=sorting
4763+
options=_pc().SortOptions(sort_keys=sorting, **kwargs)
47094764
)
47104765
return self.take(indices)
47114766

‎python/pyarrow/tests/test_array.py‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2916,6 +2916,41 @@ def test_struct_array_field():
29162916
a.field(invalid_name)
29172917

29182918

2919+
def test_struct_array_flattened_field():
2920+
ty = pa.struct([pa.field('x', pa.int16()),
2921+
pa.field('y', pa.float32())])
2922+
a = pa.array([(1, 2.5), (3, 4.5), (5, 6.5)], type=ty,
2923+
mask=pa.array([False, True, False]))
2924+
2925+
x0 = a._flattened_field(0)
2926+
y0 = a._flattened_field(1)
2927+
x1 = a._flattened_field(-2)
2928+
y1 = a._flattened_field(-1)
2929+
x2 = a._flattened_field('x')
2930+
y2 = a._flattened_field('y')
2931+
2932+
assert isinstance(x0, pa.lib.Int16Array)
2933+
assert isinstance(y1, pa.lib.FloatArray)
2934+
assert x0.equals(pa.array([1, None, 5], type=pa.int16()))
2935+
assert y0.equals(pa.array([2.5, None, 6.5], type=pa.float32()))
2936+
assert x0.equals(x1)
2937+
assert x0.equals(x2)
2938+
assert y0.equals(y1)
2939+
assert y0.equals(y2)
2940+
2941+
for invalid_index in [None, pa.int16()]:
2942+
with pytest.raises(TypeError):
2943+
a._flattened_field(invalid_index)
2944+
2945+
for invalid_index in [3, -3]:
2946+
with pytest.raises(IndexError):
2947+
a._flattened_field(invalid_index)
2948+
2949+
for invalid_name in ['z', '']:
2950+
with pytest.raises(KeyError):
2951+
a._flattened_field(invalid_name)
2952+
2953+
29192954
def test_empty_cast():
29202955
types = [
29212956
pa.null(),
@@ -3287,3 +3322,57 @@ def test_to_pandas_timezone():
32873322
arr = pa.chunked_array([arr])
32883323
s = arr.to_pandas()
32893324
assert s.dt.tz is not None
3325+
3326+
3327+
def test_array_sort():
3328+
arr = pa.array([5, 7, 35], type=pa.int64())
3329+
sorted_arr = arr.sort("descending")
3330+
assert sorted_arr.to_pylist() == [35, 7, 5]
3331+
3332+
arr = pa.chunked_array([[1, 2, 3], [4, 5, 6]])
3333+
sorted_arr = arr.sort("descending")
3334+
assert sorted_arr.to_pylist() == [6, 5, 4, 3, 2, 1]
3335+
3336+
arr = pa.array([5, 7, 35, None], type=pa.int64())
3337+
sorted_arr = arr.sort("descending", null_placement="at_end")
3338+
assert sorted_arr.to_pylist() == [35, 7, 5, None]
3339+
sorted_arr = arr.sort("descending", null_placement="at_start")
3340+
assert sorted_arr.to_pylist() == [None, 35, 7, 5]
3341+
3342+
3343+
def test_struct_array_sort():
3344+
arr = pa.StructArray.from_arrays([
3345+
pa.array([5, 7, 7, 35], type=pa.int64()),
3346+
pa.array(["foo", "car", "bar", "foobar"])
3347+
], names=["a", "b"])
3348+
3349+
sorted_arr = arr.sort("descending", by="a")
3350+
assert sorted_arr.to_pylist() == [
3351+
{"a": 35, "b": "foobar"},
3352+
{"a": 7, "b": "car"},
3353+
{"a": 7, "b": "bar"},
3354+
{"a": 5, "b": "foo"},
3355+
]
3356+
3357+
arr_with_nulls = pa.StructArray.from_arrays([
3358+
pa.array([5, 7, 7, 35], type=pa.int64()),
3359+
pa.array(["foo", "car", "bar", "foobar"])
3360+
], names=["a", "b"], mask=pa.array([False, False, True, False]))
3361+
3362+
sorted_arr = arr_with_nulls.sort(
3363+
"descending", by="a", null_placement="at_start")
3364+
assert sorted_arr.to_pylist() == [
3365+
None,
3366+
{"a": 35, "b": "foobar"},
3367+
{"a": 7, "b": "car"},
3368+
{"a": 5, "b": "foo"},
3369+
]
3370+
3371+
sorted_arr = arr_with_nulls.sort(
3372+
"descending", by="a", null_placement="at_end")
3373+
assert sorted_arr.to_pylist() == [
3374+
{"a": 35, "b": "foobar"},
3375+
{"a": 7, "b": "car"},
3376+
{"a": 5, "b": "foo"},
3377+
None
3378+
]

‎python/pyarrow/tests/test_table.py‎

Lines changed: 52 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -2025,23 +2025,6 @@ def sorted_by_keys(d):
20252025
}
20262026

20272027

2028-
def test_table_sort_by():
2029-
table = pa.table([
2030-
pa.array([3, 1, 4, 2, 5]),
2031-
pa.array(["b", "a", "b", "a", "c"]),
2032-
], names=["values", "keys"])
2033-
2034-
assert table.sort_by("values").to_pydict() == {
2035-
"keys": ["a", "a", "b", "b", "c"],
2036-
"values": [1, 2, 3, 4, 5]
2037-
}
2038-
2039-
assert table.sort_by([("values", "descending")]).to_pydict() == {
2040-
"keys": ["c", "b", "b", "a", "a"],
2041-
"values": [5, 4, 3, 2, 1]
2042-
}
2043-
2044-
20452028
def test_table_to_recordbatchreader():
20462029
table = pa.Table.from_pydict({'x': [1, 2, 3]})
20472030
reader = table.to_reader()
@@ -2204,3 +2187,55 @@ def test_table_cast_invalid():
22042187

22052188
table = pa.table({'a': [None, 1], 'b': [False, True]})
22062189
assert table.cast(new_schema).schema == new_schema
2190+
2191+
2192+
def test_table_sort_by():
2193+
table = pa.table([
2194+
pa.array([3, 1, 4, 2, 5]),
2195+
pa.array(["b", "a", "b", "a", "c"]),
2196+
], names=["values", "keys"])
2197+
2198+
assert table.sort_by("values").to_pydict() == {
2199+
"keys": ["a", "a", "b", "b", "c"],
2200+
"values": [1, 2, 3, 4, 5]
2201+
}
2202+
2203+
assert table.sort_by([("values", "descending")]).to_pydict() == {
2204+
"keys": ["c", "b", "b", "a", "a"],
2205+
"values": [5, 4, 3, 2, 1]
2206+
}
2207+
2208+
tab = pa.Table.from_arrays([
2209+
pa.array([5, 7, 7, 35], type=pa.int64()),
2210+
pa.array(["foo", "car", "bar", "foobar"])
2211+
], names=["a", "b"])
2212+
2213+
sorted_tab = tab.sort_by([("a", "descending")])
2214+
sorted_tab_dict = sorted_tab.to_pydict()
2215+
assert sorted_tab_dict["a"] == [35, 7, 7, 5]
2216+
assert sorted_tab_dict["b"] == ["foobar", "car", "bar", "foo"]
2217+
2218+
sorted_tab = tab.sort_by([("a", "ascending")])
2219+
sorted_tab_dict = sorted_tab.to_pydict()
2220+
assert sorted_tab_dict["a"] == [5, 7, 7, 35]
2221+
assert sorted_tab_dict["b"] == ["foo", "car", "bar", "foobar"]
2222+
2223+
2224+
def test_record_batch_sort():
2225+
rb = pa.RecordBatch.from_arrays([
2226+
pa.array([7, 35, 7, 5], type=pa.int64()),
2227+
pa.array([4, 1, 3, 2], type=pa.int64()),
2228+
pa.array(["foo", "car", "bar", "foobar"])
2229+
], names=["a", "b", "c"])
2230+
2231+
sorted_rb = rb.sort_by([("a", "descending"), ("b", "descending")])
2232+
sorted_rb_dict = sorted_rb.to_pydict()
2233+
assert sorted_rb_dict["a"] == [35, 7, 7, 5]
2234+
assert sorted_rb_dict["b"] == [1, 4, 3, 2]
2235+
assert sorted_rb_dict["c"] == ["car", "foo", "bar", "foobar"]
2236+
2237+
sorted_rb = rb.sort_by([("a", "ascending"), ("b", "ascending")])
2238+
sorted_rb_dict = sorted_rb.to_pydict()
2239+
assert sorted_rb_dict["a"] == [5, 7, 7, 35]
2240+
assert sorted_rb_dict["b"] == [2, 3, 4, 1]
2241+
assert sorted_rb_dict["c"] == ["foobar", "bar", "foo", "car"]

0 commit comments

Comments
 (0)