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

deltacat/compute/compactor_v2/model/merge_input.py

@@ -0,0 +1,143 @@
+from __future__ import annotations
+
+from ray.types import ObjectRef
+from typing import Dict, List, Optional, Any
+from deltacat.utils.metrics import MetricsConfig
+from deltacat.utils.common import ReadKwargsProvider
+from deltacat.io.object_store import IObjectStore
+from deltacat.storage import (
+    Partition,
+    SortKey,
+    interface as unimplemented_deltacat_storage,
+)
+from deltacat.compute.compactor_v2.constants import (
+    DROP_DUPLICATES,
+    MAX_RECORDS_PER_COMPACTED_FILE,
+)
+from deltacat.types.media import ContentType
+from deltacat.compute.compactor.model.round_completion_info import RoundCompletionInfo
+from deltacat.compute.compactor.model.delta_file_envelope import DeltaFileEnvelopeGroups
+
+
+class MergeInput(Dict):
+    @staticmethod
+    def of(
+        dfe_groups_refs: List[ObjectRef[DeltaFileEnvelopeGroups]],
+        write_to_partition: Partition,
+        compacted_file_content_type: ContentType,
+        primary_keys: List[str],
+        hash_group_index: int,
+        num_hash_groups: int,
+        hash_bucket_count: int,
+        drop_duplicates: Optional[bool] = DROP_DUPLICATES,
+        sort_keys: Optional[List[SortKey]] = None,
+        merge_task_index: Optional[int] = 0,
+        max_records_per_output_file: Optional[int] = MAX_RECORDS_PER_COMPACTED_FILE,
+        enable_profiler: Optional[bool] = False,
+        metrics_config: Optional[MetricsConfig] = None,
+        s3_table_writer_kwargs: Optional[Dict[str, Any]] = None,
+        read_kwargs_provider: Optional[ReadKwargsProvider] = None,
+        round_completion_info: Optional[RoundCompletionInfo] = None,
+        object_store: Optional[IObjectStore] = None,
+        deltacat_storage=unimplemented_deltacat_storage,
+        deltacat_storage_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> MergeInput:
+
+        result = MergeInput()
+        result["dfe_groups_refs"] = dfe_groups_refs
+        result["write_to_partition"] = write_to_partition
+        result["compacted_file_content_type"] = compacted_file_content_type
+        result["primary_keys"] = primary_keys
+        result["hash_group_index"] = hash_group_index
+        result["num_hash_groups"] = num_hash_groups
+        result["hash_bucket_count"] = hash_bucket_count
+        result["drop_duplicates"] = drop_duplicates
+        result["sort_keys"] = sort_keys
+        result["merge_task_index"] = merge_task_index
+        result["max_records_per_output_file"] = max_records_per_output_file
+        result["enable_profiler"] = enable_profiler
+        result["metrics_config"] = metrics_config
+        result["s3_table_writer_kwargs"] = s3_table_writer_kwargs or {}
+        result["read_kwargs_provider"] = read_kwargs_provider
+        result["round_completion_info"] = round_completion_info
+        result["object_store"] = object_store
+        result["deltacat_storage"] = deltacat_storage
+        result["deltacat_storage_kwargs"] = deltacat_storage_kwargs or {}
+
+        return result
+
+    @property
+    def dfe_groups_refs(self) -> List[ObjectRef[DeltaFileEnvelopeGroups]]:
+        return self["dfe_groups_refs"]
+
+    @property
+    def write_to_partition(self) -> Partition:
+        return self["write_to_partition"]
+
+    @property
+    def compacted_file_content_type(self) -> ContentType:
+        return self["compacted_file_content_type"]
+
+    @property
+    def primary_keys(self) -> List[str]:
+        return self["primary_keys"]
+
+    @property
+    def hash_group_index(self) -> int:
+        return self["hash_group_index"]
+
+    @property
+    def num_hash_groups(self) -> int:
+        return self["num_hash_groups"]
+
+    @property
+    def hash_bucket_count(self) -> int:
+        return self["hash_bucket_count"]
+
+    @property
+    def drop_duplicates(self) -> int:
+        return self["drop_duplicates"]
+
+    @property
+    def sort_keys(self) -> Optional[List[SortKey]]:
+        return self.get("sort_keys")
+
+    @property
+    def merge_task_index(self) -> int:
+        return self.get("merge_task_index")
+
+    @property
+    def max_records_per_output_file(self) -> int:
+        return self.get("max_records_per_output_file")
+
+    @property
+    def enable_profiler(self) -> bool:
+        return self.get("enable_profiler")
+
+    @property
+    def metrics_config(self) -> Optional[MetricsConfig]:
+        return self.get("metrics_config")
+
+    @property
+    def s3_table_writer_kwargs(self) -> Optional[Dict[str, Any]]:
+        return self.get("s3_table_writer_kwargs")
+
+    @property
+    def read_kwargs_provider(self) -> Optional[ReadKwargsProvider]:
+        return self.get("read_kwargs_provider")
+
+    @property
+    def round_completion_info(self) -> Optional[RoundCompletionInfo]:
+        return self.get("round_completion_info")
+
+    @property
+    def object_store(self) -> Optional[IObjectStore]:
+        return self.get("object_store")
+
+    @property
+    def deltacat_storage(self) -> unimplemented_deltacat_storage:
+        return self["deltacat_storage"]
+
+    @property
+    def deltacat_storage_kwargs(self) -> Optional[Dict[str, Any]]:
+        return self.get("deltacat_storage_kwargs")

deltacat/compute/compactor_v2/model/merge_result.py

@@ -0,0 +1,12 @@
+from typing import NamedTuple, List
+from deltacat.compute.compactor.model.materialize_result import MaterializeResult
+
+import numpy as np
+
+
+class MergeResult(NamedTuple):
+    materialize_results: List[MaterializeResult]
+    deduped_record_count: np.int64
+    peak_memory_usage_bytes: np.double
+    telemetry_time_in_seconds: np.double
+    task_completed_at: np.double

deltacat/compute/compactor_v2/steps/hash_bucket.py

@@ -0,0 +1,203 @@
+import importlib
+import logging
+import time
+from contextlib import nullcontext
+from typing import List, Optional, Tuple
+from deltacat.compute.compactor_v2.model.hash_bucket_input import HashBucketInput
+import numpy as np
+import pyarrow as pa
+import ray
+from deltacat import logs
+from deltacat.compute.compactor import (
+    DeltaAnnotated,
+    DeltaFileEnvelope,
+)
+from deltacat.compute.compactor.model.delta_file_envelope import DeltaFileEnvelopeGroups
+from deltacat.compute.compactor_v2.model.hash_bucket_result import HashBucketResult
+from deltacat.compute.compactor_v2.utils.primary_key_index import (
+    group_hash_bucket_indices,
+    group_by_pk_hash_bucket,
+)
+from deltacat.storage import interface as unimplemented_deltacat_storage
+from deltacat.types.media import StorageType
+from deltacat.utils.ray_utils.runtime import (
+    get_current_ray_task_id,
+    get_current_ray_worker_id,
+)
+from deltacat.utils.common import ReadKwargsProvider
+from deltacat.utils.performance import timed_invocation
+from deltacat.utils.metrics import emit_timer_metrics
+from deltacat.utils.resources import get_current_node_peak_memory_usage_in_bytes
+
+if importlib.util.find_spec("memray"):
+    import memray
+
+logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
+
+
+def _read_delta_file_envelopes(
+    annotated_delta: DeltaAnnotated,
+    read_kwargs_provider: Optional[ReadKwargsProvider],
+    deltacat_storage=unimplemented_deltacat_storage,
+    deltacat_storage_kwargs: Optional[dict] = None,
+) -> Tuple[Optional[List[DeltaFileEnvelope]], int, int]:
+
+    tables = deltacat_storage.download_delta(
+        annotated_delta,
+        max_parallelism=1,
+        file_reader_kwargs_provider=read_kwargs_provider,
+        storage_type=StorageType.LOCAL,
+        **deltacat_storage_kwargs,
+    )
+    annotations = annotated_delta.annotations
+    assert (
+        len(tables) == len(annotations),
+        f"Unexpected Error: Length of downloaded delta manifest tables "
+        f"({len(tables)}) doesn't match the length of delta manifest "
+        f"annotations ({len(annotations)}).",
+    )
+    if not tables:
+        return None, 0, 0
+
+    delta_stream_position = annotations[0].annotation_stream_position
+    delta_type = annotations[0].annotation_delta_type
+
+    for annotation in annotations:
+        assert annotation.annotation_stream_position == delta_stream_position, (
+            f"Annotation stream position does not match - {annotation.annotation_stream_position} "
+            f"!= {delta_stream_position}"
+        )
+        assert annotation.annotation_delta_type == delta_type, (
+            f"Annotation delta type does not match - {annotation.annotation_delta_type} "
+            f"!= {delta_type}"
+        )
+
+    delta_file_envelopes = []
+    table = pa.concat_tables(tables)
+    total_record_count = len(table)
+    total_size_bytes = int(table.nbytes)
+
+    delta_file = DeltaFileEnvelope.of(
+        stream_position=delta_stream_position,
+        delta_type=delta_type,
+        table=table,
+    )
+    delta_file_envelopes.append(delta_file)
+    return delta_file_envelopes, total_record_count, total_size_bytes
+
+
+def _group_file_records_by_pk_hash_bucket(
+    annotated_delta: DeltaAnnotated,
+    num_hash_buckets: int,
+    primary_keys: List[str],
+    read_kwargs_provider: Optional[ReadKwargsProvider] = None,
+    deltacat_storage=unimplemented_deltacat_storage,
+    deltacat_storage_kwargs: Optional[dict] = None,
+) -> Tuple[Optional[DeltaFileEnvelopeGroups], int, int]:
+    # read input parquet s3 objects into a list of delta file envelopes
+    (
+        delta_file_envelopes,
+        total_record_count,
+        total_size_bytes,
+    ) = _read_delta_file_envelopes(
+        annotated_delta,
+        read_kwargs_provider,
+        deltacat_storage,
+        deltacat_storage_kwargs,
+    )
+
+    if delta_file_envelopes is None:
+        return None, 0, 0
+
+    logger.info(f"Read all delta file envelopes: {len(delta_file_envelopes)}")
+
+    # group the data by primary key hash value
+    hb_to_delta_file_envelopes = np.empty([num_hash_buckets], dtype="object")
+    for dfe in delta_file_envelopes:
+        logger.info("Grouping by pk hash bucket")
+        start = time.monotonic()
+        hash_bucket_to_table = group_by_pk_hash_bucket(
+            dfe.table,
+            num_hash_buckets,
+            primary_keys,
+        )
+        group_end = time.monotonic()
+        logger.info(f"Grouping took: {group_end - start}")
+        for hb, table in enumerate(hash_bucket_to_table):
+            if table:
+                if hb_to_delta_file_envelopes[hb] is None:
+                    hb_to_delta_file_envelopes[hb] = []
+                hb_to_delta_file_envelopes[hb].append(
+                    DeltaFileEnvelope.of(
+                        stream_position=dfe.stream_position,
+                        file_index=dfe.file_index,
+                        delta_type=dfe.delta_type,
+                        table=table,
+                    )
+                )
+    return hb_to_delta_file_envelopes, total_record_count, total_size_bytes
+
+
+def _timed_hash_bucket(input: HashBucketInput):
+    task_id = get_current_ray_task_id()
+    worker_id = get_current_ray_worker_id()
+    with memray.Tracker(
+        f"hash_bucket_{worker_id}_{task_id}.bin"
+    ) if input.enable_profiler else nullcontext():
+        (
+            delta_file_envelope_groups,
+            total_record_count,
+            total_size_bytes,
+        ) = _group_file_records_by_pk_hash_bucket(
+            annotated_delta=input.annotated_delta,
+            num_hash_buckets=input.num_hash_buckets,
+            primary_keys=input.primary_keys,
+            read_kwargs_provider=input.read_kwargs_provider,
+            deltacat_storage=input.deltacat_storage,
+            deltacat_storage_kwargs=input.deltacat_storage_kwargs,
+        )
+        hash_bucket_group_to_obj_id_tuple = group_hash_bucket_indices(
+            hash_bucket_object_groups=delta_file_envelope_groups,
+            num_buckets=input.num_hash_buckets,
+            num_groups=input.num_hash_groups,
+            object_store=input.object_store,
+        )
+
+        peak_memory_usage_bytes = get_current_node_peak_memory_usage_in_bytes()
+        return HashBucketResult(
+            hash_bucket_group_to_obj_id_tuple,
+            np.int64(total_size_bytes),
+            np.int64(total_record_count),
+            np.double(peak_memory_usage_bytes),
+            np.double(0.0),
+            np.double(time.time()),
+        )
+
+
+@ray.remote
+def hash_bucket(input: HashBucketInput) -> HashBucketResult:
+
+    logger.info(f"Starting hash bucket task...")
+    hash_bucket_result, duration = timed_invocation(
+        func=_timed_hash_bucket, input=input
+    )
+
+    emit_metrics_time = 0.0
+    if input.metrics_config:
+        emit_result, latency = timed_invocation(
+            func=emit_timer_metrics,
+            metrics_name="hash_bucket",
+            value=duration,
+            metrics_config=input.metrics_config,
+        )
+        emit_metrics_time = latency
+
+    logger.info(f"Finished hash bucket task...")
+    return HashBucketResult(
+        hash_bucket_result[0],
+        hash_bucket_result[1],
+        hash_bucket_result[2],
+        hash_bucket_result[3],
+        np.double(emit_metrics_time),
+        hash_bucket_result[5],
+    )