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

deltacat/compute/compactor/model/compaction_session_audit_info.py

@@ -1,5 +1,6 @@
 # Allow classes to use self-referencing Type hints in Python 3.7.
 from __future__ import annotations
+import pyarrow as pa
 import logging
 from deltacat import logs
 from typing import List, Union
@@ -18,6 +19,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 +54,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 +70,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 +144,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 +175,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 +251,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 +377,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 +391,42 @@ 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")
+
+    @property
+    def pyarrow_version(self) -> str:
+        """
+        The version of PyArrow used.
+        """
+        return self.get("pyarrowVersion")
+
     # Setters follow
 
     def set_audit_url(self, audit_url: str) -> CompactionSessionAuditInfo:
@@ -428,6 +511,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 +541,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 +607,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 +711,42 @@ 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
+
+    def set_pyarrow_version(self, value: str) -> CompactionSessionAuditInfo:
+        self["pyarrowVersion"] = value
+        return self
+
     # High level methods to save stats
     def save_step_stats(
         self,
@@ -673,7 +816,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 +843,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,11 +867,13 @@ 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,
                 ]
             )
         )
 
+        self.set_pyarrow_version(pa.__version__)
         self.set_telemetry_time_in_seconds(total_telemetry_time)

deltacat/compute/compactor/model/delta_annotated.py

@@ -2,7 +2,9 @@
 from __future__ import annotations
 
 import logging
-from types import FunctionType
+import copy
+from deltacat.types.media import ContentType, ContentEncoding
+from deltacat.types.partial_download import PartialParquetParameters
 from typing import Callable, List, Optional, Union
 
 from deltacat import logs
@@ -64,7 +66,9 @@ class DeltaAnnotated(Delta):
         annotated_deltas: List[DeltaAnnotated],
         min_delta_bytes: float,
         min_file_counts: Optional[Union[int, float]] = float("inf"),
-        estimation_function: Optional[Callable] = None,
+        estimation_function: Optional[
+            Callable[[ManifestEntry], float]
+        ] = lambda entry: entry.meta.content_length,
     ) -> List[DeltaAnnotated]:
         """
         Simple greedy algorithm to split/merge 1 or more annotated deltas into
@@ -76,11 +80,21 @@ class DeltaAnnotated(Delta):
         of bytes at rest for the associated object. Returns the list of annotated
         delta groups.
         """
-        groups = []
+        split_annotated_deltas: List[DeltaAnnotated] = []
+        groups: List[DeltaAnnotated] = []
         new_da = DeltaAnnotated()
         new_da_bytes = 0
         da_group_entry_count = 0
-        for src_da in annotated_deltas:
+
+        for delta_annotated in annotated_deltas:
+            split_annotated_deltas.extend(DeltaAnnotated._split_single(delta_annotated))
+
+        logger.info(
+            f"Split the {len(annotated_deltas)} annotated deltas "
+            f"into {len(split_annotated_deltas)} groups."
+        )
+
+        for src_da in split_annotated_deltas:
             src_da_annotations = src_da.annotations
             src_da_entries = src_da.manifest.entries
             assert (
@@ -105,11 +119,7 @@ class DeltaAnnotated(Delta):
                     src_da, new_da, src_entry, src_da_annotations[i]
                 )
                 # TODO: Fetch s3_obj["Size"] if entry content length undefined?
-                estimated_new_da_bytes = (
-                    estimation_function(src_entry.meta.content_length)
-                    if type(estimation_function) is FunctionType
-                    else src_entry.meta.content_length
-                )
+                estimated_new_da_bytes = estimation_function(src_entry)
                 new_da_bytes += estimated_new_da_bytes
                 da_group_entry_count += 1
                 if (
@@ -132,6 +142,7 @@ class DeltaAnnotated(Delta):
                     da_group_entry_count = 0
         if new_da:
             groups.append(new_da)
+
         return groups
 
     @staticmethod
@@ -207,3 +218,78 @@ class DeltaAnnotated(Delta):
                 dst_da.type = None
             entries.append(src_entry)
             dst_da.annotations.append(src_annotation)
+
+    @staticmethod
+    def _split_single(delta_annotated: DeltaAnnotated) -> List[DeltaAnnotated]:
+        """
+        Split a single delta annotated into multiple granular
+        annotated entries. Note that split is not always guaranteed.
+
+        Note: Currently we are only able to split the Parquet File downloads.
+        """
+
+        result = []
+
+        if (
+            delta_annotated.meta
+            and delta_annotated.manifest
+            and delta_annotated.meta.content_type == ContentType.PARQUET
+            and delta_annotated.meta.content_encoding == ContentEncoding.IDENTITY
+        ):
+            # we split by row groups
+            for entry_index, entry in enumerate(delta_annotated.manifest.entries):
+                input_split_params = None
+                if entry.meta and entry.meta.content_type_parameters:
+                    for type_params in entry.meta.content_type_parameters:
+                        if (
+                            isinstance(type_params, PartialParquetParameters)
+                            and type_params.num_row_groups > 1
+                            and type_params.pq_metadata
+                        ):
+                            input_split_params = type_params
+                            break
+
+                if input_split_params:
+                    logger.info(
+                        f"Splitting input file with URI: {entry.uri} into "
+                        f"different {input_split_params.num_row_groups} entries"
+                    )
+
+                    for rg in input_split_params.row_groups_to_download:
+                        new_da = DeltaAnnotated()
+                        new_entry_dict = copy.deepcopy(entry)
+                        new_entry = ManifestEntry(new_entry_dict)
+
+                        row_group_meta = input_split_params.pq_metadata.row_group(rg)
+
+                        new_partial_params = PartialParquetParameters.of(
+                            row_groups_to_download=[rg],
+                            num_row_groups=1,
+                            num_rows=row_group_meta.num_rows,
+                            in_memory_size_bytes=row_group_meta.total_byte_size,
+                            pq_metadata=input_split_params.pq_metadata,
+                        )
+
+                        new_entry.meta.content_type_parameters = [new_partial_params]
+                        for type_params in entry.meta.content_type_parameters:
+                            if not isinstance(type_params, PartialParquetParameters):
+                                new_entry.meta.content_type_parameters.append(
+                                    type_params
+                                )
+
+                        DeltaAnnotated._append_annotated_entry(
+                            delta_annotated,
+                            new_da,
+                            new_entry,
+                            delta_annotated.annotations[entry_index],
+                        )
+
+                        result.append(new_da)
+                else:
+                    return [delta_annotated]
+
+        logger.info(
+            f"Split was not performed on the delta with locator: {delta_annotated.locator}"
+        )
+
+        return [delta_annotated]

deltacat/compute/compactor/model/delta_file_envelope.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 import numpy as np
+import pyarrow as pa
 
 from deltacat.storage import DeltaType, LocalTable
 
@@ -37,8 +38,6 @@ class DeltaFileEnvelope(dict):
         """
         if stream_position is None:
             raise ValueError("Missing delta file envelope stream position.")
-        if file_index is None:
-            raise ValueError("Missing delta file envelope file index.")
         if delta_type is None:
             raise ValueError("Missing Delta file envelope delta type.")
         if table is None:
@@ -75,3 +74,16 @@ class DeltaFileEnvelope(dict):
     @property
     def file_record_count(self) -> int:
         return self["file_record_count"]
+
+    @property
+    def table_size_bytes(self) -> int:
+        if isinstance(self.table, pa.Table):
+            return self.table.nbytes
+        else:
+            raise ValueError(
+                f"Table type: {type(self.table)} not for supported for size method."
+            )
+
+    @property
+    def table_num_rows(self) -> int:
+        return len(self.table)

deltacat/compute/compactor/model/round_completion_info.py

@@ -1,6 +1,7 @@
 # Allow classes to use self-referencing Type hints in Python 3.7.
 from __future__ import annotations
 
+from typing import Tuple
 from deltacat.storage import DeltaLocator, PartitionLocator
 from deltacat.compute.compactor.model.pyarrow_write_result import PyArrowWriteResult
 from deltacat.compute.compactor.model.compaction_session_audit_info import (
@@ -40,9 +41,11 @@ class RoundCompletionInfo(dict):
         compacted_delta_locator: DeltaLocator,
         compacted_pyarrow_write_result: PyArrowWriteResult,
         sort_keys_bit_width: int,
-        rebase_source_partition_locator: Optional[PartitionLocator],
+        rebase_source_partition_locator: Optional[PartitionLocator] = None,
         manifest_entry_copied_by_reference_ratio: Optional[float] = None,
         compaction_audit_url: Optional[str] = None,
+        hash_bucket_count: Optional[int] = None,
+        hb_index_to_entry_range: Optional[Dict[int, Tuple[int, int]]] = None,
     ) -> RoundCompletionInfo:
 
         rci = RoundCompletionInfo()
@@ -55,6 +58,8 @@ class RoundCompletionInfo(dict):
             "manifestEntryCopiedByReferenceRatio"
         ] = manifest_entry_copied_by_reference_ratio
         rci["compactionAuditUrl"] = compaction_audit_url
+        rci["hashBucketCount"] = hash_bucket_count
+        rci["hbIndexToEntryRange"] = hb_index_to_entry_range
         return rci
 
     @property
@@ -97,3 +102,14 @@ class RoundCompletionInfo(dict):
     @property
     def manifest_entry_copied_by_reference_ratio(self) -> Optional[float]:
         return self["manifestEntryCopiedByReferenceRatio"]
+
+    @property
+    def hash_bucket_count(self) -> Optional[int]:
+        return self["hashBucketCount"]
+
+    @property
+    def hb_index_to_entry_range(self) -> Optional[Dict[int, Tuple[int, int]]]:
+        """
+        The start index is inclusive and end index is exclusive by default.
+        """
+        return self["hbIndexToEntryRange"]

deltacat/compute/compactor/repartition_session.py

@@ -54,6 +54,7 @@ def repartition(
     pg_config: Optional[PlacementGroupConfig] = None,
     list_deltas_kwargs: Optional[Dict[str, Any]] = None,
     read_kwargs_provider: Optional[ReadKwargsProvider] = None,
+    s3_table_writer_kwargs: Optional[Dict[str, Any]] = None,
     s3_client_kwargs: Optional[Dict[str, Any]] = None,
     deltacat_storage=unimplemented_deltacat_storage,
     **kwargs,
@@ -91,7 +92,7 @@ def repartition(
             source_partition_locator.partition_values,
         ).stream_position,
         deltacat_storage,
-        **list_deltas_kwargs,
+        list_deltas_kwargs,
     )
 
     uniform_deltas = []
@@ -131,6 +132,7 @@ def repartition(
         enable_profiler=enable_profiler,
         metrics_config=metrics_config,
         read_kwargs_provider=read_kwargs_provider,
+        s3_table_writer_kwargs=s3_table_writer_kwargs,
         repartitioned_file_content_type=repartitioned_file_content_type,
         deltacat_storage=deltacat_storage,
     )
@@ -162,6 +164,7 @@ def repartition(
         source_partition_locator,
         sort_keys,
         deltacat_storage,
+        deltacat_storage_kwargs={},
     )
     repartition_completion_info = RoundCompletionInfo.of(
         last_stream_position_to_compact,

deltacat/compute/compactor/steps/dedupe.py

@@ -107,20 +107,21 @@ def _timed_dedupe(
     dedupe_task_index: int,
     enable_profiler: bool,
     object_store: Optional[IObjectStore],
+    **kwargs,
 ):
     task_id = get_current_ray_task_id()
     worker_id = get_current_ray_worker_id()
     with memray.Tracker(
         f"dedupe_{worker_id}_{task_id}.bin"
     ) if enable_profiler else nullcontext():
-        # TODO (pdames): mitigate risk of running out of memory here in cases of
-        #  severe skew of primary key updates in deltas
+        # TODO (pdames): mitigate risk of running out of memory here in cases of severe skew of primary key updates in deltas
         logger.info(
             f"[Dedupe task {dedupe_task_index}] Getting delta file envelope "
             f"groups for {len(object_ids)} object refs..."
         )
-
-        delta_file_envelope_groups_list = object_store.get_many(object_ids)
+        delta_file_envelope_groups_list: List[object] = object_store.get_many(
+            object_ids
+        )
         hb_index_to_delta_file_envelopes_list = defaultdict(list)
         for delta_file_envelope_groups in delta_file_envelope_groups_list:
             for hb_idx, dfes in enumerate(delta_file_envelope_groups):
@@ -171,7 +172,8 @@ def _timed_dedupe(
 
             hb_table_record_count = len(table)
             table, drop_time = timed_invocation(
-                func=_drop_duplicates_by_primary_key_hash, table=table
+                func=_drop_duplicates_by_primary_key_hash,
+                table=table,
             )
             deduped_record_count = hb_table_record_count - len(table)
             total_deduped_records += deduped_record_count
@@ -227,7 +229,6 @@ def _timed_dedupe(
         )
 
         peak_memory_usage_bytes = get_current_node_peak_memory_usage_in_bytes()
-
         return DedupeResult(
             mat_bucket_to_dd_idx_obj_id,
             np.int64(total_deduped_records),
@@ -246,6 +247,7 @@ def dedupe(
     enable_profiler: bool,
     metrics_config: MetricsConfig,
     object_store: Optional[IObjectStore],
+    **kwargs,
 ) -> DedupeResult:
     logger.info(f"[Dedupe task {dedupe_task_index}] Starting dedupe task...")
     dedupe_result, duration = timed_invocation(
@@ -256,6 +258,7 @@ def dedupe(
         dedupe_task_index=dedupe_task_index,
         enable_profiler=enable_profiler,
         object_store=object_store,
+        **kwargs,
     )
 
     emit_metrics_time = 0.0