deltacat 1.1.8__py3-none-any.whl → 1.1.10__py3-none-any.whl

deltacat/tests/compute/compact_partition_rebase_test_cases.py

@@ -0,0 +1,88 @@
+import pyarrow as pa
+from deltacat.tests.compute.test_util_common import (
+    PartitionKey,
+    PartitionKeyType,
+)
+from deltacat.tests.compute.test_util_constant import (
+    DEFAULT_MAX_RECORDS_PER_FILE,
+    DEFAULT_HASH_BUCKET_COUNT,
+)
+from dataclasses import dataclass
+
+
+from deltacat.storage import (
+    DeltaType,
+)
+
+from deltacat.compute.compactor.model.compactor_version import CompactorVersion
+
+from deltacat.storage.model.sort_key import SortKey
+
+from deltacat.tests.compute.compact_partition_test_cases import (
+    BaseCompactorTestCase,
+    with_compactor_version_func_test_param,
+)
+
+
+@dataclass(frozen=True)
+class RebaseCompactionTestCaseParams(BaseCompactorTestCase):
+    """
+    A pytest parameterized test case for the `compact_partition` function with rebase compaction.
+
+    Args:
+        * (inherited from CompactorTestCase): see CompactorTestCase docstring for details
+        rebase_expected_compact_partition_result: pa.Table - expected table after rebase compaction runs. An output that is asserted on in Rebase unit tests
+    """
+
+    rebase_expected_compact_partition_result: pa.Table
+
+
+REBASE_TEST_CASES = {
+    "1-rebase-sanity": RebaseCompactionTestCaseParams(
+        primary_keys={"pk_col_1"},
+        sort_keys=[
+            SortKey.of(key_name="sk_col_1"),
+            SortKey.of(key_name="sk_col_2"),
+        ],
+        partition_keys=[PartitionKey.of("region_id", PartitionKeyType.INT)],
+        partition_values=["1"],
+        input_deltas=pa.Table.from_arrays(
+            [
+                pa.array([str(i) for i in range(10)]),
+                pa.array([i for i in range(0, 10)]),
+                pa.array(["foo"] * 10),
+                pa.array([i / 10 for i in range(10, 20)]),
+            ],
+            names=["pk_col_1", "sk_col_1", "sk_col_2", "col_1"],
+        ),
+        input_deltas_delta_type=DeltaType.UPSERT,
+        rebase_expected_compact_partition_result=pa.Table.from_arrays(
+            [
+                pa.array([str(i) for i in range(10)]),
+                pa.array([i for i in range(0, 10)]),
+                pa.array(["foo"] * 10),
+                pa.array([i / 10 for i in range(10, 20)]),
+            ],
+            names=["pk_col_1", "sk_col_1", "sk_col_2", "col_1"],
+        ),
+        expected_terminal_compact_partition_result=pa.Table.from_arrays(
+            [
+                pa.array([str(i) for i in range(10)]),
+                pa.array([i for i in range(20, 30)]),
+                pa.array(["foo"] * 10),
+                pa.array([i / 10 for i in range(40, 50)]),
+            ],
+            names=["pk_col_1", "sk_col_1", "sk_col_2", "col_1"],
+        ),
+        expected_terminal_exception=None,
+        expected_terminal_exception_message=None,
+        do_create_placement_group=False,
+        records_per_compacted_file=DEFAULT_MAX_RECORDS_PER_FILE,
+        hash_bucket_count=DEFAULT_HASH_BUCKET_COUNT,
+        read_kwargs_provider=None,
+        drop_duplicates=True,
+        skip_enabled_compact_partition_drivers=[CompactorVersion.V1],
+    ),
+}
+
+REBASE_TEST_CASES = with_compactor_version_func_test_param(REBASE_TEST_CASES)

deltacat/tests/compute/compact_partition_rebase_then_incremental_test_cases.py

@@ -37,6 +37,7 @@ from deltacat.tests.compute.compact_partition_test_cases import (
     EMPTY_UTSV_PATH,
 )
 from deltacat.storage import DeleteParameters
+from deltacat.exceptions import ValidationError
 
 
 @dataclass(frozen=True)
@@ -1538,8 +1539,8 @@ REBASE_THEN_INCREMENTAL_DELETE_DELTA_TYPE_TEST_CASES = {
                 ]
             ),
         ),
-        expected_terminal_exception=AssertionError,
-        expected_terminal_exception_message="Delete type deltas are required to have delete parameters defined",
+        expected_terminal_exception=ValidationError,
+        expected_terminal_exception_message="One of the assertions in DeltaCAT has failed",
         do_create_placement_group=True,
         records_per_compacted_file=DEFAULT_MAX_RECORDS_PER_FILE,
         hash_bucket_count=DEFAULT_HASH_BUCKET_COUNT,

deltacat/tests/compute/compact_partition_test_cases.py

@@ -28,6 +28,8 @@ from deltacat.compute.compactor.model.compactor_version import CompactorVersion
 
 from deltacat.storage.model.sort_key import SortKey
 
+from deltacat.exceptions import ValidationError
+
 ZERO_VALUED_SORT_KEY, ZERO_VALUED_PARTITION_VALUES_PARAM = [], []
 ZERO_VALUED_PARTITION_KEYS_PARAM = None
 ZERO_VALUED_PRIMARY_KEY = {}
@@ -570,8 +572,8 @@ INCREMENTAL_TEST_CASES: Dict[str, IncrementalCompactionTestCaseParams] = {
             ],
             names=["pk_col_1", "sk_col_1"],
         ),
-        expected_terminal_exception=AssertionError,
-        expected_terminal_exception_message="hash_bucket_count is a required arg for compactor v2",
+        expected_terminal_exception=ValidationError,
+        expected_terminal_exception_message="One of the assertions in DeltaCAT has failed",
         do_create_placement_group=False,
         records_per_compacted_file=DEFAULT_MAX_RECORDS_PER_FILE,
         hash_bucket_count=None,

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

@@ -0,0 +1,209 @@
+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_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

deltacat/tests/compute/compactor_v2/test_compaction_session.py

@@ -1,88 +1,255 @@
-import unittest
-import sqlite3
+from typing import Dict, Any
 import ray
 import os
-from unittest.mock import patch
+import pytest
+import boto3
+from deltacat.compute.compactor.model.compaction_session_audit_info import (
+    CompactionSessionAuditInfo,
+)
+from boto3.resources.base import ServiceResource
 import deltacat.tests.local_deltacat_storage as ds
 from deltacat.types.media import ContentType
-from deltacat.compute.compactor_v2.compaction_session import compact_partition
+from deltacat.compute.compactor_v2.compaction_session import (
+    compact_partition,
+)
 from deltacat.compute.compactor.model.compact_partition_params import (
     CompactPartitionParams,
 )
-from deltacat.utils.common import current_time_ms
-from deltacat.tests.test_utils.pyarrow import stage_partition_from_file_paths
+from deltacat.tests.test_utils.utils import read_s3_contents
+from deltacat.tests.compute.test_util_constant import (
+    TEST_S3_RCF_BUCKET_NAME,
+)
+from deltacat.tests.compute.test_util_common import get_rcf
+from deltacat.tests.test_utils.pyarrow import (
+    stage_partition_from_file_paths,
+    commit_delta_to_staged_partition,
+    commit_delta_to_partition,
+)
+from moto import mock_s3
+
+DATABASE_FILE_PATH_KEY, DATABASE_FILE_PATH_VALUE = (
+    "db_file_path",
+    "deltacat/tests/local_deltacat_storage/db_test.sqlite",
+)
+
+
+@pytest.fixture(autouse=True, scope="module")
+def setup_ray_cluster():
+    ray.init(local_mode=True, ignore_reinit_error=True)
+    yield
+    ray.shutdown()
+
+
+@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(scope="module")
+def s3_resource(mock_aws_credential):
+    with mock_s3():
+        yield boto3.resource("s3")
 
-class TestCompactionSession(unittest.TestCase):
+
+@pytest.fixture(autouse=True, scope="module")
+def setup_compaction_artifacts_s3_bucket(s3_resource: ServiceResource):
+    s3_resource.create_bucket(
+        ACL="authenticated-read",
+        Bucket=TEST_S3_RCF_BUCKET_NAME,
+    )
+    yield
+
+
+@pytest.fixture(scope="function")
+def local_deltacat_storage_kwargs(request: pytest.FixtureRequest):
+    kwargs_for_local_deltacat_storage: Dict[str, Any] = {
+        DATABASE_FILE_PATH_KEY: DATABASE_FILE_PATH_VALUE,
+    }
+    yield kwargs_for_local_deltacat_storage
+    if os.path.exists(DATABASE_FILE_PATH_VALUE):
+        os.remove(DATABASE_FILE_PATH_VALUE)
+
+
+class TestCompactionSession:
     """
     This class adds specific tests that aren't part of the parametrized test suite.
     """
 
-    DB_FILE_PATH = f"{current_time_ms()}.db"
     NAMESPACE = "compact_partition_v2_namespace"
+    BACKFILL_FILE_PATH = (
+        "deltacat/tests/compute/compactor_v2/data/backfill_source_date_pk.csv"
+    )
+    INCREMENTAL_FILE_PATH = (
+        "deltacat/tests/compute/compactor_v2/data/incremental_source_date_pk.csv"
+    )
+
+    def test_compact_partition_when_no_input_deltas_to_compact(
+        self, local_deltacat_storage_kwargs
+    ):
+        # setup
+        staged_source = stage_partition_from_file_paths(
+            self.NAMESPACE, ["test"], **local_deltacat_storage_kwargs
+        )
+        source_partition = ds.commit_partition(
+            staged_source, **local_deltacat_storage_kwargs
+        )
 
-    @classmethod
-    def setUpClass(cls):
-        ray.init(local_mode=True, ignore_reinit_error=True)
+        staged_dest = stage_partition_from_file_paths(
+            self.NAMESPACE, ["destination"], **local_deltacat_storage_kwargs
+        )
+        dest_partition = ds.commit_partition(
+            staged_dest, **local_deltacat_storage_kwargs
+        )
 
-        con = sqlite3.connect(cls.DB_FILE_PATH)
-        cur = con.cursor()
-        cls.kwargs = {ds.SQLITE_CON_ARG: con, ds.SQLITE_CUR_ARG: cur}
-        cls.deltacat_storage_kwargs = {ds.DB_FILE_PATH_ARG: cls.DB_FILE_PATH}
+        # action
+        rcf_url = compact_partition(
+            CompactPartitionParams.of(
+                {
+                    "compaction_artifact_s3_bucket": TEST_S3_RCF_BUCKET_NAME,
+                    "compacted_file_content_type": ContentType.PARQUET,
+                    "dd_max_parallelism_ratio": 1.0,
+                    "deltacat_storage": ds,
+                    "deltacat_storage_kwargs": local_deltacat_storage_kwargs,
+                    "destination_partition_locator": dest_partition.locator,
+                    "drop_duplicates": True,
+                    "hash_bucket_count": 2,
+                    "last_stream_position_to_compact": source_partition.stream_position,
+                    "list_deltas_kwargs": {
+                        **local_deltacat_storage_kwargs,
+                        **{"equivalent_table_types": []},
+                    },
+                    "primary_keys": ["pk"],
+                    "rebase_source_partition_locator": None,
+                    "rebase_source_partition_high_watermark": None,
+                    "records_per_compacted_file": 4000,
+                    "s3_client_kwargs": {},
+                    "source_partition_locator": source_partition.locator,
+                }
+            )
+        )
 
-        super().setUpClass()
+        # verify that no RCF is written
+        assert rcf_url is None
 
-    @classmethod
-    def doClassCleanups(cls) -> None:
-        os.remove(cls.DB_FILE_PATH)
-        ray.shutdown()
-        super().tearDownClass()
+    def test_compact_partition_when_rcf_was_written_by_past_commit(
+        self, s3_resource, local_deltacat_storage_kwargs
+    ):
+        """
+        Backward compatibility test for when a RCF was written by a previous commit.
+        """
 
-    @patch("deltacat.compute.compactor_v2.compaction_session.rcf")
-    @patch("deltacat.compute.compactor_v2.compaction_session.s3_utils")
-    def test_compact_partition_when_no_input_deltas_to_compact(self, s3_utils, rcf_url):
         # setup
-        rcf_url.read_round_completion_file.return_value = None
         staged_source = stage_partition_from_file_paths(
-            self.NAMESPACE, ["test"], **self.deltacat_storage_kwargs
+            self.NAMESPACE, ["source"], **local_deltacat_storage_kwargs
         )
-        source_partition = ds.commit_partition(
-            staged_source, **self.deltacat_storage_kwargs
+
+        source_delta = commit_delta_to_staged_partition(
+            staged_source, [self.BACKFILL_FILE_PATH], **local_deltacat_storage_kwargs
         )
 
         staged_dest = stage_partition_from_file_paths(
-            self.NAMESPACE, ["destination"], **self.deltacat_storage_kwargs
+            self.NAMESPACE, ["destination"], **local_deltacat_storage_kwargs
         )
         dest_partition = ds.commit_partition(
-            staged_dest, **self.deltacat_storage_kwargs
+            staged_dest, **local_deltacat_storage_kwargs
         )
 
         # action
         rcf_url = compact_partition(
             CompactPartitionParams.of(
                 {
-                    "compaction_artifact_s3_bucket": "test_bucket",
+                    "compaction_artifact_s3_bucket": TEST_S3_RCF_BUCKET_NAME,
                     "compacted_file_content_type": ContentType.PARQUET,
                     "dd_max_parallelism_ratio": 1.0,
                     "deltacat_storage": ds,
-                    "deltacat_storage_kwargs": self.deltacat_storage_kwargs,
+                    "deltacat_storage_kwargs": local_deltacat_storage_kwargs,
                     "destination_partition_locator": dest_partition.locator,
                     "drop_duplicates": True,
                     "hash_bucket_count": 1,
-                    "last_stream_position_to_compact": source_partition.stream_position,
+                    "last_stream_position_to_compact": source_delta.stream_position,
                     "list_deltas_kwargs": {
-                        **self.deltacat_storage_kwargs,
+                        **local_deltacat_storage_kwargs,
                         **{"equivalent_table_types": []},
                     },
                     "primary_keys": [],
+                    "rebase_source_partition_locator": source_delta.partition_locator,
+                    "rebase_source_partition_high_watermark": None,
+                    "records_per_compacted_file": 4000,
+                    "s3_client_kwargs": {},
+                    "source_partition_locator": source_delta.partition_locator,
+                }
+            )
+        )
+
+        bucket, backfill_key1, backfill_key2 = rcf_url.strip("s3://").split("/")
+        assert bucket == TEST_S3_RCF_BUCKET_NAME
+
+        # Now delete the RCF at new location and copy it to old location
+        # Copy the RCF from rcf_url to another location
+        s3_resource.Object(TEST_S3_RCF_BUCKET_NAME, f"{backfill_key1}.json").copy_from(
+            CopySource=f"{TEST_S3_RCF_BUCKET_NAME}/{backfill_key1}/{backfill_key2}"
+        )
+
+        s3_resource.Object(
+            TEST_S3_RCF_BUCKET_NAME, f"{backfill_key1}/{backfill_key2}"
+        ).delete()
+
+        # Now run an incremental compaction and verify if the previous RCF was read properly.
+
+        new_source_delta = commit_delta_to_partition(
+            source_delta.partition_locator,
+            [self.INCREMENTAL_FILE_PATH],
+            **local_deltacat_storage_kwargs,
+        )
+
+        new_rcf_url = compact_partition(
+            CompactPartitionParams.of(
+                {
+                    "compaction_artifact_s3_bucket": TEST_S3_RCF_BUCKET_NAME,
+                    "compacted_file_content_type": ContentType.PARQUET,
+                    "dd_max_parallelism_ratio": 1.0,
+                    "deltacat_storage": ds,
+                    "deltacat_storage_kwargs": local_deltacat_storage_kwargs,
+                    "destination_partition_locator": dest_partition.locator,
+                    "drop_duplicates": True,
+                    "hash_bucket_count": 1,
+                    "last_stream_position_to_compact": new_source_delta.stream_position,
+                    "list_deltas_kwargs": {
+                        **local_deltacat_storage_kwargs,
+                        **{"equivalent_table_types": []},
+                    },
+                    "primary_keys": ["pk"],
                     "rebase_source_partition_locator": None,
                     "rebase_source_partition_high_watermark": None,
                     "records_per_compacted_file": 4000,
                     "s3_client_kwargs": {},
-                    "source_partition_locator": source_partition.locator,
+                    "source_partition_locator": new_source_delta.partition_locator,
                 }
             )
         )
 
-        # verify that no RCF is written
-        self.assertIsNone(rcf_url)
+        new_bucket, incremental_key1, incremental_key2 = new_rcf_url.strip(
+            "s3://"
+        ).split("/")
+
+        assert new_bucket == TEST_S3_RCF_BUCKET_NAME
+        assert backfill_key1 == incremental_key1
+        assert backfill_key2 != incremental_key2
+
+        rcf = get_rcf(s3_resource, new_rcf_url)
+
+        _, compaction_audit_key = rcf.compaction_audit_url.strip("s3://").split("/", 1)
+        compaction_audit = CompactionSessionAuditInfo(
+            **read_s3_contents(
+                s3_resource, TEST_S3_RCF_BUCKET_NAME, compaction_audit_key
+            )
+        )
+
+        # as it should be running incremental
+        assert compaction_audit.uniform_deltas_created == 1
+        assert compaction_audit.input_records == 6