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

deltacat/tests/compute/test_compact_partition_params.py

@@ -1,6 +1,7 @@
 import json
-
+import tempfile
 import unittest
+import uuid
 
 
 class TestCompactPartitionParams(unittest.TestCase):
@@ -8,9 +9,14 @@ class TestCompactPartitionParams(unittest.TestCase):
     def setUpClass(cls):
         from deltacat.types.media import ContentType
         from deltacat.utils.metrics import MetricsConfig, MetricsTarget
+        from deltacat.catalog import CatalogProperties
+
+        # Create a temporary catalog for testing
+        tmpdir = tempfile.mkdtemp()
+        cls.test_catalog = CatalogProperties(root=tmpdir)
 
         cls.VALID_COMPACT_PARTITION_PARAMS = {
-            "compaction_artifact_s3_bucket": "foobar",
+            "catalog": cls.test_catalog,
             "compacted_file_content_type": ContentType.PARQUET,
             "deltacat_storage": "foobar",
             "destination_partition_locator": {
@@ -23,15 +29,16 @@ class TestCompactPartitionParams(unittest.TestCase):
                         "tableVersion": "1",
                     },
                     "streamId": "foobar",
-                    "storageType": "fooType",
+                    "format": "fooType",
                 },
                 "partitionValues": [],
-                "partitionId": None,
+                "partitionId": str(uuid.uuid4()),
             },
             "hash_bucket_count": 200,
             "last_stream_position_to_compact": 168000000000,
             "list_deltas_kwargs": {"equivalent_table_types": []},
             "primary_keys": {"id"},
+            "all_column_names": ["id", "foo", "bar", "baz"],
             "properties": {
                 "parent_stream_position": "1688000000000",
             },
@@ -47,12 +54,12 @@ class TestCompactPartitionParams(unittest.TestCase):
                         "table_version": "1",
                     },
                     "streamId": "foobar",
-                    "storageType": "fooType",
+                    "format": "fooType",
                 },
                 "partitionValues": [],
                 "partitionId": "79612ea39ac5493eae925abe60767d42",
             },
-            "s3_table_writer_kwargs": {
+            "table_writer_kwargs": {
                 "version": "1.0",
                 "flavor": "foobar",
                 "coerce_timestamps": "ms",
@@ -67,7 +74,7 @@ class TestCompactPartitionParams(unittest.TestCase):
                         "tableVersion": "2",
                     },
                     "streamId": "foobar",
-                    "storageType": "fooType",
+                    "format": "fooType",
                 },
                 "partitionValues": [],
                 "partitionId": "79612ea39ac5493eae925abe60767d42",
@@ -103,10 +110,8 @@ class TestCompactPartitionParams(unittest.TestCase):
             json.loads(serialized_params)["compacted_file_content_type"]
             == params.compacted_file_content_type
         )
-        assert (
-            json.loads(serialized_params)["compaction_artifact_s3_bucket"]
-            == params.compaction_artifact_s3_bucket
-        )
+        catalog_json = json.loads(serialized_params)["catalog"]
+        assert catalog_json["_root"] == params.catalog.root
         assert (
             json.loads(serialized_params)["hash_bucket_count"]
             == params.hash_bucket_count

deltacat/tests/compute/test_compact_partition_rebase.py

@@ -1,43 +1,40 @@
-import ray
+import tempfile
 import os
-from moto import mock_s3
+from typing import Any, Callable, Dict, List, Optional, Set
 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,
+    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 (
-    create_src_w_deltas_destination_rebase_w_deltas_strategy,
-)
 from deltacat.tests.compute.compact_partition_rebase_test_cases import (
     REBASE_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,
+    metastore,
 )
-from deltacat.types.media import ContentType
 from deltacat.compute.compactor.model.compact_partition_params import (
     CompactPartitionParams,
 )
@@ -48,11 +45,6 @@ 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",
-)
-
 
 """
 MODULE scoped fixtures
@@ -66,54 +58,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 +148,13 @@ def enable_bucketing_spec_validation(monkeypatch):
     ],
     ids=[test_name for test_name in REBASE_TEST_CASES],
 )
-def test_compact_partition_rebase_same_source_and_destination(
+def test_compact_partition_rebase_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],
     input_deltas_delta_type: str,
@@ -225,21 +173,20 @@ def test_compact_partition_rebase_same_source_and_destination(
     compact_partition_func: Callable,
     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 the scenario where source partition locator == destination partition locator,
     but rebase source partition locator is different.
     This scenario could occur when hash bucket count changes.
+
+    This version uses the main metastore implementation instead of local storage.
     """
     partition_keys = partition_keys_param
     (
         source_table_stream,
         _,
         rebased_table_stream,
-    ) = create_src_w_deltas_destination_rebase_w_deltas_strategy(
-        primary_keys,
+    ) = create_src_w_deltas_destination_rebase_w_deltas_strategy_main(
         sort_keys,
         partition_keys,
         input_deltas_param,
@@ -247,14 +194,31 @@ def test_compact_partition_rebase_same_source_and_destination(
         partition_values_param,
         ds_mock_kwargs,
     )
-    source_partition: Partition = ds.get_partition(
+
+    # 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:
+        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))
+            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,
     )
-    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,
+    )
+    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,
         **ds_mock_kwargs,
     )
     num_workers, worker_instance_cpu = DEFAULT_NUM_WORKERS, DEFAULT_WORKER_INSTANCE_CPUS
@@ -269,10 +233,10 @@ def test_compact_partition_rebase_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,
@@ -284,11 +248,11 @@ def test_compact_partition_rebase_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,
                 "drop_duplicates": drop_duplicates_param,
@@ -305,16 +269,14 @@ def test_compact_partition_rebase_same_source_and_destination(
         object_store_put_many_spy = mocker.spy(FileObjectStore, "put_many")
 
         # 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")
         )
 
-        # 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():
@@ -325,8 +287,12 @@ def test_compact_partition_rebase_same_source_and_destination(
                 == round_completion_info.compacted_pyarrow_write_result.files
             )
 
-        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
@@ -336,13 +302,17 @@ def test_compact_partition_rebase_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"),
+            )
         )
         assert (
             compacted_delta_locator.stream_position == last_stream_position_to_compact
         ), "Compacted delta locator must be equal to last stream position"
-        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)
@@ -351,7 +321,7 @@ def test_compact_partition_rebase_same_source_and_destination(
         if primary_keys:
             sorting_cols.extend([(val, "ascending") for val in primary_keys])
         if sort_keys:
-            sorting_cols.extend(sort_keys)
+            sorting_cols.extend([pa_key for key in sort_keys for pa_key in key.arrow])
 
         rebase_expected_compact_partition_result = (
             rebase_expected_compact_partition_result.combine_chunks().sort_by(