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

deltacat/compute/compactor/model/compact_partition_params.py

@@ -1,10 +1,29 @@
 from __future__ import annotations
-
+import importlib
 import copy
 import json
 from typing import Any, Dict, List, Optional
-
+from deltacat.io.object_store import IObjectStore
+from deltacat.utils.common import ReadKwargsProvider
 from deltacat.types.media import ContentType
+from deltacat.utils.placement import PlacementGroupConfig
+from deltacat.io.ray_plasma_object_store import RayPlasmaObjectStore
+from deltacat.storage import (
+    interface as unimplemented_deltacat_storage,
+    PartitionLocator,
+    SortKey,
+)
+from deltacat.compute.compactor_v2.constants import (
+    MAX_RECORDS_PER_COMPACTED_FILE,
+    MIN_DELTA_BYTES_IN_BATCH,
+    MIN_FILES_IN_BATCH,
+    AVERAGE_RECORD_SIZE_BYTES,
+    TASK_MAX_PARALLELISM,
+    DROP_DUPLICATES,
+)
+from deltacat.constants import PYARROW_INFLATION_MULTIPLIER
+from deltacat.compute.compactor.utils.sort_key import validate_sort_keys
+from deltacat.utils.metrics import MetricsConfig
 
 
 class CompactPartitionParams(dict):
@@ -16,104 +35,324 @@ class CompactPartitionParams(dict):
     def of(params: Optional[Dict]) -> CompactPartitionParams:
         if params is None:
             params = {}
-        compact_partition_params = CompactPartitionParams()
-        compact_partition_params["destination_partition_locator"] = params.get(
-            "destination_partition_locator"
-        )
-        compact_partition_params["last_stream_position_to_compact"] = params.get(
-            "last_stream_position_to_compact"
+
+        assert (
+            params.get("destination_partition_locator") is not None
+        ), "destination_partition_locator is a required arg"
+        assert (
+            params.get("last_stream_position_to_compact") is not None
+        ), "last_stream_position_to_compact is a required arg"
+        assert (
+            params.get("source_partition_locator") is not None
+        ), "source_partition_locator is a required arg"
+        assert (
+            params.get("compaction_artifact_s3_bucket") is not None
+        ), "compaction_artifact_s3_bucket is a required arg"
+
+        result = CompactPartitionParams(params)
+
+        # TODO: move defaults to single file
+        result.records_per_compacted_file = params.get(
+            "records_per_compacted_file", MAX_RECORDS_PER_COMPACTED_FILE
         )
-        compact_partition_params["source_partition_locator"] = params.get(
-            "source_partition_locator"
+        result.compacted_file_content_type = params.get(
+            "compacted_file_content_type", ContentType.PARQUET
         )
-        compact_partition_params["primary_keys"] = params.get("primary_keys")
-        compact_partition_params["rebase_source_partition_locator"] = params.get(
-            "rebase_source_partition_locator"
+        result.object_store = params.get("object_store", RayPlasmaObjectStore())
+        result.enable_profiler = params.get("enable_profiler", False)
+        result.deltacat_storage = params.get(
+            "deltacat_storage", unimplemented_deltacat_storage
         )
-        compact_partition_params["rebase_source_partition_high_watermark"] = params.get(
-            "rebase_source_partition_high_watermark"
+        result.s3_client_kwargs = params.get("s3_client_kwargs", {})
+        result.deltacat_storage_kwargs = params.get("deltacat_storage_kwargs", {})
+        result.list_deltas_kwargs = params.get("list_deltas_kwargs", {})
+        result.s3_table_writer_kwargs = params.get("s3_table_writer_kwargs", {})
+        result.bit_width_of_sort_keys = validate_sort_keys(
+            result.source_partition_locator,
+            result.sort_keys,
+            result.deltacat_storage,
+            result.deltacat_storage_kwargs,
         )
-        compact_partition_params["hash_bucket_count"] = params.get("hash_bucket_count")
-        compact_partition_params["deltacat_storage"] = params.get("deltacat_storage")
-        compact_partition_params["compaction_artifact_s3_bucket"] = params.get(
-            "compaction_artifact_s3_bucket"
+        result.task_max_parallelism = params.get(
+            "task_max_parallelism", TASK_MAX_PARALLELISM
         )
-        compact_partition_params["properties"] = params.get("properties")
-        compact_partition_params["compacted_file_content_type"] = params.get(
-            "compacted_file_content_type"
+        result.min_files_in_batch = params.get("min_files_in_batch", MIN_FILES_IN_BATCH)
+        result.min_delta_bytes_in_batch = params.get(
+            "min_delta_bytes_in_batch", MIN_DELTA_BYTES_IN_BATCH
         )
-        compact_partition_params["list_deltas_kwargs"] = params.get(
-            "list_deltas_kwargs"
+        result.previous_inflation = params.get(
+            "previous_inflation", PYARROW_INFLATION_MULTIPLIER
         )
-        compact_partition_params["pg_config"] = params.get("pg_config")
-        compact_partition_params["read_kwargs_provider"] = params.get(
-            "read_kwargs_provider"
+        result.average_record_size_bytes = params.get(
+            "average_record_size_bytes", AVERAGE_RECORD_SIZE_BYTES
         )
-        compact_partition_params["s3_table_writer_kwargs"] = params.get(
-            "s3_table_writer_kwargs"
+        result.hash_group_count = params.get(
+            "hash_group_count", result.hash_bucket_count
         )
-        return compact_partition_params
+        result.drop_duplicates = params.get("drop_duplicates", DROP_DUPLICATES)
+
+        if not importlib.util.find_spec("memray"):
+            result.enable_profiler = False
+
+        if result.primary_keys:
+            result.primary_keys = sorted(result.primary_keys)
+
+        # assertions
+        assert (
+            result.source_partition_locator.partition_values
+            == result.destination_partition_locator.partition_values
+        ), "Source and destination partitions values must be equal"
+
+        assert (
+            result.records_per_compacted_file and result.records_per_compacted_file >= 1
+        ), "Max records per output file must be a positive value"
+
+        return result
 
     @property
-    def destination_partition_locator(self) -> Optional[dict]:
-        return self["destination_partition_locator"]
+    def destination_partition_locator(self) -> PartitionLocator:
+        val = self["destination_partition_locator"]
+        if not isinstance(val, PartitionLocator):
+            val = PartitionLocator(val)
+
+        return val
+
+    @destination_partition_locator.setter
+    def destination_partition_locator(self, locator: PartitionLocator) -> None:
+        self["destination_partition_locator"] = locator
 
     @property
-    def last_stream_position_to_compact(self) -> Optional[int]:
+    def last_stream_position_to_compact(self) -> int:
         return self["last_stream_position_to_compact"]
 
+    @last_stream_position_to_compact.setter
+    def last_stream_position_to_compact(self, stream_position: int) -> None:
+        self["last_stream_position_to_compact"] = stream_position
+
     @property
-    def source_partition_locator(self) -> Optional[dict]:
-        return self["source_partition_locator"]
+    def source_partition_locator(self) -> PartitionLocator:
+        val = self["source_partition_locator"]
+        if not isinstance(val, PartitionLocator):
+            val = PartitionLocator(val)
+        return val
+
+    @source_partition_locator.setter
+    def source_partition_locator(self, locator: PartitionLocator) -> None:
+        self["source_partition_locator"] = locator
 
     @property
-    def primary_keys(self) -> Optional[List[str]]:
-        return list(self["primary_keys"])
+    def compaction_artifact_s3_bucket(self) -> str:
+        return self["compaction_artifact_s3_bucket"]
+
+    @compaction_artifact_s3_bucket.setter
+    def compaction_artifact_s3_bucket(self, s3_bucket: str) -> None:
+        self["compaction_artifact_s3_bucket"] = s3_bucket
 
     @property
-    def rebase_source_partition_locator(self) -> Optional[dict]:
-        return self["rebase_source_partition_locator"]
+    def deltacat_storage(self) -> unimplemented_deltacat_storage:
+        return self["deltacat_storage"]
+
+    @deltacat_storage.setter
+    def deltacat_storage(self, storage: unimplemented_deltacat_storage) -> None:
+        self["deltacat_storage"] = storage
 
     @property
-    def rebase_source_partition_high_watermark(self) -> Optional[int]:
-        return self["rebase_source_partition_high_watermark"]
+    def object_store(self) -> IObjectStore:
+        return self["object_store"]
+
+    @object_store.setter
+    def object_store(self, obj_store: IObjectStore) -> None:
+        self["object_store"] = obj_store
 
     @property
-    def hash_bucket_count(self) -> Optional[int]:
-        return self["hash_bucket_count"]
+    def compacted_file_content_type(self) -> ContentType:
+        return self["compacted_file_content_type"]
+
+    @compacted_file_content_type.setter
+    def compacted_file_content_type(self, content_type: ContentType) -> None:
+        self["compacted_file_content_type"] = content_type
 
     @property
-    def deltacat_storage(self) -> Optional[str]:
-        return self["deltacat_storage"]
+    def task_max_parallelism(self) -> int:
+        return self["task_max_parallelism"]
+
+    @task_max_parallelism.setter
+    def task_max_parallelism(self, max_parallelism: int) -> None:
+        self["task_max_parallelism"] = max_parallelism
 
     @property
-    def compaction_artifact_s3_bucket(self) -> Optional[str]:
-        return self["compaction_artifact_s3_bucket"]
+    def average_record_size_bytes(self) -> float:
+        return self["average_record_size_bytes"]
+
+    @average_record_size_bytes.setter
+    def average_record_size_bytes(self, average_record_size_bytes: float) -> None:
+        self["average_record_size_bytes"] = average_record_size_bytes
 
     @property
-    def properties(self) -> Optional[Dict[str, str]]:
-        return self["properties"]
+    def min_files_in_batch(self) -> float:
+        return self["min_files_in_batch"]
+
+    @min_files_in_batch.setter
+    def min_files_in_batch(self, min_files_in_batch: float) -> None:
+        self["min_files_in_batch"] = min_files_in_batch
 
     @property
-    def compacted_file_content_type(self) -> Optional[ContentType]:
-        return self["compacted_file_content_type"]
+    def min_delta_bytes_in_batch(self) -> float:
+        return self["min_delta_bytes_in_batch"]
+
+    @min_delta_bytes_in_batch.setter
+    def min_delta_bytes_in_batch(self, min_delta_bytes_in_batch: float) -> None:
+        self["min_delta_bytes_in_batch"] = min_delta_bytes_in_batch
 
     @property
-    def list_deltas_kwargs(self) -> Optional[dict]:
-        return self["list_deltas_kwargs"]
+    def previous_inflation(self) -> float:
+        return self["previous_inflation"]
+
+    @previous_inflation.setter
+    def previous_inflation(self, previous_inflation: float) -> None:
+        self["previous_inflation"] = previous_inflation
 
     @property
-    def pg_config(self) -> Optional[Any]:
-        return self["pg_config"]
+    def enable_profiler(self) -> bool:
+        return self["enable_profiler"]
+
+    @enable_profiler.setter
+    def enable_profiler(self, value: bool) -> None:
+        self["enable_profiler"] = value
 
     @property
-    def read_kwargs_provider(self) -> Optional[Any]:
-        return self["read_kwargs_provider"]
+    def list_deltas_kwargs(self) -> dict:
+        return self["list_deltas_kwargs"]
+
+    @list_deltas_kwargs.setter
+    def list_deltas_kwargs(self, kwargs: dict) -> None:
+        self["list_deltas_kwargs"] = kwargs
 
     @property
-    def s3_table_writer_kwargs(self) -> Optional[Any]:
+    def s3_table_writer_kwargs(self) -> dict:
         return self["s3_table_writer_kwargs"]
 
+    @s3_table_writer_kwargs.setter
+    def s3_table_writer_kwargs(self, kwargs: dict) -> None:
+        self["s3_table_writer_kwargs"] = kwargs
+
+    @property
+    def deltacat_storage_kwargs(self) -> dict:
+        return self["deltacat_storage_kwargs"]
+
+    @deltacat_storage_kwargs.setter
+    def deltacat_storage_kwargs(self, kwargs: dict) -> None:
+        self["deltacat_storage_kwargs"] = kwargs
+
+    @property
+    def s3_client_kwargs(self) -> dict:
+        return self["s3_client_kwargs"]
+
+    @s3_client_kwargs.setter
+    def s3_client_kwargs(self, kwargs: dict) -> None:
+        self["s3_client_kwargs"] = kwargs
+
+    @property
+    def records_per_compacted_file(self) -> int:
+        return self["records_per_compacted_file"]
+
+    @records_per_compacted_file.setter
+    def records_per_compacted_file(self, count: int) -> None:
+        self["records_per_compacted_file"] = count
+
+    @property
+    def drop_duplicates(self) -> bool:
+        return self["drop_duplicates"]
+
+    @drop_duplicates.setter
+    def drop_duplicates(self, value: bool):
+        self["drop_duplicates"] = value
+
+    @property
+    def bit_width_of_sort_keys(self) -> int:
+        return self["bit_width_of_sort_keys"]
+
+    @bit_width_of_sort_keys.setter
+    def bit_width_of_sort_keys(self, width: int) -> None:
+        self["bit_width_of_sort_keys"] = width
+
+    @property
+    def hash_bucket_count(self) -> Optional[int]:
+        return self.get("hash_bucket_count")
+
+    @hash_bucket_count.setter
+    def hash_bucket_count(self, count: int) -> None:
+        self["hash_bucket_count"] = count
+
+    @property
+    def hash_group_count(self) -> int:
+        return self["hash_group_count"]
+
+    @hash_group_count.setter
+    def hash_group_count(self, count: int) -> None:
+        self["hash_group_count"] = count
+
+    @property
+    def primary_keys(self) -> Optional[List[str]]:
+        return self.get("primary_keys")
+
+    @primary_keys.setter
+    def primary_keys(self, keys: List[str]) -> None:
+        self["primary_keys"] = keys
+
+    @property
+    def rebase_source_partition_locator(self) -> Optional[PartitionLocator]:
+        val = self.get("rebase_source_partition_locator")
+
+        if val and not isinstance(val, PartitionLocator):
+            val = PartitionLocator(val)
+
+        return val
+
+    @rebase_source_partition_locator.setter
+    def rebase_source_partition_locator(self, locator: PartitionLocator) -> None:
+        self["rebase_source_partition_locator"] = locator
+
+    @property
+    def rebase_source_partition_high_watermark(self) -> Optional[int]:
+        return self.get("rebase_source_partition_high_watermark")
+
+    @rebase_source_partition_high_watermark.setter
+    def rebase_source_partition_high_watermark(self, high_watermark: int) -> None:
+        self["rebase_source_partition_high_watermark"] = high_watermark
+
+    @property
+    def pg_config(self) -> Optional[PlacementGroupConfig]:
+        return self.get("pg_config")
+
+    @pg_config.setter
+    def pg_config(self, config: PlacementGroupConfig) -> None:
+        self["pg_config"] = config
+
+    @property
+    def read_kwargs_provider(self) -> Optional[ReadKwargsProvider]:
+        return self.get("read_kwargs_provider")
+
+    @read_kwargs_provider.setter
+    def read_kwargs_provider(self, kwargs_provider: ReadKwargsProvider) -> None:
+        self["read_kwargs_provider"] = kwargs_provider
+
+    @property
+    def sort_keys(self) -> Optional[List[SortKey]]:
+        return self.get("sort_keys")
+
+    @sort_keys.setter
+    def sort_keys(self, keys: List[SortKey]) -> None:
+        self["sort_keys"] = keys
+
+    @property
+    def metrics_config(self) -> Optional[MetricsConfig]:
+        return self.get("metrics_config")
+
+    @metrics_config.setter
+    def metrics_config(self, config: MetricsConfig) -> None:
+        self["metrics_config"] = config
+
     @staticmethod
     def json_handler_for_compact_partition_params(obj):
         """