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

deltacat/compute/compactor_v2/model/hash_bucket_input.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+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 interface as unimplemented_deltacat_storage
+from deltacat.compute.compactor import DeltaAnnotated
+
+
+class HashBucketInput(Dict):
+    @staticmethod
+    def of(
+        annotated_delta: DeltaAnnotated,
+        primary_keys: List[str],
+        num_hash_buckets: int,
+        num_hash_groups: int,
+        enable_profiler: Optional[bool] = False,
+        metrics_config: Optional[MetricsConfig] = None,
+        read_kwargs_provider: Optional[ReadKwargsProvider] = None,
+        object_store: Optional[IObjectStore] = None,
+        deltacat_storage=unimplemented_deltacat_storage,
+        deltacat_storage_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> HashBucketInput:
+
+        result = HashBucketInput()
+        result["annotated_delta"] = annotated_delta
+        result["primary_keys"] = primary_keys
+        result["num_hash_buckets"] = num_hash_buckets
+        result["num_hash_groups"] = num_hash_groups
+        result["enable_profiler"] = enable_profiler
+        result["metrics_config"] = metrics_config
+        result["read_kwargs_provider"] = read_kwargs_provider
+        result["object_store"] = object_store
+        result["deltacat_storage"] = deltacat_storage
+        result["deltacat_storage_kwargs"] = deltacat_storage_kwargs or {}
+
+        return result
+
+    @property
+    def annotated_delta(self) -> DeltaAnnotated:
+        return self["annotated_delta"]
+
+    @property
+    def primary_keys(self) -> List[str]:
+        return self["primary_keys"]
+
+    @property
+    def num_hash_buckets(self) -> int:
+        return self["num_hash_buckets"]
+
+    @property
+    def num_hash_groups(self) -> int:
+        return self["num_hash_groups"]
+
+    @property
+    def enable_profiler(self) -> Optional[bool]:
+        return self.get("enable_profiler")
+
+    @property
+    def metrics_config(self) -> Optional[MetricsConfig]:
+        return self.get("metrics_config")
+
+    @property
+    def read_kwargs_provider(self) -> Optional[ReadKwargsProvider]:
+        return self.get("read_kwargs_provider")
+
+    @property
+    def object_store(self) -> Optional[IObjectStore]:
+        return self.get("object_store")
+
+    @property
+    def deltacat_storage(self) -> unimplemented_deltacat_storage:
+        return self.get("deltacat_storage")
+
+    @property
+    def deltacat_storage_kwargs(self) -> Optional[Dict[str, Any]]:
+        return self.get("deltacat_storage_kwargs")

deltacat/compute/compactor_v2/model/hash_bucket_result.py

@@ -0,0 +1,12 @@
+from typing import NamedTuple
+
+import numpy as np
+
+
+class HashBucketResult(NamedTuple):
+    hash_bucket_group_to_obj_id_tuple: np.ndarray
+    hb_size_bytes: np.int64
+    hb_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/model/merge_input.py

@@ -0,0 +1,127 @@
+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.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,
+        sort_keys: Optional[List[SortKey]] = None,
+        merge_task_index: Optional[int] = 0,
+        max_records_per_output_file: Optional[int] = 4_000_000,
+        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["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 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],
+    )

deltacat/compute/compactor_v2/steps/merge.py

@@ -0,0 +1,41 @@
+import logging
+from deltacat.compute.compactor_v2.model.merge_input import MergeInput
+import numpy as np
+import ray
+from deltacat import logs
+from deltacat.compute.compactor_v2.model.merge_result import MergeResult
+from deltacat.utils.performance import timed_invocation
+from deltacat.utils.metrics import emit_timer_metrics
+
+logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
+
+
+def _timed_merge(input: MergeInput) -> MergeResult:
+    # TODO: Implementation goes here
+    pass
+
+
+@ray.remote
+def merge(input: MergeInput) -> MergeResult:
+
+    logger.info(f"Starting merge task...")
+    merge_result, duration = timed_invocation(func=_timed_merge, input=input)
+
+    emit_metrics_time = 0.0
+    if input.metrics_config:
+        emit_result, latency = timed_invocation(
+            func=emit_timer_metrics,
+            metrics_name="merge",
+            value=duration,
+            metrics_config=input.metrics_config,
+        )
+        emit_metrics_time = latency
+
+    logger.info(f"Finished merge task...")
+    return MergeResult(
+        merge_result[0],
+        merge_result[1],
+        merge_result[2],
+        np.double(emit_metrics_time),
+        merge_result[4],
+    )

deltacat/compute/compactor_v2/utils/content_type_params.py

@@ -0,0 +1,37 @@
+from deltacat.storage import (
+    Delta,
+    interface as unimplemented_deltacat_storage,
+)
+from typing import Dict, Optional, Any
+from deltacat.types.media import TableType
+from deltacat.types.media import ContentType
+from deltacat.types.partial_download import PartialParquetParameters
+
+
+def append_content_type_params(
+    delta: Delta,
+    entry_index: int,
+    deltacat_storage: unimplemented_deltacat_storage,
+    deltacat_storage_kwargs: Optional[Dict[str, Any]] = {},
+):
+
+    entry = delta.manifest.entries[entry_index]
+
+    if entry.meta.content_type == ContentType.PARQUET:
+        # Set partial download content type parameters to allow
+        # row group level rebatching for parquet.
+        pq_file = deltacat_storage.download_delta_manifest_entry(
+            delta.locator,
+            entry_index=entry_index,
+            table_type=TableType.PYARROW_PARQUET,
+            **deltacat_storage_kwargs,
+        )
+
+        if not entry.meta.content_type_parameters:
+            entry.meta.content_type_parameters = []
+
+        entry.meta.content_type_parameters.append(
+            PartialParquetParameters.of(pq_metadata=pq_file.metadata)
+        )
+
+    return entry

deltacat/compute/compactor_v2/utils/io.py

@@ -0,0 +1,149 @@
+import logging
+import functools
+from deltacat.constants import PYARROW_INFLATION_MULTIPLIER
+from deltacat.storage import (
+    PartitionLocator,
+    Delta,
+    interface as unimplemented_deltacat_storage,
+)
+from deltacat import logs
+from deltacat.compute.compactor.utils import io as io_v1
+from deltacat.compute.compactor import DeltaAnnotated
+from typing import Dict, List, Optional, Any
+from deltacat.compute.compactor_v2.constants import (
+    MIN_FILES_IN_BATCH,
+    MIN_DELTA_BYTES_IN_BATCH,
+)
+from deltacat.compute.compactor.model.compaction_session_audit_info import (
+    CompactionSessionAuditInfo,
+)
+from deltacat.compute.compactor_v2.utils.task_options import (
+    estimate_manifest_entry_size_bytes,
+)
+from deltacat.compute.compactor_v2.utils.content_type_params import (
+    append_content_type_params,
+)
+
+logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
+
+
+def discover_deltas(
+    source_partition_locator: PartitionLocator,
+    last_stream_position_to_compact: int,
+    rebase_source_partition_locator: Optional[PartitionLocator] = None,
+    rebase_source_partition_high_watermark: Optional[int] = None,
+    rcf_high_watermark: Optional[int] = None,
+    deltacat_storage=unimplemented_deltacat_storage,
+    deltacat_storage_kwargs: Optional[Dict[str, Any]] = {},
+    list_deltas_kwargs: Optional[Dict[str, Any]] = {},
+) -> List[Delta]:
+
+    previous_compacted_high_watermark = (
+        rebase_source_partition_high_watermark or rcf_high_watermark
+    )
+
+    delta_source_partition_locator = (
+        rebase_source_partition_locator or source_partition_locator
+    )
+
+    result = []
+
+    delta_source_incremental_deltas = io_v1._discover_deltas(
+        delta_source_partition_locator,
+        previous_compacted_high_watermark,
+        last_stream_position_to_compact,
+        deltacat_storage,
+        deltacat_storage_kwargs,
+        list_deltas_kwargs,
+    )
+
+    result.extend(delta_source_incremental_deltas)
+
+    logger.info(
+        f"Length of input deltas from delta source table is {len(delta_source_incremental_deltas)}"
+        f" from ({previous_compacted_high_watermark}, {last_stream_position_to_compact}]"
+    )
+
+    if rebase_source_partition_locator:
+        previous_compacted_deltas = io_v1._discover_deltas(
+            source_partition_locator,
+            None,
+            None,
+            deltacat_storage,
+            deltacat_storage_kwargs,
+            list_deltas_kwargs,
+        )
+
+        result.extend(previous_compacted_deltas)
+
+        logger.info(
+            f"Length of input deltas from previous compacted table is {len(previous_compacted_deltas)}"
+            f" from ({None}, {None}]"
+        )
+
+    return result
+
+
+def create_uniform_input_deltas(
+    input_deltas: List[Delta],
+    hash_bucket_count: int,
+    compaction_audit: CompactionSessionAuditInfo,
+    min_delta_bytes: Optional[float] = MIN_DELTA_BYTES_IN_BATCH,
+    min_file_counts: Optional[float] = MIN_FILES_IN_BATCH,
+    previous_inflation: Optional[float] = PYARROW_INFLATION_MULTIPLIER,
+    deltacat_storage=unimplemented_deltacat_storage,
+    deltacat_storage_kwargs: Optional[Dict[str, Any]] = {},
+) -> List[DeltaAnnotated]:
+
+    delta_bytes = 0
+    delta_manifest_entries = 0
+    estimated_da_bytes = 0
+    input_da_list = []
+
+    for delta in input_deltas:
+        manifest_entries = delta.manifest.entries
+        delta_manifest_entries += len(manifest_entries)
+        for entry_index in range(len(manifest_entries)):
+            entry = append_content_type_params(
+                delta=delta,
+                entry_index=entry_index,
+                deltacat_storage=deltacat_storage,
+                deltacat_storage_kwargs=deltacat_storage_kwargs,
+            )
+
+            delta_bytes += entry.meta.content_length
+            estimated_da_bytes += estimate_manifest_entry_size_bytes(
+                entry=entry, previous_inflation=previous_inflation
+            )
+
+        delta_annotated = DeltaAnnotated.of(delta)
+        input_da_list.append(delta_annotated)
+
+    logger.info(f"Input deltas to compact this round: " f"{len(input_da_list)}")
+    logger.info(f"Input delta bytes to compact: {delta_bytes}")
+    logger.info(f"Input delta files to compact: {delta_manifest_entries}")
+
+    if not input_da_list:
+        raise RuntimeError("No input deltas to compact!")
+
+    size_estimation_function = functools.partial(
+        estimate_manifest_entry_size_bytes, previous_inflation=previous_inflation
+    )
+
+    rebatched_da_list = DeltaAnnotated.rebatch(
+        input_da_list,
+        min_delta_bytes=min_delta_bytes,
+        min_file_counts=min_file_counts,
+        estimation_function=size_estimation_function,
+    )
+
+    compaction_audit.set_input_size_bytes(delta_bytes)
+    compaction_audit.set_input_file_count(delta_manifest_entries)
+    compaction_audit.set_estimated_in_memory_size_bytes_during_discovery(
+        estimated_da_bytes
+    )
+
+    logger.info(f"Hash bucket count: {hash_bucket_count}")
+    logger.info(f"Input uniform delta count: {len(rebatched_da_list)}")
+
+    return rebatched_da_list