refactor partition_summary_limit into SnapshotSummaryCollector constructor

Author: Yingjian Wu

pyiceberg/table/snapshots.py

@@ -272,13 +272,10 @@ class SnapshotSummaryCollector:
     partition_metrics: DefaultDict[str, UpdateMetrics]
     max_changed_partitions_for_summaries: int
 
-    def __init__(self) -> None:
+    def __init__(self, partition_summary_limit: int = 0) -> None:
         self.metrics = UpdateMetrics()
         self.partition_metrics = defaultdict(UpdateMetrics)
-        self.max_changed_partitions_for_summaries = 0
-
-    def set_partition_summary_limit(self, limit: int) -> None:
-        self.max_changed_partitions_for_summaries = limit
+        self.max_changed_partitions_for_summaries = partition_summary_limit
 
     def add_file(self, data_file: DataFile, schema: Schema, partition_spec: PartitionSpec = UNPARTITIONED_PARTITION_SPEC) -> None:
         self.metrics.add_file(data_file)

pyiceberg/table/update/snapshot.py

@@ -203,13 +203,12 @@ def _write_delete_manifest() -> List[ManifestFile]:
     def _summary(self, snapshot_properties: Dict[str, str] = EMPTY_DICT) -> Summary:
         from pyiceberg.table import TableProperties
 
-        ssc = SnapshotSummaryCollector()
         partition_summary_limit = int(
             self._transaction.table_metadata.properties.get(
                 TableProperties.WRITE_PARTITION_SUMMARY_LIMIT, TableProperties.WRITE_PARTITION_SUMMARY_LIMIT_DEFAULT
             )
         )
-        ssc.set_partition_summary_limit(partition_summary_limit)
+        ssc = SnapshotSummaryCollector(partition_summary_limit=partition_summary_limit)
 
         for data_file in self._added_data_files:
             ssc.add_file(

tests/table/test_snapshots.py

@@ -201,8 +201,27 @@ def test_snapshot_summary_collector_with_partition() -> None:
         "changed-partition-count": "2",
     }
 
+
+@pytest.mark.integration
+def test_snapshot_summary_collector_with_partition_limit_in_constructor() -> None:
+    # Given
+    partition_summary_limit = 10
+    ssc = SnapshotSummaryCollector(partition_summary_limit=partition_summary_limit)
+
+    assert ssc.build() == {}
+    schema = Schema(
+        NestedField(field_id=1, name="bool_field", field_type=BooleanType(), required=False),
+        NestedField(field_id=2, name="string_field", field_type=StringType(), required=False),
+        NestedField(field_id=3, name="int_field", field_type=IntegerType(), required=False),
+    )
+    spec = PartitionSpec(PartitionField(source_id=3, field_id=1001, transform=IdentityTransform(), name="int_field"))
+    data_file_1 = DataFile(content=DataFileContent.DATA, record_count=100, file_size_in_bytes=1234, partition=Record(int_field=1))
+    data_file_2 = DataFile(content=DataFileContent.DATA, record_count=200, file_size_in_bytes=4321, partition=Record(int_field=2))
+
     # When
-    ssc.set_partition_summary_limit(10)
+    ssc.add_file(data_file=data_file_1, schema=schema, partition_spec=spec)
+    ssc.remove_file(data_file=data_file_1, schema=schema, partition_spec=spec)
+    ssc.remove_file(data_file=data_file_2, schema=schema, partition_spec=spec)
 
     # Then
     assert ssc.build() == {