Update ffi examples

Author: timsaucer

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

@@ -45,16 +45,16 @@ def test_ffi_aggregate_register():
 
     result = ctx.sql("select my_custom_sum(a) from test_table group by b").collect()
 
-    assert result
+    assert len(result) == 2
     assert result[0].num_columns == 1
 
-    # Normalizing table registration in _normalize_table_provider feeds the Rust layer
-    # an actual TableProvider, so collect() emits the grouped rows in a single record batch
-    # instead of two separate batches.
-    aggregates = pa.concat_arrays([batch.column(0) for batch in result])
+    result = [r.column(0) for r in result]
+    expected = [
+        pa.array([3], type=pa.int64()),
+        pa.array([3], type=pa.int64()),
+    ]
 
-    assert len(aggregates) == 2
-    assert aggregates.to_pylist() == [3, 3]
+    assert result == expected
 
 
 def test_ffi_aggregate_call_directly():
@@ -65,13 +65,13 @@ def test_ffi_aggregate_call_directly():
         ctx.table("test_table").aggregate([col("b")], [my_udaf(col("a"))]).collect()
     )
 
-    # Normalizing table registration in _normalize_table_provider feeds the Rust layer
-    # an actual TableProvider, so collect() emits the grouped rows in a single record batch
-    # instead of two separate batches.
-    assert result
+    assert len(result) == 2
     assert result[0].num_columns == 2
 
-    aggregates = pa.concat_arrays([batch.column(1) for batch in result])
+    result = [r.column(1) for r in result]
+    expected = [
+        pa.array([3], type=pa.int64()),
+        pa.array([3], type=pa.int64()),
+    ]
 
-    assert len(aggregates) == 2
-    assert aggregates.to_pylist() == [3, 3]
+    assert result == expected

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

@@ -36,9 +36,9 @@ def test_catalog_provider():
 
     my_catalog_schemas = my_catalog.names()
     assert expected_schema_name in my_catalog_schemas
-    my_database = my_catalog.database(expected_schema_name)
-    assert expected_table_name in my_database.names()
-    my_table = my_database.table(expected_table_name)
+    my_schema = my_catalog.schema(expected_schema_name)
+    assert expected_table_name in my_schema.names()
+    my_table = my_schema.table(expected_table_name)
     assert expected_table_columns == my_table.schema.names
 
     result = ctx.table(

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

@@ -18,14 +18,14 @@
 from __future__ import annotations
 
 import pyarrow as pa
-from datafusion import SessionContext, Table
+from datafusion import SessionContext
 from datafusion_ffi_example import MyTableProvider
 
 
 def test_table_loading():
     ctx = SessionContext()
     table = MyTableProvider(3, 2, 4)
-    ctx.register_table("t", Table.from_capsule(table.__datafusion_table_provider__()))
+    ctx.register_table("t", table)
     result = ctx.table("t").collect()
 
     assert len(result) == 4