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

deltacat/tests/compute/test_compact_partition_rebase_then_incremental.py

@@ -1,45 +1,42 @@
-import ray
-import os
-from moto import mock_s3
+import tempfile
+from typing import Any, Dict, List, Optional, Set, Tuple, Callable
+import uuid
 import pytest
-import boto3
-from boto3.resources.base import ServiceResource
+
 import pyarrow as pa
-from deltacat.io.ray_plasma_object_store import RayPlasmaObjectStore
+import ray
+import pandas as pd
+
+from deltacat.io.file_object_store import FileObjectStore
 from pytest_benchmark.fixture import BenchmarkFixture
 
 from deltacat.tests.compute.test_util_constant import (
     BASE_TEST_SOURCE_NAMESPACE,
     BASE_TEST_SOURCE_TABLE_NAME,
     BASE_TEST_SOURCE_TABLE_VERSION,
-    TEST_S3_RCF_BUCKET_NAME,
     DEFAULT_NUM_WORKERS,
     DEFAULT_WORKER_INSTANCE_CPUS,
 )
 from deltacat.compute.compactor.model.compactor_version import CompactorVersion
 from deltacat.tests.compute.test_util_common import (
-    get_rcf,
-)
-from deltacat.tests.test_utils.utils import read_s3_contents
-from deltacat.tests.compute.test_util_common import (
-    get_compacted_delta_locator_from_rcf,
-)
-from deltacat.tests.compute.test_util_create_table_deltas_repo import (
-    create_incremental_deltas_on_source_table,
-)
-from deltacat.tests.compute.test_util_create_table_deltas_repo import (
-    create_src_w_deltas_destination_rebase_w_deltas_strategy,
+    create_src_w_deltas_destination_rebase_w_deltas_strategy_main,
+    create_incremental_deltas_on_source_table_main,
+    get_rci_from_partition,
+    read_audit_file,
+    PartitionKey,
+    get_compacted_delta_locator_from_partition,
 )
 from deltacat.tests.compute.compact_partition_rebase_then_incremental_test_cases import (
     REBASE_THEN_INCREMENTAL_TEST_CASES,
 )
-from typing import Any, Callable, Dict, List, Optional, Set, Tuple
+
 from deltacat.types.media import StorageType
 from deltacat.storage import (
     DeltaType,
     DeltaLocator,
     Partition,
     PartitionLocator,
+    metastore,
 )
 from deltacat.types.media import ContentType
 from deltacat.compute.compactor.model.compact_partition_params import (
@@ -64,29 +61,24 @@ def setup_ray_cluster():
     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
-
+"""
+FUNCTION scoped fixtures
+"""
 
-@pytest.fixture(scope="module")
-def s3_resource(mock_aws_credential):
-    with mock_s3():
-        yield boto3.resource("s3")
 
+@pytest.fixture(autouse=True, scope="function")
+def enable_bucketing_spec_validation(monkeypatch):
+    """
+    Enable the bucketing spec validation for all tests.
+    This will help catch hash bucket drift in testing.
+    """
+    import deltacat.compute.compactor_v2.steps.merge
 
-@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,
+    monkeypatch.setattr(
+        deltacat.compute.compactor_v2.steps.merge,
+        "BUCKETING_SPEC_COMPLIANCE_PROFILE",
+        "ASSERT",
     )
-    yield
 
 
 @pytest.mark.parametrize(
@@ -162,13 +154,12 @@ def setup_compaction_artifacts_s3_bucket(s3_resource: ServiceResource):
     ],
     ids=[test_name for test_name in REBASE_THEN_INCREMENTAL_TEST_CASES],
 )
-def test_compact_partition_rebase_then_incremental(
-    s3_resource: ServiceResource,
-    local_deltacat_storage_kwargs: Dict[str, Any],
+def test_compact_partition_rebase_then_incremental_main(
+    main_deltacat_storage_kwargs: Dict[str, Any],
     test_name: str,
     primary_keys: Set[str],
     sort_keys: List[Optional[Any]],
-    partition_keys_param: Optional[List[Any]],
+    partition_keys_param: Optional[List[PartitionKey]],
     partition_values_param: List[Optional[str]],
     input_deltas_param: List[pa.Array],
     input_deltas_delta_type: str,
@@ -188,9 +179,15 @@ def test_compact_partition_rebase_then_incremental(
     compact_partition_func: Callable,
     benchmark: BenchmarkFixture,
 ):
-    import deltacat.tests.local_deltacat_storage as ds
+    ds_mock_kwargs = main_deltacat_storage_kwargs
+    """
+    This test performs rebase compaction first, then incremental compaction on the same data.
+    This tests the scenario where we first do a rebase (with different source/destination partitions)
+    and then follow up with incremental compaction using the result of the rebase.
+
+    This version uses the main metastore implementation instead of local storage.
+    """
 
-    ds_mock_kwargs = local_deltacat_storage_kwargs
     """
     REBASE
     """
@@ -199,7 +196,7 @@ def test_compact_partition_rebase_then_incremental(
         source_table_stream,
         destination_table_stream,
         rebased_table_stream,
-    ) = create_src_w_deltas_destination_rebase_w_deltas_strategy(
+    ) = create_src_w_deltas_destination_rebase_w_deltas_strategy_main(
         sort_keys,
         partition_keys,
         input_deltas_param,
@@ -207,19 +204,48 @@ def test_compact_partition_rebase_then_incremental(
         partition_values_param,
         ds_mock_kwargs,
     )
-    source_partition: Partition = ds.get_partition(
+
+    # Convert partition values for partition lookup (same as in other helper functions)
+    converted_partition_values_for_lookup = partition_values_param
+    if partition_values_param and partition_keys:
+        converted_partition_values_for_lookup = []
+        for i, (value, pk) in enumerate(zip(partition_values_param, partition_keys)):
+            if pk.key_type.value == "int":  # Use .value to get string representation
+                converted_partition_values_for_lookup.append(int(value))
+            elif pk.key_type.value == "timestamp":
+                # Handle timestamp partition values
+                if isinstance(value, str) and "T" in value and value.endswith("Z"):
+                    ts = pd.to_datetime(value)
+                    # Convert to microseconds since epoch for PyArrow timestamp[us]
+                    converted_partition_values_for_lookup.append(
+                        int(ts.timestamp() * 1_000_000)
+                    )
+                else:
+                    converted_partition_values_for_lookup.append(value)
+            else:
+                converted_partition_values_for_lookup.append(value)
+
+    source_partition: Partition = metastore.get_partition(
         source_table_stream.locator,
-        partition_values_param,
+        converted_partition_values_for_lookup,
         **ds_mock_kwargs,
     )
+    # Generate a destination partition ID based on the source partition
+    destination_partition_id = str(uuid.uuid4())
     destination_partition_locator: PartitionLocator = PartitionLocator.of(
         destination_table_stream.locator,
-        partition_values_param,
-        None,
+        converted_partition_values_for_lookup,
+        destination_partition_id,
+    )
+    all_column_names = metastore.get_table_version_column_names(
+        destination_partition_locator.namespace,
+        destination_partition_locator.table_name,
+        destination_partition_locator.table_version,
+        **ds_mock_kwargs,
     )
-    rebased_partition: Partition = ds.get_partition(
+    rebased_partition: Partition = metastore.get_partition(
         rebased_table_stream.locator,
-        partition_values_param,
+        converted_partition_values_for_lookup,
         **ds_mock_kwargs,
     )
     num_workers, worker_instance_cpu = DEFAULT_NUM_WORKERS, DEFAULT_WORKER_INSTANCE_CPUS
@@ -230,151 +256,195 @@ def test_compact_partition_rebase_then_incremental(
         pgm = PlacementGroupManager(
             1, total_cpus, worker_instance_cpu, memory_per_bundle=4000000
         ).pgs[0]
-    compact_partition_params = 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": ds_mock_kwargs,
-            "destination_partition_locator": destination_partition_locator,
-            "hash_bucket_count": hash_bucket_count_param,
-            "last_stream_position_to_compact": source_partition.stream_position,
-            "list_deltas_kwargs": {**ds_mock_kwargs, **{"equivalent_table_types": []}},
-            "object_store": RayPlasmaObjectStore(),
-            "pg_config": pgm,
-            "primary_keys": primary_keys,
-            "read_kwargs_provider": read_kwargs_provider_param,
-            "rebase_source_partition_locator": source_partition.locator,
-            "records_per_compacted_file": records_per_compacted_file_param,
-            "s3_client_kwargs": {},
-            "source_partition_locator": rebased_partition.locator,
-            "sort_keys": sort_keys if sort_keys else None,
-        }
-    )
-    # execute
-    rcf_file_s3_uri = benchmark(compact_partition_func, compact_partition_params)
-    compacted_delta_locator: DeltaLocator = get_compacted_delta_locator_from_rcf(
-        s3_resource, rcf_file_s3_uri
-    )
-    tables = ds.download_delta(
-        compacted_delta_locator, storage_type=StorageType.LOCAL, **ds_mock_kwargs
-    )
-    actual_rebase_compacted_table = pa.concat_tables(tables)
-    # if no primary key is specified then sort by sort_key for consistent assertion
-    sorting_cols: List[Any] = (
-        [(val, "ascending") for val in primary_keys]
-        if primary_keys
-        else [pa_key for key in sort_keys for pa_key in key.arrow]
-        if sort_keys
-        else []
-    )
-    rebase_expected_compact_partition_result = (
-        rebase_expected_compact_partition_result.combine_chunks().sort_by(sorting_cols)
-    )
-    actual_rebase_compacted_table = (
-        actual_rebase_compacted_table.combine_chunks().sort_by(sorting_cols)
-    )
-    assert actual_rebase_compacted_table.equals(
-        rebase_expected_compact_partition_result
-    ), f"{actual_rebase_compacted_table} does not match {rebase_expected_compact_partition_result}"
-    """
-    INCREMENTAL
-    """
-    (
-        source_partition_locator_w_deltas,
-        new_delta,
-        incremental_delta_length,
-        has_delete_deltas,
-    ) = create_incremental_deltas_on_source_table(
-        BASE_TEST_SOURCE_NAMESPACE,
-        BASE_TEST_SOURCE_TABLE_NAME,
-        BASE_TEST_SOURCE_TABLE_VERSION,
-        source_table_stream,
-        partition_values_param,
-        incremental_deltas,
-        ds_mock_kwargs,
-    )
-    compact_partition_params = 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": ds_mock_kwargs,
-            "destination_partition_locator": compacted_delta_locator.partition_locator,
-            "drop_duplicates": drop_duplicates_param,
-            "hash_bucket_count": hash_bucket_count_param,
-            "last_stream_position_to_compact": new_delta.stream_position,
-            "list_deltas_kwargs": {**ds_mock_kwargs, **{"equivalent_table_types": []}},
-            "object_store": RayPlasmaObjectStore(),
-            "pg_config": pgm,
-            "primary_keys": primary_keys,
-            "read_kwargs_provider": read_kwargs_provider_param,
-            "rebase_source_partition_locator": None,
-            "rebase_source_partition_high_watermark": None,
-            "records_per_compacted_file": records_per_compacted_file_param,
-            "s3_client_kwargs": {},
-            "source_partition_locator": source_partition_locator_w_deltas,
-            "sort_keys": sort_keys if sort_keys else None,
-        }
-    )
-    if expected_terminal_exception:
-        with pytest.raises(expected_terminal_exception) as exc_info:
-            compact_partition_func(compact_partition_params)
-        assert expected_terminal_exception_message in str(exc_info.value)
-        return
-    rcf_file_s3_uri = compact_partition_func(compact_partition_params)
-    round_completion_info = get_rcf(s3_resource, rcf_file_s3_uri)
-    compacted_delta_locator_incremental: DeltaLocator = (
-        round_completion_info.compacted_delta_locator
-    )
-    # assert if RCF covers all files
-    if compactor_version != CompactorVersion.V1.value:
-        previous_end = None
-        for start, end in round_completion_info.hb_index_to_entry_range.values():
-            assert (previous_end is None and start == 0) or start == previous_end
-            previous_end = end
-        assert (
-            previous_end == round_completion_info.compacted_pyarrow_write_result.files
-        )
 
-    audit_bucket, audit_key = round_completion_info.compaction_audit_url.replace(
-        "s3://", ""
-    ).split("/", 1)
-    compaction_audit_obj: dict = read_s3_contents(s3_resource, audit_bucket, audit_key)
-    compaction_audit: CompactionSessionAuditInfo = CompactionSessionAuditInfo(
-        **compaction_audit_obj
-    )
+    with tempfile.TemporaryDirectory() as test_dir:
+        # Extract catalog from storage kwargs
+        catalog = ds_mock_kwargs.get("inner")
 
-    tables = ds.download_delta(
-        compacted_delta_locator_incremental,
-        storage_type=StorageType.LOCAL,
-        **ds_mock_kwargs,
-    )
-    actual_compacted_table = pa.concat_tables(tables)
-    expected_terminal_compact_partition_result = (
-        expected_terminal_compact_partition_result.combine_chunks().sort_by(
-            sorting_cols
+        compact_partition_params = CompactPartitionParams.of(
+            {
+                "catalog": catalog,
+                "compacted_file_content_type": ContentType.PARQUET,
+                "dd_max_parallelism_ratio": 1.0,
+                "deltacat_storage": metastore,
+                "deltacat_storage_kwargs": ds_mock_kwargs,
+                "destination_partition_locator": destination_partition_locator,
+                "hash_bucket_count": hash_bucket_count_param,
+                "last_stream_position_to_compact": source_partition.stream_position,
+                "list_deltas_kwargs": {
+                    **ds_mock_kwargs,
+                    **{"equivalent_table_types": []},
+                },
+                "object_store": FileObjectStore(test_dir),
+                "original_fields": {
+                    "pk_col_1",
+                    "pk_col_2",
+                    "sk_col_1",
+                    "sk_col_2",
+                    "col_1",
+                    "col_2",
+                    "region_id",
+                },
+                "pg_config": pgm,
+                "primary_keys": primary_keys,
+                "all_column_names": all_column_names,
+                "read_kwargs_provider": read_kwargs_provider_param,
+                "rebase_source_partition_locator": source_partition.locator,
+                "records_per_compacted_file": records_per_compacted_file_param,
+                "source_partition_locator": rebased_partition.locator,
+                "sort_keys": sort_keys if sort_keys else None,
+            }
         )
-    )
-    actual_compacted_table = actual_compacted_table.combine_chunks().sort_by(
-        sorting_cols
-    )
-    # NOTE: if delete type-deltas are present this relationship no longer holds true
-    if not has_delete_deltas:
-        assert compaction_audit.input_records == (
-            incremental_delta_length if incremental_deltas else 0
-        ) + len(actual_rebase_compacted_table), (
-            " Total input records must be equal to incremental deltas"
-            " + previous compacted table size"
+        # execute
+        benchmark(compact_partition_func, compact_partition_params)
+        compacted_delta_locator: DeltaLocator = (
+            get_compacted_delta_locator_from_partition(
+                destination_partition_locator,
+                metastore,
+                catalog=catalog,
+            )
+        )
+        tables = metastore.download_delta(
+            compacted_delta_locator,
+            storage_type=StorageType.LOCAL,
+            **ds_mock_kwargs,
+        )
+        actual_rebase_compacted_table = pa.concat_tables(tables)
+        all_column_names = metastore.get_table_version_column_names(
+            destination_partition_locator.namespace,
+            destination_partition_locator.table_name,
+            destination_partition_locator.table_version,
+            **ds_mock_kwargs,
+        )
+        # if no primary key is specified then sort by sort_key for consistent assertion
+        sorting_cols: List[Any] = (
+            [(val, "ascending") for val in primary_keys]
+            if primary_keys
+            else [pa_key for key in sort_keys for pa_key in key.arrow]
+            if sort_keys
+            else []
+        )
+        rebase_expected_compact_partition_result = (
+            rebase_expected_compact_partition_result.combine_chunks().sort_by(
+                sorting_cols
+            )
+        )
+        actual_rebase_compacted_table = (
+            actual_rebase_compacted_table.combine_chunks().sort_by(sorting_cols)
+        )
+        assert actual_rebase_compacted_table.equals(
+            rebase_expected_compact_partition_result
+        ), f"{actual_rebase_compacted_table} does not match {rebase_expected_compact_partition_result}"
+
+        """
+        INCREMENTAL
+        """
+        (
+            source_partition_locator_w_deltas,
+            new_delta,
+            incremental_delta_length,
+            has_delete_deltas,
+        ) = create_incremental_deltas_on_source_table_main(
+            BASE_TEST_SOURCE_NAMESPACE,
+            BASE_TEST_SOURCE_TABLE_NAME,
+            BASE_TEST_SOURCE_TABLE_VERSION,
+            source_table_stream,
+            partition_values_param,
+            incremental_deltas,
+            ds_mock_kwargs,
         )
 
-    assert actual_compacted_table.equals(
-        expected_terminal_compact_partition_result
-    ), f"{actual_compacted_table} does not match {expected_terminal_compact_partition_result}"
+        # Handle empty incremental deltas case
+        if new_delta is None:
+            # For empty incremental deltas, the expected result should be the same as rebase result
+            # Skip incremental compaction and just verify the rebase result
+            actual_compact_partition_result = actual_rebase_compacted_table
+            compaction_audit = None
+        else:
+            # Perform incremental compaction when there are actual deltas
+            last_stream_position = new_delta.stream_position
 
-    if assert_compaction_audit is not None:
-        if not assert_compaction_audit(compactor_version, compaction_audit):
-            assert False, "Compaction audit assertion failed"
-    return
+            compact_partition_params = CompactPartitionParams.of(
+                {
+                    "catalog": catalog,
+                    "compacted_file_content_type": ContentType.PARQUET,
+                    "dd_max_parallelism_ratio": 1.0,
+                    "deltacat_storage": metastore,
+                    "deltacat_storage_kwargs": ds_mock_kwargs,
+                    "destination_partition_locator": compacted_delta_locator.partition_locator,
+                    "drop_duplicates": drop_duplicates_param,
+                    "hash_bucket_count": hash_bucket_count_param,
+                    "last_stream_position_to_compact": last_stream_position,
+                    "list_deltas_kwargs": {
+                        **ds_mock_kwargs,
+                        **{"equivalent_table_types": []},
+                    },
+                    "object_store": FileObjectStore(test_dir),
+                    "original_fields": {
+                        "pk_col_1",
+                        "pk_col_2",
+                        "sk_col_1",
+                        "sk_col_2",
+                        "col_1",
+                        "col_2",
+                        "region_id",
+                    },
+                    "pg_config": pgm,
+                    "primary_keys": primary_keys,
+                    "all_column_names": all_column_names,
+                    "read_kwargs_provider": read_kwargs_provider_param,
+                    "rebase_source_partition_locator": None,
+                    "rebase_source_partition_high_watermark": None,
+                    "records_per_compacted_file": records_per_compacted_file_param,
+                    "source_partition_locator": source_partition_locator_w_deltas,
+                    "sort_keys": sort_keys if sort_keys else None,
+                }
+            )
+            if expected_terminal_exception:
+                with pytest.raises(expected_terminal_exception) as exc_info:
+                    compact_partition_func(compact_partition_params)
+                assert expected_terminal_exception_message in str(exc_info.value)
+                return
+            compact_partition_func(compact_partition_params)
+            # assert
+            compacted_delta_locator: DeltaLocator = (
+                get_compacted_delta_locator_from_partition(
+                    destination_partition_locator, metastore, catalog=catalog
+                )
+            )
+            tables = metastore.download_delta(
+                compacted_delta_locator,
+                storage_type=StorageType.LOCAL,
+                **ds_mock_kwargs,
+            )
+            actual_compact_partition_result = pa.concat_tables(tables)
+
+            # Get compaction audit for verification if needed
+            round_completion_info = get_rci_from_partition(
+                destination_partition_locator, metastore, catalog=catalog
+            )
+            # Get catalog root for audit file resolution
+            catalog_root = catalog.root
+
+            compaction_audit_obj: dict = read_audit_file(
+                round_completion_info.compaction_audit_url, catalog_root
+            )
+            compaction_audit = CompactionSessionAuditInfo(**compaction_audit_obj)
+
+        # Verify the final result
+        actual_compact_partition_result = (
+            actual_compact_partition_result.combine_chunks().sort_by(sorting_cols)
+        )
+        expected_terminal_compact_partition_result = (
+            expected_terminal_compact_partition_result.combine_chunks().sort_by(
+                sorting_cols
+            )
+        )
+        assert actual_compact_partition_result.equals(
+            expected_terminal_compact_partition_result
+        ), f"{actual_compact_partition_result} does not match {expected_terminal_compact_partition_result}"
+
+        if assert_compaction_audit is not None and compaction_audit is not None:
+            if not assert_compaction_audit(compactor_version, compaction_audit):
+                pytest.fail("Compaction audit assertion failed")
+        return