From 5ceb80d6b3d2c99a3b2694215d86b65e75b94030 Mon Sep 17 00:00:00 2001 From: Kevin Liu Date: Fri, 8 Mar 2024 08:25:53 -0800 Subject: [PATCH] rewrite set_properties --- tests/integration/test_writes.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/tests/integration/test_writes.py b/tests/integration/test_writes.py index 51c60dcadc..e653bc868a 100644 --- a/tests/integration/test_writes.py +++ b/tests/integration/test_writes.py @@ -36,8 +36,8 @@ from pyiceberg.catalog.sql import SqlCatalog from pyiceberg.exceptions import NoSuchTableError from pyiceberg.schema import Schema +from pyiceberg.table import Table, TableProperties, _dataframe_to_data_files from pyiceberg.typedef import Properties -from pyiceberg.table import SetPropertiesUpdate, Table, TableProperties, _dataframe_to_data_files from pyiceberg.types import ( BinaryType, BooleanType, @@ -391,11 +391,6 @@ def get_data_files_count(identifier: str) -> int: """ ).count() - def set_table_properties(tbl: Table, properties: Properties) -> Table: - with tbl.transaction() as transaction: - transaction._apply((SetPropertiesUpdate(updates=properties),)) - return tbl - # writes 1 data file since the table is smaller than default target file size assert arrow_table_with_null.nbytes < TableProperties.WRITE_TARGET_FILE_SIZE_BYTES_DEFAULT tbl.overwrite(arrow_table_with_null) @@ -409,7 +404,7 @@ def set_table_properties(tbl: Table, properties: Properties) -> Table: # writes multiple data files once target file size is overridden target_file_size = arrow_table_with_null.nbytes - tbl = set_table_properties(tbl, {TableProperties.WRITE_TARGET_FILE_SIZE_BYTES: str(target_file_size)}) + tbl = tbl.transaction().set_properties({TableProperties.WRITE_TARGET_FILE_SIZE_BYTES: str(target_file_size)}).commit_transaction() assert str(target_file_size) == tbl.properties.get(TableProperties.WRITE_TARGET_FILE_SIZE_BYTES) assert target_file_size < bigger_arrow_tbl.nbytes tbl.overwrite(bigger_arrow_tbl) @@ -417,7 +412,7 @@ def set_table_properties(tbl: Table, properties: Properties) -> Table: # writes half the number of data files when target file size doubles target_file_size = arrow_table_with_null.nbytes * 2 - tbl = set_table_properties(tbl, {TableProperties.WRITE_TARGET_FILE_SIZE_BYTES: str(target_file_size)}) + tbl = tbl.transaction().set_properties({TableProperties.WRITE_TARGET_FILE_SIZE_BYTES: str(target_file_size)}).commit_transaction() assert str(target_file_size) == tbl.properties.get(TableProperties.WRITE_TARGET_FILE_SIZE_BYTES) assert target_file_size < bigger_arrow_tbl.nbytes tbl.overwrite(bigger_arrow_tbl)