jorisvandenbossche commented on code in PR #14781:
URL: https://github.com/apache/arrow/pull/14781#discussion_r1042008293
##########
python/pyarrow/tests/test_table.py:
##########
@@ -2204,3 +2204,100 @@ def test_table_cast_invalid():
table = pa.table({'a': [None, 1], 'b': [False, True]})
assert table.cast(new_schema).schema == new_schema
+
+
+def test_table_sort():
+ tab = pa.Table.from_arrays([
+ pa.array([5, 7, 7, 35], type=pa.int64()),
+ pa.array(["foo", "car", "bar", "foobar"])
+ ], names=["a", "b"])
+
+ sorted_tab = tab.sort_by([("a", "descending")])
+ sorted_tab_dict = sorted_tab.to_pydict()
+ assert sorted_tab_dict["a"] == [35, 7, 7, 5]
+ assert sorted_tab_dict["b"] == ["foobar", "car", "bar", "foo"]
+
+ sorted_tab = tab.sort_by([("a", "ascending")])
+ sorted_tab_dict = sorted_tab.to_pydict()
+ assert sorted_tab_dict["a"] == [5, 7, 7, 35]
+ assert sorted_tab_dict["b"] == ["foo", "car", "bar", "foobar"]
+
+
+def test_record_batch_sort():
+ rb = pa.RecordBatch.from_arrays([
+ pa.array([7, 35, 7, 5], type=pa.int64()),
+ pa.array([4, 1, 3, 2], type=pa.int64()),
+ pa.array(["foo", "car", "bar", "foobar"])
+ ], names=["a", "b", "c"])
+
+ sorted_rb = rb.sort_by([("a", "descending"), ("b", "descending")])
+ sorted_rb_dict = sorted_rb.to_pydict()
+ assert sorted_rb_dict["a"] == [35, 7, 7, 5]
+ assert sorted_rb_dict["b"] == [1, 4, 3, 2]
+ assert sorted_rb_dict["c"] == ["car", "foo", "bar", "foobar"]
+
+ sorted_rb = rb.sort_by([("a", "ascending"), ("b", "ascending")])
+ sorted_rb_dict = sorted_rb.to_pydict()
+ assert sorted_rb_dict["a"] == [5, 7, 7, 35]
+ assert sorted_rb_dict["b"] == [2, 3, 4, 1]
+ assert sorted_rb_dict["c"] == ["foobar", "bar", "foo", "car"]
+
+ # test multi-key record batch sorter (> 8 sort keys)
Review Comment:
@Jedi18 thanks for the clarification! Since this stripped down version of
the PR actually doesn't touch the RecordBatch sort implementation, it might not
be needed to add those tests (I assume those two variants of sorters are also
already tested in C++)?
--
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]
