deltacat 2.0.0b11__py3-none-any.whl → 2.0.0b12__py3-none-any.whl

deltacat/compute/compactor/utils/round_completion_file.py

@@ -1,97 +0,0 @@
-import json
-import logging
-from typing import Dict, Any
-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
-from deltacat.utils.metrics import metrics
-
-logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
-
-
-def get_round_completion_file_s3_url(
-    bucket: str,
-    source_partition_locator: PartitionLocator,
-    destination_partition_locator: Optional[PartitionLocator] = None,
-) -> str:
-
-    base_url = source_partition_locator.path(f"s3://{bucket}")
-    if destination_partition_locator:
-        base_url = destination_partition_locator.path(
-            f"s3://{bucket}/{source_partition_locator.hexdigest()}"
-        )
-
-    return f"{base_url}.json"
-
-
-@metrics
-def read_round_completion_file(
-    bucket: str,
-    source_partition_locator: PartitionLocator,
-    destination_partition_locator: Optional[PartitionLocator] = None,
-    **s3_client_kwargs: Optional[Dict[str, Any]],
-) -> RoundCompletionInfo:
-
-    all_uris = []
-    if destination_partition_locator:
-        round_completion_file_url_with_destination = get_round_completion_file_s3_url(
-            bucket,
-            source_partition_locator,
-            destination_partition_locator,
-        )
-        all_uris.append(round_completion_file_url_with_destination)
-
-    # Note: we read from RCF at two different URI for backward
-    # compatibility reasons.
-    round_completion_file_url_prev = get_round_completion_file_s3_url(
-        bucket,
-        source_partition_locator,
-    )
-
-    all_uris.append(round_completion_file_url_prev)
-
-    round_completion_info = None
-
-    for rcf_uri in all_uris:
-        logger.info(f"Reading round completion file from: {rcf_uri}")
-        result = s3_utils.download(rcf_uri, False, **s3_client_kwargs)
-        if result:
-            json_str = result["Body"].read().decode("utf-8")
-            round_completion_info = RoundCompletionInfo(json.loads(json_str))
-            logger.info(f"Read round completion info: {round_completion_info}")
-            break
-        else:
-            logger.warning(f"Round completion file not present at {rcf_uri}")
-
-    return round_completion_info
-
-
-@metrics
-def write_round_completion_file(
-    bucket: Optional[str],
-    source_partition_locator: Optional[PartitionLocator],
-    destination_partition_locator: Optional[PartitionLocator],
-    round_completion_info: RoundCompletionInfo,
-    completion_file_s3_url: Optional[str] = None,
-    **s3_client_kwargs: Optional[Dict[str, Any]],
-) -> str:
-    if bucket is None and completion_file_s3_url is None:
-        raise AssertionError("Either bucket or completion_file_s3_url must be passed")
-
-    logger.info(f"writing round completion file contents: {round_completion_info}")
-    if completion_file_s3_url is None:
-        completion_file_s3_url = get_round_completion_file_s3_url(
-            bucket,
-            source_partition_locator,
-            destination_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)),
-        **s3_client_kwargs,
-    )
-    logger.info(f"round completion file written to: {completion_file_s3_url}")
-    return completion_file_s3_url

deltacat/compute/merge_on_read/__init__.py

@@ -1,4 +0,0 @@
-from deltacat.types.media import DistributedDatasetType
-from deltacat.compute.merge_on_read.daft import merge as daft_merge
-
-MERGE_FUNC_BY_DISTRIBUTED_DATASET_TYPE = {DistributedDatasetType.DAFT.value: daft_merge}

deltacat/compute/merge_on_read/daft.py

@@ -1,40 +0,0 @@
-import logging
-from deltacat.compute.merge_on_read.model.merge_on_read_params import MergeOnReadParams
-from deltacat.storage.model.types import DistributedDataset
-from deltacat.types.media import TableType, DistributedDatasetType
-from deltacat.compute.merge_on_read.utils.delta import create_df_from_all_deltas
-from deltacat import logs
-
-logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
-
-
-def merge(params: MergeOnReadParams, **kwargs) -> DistributedDataset:
-    """
-    Merges the given deltas and returns the result as distributed dataframe.
-    It reads the deltas into the Daft dataframe and leverages operations supported
-    by Daft to perform an efficient merge using Ray cluster.
-
-    TODO(raghumdani): Perform actual merge.
-    """
-
-    delta_dfs = create_df_from_all_deltas(
-        deltas=params.deltas,
-        table_type=TableType.PYARROW,
-        distributed_dataset_type=DistributedDatasetType.DAFT,
-        reader_kwargs=params.reader_kwargs,
-        deltacat_storage=params.deltacat_storage,
-        deltacat_storage_kwargs=params.deltacat_storage_kwargs,
-        **kwargs,
-    )
-
-    logger.info(f"Merging {len(delta_dfs)} delta dfs...")
-
-    # TODO: This code should be optimized from daft side
-    result = None
-    for df in delta_dfs:
-        if result is None:
-            result = df
-        else:
-            result = result.concat(df)
-
-    return result

deltacat/compute/merge_on_read/model/merge_on_read_params.py

@@ -1,66 +0,0 @@
-from __future__ import annotations
-from typing import Optional, Dict, List, Union, Any
-from deltacat.storage import (
-    Delta,
-    DeltaLocator,
-    interface as unimplemented_deltacat_storage,
-)
-
-
-class MergeOnReadParams(dict):
-    """
-    This class represents the parameters passed to compact_partition (deltacat/compute/compactor/compaction_session.py)
-    """
-
-    @staticmethod
-    def of(params: Optional[Dict]) -> MergeOnReadParams:
-        params = {} if params is None else params
-
-        result = MergeOnReadParams(params)
-        assert result.deltas is not None, "deltas is a required arg"
-
-        result.deltacat_storage = params.get(
-            "deltacat_storage", unimplemented_deltacat_storage
-        )
-        result.reader_kwargs = params.get("reader_kwargs", {})
-        result.deltacat_storage_kwargs = params.get("deltacat_storage_kwargs", {})
-
-        return result
-
-    @property
-    def deltas(self) -> List[Union[Delta, DeltaLocator]]:
-        """
-        The list of deltas to compact in-memory.
-        """
-        return self["deltas"]
-
-    @deltas.setter
-    def deltas(self, to_set: List[Union[Delta, DeltaLocator]]) -> None:
-        self["deltas"] = to_set
-
-    @property
-    def reader_kwargs(self) -> Dict[Any, Any]:
-        """
-        The key word arguments to be passed to the reader.
-        """
-        return self["reader_kwargs"]
-
-    @reader_kwargs.setter
-    def reader_kwargs(self, kwargs: Dict[Any, Any]) -> None:
-        self["reader_kwargs"] = kwargs
-
-    @property
-    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 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

deltacat/compute/merge_on_read/utils/delta.py

@@ -1,42 +0,0 @@
-from typing import List, Dict, Any, Optional, Union
-from deltacat.storage.model.delta import Delta, DeltaLocator
-from deltacat.storage.model.types import DistributedDataset
-from deltacat.storage import (
-    interface as unimplemented_deltacat_storage,
-)
-from deltacat.types.media import TableType, StorageType, DistributedDatasetType
-
-
-def create_df_from_all_deltas(
-    deltas: List[Union[Delta, DeltaLocator]],
-    table_type: TableType,
-    distributed_dataset_type: DistributedDatasetType,
-    reader_kwargs: Optional[Dict[Any, Any]] = None,
-    deltacat_storage=unimplemented_deltacat_storage,
-    deltacat_storage_kwargs: Optional[Dict[Any, Any]] = None,
-    *args,
-    **kwargs
-) -> List[DistributedDataset]:  # type: ignore
-    """
-    This method creates a distributed dataset for each delta and returns their references.
-    """
-
-    if reader_kwargs is None:
-        reader_kwargs = {}
-    if deltacat_storage_kwargs is None:
-        deltacat_storage_kwargs = {}
-
-    df_list = []
-
-    for delta in deltas:
-        df = deltacat_storage.download_delta(
-            delta_like=delta,
-            table_type=table_type,
-            distributed_dataset_type=distributed_dataset_type,
-            storage_type=StorageType.DISTRIBUTED,
-            **reader_kwargs,
-            **deltacat_storage_kwargs
-        )
-        df_list.append(df)
-
-    return df_list

deltacat/tests/compute/compactor/utils/test_round_completion_file.py

@@ -1,231 +0,0 @@
-import pytest
-import os
-from moto import mock_s3
-import boto3
-from boto3.resources.base import ServiceResource
-from deltacat.compute.compactor.utils.round_completion_file import (
-    read_round_completion_file,
-    write_round_completion_file,
-)
-from deltacat.tests.compute.test_util_common import get_test_partition_locator
-from deltacat.compute.compactor import RoundCompletionInfo
-
-RCF_BUCKET_NAME = "rcf-bucket"
-
-
-@pytest.fixture(autouse=True, scope="module")
-def mock_aws_credential():
-    os.environ["AWS_ACCESS_KEY_ID"] = "testing"
-    os.environ["AWS_SECRET_ACCESS_ID"] = "testing"
-    os.environ["AWS_SECURITY_TOKEN"] = "testing"
-    os.environ["AWS_SESSION_TOKEN"] = "testing"
-    os.environ["AWS_DEFAULT_REGION"] = "us-east-1"
-    yield
-
-
-@pytest.fixture(autouse=True, scope="module")
-def s3_resource(mock_aws_credential):
-    with mock_s3():
-        yield boto3.resource("s3")
-
-
-@pytest.fixture(autouse=True, scope="function")
-def setup_compaction_artifacts_s3_bucket(s3_resource: ServiceResource):
-    s3_resource.create_bucket(
-        ACL="authenticated-read",
-        Bucket=RCF_BUCKET_NAME,
-    )
-    yield
-    s3_resource.Bucket(RCF_BUCKET_NAME).objects.all().delete()
-
-
-class TestReadWriteRoundCompletionFile:
-    def test_read_when_rcf_written_without_destination(self):
-        """
-        This test case tests the backward compatibility by successfully
-        reading the previously written rcf.
-        """
-
-        source_locator = get_test_partition_locator("source")
-        destination_locator = get_test_partition_locator("destination")
-
-        expected_rcf = RoundCompletionInfo.of(
-            high_watermark=122,
-            compacted_delta_locator={},
-            compacted_pyarrow_write_result={},
-            sort_keys_bit_width=12,
-        )
-
-        rcf_url = write_round_completion_file(
-            RCF_BUCKET_NAME, source_locator, None, expected_rcf
-        )
-
-        rcf = read_round_completion_file(
-            RCF_BUCKET_NAME, source_locator, destination_locator
-        )
-
-        assert (
-            rcf_url == "s3://rcf-bucket/f9829af39770d904dbb811bd8f4e886dd307f507.json"
-        )
-        assert rcf == expected_rcf
-
-    def test_read_when_rcf_written_with_destination(self):
-        """
-        This test case tests the backward compatibility by successfully
-        reading the previously written rcf.
-        """
-
-        source_locator = get_test_partition_locator("source")
-        destination_locator = get_test_partition_locator("destination")
-
-        expected_rcf = RoundCompletionInfo.of(
-            high_watermark=122,
-            compacted_delta_locator={},
-            compacted_pyarrow_write_result={},
-            sort_keys_bit_width=12,
-        )
-
-        rcf_url = write_round_completion_file(
-            RCF_BUCKET_NAME, source_locator, destination_locator, expected_rcf
-        )
-
-        rcf = read_round_completion_file(
-            RCF_BUCKET_NAME, source_locator, destination_locator
-        )
-
-        assert (
-            rcf_url
-            == "s3://rcf-bucket/f9829af39770d904dbb811bd8f4e886dd307f507/e9939deadc091b3289a2eb0ca56b1ba86b9892f4.json"
-        )
-        assert rcf == expected_rcf
-
-    def test_read_without_destination_when_rcf_written_with_destination(self):
-        """
-        This test case tests the backward compatibility by successfully
-        reading the previously written rcf.
-        """
-
-        source_locator = get_test_partition_locator("source")
-        destination_locator = get_test_partition_locator("destination")
-
-        expected_rcf = RoundCompletionInfo.of(
-            high_watermark=122,
-            compacted_delta_locator={},
-            compacted_pyarrow_write_result={},
-            sort_keys_bit_width=12,
-        )
-
-        write_round_completion_file(
-            RCF_BUCKET_NAME, source_locator, destination_locator, expected_rcf
-        )
-
-        rcf = read_round_completion_file(RCF_BUCKET_NAME, source_locator, None)
-
-        assert rcf is None
-
-    def test_read_without_destination_when_rcf_written_without_destination(self):
-        """
-        This test case tests the backward compatibility by successfully
-        reading the previously written rcf.
-        """
-
-        source_locator = get_test_partition_locator("source")
-
-        expected_rcf = RoundCompletionInfo.of(
-            high_watermark=122,
-            compacted_delta_locator={},
-            compacted_pyarrow_write_result={},
-            sort_keys_bit_width=12,
-        )
-
-        write_round_completion_file(RCF_BUCKET_NAME, source_locator, None, expected_rcf)
-
-        rcf = read_round_completion_file(RCF_BUCKET_NAME, source_locator, None)
-
-        assert rcf == expected_rcf
-
-    def test_read_when_rcf_written_both_with_and_without_destination(self):
-        """
-        This test case tests the backward compatibility by successfully
-        reading the previously written rcf.
-        """
-
-        source_locator = get_test_partition_locator("source")
-        destination_locator = get_test_partition_locator("destination")
-
-        expected_rcf = RoundCompletionInfo.of(
-            high_watermark=122,
-            compacted_delta_locator={},
-            compacted_pyarrow_write_result={},
-            sort_keys_bit_width=12,
-        )
-
-        expected_rcf_2 = RoundCompletionInfo.of(
-            high_watermark=1223,
-            compacted_delta_locator={},
-            compacted_pyarrow_write_result={},
-            sort_keys_bit_width=1233,
-        )
-
-        write_round_completion_file(RCF_BUCKET_NAME, source_locator, None, expected_rcf)
-
-        write_round_completion_file(
-            RCF_BUCKET_NAME, source_locator, destination_locator, expected_rcf_2
-        )
-
-        rcf = read_round_completion_file(
-            RCF_BUCKET_NAME, source_locator, destination_locator
-        )
-
-        assert rcf == expected_rcf_2
-
-    def test_read_when_none_destination_partition_id(self):
-
-        source_locator = get_test_partition_locator("source")
-        destination_locator = get_test_partition_locator(None)
-
-        expected_rcf = RoundCompletionInfo.of(
-            high_watermark=122,
-            compacted_delta_locator={},
-            compacted_pyarrow_write_result={},
-            sort_keys_bit_width=12,
-        )
-
-        write_round_completion_file(
-            RCF_BUCKET_NAME, source_locator, destination_locator, expected_rcf
-        )
-
-        rcf = read_round_completion_file(
-            RCF_BUCKET_NAME, source_locator, destination_locator
-        )
-
-        assert rcf == expected_rcf
-
-    def test_write_when_custom_url_is_passed(self):
-        """
-        This test case tests the backward compatibility by successfully
-        reading the previously written rcf.
-        """
-
-        source_locator = get_test_partition_locator("source")
-
-        expected_rcf = RoundCompletionInfo.of(
-            high_watermark=122,
-            compacted_delta_locator={},
-            compacted_pyarrow_write_result={},
-            sort_keys_bit_width=12,
-        )
-
-        completion_file_s3_url = f"s3://{RCF_BUCKET_NAME}/test.json"
-        rcf_url = write_round_completion_file(
-            RCF_BUCKET_NAME,
-            source_locator,
-            None,
-            expected_rcf,
-            completion_file_s3_url=completion_file_s3_url,
-        )
-
-        rcf = read_round_completion_file(RCF_BUCKET_NAME, source_locator, None)
-
-        assert rcf_url == completion_file_s3_url
-        assert rcf is None