@@ -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-
20452028def 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