21
21
import pyarrow as pa
22
22
import pytest
23
23
24
- from pyiceberg import schema
25
24
from pyiceberg .exceptions import ResolveError , ValidationError
26
25
from pyiceberg .schema import (
27
26
Accessor ,
28
27
Schema ,
29
28
build_position_accessors ,
29
+ index_by_id ,
30
+ index_by_name ,
30
31
promote ,
31
32
prune_columns ,
32
33
sanitize_column_names ,
@@ -90,15 +91,15 @@ def test_schema_str(table_schema_simple: Schema) -> None:
90
91
91
92
def test_schema_repr_single_field () -> None :
92
93
"""Test schema representation"""
93
- actual = repr (schema . Schema (NestedField (field_id = 1 , name = "foo" , field_type = StringType ()), schema_id = 1 ))
94
+ actual = repr (Schema (NestedField (field_id = 1 , name = "foo" , field_type = StringType ()), schema_id = 1 ))
94
95
expected = "Schema(NestedField(field_id=1, name='foo', field_type=StringType(), required=False), schema_id=1, identifier_field_ids=[])"
95
96
assert expected == actual
96
97
97
98
98
99
def test_schema_repr_two_fields () -> None :
99
100
"""Test schema representation"""
100
101
actual = repr (
101
- schema . Schema (
102
+ Schema (
102
103
NestedField (field_id = 1 , name = "foo" , field_type = StringType ()),
103
104
NestedField (field_id = 2 , name = "bar" , field_type = IntegerType (), required = False ),
104
105
schema_id = 1 ,
@@ -111,7 +112,7 @@ def test_schema_repr_two_fields() -> None:
111
112
def test_schema_raise_on_duplicate_names () -> None :
112
113
"""Test schema representation"""
113
114
with pytest .raises (ValueError ) as exc_info :
114
- schema . Schema (
115
+ Schema (
115
116
NestedField (field_id = 1 , name = "foo" , field_type = StringType (), required = False ),
116
117
NestedField (field_id = 2 , name = "bar" , field_type = IntegerType (), required = True ),
117
118
NestedField (field_id = 3 , name = "baz" , field_type = BooleanType (), required = False ),
@@ -125,7 +126,7 @@ def test_schema_raise_on_duplicate_names() -> None:
125
126
126
127
def test_schema_index_by_id_visitor (table_schema_nested : Schema ) -> None :
127
128
"""Test index_by_id visitor function"""
128
- index = schema . index_by_id (table_schema_nested )
129
+ index = index_by_id (table_schema_nested )
129
130
assert index == {
130
131
1 : NestedField (field_id = 1 , name = "foo" , field_type = StringType (), required = False ),
131
132
2 : NestedField (field_id = 2 , name = "bar" , field_type = IntegerType (), required = True ),
@@ -198,7 +199,7 @@ def test_schema_index_by_id_visitor(table_schema_nested: Schema) -> None:
198
199
199
200
def test_schema_index_by_name_visitor (table_schema_nested : Schema ) -> None :
200
201
"""Test index_by_name visitor function"""
201
- table_schema_nested = schema . Schema (
202
+ table_schema_nested = Schema (
202
203
NestedField (field_id = 1 , name = "foo" , field_type = StringType (), required = False ),
203
204
NestedField (field_id = 2 , name = "bar" , field_type = IntegerType (), required = True ),
204
205
NestedField (field_id = 3 , name = "baz" , field_type = BooleanType (), required = False ),
@@ -245,7 +246,7 @@ def test_schema_index_by_name_visitor(table_schema_nested: Schema) -> None:
245
246
schema_id = 1 ,
246
247
identifier_field_ids = [2 ],
247
248
)
248
- index = schema . index_by_name (table_schema_nested )
249
+ index = index_by_name (table_schema_nested )
249
250
assert index == {
250
251
"foo" : 1 ,
251
252
"bar" : 2 ,
@@ -301,7 +302,7 @@ def test_schema_find_column_name_by_id(table_schema_simple: Schema) -> None:
301
302
302
303
def test_schema_find_field_by_id (table_schema_simple : Schema ) -> None :
303
304
"""Test finding a column using its field ID"""
304
- index = schema . index_by_id (table_schema_simple )
305
+ index = index_by_id (table_schema_simple )
305
306
306
307
column1 = index [1 ]
307
308
assert isinstance (column1 , NestedField )
@@ -324,15 +325,15 @@ def test_schema_find_field_by_id(table_schema_simple: Schema) -> None:
324
325
325
326
def test_schema_find_field_by_id_raise_on_unknown_field (table_schema_simple : Schema ) -> None :
326
327
"""Test raising when the field ID is not found among columns"""
327
- index = schema . index_by_id (table_schema_simple )
328
+ index = index_by_id (table_schema_simple )
328
329
with pytest .raises (Exception ) as exc_info :
329
330
_ = index [4 ]
330
331
assert str (exc_info .value ) == "4"
331
332
332
333
333
334
def test_schema_find_field_type_by_id (table_schema_simple : Schema ) -> None :
334
335
"""Test retrieving a columns' type using its field ID"""
335
- index = schema . index_by_id (table_schema_simple )
336
+ index = index_by_id (table_schema_simple )
336
337
assert index [1 ] == NestedField (field_id = 1 , name = "foo" , field_type = StringType (), required = False )
337
338
assert index [2 ] == NestedField (field_id = 2 , name = "bar" , field_type = IntegerType (), required = True )
338
339
assert index [3 ] == NestedField (field_id = 3 , name = "baz" , field_type = BooleanType (), required = False )
@@ -341,7 +342,7 @@ def test_schema_find_field_type_by_id(table_schema_simple: Schema) -> None:
341
342
def test_index_by_id_schema_visitor_raise_on_unregistered_type () -> None :
342
343
"""Test raising a NotImplementedError when an invalid type is provided to the index_by_id function"""
343
344
with pytest .raises (NotImplementedError ) as exc_info :
344
- schema . index_by_id ("foo" ) # type: ignore
345
+ index_by_id ("foo" ) # type: ignore
345
346
assert "Cannot visit non-type: foo" in str (exc_info .value )
346
347
347
348
0 commit comments