Refactor aggregate tests to simplify result assertions and improve readability

Author: kosiew

examples/datafusion-ffi-example/python/tests/_test_aggregate_udf.py

@@ -45,16 +45,13 @@ def test_ffi_aggregate_register():
 
     result = ctx.sql("select my_custom_sum(a) from test_table group by b").collect()
 
-    assert len(result) == 2
+    assert result
     assert result[0].num_columns == 1
 
-    result = [r.column(0) for r in result]
-    expected = [
-        pa.array([3], type=pa.int64()),
-        pa.array([3], type=pa.int64()),
-    ]
+    aggregates = pa.concat_arrays([batch.column(0) for batch in result])
 
-    assert result == expected
+    assert len(aggregates) == 2
+    assert aggregates.to_pylist() == [3, 3]
 
 
 def test_ffi_aggregate_call_directly():
@@ -65,13 +62,10 @@ def test_ffi_aggregate_call_directly():
         ctx.table("test_table").aggregate([col("b")], [my_udaf(col("a"))]).collect()
     )
 
-    assert len(result) == 2
+    assert result
     assert result[0].num_columns == 2
 
-    result = [r.column(1) for r in result]
-    expected = [
-        pa.array([3], type=pa.int64()),
-        pa.array([3], type=pa.int64()),
-    ]
+    aggregates = pa.concat_arrays([batch.column(1) for batch in result])
 
-    assert result == expected
+    assert len(aggregates) == 2
+    assert aggregates.to_pylist() == [3, 3]