deltacat 1.1.17__py3-none-any.whl → 1.1.19__py3-none-any.whl

deltacat/compute/compactor_v2/utils/task_options.py

@@ -4,11 +4,8 @@ from deltacat import logs
 from deltacat.compute.compactor_v2.model.merge_file_group import (
     LocalMergeFileGroupsProvider,
 )
-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
@@ -16,50 +13,25 @@ from deltacat.compute.compactor.model.round_completion_info import RoundCompleti
 from deltacat.compute.compactor_v2.utils.primary_key_index import (
     hash_group_index_to_hash_bucket_indices,
 )
-from deltacat.compute.compactor_v2.constants import (
-    PARQUET_TO_PYARROW_INFLATION,
+from deltacat.compute.resource_estimation.manifest import (
+    estimate_manifest_entry_num_rows,
+    estimate_manifest_entry_size_bytes,
+    estimate_manifest_entry_column_size_bytes,
+)
+from deltacat.compute.resource_estimation.model import (
+    EstimateResourcesParams,
+    OperationType,
 )
 from deltacat.exceptions import RetryableError
 
 logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
 
 
-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 * PARQUET_TO_PYARROW_INFLATION
-
-
-def get_task_options(
-    cpu: float, memory: float, ray_custom_resources: Optional[Dict] = None
+def _get_task_options(
+    cpu: float,
+    memory: float,
+    ray_custom_resources: Optional[Dict] = None,
+    scheduling_strategy: str = "SPREAD",
 ) -> Dict:
 
     # NOTE: With DEFAULT scheduling strategy in Ray 2.20.0, autoscaler does
@@ -67,7 +39,11 @@ def get_task_options(
     # 20 tasks get scheduled out of 100 tasks in queue. Hence, we use SPREAD
     # which is also ideal for merge and hash bucket tasks.
     # https://docs.ray.io/en/latest/ray-core/scheduling/index.html
-    task_opts = {"num_cpus": cpu, "memory": memory, "scheduling_strategy": "SPREAD"}
+    task_opts = {
+        "num_cpus": cpu,
+        "memory": memory,
+        "scheduling_strategy": scheduling_strategy,
+    }
 
     if ray_custom_resources:
         task_opts["resources"] = ray_custom_resources
@@ -81,53 +57,110 @@ def get_task_options(
     return task_opts
 
 
-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 * PARQUET_TO_PYARROW_INFLATION
+def _get_merge_task_options(
+    index: int,
+    hb_group_idx: int,
+    data_size: float,
+    pk_size_bytes: float,
+    num_rows: int,
+    num_hash_groups: int,
+    total_memory_buffer_percentage: int,
+    incremental_index_array_size: int,
+    debug_memory_params: Dict[str, Any],
+    ray_custom_resources: Optional[Dict],
+    estimate_resources_params: EstimateResourcesParams,
+    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] = {},
+    memory_logs_enabled: Optional[bool] = None,
+) -> Dict[str, Any]:
+    if (
+        round_completion_info
+        and compacted_delta_manifest
+        and round_completion_info.hb_index_to_entry_range
+    ):
 
-    return entry.meta.content_length * previous_inflation
+        previous_inflation = (
+            round_completion_info.compacted_pyarrow_write_result.pyarrow_bytes
+            / round_completion_info.compacted_pyarrow_write_result.file_bytes
+        )
+        debug_memory_params["previous_inflation"] = previous_inflation
 
+        average_record_size = (
+            round_completion_info.compacted_pyarrow_write_result.pyarrow_bytes
+            / round_completion_info.compacted_pyarrow_write_result.records
+        )
+        debug_memory_params["average_record_size"] = average_record_size
 
-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
+        iterable = hash_group_index_to_hash_bucket_indices(
+            hb_group_idx, round_completion_info.hash_bucket_count, num_hash_groups
+        )
 
-    type_params = _get_parquet_type_params_if_exist(entry=entry)
+        for hb_idx in iterable:
+            if round_completion_info.hb_index_to_entry_range.get(str(hb_idx)) is None:
+                continue
 
-    if type_params:
-        return type_params.num_rows
+            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]
 
-    total_size_bytes = estimate_manifest_entry_size_bytes(
-        entry=entry, previous_inflation=previous_inflation, **kwargs
-    )
+                current_entry_size = estimate_manifest_entry_size_bytes(
+                    entry=entry,
+                    operation_type=OperationType.PYARROW_DOWNLOAD,
+                    estimate_resources_params=estimate_resources_params,
+                )
+                current_entry_rows = estimate_manifest_entry_num_rows(
+                    entry=entry,
+                    operation_type=OperationType.PYARROW_DOWNLOAD,
+                    estimate_resources_params=estimate_resources_params,
+                )
 
-    return int(total_size_bytes / average_record_size_bytes)
+                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,
+                        operation_type=OperationType.PYARROW_DOWNLOAD,
+                        estimate_resources_params=estimate_resources_params,
+                    )
 
-def estimate_manifest_entry_column_size_bytes(
-    entry: ManifestEntry, columns: Optional[List[str]] = None
-) -> Optional[float]:
-    if not columns:
-        return 0
+                    if pk_size is None:
+                        pk_size_bytes += current_entry_size
+                    else:
+                        pk_size_bytes += pk_size
 
-    type_params = _get_parquet_type_params_if_exist(entry=entry)
+    # total data downloaded + primary key hash column + pyarrow-to-numpy conversion
+    # + primary key column + hashlib inefficiency + dict size for merge + incremental index array size
+    total_memory = (
+        data_size
+        + pk_size_bytes
+        + pk_size_bytes
+        + num_rows * 20
+        + num_rows * 20
+        + num_rows * 20
+        + incremental_index_array_size
+    )
+    debug_memory_params["data_size"] = data_size
+    debug_memory_params["num_rows"] = num_rows
+    debug_memory_params["pk_size_bytes"] = pk_size_bytes
+    debug_memory_params["incremental_index_array_size"] = incremental_index_array_size
+    debug_memory_params["total_memory"] = total_memory
+    debug_memory_params["estimate_resources_params"] = estimate_resources_params
 
-    if type_params and type_params.pq_metadata:
-        return _calculate_parquet_column_size(type_params=type_params, columns=columns)
+    total_memory = total_memory * (1 + total_memory_buffer_percentage / 100.0)
+    debug_memory_params["total_memory_with_buffer"] = total_memory
+    logger.debug_conditional(
+        f"[Merge task {index}]: Params used for calculating merge memory: {debug_memory_params}",
+        memory_logs_enabled,
+    )
 
-    return None
+    return _get_task_options(0.01, total_memory, ray_custom_resources)
 
 
 def hash_bucket_resource_options_provider(
@@ -136,6 +169,7 @@ def hash_bucket_resource_options_provider(
     previous_inflation: float,
     average_record_size_bytes: float,
     total_memory_buffer_percentage: int,
+    estimate_resources_params: EstimateResourcesParams,
     primary_keys: List[str] = None,
     ray_custom_resources: Optional[Dict] = None,
     memory_logs_enabled: Optional[bool] = None,
@@ -153,19 +187,25 @@ def hash_bucket_resource_options_provider(
 
     for entry in item.manifest.entries:
         entry_size = estimate_manifest_entry_size_bytes(
-            entry=entry, previous_inflation=previous_inflation
+            entry=entry,
+            operation_type=OperationType.PYARROW_DOWNLOAD,
+            estimate_resources_params=estimate_resources_params,
+            **kwargs,
         )
         num_rows += estimate_manifest_entry_num_rows(
             entry=entry,
-            previous_inflation=previous_inflation,
-            average_record_size_bytes=average_record_size_bytes,
+            operation_type=OperationType.PYARROW_DOWNLOAD,
+            estimate_resources_params=estimate_resources_params,
+            **kwargs,
         )
         size_bytes += entry_size
 
         if primary_keys:
             pk_size = estimate_manifest_entry_column_size_bytes(
                 entry=entry,
+                operation_type=OperationType.PYARROW_DOWNLOAD,
                 columns=primary_keys,
+                estimate_resources_params=estimate_resources_params,
             )
 
             if pk_size is None:
@@ -187,6 +227,7 @@ def hash_bucket_resource_options_provider(
     debug_memory_params["num_rows"] = num_rows
     debug_memory_params["total_pk_size"] = total_pk_size
     debug_memory_params["total_memory"] = total_memory
+    debug_memory_params["estimate_resources_params"] = estimate_resources_params
 
     debug_memory_params["previous_inflation"] = previous_inflation
     debug_memory_params["average_record_size_bytes"] = average_record_size_bytes
@@ -199,7 +240,7 @@ def hash_bucket_resource_options_provider(
         memory_logs_enabled,
     )
 
-    return get_task_options(0.01, total_memory, ray_custom_resources)
+    return _get_task_options(0.01, total_memory, ray_custom_resources)
 
 
 def merge_resource_options_provider(
@@ -209,6 +250,7 @@ def merge_resource_options_provider(
     hash_group_size_bytes: Dict[int, int],
     hash_group_num_rows: Dict[int, int],
     total_memory_buffer_percentage: int,
+    estimate_resources_params: EstimateResourcesParams,
     round_completion_info: Optional[RoundCompletionInfo] = None,
     compacted_delta_manifest: Optional[Manifest] = None,
     ray_custom_resources: Optional[Dict] = None,
@@ -230,7 +272,7 @@ def merge_resource_options_provider(
     pk_size_bytes = data_size
     incremental_index_array_size = num_rows * 4
 
-    return get_merge_task_options(
+    return _get_merge_task_options(
         index,
         hb_group_idx,
         data_size,
@@ -247,6 +289,7 @@ def merge_resource_options_provider(
         deltacat_storage=deltacat_storage,
         deltacat_storage_kwargs=deltacat_storage_kwargs,
         memory_logs_enabled=memory_logs_enabled,
+        estimate_resources_params=estimate_resources_params,
     )
 
 
@@ -254,6 +297,7 @@ def local_merge_resource_options_provider(
     estimated_da_size: float,
     estimated_num_rows: int,
     total_memory_buffer_percentage: int,
+    estimate_resources_params: EstimateResourcesParams,
     round_completion_info: Optional[RoundCompletionInfo] = None,
     compacted_delta_manifest: Optional[Manifest] = None,
     ray_custom_resources: Optional[Dict] = None,
@@ -270,7 +314,7 @@ def local_merge_resource_options_provider(
     pk_size_bytes = estimated_da_size
     incremental_index_array_size = estimated_num_rows * 4
 
-    return get_merge_task_options(
+    return _get_merge_task_options(
         index=index,
         hb_group_idx=hb_group_idx,
         data_size=estimated_da_size,
@@ -287,104 +331,5 @@ def local_merge_resource_options_provider(
         deltacat_storage=deltacat_storage,
         deltacat_storage_kwargs=deltacat_storage_kwargs,
         memory_logs_enabled=memory_logs_enabled,
+        estimate_resources_params=estimate_resources_params,
     )
-
-
-def get_merge_task_options(
-    index: int,
-    hb_group_idx: int,
-    data_size: float,
-    pk_size_bytes: float,
-    num_rows: int,
-    num_hash_groups: int,
-    total_memory_buffer_percentage: int,
-    incremental_index_array_size: int,
-    debug_memory_params: Dict[str, Any],
-    ray_custom_resources: Optional[Dict],
-    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] = {},
-    memory_logs_enabled: Optional[bool] = None,
-) -> Dict[str, Any]:
-    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
-        )
-        debug_memory_params["previous_inflation"] = previous_inflation
-
-        average_record_size = (
-            round_completion_info.compacted_pyarrow_write_result.pyarrow_bytes
-            / round_completion_info.compacted_pyarrow_write_result.records
-        )
-        debug_memory_params["average_record_size"] = average_record_size
-
-        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:
-            if round_completion_info.hb_index_to_entry_range.get(str(hb_idx)) is None:
-                continue
-
-            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 + pyarrow-to-numpy conversion
-    # + primary key column + hashlib inefficiency + dict size for merge + incremental index array size
-    total_memory = (
-        data_size
-        + pk_size_bytes
-        + pk_size_bytes
-        + num_rows * 20
-        + num_rows * 20
-        + num_rows * 20
-        + incremental_index_array_size
-    )
-    debug_memory_params["data_size"] = data_size
-    debug_memory_params["num_rows"] = num_rows
-    debug_memory_params["pk_size_bytes"] = pk_size_bytes
-    debug_memory_params["incremental_index_array_size"] = incremental_index_array_size
-    debug_memory_params["total_memory"] = total_memory
-
-    total_memory = total_memory * (1 + total_memory_buffer_percentage / 100.0)
-    debug_memory_params["total_memory_with_buffer"] = total_memory
-    logger.debug_conditional(
-        f"[Merge task {index}]: Params used for calculating merge memory: {debug_memory_params}",
-        memory_logs_enabled,
-    )
-
-    return get_task_options(0.01, total_memory, ray_custom_resources)

deltacat/compute/resource_estimation/__init__.py

@@ -0,0 +1,27 @@
+from deltacat.compute.resource_estimation.model import (
+    ResourceEstimationMethod,
+    EstimatedResources,
+    Statistics,
+    EstimateResourcesParams,
+    OperationType,
+)
+from deltacat.compute.resource_estimation.manifest import (
+    estimate_manifest_entry_column_size_bytes,
+    estimate_manifest_entry_num_rows,
+    estimate_manifest_entry_size_bytes,
+)
+from deltacat.compute.resource_estimation.delta import (
+    estimate_resources_required_to_process_delta,
+)
+
+__all__ = [
+    "ResourceEstimationMethod",
+    "EstimatedResources",
+    "EstimateResourcesParams",
+    "Statistics",
+    "estimate_resources_required_to_process_delta",
+    "estimate_manifest_entry_size_bytes",
+    "estimate_manifest_entry_num_rows",
+    "estimate_manifest_entry_column_size_bytes",
+    "OperationType",
+]