deltacat 0.2.10__py3-none-any.whl → 1.0.0__py3-none-any.whl

deltacat/compute/compactor_v2/compaction_session.py

@@ -6,18 +6,24 @@ import logging
 import ray
 import time
 import json
+
+from deltacat.compute.compactor_v2.model.merge_file_group import (
+    RemoteMergeFileGroupsProvider,
+)
+from deltacat.compute.compactor_v2.model.hash_bucket_input import HashBucketInput
+
+from deltacat.compute.compactor_v2.model.merge_input import MergeInput
+
 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 import PyArrowWriteResult, RoundCompletionInfo
 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.compute.compactor_v2.utils.merge import (
+    generate_local_merge_input,
+)
 from deltacat.storage import (
     Delta,
     DeltaLocator,
@@ -210,107 +216,6 @@ def _execute_compaction(
         logger.info("No input deltas found to compact.")
         return None, None, None
 
-    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,
-        ray_custom_resources=params.ray_custom_resources,
-    )
-
-    hb_start = time.monotonic()
-
-    def hash_bucket_input_provider(index, item):
-        return {
-            "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(
@@ -325,60 +230,176 @@ def _execute_compaction(
         **params.deltacat_storage_kwargs,
     )
 
-    # BSP Step 2: Merge
-    merge_options_provider = functools.partial(
+    hb_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_manifest=previous_compacted_delta_manifest,
+        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,
-        deltacat_storage=params.deltacat_storage,
-        deltacat_storage_kwargs=params.deltacat_storage_kwargs,
         ray_custom_resources=params.ray_custom_resources,
     )
 
-    def merge_input_provider(index, item):
-        return {
-            "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_bucket_count=params.hash_bucket_count,
-                drop_duplicates=params.drop_duplicates,
-                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,
-            )
-        }
+    total_input_records_count = np.int64(0)
+    total_hb_record_count = np.int64(0)
+    telemetry_time_hb = 0
+    if params.hash_bucket_count == 1:
+        merge_start = time.monotonic()
+        local_merge_input = generate_local_merge_input(
+            params, uniform_deltas, compacted_partition, round_completion_info
+        )
+        local_merge_result = ray.get(mg.merge.remote(local_merge_input))
+        total_input_records_count += local_merge_result.input_record_count
+        merge_results = [local_merge_result]
+        merge_invoke_end = time.monotonic()
+    else:
+        hb_start = time.monotonic()
+
+        def hash_bucket_input_provider(index, item):
+            return {
+                "input": HashBucketInput.of(
+                    item,
+                    primary_keys=params.primary_keys,
+                    hb_task_index=index,
+                    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,
+                )
+            }
+
+        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_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,
+        )
 
-    merge_start = time.monotonic()
+        hb_invoke_end = 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,
-    )
+        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,
+        )
+
+        hb_data_processed_size_bytes = 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_input_records_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_input_records_count} hash bucket records from hash bucketing step..."
+        )
+
+        total_hb_record_count = total_input_records_count
+        compaction_audit.set_hash_bucket_processed_size_bytes(
+            hb_data_processed_size_bytes.item()
+        )
+
+        # 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_manifest=previous_compacted_delta_manifest,
+            primary_keys=params.primary_keys,
+            deltacat_storage=params.deltacat_storage,
+            deltacat_storage_kwargs=params.deltacat_storage_kwargs,
+            ray_custom_resources=params.ray_custom_resources,
+        )
+
+        def merge_input_provider(index, item):
+            return {
+                "input": MergeInput.of(
+                    merge_file_groups_provider=RemoteMergeFileGroupsProvider(
+                        hash_group_index=item[0],
+                        dfe_groups_refs=item[1],
+                        hash_bucket_count=params.hash_bucket_count,
+                        num_hash_groups=params.hash_group_count,
+                        object_store=params.object_store,
+                    ),
+                    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,
+                    drop_duplicates=params.drop_duplicates,
+                    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)
 
-    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()
@@ -387,6 +408,8 @@ def _execute_compaction(
     total_dd_record_count = sum([ddr.deduped_record_count for ddr in merge_results])
     logger.info(f"Deduped {total_dd_record_count} records...")
 
+    compaction_audit.set_input_records(total_input_records_count.item())
+
     telemetry_time_merge = compaction_audit.save_step_stats(
         CompactionSessionAuditInfo.MERGE_STEP_NAME,
         merge_results,

deltacat/compute/compactor_v2/model/hash_bucket_input.py

@@ -15,6 +15,7 @@ class HashBucketInput(Dict):
         primary_keys: List[str],
         num_hash_buckets: int,
         num_hash_groups: int,
+        hb_task_index: Optional[int] = 0,
         enable_profiler: Optional[bool] = False,
         metrics_config: Optional[MetricsConfig] = None,
         read_kwargs_provider: Optional[ReadKwargsProvider] = None,
@@ -26,6 +27,7 @@ class HashBucketInput(Dict):
         result = HashBucketInput()
         result["annotated_delta"] = annotated_delta
         result["primary_keys"] = primary_keys
+        result["hb_task_index"] = hb_task_index
         result["num_hash_buckets"] = num_hash_buckets
         result["num_hash_groups"] = num_hash_groups
         result["enable_profiler"] = enable_profiler
@@ -45,6 +47,10 @@ class HashBucketInput(Dict):
     def primary_keys(self) -> List[str]:
         return self["primary_keys"]
 
+    @property
+    def hb_task_index(self) -> List[str]:
+        return self["hb_task_index"]
+
     @property
     def num_hash_buckets(self) -> int:
         return self["num_hash_buckets"]

deltacat/compute/compactor_v2/model/merge_file_group.py

@@ -0,0 +1,213 @@
+# Allow classes to use self-referencing Type hints in Python 3.7.
+from __future__ import annotations
+
+import logging
+import time
+from abc import ABC, abstractmethod
+from collections import defaultdict
+
+from deltacat.utils.common import ReadKwargsProvider
+from ray.types import ObjectRef
+
+from deltacat.compute.compactor.model.delta_file_envelope import DeltaFileEnvelopeGroups
+from deltacat.compute.compactor_v2.utils.delta import read_delta_file_envelopes
+
+from deltacat.compute.compactor_v2.utils.primary_key_index import (
+    hash_group_index_to_hash_bucket_indices,
+)
+
+from deltacat.storage import interface as unimplemented_deltacat_storage
+
+from deltacat.io.object_store import IObjectStore
+
+from deltacat import logs
+
+from deltacat.compute.compactor import DeltaFileEnvelope, DeltaAnnotated
+
+from typing import List, Optional
+
+logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
+
+
+class MergeFileGroup(dict):
+    @staticmethod
+    def of(hb_index: int, dfe_groups: Optional[List[List[DeltaFileEnvelope]]] = None):
+        """
+        Creates a container with delta file envelope groupings and other
+        additional properties used primarily for the merging step.
+
+        Args:
+            hb_index: This signifies the hash bucket index corresponding to the envelope delta file groups.
+            dfe_groups: A list of delta file envelope groups.
+                If not present, the provided hash bucket index is a copy by reference candidate during the merge step.
+
+        Returns:
+            A dict
+
+        """
+        d = MergeFileGroup()
+        d["hb_index"] = hb_index
+        d["dfe_groups"] = dfe_groups
+        return d
+
+    @property
+    def dfe_groups(self) -> Optional[List[List[DeltaFileEnvelope]]]:
+        return self["dfe_groups"]
+
+    @property
+    def hb_index(self) -> int:
+        return self["hb_index"]
+
+
+class MergeFileGroupsProvider(ABC):
+    @abstractmethod
+    def create(self) -> List[MergeFileGroup]:
+        """
+        Creates a list of merge file groups.
+
+        Returns: a list of merge file groups.
+
+        """
+        raise NotImplementedError("Method not implemented")
+
+    @property
+    @abstractmethod
+    def hash_group_index(self):
+        raise NotImplementedError("Method not implemented")
+
+
+class LocalMergeFileGroupsProvider(MergeFileGroupsProvider):
+    """
+    A factory class for producing merge file groups given local delta file envelopes.
+    """
+
+    LOCAL_HASH_BUCKET_INDEX = 0
+    LOCAL_HASH_GROUP_INDEX = 0
+
+    def __init__(
+        self,
+        uniform_deltas: List[DeltaAnnotated],
+        read_kwargs_provider: Optional[ReadKwargsProvider],
+        deltacat_storage=unimplemented_deltacat_storage,
+        deltacat_storage_kwargs: Optional[dict] = None,
+    ):
+        self._deltas = uniform_deltas
+        self._read_kwargs_provider = read_kwargs_provider
+        self._deltacat_storage = deltacat_storage
+        self._deltacat_storage_kwargs = deltacat_storage_kwargs
+        self._loaded_deltas = False
+
+    def _read_deltas_locally(self):
+        local_dfe_list = []
+        input_records_count = 0
+        uniform_deltas = self._deltas
+        logger.info(f"Getting {len(uniform_deltas)} DFE Tasks.")
+        dfe_start = time.monotonic()
+        for annotated_delta in uniform_deltas:
+            (
+                delta_file_envelopes,
+                total_record_count,
+                total_size_bytes,
+            ) = read_delta_file_envelopes(
+                annotated_delta,
+                self._read_kwargs_provider,
+                self._deltacat_storage,
+                self._deltacat_storage_kwargs,
+            )
+            if delta_file_envelopes:
+                local_dfe_list.extend(delta_file_envelopes)
+                input_records_count += total_record_count
+        dfe_end = time.monotonic()
+        logger.info(
+            f"Retrieved {len(local_dfe_list)} DFE Tasks in {dfe_end - dfe_start}s."
+        )
+
+        self._dfe_groups = [local_dfe_list] if len(local_dfe_list) > 0 else None
+        self._loaded_deltas = True
+
+    def create(self) -> List[MergeFileGroup]:
+        if not self._loaded_deltas:
+            self._read_deltas_locally()
+
+        # Since hash bucketing is skipped for local merges, we use a fixed index here.
+        return [
+            MergeFileGroup.of(
+                hb_index=LocalMergeFileGroupsProvider.LOCAL_HASH_BUCKET_INDEX,
+                dfe_groups=self._dfe_groups,
+            )
+        ]
+
+    @property
+    def hash_group_index(self):
+        return LocalMergeFileGroupsProvider.LOCAL_HASH_GROUP_INDEX
+
+
+class RemoteMergeFileGroupsProvider(MergeFileGroupsProvider):
+    """
+    A factory class for producing merge file groups given delta file envelope object refs
+        and hash bucketing parameters. Delta file envelopes are pulled from the object store
+        remotely and loaded with in-memory pyarrow tables.
+    """
+
+    def __init__(
+        self,
+        hash_group_index: int,
+        dfe_groups_refs: List[ObjectRef[DeltaFileEnvelopeGroups]],
+        hash_bucket_count: int,
+        num_hash_groups: int,
+        object_store: IObjectStore,
+    ):
+        self.hash_bucket_count = hash_bucket_count
+        self.num_hash_groups = num_hash_groups
+        self.object_store = object_store
+        self._hash_group_index = hash_group_index
+        self._dfe_groups_refs = dfe_groups_refs
+        self._dfe_groups = []
+        self._loaded_from_object_store = False
+
+    def _load_deltas_from_object_store(self):
+        delta_file_envelope_groups_list = self.object_store.get_many(
+            self._dfe_groups_refs
+        )
+        hb_index_to_delta_file_envelopes_list = defaultdict(list)
+        for delta_file_envelope_groups in delta_file_envelope_groups_list:
+            assert self.hash_bucket_count == len(delta_file_envelope_groups), (
+                f"The hash bucket count must match the dfe size as {self.hash_bucket_count}"
+                f" != {len(delta_file_envelope_groups)}"
+            )
+
+            for hb_idx, dfes in enumerate(delta_file_envelope_groups):
+                if dfes:
+                    hb_index_to_delta_file_envelopes_list[hb_idx].append(dfes)
+        valid_hb_indices_iterable = hash_group_index_to_hash_bucket_indices(
+            self.hash_group_index, self.hash_bucket_count, self.num_hash_groups
+        )
+
+        total_dfes_found = 0
+        dfe_list_groups = []
+        for hb_idx in valid_hb_indices_iterable:
+            dfe_list = hb_index_to_delta_file_envelopes_list.get(hb_idx)
+            if dfe_list:
+                total_dfes_found += 1
+                dfe_list_groups.append(
+                    MergeFileGroup.of(hb_index=hb_idx, dfe_groups=dfe_list)
+                )
+            else:
+                dfe_list_groups.append(MergeFileGroup.of(hb_index=hb_idx))
+
+        assert total_dfes_found == len(hb_index_to_delta_file_envelopes_list), (
+            "The total dfe list does not match the input dfes from hash bucket as "
+            f"{total_dfes_found} != {len(hb_index_to_delta_file_envelopes_list)}"
+        )
+        self._dfe_groups = dfe_list_groups
+        self._loaded_from_object_store = True
+
+    def create(self) -> List[MergeFileGroup]:
+        if not self._loaded_from_object_store:
+            self._load_deltas_from_object_store()
+
+        return self._dfe_groups
+
+    @property
+    def hash_group_index(self):
+        return self._hash_group_index

deltacat/compute/compactor_v2/model/merge_input.py

@@ -1,7 +1,10 @@
 from __future__ import annotations
 
-from ray.types import ObjectRef
 from typing import Dict, List, Optional, Any
+
+from deltacat.compute.compactor_v2.model.merge_file_group import (
+    MergeFileGroupsProvider,
+)
 from deltacat.utils.metrics import MetricsConfig
 from deltacat.utils.common import ReadKwargsProvider
 from deltacat.io.object_store import IObjectStore
@@ -16,19 +19,15 @@ from deltacat.compute.compactor_v2.constants import (
 )
 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]],
+        merge_file_groups_provider: MergeFileGroupsProvider,
         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,
@@ -44,13 +43,10 @@ class MergeInput(Dict):
     ) -> MergeInput:
 
         result = MergeInput()
-        result["dfe_groups_refs"] = dfe_groups_refs
+        result["merge_file_groups_provider"] = merge_file_groups_provider
         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
@@ -67,8 +63,8 @@ class MergeInput(Dict):
         return result
 
     @property
-    def dfe_groups_refs(self) -> List[ObjectRef[DeltaFileEnvelopeGroups]]:
-        return self["dfe_groups_refs"]
+    def merge_file_groups_provider(self) -> MergeFileGroupsProvider:
+        return self["merge_file_groups_provider"]
 
     @property
     def write_to_partition(self) -> Partition:
@@ -82,18 +78,6 @@ class MergeInput(Dict):
     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"]