22
22
from pyarrow .compute import field
23
23
24
24
from pyarrow .acero import (
25
- TableSourceNodeOptions ,
26
25
Declaration ,
26
+ TableSourceNodeOptions ,
27
27
FilterNodeOptions ,
28
28
ProjectNodeOptions ,
29
29
AggregateNodeOptions ,
30
+ OrderByNodeOptions ,
30
31
HashJoinNodeOptions ,
31
- OrderByNodeOptions
32
32
)
33
33
34
34
try :
@@ -122,8 +122,7 @@ def test_project(table_source):
122
122
# provide name
123
123
decl = Declaration .from_sequence ([
124
124
table_source ,
125
- Declaration ("project" , ProjectNodeOptions (
126
- [pc .multiply (field ("a" ), 2 )], ["a2" ]))
125
+ Declaration ("project" , ProjectNodeOptions ([pc .multiply (field ("a" ), 2 )], ["a2" ]))
127
126
])
128
127
result = decl .to_table ()
129
128
assert result .schema .names == ["a2" ]
@@ -145,8 +144,7 @@ def test_project(table_source):
145
144
def test_aggregate_scalar (table_source ):
146
145
decl = Declaration .from_sequence ([
147
146
table_source ,
148
- Declaration ("aggregate" , AggregateNodeOptions (
149
- [("a" , "sum" , None , "a_sum" )]))
147
+ Declaration ("aggregate" , AggregateNodeOptions ([("a" , "sum" , None , "a_sum" )]))
150
148
])
151
149
result = decl .to_table ()
152
150
assert result .schema .names == ["a_sum" ]
@@ -245,30 +243,26 @@ def test_order_by():
245
243
table_source = Declaration ("table_source" , TableSourceNodeOptions (table ))
246
244
247
245
ord_opts = OrderByNodeOptions ([("b" , "ascending" )])
248
- decl = Declaration .from_sequence (
249
- [table_source , Declaration ("order_by" , ord_opts )])
246
+ decl = Declaration .from_sequence ([table_source , Declaration ("order_by" , ord_opts )])
250
247
result = decl .to_table ()
251
248
expected = pa .table ({"a" : [1 , 4 , 2 , 3 ], "b" : [1 , 2 , 3 , None ]})
252
249
assert result .equals (expected )
253
250
254
251
ord_opts = OrderByNodeOptions ([(field ("b" ), "descending" )])
255
- decl = Declaration .from_sequence (
256
- [table_source , Declaration ("order_by" , ord_opts )])
252
+ decl = Declaration .from_sequence ([table_source , Declaration ("order_by" , ord_opts )])
257
253
result = decl .to_table ()
258
254
expected = pa .table ({"a" : [2 , 4 , 1 , 3 ], "b" : [3 , 2 , 1 , None ]})
259
255
assert result .equals (expected )
260
256
261
257
ord_opts = OrderByNodeOptions ([(1 , "descending" )], null_placement = "at_start" )
262
- decl = Declaration .from_sequence (
263
- [table_source , Declaration ("order_by" , ord_opts )])
258
+ decl = Declaration .from_sequence ([table_source , Declaration ("order_by" , ord_opts )])
264
259
result = decl .to_table ()
265
260
expected = pa .table ({"a" : [3 , 2 , 4 , 1 ], "b" : [None , 3 , 2 , 1 ]})
266
261
assert result .equals (expected )
267
262
268
263
# emtpy ordering
269
264
ord_opts = OrderByNodeOptions ([])
270
- decl = Declaration .from_sequence (
271
- [table_source , Declaration ("order_by" , ord_opts )])
265
+ decl = Declaration .from_sequence ([table_source , Declaration ("order_by" , ord_opts )])
272
266
with pytest .raises (
273
267
ValueError , match = "`ordering` must be an explicit non-empty ordering"
274
268
):
@@ -283,11 +277,9 @@ def test_order_by():
283
277
284
278
def test_hash_join ():
285
279
left = pa .table ({'key' : [1 , 2 , 3 ], 'a' : [4 , 5 , 6 ]})
286
- left_source = Declaration (
287
- "table_source" , options = TableSourceNodeOptions (left ))
280
+ left_source = Declaration ("table_source" , options = TableSourceNodeOptions (left ))
288
281
right = pa .table ({'key' : [2 , 3 , 4 ], 'b' : [4 , 5 , 6 ]})
289
- right_source = Declaration (
290
- "table_source" , options = TableSourceNodeOptions (right ))
282
+ right_source = Declaration ("table_source" , options = TableSourceNodeOptions (right ))
291
283
292
284
# inner join
293
285
join_opts = HashJoinNodeOptions ("inner" , left_keys = "key" , right_keys = "key" )
0 commit comments