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

deltacat/compute/compactor/utils/io.py

@@ -10,11 +10,12 @@ from deltacat.constants import (
 from deltacat.storage import (
     PartitionLocator,
     Delta,
+    ManifestEntry,
     interface as unimplemented_deltacat_storage,
 )
 from deltacat import logs
 from deltacat.compute.compactor import DeltaAnnotated
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Union, Any
 from deltacat.compute.compactor import HighWatermark
 from deltacat.compute.compactor.model.compaction_session_audit_info import (
     CompactionSessionAuditInfo,
@@ -31,23 +32,30 @@ def discover_deltas(
     rebase_source_partition_locator: Optional[PartitionLocator],
     rebase_source_partition_high_watermark: Optional[int],
     deltacat_storage=unimplemented_deltacat_storage,
-    **kwargs,
+    deltacat_storage_kwargs: Optional[Dict[str, Any]] = {},
+    list_deltas_kwargs: Optional[Dict[str, Any]] = {},
 ) -> Tuple[List[Delta], int]:
-
+    if deltacat_storage_kwargs is None:
+        deltacat_storage_kwargs = {}
     # Source One: new deltas from uncompacted table for incremental compaction or deltas from compacted table for rebase
-    input_deltas = _discover_deltas(
-        source_partition_locator,
+    start_position_exclusive = (
         high_watermark.get(source_partition_locator)
         if isinstance(high_watermark, dict)
-        else high_watermark,
+        else high_watermark
+    )
+    input_deltas = _discover_deltas(
+        source_partition_locator,
+        start_position_exclusive,
         last_stream_position_to_compact
         if not rebase_source_partition_locator
         else deltacat_storage.get_partition(
             source_partition_locator.stream_locator,
             source_partition_locator.partition_values,
+            **deltacat_storage_kwargs,
         ).stream_position,
         deltacat_storage,
-        **kwargs,
+        deltacat_storage_kwargs,
+        list_deltas_kwargs,
     )
 
     # Source Two: delta from compacted table for incremental compaction or new deltas from uncompacted table for rebase
@@ -56,6 +64,7 @@ def discover_deltas(
         compacted_partition = deltacat_storage.get_partition(
             compacted_partition_locator.stream_locator,
             compacted_partition_locator.partition_values,
+            **deltacat_storage_kwargs,
         )
         previous_last_stream_position_compacted = (
             compacted_partition.stream_position if compacted_partition else -1
@@ -67,7 +76,8 @@ def discover_deltas(
                 None,
                 previous_last_stream_position_compacted,
                 deltacat_storage,
-                **kwargs,
+                deltacat_storage_kwargs,
+                list_deltas_kwargs,
             )
         logger.info(
             f"Length of input deltas from uncompacted table {len(input_deltas)} up to {last_stream_position_to_compact},"
@@ -80,7 +90,8 @@ def discover_deltas(
             rebase_source_partition_high_watermark,
             last_stream_position_to_compact,
             deltacat_storage,
-            **kwargs,
+            deltacat_storage_kwargs,
+            list_deltas_kwargs,
         )
         logger.info(
             f"Length of input deltas from uncompacted table {len(input_deltas_new)} up to {last_stream_position_to_compact},"
@@ -99,6 +110,8 @@ def limit_input_deltas(
     input_deltas_stats: Dict[int, DeltaStats],
     compaction_audit: CompactionSessionAuditInfo,
     deltacat_storage=unimplemented_deltacat_storage,
+    deltacat_storage_kwargs: Optional[Dict[str, Any]] = None,
+    **kwargs,
 ) -> Tuple[List[DeltaAnnotated], int, HighWatermark, bool]:
     # TODO (pdames): when row counts are available in metadata, use them
     #  instead of bytes - memory consumption depends more on number of
@@ -108,6 +121,8 @@ def limit_input_deltas(
     # this assumption could be removed, but we'd still need to know the max
     # resources we COULD get for this cluster, and the amount of memory
     # available per CPU should remain fixed across the cluster.
+    if deltacat_storage_kwargs is None:
+        deltacat_storage_kwargs = {}
     worker_cpus = int(cluster_resources["CPU"])
     worker_obj_store_mem = float(cluster_resources["object_store_memory"])
     logger.info(f"Total worker object store memory: {worker_obj_store_mem}")
@@ -135,7 +150,7 @@ def limit_input_deltas(
         for stream_pos, delta_stats in input_deltas_stats.items()
     }
     for delta in input_deltas:
-        manifest = deltacat_storage.get_delta_manifest(delta)
+        manifest = deltacat_storage.get_delta_manifest(delta, **deltacat_storage_kwargs)
         delta.manifest = manifest
         position = delta.stream_position
         delta_stats = input_deltas_stats.get(delta.stream_position, DeltaStats())
@@ -258,6 +273,8 @@ def fit_input_deltas(
     compaction_audit: CompactionSessionAuditInfo,
     hash_bucket_count: Optional[int],
     deltacat_storage=unimplemented_deltacat_storage,
+    deltacat_storage_kwargs: Optional[Dict[str, Any]] = None,
+    **kwargs,
 ) -> Tuple[List[DeltaAnnotated], int, HighWatermark, bool]:
     """
     This method tries to fit all the input deltas to run into the existing cluster. Contrary to
@@ -277,6 +294,8 @@ def fit_input_deltas(
         Tuple of list of annotated deltas, recommended hash bucket count, high watermark,
             and whether multiple rounds are required (which is always False)
     """
+    if deltacat_storage_kwargs is None:
+        deltacat_storage_kwargs = {}
     worker_cpus = int(cluster_resources["CPU"])
     total_memory = float(cluster_resources["memory"])
     high_watermark = HighWatermark()
@@ -306,8 +325,8 @@ def fit_input_deltas(
     # We assume that the cluster is capable of distributing all tasks
     # correctly. Hence, the correct in-memory size will be in the ratio of
     # in-disk size.
-    def estimate_size(content_length):
-        return (content_length * 1.0 / delta_bytes) * total_memory
+    def estimate_size(manifest_entry: ManifestEntry):
+        return (manifest_entry.meta.content_length * 1.0 / delta_bytes) * total_memory
 
     # Assuming each CPU consumes equal amount of memory
     min_delta_bytes = total_memory / worker_cpus
@@ -341,18 +360,16 @@ def _discover_deltas(
     start_position_exclusive: Optional[int],
     end_position_inclusive: int,
     deltacat_storage=unimplemented_deltacat_storage,
-    **kwargs,
+    deltacat_storage_kwargs: Optional[Dict[str, Any]] = {},
+    list_deltas_kwargs: Optional[Dict[str, Any]] = {},
 ) -> List[Delta]:
-    stream_locator = source_partition_locator.stream_locator
-    namespace = stream_locator.namespace
-    table_name = stream_locator.table_name
-    table_version = stream_locator.table_version
-    partition_values = source_partition_locator.partition_values
-    deltas_list_result = deltacat_storage.list_deltas(
-        namespace=namespace,
-        table_name=table_name,
-        partition_values=partition_values,
-        table_version=table_version,
+    if deltacat_storage_kwargs is None:
+        deltacat_storage_kwargs = {}
+
+    kwargs = {**deltacat_storage_kwargs, **list_deltas_kwargs}
+
+    deltas_list_result = deltacat_storage.list_partition_deltas(
+        partition_like=source_partition_locator,
         first_stream_position=start_position_exclusive,
         last_stream_position=end_position_inclusive,
         ascending_order=True,

deltacat/compute/compactor/utils/primary_key_index.py

@@ -6,26 +6,12 @@ import pyarrow as pa
 from ray.types import ObjectRef
 
 from deltacat import logs
-from deltacat.aws import s3u
-from deltacat.compute.compactor import (
-    PrimaryKeyIndexVersionLocator,
-)
 from deltacat.compute.compactor.utils import system_columns as sc
 from deltacat.io.object_store import IObjectStore
 
 logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
 
-
-def delete_primary_key_index_version(
-    s3_bucket: str, pki_version_locator: PrimaryKeyIndexVersionLocator
-) -> None:
-
-    logger.info(f"Deleting primary key index: {pki_version_locator}")
-    s3u.delete_files_by_prefix(
-        s3_bucket,
-        pki_version_locator.primary_key_index_version_root_path,
-    )
-    logger.info(f"Primary key index deleted: {pki_version_locator}")
+# TODO: Deprecate this module in the favor of compactor_v2
 
 
 def group_record_indices_by_hash_bucket(

deltacat/compute/compactor/utils/sort_key.py

@@ -0,0 +1,57 @@
+import pyarrow as pa
+from typing import List
+from deltacat.storage import PartitionLocator, SortKey
+
+MAX_SORT_KEYS_BIT_WIDTH = 256
+
+
+def validate_sort_keys(
+    source_partition_locator: PartitionLocator,
+    sort_keys: List[SortKey],
+    deltacat_storage,
+    deltacat_storage_kwargs,
+    **kwargs,
+) -> int:
+    """
+    Validates the input sort keys to ensure that they are unique, are using
+    a valid sort key model, are all fixed-width data types, and that the
+    sum of bit widths across sort key data types is less-than-or-equal-to
+    256. Returns the sum of bit widths across all sort keys.
+    """
+    if deltacat_storage_kwargs is None:
+        deltacat_storage_kwargs = {}
+    total_sort_keys_bit_width = 0
+    if sort_keys:
+        sort_key_names = [key.key_name for key in sort_keys]
+        assert len(sort_key_names) == len(
+            set(sort_key_names)
+        ), f"Sort key names must be unique: {sort_key_names}"
+        stream_locator = source_partition_locator.stream_locator
+        table_version_schema = deltacat_storage.get_table_version_schema(
+            stream_locator.namespace,
+            stream_locator.table_name,
+            stream_locator.table_version,
+            **deltacat_storage_kwargs,
+        )
+        if isinstance(table_version_schema, pa.Schema):
+            for sort_key_name in sort_key_names:
+                pa_field: pa.Field = pa.Schema.field(sort_key_name)
+                pa_type: pa.DataType = pa_field.type
+                try:
+                    total_sort_keys_bit_width += pa_type.bit_width
+                    if total_sort_keys_bit_width > MAX_SORT_KEYS_BIT_WIDTH:
+                        raise ValueError(
+                            f"Total length of sort keys "
+                            f"({total_sort_keys_bit_width}) is greater "
+                            f"than the max supported bit width for all "
+                            f"sort keys ({MAX_SORT_KEYS_BIT_WIDTH})"
+                        )
+                except ValueError as e:
+                    raise ValueError(
+                        f"Unable to get bit width of sort key: {pa_field}. "
+                        f"Please ensure that all sort keys are fixed-size "
+                        f"PyArrow data types."
+                    ) from e
+        else:
+            total_sort_keys_bit_width = MAX_SORT_KEYS_BIT_WIDTH
+    return total_sort_keys_bit_width

deltacat/compute/compactor/utils/system_columns.py

@@ -22,6 +22,13 @@ _PK_HASH_COLUMN_FIELD = pa.field(
     _PK_HASH_COLUMN_TYPE,
 )
 
+_PK_HASH_STRING_COLUMN_NAME = _get_sys_col_name("hash_str")
+_PK_HASH_STRING_COLUMN_TYPE = pa.string()
+_PK_HASH_STRING_COLUMN_FIELD = pa.field(
+    _PK_HASH_STRING_COLUMN_NAME,
+    _PK_HASH_STRING_COLUMN_TYPE,
+)
+
 _DEDUPE_TASK_IDX_COLUMN_NAME = _get_sys_col_name("dedupe_task_idx")
 _DEDUPE_TASK_IDX_COLUMN_TYPE = pa.int32()
 _DEDUPE_TASK_IDX_COLUMN_FIELD = pa.field(
@@ -36,6 +43,12 @@ _PARTITION_STREAM_POSITION_COLUMN_FIELD = pa.field(
     _PARTITION_STREAM_POSITION_COLUMN_TYPE,
 )
 
+_HASH_BUCKET_IDX_COLUMN_NAME = _get_sys_col_name("hash_bucket_idx")
+_HASH_BUCKET_IDX_COLUMN_TYPE = pa.int32()
+_HASH_BUCKET_IDX_COLUMN_FIELD = pa.field(
+    _HASH_BUCKET_IDX_COLUMN_NAME, _HASH_BUCKET_IDX_COLUMN_TYPE
+)
+
 _ORDERED_FILE_IDX_COLUMN_NAME = _get_sys_col_name("file_index")
 _ORDERED_FILE_IDX_COLUMN_TYPE = pa.int32()
 _ORDERED_FILE_IDX_COLUMN_FIELD = pa.field(
@@ -76,10 +89,18 @@ def get_pk_hash_column_array(obj) -> Union[pa.Array, pa.ChunkedArray]:
     return pa.array(obj, _PK_HASH_COLUMN_TYPE)
 
 
+def get_pk_hash_string_column_array(obj) -> Union[pa.Array, pa.ChunkedArray]:
+    return pa.array(obj, _PK_HASH_STRING_COLUMN_TYPE)
+
+
 def pk_hash_column_np(table: pa.Table) -> np.ndarray:
     return table[_PK_HASH_COLUMN_NAME].to_numpy()
 
 
+def pk_hash_string_column_np(table: pa.Table) -> np.ndarray:
+    return table[_PK_HASH_STRING_COLUMN_NAME].to_numpy()
+
+
 def pk_hash_column(table: pa.Table) -> pa.ChunkedArray:
     return table[_PK_HASH_COLUMN_NAME]
 
@@ -143,6 +164,10 @@ def get_delta_type_column_array(obj) -> Union[pa.Array, pa.ChunkedArray]:
     )
 
 
+def get_hash_bucket_idx_column_array(obj) -> Union[pa.Array, pa.ChunkedArray]:
+    return pa.array(obj, _HASH_BUCKET_IDX_COLUMN_TYPE)
+
+
 def get_is_source_column_array(obj) -> Union[pa.Array, pa.ChunkedArray]:
     return pa.array(
         obj,
@@ -232,6 +257,24 @@ def append_pk_hash_column(table: pa.Table, pk_hashes) -> pa.Table:
     return table
 
 
+def append_pk_hash_string_column(table: pa.Table, pk_hashes) -> pa.Table:
+
+    table = table.append_column(
+        _PK_HASH_STRING_COLUMN_FIELD, get_pk_hash_string_column_array(pk_hashes)
+    )
+    return table
+
+
+def append_hash_bucket_idx_col(table: pa.Table, hash_bucket_indexes) -> pa.Table:
+
+    table = table.append_column(
+        _HASH_BUCKET_IDX_COLUMN_FIELD,
+        get_hash_bucket_idx_column_array(hash_bucket_indexes),
+    )
+
+    return table
+
+
 def append_record_idx_col(table: pa.Table, ordered_record_indices) -> pa.Table:
 
     table = table.append_column(