Skip to content

Commit de47590

Browse files
authored
fix: support MonthTransform for partitioning (#1176)
* fix: support MonthTransform for partitioning * test: add tests for generating default names for other transforms * fix: delete duplicate test case * chore: run formatting
1 parent e545a05 commit de47590

File tree

2 files changed

+67
-0
lines changed

2 files changed

+67
-0
lines changed

pyiceberg/partitioning.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
DayTransform,
4747
HourTransform,
4848
IdentityTransform,
49+
MonthTransform,
4950
Transform,
5051
TruncateTransform,
5152
UnknownTransform,
@@ -359,6 +360,8 @@ def _visit_partition_field(schema: Schema, field: PartitionField, visitor: Parti
359360
return visitor.day(field.field_id, source_name, field.source_id)
360361
elif isinstance(transform, HourTransform):
361362
return visitor.hour(field.field_id, source_name, field.source_id)
363+
elif isinstance(transform, MonthTransform):
364+
return visitor.month(field.field_id, source_name, field.source_id)
362365
elif isinstance(transform, YearTransform):
363366
return visitor.year(field.field_id, source_name, field.source_id)
364367
elif isinstance(transform, VoidTransform):

tests/integration/test_partition_evolution.py

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,14 @@ def test_add_year(catalog: Catalog) -> None:
9292
_validate_new_partition_fields(table, 1000, 1, 1000, PartitionField(2, 1000, YearTransform(), "year_transform"))
9393

9494

95+
@pytest.mark.integration
96+
@pytest.mark.parametrize("catalog", [pytest.lazy_fixture("session_catalog_hive"), pytest.lazy_fixture("session_catalog")])
97+
def test_add_year_generates_default_name(catalog: Catalog) -> None:
98+
table = _table(catalog)
99+
table.update_spec().add_field("event_ts", YearTransform()).commit()
100+
_validate_new_partition_fields(table, 1000, 1, 1000, PartitionField(2, 1000, YearTransform(), "event_ts_year"))
101+
102+
95103
@pytest.mark.integration
96104
@pytest.mark.parametrize("catalog", [pytest.lazy_fixture("session_catalog_hive"), pytest.lazy_fixture("session_catalog")])
97105
def test_add_month(catalog: Catalog) -> None:
@@ -100,6 +108,14 @@ def test_add_month(catalog: Catalog) -> None:
100108
_validate_new_partition_fields(table, 1000, 1, 1000, PartitionField(2, 1000, MonthTransform(), "month_transform"))
101109

102110

111+
@pytest.mark.integration
112+
@pytest.mark.parametrize("catalog", [pytest.lazy_fixture("session_catalog_hive"), pytest.lazy_fixture("session_catalog")])
113+
def test_add_month_generates_default_name(catalog: Catalog) -> None:
114+
table = _table(catalog)
115+
table.update_spec().add_field("event_ts", MonthTransform()).commit()
116+
_validate_new_partition_fields(table, 1000, 1, 1000, PartitionField(2, 1000, MonthTransform(), "event_ts_month"))
117+
118+
103119
@pytest.mark.integration
104120
@pytest.mark.parametrize("catalog", [pytest.lazy_fixture("session_catalog_hive"), pytest.lazy_fixture("session_catalog")])
105121
def test_add_day(catalog: Catalog) -> None:
@@ -108,6 +124,14 @@ def test_add_day(catalog: Catalog) -> None:
108124
_validate_new_partition_fields(table, 1000, 1, 1000, PartitionField(2, 1000, DayTransform(), "day_transform"))
109125

110126

127+
@pytest.mark.integration
128+
@pytest.mark.parametrize("catalog", [pytest.lazy_fixture("session_catalog_hive"), pytest.lazy_fixture("session_catalog")])
129+
def test_add_day_generates_default_name(catalog: Catalog) -> None:
130+
table = _table(catalog)
131+
table.update_spec().add_field("event_ts", DayTransform()).commit()
132+
_validate_new_partition_fields(table, 1000, 1, 1000, PartitionField(2, 1000, DayTransform(), "event_ts_day"))
133+
134+
111135
@pytest.mark.integration
112136
@pytest.mark.parametrize("catalog", [pytest.lazy_fixture("session_catalog_hive"), pytest.lazy_fixture("session_catalog")])
113137
def test_add_hour(catalog: Catalog) -> None:
@@ -116,6 +140,14 @@ def test_add_hour(catalog: Catalog) -> None:
116140
_validate_new_partition_fields(table, 1000, 1, 1000, PartitionField(2, 1000, HourTransform(), "hour_transform"))
117141

118142

143+
@pytest.mark.integration
144+
@pytest.mark.parametrize("catalog", [pytest.lazy_fixture("session_catalog_hive"), pytest.lazy_fixture("session_catalog")])
145+
def test_add_hour_generates_default_name(catalog: Catalog) -> None:
146+
table = _table(catalog)
147+
table.update_spec().add_field("event_ts", HourTransform()).commit()
148+
_validate_new_partition_fields(table, 1000, 1, 1000, PartitionField(2, 1000, HourTransform(), "event_ts_hour"))
149+
150+
119151
@pytest.mark.integration
120152
@pytest.mark.parametrize("catalog", [pytest.lazy_fixture("session_catalog_hive"), pytest.lazy_fixture("session_catalog")])
121153
def test_add_bucket(catalog: Catalog, table_schema_simple: Schema) -> None:
@@ -124,6 +156,14 @@ def test_add_bucket(catalog: Catalog, table_schema_simple: Schema) -> None:
124156
_validate_new_partition_fields(simple_table, 1000, 1, 1000, PartitionField(1, 1000, BucketTransform(12), "bucket_transform"))
125157

126158

159+
@pytest.mark.integration
160+
@pytest.mark.parametrize("catalog", [pytest.lazy_fixture("session_catalog_hive"), pytest.lazy_fixture("session_catalog")])
161+
def test_add_bucket_generates_default_name(catalog: Catalog, table_schema_simple: Schema) -> None:
162+
simple_table = _create_table_with_schema(catalog, table_schema_simple, "1")
163+
simple_table.update_spec().add_field("foo", BucketTransform(12)).commit()
164+
_validate_new_partition_fields(simple_table, 1000, 1, 1000, PartitionField(1, 1000, BucketTransform(12), "foo_bucket_12"))
165+
166+
127167
@pytest.mark.integration
128168
@pytest.mark.parametrize("catalog", [pytest.lazy_fixture("session_catalog_hive"), pytest.lazy_fixture("session_catalog")])
129169
def test_add_truncate(catalog: Catalog, table_schema_simple: Schema) -> None:
@@ -134,6 +174,14 @@ def test_add_truncate(catalog: Catalog, table_schema_simple: Schema) -> None:
134174
)
135175

136176

177+
@pytest.mark.integration
178+
@pytest.mark.parametrize("catalog", [pytest.lazy_fixture("session_catalog_hive"), pytest.lazy_fixture("session_catalog")])
179+
def test_add_truncate_generates_default_name(catalog: Catalog, table_schema_simple: Schema) -> None:
180+
simple_table = _create_table_with_schema(catalog, table_schema_simple, "1")
181+
simple_table.update_spec().add_field("foo", TruncateTransform(1)).commit()
182+
_validate_new_partition_fields(simple_table, 1000, 1, 1000, PartitionField(1, 1000, TruncateTransform(1), "foo_trunc_1"))
183+
184+
137185
@pytest.mark.integration
138186
@pytest.mark.parametrize("catalog", [pytest.lazy_fixture("session_catalog_hive"), pytest.lazy_fixture("session_catalog")])
139187
def test_multiple_adds(catalog: Catalog) -> None:
@@ -152,6 +200,22 @@ def test_multiple_adds(catalog: Catalog) -> None:
152200
)
153201

154202

203+
@pytest.mark.integration
204+
@pytest.mark.parametrize("catalog", [pytest.lazy_fixture("session_catalog_hive"), pytest.lazy_fixture("session_catalog")])
205+
def test_add_void(catalog: Catalog, table_schema_simple: Schema) -> None:
206+
simple_table = _create_table_with_schema(catalog, table_schema_simple, "1")
207+
simple_table.update_spec().add_field("foo", VoidTransform(), "void_transform").commit()
208+
_validate_new_partition_fields(simple_table, 1000, 1, 1000, PartitionField(1, 1000, VoidTransform(), "void_transform"))
209+
210+
211+
@pytest.mark.integration
212+
@pytest.mark.parametrize("catalog", [pytest.lazy_fixture("session_catalog_hive"), pytest.lazy_fixture("session_catalog")])
213+
def test_add_void_generates_default_name(catalog: Catalog, table_schema_simple: Schema) -> None:
214+
simple_table = _create_table_with_schema(catalog, table_schema_simple, "1")
215+
simple_table.update_spec().add_field("foo", VoidTransform()).commit()
216+
_validate_new_partition_fields(simple_table, 1000, 1, 1000, PartitionField(1, 1000, VoidTransform(), "foo_null"))
217+
218+
155219
@pytest.mark.integration
156220
@pytest.mark.parametrize("catalog", [pytest.lazy_fixture("session_catalog_hive"), pytest.lazy_fixture("session_catalog")])
157221
def test_add_hour_to_day(catalog: Catalog) -> None:

0 commit comments

Comments
 (0)