deltacat 0.1.18b14__py3-none-any.whl → 0.1.18b15__py3-none-any.whl

deltacat/compute/compactor/model/compact_partition_params.py

@@ -1,10 +1,28 @@
 from __future__ import annotations
-
+import importlib
 import copy
 import json
 from typing import Any, Dict, List, Optional
-
+from deltacat.io.object_store import IObjectStore
+from deltacat.utils.common import ReadKwargsProvider
 from deltacat.types.media import ContentType
+from deltacat.utils.placement import PlacementGroupConfig
+from deltacat.io.ray_plasma_object_store import RayPlasmaObjectStore
+from deltacat.storage import (
+    interface as unimplemented_deltacat_storage,
+    PartitionLocator,
+    SortKey,
+)
+from deltacat.compute.compactor_v2.constants import (
+    MAX_RECORDS_PER_COMPACTED_FILE,
+    MIN_DELTA_BYTES_IN_BATCH,
+    MIN_FILES_IN_BATCH,
+    AVERAGE_RECORD_SIZE_BYTES,
+    TASK_MAX_PARALLELISM,
+)
+from deltacat.constants import PYARROW_INFLATION_MULTIPLIER
+from deltacat.compute.compactor.utils.sort_key import validate_sort_keys
+from deltacat.utils.metrics import MetricsConfig
 
 
 class CompactPartitionParams(dict):
@@ -16,104 +34,315 @@ class CompactPartitionParams(dict):
     def of(params: Optional[Dict]) -> CompactPartitionParams:
         if params is None:
             params = {}
-        compact_partition_params = CompactPartitionParams()
-        compact_partition_params["destination_partition_locator"] = params.get(
-            "destination_partition_locator"
-        )
-        compact_partition_params["last_stream_position_to_compact"] = params.get(
-            "last_stream_position_to_compact"
+
+        assert (
+            params.get("destination_partition_locator") is not None
+        ), "destination_partition_locator is a required arg"
+        assert (
+            params.get("last_stream_position_to_compact") is not None
+        ), "last_stream_position_to_compact is a required arg"
+        assert (
+            params.get("source_partition_locator") is not None
+        ), "source_partition_locator is a required arg"
+        assert (
+            params.get("compaction_artifact_s3_bucket") is not None
+        ), "compaction_artifact_s3_bucket is a required arg"
+
+        result = CompactPartitionParams(params)
+
+        # TODO: move defaults to single file
+        result.records_per_compacted_file = params.get(
+            "records_per_compacted_file", MAX_RECORDS_PER_COMPACTED_FILE
         )
-        compact_partition_params["source_partition_locator"] = params.get(
-            "source_partition_locator"
+        result.compacted_file_content_type = params.get(
+            "compacted_file_content_type", ContentType.PARQUET
         )
-        compact_partition_params["primary_keys"] = params.get("primary_keys")
-        compact_partition_params["rebase_source_partition_locator"] = params.get(
-            "rebase_source_partition_locator"
+        result.object_store = params.get("object_store", RayPlasmaObjectStore())
+        result.enable_profiler = params.get("enable_profiler", False)
+        result.deltacat_storage = params.get(
+            "deltacat_storage", unimplemented_deltacat_storage
         )
-        compact_partition_params["rebase_source_partition_high_watermark"] = params.get(
-            "rebase_source_partition_high_watermark"
+        result.s3_client_kwargs = params.get("s3_client_kwargs", {})
+        result.deltacat_storage_kwargs = params.get("deltacat_storage_kwargs", {})
+        result.list_deltas_kwargs = params.get("list_deltas_kwargs", {})
+        result.s3_table_writer_kwargs = params.get("s3_table_writer_kwargs", {})
+        result.bit_width_of_sort_keys = validate_sort_keys(
+            result.source_partition_locator,
+            result.sort_keys,
+            result.deltacat_storage,
+            result.deltacat_storage_kwargs,
         )
-        compact_partition_params["hash_bucket_count"] = params.get("hash_bucket_count")
-        compact_partition_params["deltacat_storage"] = params.get("deltacat_storage")
-        compact_partition_params["compaction_artifact_s3_bucket"] = params.get(
-            "compaction_artifact_s3_bucket"
+        result.task_max_parallelism = params.get(
+            "task_max_parallelism", TASK_MAX_PARALLELISM
         )
-        compact_partition_params["properties"] = params.get("properties")
-        compact_partition_params["compacted_file_content_type"] = params.get(
-            "compacted_file_content_type"
+        result.min_files_in_batch = params.get("min_files_in_batch", MIN_FILES_IN_BATCH)
+        result.min_delta_bytes_in_batch = params.get(
+            "min_delta_bytes_in_batch", MIN_DELTA_BYTES_IN_BATCH
         )
-        compact_partition_params["list_deltas_kwargs"] = params.get(
-            "list_deltas_kwargs"
+        result.previous_inflation = params.get(
+            "previous_inflation", PYARROW_INFLATION_MULTIPLIER
         )
-        compact_partition_params["pg_config"] = params.get("pg_config")
-        compact_partition_params["read_kwargs_provider"] = params.get(
-            "read_kwargs_provider"
+        result.average_record_size_bytes = params.get(
+            "average_record_size_bytes", AVERAGE_RECORD_SIZE_BYTES
         )
-        compact_partition_params["s3_table_writer_kwargs"] = params.get(
-            "s3_table_writer_kwargs"
+        result.hash_group_count = params.get(
+            "hash_group_count", result.hash_bucket_count
         )
-        return compact_partition_params
+
+        if not importlib.util.find_spec("memray"):
+            result.enable_profiler = False
+
+        if result.primary_keys:
+            result.primary_keys = sorted(result.primary_keys)
+
+        # assertions
+        assert (
+            result.source_partition_locator.partition_values
+            == result.destination_partition_locator.partition_values
+        ), "Source and destination partitions values must be equal"
+
+        assert (
+            result.records_per_compacted_file and result.records_per_compacted_file >= 1
+        ), "Max records per output file must be a positive value"
+
+        return result
 
     @property
-    def destination_partition_locator(self) -> Optional[dict]:
-        return self["destination_partition_locator"]
+    def destination_partition_locator(self) -> PartitionLocator:
+        val = self["destination_partition_locator"]
+        if not isinstance(val, PartitionLocator):
+            val = PartitionLocator(val)
+
+        return val
+
+    @destination_partition_locator.setter
+    def destination_partition_locator(self, locator: PartitionLocator) -> None:
+        self["destination_partition_locator"] = locator
 
     @property
-    def last_stream_position_to_compact(self) -> Optional[int]:
+    def last_stream_position_to_compact(self) -> int:
         return self["last_stream_position_to_compact"]
 
+    @last_stream_position_to_compact.setter
+    def last_stream_position_to_compact(self, stream_position: int) -> None:
+        self["last_stream_position_to_compact"] = stream_position
+
     @property
-    def source_partition_locator(self) -> Optional[dict]:
-        return self["source_partition_locator"]
+    def source_partition_locator(self) -> PartitionLocator:
+        val = self["source_partition_locator"]
+        if not isinstance(val, PartitionLocator):
+            val = PartitionLocator(val)
+        return val
+
+    @source_partition_locator.setter
+    def source_partition_locator(self, locator: PartitionLocator) -> None:
+        self["source_partition_locator"] = locator
 
     @property
-    def primary_keys(self) -> Optional[List[str]]:
-        return list(self["primary_keys"])
+    def compaction_artifact_s3_bucket(self) -> str:
+        return self["compaction_artifact_s3_bucket"]
+
+    @compaction_artifact_s3_bucket.setter
+    def compaction_artifact_s3_bucket(self, s3_bucket: str) -> None:
+        self["compaction_artifact_s3_bucket"] = s3_bucket
 
     @property
-    def rebase_source_partition_locator(self) -> Optional[dict]:
-        return self["rebase_source_partition_locator"]
+    def deltacat_storage(self) -> unimplemented_deltacat_storage:
+        return self["deltacat_storage"]
+
+    @deltacat_storage.setter
+    def deltacat_storage(self, storage: unimplemented_deltacat_storage) -> None:
+        self["deltacat_storage"] = storage
 
     @property
-    def rebase_source_partition_high_watermark(self) -> Optional[int]:
-        return self["rebase_source_partition_high_watermark"]
+    def object_store(self) -> IObjectStore:
+        return self["object_store"]
+
+    @object_store.setter
+    def object_store(self, obj_store: IObjectStore) -> None:
+        self["object_store"] = obj_store
 
     @property
-    def hash_bucket_count(self) -> Optional[int]:
-        return self["hash_bucket_count"]
+    def compacted_file_content_type(self) -> ContentType:
+        return self["compacted_file_content_type"]
+
+    @compacted_file_content_type.setter
+    def compacted_file_content_type(self, content_type: ContentType) -> None:
+        self["compacted_file_content_type"] = content_type
 
     @property
-    def deltacat_storage(self) -> Optional[str]:
-        return self["deltacat_storage"]
+    def task_max_parallelism(self) -> int:
+        return self["task_max_parallelism"]
+
+    @task_max_parallelism.setter
+    def task_max_parallelism(self, max_parallelism: int) -> None:
+        self["task_max_parallelism"] = max_parallelism
 
     @property
-    def compaction_artifact_s3_bucket(self) -> Optional[str]:
-        return self["compaction_artifact_s3_bucket"]
+    def average_record_size_bytes(self) -> float:
+        return self["average_record_size_bytes"]
+
+    @average_record_size_bytes.setter
+    def average_record_size_bytes(self, average_record_size_bytes: float) -> None:
+        self["average_record_size_bytes"] = average_record_size_bytes
 
     @property
-    def properties(self) -> Optional[Dict[str, str]]:
-        return self["properties"]
+    def min_files_in_batch(self) -> float:
+        return self["min_files_in_batch"]
+
+    @min_files_in_batch.setter
+    def min_files_in_batch(self, min_files_in_batch: float) -> None:
+        self["min_files_in_batch"] = min_files_in_batch
 
     @property
-    def compacted_file_content_type(self) -> Optional[ContentType]:
-        return self["compacted_file_content_type"]
+    def min_delta_bytes_in_batch(self) -> float:
+        return self["min_files_in_batch"]
+
+    @min_delta_bytes_in_batch.setter
+    def min_delta_bytes_in_batch(self, min_delta_bytes_in_batch: float) -> None:
+        self["min_delta_bytes_in_batch"] = min_delta_bytes_in_batch
 
     @property
-    def list_deltas_kwargs(self) -> Optional[dict]:
-        return self["list_deltas_kwargs"]
+    def previous_inflation(self) -> float:
+        return self["previous_inflation"]
+
+    @previous_inflation.setter
+    def previous_inflation(self, previous_inflation: float) -> None:
+        self["previous_inflation"] = previous_inflation
 
     @property
-    def pg_config(self) -> Optional[Any]:
-        return self["pg_config"]
+    def enable_profiler(self) -> bool:
+        return self["enable_profiler"]
+
+    @enable_profiler.setter
+    def enable_profiler(self, value: bool) -> None:
+        self["enable_profiler"] = value
 
     @property
-    def read_kwargs_provider(self) -> Optional[Any]:
-        return self["read_kwargs_provider"]
+    def list_deltas_kwargs(self) -> dict:
+        return self["list_deltas_kwargs"]
+
+    @list_deltas_kwargs.setter
+    def list_deltas_kwargs(self, kwargs: dict) -> None:
+        self["list_deltas_kwargs"] = kwargs
 
     @property
-    def s3_table_writer_kwargs(self) -> Optional[Any]:
+    def s3_table_writer_kwargs(self) -> dict:
         return self["s3_table_writer_kwargs"]
 
+    @s3_table_writer_kwargs.setter
+    def s3_table_writer_kwargs(self, kwargs: dict) -> None:
+        self["s3_table_writer_kwargs"] = kwargs
+
+    @property
+    def deltacat_storage_kwargs(self) -> dict:
+        return self["deltacat_storage_kwargs"]
+
+    @deltacat_storage_kwargs.setter
+    def deltacat_storage_kwargs(self, kwargs: dict) -> None:
+        self["deltacat_storage_kwargs"] = kwargs
+
+    @property
+    def s3_client_kwargs(self) -> dict:
+        return self["s3_client_kwargs"]
+
+    @s3_client_kwargs.setter
+    def s3_client_kwargs(self, kwargs: dict) -> None:
+        self["s3_client_kwargs"] = kwargs
+
+    @property
+    def records_per_compacted_file(self) -> int:
+        return self["records_per_compacted_file"]
+
+    @records_per_compacted_file.setter
+    def records_per_compacted_file(self, count: int) -> None:
+        self["records_per_compacted_file"] = count
+
+    @property
+    def bit_width_of_sort_keys(self) -> int:
+        return self["bit_width_of_sort_keys"]
+
+    @bit_width_of_sort_keys.setter
+    def bit_width_of_sort_keys(self, width: int) -> None:
+        self["bit_width_of_sort_keys"] = width
+
+    @property
+    def hash_bucket_count(self) -> Optional[int]:
+        return self.get("hash_bucket_count")
+
+    @hash_bucket_count.setter
+    def hash_bucket_count(self, count: int) -> None:
+        self["hash_bucket_count"] = count
+
+    @property
+    def hash_group_count(self) -> int:
+        return self["hash_group_count"]
+
+    @hash_group_count.setter
+    def hash_group_count(self, count: int) -> None:
+        self["hash_group_count"] = count
+
+    @property
+    def primary_keys(self) -> Optional[List[str]]:
+        return self.get("primary_keys")
+
+    @primary_keys.setter
+    def primary_keys(self, keys: List[str]) -> None:
+        self["primary_keys"] = keys
+
+    @property
+    def rebase_source_partition_locator(self) -> Optional[PartitionLocator]:
+        val = self.get("rebase_source_partition_locator")
+
+        if val and not isinstance(val, PartitionLocator):
+            val = PartitionLocator(val)
+
+        return val
+
+    @rebase_source_partition_locator.setter
+    def rebase_source_partition_locator(self, locator: PartitionLocator) -> None:
+        self["rebase_source_partition_locator"] = locator
+
+    @property
+    def rebase_source_partition_high_watermark(self) -> Optional[int]:
+        return self.get("rebase_source_partition_high_watermark")
+
+    @rebase_source_partition_high_watermark.setter
+    def rebase_source_partition_high_watermark(self, high_watermark: int) -> None:
+        self["rebase_source_partition_high_watermark"] = high_watermark
+
+    @property
+    def pg_config(self) -> Optional[PlacementGroupConfig]:
+        return self.get("pg_config")
+
+    @pg_config.setter
+    def pg_config(self, config: PlacementGroupConfig) -> None:
+        self["pg_config"] = config
+
+    @property
+    def read_kwargs_provider(self) -> Optional[ReadKwargsProvider]:
+        return self.get("read_kwargs_provider")
+
+    @read_kwargs_provider.setter
+    def read_kwargs_provider(self, kwargs_provider: ReadKwargsProvider) -> None:
+        self["read_kwargs_provider"] = kwargs_provider
+
+    @property
+    def sort_keys(self) -> Optional[List[SortKey]]:
+        return self.get("sort_keys")
+
+    @sort_keys.setter
+    def sort_keys(self, keys: List[SortKey]) -> None:
+        self["sort_keys"] = keys
+
+    @property
+    def metrics_config(self) -> Optional[MetricsConfig]:
+        return self.get("metrics_config")
+
+    @metrics_config.setter
+    def metrics_config(self, config: MetricsConfig) -> None:
+        self["metrics_config"] = config
+
     @staticmethod
     def json_handler_for_compact_partition_params(obj):
         """

deltacat/compute/compactor/model/compaction_session_audit_info.py

@@ -18,6 +18,7 @@ class CompactionSessionAuditInfo(dict):
     DEDUPE_STEP_NAME = "dedupe"
     MATERIALIZE_STEP_NAME = "materialize"
     HASH_BUCKET_STEP_NAME = "hashBucket"
+    MERGE_STEP_NAME = "merge"
 
     def __init__(self, deltacat_version: str, audit_url: str):
         self.set_deltacat_version(deltacat_version)
@@ -52,7 +53,7 @@ class CompactionSessionAuditInfo(dict):
     @property
     def uniform_deltas_created(self) -> int:
         """
-        The total number of unitform deltas fed into the hash bucket step.
+        The total number of uniform deltas fed into the hash bucket step.
         """
         return self.get("uniformDeltasCreated")
 
@@ -68,7 +69,7 @@ class CompactionSessionAuditInfo(dict):
     @property
     def input_size_bytes(self) -> float:
         """
-        The on-disk size in bytes of the input.
+        The on-disk size in bytes of the input. Analogous to bytes scanned
         """
         return self.get("inputSizeBytes")
 
@@ -142,6 +143,15 @@ class CompactionSessionAuditInfo(dict):
         """
         return self.get("materializeTaskPeakMemoryUsedBytes")
 
+    @property
+    def peak_memory_used_bytes_per_merge_task(self) -> float:
+        """
+        The peak memory used by a single merge python process. Note
+        that results may be max of merge, and hash bucketing as
+        processes are reused by Ray to run all compaction steps.
+        """
+        return self.get("mergeTaskPeakMemoryUsedBytes")
+
     @property
     def hash_bucket_post_object_store_memory_used_bytes(self) -> float:
         """
@@ -164,6 +174,13 @@ class CompactionSessionAuditInfo(dict):
         """
         return self.get("materializePostObjectStoreMemoryUsedBytes")
 
+    @property
+    def merge_post_object_store_memory_used_bytes(self) -> float:
+        """
+        The total object store memory used after merge step.
+        """
+        return self.get("mergePostObjectStoreMemoryUsedBytes")
+
     @property
     def materialize_buckets(self) -> int:
         """
@@ -233,11 +250,33 @@ class CompactionSessionAuditInfo(dict):
     @property
     def materialize_result_wait_time_in_seconds(self) -> float:
         """
-        The time it takes ray.get() to resolve after the last hash bucket task has completed.
+        The time it takes ray.get() to resolve after the last materialize task has completed.
         This value may not be accurate at less than 1 second precision.
         """
         return self.get("materializeResultWaitTimeInSeconds")
 
+    @property
+    def merge_result_wait_time_in_seconds(self) -> float:
+        """
+        The time it takes ray.get() to resolve after the last task has completed.
+        This value may not be accurate at less than 1 second precision.
+        """
+        return self.get("mergeResultWaitTimeInSeconds")
+
+    @property
+    def merge_time_in_seconds(self) -> float:
+        """
+        The time taken by merge step. This includes all merge tasks.
+        """
+        return self.get("mergeTimeInSeconds")
+
+    @property
+    def merge_invoke_time_in_seconds(self) -> float:
+        """
+        The time taken to invoke all merge tasks.
+        """
+        return self.get("mergeInvokeTimeInSeconds")
+
     @property
     def delta_discovery_time_in_seconds(self) -> float:
         """
@@ -337,6 +376,13 @@ class CompactionSessionAuditInfo(dict):
         """
         return self.get("materializeResultSize")
 
+    @property
+    def merge_result_size(self) -> float:
+        """
+        The size of the results returned by merge step.
+        """
+        return self.get("mergeResultSize")
+
     @property
     def peak_memory_used_bytes_by_compaction_session_process(self) -> float:
         """
@@ -344,6 +390,35 @@ class CompactionSessionAuditInfo(dict):
         """
         return self.get("peakMemoryUsedBytesCompactionSessionProcess")
 
+    @property
+    def estimated_in_memory_size_bytes_during_discovery(self) -> float:
+        """
+        The estimated in-memory size during the discovery. This can be used
+        to determine the accuracy of memory estimation logic.
+        """
+        return self.get("estimatedInMemorySizeBytesDuringDiscovery")
+
+    @property
+    def hash_bucket_processed_size_bytes(self) -> int:
+        """
+        The total size of the input data processed during hash bucket
+        """
+        return self.get("hashBucketProcessedSizeBytes")
+
+    @property
+    def total_cpu_seconds(self) -> float:
+        """
+        Total number of vCPUs provisioned in the cluster weighted over time.
+        """
+        return self.get("totalCPUSeconds")
+
+    @property
+    def used_cpu_seconds(self) -> float:
+        """
+        Total used vCPU in the cluster weighted over time.
+        """
+        return self.get("usedCPUSeconds")
+
     # Setters follow
 
     def set_audit_url(self, audit_url: str) -> CompactionSessionAuditInfo:
@@ -428,6 +503,12 @@ class CompactionSessionAuditInfo(dict):
         ] = peak_memory_used_bytes_per_materialize_task
         return self
 
+    def set_peak_memory_used_bytes_per_merge_task(
+        self, peak_memory_used_bytes: float
+    ) -> CompactionSessionAuditInfo:
+        self["mergeTaskPeakMemoryUsedBytes"] = peak_memory_used_bytes
+        return self
+
     def set_hash_bucket_post_object_store_memory_used_bytes(
         self, object_store_memory_used_bytes_by_hb: float
     ) -> CompactionSessionAuditInfo:
@@ -452,6 +533,12 @@ class CompactionSessionAuditInfo(dict):
         ] = object_store_memory_used_bytes_by_dedupe
         return self
 
+    def set_merge_post_object_store_memory_used_bytes(
+        self, object_store_memory_used_bytes: float
+    ) -> CompactionSessionAuditInfo:
+        self["mergePostObjectStoreMemoryUsedBytes"] = object_store_memory_used_bytes
+        return self
+
     def set_materialize_buckets(
         self, materialize_buckets: int
     ) -> CompactionSessionAuditInfo:
@@ -512,6 +599,24 @@ class CompactionSessionAuditInfo(dict):
         self.get["materializeResultWaitTimeInSeconds"] = wait_time
         return self
 
+    def set_merge_time_in_seconds(
+        self, time_in_seconds: float
+    ) -> CompactionSessionAuditInfo:
+        self["mergeTimeInSeconds"] = time_in_seconds
+        return self
+
+    def set_merge_invoke_time_in_seconds(
+        self, invoke_time: float
+    ) -> CompactionSessionAuditInfo:
+        self["mergeInvokeTimeInSeconds"] = invoke_time
+        return self
+
+    def set_merge_result_wait_time_in_seconds(
+        self, wait_time: float
+    ) -> CompactionSessionAuditInfo:
+        self.get["mergeResultWaitTimeInSeconds"] = wait_time
+        return self
+
     def set_delta_discovery_time_in_seconds(
         self, delta_discovery_time_in_seconds: float
     ) -> CompactionSessionAuditInfo:
@@ -598,12 +703,38 @@ class CompactionSessionAuditInfo(dict):
         self["materializeResultSize"] = materialize_result_size_bytes
         return self
 
+    def set_merge_result_size_bytes(
+        self, merge_result_size_bytes: float
+    ) -> CompactionSessionAuditInfo:
+        self["mergeResultSize"] = merge_result_size_bytes
+        return self
+
     def set_peak_memory_used_bytes_by_compaction_session_process(
         self, peak_memory: float
     ) -> CompactionSessionAuditInfo:
         self["peakMemoryUsedBytesCompactionSessionProcess"] = peak_memory
         return self
 
+    def set_estimated_in_memory_size_bytes_during_discovery(
+        self, memory: float
+    ) -> CompactionSessionAuditInfo:
+        self["estimatedInMemorySizeBytesDuringDiscovery"] = memory
+        return self
+
+    def set_hash_bucket_processed_size_bytes(
+        self, size: int
+    ) -> CompactionSessionAuditInfo:
+        self["hashBucketProcessedSizeBytes"] = size
+        return self
+
+    def set_total_cpu_seconds(self, value: float) -> CompactionSessionAuditInfo:
+        self["totalCPUSeconds"] = value
+        return self
+
+    def set_used_cpu_seconds(self, value: float) -> CompactionSessionAuditInfo:
+        self["usedCPUSeconds"] = value
+        return self
+
     # High level methods to save stats
     def save_step_stats(
         self,
@@ -673,7 +804,10 @@ class CompactionSessionAuditInfo(dict):
         )
 
         total_count_of_src_dfl_not_touched = sum(
-            m.referenced_pyarrow_write_result.files for m in mat_results
+            m.referenced_pyarrow_write_result.files
+            if m.referenced_pyarrow_write_result
+            else 0
+            for m in mat_results
         )
 
         logger.info(
@@ -697,10 +831,16 @@ class CompactionSessionAuditInfo(dict):
         )
 
         untouched_file_record_count = sum(
-            m.referenced_pyarrow_write_result.records for m in mat_results
+            m.referenced_pyarrow_write_result.records
+            if m.referenced_pyarrow_write_result
+            else 0
+            for m in mat_results
         )
         untouched_file_size_bytes = sum(
-            m.referenced_pyarrow_write_result.file_bytes for m in mat_results
+            m.referenced_pyarrow_write_result.file_bytes
+            if m.referenced_pyarrow_write_result
+            else 0
+            for m in mat_results
         )
 
         self.set_untouched_file_count(total_count_of_src_dfl_not_touched)
@@ -715,9 +855,10 @@ class CompactionSessionAuditInfo(dict):
         self.set_peak_memory_used_bytes_per_task(
             max(
                 [
-                    self.peak_memory_used_bytes_per_hash_bucket_task,
-                    self.peak_memory_used_bytes_per_dedupe_task,
-                    self.peak_memory_used_bytes_per_materialize_task,
+                    self.peak_memory_used_bytes_per_hash_bucket_task or 0,
+                    self.peak_memory_used_bytes_per_dedupe_task or 0,
+                    self.peak_memory_used_bytes_per_materialize_task or 0,
+                    self.peak_memory_used_bytes_per_merge_task or 0,
                 ]
             )
         )