deltacat 1.1.36__py3-none-any.whl → 2.0.0b2__py3-none-any.whl

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

@@ -1,12 +1,11 @@
 import unittest
 from unittest import mock
-from deltacat.tests.test_utils.constants import TEST_UPSERT_DELTA
-from typing import Any, Dict
 
-DATABASE_FILE_PATH_KEY, DATABASE_FILE_PATH_VALUE = (
-    "db_file_path",
-    "deltacat/tests/local_deltacat_storage/db_test.sqlite",
+from deltacat.tests.compute.conftest import (
+    create_local_deltacat_storage_file,
+    clean_up_local_deltacat_storage_file,
 )
+from deltacat.tests.test_utils.constants import TEST_UPSERT_DELTA
 
 
 class TestFitInputDeltas(unittest.TestCase):
@@ -19,9 +18,7 @@ class TestFitInputDeltas(unittest.TestCase):
             CompactionSessionAuditInfo,
         )
 
-        cls.kwargs_for_local_deltacat_storage: Dict[str, Any] = {
-            DATABASE_FILE_PATH_KEY: DATABASE_FILE_PATH_VALUE,
-        }
+        cls.kwargs_for_local_deltacat_storage = create_local_deltacat_storage_file()
 
         cls.COMPACTION_AUDIT = CompactionSessionAuditInfo("1.0", "2.3", "test")
 
@@ -30,6 +27,7 @@ class TestFitInputDeltas(unittest.TestCase):
     @classmethod
     def tearDownClass(cls) -> None:
         cls.module_patcher.stop()
+        clean_up_local_deltacat_storage_file(cls.kwargs_for_local_deltacat_storage)
 
     def test_sanity(self):
         from deltacat.compute.compactor.utils import io

deltacat/tests/compute/compactor_v2/test_compaction_session.py

@@ -1,14 +1,10 @@
-from typing import Dict, Any
 import ray
 import os
-import pyarrow as pa
 import pytest
 import boto3
-import json
 from deltacat.compute.compactor.model.compaction_session_audit_info import (
     CompactionSessionAuditInfo,
 )
-from deltacat.exceptions import ValidationError
 from boto3.resources.base import ServiceResource
 import deltacat.tests.local_deltacat_storage as ds
 from deltacat.types.media import ContentType
@@ -31,11 +27,6 @@ from deltacat.tests.test_utils.pyarrow import (
 )
 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():
@@ -69,38 +60,6 @@ def setup_compaction_artifacts_s3_bucket(s3_resource: ServiceResource):
     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)
-
-
-@pytest.fixture(scope="function")
-def disable_sha1(monkeypatch):
-    import deltacat.compute.compactor_v2.utils.primary_key_index
-
-    monkeypatch.setattr(
-        deltacat.compute.compactor_v2.utils.primary_key_index,
-        "SHA1_HASHING_FOR_MEMORY_OPTIMIZATION_DISABLED",
-        True,
-    )
-
-
-@pytest.fixture(scope="function")
-def enable_bucketing_spec_validation(monkeypatch):
-    import deltacat.compute.compactor_v2.steps.merge
-
-    monkeypatch.setattr(
-        deltacat.compute.compactor_v2.steps.merge,
-        "BUCKETING_SPEC_COMPLIANCE_PROFILE",
-        "ASSERT",
-    )
-
-
 class TestCompactionSession:
     """
     This class adds specific tests that aren't part of the parametrized test suite.
@@ -581,428 +540,3 @@ class TestCompactionSession:
                 }
             )
         )
-
-    def test_compact_partition_when_incremental_pk_hash_is_over_2gb(
-        self, s3_resource, local_deltacat_storage_kwargs, disable_sha1
-    ):
-        """
-        A test case which ensures the compaction succeeds even if the incremental
-        arrow table size is over 2GB. It is added to prevent ArrowCapacityError
-        when running is_in operation during merge.
-
-        Note that we set SHA1_HASHING_FOR_MEMORY_OPTIMIZATION_DISABLED to bypass sha1 hashing
-        which truncates the lengths of pk strings when deduping.
-        """
-        # setup
-        staged_source = stage_partition_from_file_paths(
-            self.NAMESPACE, ["source"], **local_deltacat_storage_kwargs
-        )
-        # we create chunked array to avoid ArrowCapacityError
-        chunked_pk_array = pa.chunked_array([["13bytesstring"], ["12bytestring"]])
-        table = pa.table([chunked_pk_array], names=["pk"])
-        source_delta = commit_delta_to_staged_partition(
-            staged_source, pa_table=table, **local_deltacat_storage_kwargs
-        )
-
-        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
-        )
-
-        # rebase first
-        rebase_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": source_delta.stream_position,
-                    "list_deltas_kwargs": {
-                        **local_deltacat_storage_kwargs,
-                        **{"equivalent_table_types": []},
-                    },
-                    "primary_keys": ["pk"],
-                    "rebase_source_partition_locator": source_delta.partition_locator,
-                    "rebase_source_partition_high_watermark": source_delta.stream_position,
-                    "records_per_compacted_file": 4000,
-                    "s3_client_kwargs": {},
-                    "source_partition_locator": source_delta.partition_locator,
-                    "resource_estimation_method": ResourceEstimationMethod.PREVIOUS_INFLATION,
-                }
-            )
-        )
-
-        rebased_rcf = get_rcf(s3_resource, rebase_url)
-
-        assert rebased_rcf.compacted_pyarrow_write_result.files == 1
-        assert rebased_rcf.compacted_pyarrow_write_result.records == 2
-
-        # Run incremental with a small delta on source
-        chunked_pk_array = pa.chunked_array(
-            [["13bytesstring" * 95_000_000], ["12bytestring" * 95_000_000]]
-        )  # 2.3GB
-        table = pa.table([chunked_pk_array], names=["pk"])
-
-        incremental_source_delta = commit_delta_to_partition(
-            source_delta.partition_locator,
-            pa_table=table,
-            **local_deltacat_storage_kwargs,
-        )
-        assert (
-            incremental_source_delta.partition_locator == source_delta.partition_locator
-        ), "source partition locator should not change"
-        dest_partition = ds.get_partition(
-            dest_partition.stream_locator,
-            dest_partition.partition_values,
-            **local_deltacat_storage_kwargs,
-        )
-
-        assert (
-            dest_partition.locator
-            == rebased_rcf.compacted_delta_locator.partition_locator
-        ), "The new destination partition should be same as compacted partition"
-
-        # Run incremental
-        incremental_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": incremental_source_delta.stream_position,
-                    "list_deltas_kwargs": {
-                        **local_deltacat_storage_kwargs,
-                        **{"equivalent_table_types": []},
-                    },
-                    "primary_keys": ["pk"],
-                    "records_per_compacted_file": 4000,
-                    "s3_client_kwargs": {},
-                    "source_partition_locator": incremental_source_delta.partition_locator,
-                    "resource_estimation_method": ResourceEstimationMethod.PREVIOUS_INFLATION,
-                }
-            )
-        )
-
-        incremental_rcf = get_rcf(s3_resource, incremental_url)
-
-        assert incremental_rcf.compacted_pyarrow_write_result.files == 1
-        assert (
-            incremental_rcf.compacted_pyarrow_write_result.pyarrow_bytes >= 2300000000
-        )
-        assert incremental_rcf.compacted_pyarrow_write_result.records == 4
-
-    def test_compact_partition_when_bucket_spec_validation_fails(
-        self,
-        s3_resource,
-        local_deltacat_storage_kwargs,
-        enable_bucketing_spec_validation,
-    ):
-        """
-        A test case which asserts the bucketing spec validation throws an assertion error
-        when the validation has failed.
-        """
-
-        # setup
-        staged_source = stage_partition_from_file_paths(
-            self.NAMESPACE, ["source"], **local_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"], **local_deltacat_storage_kwargs
-        )
-        dest_partition = ds.commit_partition(
-            staged_dest, **local_deltacat_storage_kwargs
-        )
-
-        # 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": 4,
-                    "last_stream_position_to_compact": source_delta.stream_position,
-                    "list_deltas_kwargs": {
-                        **local_deltacat_storage_kwargs,
-                        **{"equivalent_table_types": []},
-                    },
-                    "primary_keys": ["pk"],
-                    "rebase_source_partition_locator": source_delta.partition_locator,
-                    "rebase_source_partition_high_watermark": source_delta.stream_position,
-                    "records_per_compacted_file": 1,
-                    "s3_client_kwargs": {},
-                    "source_partition_locator": source_delta.partition_locator,
-                }
-            )
-        )
-
-        backfill_rcf = get_rcf(s3_resource, rcf_url)
-        bucket, backfill_key1, backfill_key2 = rcf_url.strip("s3://").split("/")
-        # Move the records to different hash buckets to simulate a validation failure.
-        backfill_rcf["hbIndexToEntryRange"] = {"1": [0, 3]}
-        s3_resource.Bucket(bucket).put_object(
-            Key=f"{backfill_key1}/{backfill_key2}", Body=json.dumps(backfill_rcf)
-        )
-
-        # 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_destination_partition = ds.get_partition(
-            dest_partition.stream_locator, [], **local_deltacat_storage_kwargs
-        )
-
-        with pytest.raises(ValidationError) as excinfo:
-            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": new_destination_partition.locator,
-                        "drop_duplicates": True,
-                        "hash_bucket_count": 4,
-                        "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": new_source_delta.partition_locator,
-                    }
-                )
-            )
-
-        assert (
-            "Hash bucket drift detected at index: 0. Expected hash bucket index to be 1 but found 0"
-            in str(excinfo.value)
-        )
-
-    def test_compact_partition_when_bucket_spec_validation_fails_but_env_variable_disabled(
-        self,
-        s3_resource,
-        local_deltacat_storage_kwargs,
-    ):
-        """
-        A test case which asserts even if bucketing spec validation fails, compaction doesn't
-        throw an error if the feature is not enabled.
-        """
-
-        # setup
-        staged_source = stage_partition_from_file_paths(
-            self.NAMESPACE, ["source"], **local_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"], **local_deltacat_storage_kwargs
-        )
-        dest_partition = ds.commit_partition(
-            staged_dest, **local_deltacat_storage_kwargs
-        )
-
-        # 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": 4,
-                    "last_stream_position_to_compact": source_delta.stream_position,
-                    "list_deltas_kwargs": {
-                        **local_deltacat_storage_kwargs,
-                        **{"equivalent_table_types": []},
-                    },
-                    "primary_keys": ["pk"],
-                    "rebase_source_partition_locator": source_delta.partition_locator,
-                    "rebase_source_partition_high_watermark": source_delta.stream_position,
-                    "records_per_compacted_file": 1,
-                    "s3_client_kwargs": {},
-                    "source_partition_locator": source_delta.partition_locator,
-                }
-            )
-        )
-
-        backfill_rcf = get_rcf(s3_resource, rcf_url)
-        bucket, backfill_key1, backfill_key2 = rcf_url.strip("s3://").split("/")
-        # Move the records to different hash buckets to simulate a validation failure.
-        backfill_rcf["hbIndexToEntryRange"] = {"1": [0, 3]}
-        s3_resource.Bucket(bucket).put_object(
-            Key=f"{backfill_key1}/{backfill_key2}", Body=json.dumps(backfill_rcf)
-        )
-
-        # 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_destination_partition = ds.get_partition(
-            dest_partition.stream_locator, [], **local_deltacat_storage_kwargs
-        )
-
-        new_rcf = 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": new_destination_partition.locator,
-                    "drop_duplicates": True,
-                    "hash_bucket_count": 4,
-                    "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": new_source_delta.partition_locator,
-                }
-            )
-        )
-
-        incremental_rcf = get_rcf(s3_resource, new_rcf)
-        assert incremental_rcf.hash_bucket_count == 4
-        assert len(incremental_rcf.hb_index_to_entry_range) == 2
-
-    def test_compact_partition_when_bucket_spec_validation_succeeds(
-        self,
-        s3_resource,
-        local_deltacat_storage_kwargs,
-        enable_bucketing_spec_validation,
-    ):
-        """
-        A test case which asserts the bucketing spec validation does not throw
-        and error when the validation succeeds.
-        """
-
-        # setup
-        staged_source = stage_partition_from_file_paths(
-            self.NAMESPACE, ["source"], **local_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"], **local_deltacat_storage_kwargs
-        )
-        dest_partition = ds.commit_partition(
-            staged_dest, **local_deltacat_storage_kwargs
-        )
-
-        # 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": 4,
-                    "last_stream_position_to_compact": source_delta.stream_position,
-                    "list_deltas_kwargs": {
-                        **local_deltacat_storage_kwargs,
-                        **{"equivalent_table_types": []},
-                    },
-                    "primary_keys": ["pk"],
-                    "rebase_source_partition_locator": source_delta.partition_locator,
-                    "rebase_source_partition_high_watermark": source_delta.stream_position,
-                    "records_per_compacted_file": 1,
-                    "s3_client_kwargs": {},
-                    "source_partition_locator": source_delta.partition_locator,
-                }
-            )
-        )
-
-        rcf = get_rcf(s3_resource, rcf_url)
-        assert rcf.hash_bucket_count == 4
-
-        # 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_destination_partition = ds.get_partition(
-            dest_partition.stream_locator, [], **local_deltacat_storage_kwargs
-        )
-
-        new_uri = 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": new_destination_partition.locator,
-                    "drop_duplicates": True,
-                    "hash_bucket_count": 4,
-                    "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": new_source_delta.partition_locator,
-                }
-            )
-        )
-
-        rcf = get_rcf(s3_resource, new_uri)
-        assert rcf.hash_bucket_count == 4