Skip to content

Commit e63bedf

Browse files
kevinjqliukevinjqliu
committed
fix test_schema
1 parent 6cf4a51 commit e63bedf

1 file changed

Lines changed: 49 additions & 48 deletions

File tree

‎tests/test_schema.py‎

Lines changed: 49 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
prune_columns,
3333
sanitize_column_names,
3434
)
35+
from pyiceberg.table import Table, Transaction
3536
from pyiceberg.table.update.schema import UpdateSchema
3637
from pyiceberg.typedef import EMPTY_DICT, StructProtocol
3738
from pyiceberg.types import (
@@ -927,14 +928,14 @@ def primitive_fields() -> List[NestedField]:
927928
]
928929

929930

930-
def test_add_top_level_primitives(primitive_fields: List[NestedField]) -> None:
931+
def test_add_top_level_primitives(primitive_fields: List[NestedField], table_v2: Table) -> None:
931932
for primitive_field in primitive_fields:
932933
new_schema = Schema(primitive_field)
933-
applied = UpdateSchema(transaction=None, schema=Schema()).union_by_name(new_schema)._apply() # type: ignore
934+
applied = UpdateSchema(transaction=Transaction(table_v2), schema=Schema()).union_by_name(new_schema)._apply()
934935
assert applied == new_schema
935936

936937

937-
def test_add_top_level_list_of_primitives(primitive_fields: NestedField) -> None:
938+
def test_add_top_level_list_of_primitives(primitive_fields: NestedField, table_v2: Table) -> None:
938939
for primitive_type in TEST_PRIMITIVE_TYPES:
939940
new_schema = Schema(
940941
NestedField(
@@ -944,11 +945,11 @@ def test_add_top_level_list_of_primitives(primitive_fields: NestedField) -> None
944945
required=False,
945946
)
946947
)
947-
applied = UpdateSchema(transaction=None, schema=Schema()).union_by_name(new_schema)._apply() # type: ignore
948+
applied = UpdateSchema(transaction=Transaction(table_v2), schema=Schema()).union_by_name(new_schema)._apply()
948949
assert applied.as_struct() == new_schema.as_struct()
949950

950951

951-
def test_add_top_level_map_of_primitives(primitive_fields: NestedField) -> None:
952+
def test_add_top_level_map_of_primitives(primitive_fields: NestedField, table_v2: Table) -> None:
952953
for primitive_type in TEST_PRIMITIVE_TYPES:
953954
new_schema = Schema(
954955
NestedField(
@@ -960,11 +961,11 @@ def test_add_top_level_map_of_primitives(primitive_fields: NestedField) -> None:
960961
required=False,
961962
)
962963
)
963-
applied = UpdateSchema(transaction=None, schema=Schema()).union_by_name(new_schema)._apply() # type: ignore
964+
applied = UpdateSchema(transaction=Transaction(table_v2), schema=Schema()).union_by_name(new_schema)._apply()
964965
assert applied.as_struct() == new_schema.as_struct()
965966

966967

967-
def test_add_top_struct_of_primitives(primitive_fields: NestedField) -> None:
968+
def test_add_top_struct_of_primitives(primitive_fields: NestedField, table_v2: Table) -> None:
968969
for primitive_type in TEST_PRIMITIVE_TYPES:
969970
new_schema = Schema(
970971
NestedField(
@@ -974,11 +975,11 @@ def test_add_top_struct_of_primitives(primitive_fields: NestedField) -> None:
974975
required=False,
975976
)
976977
)
977-
applied = UpdateSchema(transaction=None, schema=Schema()).union_by_name(new_schema)._apply() # type: ignore
978+
applied = UpdateSchema(transaction=Transaction(table_v2), schema=Schema()).union_by_name(new_schema)._apply()
978979
assert applied.as_struct() == new_schema.as_struct()
979980

980981

981-
def test_add_nested_primitive(primitive_fields: NestedField) -> None:
982+
def test_add_nested_primitive(primitive_fields: NestedField, table_v2: Table) -> None:
982983
for primitive_type in TEST_PRIMITIVE_TYPES:
983984
current_schema = Schema(NestedField(field_id=1, name="aStruct", field_type=StructType(), required=False))
984985
new_schema = Schema(
@@ -989,7 +990,7 @@ def test_add_nested_primitive(primitive_fields: NestedField) -> None:
989990
required=False,
990991
)
991992
)
992-
applied = UpdateSchema(None, None, schema=current_schema).union_by_name(new_schema)._apply() # type: ignore
993+
applied = UpdateSchema(transaction=Transaction(table_v2), schema=current_schema).union_by_name(new_schema)._apply()
993994
assert applied.as_struct() == new_schema.as_struct()
994995

995996

@@ -1002,18 +1003,18 @@ def _primitive_fields(types: List[PrimitiveType], start_id: int = 0) -> List[Nes
10021003
return fields
10031004

10041005

1005-
def test_add_nested_primitives(primitive_fields: NestedField) -> None:
1006+
def test_add_nested_primitives(primitive_fields: NestedField, table_v2: Table) -> None:
10061007
current_schema = Schema(NestedField(field_id=1, name="aStruct", field_type=StructType(), required=False))
10071008
new_schema = Schema(
10081009
NestedField(
10091010
field_id=1, name="aStruct", field_type=StructType(*_primitive_fields(TEST_PRIMITIVE_TYPES, 2)), required=False
10101011
)
10111012
)
1012-
applied = UpdateSchema(transaction=None, schema=current_schema).union_by_name(new_schema)._apply() # type: ignore
1013+
applied = UpdateSchema(transaction=Transaction(table_v2), schema=current_schema).union_by_name(new_schema)._apply()
10131014
assert applied.as_struct() == new_schema.as_struct()
10141015

10151016

1016-
def test_add_nested_lists(primitive_fields: NestedField) -> None:
1017+
def test_add_nested_lists(primitive_fields: NestedField, table_v2: Table) -> None:
10171018
new_schema = Schema(
10181019
NestedField(
10191020
field_id=1,
@@ -1050,11 +1051,11 @@ def test_add_nested_lists(primitive_fields: NestedField) -> None:
10501051
required=False,
10511052
)
10521053
)
1053-
applied = UpdateSchema(transaction=None, schema=Schema()).union_by_name(new_schema)._apply() # type: ignore
1054+
applied = UpdateSchema(transaction=Transaction(table_v2), schema=Schema()).union_by_name(new_schema)._apply()
10541055
assert applied.as_struct() == new_schema.as_struct()
10551056

10561057

1057-
def test_add_nested_struct(primitive_fields: NestedField) -> None:
1058+
def test_add_nested_struct(primitive_fields: NestedField, table_v2: Table) -> None:
10581059
new_schema = Schema(
10591060
NestedField(
10601061
field_id=1,
@@ -1100,11 +1101,11 @@ def test_add_nested_struct(primitive_fields: NestedField) -> None:
11001101
required=False,
11011102
)
11021103
)
1103-
applied = UpdateSchema(transaction=None, schema=Schema()).union_by_name(new_schema)._apply() # type: ignore
1104+
applied = UpdateSchema(transaction=Transaction(table_v2), schema=Schema()).union_by_name(new_schema)._apply()
11041105
assert applied.as_struct() == new_schema.as_struct()
11051106

11061107

1107-
def test_add_nested_maps(primitive_fields: NestedField) -> None:
1108+
def test_add_nested_maps(primitive_fields: NestedField, table_v2: Table) -> None:
11081109
new_schema = Schema(
11091110
NestedField(
11101111
field_id=1,
@@ -1143,11 +1144,11 @@ def test_add_nested_maps(primitive_fields: NestedField) -> None:
11431144
required=False,
11441145
)
11451146
)
1146-
applied = UpdateSchema(transaction=None, schema=Schema()).union_by_name(new_schema)._apply() # type: ignore
1147+
applied = UpdateSchema(transaction=Transaction(table_v2), schema=Schema()).union_by_name(new_schema)._apply()
11471148
assert applied.as_struct() == new_schema.as_struct()
11481149

11491150

1150-
def test_detect_invalid_top_level_list() -> None:
1151+
def test_detect_invalid_top_level_list(table_v2: Table) -> None:
11511152
current_schema = Schema(
11521153
NestedField(
11531154
field_id=1,
@@ -1166,10 +1167,10 @@ def test_detect_invalid_top_level_list() -> None:
11661167
)
11671168

11681169
with pytest.raises(ValidationError, match="Cannot change column type: aList.element: string -> double"):
1169-
_ = UpdateSchema(transaction=None, schema=current_schema).union_by_name(new_schema)._apply() # type: ignore
1170+
_ = UpdateSchema(transaction=Transaction(table_v2), schema=current_schema).union_by_name(new_schema)._apply()
11701171

11711172

1172-
def test_detect_invalid_top_level_maps() -> None:
1173+
def test_detect_invalid_top_level_maps(table_v2: Table) -> None:
11731174
current_schema = Schema(
11741175
NestedField(
11751176
field_id=1,
@@ -1188,68 +1189,68 @@ def test_detect_invalid_top_level_maps() -> None:
11881189
)
11891190

11901191
with pytest.raises(ValidationError, match="Cannot change column type: aMap.key: string -> uuid"):
1191-
_ = UpdateSchema(transaction=None, schema=current_schema).union_by_name(new_schema)._apply() # type: ignore
1192+
_ = UpdateSchema(transaction=Transaction(table_v2), schema=current_schema).union_by_name(new_schema)._apply()
11921193

11931194

1194-
def test_allow_double_to_float() -> None:
1195+
def test_allow_double_to_float(table_v2: Table) -> None:
11951196
current_schema = Schema(NestedField(field_id=1, name="aCol", field_type=DoubleType(), required=False))
11961197
new_schema = Schema(NestedField(field_id=1, name="aCol", field_type=FloatType(), required=False))
11971198

1198-
applied = UpdateSchema(transaction=None, schema=current_schema).union_by_name(new_schema)._apply() # type: ignore
1199+
applied = UpdateSchema(transaction=Transaction(table_v2), schema=current_schema).union_by_name(new_schema)._apply()
11991200

12001201
assert applied.as_struct() == current_schema.as_struct()
12011202
assert len(applied.fields) == 1
12021203
assert isinstance(applied.fields[0].field_type, DoubleType)
12031204

12041205

1205-
def test_promote_float_to_double() -> None:
1206+
def test_promote_float_to_double(table_v2: Table) -> None:
12061207
current_schema = Schema(NestedField(field_id=1, name="aCol", field_type=FloatType(), required=False))
12071208
new_schema = Schema(NestedField(field_id=1, name="aCol", field_type=DoubleType(), required=False))
12081209

1209-
applied = UpdateSchema(transaction=None, schema=current_schema).union_by_name(new_schema)._apply() # type: ignore
1210+
applied = UpdateSchema(transaction=Transaction(table_v2), schema=current_schema).union_by_name(new_schema)._apply()
12101211

12111212
assert applied.as_struct() == new_schema.as_struct()
12121213
assert len(applied.fields) == 1
12131214
assert isinstance(applied.fields[0].field_type, DoubleType)
12141215

12151216

1216-
def test_allow_long_to_int() -> None:
1217+
def test_allow_long_to_int(table_v2: Table) -> None:
12171218
current_schema = Schema(NestedField(field_id=1, name="aCol", field_type=LongType(), required=False))
12181219
new_schema = Schema(NestedField(field_id=1, name="aCol", field_type=IntegerType(), required=False))
12191220

1220-
applied = UpdateSchema(transaction=None, schema=current_schema).union_by_name(new_schema)._apply() # type: ignore
1221+
applied = UpdateSchema(transaction=Transaction(table_v2), schema=current_schema).union_by_name(new_schema)._apply()
12211222

12221223
assert applied.as_struct() == current_schema.as_struct()
12231224
assert len(applied.fields) == 1
12241225
assert isinstance(applied.fields[0].field_type, LongType)
12251226

12261227

1227-
def test_promote_int_to_long() -> None:
1228+
def test_promote_int_to_long(table_v2: Table) -> None:
12281229
current_schema = Schema(NestedField(field_id=1, name="aCol", field_type=IntegerType(), required=False))
12291230
new_schema = Schema(NestedField(field_id=1, name="aCol", field_type=LongType(), required=False))
12301231

1231-
applied = UpdateSchema(transaction=None, schema=current_schema).union_by_name(new_schema)._apply() # type: ignore
1232+
applied = UpdateSchema(transaction=Transaction(table_v2), schema=current_schema).union_by_name(new_schema)._apply()
12321233

12331234
assert applied.as_struct() == new_schema.as_struct()
12341235
assert len(applied.fields) == 1
12351236
assert isinstance(applied.fields[0].field_type, LongType)
12361237

12371238

1238-
def test_detect_invalid_promotion_string_to_float() -> None:
1239+
def test_detect_invalid_promotion_string_to_float(table_v2: Table) -> None:
12391240
current_schema = Schema(NestedField(field_id=1, name="aCol", field_type=StringType(), required=False))
12401241
new_schema = Schema(NestedField(field_id=1, name="aCol", field_type=FloatType(), required=False))
12411242

12421243
with pytest.raises(ValidationError, match="Cannot change column type: aCol: string -> float"):
1243-
_ = UpdateSchema(transaction=None, schema=current_schema).union_by_name(new_schema)._apply() # type: ignore
1244+
_ = UpdateSchema(transaction=Transaction(table_v2), schema=current_schema).union_by_name(new_schema)._apply()
12441245

12451246

12461247
# decimal(P,S) Fixed-point decimal; precision P, scale S -> Scale is fixed [1],
12471248
# precision must be 38 or less
1248-
def test_type_promote_decimal_to_fixed_scale_with_wider_precision() -> None:
1249+
def test_type_promote_decimal_to_fixed_scale_with_wider_precision(table_v2: Table) -> None:
12491250
current_schema = Schema(NestedField(field_id=1, name="aCol", field_type=DecimalType(precision=20, scale=1), required=False))
12501251
new_schema = Schema(NestedField(field_id=1, name="aCol", field_type=DecimalType(precision=22, scale=1), required=False))
12511252

1252-
applied = UpdateSchema(transaction=None, schema=current_schema).union_by_name(new_schema)._apply() # type: ignore
1253+
applied = UpdateSchema(transaction=Transaction(table_v2), schema=current_schema).union_by_name(new_schema)._apply()
12531254

12541255
assert applied.as_struct() == new_schema.as_struct()
12551256
assert len(applied.fields) == 1
@@ -1260,7 +1261,7 @@ def test_type_promote_decimal_to_fixed_scale_with_wider_precision() -> None:
12601261
assert decimal_type.scale == 1
12611262

12621263

1263-
def test_add_nested_structs(primitive_fields: NestedField) -> None:
1264+
def test_add_nested_structs(primitive_fields: NestedField, table_v2: Table) -> None:
12641265
schema = Schema(
12651266
NestedField(
12661267
field_id=1,
@@ -1317,7 +1318,7 @@ def test_add_nested_structs(primitive_fields: NestedField) -> None:
13171318
required=False,
13181319
)
13191320
)
1320-
applied = UpdateSchema(transaction=None, schema=schema).union_by_name(new_schema)._apply() # type: ignore
1321+
applied = UpdateSchema(transaction=Transaction(table_v2), schema=schema).union_by_name(new_schema)._apply()
13211322

13221323
expected = Schema(
13231324
NestedField(
@@ -1352,15 +1353,15 @@ def test_add_nested_structs(primitive_fields: NestedField) -> None:
13521353
assert applied.as_struct() == expected.as_struct()
13531354

13541355

1355-
def test_replace_list_with_primitive() -> None:
1356+
def test_replace_list_with_primitive(table_v2: Table) -> None:
13561357
current_schema = Schema(NestedField(field_id=1, name="aCol", field_type=ListType(element_id=2, element_type=StringType())))
13571358
new_schema = Schema(NestedField(field_id=1, name="aCol", field_type=StringType()))
13581359

13591360
with pytest.raises(ValidationError, match="Cannot change column type: list<string> is not a primitive"):
1360-
_ = UpdateSchema(transaction=None, schema=current_schema).union_by_name(new_schema)._apply() # type: ignore
1361+
_ = UpdateSchema(transaction=Transaction(table_v2), schema=current_schema).union_by_name(new_schema)._apply()
13611362

13621363

1363-
def test_mirrored_schemas() -> None:
1364+
def test_mirrored_schemas(table_v2: Table) -> None:
13641365
current_schema = Schema(
13651366
NestedField(9, "struct1", StructType(NestedField(8, "string1", StringType(), required=False)), required=False),
13661367
NestedField(6, "list1", ListType(element_id=7, element_type=StringType(), element_required=False), required=False),
@@ -1380,12 +1381,12 @@ def test_mirrored_schemas() -> None:
13801381
NestedField(9, "string6", StringType(), required=False),
13811382
)
13821383

1383-
applied = UpdateSchema(transaction=None, schema=current_schema).union_by_name(mirrored_schema)._apply() # type: ignore
1384+
applied = UpdateSchema(transaction=Transaction(table_v2), schema=current_schema).union_by_name(mirrored_schema)._apply()
13841385

13851386
assert applied.as_struct() == current_schema.as_struct()
13861387

13871388

1388-
def test_add_new_top_level_struct() -> None:
1389+
def test_add_new_top_level_struct(table_v2: Table) -> None:
13891390
current_schema = Schema(
13901391
NestedField(
13911392
1,
@@ -1432,12 +1433,12 @@ def test_add_new_top_level_struct() -> None:
14321433
),
14331434
)
14341435

1435-
applied = UpdateSchema(transaction=None, schema=current_schema).union_by_name(observed_schema)._apply() # type: ignore
1436+
applied = UpdateSchema(transaction=Transaction(table_v2), schema=current_schema).union_by_name(observed_schema)._apply()
14361437

14371438
assert applied.as_struct() == observed_schema.as_struct()
14381439

14391440

1440-
def test_append_nested_struct() -> None:
1441+
def test_append_nested_struct(table_v2: Table) -> None:
14411442
current_schema = Schema(
14421443
NestedField(
14431444
field_id=1,
@@ -1511,12 +1512,12 @@ def test_append_nested_struct() -> None:
15111512
)
15121513
)
15131514

1514-
applied = UpdateSchema(transaction=None, schema=current_schema).union_by_name(observed_schema)._apply() # type: ignore
1515+
applied = UpdateSchema(transaction=Transaction(table_v2), schema=current_schema).union_by_name(observed_schema)._apply()
15151516

15161517
assert applied.as_struct() == observed_schema.as_struct()
15171518

15181519

1519-
def test_append_nested_lists() -> None:
1520+
def test_append_nested_lists(table_v2: Table) -> None:
15201521
current_schema = Schema(
15211522
NestedField(
15221523
field_id=1,
@@ -1576,7 +1577,7 @@ def test_append_nested_lists() -> None:
15761577
required=False,
15771578
)
15781579
)
1579-
union = UpdateSchema(transaction=None, schema=current_schema).union_by_name(observed_schema)._apply() # type: ignore
1580+
union = UpdateSchema(transaction=Transaction(table_v2), schema=current_schema).union_by_name(observed_schema)._apply()
15801581

15811582
expected = Schema(
15821583
NestedField(
@@ -1617,7 +1618,7 @@ def test_append_nested_lists() -> None:
16171618
assert union.as_struct() == expected.as_struct()
16181619

16191620

1620-
def test_union_with_pa_schema(primitive_fields: NestedField) -> None:
1621+
def test_union_with_pa_schema(primitive_fields: NestedField, table_v2: Table) -> None:
16211622
base_schema = Schema(NestedField(field_id=1, name="foo", field_type=StringType(), required=True))
16221623

16231624
pa_schema = pa.schema(
@@ -1628,7 +1629,7 @@ def test_union_with_pa_schema(primitive_fields: NestedField) -> None:
16281629
]
16291630
)
16301631

1631-
new_schema = UpdateSchema(transaction=None, schema=base_schema).union_by_name(pa_schema)._apply() # type: ignore
1632+
new_schema = UpdateSchema(transaction=Transaction(table_v2), schema=base_schema).union_by_name(pa_schema)._apply()
16321633

16331634
expected_schema = Schema(
16341635
NestedField(field_id=1, name="foo", field_type=StringType(), required=True),

0 commit comments

Comments
 (0)