Add enforce_primary_key_uniqueness support to converter

Author: yiqinzhu

deltacat/compute/converter/converter_session.py

@@ -10,11 +10,11 @@
 )
 import logging
 from deltacat import logs
-from collections import defaultdict
 from deltacat.compute.converter.model.converter_session_params import (
     ConverterSessionParams,
 )
 
+
 from deltacat.compute.converter.constants import DEFAULT_MAX_PARALLEL_DATA_FILE_DOWNLOAD
 from deltacat.compute.converter.steps.convert import convert
 from deltacat.compute.converter.model.convert_input import ConvertInput
@@ -23,13 +23,16 @@
     parquet_files_dict_to_iceberg_data_files,
 )
 from deltacat.compute.converter.utils.converter_session_utils import (
-    check_data_files_sequence_number,
     construct_iceberg_table_prefix,
 )
 from deltacat.compute.converter.pyiceberg.replace_snapshot import (
     commit_overwrite_snapshot,
+    commit_append_snapshot,
 )
 from deltacat.compute.converter.pyiceberg.catalog import load_table
+from deltacat.compute.converter.utils.converter_session_utils import (
+    group_all_files_to_each_bucket,
+)
 
 logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
 
@@ -44,33 +47,15 @@ def converter_session(params: ConverterSessionParams, **kwargs):
     catalog = params.catalog
     table_name = params.iceberg_table_name
     iceberg_table = load_table(catalog, table_name)
+    enforce_primary_key_uniqueness = params.enforce_primary_key_uniqueness
     data_file_dict, equality_delete_dict, pos_delete_dict = fetch_all_bucket_files(
         iceberg_table
     )
-
-    # files_for_each_bucket contains the following files list:
-    # {partition_value: [(equality_delete_files_list, data_files_list, pos_delete_files_list)]
-    files_for_each_bucket = defaultdict(tuple)
-    for k, v in data_file_dict.items():
-        logger.info(f"data_file: k, v:{k, v}")
-    for k, v in equality_delete_dict.items():
-        logger.info(f"equality_delete_file: k, v:{k, v}")
-    for partition_value, equality_delete_file_list in equality_delete_dict.items():
-        (
-            result_equality_delete_file,
-            result_data_file,
-        ) = check_data_files_sequence_number(
-            data_files_list=data_file_dict[partition_value],
-            equality_delete_files_list=equality_delete_dict[partition_value],
-        )
-        logger.info(f"result_data_file:{result_data_file}")
-        logger.info(f"result_equality_delete_file:{result_equality_delete_file}")
-        files_for_each_bucket[partition_value] = (
-            result_data_file,
-            result_equality_delete_file,
-            [],
-        )
-
+    convert_input_files_for_all_buckets = group_all_files_to_each_bucket(
+        data_file_dict=data_file_dict,
+        equality_delete_dict=equality_delete_dict,
+        pos_delete_dict=pos_delete_dict,
+    )
     iceberg_warehouse_bucket_name = params.iceberg_warehouse_bucket_name
     iceberg_namespace = params.iceberg_namespace
     iceberg_table_warehouse_prefix = construct_iceberg_table_prefix(
@@ -116,6 +101,7 @@ def convert_input_provider(index, item):
                 iceberg_table_warehouse_prefix=iceberg_table_warehouse_prefix,
                 identifier_fields=identifier_fields,
                 compact_small_files=compact_small_files,
+                enforce_primary_key_uniqueness=enforce_primary_key_uniqueness,
                 position_delete_for_multiple_data_files=position_delete_for_multiple_data_files,
                 max_parallel_data_file_download=max_parallel_data_file_download,
             )
@@ -125,7 +111,7 @@ def convert_input_provider(index, item):
     # Assuming that memory consume by each bucket doesn't exceed one node's memory limit.
     # TODO: Add split mechanism to split large buckets
     convert_tasks_pending = invoke_parallel(
-        items=files_for_each_bucket.items(),
+        items=convert_input_files_for_all_buckets.items(),
         ray_task=convert,
         max_parallelism=task_max_parallelism,
         options_provider=convert_options_provider,
@@ -143,9 +129,16 @@ def convert_input_provider(index, item):
         table_metadata=iceberg_table.metadata,
         files_dict_list=to_be_added_files_dict_list,
     )
-    commit_overwrite_snapshot(
-        iceberg_table=iceberg_table,
-        # equality_delete_files + data file that all rows are deleted
-        to_be_deleted_files_list=to_be_deleted_files_list[0],
-        new_position_delete_files=new_position_delete_files,
-    )
+    print(f"debug_to_be_deleted_files:{to_be_deleted_files_list}")
+    if not to_be_deleted_files_list:
+        commit_append_snapshot(
+            iceberg_table=iceberg_table,
+            new_position_delete_files=new_position_delete_files,
+        )
+    else:
+        commit_overwrite_snapshot(
+            iceberg_table=iceberg_table,
+            # equality_delete_files + data file that all rows are deleted
+            to_be_deleted_files_list=to_be_deleted_files_list,
+            new_position_delete_files=new_position_delete_files,
+        )

deltacat/compute/converter/model/convert_input.py

@@ -1,26 +1,29 @@
 from __future__ import annotations
 from typing import Dict, List
+from deltacat.compute.converter.model.convert_input_files import ConvertInputFiles
 
 
 class ConvertInput(Dict):
     @staticmethod
     def of(
-        files_for_each_bucket,
+        convert_input_files,
         convert_task_index,
         iceberg_table_warehouse_prefix,
         identifier_fields,
         compact_small_files,
+        enforce_primary_key_uniqueness,
         position_delete_for_multiple_data_files,
         max_parallel_data_file_download,
         s3_file_system,
     ) -> ConvertInput:
 
         result = ConvertInput()
-        result["files_for_each_bucket"] = files_for_each_bucket
+        result["convert_input_files"] = convert_input_files
         result["convert_task_index"] = convert_task_index
         result["identifier_fields"] = identifier_fields
         result["iceberg_table_warehouse_prefix"] = iceberg_table_warehouse_prefix
         result["compact_small_files"] = compact_small_files
+        result["enforce_primary_key_uniqueness"] = enforce_primary_key_uniqueness
         result[
             "position_delete_for_multiple_data_files"
         ] = position_delete_for_multiple_data_files
@@ -30,8 +33,8 @@ def of(
         return result
 
     @property
-    def files_for_each_bucket(self) -> tuple:
-        return self["files_for_each_bucket"]
+    def convert_input_files(self) -> ConvertInputFiles:
+        return self["convert_input_files"]
 
     @property
     def identifier_fields(self) -> List[str]:
@@ -49,6 +52,10 @@ def iceberg_table_warehouse_prefix(self) -> str:
     def compact_small_files(self) -> bool:
         return self["compact_small_files"]
 
+    @property
+    def enforce_primary_key_uniqueness(self) -> bool:
+        return self["enforce_primary_key_uniqueness"]
+
     @property
     def position_delete_for_multiple_data_files(self) -> bool:
         return self["position_delete_for_multiple_data_files"]

deltacat/compute/converter/model/convert_input_files.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+from typing import Dict
+
+
+class ConvertInputFiles(Dict):
+    @staticmethod
+    def of(
+        partition_value,
+        all_data_files_for_dedupe=None,
+        applicable_data_files=None,
+        applicable_equality_delete_files=None,
+        existing_position_delete_files=None,
+    ) -> ConvertInputFiles:
+
+        result = ConvertInputFiles()
+        result["partition_value"] = partition_value
+        result["all_data_files_for_dedupe"] = all_data_files_for_dedupe
+        result["applicable_data_files"] = applicable_data_files
+        result["applicable_equality_delete_files"] = applicable_equality_delete_files
+        result["existing_position_delete_files"] = existing_position_delete_files
+        return result
+
+    @property
+    def partition_value(self):
+        return self["partition_value"]
+
+    @property
+    def all_data_files_for_dedupe(self):
+        return self["all_data_files_for_dedupe"]
+
+    @property
+    def applicable_data_files(self):
+        return self["applicable_data_files"]
+
+    @property
+    def applicable_equality_delete_files(self):
+        return self["applicable_equality_delete_files"]
+
+    @property
+    def existing_position_delete_files(self):
+        return self["existing_position_delete_files"]
+
+    @partition_value.setter
+    def partition_value(self, partition_value):
+        self["partition_value"] = partition_value
+
+    @all_data_files_for_dedupe.setter
+    def all_data_files_for_dedupe(self, all_data_files_for_dedupe):
+        self["all_data_files_for_dedupe"] = all_data_files_for_dedupe
+
+    @applicable_data_files.setter
+    def applicable_data_files(self, applicable_data_files):
+        self["applicable_data_files"] = applicable_data_files
+
+    @applicable_equality_delete_files.setter
+    def applicable_equality_delete_files(self, applicable_equality_delete_files):
+        self["applicable_equality_delete_files"] = applicable_equality_delete_files
+
+    @existing_position_delete_files.setter
+    def existing_position_delete_files(self, existing_position_delete_files):
+        self["existing_position_delete_files"] = existing_position_delete_files

deltacat/compute/converter/model/converter_session_params.py

@@ -23,6 +23,9 @@ def of(params: Optional[Dict]) -> ConverterSessionParams:
         ), "iceberg_namespace is a required arg"
         result = ConverterSessionParams(params)
 
+        result.enforce_primary_key_uniqueness = params.get(
+            "enforce_primary_key_uniqueness", False
+        )
         result.compact_small_files = params.get("compact_small_files", False)
 
         # For Iceberg v3 spec, option to produce delete vector that can establish 1:1 mapping with data files.
@@ -51,6 +54,14 @@ def iceberg_warehouse_bucket_name(self) -> str:
     def iceberg_namespace(self) -> str:
         return self["iceberg_namespace"]
 
+    @property
+    def enforce_primary_key_uniqueness(self) -> bool:
+        return self["enforce_primary_key_uniqueness"]
+
+    @enforce_primary_key_uniqueness.setter
+    def enforce_primary_key_uniqueness(self, enforce_primary_key_uniqueness) -> None:
+        self["compact_small_files"] = enforce_primary_key_uniqueness
+
     @property
     def compact_small_files(self) -> bool:
         return self["compact_small_files"]

deltacat/compute/converter/pyiceberg/replace_snapshot.py

@@ -159,16 +159,41 @@ def commit_overwrite_snapshot(
 ):
     commit_uuid = uuid.uuid4()
     with iceberg_table.transaction() as tx:
+        print(
+            f"iceberg_table.metadata.name_mapping:{iceberg_table.metadata.name_mapping()}"
+        )
         if iceberg_table.metadata.name_mapping() is None:
-            iceberg_table.set_properties(
+            tx.set_properties(
                 **{
-                    "schema.name-mapping.default": iceberg_table.table_metadata.schema().name_mapping.model_dump_json()
+                    "schema.name-mapping.default": tx.table_metadata.schema().name_mapping.model_dump_json()
                 }
             )
         with tx.update_snapshot().overwrite(
             commit_uuid=commit_uuid
         ) as overwrite_snapshot:
-            for data_file in new_position_delete_files:
-                overwrite_snapshot.append_data_file(data_file)
-            for original_data_file in to_be_deleted_files_list:
-                overwrite_snapshot.delete_data_file(original_data_file)
+            if new_position_delete_files:
+                for data_file in new_position_delete_files:
+                    overwrite_snapshot.append_data_file(data_file)
+            if to_be_deleted_files_list:
+                for original_data_file in to_be_deleted_files_list:
+                    overwrite_snapshot.delete_data_file(original_data_file)
+
+
+def commit_append_snapshot(iceberg_table, new_position_delete_files):
+    commit_uuid = uuid.uuid4()
+    with iceberg_table.transaction() as tx:
+        print(
+            f"iceberg_table.metadata.name_mapping:{iceberg_table.metadata.name_mapping()}"
+        )
+        if iceberg_table.metadata.name_mapping() is None:
+            tx.set_properties(
+                **{
+                    "schema.name-mapping.default": tx.table_metadata.schema().name_mapping.model_dump_json()
+                }
+            )
+        with tx.update_snapshot().fast_append(
+            commit_uuid=commit_uuid
+        ) as append_snapshot:
+            if new_position_delete_files:
+                for data_file in new_position_delete_files:
+                    append_snapshot.append_data_file(data_file)