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

deltacat/tests/compute/compactor_v2/steps/test_hash_bucket.py

@@ -0,0 +1,199 @@
+import unittest
+import sqlite3
+import ray
+import os
+from collections import defaultdict
+from deltacat.compute.compactor import DeltaAnnotated
+import deltacat.tests.local_deltacat_storage as ds
+from deltacat.io.ray_plasma_object_store import RayPlasmaObjectStore
+from deltacat.compute.compactor_v2.model.hash_bucket_input import HashBucketInput
+from deltacat.compute.compactor_v2.model.hash_bucket_result import HashBucketResult
+from deltacat.compute.compactor_v2.steps.hash_bucket import hash_bucket
+from deltacat.utils.common import current_time_ms
+from deltacat.tests.test_utils.pyarrow import create_delta_from_csv_file
+
+
+class TestHashBucket(unittest.TestCase):
+    HASH_BUCKET_NAMESPACE = "test_hash_bucket"
+    DB_FILE_PATH = f"{current_time_ms()}.db"
+    STRING_PK_FILE_PATH = (
+        "deltacat/tests/compute/compactor_v2/steps/data/string_pk_table.csv"
+    )
+    DATE_PK_FILE_PATH = (
+        "deltacat/tests/compute/compactor_v2/steps/data/date_pk_table.csv"
+    )
+    MULTIPLE_PK_FILE_PATH = (
+        "deltacat/tests/compute/compactor_v2/steps/data/multiple_pk_table.csv"
+    )
+    NO_PK_FILE_PATH = "deltacat/tests/compute/compactor_v2/steps/data/no_pk_table.csv"
+
+    @classmethod
+    def setUpClass(cls):
+        ray.init(local_mode=True, ignore_reinit_error=True)
+
+        con = sqlite3.connect(cls.DB_FILE_PATH)
+        cur = con.cursor()
+        cls.kwargs = {ds.SQLITE_CON_ARG: con, ds.SQLITE_CUR_ARG: cur}
+        cls.deltacat_storage_kwargs = {ds.DB_FILE_PATH_ARG: cls.DB_FILE_PATH}
+
+        super().setUpClass()
+
+    @classmethod
+    def doClassCleanups(cls) -> None:
+        os.remove(cls.DB_FILE_PATH)
+
+    def test_single_string_pk_correctly_hashes(self):
+        # setup
+        delta = create_delta_from_csv_file(
+            self.HASH_BUCKET_NAMESPACE, [self.STRING_PK_FILE_PATH], **self.kwargs
+        )
+
+        annotated_delta = DeltaAnnotated.of(delta)
+        object_store = RayPlasmaObjectStore()
+        hb_input = HashBucketInput.of(
+            annotated_delta=annotated_delta,
+            primary_keys=["pk"],
+            num_hash_buckets=3,
+            num_hash_groups=2,
+            deltacat_storage=ds,
+            deltacat_storage_kwargs=self.deltacat_storage_kwargs,
+            object_store=object_store,
+        )
+
+        # action
+        hb_result_promise = hash_bucket.remote(hb_input)
+        hb_result: HashBucketResult = ray.get(hb_result_promise)
+
+        # assert
+        # PK hash column is also persisted.
+        self._validate_hash_bucket_result(
+            hb_result,
+            record_count=6,
+            num_hash_buckets=3,
+            num_columns=3,
+            object_store=object_store,
+        )
+
+    def test_single_date_pk_correctly_hashes(self):
+        # setup
+        delta = create_delta_from_csv_file(
+            self.HASH_BUCKET_NAMESPACE, [self.DATE_PK_FILE_PATH], **self.kwargs
+        )
+
+        annotated_delta = DeltaAnnotated.of(delta)
+        object_store = RayPlasmaObjectStore()
+        hb_input = HashBucketInput.of(
+            annotated_delta=annotated_delta,
+            primary_keys=["pk"],
+            num_hash_buckets=2,
+            num_hash_groups=1,
+            deltacat_storage=ds,
+            deltacat_storage_kwargs=self.deltacat_storage_kwargs,
+            object_store=object_store,
+        )
+
+        # action
+        hb_result_promise = hash_bucket.remote(hb_input)
+        hb_result: HashBucketResult = ray.get(hb_result_promise)
+
+        # assert
+        self._validate_hash_bucket_result(
+            hb_result,
+            record_count=7,
+            num_hash_buckets=2,
+            num_columns=3,
+            object_store=object_store,
+        )
+
+    def test_no_pk_does_not_hash(self):
+        # setup
+        delta = create_delta_from_csv_file(
+            self.HASH_BUCKET_NAMESPACE, [self.NO_PK_FILE_PATH], **self.kwargs
+        )
+
+        annotated_delta = DeltaAnnotated.of(delta)
+        object_store = RayPlasmaObjectStore()
+        hb_input = HashBucketInput.of(
+            annotated_delta=annotated_delta,
+            primary_keys=[],
+            num_hash_buckets=2,
+            num_hash_groups=1,
+            deltacat_storage=ds,
+            deltacat_storage_kwargs=self.deltacat_storage_kwargs,
+            object_store=object_store,
+        )
+
+        # action
+        hb_result_promise = hash_bucket.remote(hb_input)
+        hb_result: HashBucketResult = ray.get(hb_result_promise)
+
+        # assert
+        self._validate_hash_bucket_result(
+            hb_result,
+            record_count=6,
+            num_hash_buckets=2,
+            num_columns=3,
+            object_store=object_store,
+        )
+
+    def test_multiple_pk_correctly_hashes(self):
+        # setup
+        delta = create_delta_from_csv_file(
+            self.HASH_BUCKET_NAMESPACE, [self.MULTIPLE_PK_FILE_PATH], **self.kwargs
+        )
+
+        annotated_delta = DeltaAnnotated.of(delta)
+        object_store = RayPlasmaObjectStore()
+        hb_input = HashBucketInput.of(
+            annotated_delta=annotated_delta,
+            primary_keys=["pk1", "pk2"],
+            num_hash_buckets=2,
+            num_hash_groups=1,
+            deltacat_storage=ds,
+            deltacat_storage_kwargs=self.deltacat_storage_kwargs,
+            object_store=object_store,
+        )
+
+        # action
+        hb_result_promise = hash_bucket.remote(hb_input)
+        hb_result: HashBucketResult = ray.get(hb_result_promise)
+
+        # assert
+        self._validate_hash_bucket_result(
+            hb_result,
+            record_count=6,
+            num_hash_buckets=2,
+            num_columns=4,
+            object_store=object_store,
+        )
+
+    def _validate_hash_bucket_result(
+        self,
+        hb_result: HashBucketResult,
+        record_count: int,
+        num_hash_buckets: int,
+        num_columns: int,
+        object_store,
+    ):
+
+        self.assertEqual(hb_result.hb_record_count, record_count)
+        self.assertIsNotNone(hb_result)
+        self.assertIsNotNone(hb_result.peak_memory_usage_bytes)
+        self.assertIsNotNone(hb_result.task_completed_at)
+        self.assertIsNotNone(hb_result.telemetry_time_in_seconds)
+
+        hb_index_to_dfes = defaultdict(list)
+        total_records_in_result = 0
+        for _, object_id in enumerate(hb_result.hash_bucket_group_to_obj_id_tuple):
+            if object_id:
+                obj = object_store.get(object_id[0])
+                for hb_idx, dfes in enumerate(obj):
+                    if dfes is not None:
+                        hb_index_to_dfes[hb_idx].extend(dfes)
+                        for dfe in dfes:
+                            self.assertIsNotNone(dfe)
+                            total_records_in_result += len(dfe.table)
+                            self.assertEqual(num_columns, len(dfe.table.column_names))
+
+        self.assertTrue(len(hb_index_to_dfes) <= num_hash_buckets)
+        self.assertEqual(total_records_in_result, record_count)

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,348 @@
+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,
+    validation_callback_func,  # use and implement func and func_kwargs if you want to run additional validations apart from the ones in the test
+    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,
+    )
+    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)
+    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