22import pyarrow as pa
33import pytest
44
5+
56def test_register_filtered_dataframe ():
67 ctx = SessionContext ()
78
8- data = {
9- "a" : [1 , 2 , 3 , 4 , 5 ],
10- "b" : [10 , 20 , 30 , 40 , 50 ]
11- }
9+ data = {"a" : [1 , 2 , 3 , 4 , 5 ], "b" : [10 , 20 , 30 , 40 , 50 ]}
1210
1311 df = ctx .from_pydict (data , "my_table" )
1412
1513 df_filtered = df .filter (col ("a" ) > literal (2 ))
1614 view = df_filtered .into_view ()
17-
15+
1816 assert view .kind == "view"
19-
17+
2018 ctx .register_table ("view1" , view )
2119
2220 df_view = ctx .sql ("SELECT * FROM view1" )
@@ -25,19 +23,14 @@ def test_register_filtered_dataframe():
2523
2624 result_dicts = [batch .to_pydict () for batch in filtered_results ]
2725
28- expected_results = [
29- {"a" : [3 , 4 , 5 ], "b" : [30 , 40 , 50 ]}
30- ]
26+ expected_results = [{"a" : [3 , 4 , 5 ], "b" : [30 , 40 , 50 ]}]
3127
3228 assert result_dicts == expected_results
33-
29+
3430 df_results = df .collect ()
3531
3632 df_result_dicts = [batch .to_pydict () for batch in df_results ]
3733
38- expected_df_results = [
39- {"a" : [1 , 2 , 3 , 4 , 5 ], "b" : [10 , 20 , 30 , 40 , 50 ]}
40- ]
34+ expected_df_results = [{"a" : [1 , 2 , 3 , 4 , 5 ], "b" : [10 , 20 , 30 , 40 , 50 ]}]
4135
4236 assert df_result_dicts == expected_df_results
43-
0 commit comments