deltacat 0.1.18b1__py3-none-any.whl → 0.1.18b3__py3-none-any.whl

deltacat/compute/compactor/steps/repartition.py

@@ -0,0 +1,210 @@
+import importlib
+import logging
+from contextlib import nullcontext
+import pyarrow.compute as pc
+import pyarrow as pa
+from typing import List, Optional
+from deltacat.types.media import StorageType, ContentType
+import ray
+from deltacat import logs
+from deltacat.compute.compactor import DeltaAnnotated
+from deltacat.compute.compactor.model.repartition_result import RepartitionResult
+from deltacat.storage import interface as unimplemented_deltacat_storage
+from deltacat.storage import Partition
+from deltacat.utils.ray_utils.runtime import (
+    get_current_ray_task_id,
+    get_current_ray_worker_id,
+)
+from deltacat.utils.common import ReadKwargsProvider
+from deltacat.utils.performance import timed_invocation
+from deltacat.utils.metrics import emit_timer_metrics, MetricsConfig
+from deltacat.storage import Delta
+from enum import Enum
+
+if importlib.util.find_spec("memray"):
+    import memray
+
+logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
+
+"""
+Similar to Spark (https://sparkbyexamples.com/spark/spark-partitioning-understanding/), where
+partition helps in localizing the data and reduce the data shuffling across the network nodes reducing network latency
+which is a major component of the transformation operation thereby reducing the time of completion.
+Deltacat with Ray can support different partitioning strategies to reduce the data movement either across network or between compute and storage
+Note that the term partition here is different from the term used in catalog
+Type of Partition:
+Range Partition: It assigns rows to partitions based on column values falling within a given range, e.g., repartition(column="last_updated", ranges=['2023-01-01', '2023-02-01', '2023-03-01']), data will be split into 4 files
+Hash Partition: Hash Partitioning attempts to spread the data evenly across various partitions based on the key, e.g., repartition(column="last_updated", num_partitions=10), data will be split into 10 files evenly
+"""
+
+
+class RepartitionType(str, Enum):
+    RANGE = "range"
+    HASH = "hash"
+
+
+def generate_unique_name(base_name: str, existing_names: List[str]):
+    counter = 1
+    while base_name + str(counter) in existing_names:
+        counter += 1
+    return base_name + str(counter)
+
+
+def repartition_range(
+    tables: List[pa.Table],
+    destination_partition: Partition,
+    repartition_args: dict,
+    max_records_per_output_file: int,
+    repartitioned_file_content_type: ContentType = ContentType.PARQUET,
+    deltacat_storage=unimplemented_deltacat_storage,
+):
+    """
+    Repartitions a list of Arrow tables based on specified ranges and stores the repartitioned tables.
+
+    Args:
+        tables (List[pa.Table]): List of tables to be repartitioned.
+        destination_partition (Partition): The partition to store the repartitioned tables.
+        repartition_args (dict): Arguments for repartitioning. Must include a "column" key for the column to partition
+                                 on and a "ranges" key for a list of partition range values.
+        max_records_per_output_file (int): Maximum number of records per output file.
+        repartitioned_file_content_type (ContentType, optional): The content type of the repartitioned files.
+                                Defaults to ContentType.PARQUET.
+        deltacat_storage (unimplemented): Storage where to put the repartitioned data.
+
+    Raises:
+        ValueError: If no partition ranges are specified or if the column to partition on does not exist in the tables.
+
+    Returns:
+        RepartitionResult: Contains a list of the stored deltas for each partition range.
+    Examples:
+        repartition_args['ranges']= [x, y, z], The tables will be repartitioned into 4 files, i.e., (-inf, x], (x, y], (y, z], (z, inf)
+
+    Note:
+        The function assumes that the tables all share the same schema. If the column to partition on does not exist
+        in the tables, an error will be raised. For each partition range, a new file is created. This could result in
+        more output files than input files.
+    """
+    column: str = repartition_args["column"]
+    partition_ranges: List = repartition_args["ranges"]
+    if len(partition_ranges) == 0:
+        raise ValueError("No partition ranges specified")
+    # check if the column exists in the table
+    # TODO: design a better way to handle the case when the column does not exist in the table, e.g., backfill + repartition by stream position + file id
+    if not all(column in table.column_names for table in tables):
+        raise ValueError(f"Column {column} does not exist in the table")
+    partition_ranges.sort()
+    partition_ranges = [-float("Inf")] + partition_ranges + [float("Inf")]
+    partitioned_tables_list = [[] for _ in range(len(partition_ranges) - 1)]
+
+    total_record_count = 0
+    col_name_int64 = f"{column}_int64"
+    col_name_int64 = generate_unique_name(col_name_int64, tables[0].schema.names)
+    for table in tables:
+        total_record_count += len(table)
+        table_new = table.add_column(
+            0,
+            pa.field(col_name_int64, pa.int64()),
+            pc.cast(table[column], pa.int64()),
+        )
+        # Iterate over pairs of values in partition_ranges
+        for i, (lower_limit, upper_limit) in enumerate(
+            zip(partition_ranges[:-1], partition_ranges[1:]), start=0
+        ):
+            # Add the table filtered by the lower and upper limits to partitioned_tables_list
+            partitioned_tables_list[i].append(
+                table_new.filter(
+                    (pc.field(col_name_int64) > pc.scalar(lower_limit))
+                    & (pc.field(col_name_int64) <= pc.scalar(upper_limit))
+                )
+            )
+    partition_table_length = 0
+    # After re-grouping the tables by specified ranges, for each group, we need concat and stage the tables
+    partition_deltas: List[Delta] = []
+    for partition_tables in partitioned_tables_list:
+        if len(partition_tables) > 0:
+            partition_table: pa.Table = pa.concat_tables(partition_tables)
+            if len(partition_table) > 0:
+                partition_table_length += len(partition_table)
+                partition_delta: Delta = deltacat_storage.stage_delta(
+                    partition_table,
+                    destination_partition,
+                    max_records_per_entry=max_records_per_output_file,
+                    content_type=repartitioned_file_content_type,
+                )
+                partition_deltas.append(partition_delta)
+
+    assert (
+        partition_table_length == total_record_count
+    ), f"Repartitioned table should have the same number of records {partition_table_length} as the original table {total_record_count}"
+    return RepartitionResult(
+        range_deltas=partition_deltas,
+    )
+
+
+def _timed_repartition(
+    annotated_delta: DeltaAnnotated,
+    destination_partition: Partition,
+    repartition_type: RepartitionType,
+    repartition_args: dict,
+    max_records_per_output_file: int,
+    enable_profiler: bool,
+    read_kwargs_provider: Optional[ReadKwargsProvider],
+    repartitioned_file_content_type: ContentType = ContentType.PARQUET,
+    deltacat_storage=unimplemented_deltacat_storage,
+) -> RepartitionResult:
+    task_id = get_current_ray_task_id()
+    worker_id = get_current_ray_worker_id()
+    with memray.Tracker(
+        f"repartition_{worker_id}_{task_id}.bin"
+    ) if enable_profiler else nullcontext():
+        tables: List[pa.Table] = deltacat_storage.download_delta(
+            annotated_delta,
+            storage_type=StorageType.LOCAL,
+            file_reader_kwargs_provider=read_kwargs_provider,
+        )
+        if repartition_type == RepartitionType.RANGE:
+            return repartition_range(
+                tables=tables,
+                destination_partition=destination_partition,
+                repartition_args=repartition_args,
+                max_records_per_output_file=max_records_per_output_file,
+                repartitioned_file_content_type=repartitioned_file_content_type,
+                deltacat_storage=deltacat_storage,
+            )
+        else:
+            raise NotImplementedError(
+                f"Repartition type {repartition_type} is not supported."
+            )
+
+
+@ray.remote
+def repartition(
+    annotated_delta: DeltaAnnotated,
+    destination_partition: Partition,
+    repartition_type: RepartitionType,
+    repartition_args: dict,
+    max_records_per_output_file: int,
+    enable_profiler: bool,
+    metrics_config: Optional[MetricsConfig],
+    read_kwargs_provider: Optional[ReadKwargsProvider],
+    repartitioned_file_content_type: ContentType = ContentType.PARQUET,
+    deltacat_storage=unimplemented_deltacat_storage,
+) -> RepartitionResult:
+    logger.info(f"Starting repartition task...")
+    repartition_result, duration = timed_invocation(
+        func=_timed_repartition,
+        annotated_delta=annotated_delta,
+        destination_partition=destination_partition,
+        repartition_type=repartition_type,
+        repartition_args=repartition_args,
+        max_records_per_output_file=max_records_per_output_file,
+        enable_profiler=enable_profiler,
+        read_kwargs_provider=read_kwargs_provider,
+        repartitioned_file_content_type=repartitioned_file_content_type,
+        deltacat_storage=deltacat_storage,
+    )
+    if metrics_config:
+        emit_timer_metrics(
+            metrics_name="repartition", value=duration, metrics_config=metrics_config
+        )
+    return repartition_result

deltacat/compute/compactor/utils/io.py

@@ -1,7 +1,11 @@
 import logging
 import math
 from deltacat.compute.stats.models.delta_stats import DeltaStats
-from deltacat.constants import PYARROW_INFLATION_MULTIPLIER, BYTES_PER_MEBIBYTE
+from deltacat.constants import (
+    PYARROW_INFLATION_MULTIPLIER,
+    BYTES_PER_MEBIBYTE,
+    MEMORY_TO_HASH_BUCKET_COUNT_RATIO,
+)
 
 from deltacat.storage import (
     PartitionLocator,
@@ -84,54 +88,6 @@ def discover_deltas(
     return input_deltas, previous_last_stream_position_compacted
 
 
-def _discover_deltas(
-    source_partition_locator: PartitionLocator,
-    start_position_exclusive: Optional[int],
-    end_position_inclusive: int,
-    deltacat_storage=unimplemented_deltacat_storage,
-    **kwargs,
-) -> 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,
-        first_stream_position=start_position_exclusive,
-        last_stream_position=end_position_inclusive,
-        ascending_order=True,
-        include_manifest=True,
-        **kwargs,
-    )
-    deltas = deltas_list_result.all_items()
-    if not deltas:
-        raise RuntimeError(
-            f"Unexpected Error: Couldn't find any deltas to "
-            f"compact in delta stream position range "
-            f"('{start_position_exclusive}', "
-            f"'{end_position_inclusive}']. Source partition: "
-            f"{source_partition_locator}"
-        )
-    if start_position_exclusive == deltas[0].stream_position:
-        first_delta = deltas.pop(0)
-        logger.info(
-            f"Removed exclusive start delta w/ expected stream "
-            f"position '{start_position_exclusive}' from deltas to "
-            f"compact: {first_delta}"
-        )
-    logger.info(
-        f"Count of deltas to compact in delta stream "
-        f"position range ('{start_position_exclusive}', "
-        f"'{end_position_inclusive}']: {len(deltas)}. Source "
-        f"partition: '{source_partition_locator}'"
-    )
-    return deltas
-
-
 def limit_input_deltas(
     input_deltas: List[Delta],
     cluster_resources: Dict[str, float],
@@ -285,3 +241,129 @@ def limit_input_deltas(
     logger.info(f"Input uniform delta count: {len(rebatched_da_list)}")
 
     return rebatched_da_list, hash_bucket_count, high_watermark, require_multiple_rounds
+
+
+def fit_input_deltas(
+    input_deltas: List[Delta],
+    cluster_resources: Dict[str, float],
+    hash_bucket_count: Optional[int],
+    deltacat_storage=unimplemented_deltacat_storage,
+) -> Tuple[List[DeltaAnnotated], int, HighWatermark, bool]:
+    """
+    This method tries to fit all the input deltas to run into the existing cluster. Contrary to
+    'limit_input_deltas', it will not fail if the current cluster cannot run the compaction reliably.
+    It is the responsibility of the caller to ensure that they pass enough resources for the job to execute.
+
+    Note: There is a possibility that individual file could be very large, which makes the deltas non uniform.
+    In such scenarios, it is advisable to allocate multiple vCPUs to the tasks to ensure parallelism.
+
+    Args:
+        input_deltas: The input deltas to be normalized.
+        cluster_resources: Total available resources in the cluster.
+        hash_bucket_count: The hash bucket count.
+        deltacat_storage: An implementation of the DeltaCAT storage interface.
+
+    Returns:
+        Tuple of list of annotated deltas, recommended hash bucket count, high watermark,
+            and whether multiple rounds are required (which is always False)
+    """
+    worker_cpus = int(cluster_resources["CPU"])
+    total_memory = float(cluster_resources["memory"])
+    high_watermark = HighWatermark()
+    annotated_input_da_list = []
+    delta_bytes = 0
+    total_files = 0
+
+    if not input_deltas:
+        raise AssertionError("No input deltas found!")
+
+    for delta in input_deltas:
+        manifest_entries = delta.manifest.entries
+        position = delta.stream_position
+
+        for entry in manifest_entries:
+            delta_bytes += entry.meta.content_length
+
+        total_files += len(manifest_entries)
+
+        high_watermark.set(
+            delta.locator.partition_locator,
+            max(position, high_watermark.get(delta.locator.partition_locator)),
+        )
+        delta_annotated = DeltaAnnotated.of(delta)
+        annotated_input_da_list.append(delta_annotated)
+
+    # 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
+
+    # Assuming each CPU consumes equal amount of memory
+    min_delta_bytes = total_memory / worker_cpus
+    rebatched_da_list = DeltaAnnotated.rebatch(
+        annotated_deltas=annotated_input_da_list,
+        min_delta_bytes=min_delta_bytes,
+        estimation_function=estimate_size,
+    )
+
+    # Recommended hash buckets based on the experiments performed
+    # using S3 input for optimal throughput.
+    if hash_bucket_count is None:
+        hash_bucket_count = int(
+            math.ceil(total_memory / MEMORY_TO_HASH_BUCKET_COUNT_RATIO)
+        )
+
+    logger.info(
+        f"Input delta bytes: {delta_bytes}, Total files: {total_files}, The worker_cpus: {worker_cpus}, "
+        f" total_memory: {total_memory}, and hash_bucket_count: {hash_bucket_count}"
+    )
+    return rebatched_da_list, hash_bucket_count, high_watermark, False
+
+
+def _discover_deltas(
+    source_partition_locator: PartitionLocator,
+    start_position_exclusive: Optional[int],
+    end_position_inclusive: int,
+    deltacat_storage=unimplemented_deltacat_storage,
+    **kwargs,
+) -> 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,
+        first_stream_position=start_position_exclusive,
+        last_stream_position=end_position_inclusive,
+        ascending_order=True,
+        include_manifest=True,
+        **kwargs,
+    )
+    deltas = deltas_list_result.all_items()
+    if not deltas:
+        raise RuntimeError(
+            f"Unexpected Error: Couldn't find any deltas to "
+            f"compact in delta stream position range "
+            f"('{start_position_exclusive}', "
+            f"'{end_position_inclusive}']. Source partition: "
+            f"{source_partition_locator}"
+        )
+    if start_position_exclusive == deltas[0].stream_position:
+        first_delta = deltas.pop(0)
+        logger.info(
+            f"Removed exclusive start delta w/ expected stream "
+            f"position '{start_position_exclusive}' from deltas to "
+            f"compact: {first_delta}"
+        )
+    logger.info(
+        f"Count of deltas to compact in delta stream "
+        f"position range ('{start_position_exclusive}', "
+        f"'{end_position_inclusive}']: {len(deltas)}. Source "
+        f"partition: '{source_partition_locator}'"
+    )
+    return deltas

deltacat/compute/compactor/utils/round_completion_file.py

@@ -4,6 +4,8 @@ import logging
 from deltacat import logs
 from deltacat.compute.compactor import RoundCompletionInfo
 from deltacat.storage import PartitionLocator
+from deltacat.aws import s3u as s3_utils
+from typing import Optional
 
 logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
 
@@ -20,8 +22,6 @@ def read_round_completion_file(
     bucket: str, source_partition_locator: PartitionLocator
 ) -> RoundCompletionInfo:
 
-    from deltacat.aws import s3u as s3_utils
-
     round_completion_file_url = get_round_completion_file_s3_url(
         bucket,
         source_partition_locator,
@@ -37,21 +37,19 @@ def read_round_completion_file(
 
 
 def write_round_completion_file(
-    bucket: str,
-    source_partition_locator: PartitionLocator,
+    bucket: Optional[str],
+    source_partition_locator: Optional[PartitionLocator],
     round_completion_info: RoundCompletionInfo,
+    completion_file_s3_url: str = None,
 ) -> str:
 
-    from deltacat.aws import s3u as s3_utils
-
     logger.info(f"writing round completion file contents: {round_completion_info}")
-    round_completion_file_s3_url = get_round_completion_file_s3_url(
-        bucket,
-        source_partition_locator,
-    )
-    logger.info(f"writing round completion file to: {round_completion_file_s3_url}")
-    s3_utils.upload(
-        round_completion_file_s3_url, str(json.dumps(round_completion_info))
-    )
-    logger.info(f"round completion file written to: {round_completion_file_s3_url}")
-    return round_completion_file_s3_url
+    if completion_file_s3_url is None:
+        completion_file_s3_url = get_round_completion_file_s3_url(
+            bucket,
+            source_partition_locator,
+        )
+    logger.info(f"writing round completion file to: {completion_file_s3_url}")
+    s3_utils.upload(completion_file_s3_url, str(json.dumps(round_completion_info)))
+    logger.info(f"round completion file written to: {completion_file_s3_url}")
+    return completion_file_s3_url

deltacat/constants.py

@@ -52,3 +52,5 @@ PYARROW_INFLATION_MULTIPLIER_ALL_COLUMNS = 6
 PRIMARY_KEY_INDEX_WRITE_BOTO3_CONFIG = {
     "retries": {"max_attempts": 25, "mode": "standard"}
 }
+
+MEMORY_TO_HASH_BUCKET_COUNT_RATIO = 0.0512 * BYTES_PER_TEBIBYTE

deltacat/storage/interface.py

@@ -443,7 +443,7 @@ def get_partition(
 
 
 def stage_delta(
-    data: Union[LocalTable, LocalDataset, DistributedDataset],
+    data: Union[LocalTable, LocalDataset, DistributedDataset, Manifest],
     partition: Partition,
     delta_type: DeltaType = DeltaType.UPSERT,
     max_records_per_entry: Optional[int] = None,

deltacat/storage/model/types.py

@@ -1,15 +1,23 @@
 from enum import Enum
-from typing import Any, List, Union
+from typing import List, Union, Any
 
 import numpy as np
 import pandas as pd
 import pyarrow as pa
+import pkg_resources
 from ray.data._internal.arrow_block import ArrowRow
 from ray.data.dataset import Dataset
 
 LocalTable = Union[pa.Table, pd.DataFrame, np.ndarray]
 LocalDataset = List[LocalTable]
-DistributedDataset = Dataset[Union[ArrowRow, np.ndarray, Any]]
+# Starting Ray 2.5.0, Dataset follows a strict mode (https://docs.ray.io/en/latest/data/faq.html#migrating-to-strict-mode),
+# and generic annotation is removed. So add a version checker to determine whether to use the old or new definition.
+ray_version = pkg_resources.parse_version(pkg_resources.get_distribution("ray").version)
+change_version = pkg_resources.parse_version("2.5.0")
+if ray_version < change_version:
+    DistributedDataset = Dataset[Union[ArrowRow, np.ndarray, Any]]
+else:
+    DistributedDataset = Dataset
 
 
 class DeltaType(str, Enum):

deltacat/tests/compactor/utils/test_io.py

@@ -0,0 +1,69 @@
+import unittest
+from unittest import mock
+from deltacat.tests.test_utils.constants import TEST_DELTA
+
+
+class TestFitInputDeltas(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.module_patcher = mock.patch.dict("sys.modules", {"ray": mock.MagicMock()})
+        cls.module_patcher.start()
+
+        super().setUpClass()
+
+    def test_sanity(self):
+        from deltacat.compute.compactor.utils import io
+
+        (
+            delta_list,
+            hash_bucket_count,
+            high_watermark,
+            require_multiple_rounds,
+        ) = io.fit_input_deltas([TEST_DELTA], {"CPU": 1, "memory": 20000000}, None)
+
+        self.assertIsNotNone(hash_bucket_count)
+        self.assertTrue(1, len(delta_list))
+        self.assertIsNotNone(high_watermark)
+        self.assertFalse(require_multiple_rounds)
+
+    def test_when_hash_bucket_count_overridden(self):
+        from deltacat.compute.compactor.utils import io
+
+        (
+            delta_list,
+            hash_bucket_count,
+            high_watermark,
+            require_multiple_rounds,
+        ) = io.fit_input_deltas([TEST_DELTA], {"CPU": 1, "memory": 20000000}, 20)
+
+        self.assertEqual(20, hash_bucket_count)
+        self.assertEqual(1, len(delta_list))
+        self.assertIsNotNone(high_watermark)
+        self.assertFalse(require_multiple_rounds)
+
+    def test_when_not_enough_memory_splits_manifest_entries(self):
+        from deltacat.compute.compactor.utils import io
+
+        (
+            delta_list,
+            hash_bucket_count,
+            high_watermark,
+            require_multiple_rounds,
+        ) = io.fit_input_deltas([TEST_DELTA], {"CPU": 2, "memory": 10}, 20)
+
+        self.assertIsNotNone(hash_bucket_count)
+        self.assertTrue(2, len(delta_list))
+        self.assertIsNotNone(high_watermark)
+        self.assertFalse(require_multiple_rounds)
+
+    def test_when_no_input_deltas(self):
+        from deltacat.compute.compactor.utils import io
+
+        with self.assertRaises(AssertionError):
+            io.fit_input_deltas([], {"CPU": 100, "memory": 20000.0}, None)
+
+    def test_when_cpu_resources_is_not_passed(self):
+        from deltacat.compute.compactor.utils import io
+
+        with self.assertRaises(KeyError):
+            io.fit_input_deltas([], {}, None)