33import pyarrow as pa
44import pyarrow .compute as pc
55import pytest
6-
7- from vortex .arrow .expression import arrow_to_vortex , _schema_for_substrait
6+ from vortex .arrow .expression import _schema_for_substrait , arrow_to_vortex
87
98
109class TestSchemaForSubstrait :
@@ -21,29 +20,35 @@ def test_binary_view_mapped_to_binary(self):
2120 assert result .field ("col" ).type == pa .binary ()
2221
2322 def test_other_types_unchanged (self ):
24- schema = pa .schema ([
25- ("int_col" , pa .int64 ()),
26- ("str_col" , pa .string ()),
27- ("bin_col" , pa .binary ()),
28- ("float_col" , pa .float64 ()),
29- ])
23+ schema = pa .schema (
24+ [
25+ ("int_col" , pa .int64 ()),
26+ ("str_col" , pa .string ()),
27+ ("bin_col" , pa .binary ()),
28+ ("float_col" , pa .float64 ()),
29+ ]
30+ )
3031 result = _schema_for_substrait (schema )
3132 assert result == schema
3233
3334 def test_mixed_schema (self ):
34- schema = pa .schema ([
35- ("sv" , pa .string_view ()),
36- ("bv" , pa .binary_view ()),
37- ("s" , pa .string ()),
38- ("i" , pa .int64 ()),
39- ])
35+ schema = pa .schema (
36+ [
37+ ("sv" , pa .string_view ()),
38+ ("bv" , pa .binary_view ()),
39+ ("s" , pa .string ()),
40+ ("i" , pa .int64 ()),
41+ ]
42+ )
4043 result = _schema_for_substrait (schema )
41- expected = pa .schema ([
42- ("sv" , pa .string ()),
43- ("bv" , pa .binary ()),
44- ("s" , pa .string ()),
45- ("i" , pa .int64 ()),
46- ])
44+ expected = pa .schema (
45+ [
46+ ("sv" , pa .string ()),
47+ ("bv" , pa .binary ()),
48+ ("s" , pa .string ()),
49+ ("i" , pa .int64 ()),
50+ ]
51+ )
4752 assert result == expected
4853
4954
@@ -69,19 +74,24 @@ def test_string_view_comparison_expression(self):
6974 assert vortex_expr is not None
7075
7176 def test_mixed_view_and_regular_types (self ):
72- schema = pa .schema ([
73- ("id" , pa .int64 ()),
74- ("name" , pa .string_view ()),
75- ("data" , pa .binary_view ()),
76- ])
77+ schema = pa .schema (
78+ [
79+ ("id" , pa .int64 ()),
80+ ("name" , pa .string_view ()),
81+ ("data" , pa .binary_view ()),
82+ ]
83+ )
7784 expr = (pc .field ("id" ) > 10 ) & (pc .field ("name" ) == "test" )
7885 vortex_expr = arrow_to_vortex (expr , schema )
7986 assert vortex_expr is not None
8087
81- @pytest .mark .parametrize ("view_type,value" , [
82- (pa .string_view (), "test" ),
83- (pa .binary_view (), b"test" ),
84- ])
88+ @pytest .mark .parametrize (
89+ "view_type,value" ,
90+ [
91+ (pa .string_view (), "test" ),
92+ (pa .binary_view (), b"test" ),
93+ ],
94+ )
8595 def test_view_types_parametrized (self , view_type , value ):
8696 schema = pa .schema ([("col" , view_type )])
8797 expr = pc .field ("col" ) == value
0 commit comments