deltacat 0.1.18b14__py3-none-any.whl → 0.1.18b16__py3-none-any.whl

deltacat/tests/{compactor → compute/compactor}/utils/test_io.py

@@ -1,6 +1,12 @@
 import unittest
 from unittest import mock
 from deltacat.tests.test_utils.constants import TEST_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",
+)
 
 
 class TestFitInputDeltas(unittest.TestCase):
@@ -13,6 +19,10 @@ class TestFitInputDeltas(unittest.TestCase):
             CompactionSessionAuditInfo,
         )
 
+        cls.kwargs_for_local_deltacat_storage: Dict[str, Any] = {
+            DATABASE_FILE_PATH_KEY: DATABASE_FILE_PATH_VALUE,
+        }
+
         cls.COMPACTION_AUDIT = CompactionSessionAuditInfo("1.0", "test")
 
         super().setUpClass()
@@ -23,6 +33,7 @@ class TestFitInputDeltas(unittest.TestCase):
 
     def test_sanity(self):
         from deltacat.compute.compactor.utils import io
+        import deltacat.tests.local_deltacat_storage as ds
 
         (
             delta_list,
@@ -30,7 +41,12 @@ class TestFitInputDeltas(unittest.TestCase):
             high_watermark,
             require_multiple_rounds,
         ) = io.fit_input_deltas(
-            [TEST_DELTA], {"CPU": 1, "memory": 20000000}, self.COMPACTION_AUDIT, None
+            [TEST_DELTA],
+            {"CPU": 1, "memory": 20000000},
+            self.COMPACTION_AUDIT,
+            None,
+            ds,
+            self.kwargs_for_local_deltacat_storage,
         )
 
         self.assertIsNotNone(hash_bucket_count)
@@ -44,6 +60,7 @@ class TestFitInputDeltas(unittest.TestCase):
 
     def test_when_hash_bucket_count_overridden(self):
         from deltacat.compute.compactor.utils import io
+        import deltacat.tests.local_deltacat_storage as ds
 
         (
             delta_list,
@@ -51,7 +68,12 @@ class TestFitInputDeltas(unittest.TestCase):
             high_watermark,
             require_multiple_rounds,
         ) = io.fit_input_deltas(
-            [TEST_DELTA], {"CPU": 1, "memory": 20000000}, self.COMPACTION_AUDIT, 20
+            [TEST_DELTA],
+            {"CPU": 1, "memory": 20000000},
+            self.COMPACTION_AUDIT,
+            20,
+            ds,
+            self.kwargs_for_local_deltacat_storage,
         )
 
         self.assertEqual(20, hash_bucket_count)
@@ -61,6 +83,7 @@ class TestFitInputDeltas(unittest.TestCase):
 
     def test_when_not_enough_memory_splits_manifest_entries(self):
         from deltacat.compute.compactor.utils import io
+        import deltacat.tests.local_deltacat_storage as ds
 
         (
             delta_list,
@@ -68,7 +91,12 @@ class TestFitInputDeltas(unittest.TestCase):
             high_watermark,
             require_multiple_rounds,
         ) = io.fit_input_deltas(
-            [TEST_DELTA], {"CPU": 2, "memory": 10}, self.COMPACTION_AUDIT, 20
+            [TEST_DELTA],
+            {"CPU": 2, "memory": 10},
+            self.COMPACTION_AUDIT,
+            20,
+            ds,
+            self.kwargs_for_local_deltacat_storage,
         )
 
         self.assertIsNotNone(hash_bucket_count)
@@ -78,14 +106,28 @@ class TestFitInputDeltas(unittest.TestCase):
 
     def test_when_no_input_deltas(self):
         from deltacat.compute.compactor.utils import io
+        import deltacat.tests.local_deltacat_storage as ds
 
         with self.assertRaises(AssertionError):
             io.fit_input_deltas(
-                [], {"CPU": 100, "memory": 20000.0}, self.COMPACTION_AUDIT, None
+                [],
+                {"CPU": 100, "memory": 20000.0},
+                self.COMPACTION_AUDIT,
+                None,
+                ds,
+                self.kwargs_for_local_deltacat_storage,
             )
 
     def test_when_cpu_resources_is_not_passed(self):
         from deltacat.compute.compactor.utils import io
+        import deltacat.tests.local_deltacat_storage as ds
 
         with self.assertRaises(KeyError):
-            io.fit_input_deltas([], {}, self.COMPACTION_AUDIT, None)
+            io.fit_input_deltas(
+                [],
+                {},
+                self.COMPACTION_AUDIT,
+                None,
+                ds,
+                self.kwargs_for_local_deltacat_storage,
+            )

deltacat/tests/{compactor → compute}/test_compact_partition_params.py

@@ -27,7 +27,7 @@ class TestCompactPartitionParams(unittest.TestCase):
                 "partitionValues": [],
                 "partitionId": None,
             },
-            "hash_bucket_count": None,
+            "hash_bucket_count": 200,
             "last_stream_position_to_compact": 168000000000,
             "list_deltas_kwargs": {"equivalent_table_types": []},
             "primary_keys": {"id"},
@@ -75,41 +75,16 @@ class TestCompactPartitionParams(unittest.TestCase):
 
         super().setUpClass()
 
-    def test_destination_partition_locator_is_optional(self):
-        from deltacat.compute.compactor.model.compact_partition_params import (
-            CompactPartitionParams,
-        )
-
-        params = CompactPartitionParams.of({})
-        assert params.destination_partition_locator is None
-
     def test_serialize_returns_json_string(self):
         from deltacat.compute.compactor.model.compact_partition_params import (
             CompactPartitionParams,
         )
 
         params = CompactPartitionParams.of(
-            {"destination_partition_locator": "my-partition"}
+            TestCompactPartitionParams.VALID_COMPACT_PARTITION_PARAMS
         )
         serialized_params = params.serialize()
         assert isinstance(serialized_params, str)
-        assert json.loads(serialized_params) == {
-            "compacted_file_content_type": None,
-            "compaction_artifact_s3_bucket": None,
-            "deltacat_storage": None,
-            "hash_bucket_count": None,
-            "last_stream_position_to_compact": None,
-            "list_deltas_kwargs": None,
-            "pg_config": None,
-            "primary_keys": None,
-            "properties": None,
-            "read_kwargs_provider": None,
-            "rebase_source_partition_high_watermark": None,
-            "rebase_source_partition_locator": None,
-            "s3_table_writer_kwargs": None,
-            "source_partition_locator": None,
-            "destination_partition_locator": "my-partition",
-        }
 
     def test_serialize_returns_json_string_with_all_fields(self):
         from deltacat.compute.compactor.model.compact_partition_params import (
@@ -142,7 +117,6 @@ class TestCompactPartitionParams(unittest.TestCase):
             == params.list_deltas_kwargs
         )
         assert json.loads(serialized_params)["primary_keys"] == params.primary_keys
-        assert json.loads(serialized_params)["properties"] == params.properties
         assert (
             json.loads(serialized_params)["rebase_source_partition_high_watermark"]
             == params.rebase_source_partition_high_watermark
@@ -165,7 +139,12 @@ class TestCompactPartitionParams(unittest.TestCase):
             CompactPartitionParams,
         )
 
-        params = CompactPartitionParams.of({"primary_keys": {"foo", "bar", "baz"}})
+        params = CompactPartitionParams.of(
+            {
+                **TestCompactPartitionParams.VALID_COMPACT_PARTITION_PARAMS,
+                "primary_keys": {"foo", "bar", "baz"},
+            }
+        )
         serialized_params = params.serialize()
         self.assertCountEqual(
             json.loads(serialized_params)["primary_keys"], ["foo", "bar", "baz"]
@@ -180,7 +159,12 @@ class TestCompactPartitionParams(unittest.TestCase):
             def toJSON(self) -> str:
                 return "my-json-object"
 
-        params = CompactPartitionParams.of({"compacted_file_content_type": MyObject()})
+        params = CompactPartitionParams.of(
+            {
+                **TestCompactPartitionParams.VALID_COMPACT_PARTITION_PARAMS,
+                "compacted_file_content_type": MyObject(),
+            }
+        )
         serialized_params = params.serialize()
         assert (
             json.loads(serialized_params)["compacted_file_content_type"]

deltacat/tests/compute/test_compaction_session_incremental.py

@@ -0,0 +1,363 @@
+import ray
+from moto import mock_s3
+import pytest
+import os
+import json
+import boto3
+from typing import Any, Dict, List, Optional, Set
+from boto3.resources.base import ServiceResource
+import pyarrow as pa
+from deltacat.tests.test_utils.utils import read_s3_contents
+from deltacat.tests.compute.common import (
+    setup_sort_and_partition_keys,
+    PartitionKey,
+    TEST_S3_RCF_BUCKET_NAME,
+    BASE_TEST_SOURCE_NAMESPACE,
+    BASE_TEST_SOURCE_TABLE_NAME,
+    BASE_TEST_DESTINATION_NAMESPACE,
+    BASE_TEST_DESTINATION_TABLE_NAME,
+)
+from deltacat.tests.compute.testcases import (
+    INCREMENTAL_TEST_CASES,
+)
+
+DATABASE_FILE_PATH_KEY, DATABASE_FILE_PATH_VALUE = (
+    "db_file_path",
+    "deltacat/tests/local_deltacat_storage/db_test.sqlite",
+)
+
+"""
+MODULE scoped fixtures
+"""
+
+
+@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 setup_s3_resource(mock_aws_credential):
+    with mock_s3():
+        yield boto3.resource("s3")
+
+
+@pytest.fixture(scope="module")
+def setup_compaction_artifacts_s3_bucket(setup_s3_resource: ServiceResource):
+    setup_s3_resource.create_bucket(
+        ACL="authenticated-read",
+        Bucket=TEST_S3_RCF_BUCKET_NAME,
+    )
+    yield
+
+
+"""
+FUNCTION scoped fixtures
+"""
+
+
+@pytest.fixture(scope="function")
+def teardown_local_deltacat_storage_db():
+    if os.path.exists(DATABASE_FILE_PATH_VALUE):
+        os.remove(DATABASE_FILE_PATH_VALUE)
+
+
+@pytest.fixture(scope="function")
+def setup_local_deltacat_storage_conn():
+    # 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
+
+
+def setup_incremental_source_and_destination_tables(
+    source_table_version: str,
+    destination_table_version: str,
+    primary_keys: Set[str],
+    sort_keys: Optional[List[Any]],
+    partition_keys: Optional[List[PartitionKey]],
+    column_names: List[str],
+    arrow_arrays: List[pa.Array],
+    partition_values: Optional[List[Any]],
+    ds_mock_kwargs: Optional[Dict[str, Any]],
+    source_namespace: str = BASE_TEST_SOURCE_NAMESPACE,
+    source_table_name: str = BASE_TEST_SOURCE_TABLE_NAME,
+    destination_namespace: str = BASE_TEST_DESTINATION_NAMESPACE,
+    destination_table_name: str = BASE_TEST_DESTINATION_TABLE_NAME,
+):
+    import deltacat.tests.local_deltacat_storage as ds
+    from deltacat.types.media import ContentType
+    from deltacat.storage import Partition, Stream
+
+    ds.create_namespace(source_namespace, {}, **ds_mock_kwargs)
+    ds.create_table_version(
+        source_namespace,
+        source_table_name,
+        source_table_version,
+        primary_key_column_names=list(primary_keys),
+        sort_keys=sort_keys,
+        partition_keys=partition_keys,
+        supported_content_types=[ContentType.PARQUET],
+        **ds_mock_kwargs,
+    )
+    source_table_stream: Stream = ds.get_stream(
+        namespace=source_namespace,
+        table_name=source_table_name,
+        table_version=source_table_version,
+        **ds_mock_kwargs,
+    )
+    test_table: pa.Table = pa.Table.from_arrays(arrow_arrays, names=column_names)
+    staged_partition: Partition = ds.stage_partition(
+        source_table_stream, partition_values, **ds_mock_kwargs
+    )
+    ds.commit_delta(
+        ds.stage_delta(test_table, staged_partition, **ds_mock_kwargs), **ds_mock_kwargs
+    )
+    ds.commit_partition(staged_partition, **ds_mock_kwargs)
+    # create the destination table
+    ds.create_namespace(destination_namespace, {}, **ds_mock_kwargs)
+    ds.create_table_version(
+        destination_namespace,
+        destination_table_name,
+        destination_table_version,
+        primary_key_column_names=list(primary_keys),
+        sort_keys=sort_keys,
+        partition_keys=partition_keys,
+        supported_content_types=[ContentType.PARQUET],
+        **ds_mock_kwargs,
+    )
+    destination_table_stream: Stream = ds.get_stream(
+        namespace=destination_namespace,
+        table_name=destination_table_name,
+        table_version=destination_table_version,
+        **ds_mock_kwargs,
+    )
+    source_table_stream_after_committed: Stream = ds.get_stream(
+        namespace=source_namespace,
+        table_name=source_table_name,
+        table_version=source_table_version,
+        **ds_mock_kwargs,
+    )
+    return source_table_stream_after_committed, destination_table_stream
+
+
+@pytest.mark.parametrize(
+    [
+        "test_name",
+        "source_table_version",
+        "destination_table_version",
+        "primary_keys_param",
+        "sort_keys_param",
+        "partition_keys_param",
+        "column_names_param",
+        "arrow_arrays_param",
+        "rebase_source_partition_locator_param",
+        "partition_values_param",
+        "expected_result",
+        "validation_callback_func",
+        "validation_callback_func_kwargs",
+        "do_teardown_local_deltacat_storage_db",
+        "use_prev_compacted",
+        "create_placement_group_param",
+        "records_per_compacted_file_param",
+        "hash_bucket_count_param",
+        "compact_partition_func",
+    ],
+    [
+        (
+            test_name,
+            source_table_version,
+            destination_table_version,
+            primary_keys_param,
+            sort_keys_param,
+            partition_keys_param,
+            column_names_param,
+            arrow_arrays_param,
+            rebase_source_partition_locator_param,
+            partition_values_param,
+            expected_result,
+            validation_callback_func,
+            validation_callback_func_kwargs,
+            do_teardown_local_deltacat_storage_db,
+            use_prev_compacted,
+            create_placement_group_param,
+            records_per_compacted_file_param,
+            hash_bucket_count_param,
+            compact_partition_func,
+        )
+        for test_name, (
+            source_table_version,
+            destination_table_version,
+            primary_keys_param,
+            sort_keys_param,
+            partition_keys_param,
+            column_names_param,
+            arrow_arrays_param,
+            rebase_source_partition_locator_param,
+            partition_values_param,
+            expected_result,
+            validation_callback_func,
+            validation_callback_func_kwargs,
+            do_teardown_local_deltacat_storage_db,
+            use_prev_compacted,
+            create_placement_group_param,
+            records_per_compacted_file_param,
+            hash_bucket_count_param,
+            compact_partition_func,
+        ) in INCREMENTAL_TEST_CASES.items()
+    ],
+    ids=[test_name for test_name in INCREMENTAL_TEST_CASES],
+    indirect=[],
+)
+def test_compact_partition_incremental(
+    request: pytest.FixtureRequest,
+    setup_s3_resource: ServiceResource,
+    setup_local_deltacat_storage_conn: Dict[str, Any],
+    setup_compaction_artifacts_s3_bucket: None,
+    test_name: str,
+    source_table_version: str,
+    destination_table_version: str,
+    primary_keys_param: Set[str],
+    sort_keys_param,
+    partition_keys_param,
+    column_names_param: List[str],
+    arrow_arrays_param: List[pa.Array],
+    rebase_source_partition_locator_param,
+    partition_values_param,
+    expected_result,
+    # use and implement func and func_kwargs if you want to run additional validations apart from the ones in the test
+    validation_callback_func,
+    validation_callback_func_kwargs,
+    do_teardown_local_deltacat_storage_db,
+    use_prev_compacted,
+    create_placement_group_param,
+    records_per_compacted_file_param,
+    hash_bucket_count_param,
+    compact_partition_func,
+):
+    import deltacat.tests.local_deltacat_storage as ds
+    from deltacat.types.media import ContentType
+    from deltacat.storage import (
+        PartitionLocator,
+    )
+    from deltacat.compute.compactor.model.compact_partition_params import (
+        CompactPartitionParams,
+    )
+    from deltacat.utils.placement import (
+        PlacementGroupManager,
+    )
+    from deltacat.compute.compactor import (
+        RoundCompletionInfo,
+    )
+
+    ds_mock_kwargs = setup_local_deltacat_storage_conn
+
+    # setup
+    sort_keys, partition_keys = setup_sort_and_partition_keys(
+        sort_keys_param, partition_keys_param
+    )
+    (
+        source_table_stream,
+        destination_table_stream,
+    ) = setup_incremental_source_and_destination_tables(
+        source_table_version,
+        destination_table_version,
+        primary_keys_param,
+        sort_keys,
+        partition_keys,
+        column_names_param,
+        arrow_arrays_param,
+        partition_values_param,
+        ds_mock_kwargs,
+        f"{test_name}_src_namespace",
+        f"{test_name}_table_src",
+        f"{test_name}_dest_namespace",
+        f"{test_name}_table_dest",
+    )
+    ray.shutdown()
+    ray.init(local_mode=True)
+    assert ray.is_initialized()
+    source_partition = ds.get_partition(
+        source_table_stream.locator,
+        partition_values_param,
+        **ds_mock_kwargs,
+    )
+    destination_partition_locator = PartitionLocator.of(
+        destination_table_stream.locator,
+        partition_values_param,
+        None,
+    )
+    num_workers, worker_instance_cpu = 1, 1
+    total_cpus = num_workers * worker_instance_cpu
+    pgm = None
+    if create_placement_group_param:
+        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": []}},
+            "pg_config": pgm,
+            "primary_keys": primary_keys_param,
+            "rebase_source_partition_locator": rebase_source_partition_locator_param,
+            "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,
+        }
+    )
+    # execute
+    rcf_file_s3_uri = compact_partition_func(compact_partition_params)
+    # validate
+    _, rcf_object_key = rcf_file_s3_uri.rsplit("/", 1)
+    rcf_file_output: Dict[str, Any] = read_s3_contents(
+        setup_s3_resource, TEST_S3_RCF_BUCKET_NAME, rcf_object_key
+    )
+    round_completion_info = RoundCompletionInfo(**rcf_file_output)
+    print(f"rcf_file_output: {json.dumps(rcf_file_output, indent=2)}")
+    compacted_delta_locator = round_completion_info.compacted_delta_locator
+    tables = ds.download_delta(compacted_delta_locator, **ds_mock_kwargs)
+    compacted_table = pa.concat_tables(tables)
+
+    # the compacted table may contain multiple files and chunks
+    # and order of records may be incorrect due to multiple files.
+    expected_result = expected_result.combine_chunks().sort_by(
+        [(val, "ascending") for val in primary_keys_param]
+    )
+    compacted_table = compacted_table.combine_chunks().sort_by(
+        [(val, "ascending") for val in primary_keys_param]
+    )
+
+    assert compacted_table.equals(
+        expected_result
+    ), f"{compacted_table} does not match {expected_result}"
+    if (
+        validation_callback_func is not None
+        and validation_callback_func_kwargs is not None
+    ):
+        validation_callback_func(**validation_callback_func_kwargs)
+    # https://docs.pytest.org/en/7.1.x/reference/reference.html#pytest.FixtureRequest.getfixturevalue
+    if do_teardown_local_deltacat_storage_db:
+        request.getfixturevalue("teardown_local_deltacat_storage_db")
+    return