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

deltacat/compute/compactor_v2/compaction_session.py

@@ -0,0 +1,506 @@
+import importlib
+from contextlib import nullcontext
+import numpy as np
+import functools
+import logging
+import ray
+import time
+import json
+from deltacat.aws import s3u as s3_utils
+import deltacat
+from deltacat import logs
+from deltacat.compute.compactor import (
+    PyArrowWriteResult,
+    RoundCompletionInfo,
+)
+from deltacat.compute.compactor_v2.model.merge_input import MergeInput
+from deltacat.compute.compactor_v2.model.merge_result import MergeResult
+from deltacat.compute.compactor_v2.model.hash_bucket_input import HashBucketInput
+from deltacat.compute.compactor_v2.model.hash_bucket_result import HashBucketResult
+from deltacat.compute.compactor.model.materialize_result import MaterializeResult
+from deltacat.storage import (
+    Delta,
+    DeltaLocator,
+    Partition,
+)
+from deltacat.compute.compactor.model.compact_partition_params import (
+    CompactPartitionParams,
+)
+from deltacat.utils.ray_utils.concurrency import (
+    invoke_parallel,
+    task_resource_options_provider,
+)
+from deltacat.compute.compactor_v2.steps import merge as mg
+from deltacat.compute.compactor_v2.steps import hash_bucket as hb
+from deltacat.compute.compactor_v2.utils import io
+from deltacat.compute.compactor.utils import round_completion_file as rcf
+
+from typing import List, Optional, Tuple
+from collections import defaultdict
+from deltacat.compute.compactor.model.compaction_session_audit_info import (
+    CompactionSessionAuditInfo,
+)
+from deltacat.utils.resources import (
+    get_current_node_peak_memory_usage_in_bytes,
+)
+from deltacat.compute.compactor_v2.utils.task_options import (
+    hash_bucket_resource_options_provider,
+    merge_resource_options_provider,
+)
+from deltacat.utils.resources import ClusterUtilizationOverTimeRange
+
+if importlib.util.find_spec("memray"):
+    import memray
+
+
+logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
+
+
+def compact_partition(params: CompactPartitionParams, **kwargs) -> Optional[str]:
+
+    assert (
+        params.hash_bucket_count is not None and params.hash_bucket_count >= 1
+    ), "hash_bucket_count is a required arg for compactor v2"
+
+    with memray.Tracker(
+        f"compaction_partition.bin"
+    ) if params.enable_profiler else nullcontext(), ClusterUtilizationOverTimeRange() as cluster_util:
+        (new_partition, new_rci, new_rcf_partition_locator,) = _execute_compaction(
+            params,
+            cluster_util=cluster_util,
+            **kwargs,
+        )
+
+        logger.info(
+            f"Partition-{params.source_partition_locator} -> "
+            f"Compaction session data processing completed"
+        )
+        round_completion_file_s3_url = None
+        if new_partition:
+            logger.info(f"Committing compacted partition to: {new_partition.locator}")
+            partition = params.deltacat_storage.commit_partition(
+                new_partition, **params.deltacat_storage_kwargs
+            )
+            logger.info(f"Committed compacted partition: {partition}")
+
+            round_completion_file_s3_url = rcf.write_round_completion_file(
+                params.compaction_artifact_s3_bucket,
+                new_rcf_partition_locator,
+                new_rci,
+                **params.s3_client_kwargs,
+            )
+        else:
+            logger.warn("No new partition was committed during compaction.")
+
+        logger.info(
+            f"Completed compaction session for: {params.source_partition_locator}"
+        )
+        return round_completion_file_s3_url
+
+
+def _execute_compaction(
+    params: CompactPartitionParams, **kwargs
+) -> Tuple[Optional[Partition], Optional[RoundCompletionInfo], Optional[str]]:
+
+    rcf_source_partition_locator = (
+        params.rebase_source_partition_locator or params.source_partition_locator
+    )
+
+    base_audit_url = rcf_source_partition_locator.path(
+        f"s3://{params.compaction_artifact_s3_bucket}/compaction-audit"
+    )
+    audit_url = f"{base_audit_url}.json"
+    logger.info(f"Compaction audit will be written to {audit_url}")
+    compaction_audit = CompactionSessionAuditInfo(deltacat.__version__, audit_url)
+
+    compaction_audit.set_hash_bucket_count(params.hash_bucket_count)
+
+    compaction_start = time.monotonic()
+
+    task_max_parallelism = params.task_max_parallelism
+
+    if params.pg_config:
+        logger.info(
+            "pg_config specified. Tasks will be scheduled in a placement group."
+        )
+        cluster_resources = params.pg_config.resource
+        cluster_cpus = cluster_resources["CPU"]
+        cluster_memory = cluster_resources["memory"]
+        task_max_parallelism = cluster_cpus
+        compaction_audit.set_cluster_cpu_max(cluster_cpus)
+        compaction_audit.set_total_cluster_memory_bytes(cluster_memory)
+
+    # read the results from any previously completed compaction round
+    round_completion_info = None
+    high_watermark = None
+    previous_compacted_delta = None
+
+    if not params.rebase_source_partition_locator:
+        round_completion_info = rcf.read_round_completion_file(
+            params.compaction_artifact_s3_bucket,
+            params.source_partition_locator,
+            **params.s3_client_kwargs,
+        )
+        if not round_completion_info:
+            logger.info(
+                f"Both rebase partition and round completion file not found. Performing an entire backfill on source."
+            )
+        else:
+            compacted_delta_locator = round_completion_info.compacted_delta_locator
+            previous_compacted_delta = params.deltacat_storage.get_delta(
+                namespace=compacted_delta_locator.namespace,
+                table_name=compacted_delta_locator.table_name,
+                table_version=compacted_delta_locator.table_version,
+                stream_position=compacted_delta_locator.stream_position,
+                include_manifest=True,
+                **params.deltacat_storage_kwargs,
+            )
+
+            high_watermark = round_completion_info.high_watermark
+            logger.info(f"Setting round completion high watermark: {high_watermark}")
+            assert (
+                params.hash_bucket_count == round_completion_info.hash_bucket_count
+            ), (
+                "The hash bucket count has changed. "
+                "Kindly run rebase compaction and trigger incremental again. "
+                f"Hash Bucket count in RCF={round_completion_info.hash_bucket_count} "
+                f"not equal to Hash bucket count in args={params.hash_bucket_count}."
+            )
+
+        logger.info(f"Round completion file: {round_completion_info}")
+
+    delta_discovery_start = time.monotonic()
+
+    input_deltas = io.discover_deltas(
+        params.source_partition_locator,
+        params.last_stream_position_to_compact,
+        params.rebase_source_partition_locator,
+        params.rebase_source_partition_high_watermark,
+        high_watermark,
+        params.deltacat_storage,
+        params.deltacat_storage_kwargs,
+        params.list_deltas_kwargs,
+    )
+
+    delta_discovery_end = time.monotonic()
+    compaction_audit.set_delta_discovery_time_in_seconds(
+        delta_discovery_end - delta_discovery_start
+    )
+
+    s3_utils.upload(
+        compaction_audit.audit_url,
+        str(json.dumps(compaction_audit)),
+        **params.s3_client_kwargs,
+    )
+
+    if not input_deltas:
+        logger.info("No input deltas found to compact.")
+        return None, None, None
+
+    uniform_deltas = io.create_uniform_input_deltas(
+        input_deltas=input_deltas,
+        hash_bucket_count=params.hash_bucket_count,
+        compaction_audit=compaction_audit,
+        deltacat_storage=params.deltacat_storage,
+        previous_inflation=params.previous_inflation,
+        min_delta_bytes=params.min_delta_bytes_in_batch,
+        min_file_counts=params.min_files_in_batch,
+        deltacat_storage_kwargs=params.deltacat_storage_kwargs,
+    )
+
+    compaction_audit.set_uniform_deltas_created(len(uniform_deltas))
+
+    hb_options_provider = functools.partial(
+        task_resource_options_provider,
+        pg_config=params.pg_config,
+        resource_amount_provider=hash_bucket_resource_options_provider,
+        previous_inflation=params.previous_inflation,
+        average_record_size_bytes=params.average_record_size_bytes,
+        primary_keys=params.primary_keys,
+    )
+
+    hb_start = time.monotonic()
+
+    hash_bucket_input_provider = lambda index, item: {
+        "input": HashBucketInput.of(
+            item,
+            primary_keys=params.primary_keys,
+            num_hash_buckets=params.hash_bucket_count,
+            num_hash_groups=params.hash_group_count,
+            enable_profiler=params.enable_profiler,
+            metrics_config=params.metrics_config,
+            read_kwargs_provider=params.read_kwargs_provider,
+            object_store=params.object_store,
+            deltacat_storage=params.deltacat_storage,
+            deltacat_storage_kwargs=params.deltacat_storage_kwargs,
+        )
+    }
+
+    hb_tasks_pending = invoke_parallel(
+        items=uniform_deltas,
+        ray_task=hb.hash_bucket,
+        max_parallelism=task_max_parallelism,
+        options_provider=hb_options_provider,
+        kwargs_provider=hash_bucket_input_provider,
+    )
+
+    hb_invoke_end = time.monotonic()
+
+    logger.info(f"Getting {len(hb_tasks_pending)} hash bucket results...")
+    hb_results: List[HashBucketResult] = ray.get(hb_tasks_pending)
+    logger.info(f"Got {len(hb_results)} hash bucket results.")
+    hb_end = time.monotonic()
+
+    # we use time.time() here because time.monotonic() has no reference point
+    # whereas time.time() measures epoch seconds. Hence, it will be reasonable
+    # to compare time.time()s captured in different nodes.
+    hb_results_retrieved_at = time.time()
+
+    telemetry_time_hb = compaction_audit.save_step_stats(
+        CompactionSessionAuditInfo.HASH_BUCKET_STEP_NAME,
+        hb_results,
+        hb_results_retrieved_at,
+        hb_invoke_end - hb_start,
+        hb_end - hb_start,
+    )
+
+    s3_utils.upload(
+        compaction_audit.audit_url,
+        str(json.dumps(compaction_audit)),
+        **params.s3_client_kwargs,
+    )
+
+    all_hash_group_idx_to_obj_id = defaultdict(list)
+    all_hash_group_idx_to_size_bytes = defaultdict(int)
+    all_hash_group_idx_to_num_rows = defaultdict(int)
+    hb_data_processed_size_bytes = np.int64(0)
+    total_hb_record_count = np.int64(0)
+
+    # initialize all hash groups
+    for hb_group in range(params.hash_group_count):
+        all_hash_group_idx_to_num_rows[hb_group] = 0
+        all_hash_group_idx_to_obj_id[hb_group] = []
+        all_hash_group_idx_to_size_bytes[hb_group] = 0
+
+    for hb_result in hb_results:
+        hb_data_processed_size_bytes += hb_result.hb_size_bytes
+        total_hb_record_count += hb_result.hb_record_count
+
+        for hash_group_index, object_id_size_tuple in enumerate(
+            hb_result.hash_bucket_group_to_obj_id_tuple
+        ):
+            if object_id_size_tuple:
+                all_hash_group_idx_to_obj_id[hash_group_index].append(
+                    object_id_size_tuple[0]
+                )
+                all_hash_group_idx_to_size_bytes[
+                    hash_group_index
+                ] += object_id_size_tuple[1].item()
+                all_hash_group_idx_to_num_rows[
+                    hash_group_index
+                ] += object_id_size_tuple[2].item()
+
+    logger.info(
+        f"Got {total_hb_record_count} hash bucket records from hash bucketing step..."
+    )
+
+    compaction_audit.set_input_records(total_hb_record_count.item())
+    compaction_audit.set_hash_bucket_processed_size_bytes(
+        hb_data_processed_size_bytes.item()
+    )
+
+    # create a new stream for this round
+    compacted_stream_locator = params.destination_partition_locator.stream_locator
+    compacted_stream = params.deltacat_storage.get_stream(
+        compacted_stream_locator.namespace,
+        compacted_stream_locator.table_name,
+        compacted_stream_locator.table_version,
+        **params.deltacat_storage_kwargs,
+    )
+    compacted_partition = params.deltacat_storage.stage_partition(
+        compacted_stream,
+        params.destination_partition_locator.partition_values,
+        **params.deltacat_storage_kwargs,
+    )
+
+    # BSP Step 2: Merge
+    merge_options_provider = functools.partial(
+        task_resource_options_provider,
+        pg_config=params.pg_config,
+        resource_amount_provider=merge_resource_options_provider,
+        num_hash_groups=params.hash_group_count,
+        hash_group_size_bytes=all_hash_group_idx_to_size_bytes,
+        hash_group_num_rows=all_hash_group_idx_to_num_rows,
+        round_completion_info=round_completion_info,
+        compacted_delta=previous_compacted_delta,
+        primary_keys=params.primary_keys,
+        deltacat_storage=params.deltacat_storage,
+        deltacat_storage_kwargs=params.deltacat_storage_kwargs,
+    )
+
+    merge_input_provider = lambda index, item: {
+        "input": MergeInput.of(
+            dfe_groups_refs=item[1],
+            write_to_partition=compacted_partition,
+            compacted_file_content_type=params.compacted_file_content_type,
+            primary_keys=params.primary_keys,
+            sort_keys=params.sort_keys,
+            merge_task_index=index,
+            hash_group_index=item[0],
+            num_hash_groups=params.hash_group_count,
+            max_records_per_output_file=params.records_per_compacted_file,
+            enable_profiler=params.enable_profiler,
+            metrics_config=params.metrics_config,
+            s3_table_writer_kwargs=params.s3_table_writer_kwargs,
+            read_kwargs_provider=params.read_kwargs_provider,
+            round_completion_info=round_completion_info,
+            object_store=params.object_store,
+            deltacat_storage=params.deltacat_storage,
+            deltacat_storage_kwargs=params.deltacat_storage_kwargs,
+        )
+    }
+
+    merge_start = time.monotonic()
+
+    merge_tasks_pending = invoke_parallel(
+        items=all_hash_group_idx_to_obj_id.items(),
+        ray_task=mg.merge,
+        max_parallelism=task_max_parallelism,
+        options_provider=merge_options_provider,
+        kwargs_provider=merge_input_provider,
+    )
+
+    merge_invoke_end = time.monotonic()
+    logger.info(f"Getting {len(merge_tasks_pending)} merge results...")
+    merge_results: List[MergeResult] = ray.get(merge_tasks_pending)
+    logger.info(f"Got {len(merge_results)} merge results.")
+
+    merge_results_retrieved_at = time.time()
+    merge_end = time.monotonic()
+
+    total_dd_record_count = sum([ddr.deduped_record_count for ddr in merge_results])
+    logger.info(f"Deduped {total_dd_record_count} records...")
+
+    telemetry_time_merge = compaction_audit.save_step_stats(
+        CompactionSessionAuditInfo.MERGE_STEP_NAME,
+        merge_results,
+        merge_results_retrieved_at,
+        merge_invoke_end - merge_start,
+        merge_end - merge_start,
+    )
+
+    compaction_audit.set_records_deduped(total_dd_record_count.item())
+
+    mat_results = []
+    for merge_result in merge_results:
+        mat_results.extend(merge_result.materialize_results)
+
+    mat_results: List[MaterializeResult] = sorted(
+        mat_results, key=lambda m: m.task_index
+    )
+
+    deltas = [m.delta for m in mat_results]
+
+    hb_id_to_entry_indices_range = {}
+    file_index = 0
+    previous_task_index = -1
+
+    for m in mat_results:
+        assert m.pyarrow_write_result.files >= 1, "Atleast file must be materialized"
+        assert m.task_index != previous_task_index, (
+            "Multiple materialize results found for a " f"hash bucket: {m.task_index}"
+        )
+
+        hb_id_to_entry_indices_range[str(m.task_index)] = (
+            file_index,
+            file_index + m.pyarrow_write_result.files - 1,
+        )
+
+        file_index += m.pyarrow_write_result.files
+        previous_task_index = m.task_index
+
+    s3_utils.upload(
+        compaction_audit.audit_url,
+        str(json.dumps(compaction_audit)),
+        **params.s3_client_kwargs,
+    )
+
+    mat_results = sorted(mat_results, key=lambda m: m.task_index)
+    deltas = [m.delta for m in mat_results]
+
+    # Note: An appropriate last stream position must be set
+    # to avoid correctness issue.
+    merged_delta = Delta.merge_deltas(
+        deltas,
+        stream_position=params.last_stream_position_to_compact,
+    )
+
+    record_info_msg = (
+        f"Hash bucket records: {total_hb_record_count},"
+        f" Deduped records: {total_dd_record_count}, "
+        f" Materialized records: {merged_delta.meta.record_count}"
+    )
+    logger.info(record_info_msg)
+
+    compacted_delta = params.deltacat_storage.commit_delta(
+        merged_delta,
+        properties=kwargs.get("properties", {}),
+        **params.deltacat_storage_kwargs,
+    )
+
+    logger.info(f"Committed compacted delta: {compacted_delta}")
+
+    compaction_end = time.monotonic()
+    compaction_audit.set_compaction_time_in_seconds(compaction_end - compaction_start)
+
+    new_compacted_delta_locator = DeltaLocator.of(
+        compacted_partition.locator,
+        compacted_delta.stream_position,
+    )
+
+    pyarrow_write_result = PyArrowWriteResult.union(
+        [m.pyarrow_write_result for m in mat_results]
+    )
+
+    session_peak_memory = get_current_node_peak_memory_usage_in_bytes()
+    compaction_audit.set_peak_memory_used_bytes_by_compaction_session_process(
+        session_peak_memory
+    )
+
+    compaction_audit.save_round_completion_stats(
+        mat_results, telemetry_time_hb + telemetry_time_merge
+    )
+
+    cluster_util: ClusterUtilizationOverTimeRange = kwargs.get("cluster_util")
+
+    if cluster_util:
+        compaction_audit.set_total_cpu_seconds(cluster_util.total_vcpu_seconds)
+        compaction_audit.set_used_cpu_seconds(cluster_util.used_vcpu_seconds)
+
+    s3_utils.upload(
+        compaction_audit.audit_url,
+        str(json.dumps(compaction_audit)),
+        **params.s3_client_kwargs,
+    )
+
+    new_round_completion_info = RoundCompletionInfo.of(
+        high_watermark=params.last_stream_position_to_compact,
+        compacted_delta_locator=new_compacted_delta_locator,
+        compacted_pyarrow_write_result=pyarrow_write_result,
+        sort_keys_bit_width=params.bit_width_of_sort_keys,
+        manifest_entry_copied_by_reference_ratio=compaction_audit.untouched_file_ratio,
+        compaction_audit_url=audit_url,
+        hash_bucket_count=params.hash_bucket_count,
+        hb_index_to_entry_range=hb_id_to_entry_indices_range,
+    )
+
+    logger.info(
+        f"partition-{params.source_partition_locator.partition_values},"
+        f"compacted at: {params.last_stream_position_to_compact},"
+    )
+
+    return (
+        compacted_partition,
+        new_round_completion_info,
+        rcf_source_partition_locator,
+    )

deltacat/compute/compactor_v2/constants.py

@@ -0,0 +1,34 @@
+TOTAL_BYTES_IN_SHA1_HASH = 20
+
+PK_DELIMITER = "L6kl7u5f"
+
+MAX_RECORDS_PER_COMPACTED_FILE = 4_000_000
+
+# The maximum amount of delta bytes allowed in a batch.
+# A single task will not process more than these many bytes
+# unless a single manifest entry (non-parquet) or single row
+# group (parquet) is bigger than this size.
+MIN_DELTA_BYTES_IN_BATCH = 5_000_000_000
+
+# The total number of files that can be processed in a
+# batch. Hence, if there are tiny files, this value can be
+# limited so that enough parallelism can be attained.
+MIN_FILES_IN_BATCH = float("inf")
+
+# The average record size in a table.
+AVERAGE_RECORD_SIZE_BYTES = 1000
+
+# Maximum parallelism for the tasks at each BSP step.
+# Default is the number of vCPUs in about 168
+# r5.8xlarge EC2 instances.
+TASK_MAX_PARALLELISM = 5367
+
+# The percentage of memory that needs to be estimated
+# as buffer. This value will ensure the job doesn't run out
+# of memory by considering buffer for uncertainities.
+TOTAL_MEMORY_BUFFER_PERCENTAGE = 20
+
+# The total size of records that will be hash bucketed at once
+# Since, sorting is nlogn, we ensure that is not performed
+# on a very large dataset for best performance.
+MAX_SIZE_OF_RECORD_BATCH_IN_GIB = 2 * 1024 * 1024 * 1024

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