22
22
from datafusion import SessionContext
23
23
from datafusion_ffi_example import MyCatalogProvider
24
24
25
+
25
26
def test_catalog_provider ():
26
27
ctx = SessionContext ()
27
28
28
29
my_catalog_name = "my_catalog"
29
30
expected_schema_name = "my_schema"
30
31
expected_table_name = "my_table"
31
- expected_table_columns = [' units' , ' price' ]
32
+ expected_table_columns = [" units" , " price" ]
32
33
33
34
catalog_provider = MyCatalogProvider ()
34
35
ctx .register_catalog_provider (my_catalog_name , catalog_provider )
@@ -41,12 +42,9 @@ def test_catalog_provider():
41
42
my_table = my_database .table (expected_table_name )
42
43
assert expected_table_columns == my_table .schema .names
43
44
44
- ctx .register_table (expected_table_name , my_table )
45
- expected_df = ctx .sql (f"SELECT * FROM { expected_table_name } " ).to_pandas ()
46
- assert len (expected_df ) == 5
47
- assert expected_table_columns == expected_df .columns .tolist ()
48
-
49
- result = ctx .table (f"{ my_catalog_name } .{ expected_schema_name } .{ expected_table_name } " ).collect ()
45
+ result = ctx .table (
46
+ f"{ my_catalog_name } .{ expected_schema_name } .{ expected_table_name } "
47
+ ).collect ()
50
48
assert len (result ) == 2
51
49
52
50
col0_result = [r .column (0 ) for r in result ]
@@ -60,4 +58,4 @@ def test_catalog_provider():
60
58
pa .array ([1.5 , 2.5 ], type = pa .float64 ()),
61
59
]
62
60
assert col0_result == expected_col0
63
- assert col1_result == expected_col1
61
+ assert col1_result == expected_col1
0 commit comments