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

deltacat/compute/compactor_v2/utils/io.py

@@ -0,0 +1,152 @@
+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,
+    enable_input_split: Optional[bool] = False,
+    deltacat_storage=unimplemented_deltacat_storage,
+    deltacat_storage_kwargs: Optional[Dict[str, Any]] = {},
+) -> List[DeltaAnnotated]:
+
+    delta_bytes = 0
+    delta_manifest_entries_count = 0
+    estimated_da_bytes = 0
+    input_da_list = []
+
+    for delta in input_deltas:
+        if enable_input_split:
+            append_content_type_params(
+                delta=delta,
+                deltacat_storage=deltacat_storage,
+                deltacat_storage_kwargs=deltacat_storage_kwargs,
+            )
+
+        manifest_entries = delta.manifest.entries
+        delta_manifest_entries_count += len(manifest_entries)
+
+        for entry_index in range(len(manifest_entries)):
+            entry = manifest_entries[entry_index]
+            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_count}")
+
+    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_count)
+    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

deltacat/compute/compactor_v2/utils/primary_key_index.py

@@ -0,0 +1,341 @@
+import logging
+from typing import List, Optional, Iterable
+
+import numpy as np
+import pyarrow as pa
+import pyarrow.compute as pc
+import uuid
+import hashlib
+from deltacat.compute.compactor_v2.constants import (
+    TOTAL_BYTES_IN_SHA1_HASH,
+    PK_DELIMITER,
+    MAX_SIZE_OF_RECORD_BATCH_IN_GIB,
+)
+import time
+from deltacat.compute.compactor.model.delta_file_envelope import DeltaFileEnvelope
+from deltacat import logs
+from deltacat.compute.compactor.utils import system_columns as sc
+from deltacat.io.object_store import IObjectStore
+from deltacat.utils.performance import timed_invocation
+
+logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
+
+
+def _append_sha1_hash_to_table(table: pa.Table, hash_column: pa.Array) -> pa.Table:
+    hash_column_np = hash_column.to_numpy()
+
+    result = []
+    for hash_value in hash_column_np:
+        result.append(hashlib.sha1(hash_value.encode("utf-8")).hexdigest())
+
+    return sc.append_pk_hash_string_column(table, result)
+
+
+def _is_sha1_desired(hash_columns: List[pa.Array]) -> bool:
+    total_size = 0
+    total_len = 0
+
+    for hash_column in hash_columns:
+        total_size += hash_column.nbytes
+        total_len += len(hash_column)
+
+    logger.info(
+        f"Found total length of hash column={total_len} and total_size={total_size}"
+    )
+
+    return total_size > TOTAL_BYTES_IN_SHA1_HASH * total_len
+
+
+def _append_table_by_hash_bucket(
+    pki_table: pa.Table, hash_bucket_to_table: np.ndarray
+) -> int:
+
+    hb_pk_table, sort_latency = timed_invocation(
+        lambda: pki_table.sort_by(sc._HASH_BUCKET_IDX_COLUMN_NAME)
+    )
+    logger.info(f"Sorting a pk table of length {len(pki_table)} took {sort_latency}s")
+
+    hb_pk_grouped_by, groupby_latency = timed_invocation(
+        lambda: hb_pk_table.group_by(sc._HASH_BUCKET_IDX_COLUMN_NAME).aggregate(
+            [(sc._HASH_BUCKET_IDX_COLUMN_NAME, "count")]
+        )
+    )
+
+    logger.info(
+        f"Grouping a pki table of length {len(pki_table)} took {groupby_latency}s"
+    )
+
+    group_count_array = hb_pk_grouped_by[f"{sc._HASH_BUCKET_IDX_COLUMN_NAME}_count"]
+    hb_group_array = hb_pk_grouped_by[sc._HASH_BUCKET_IDX_COLUMN_NAME]
+
+    result_len = 0
+    for i, group_count in enumerate(group_count_array):
+        hb_idx = hb_group_array[i].as_py()
+        pyarrow_table = hb_pk_table.slice(offset=result_len, length=group_count.as_py())
+        pyarrow_table = pyarrow_table.drop(
+            [sc._HASH_BUCKET_IDX_COLUMN_NAME, sc._PK_HASH_STRING_COLUMN_NAME]
+        )
+        if hash_bucket_to_table[hb_idx] is None:
+            hash_bucket_to_table[hb_idx] = []
+        hash_bucket_to_table[hb_idx].append(pyarrow_table)
+        result_len += len(pyarrow_table)
+
+    return result_len
+
+
+def _optimized_group_record_batches_by_hash_bucket(
+    pki_table: pa.Table, num_buckets: int
+):
+
+    input_table_len = len(pki_table)
+
+    hash_bucket_to_tables = np.empty([num_buckets], dtype="object")
+    hb_to_table = np.empty([num_buckets], dtype="object")
+
+    # This split will ensure that the sort is not performed on a very huge table
+    # resulting in ArrowInvalid: offset overflow while concatenating arrays
+    # Known issue with Arrow: https://github.com/apache/arrow/issues/25822
+    table_batches, to_batches_latency = timed_invocation(lambda: pki_table.to_batches())
+
+    logger.info(f"to_batches took {to_batches_latency} for {len(pki_table)} rows")
+
+    current_bytes = 0
+    record_batches = []
+    result_len = 0
+    for record_batch in table_batches:
+        current_bytes += record_batch.nbytes
+        record_batches.append(record_batch)
+        if current_bytes >= MAX_SIZE_OF_RECORD_BATCH_IN_GIB:
+            logger.info(
+                f"Total number of record batches without exceeding {MAX_SIZE_OF_RECORD_BATCH_IN_GIB} "
+                f"is {len(record_batches)} and size {current_bytes}"
+            )
+            appended_len, append_latency = timed_invocation(
+                _append_table_by_hash_bucket,
+                pa.Table.from_batches(record_batches),
+                hash_bucket_to_tables,
+            )
+            logger.info(
+                f"Appended the hash bucketed batch of {appended_len} in {append_latency}s"
+            )
+
+            result_len += appended_len
+            current_bytes = 0
+            record_batches.clear()
+
+    if record_batches:
+        appended_len, append_latency = timed_invocation(
+            _append_table_by_hash_bucket,
+            pa.Table.from_batches(record_batches),
+            hash_bucket_to_tables,
+        )
+        result_len += appended_len
+        current_bytes = 0
+        record_batches.clear()
+
+    concat_start = time.monotonic()
+    for hb, tables in enumerate(hash_bucket_to_tables):
+        if tables:
+            assert hb_to_table[hb] is None, f"The HB index is repeated {hb}"
+            hb_to_table[hb] = pa.concat_tables(tables)
+
+    concat_end = time.monotonic()
+    logger.info(
+        f"Total time taken to concat all record batches with length "
+        f"{input_table_len}: {concat_end - concat_start}s"
+    )
+
+    assert (
+        input_table_len == result_len
+    ), f"Grouping has resulted in record loss as {result_len} != {input_table_len}"
+
+    return hb_to_table
+
+
+def group_by_pk_hash_bucket(
+    table: pa.Table, num_buckets: int, primary_keys: List[str]
+) -> np.ndarray:
+    table = generate_pk_hash_column([table], primary_keys, requires_sha1=True)[0]
+
+    # group hash bucket record indices
+    result = group_record_indices_by_hash_bucket(
+        table,
+        num_buckets,
+    )
+
+    return result
+
+
+def generate_pk_hash_column(
+    tables: List[pa.Table],
+    primary_keys: Optional[List[str]] = None,
+    requires_sha1: bool = False,
+) -> List[pa.Table]:
+    """
+    Returns a new table list after generating the primary key hash if desired.
+
+    1. If there are no primary keys, each hash will be unique uuid/sha1 hex
+    2. If there are more than 0 primary keys, returns a table with pk hash column appended.
+    """
+
+    def _generate_pk_hash(table: pa.Table) -> pa.Array:
+        pk_columns = []
+        for pk_name in primary_keys:
+            pk_columns.append(pc.cast(table[pk_name], pa.string()))
+
+        pk_columns.append(PK_DELIMITER)
+        hash_column = pc.binary_join_element_wise(*pk_columns)
+        return hash_column
+
+    def _generate_uuid(table: pa.Table) -> pa.Array:
+        hash_column = pa.array(
+            [uuid.uuid4().hex for _ in range(len(table))], pa.string()
+        )
+        return hash_column
+
+    start = time.monotonic()
+
+    hash_column_list = []
+
+    can_sha1 = False
+    if primary_keys:
+        hash_column_list = [_generate_pk_hash(table) for table in tables]
+
+        can_sha1 = requires_sha1 or _is_sha1_desired(hash_column_list)
+    else:
+        hash_column_list = [_generate_uuid(table) for table in tables]
+
+    logger.info(
+        f"can_generate_sha1={can_sha1} for the table and requires_sha1={requires_sha1}"
+    )
+
+    result = []
+
+    total_len = 0
+    total_size = 0
+    for index, table in enumerate(tables):
+        if can_sha1:
+            table = _append_sha1_hash_to_table(table, hash_column_list[index])
+        else:
+            table = table.append_column(
+                sc._PK_HASH_STRING_COLUMN_FIELD, hash_column_list[index]
+            )
+
+        total_len += len(table)
+        total_size += hash_column_list[index].nbytes
+
+        result.append(table)
+
+    end = time.monotonic()
+
+    logger.info(
+        f"Took {end - start}s to generate pk hash of len: {total_len}"
+        f" for size: {total_size} bytes"
+    )
+
+    return result
+
+
+def group_record_indices_by_hash_bucket(
+    pki_table: pa.Table, num_buckets: int
+) -> np.ndarray:
+    """
+    Groups the record indices by it's corresponding hash bucket. Hence, this method may
+    create num_buckets tables as a result.
+    """
+
+    input_table_len = len(pki_table)
+
+    hash_bucket_id_col_list = np.empty([input_table_len], dtype="int32")
+    bucketing_start_time = time.monotonic()
+
+    for index, hash_value in enumerate(sc.pk_hash_string_column_np(pki_table)):
+        hash_bucket = pk_digest_to_hash_bucket_index(hash_value, num_buckets)
+        hash_bucket_id_col_list[index] = hash_bucket
+
+    pki_table = sc.append_hash_bucket_idx_col(pki_table, hash_bucket_id_col_list)
+    bucketing_end_time = time.monotonic()
+
+    logger.info(
+        f"Took {bucketing_end_time - bucketing_start_time}s to generate the "
+        f"hb index for {len(pki_table)} rows"
+    )
+
+    result, group_latency = timed_invocation(
+        _optimized_group_record_batches_by_hash_bucket,
+        pki_table=pki_table,
+        num_buckets=num_buckets,
+    )
+
+    logger.info(
+        f"Final grouping of table with {input_table_len} records took: {group_latency}s"
+    )
+
+    return result
+
+
+def group_hash_bucket_indices(
+    hash_bucket_object_groups: np.ndarray,
+    num_buckets: int,
+    num_groups: int,
+    object_store: Optional[IObjectStore] = None,
+) -> np.ndarray:
+    """
+    This method persists all tables for a given hash bucket into the object store
+    and returns the object references for each hash group.
+    """
+
+    hash_bucket_group_to_obj_id_size_tuple = np.empty([num_groups], dtype="object")
+
+    if hash_bucket_object_groups is None:
+        return hash_bucket_group_to_obj_id_size_tuple
+
+    hb_group_to_object = np.empty([num_groups], dtype="object")
+    hash_group_to_size = np.empty([num_groups], dtype="int64")
+    hash_group_to_num_rows = np.empty([num_groups], dtype="int64")
+
+    for hb_index, obj in enumerate(hash_bucket_object_groups):
+        if obj:
+            hb_group = hash_bucket_index_to_hash_group_index(hb_index, num_groups)
+            if hb_group_to_object[hb_group] is None:
+                hb_group_to_object[hb_group] = np.empty([num_buckets], dtype="object")
+                hash_group_to_size[hb_group] = np.int64(0)
+                hash_group_to_num_rows[hb_group] = np.int64(0)
+            hb_group_to_object[hb_group][hb_index] = obj
+            for dfe in obj:
+                casted_dfe: DeltaFileEnvelope = dfe
+                hash_group_to_size[hb_group] += casted_dfe.table_size_bytes
+                hash_group_to_num_rows[hb_group] += casted_dfe.table_num_rows
+
+    for hb_group, obj in enumerate(hb_group_to_object):
+        if obj is None:
+            continue
+        object_ref = object_store.put(obj)
+        hash_bucket_group_to_obj_id_size_tuple[hb_group] = (
+            object_ref,
+            hash_group_to_size[hb_group],
+            hash_group_to_num_rows[hb_group],
+        )
+        del object_ref
+    return hash_bucket_group_to_obj_id_size_tuple
+
+
+def hash_bucket_index_to_hash_group_index(hb_index: int, num_groups: int) -> int:
+    return hb_index % num_groups
+
+
+def hash_group_index_to_hash_bucket_indices(
+    hb_group: int, num_buckets: int, num_groups: int
+) -> Iterable[int]:
+
+    if hb_group > num_buckets:
+        return []
+
+    return range(hb_group, num_buckets, num_groups)
+
+
+def pk_digest_to_hash_bucket_index(digest: str, num_buckets: int) -> int:
+    """
+    Generates the hash bucket index from the given digest.
+    """
+    return int(digest, 16) % num_buckets

deltacat/compute/compactor_v2/utils/task_options.py

@@ -0,0 +1,221 @@
+from typing import Dict, Optional, List, Tuple
+from deltacat.types.media import ContentEncoding, ContentType
+from deltacat.types.partial_download import PartialParquetParameters
+from deltacat.storage import (
+    Manifest,
+    ManifestEntry,
+    interface as unimplemented_deltacat_storage,
+)
+from deltacat.compute.compactor.model.delta_annotated import DeltaAnnotated
+from deltacat.compute.compactor.model.round_completion_info import RoundCompletionInfo
+from deltacat.compute.compactor_v2.utils.primary_key_index import (
+    hash_group_index_to_hash_bucket_indices,
+)
+from deltacat.compute.compactor_v2.constants import TOTAL_MEMORY_BUFFER_PERCENTAGE
+
+
+def _get_parquet_type_params_if_exist(
+    entry: ManifestEntry,
+) -> Optional[PartialParquetParameters]:
+    if (
+        entry.meta
+        and entry.meta.content_type == ContentType.PARQUET
+        and entry.meta.content_encoding == ContentEncoding.IDENTITY
+        and entry.meta.content_type_parameters
+    ):
+        for type_params in entry.meta.content_type_parameters:
+            if isinstance(type_params, PartialParquetParameters):
+                return type_params
+    return None
+
+
+def _calculate_parquet_column_size(
+    type_params: PartialParquetParameters, columns: List[str]
+):
+    column_size = 0.0
+    for rg in type_params.row_groups_to_download:
+        columns_found = 0
+        row_group_meta = type_params.pq_metadata.row_group(rg)
+        for col in range(row_group_meta.num_columns):
+            column_meta = row_group_meta.column(col)
+            if column_meta.path_in_schema in columns:
+                columns_found += 1
+                column_size += column_meta.total_uncompressed_size
+        assert columns_found == len(columns), (
+            "Columns not found in the parquet data as "
+            f"{columns_found} != {len(columns)}"
+        )
+    return column_size
+
+
+def estimate_manifest_entry_size_bytes(
+    entry: ManifestEntry, previous_inflation: float, **kwargs
+) -> float:
+    if entry.meta.source_content_length:
+        return entry.meta.source_content_length
+
+    type_params = _get_parquet_type_params_if_exist(entry=entry)
+
+    if type_params:
+        return type_params.in_memory_size_bytes
+
+    return entry.meta.content_length * previous_inflation
+
+
+def estimate_manifest_entry_num_rows(
+    entry: ManifestEntry,
+    average_record_size_bytes: float,
+    previous_inflation: float,
+    **kwargs,
+) -> int:
+    if entry.meta.record_count:
+        return entry.meta.record_count
+
+    type_params = _get_parquet_type_params_if_exist(entry=entry)
+
+    if type_params:
+        return type_params.num_rows
+
+    total_size_bytes = estimate_manifest_entry_size_bytes(
+        entry=entry, previous_inflation=previous_inflation, **kwargs
+    )
+
+    return int(total_size_bytes / average_record_size_bytes)
+
+
+def estimate_manifest_entry_column_size_bytes(
+    entry: ManifestEntry, columns: Optional[List[str]] = None
+) -> Optional[float]:
+    if not columns:
+        return 0
+
+    type_params = _get_parquet_type_params_if_exist(entry=entry)
+
+    if type_params and type_params.pq_metadata:
+        return _calculate_parquet_column_size(type_params=type_params, columns=columns)
+
+    return None
+
+
+def hash_bucket_resource_options_provider(
+    index: int,
+    item: DeltaAnnotated,
+    previous_inflation: float,
+    average_record_size_bytes: float,
+    primary_keys: List[str] = None,
+    **kwargs,
+) -> Dict:
+    size_bytes = 0.0
+    num_rows = 0
+    total_pk_size = 0
+
+    if not item.manifest or not item.manifest.entries:
+        return {"CPU": 0.01}
+
+    for entry in item.manifest.entries:
+        entry_size = estimate_manifest_entry_size_bytes(
+            entry=entry, previous_inflation=previous_inflation
+        )
+        num_rows += estimate_manifest_entry_num_rows(
+            entry=entry,
+            previous_inflation=previous_inflation,
+            average_record_size_bytes=average_record_size_bytes,
+        )
+        size_bytes += entry_size
+
+        if primary_keys:
+            pk_size = estimate_manifest_entry_column_size_bytes(
+                entry=entry,
+                columns=primary_keys,
+            )
+
+            if pk_size is None:
+                total_pk_size += entry_size
+            else:
+                total_pk_size += pk_size
+
+    # total size + pk size + pk hash column + hash bucket index column
+    # Refer to hash_bucket step for more details.
+    total_memory = size_bytes + total_pk_size + num_rows * 20 + num_rows * 4
+
+    # Consider buffer
+    total_memory = total_memory * (1 + TOTAL_MEMORY_BUFFER_PERCENTAGE / 100.0)
+
+    return {"num_cpus": 0.01, "memory": total_memory}
+
+
+def merge_resource_options_provider(
+    index: int,
+    item: Tuple[int, List],
+    num_hash_groups: int,
+    hash_group_size_bytes: Dict[int, int],
+    hash_group_num_rows: Dict[int, int],
+    round_completion_info: Optional[RoundCompletionInfo] = None,
+    compacted_delta_manifest: Optional[Manifest] = None,
+    primary_keys: Optional[List[str]] = None,
+    deltacat_storage=unimplemented_deltacat_storage,
+    deltacat_storage_kwargs: Optional[Dict] = {},
+    **kwargs,
+) -> Dict:
+    hb_group_idx = item[0]
+
+    data_size = hash_group_size_bytes.get(hb_group_idx, 0)
+    num_rows = hash_group_num_rows.get(hb_group_idx, 0)
+
+    pk_size_bytes = 0
+
+    if (
+        round_completion_info
+        and compacted_delta_manifest
+        and round_completion_info.hb_index_to_entry_range
+    ):
+
+        previous_inflation = (
+            round_completion_info.compacted_pyarrow_write_result.pyarrow_bytes
+            / round_completion_info.compacted_pyarrow_write_result.file_bytes
+        )
+        average_record_size = (
+            round_completion_info.compacted_pyarrow_write_result.pyarrow_bytes
+            / round_completion_info.compacted_pyarrow_write_result.records
+        )
+
+        iterable = hash_group_index_to_hash_bucket_indices(
+            hb_group_idx, round_completion_info.hash_bucket_count, num_hash_groups
+        )
+
+        for hb_idx in iterable:
+            entry_start, entry_end = round_completion_info.hb_index_to_entry_range[
+                str(hb_idx)
+            ]
+            for entry_index in range(entry_start, entry_end):
+                entry = compacted_delta_manifest.entries[entry_index]
+
+                current_entry_size = estimate_manifest_entry_size_bytes(
+                    entry=entry, previous_inflation=previous_inflation
+                )
+                current_entry_rows = estimate_manifest_entry_num_rows(
+                    entry=entry,
+                    average_record_size_bytes=average_record_size,
+                    previous_inflation=previous_inflation,
+                )
+
+                data_size += current_entry_size
+                num_rows += current_entry_rows
+
+                if primary_keys:
+                    pk_size = estimate_manifest_entry_column_size_bytes(
+                        entry=entry,
+                        columns=primary_keys,
+                    )
+
+                    if pk_size is None:
+                        pk_size_bytes += current_entry_size
+                    else:
+                        pk_size_bytes += pk_size
+
+    # total data downloaded + primary key hash column + primary key column + dict size for merge
+    total_memory = data_size + pk_size_bytes + num_rows * 20 + num_rows * 20
+
+    total_memory = total_memory * (1 + TOTAL_MEMORY_BUFFER_PERCENTAGE / 100.0)
+
+    return {"num_cpus": 0.01, "memory": total_memory}