deltacat 1.1.38__py3-none-any.whl → 2.0.0__py3-none-any.whl

deltacat/tests/compute/test_compact_partition_incremental.py

@@ -1,29 +1,29 @@
-import ray
-from moto import mock_s3
-import pytest
-import os
 import logging
-import boto3
-from typing import Any, Callable, Dict, List, Optional, Set, Tuple
-from boto3.resources.base import ServiceResource
+from typing import Any, Dict, List, Optional, Set, Tuple, Callable
+import uuid
+import pytest
+
 import pyarrow as pa
+import ray
+
 from pytest_benchmark.fixture import BenchmarkFixture
 from deltacat.types.media import StorageType
 
 from deltacat.tests.compute.test_util_common import (
-    get_rcf,
+    get_rci_from_partition,
+    read_audit_file,
+    PartitionKeyType,
 )
-from deltacat.compute.compactor.model.compactor_version import CompactorVersion
-from deltacat.tests.test_utils.utils import read_s3_contents
-from deltacat.tests.compute.test_util_create_table_deltas_repo import (
-    create_src_w_deltas_destination_plus_destination,
-    add_late_deltas_to_partition,
+from deltacat.tests.compute.test_util_common import (
+    add_late_deltas_to_partition_main,
+    create_src_w_deltas_destination_plus_destination_main,
 )
+from deltacat.compute.compactor.model.compactor_version import CompactorVersion
+
 from deltacat.tests.compute.compact_partition_test_cases import (
     INCREMENTAL_TEST_CASES,
 )
 from deltacat.tests.compute.test_util_constant import (
-    TEST_S3_RCF_BUCKET_NAME,
     DEFAULT_NUM_WORKERS,
     DEFAULT_WORKER_INSTANCE_CPUS,
 )
@@ -37,6 +37,7 @@ from deltacat.storage import (
     DeltaLocator,
     Partition,
     PartitionLocator,
+    metastore,
 )
 from deltacat.types.media import ContentType
 from deltacat.compute.compactor.model.compaction_session_audit_info import (
@@ -50,11 +51,6 @@ from deltacat.utils.placement import (
 )
 from deltacat import logs
 
-DATABASE_FILE_PATH_KEY, DATABASE_FILE_PATH_VALUE = (
-    "db_file_path",
-    "deltacat/tests/local_deltacat_storage/db_test.sqlite",
-)
-
 logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
 
 
@@ -70,55 +66,11 @@ 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
-
-
-@pytest.fixture(autouse=True, scope="module")
-def cleanup_the_database_file_after_all_compaction_session_package_tests_complete():
-    # make sure the database file is deleted after all the compactor package tests are completed
-    yield
-    if os.path.exists(DATABASE_FILE_PATH_VALUE):
-        os.remove(DATABASE_FILE_PATH_VALUE)
-
-
-@pytest.fixture(scope="module")
-def s3_resource():
-    with mock_s3():
-        yield boto3.resource("s3")
-
-
-@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
-
-
 """
 FUNCTION scoped fixtures
 """
 
 
-@pytest.fixture(scope="function")
-def offer_local_deltacat_storage_kwargs(request: pytest.FixtureRequest):
-    # see deltacat/tests/local_deltacat_storage/README.md for documentation
-    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(autouse=True, scope="function")
 def enable_bucketing_spec_validation(monkeypatch):
     """
@@ -134,6 +86,11 @@ def enable_bucketing_spec_validation(monkeypatch):
     )
 
 
+@pytest.fixture(scope="function")
+def temp_dir(tmp_path):
+    return str(tmp_path)
+
+
 @pytest.mark.parametrize(
     [
         "test_name",
@@ -207,9 +164,8 @@ def enable_bucketing_spec_validation(monkeypatch):
     ],
     ids=[test_name for test_name in INCREMENTAL_TEST_CASES],
 )
-def test_compact_partition_incremental(
-    s3_resource: ServiceResource,
-    offer_local_deltacat_storage_kwargs: Dict[str, Any],
+def test_compact_partition_incremental_main(
+    main_deltacat_storage_kwargs: Dict[str, Any],
     test_name: str,
     primary_keys: Set[str],
     sort_keys: Dict[str, str],
@@ -233,9 +189,16 @@ def test_compact_partition_incremental(
     compact_partition_func: Callable,
     benchmark: BenchmarkFixture,
 ):
-    import deltacat.tests.local_deltacat_storage as ds
+    # Skip in-place compaction tests for main storage as it's not yet implemented
+    if is_inplace:
+        pytest.skip(
+            "In-place compaction not yet implemented in main storage (delta prepending limitation)"
+        )
 
-    ds_mock_kwargs: Dict[str, Any] = offer_local_deltacat_storage_kwargs
+    ds_mock_kwargs: Dict[str, Any] = main_deltacat_storage_kwargs
+
+    # Extract catalog from storage kwargs
+    catalog = ds_mock_kwargs.get("inner")
 
     # setup
     partition_keys = partition_keys_param
@@ -246,8 +209,7 @@ def test_compact_partition_incremental(
         source_table_namespace,
         source_table_name,
         source_table_version,
-    ) = create_src_w_deltas_destination_plus_destination(
-        primary_keys,
+    ) = create_src_w_deltas_destination_plus_destination_main(
         sort_keys,
         partition_keys,
         input_deltas,
@@ -256,15 +218,38 @@ def test_compact_partition_incremental(
         ds_mock_kwargs,
         is_inplace,
     )
-    source_partition: Partition = ds.get_partition(
+
+    # Convert partition values to correct types for get_partition call
+    converted_partition_values = []
+    if partition_values_param and partition_keys:
+        # partition_values_param is a single string, but we need to handle it as a list
+        partition_values_list = (
+            [partition_values_param]
+            if isinstance(partition_values_param, str)
+            else partition_values_param
+        )
+        for i, (value, pk) in enumerate(zip(partition_values_list, partition_keys)):
+            if pk.key_type == PartitionKeyType.INT:
+                converted_partition_values.append(int(value))
+            else:
+                converted_partition_values.append(value)
+    else:
+        converted_partition_values = (
+            [partition_values_param] if partition_values_param else []
+        )
+
+    source_partition: Partition = metastore.get_partition(
         source_table_stream.locator,
-        partition_values_param,
+        converted_partition_values,
+        partition_scheme_id="default_partition_scheme" if partition_keys else None,
         **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,
+        destination_partition_id,
     )
     num_workers, worker_instance_cpu = DEFAULT_NUM_WORKERS, DEFAULT_WORKER_INSTANCE_CPUS
     total_cpus: int = num_workers * worker_instance_cpu
@@ -275,12 +260,18 @@ def test_compact_partition_incremental(
         if create_placement_group_param
         else None
     )
+    all_column_names = metastore.get_table_version_column_names(
+        destination_table_stream.locator.table_locator.namespace,
+        destination_table_stream.locator.table_locator.table_name,
+        destination_table_stream.locator.table_version_locator.table_version,
+        catalog=catalog,
+    )
     compact_partition_params = CompactPartitionParams.of(
         {
-            "compaction_artifact_s3_bucket": TEST_S3_RCF_BUCKET_NAME,
+            "catalog": catalog,
             "compacted_file_content_type": ContentType.PARQUET,
             "dd_max_parallelism_ratio": 1.0,
-            "deltacat_storage": ds,
+            "deltacat_storage": metastore,
             "deltacat_storage_kwargs": ds_mock_kwargs,
             "destination_partition_locator": destination_partition_locator,
             "drop_duplicates": drop_duplicates_param,
@@ -289,11 +280,11 @@ def test_compact_partition_incremental(
             "list_deltas_kwargs": {**ds_mock_kwargs, **{"equivalent_table_types": []}},
             "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,
-            "s3_client_kwargs": {},
             "source_partition_locator": source_partition.locator,
             "sort_keys": sort_keys if sort_keys else None,
         }
@@ -304,18 +295,17 @@ def test_compact_partition_incremental(
         """
         This callable runs right before invoking the benchmark target function (compaction).
         This is needed as the benchmark module will invoke the target function multiple times
-        in a single test run, which can lead to non-idempotent behavior if RCFs are generated.
+        in a single test run, which can lead to non-idempotent behavior if RCIs are generated.
 
         Returns: args, kwargs
         """
-        s3_resource.Bucket(TEST_S3_RCF_BUCKET_NAME).objects.all().delete()
         return (compact_partition_params,), {}
 
     if add_late_deltas:
         # NOTE: In the case of in-place compaction it is plausible that new deltas may be added to the source partition during compaction
         # (so that the source_partitition.stream_position > last_stream_position_to_compact).
         # This parameter helps simulate the case to check that no late deltas are dropped even when the compacted partition is created.
-        latest_delta, _ = add_late_deltas_to_partition(
+        latest_delta, _ = add_late_deltas_to_partition_main(
             add_late_deltas, source_partition, ds_mock_kwargs
         )
     if expected_terminal_exception:
@@ -323,27 +313,28 @@ def test_compact_partition_incremental(
             compact_partition_func(compact_partition_params)
         assert expected_terminal_exception_message in str(exc_info.value)
         return
-    rcf_file_s3_uri = benchmark.pedantic(
-        compact_partition_func, setup=_incremental_compaction_setup
-    )
+    benchmark.pedantic(compact_partition_func, setup=_incremental_compaction_setup)
 
-    # validate
-    round_completion_info: RoundCompletionInfo = get_rcf(s3_resource, rcf_file_s3_uri)
+    # validate - get RoundCompletionInfo from the compacted partition
+    round_completion_info: RoundCompletionInfo = get_rci_from_partition(
+        destination_partition_locator, metastore, catalog=catalog
+    )
     compacted_delta_locator: DeltaLocator = (
         round_completion_info.compacted_delta_locator
     )
-    audit_bucket, audit_key = RoundCompletionInfo.get_audit_bucket_name_and_key(
-        round_completion_info.compaction_audit_url
-    )
 
-    compaction_audit_obj: Dict[str, Any] = read_s3_contents(
-        s3_resource, audit_bucket, audit_key
+    # Get catalog root for audit file resolution
+    catalog_root = catalog.root
+
+    compaction_audit_obj: Dict[str, Any] = read_audit_file(
+        round_completion_info.compaction_audit_url, catalog_root
     )
+
     compaction_audit: CompactionSessionAuditInfo = CompactionSessionAuditInfo(
         **compaction_audit_obj
     )
 
-    # assert if RCF covers all files
+    # assert if RCI 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():
@@ -353,7 +344,7 @@ def test_compact_partition_incremental(
             previous_end == round_completion_info.compacted_pyarrow_write_result.files
         )
 
-    tables = ds.download_delta(
+    tables = metastore.download_delta(
         compacted_delta_locator, storage_type=StorageType.LOCAL, **ds_mock_kwargs
     )
     actual_compacted_table = pa.concat_tables(tables)
@@ -387,25 +378,27 @@ def test_compact_partition_incremental(
             == destination_partition_locator.partition_values
             and source_partition.locator.stream_id
             == destination_partition_locator.stream_id
-        ), f"The source partition: {source_partition.locator.canonical_string} should match the destination partition: {destination_partition_locator.canonical_string}"
+        ), f"The source partition: {source_partition.locator} should match the destination partition: {destination_partition_locator}"
         assert (
             compacted_delta_locator.stream_id == source_partition.locator.stream_id
         ), "The compacted delta should be in the same stream as the source"
-        source_partition: Partition = ds.get_partition(
+        source_partition: Partition = metastore.get_partition(
             source_table_stream.locator,
-            partition_values_param,
+            converted_partition_values,
+            partition_scheme_id="default_partition_scheme" if partition_keys else None,
             **ds_mock_kwargs,
         )
-        compacted_partition: Optional[Partition] = ds.get_partition(
+        compacted_partition: Optional[Partition] = metastore.get_partition(
             compacted_delta_locator.stream_locator,
-            partition_values_param,
+            converted_partition_values,
+            partition_scheme_id="default_partition_scheme" if partition_keys else None,
             **ds_mock_kwargs,
         )
         assert (
             compacted_partition.state == source_partition.state == CommitState.COMMITTED
         ), f"The compacted/source table partition should be in {CommitState.COMMITTED} state and not {CommitState.DEPRECATED}"
         if add_late_deltas:
-            compacted_partition_deltas: List[Delta] = ds.list_partition_deltas(
+            compacted_partition_deltas: List[Delta] = metastore.list_partition_deltas(
                 partition_like=compacted_partition,
                 ascending_order=False,
                 **ds_mock_kwargs,

deltacat/tests/compute/test_compact_partition_multiple_rounds.py

@@ -1,43 +1,38 @@
-import ray
-import os
-from moto import mock_s3
+import tempfile
+from typing import Any, Dict, List, Optional, Set, Callable
 import pytest
-import boto3
-from boto3.resources.base import ServiceResource
 import pyarrow as pa
+import ray
+
 from deltacat.io.file_object_store import FileObjectStore
 from pytest_benchmark.fixture import BenchmarkFixture
-import tempfile
 
 from deltacat.tests.compute.test_util_constant import (
-    TEST_S3_RCF_BUCKET_NAME,
     DEFAULT_NUM_WORKERS,
     DEFAULT_WORKER_INSTANCE_CPUS,
 )
 from deltacat.tests.compute.test_util_common import (
-    get_rcf,
+    get_rci_from_partition,
+    read_audit_file,
+    PartitionKey,
+    get_compacted_delta_locator_from_partition,
 )
-from deltacat.tests.test_utils.utils import read_s3_contents
-from deltacat.compute.compactor.model.compactor_version import CompactorVersion
 from deltacat.tests.compute.test_util_common import (
-    get_compacted_delta_locator_from_rcf,
+    multiple_rounds_create_src_w_deltas_destination_rebase_w_deltas_strategy_main,
 )
+
+from deltacat.compute.compactor.model.compactor_version import CompactorVersion
 from deltacat.compute.compactor.model.compaction_session_audit_info import (
     CompactionSessionAuditInfo,
 )
-from deltacat.tests.compute.test_util_create_table_deltas_repo import (
-    multiple_rounds_create_src_w_deltas_destination_rebase_w_deltas_strategy,
-)
 from deltacat.tests.compute.compact_partition_multiple_rounds_test_cases import (
     MULTIPLE_ROUNDS_TEST_CASES,
 )
-from typing import Any, Callable, Dict, List, Optional, Set
-from deltacat.types.media import StorageType
+from deltacat.types.media import StorageType, ContentType
 from deltacat.storage import (
     DeltaLocator,
     Partition,
 )
-from deltacat.types.media import ContentType
 from deltacat.compute.compactor.model.compact_partition_params import (
     CompactPartitionParams,
 )
@@ -47,11 +42,7 @@ from deltacat.compute.compactor import (
 from deltacat.utils.placement import (
     PlacementGroupManager,
 )
-
-DATABASE_FILE_PATH_KEY, DATABASE_FILE_PATH_VALUE = (
-    "db_file_path",
-    "deltacat/tests/local_deltacat_storage/db_test.sqlite",
-)
+from deltacat.storage import metastore
 
 
 """
@@ -66,54 +57,11 @@ 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
-
-
-@pytest.fixture(autouse=True, scope="module")
-def cleanup_the_database_file_after_all_compaction_session_package_tests_complete():
-    # make sure the database file is deleted after all the compactor package tests are completed
-    if os.path.exists(DATABASE_FILE_PATH_VALUE):
-        os.remove(DATABASE_FILE_PATH_VALUE)
-
-
-@pytest.fixture(scope="module")
-def s3_resource(mock_aws_credential):
-    with mock_s3():
-        yield boto3.resource("s3")
-
-
-@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
-
-
 """
 FUNCTION scoped fixtures
 """
 
 
-@pytest.fixture(scope="function")
-def local_deltacat_storage_kwargs(request: pytest.FixtureRequest):
-    # see deltacat/tests/local_deltacat_storage/README.md for documentation
-    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(autouse=True, scope="function")
 def enable_bucketing_spec_validation(monkeypatch):
     """
@@ -199,14 +147,13 @@ def enable_bucketing_spec_validation(monkeypatch):
     ],
     ids=[test_name for test_name in MULTIPLE_ROUNDS_TEST_CASES],
 )
-def test_compact_partition_rebase_multiple_rounds_same_source_and_destination(
+def test_compact_partition_rebase_multiple_rounds_same_source_and_destination_main(
     mocker,
-    s3_resource: ServiceResource,
-    local_deltacat_storage_kwargs: Dict[str, Any],
+    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],
     expected_terminal_compact_partition_result: pa.Table,
@@ -225,38 +172,63 @@ def test_compact_partition_rebase_multiple_rounds_same_source_and_destination(
     num_rounds_param: int,
     benchmark: BenchmarkFixture,
 ):
-    import deltacat.tests.local_deltacat_storage as ds
-
-    ds_mock_kwargs = local_deltacat_storage_kwargs
+    ds_mock_kwargs = main_deltacat_storage_kwargs
     """
     This test tests different multi-round compaction rebase configurations,
-    as specified in compact_partition_multiple_rounds_test_cases.py
+    as specified in compact_partition_multiple_rounds_test_cases.py.
     These tests do not test multi-round compaction backfill, which is
     currently unsupported.
+
+    This version uses the main metastore implementation instead of local storage.
     """
     (
         source_table_stream,
         _,
         rebased_table_stream,
         _,
-    ) = multiple_rounds_create_src_w_deltas_destination_rebase_w_deltas_strategy(
-        primary_keys,
+    ) = multiple_rounds_create_src_w_deltas_destination_rebase_w_deltas_strategy_main(
         sort_keys,
         partition_keys_param,
         input_deltas_param,
         partition_values_param,
         ds_mock_kwargs,
     )
-    source_partition: Partition = ds.get_partition(
-        source_table_stream.locator,
-        partition_values_param,
+    # Convert partition values for partition lookup (same as in the helper function)
+    converted_partition_values_for_lookup = partition_values_param
+    if partition_values_param and partition_keys_param:
+        converted_partition_values_for_lookup = []
+        for i, (value, key) in enumerate(
+            zip(partition_values_param, partition_keys_param)
+        ):
+            if key.key_type == "int":
+                converted_partition_values_for_lookup.append(int(value))
+            elif key.key_type == "string":
+                converted_partition_values_for_lookup.append(str(value))
+            elif key.key_type == "timestamp":
+                converted_partition_values_for_lookup.append(
+                    value
+                )  # Keep as is for now
+            else:
+                converted_partition_values_for_lookup.append(value)
+
+    source_partition: Partition = metastore.get_partition(
+        stream_locator=source_table_stream.locator,
+        partition_values=converted_partition_values_for_lookup,
+        partition_scheme_id=source_table_stream.partition_scheme.id,
         **ds_mock_kwargs,
     )
-    rebased_partition: Partition = ds.get_partition(
-        rebased_table_stream.locator,
-        partition_values_param,
+    rebased_partition: Partition = metastore.get_partition(
+        stream_locator=rebased_table_stream.locator,
+        partition_values=converted_partition_values_for_lookup,
+        partition_scheme_id=rebased_table_stream.partition_scheme.id,
         **ds_mock_kwargs,
     )
+    all_column_names = metastore.get_table_version_column_names(
+        rebased_table_stream.locator.table_locator.namespace,
+        rebased_table_stream.locator.table_locator.table_name,
+        rebased_table_stream.locator.table_version_locator.table_version,
+        catalog=ds_mock_kwargs.get("inner"),
+    )
     total_cpus = DEFAULT_NUM_WORKERS * DEFAULT_WORKER_INSTANCE_CPUS
     pgm = None
     if create_placement_group_param:
@@ -266,10 +238,10 @@ def test_compact_partition_rebase_multiple_rounds_same_source_and_destination(
     with tempfile.TemporaryDirectory() as test_dir:
         compact_partition_params = CompactPartitionParams.of(
             {
-                "compaction_artifact_s3_bucket": TEST_S3_RCF_BUCKET_NAME,
+                "catalog": ds_mock_kwargs.get("inner"),
                 "compacted_file_content_type": ContentType.PARQUET,
                 "dd_max_parallelism_ratio": 1.0,
-                "deltacat_storage": ds,
+                "deltacat_storage": metastore,
                 "deltacat_storage_kwargs": ds_mock_kwargs,
                 "destination_partition_locator": rebased_partition.locator,
                 "hash_bucket_count": hash_bucket_count_param,
@@ -281,11 +253,11 @@ def test_compact_partition_rebase_multiple_rounds_same_source_and_destination(
                 "object_store": FileObjectStore(test_dir),
                 "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,
                 "rebase_source_partition_high_watermark": rebased_partition.stream_position,
                 "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,
                 "num_rounds": num_rounds_param,
@@ -308,23 +280,25 @@ def test_compact_partition_rebase_multiple_rounds_same_source_and_destination(
         object_store_clear_spy = mocker.spy(FileObjectStore, "clear")
 
         # execute
-        rcf_file_s3_uri = benchmark(compact_partition_func, compact_partition_params)
+        benchmark(compact_partition_func, compact_partition_params)
 
-        round_completion_info: RoundCompletionInfo = get_rcf(
-            s3_resource, rcf_file_s3_uri
-        )
-        audit_bucket, audit_key = RoundCompletionInfo.get_audit_bucket_name_and_key(
-            round_completion_info.compaction_audit_url
+        # Get RoundCompletionInfo from the compacted partition
+        round_completion_info: RoundCompletionInfo = get_rci_from_partition(
+            rebased_partition.locator, metastore, catalog=ds_mock_kwargs.get("inner")
         )
 
-        compaction_audit_obj: Dict[str, Any] = read_s3_contents(
-            s3_resource, audit_bucket, audit_key
+        # Get catalog root for audit file resolution
+        catalog = ds_mock_kwargs.get("inner")
+        catalog_root = catalog.root
+
+        compaction_audit_obj: Dict[str, Any] = read_audit_file(
+            round_completion_info.compaction_audit_url, catalog_root
         )
         compaction_audit: CompactionSessionAuditInfo = CompactionSessionAuditInfo(
             **compaction_audit_obj
         )
 
-        # assert if RCF covers all files
+        # assert if RCI covers all files
         # multiple rounds feature is only supported in V2 compactor
         previous_end = None
         for start, end in round_completion_info.hb_index_to_entry_range.values():
@@ -338,16 +312,24 @@ def test_compact_partition_rebase_multiple_rounds_same_source_and_destination(
         assert (
             execute_compaction_result_spy.call_args.args[-1] is False
         ), "Table version erroneously marked as in-place compacted!"
-        compacted_delta_locator: DeltaLocator = get_compacted_delta_locator_from_rcf(
-            s3_resource, rcf_file_s3_uri
+        compacted_delta_locator: DeltaLocator = (
+            get_compacted_delta_locator_from_partition(
+                rebased_partition.locator,
+                metastore,
+                catalog=ds_mock_kwargs.get("inner"),
+            )
         )
-        tables = ds.download_delta(
+        tables = metastore.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 sort_keys
+            [(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(