deltacat 0.1.18b3__py3-none-any.whl → 0.1.18b7__py3-none-any.whl

deltacat/io/object_store.py

@@ -0,0 +1,51 @@
+from typing import List, Any
+
+
+class IObjectStore:
+    """
+    An object store interface.
+    """
+
+    def setup(self, *args, **kwargs) -> Any:
+        ...
+
+        """
+        Sets up everything needed to run the object store.
+        """
+
+    def put(self, obj: object, *args, **kwargs) -> Any:
+        """
+        Put a single object into the object store. Depending
+        on the implementation, this method can be sync or async.
+        """
+        return self.put_many([obj])[0]
+
+    def put_many(self, objects: List[object], *args, **kwargs) -> List[Any]:
+        ...
+
+        """
+        Put many objects into the object store. It would return an ordered list
+        of object references corresponding to each object in the input.
+        """
+
+    def get(self, ref: Any, *args, **kwargs) -> object:
+        """
+        Get a single object from an object store.
+        """
+        return self.get_many([ref])[0]
+
+    def get_many(self, refs: List[Any], *args, **kwargs) -> List[object]:
+        ...
+
+        """
+        Get a list of objects from the object store. Use this method to
+        avoid multiple get calls. Note that depending on implementation it may
+        or may not return ordered results.
+        """
+
+    def clear(self, *args, **kwargs) -> bool:
+        ...
+
+        """
+        Clears the object store and all the associated data in it.
+        """

deltacat/io/ray_plasma_object_store.py

@@ -0,0 +1,23 @@
+import ray
+from ray import cloudpickle
+from deltacat.io.object_store import IObjectStore
+from typing import Any, List
+
+
+class RayPlasmaObjectStore(IObjectStore):
+    """
+    An implementation of object store that uses Ray plasma object store.
+    """
+
+    def put_many(self, objects: List[object], *args, **kwargs) -> List[Any]:
+        result = []
+        for obj in objects:
+            object_ref = ray.put(obj)
+            pickled = cloudpickle.dumps(object_ref)
+            result.append(pickled)
+
+        return result
+
+    def get_many(self, refs: List[Any], *args, **kwargs) -> List[object]:
+        loaded_refs = [cloudpickle.loads(obj_id) for obj_id in refs]
+        return ray.get(loaded_refs)

deltacat/io/redis_object_store.py

@@ -0,0 +1,114 @@
+import logging
+from ray import cloudpickle
+import time
+from deltacat.io.object_store import IObjectStore
+from typing import Any, List
+from deltacat import logs
+import uuid
+import socket
+import redis
+from collections import defaultdict
+
+
+logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
+
+
+class RedisObjectStore(IObjectStore):
+    """
+    An implementation of object store that uses Redis in memory DB.
+    """
+
+    def __init__(self) -> None:
+        self.client_cache = {}
+        self.current_ip = None
+        self.SEPARATOR = "_"
+        super().__init__()
+
+    def put(self, obj: object, *args, **kwargs) -> Any:
+        serialized = cloudpickle.dumps(obj)
+        uid = uuid.uuid4()
+        current_ip = self._get_current_ip()
+        ref = self._create_ref(uid, current_ip)
+
+        client = self._get_client_by_ip(current_ip)
+        if client.set(uid.__str__(), serialized):
+            return ref
+        else:
+            raise RuntimeError(f"Unable to write {ref} to cache")
+
+    def put_many(self, objects: List[object], *args, **kwargs) -> List[Any]:
+        input = {}
+        result = []
+        current_ip = self._get_current_ip()
+        for obj in objects:
+            serialized = cloudpickle.dumps(obj)
+            uid = uuid.uuid4()
+            ref = self._create_ref(uid, current_ip)
+            input[uid.__str__()] = serialized
+            result.append(ref)
+
+        client = self._get_client_by_ip(current_ip)
+
+        if client.mset(input):
+            return result
+        else:
+            raise RuntimeError("Unable to update cache")
+
+    def get_many(self, refs: List[Any], *args, **kwargs) -> List[object]:
+        result = []
+        uid_per_ip = defaultdict(lambda: [])
+
+        start = time.monotonic()
+        for ref in refs:
+            uid, ip = ref.split(self.SEPARATOR)
+            uid_per_ip[ip].append(uid)
+
+        for (ip, uids) in uid_per_ip.items():
+            client = self._get_client_by_ip(ip)
+            cache_result = client.mget(uids)
+            assert len(cache_result) == len(
+                uids
+            ), "Not all values were returned from cache"
+
+            total_bytes = 0
+
+            deserialize_start = time.monotonic()
+            for serialized in cache_result:
+                deserialized = cloudpickle.loads(serialized)
+                total_bytes += len(serialized)
+                result.append(deserialized)
+
+            deserialize_end = time.monotonic()
+            logger.debug(
+                f"The time taken to deserialize {total_bytes} bytes is: {deserialize_end - deserialize_start}",
+            )
+
+        end = time.monotonic()
+
+        logger.info(f"The total time taken to read all objects is: {end - start}")
+
+        return result
+
+    def get(self, ref: Any, *args, **kwargs) -> object:
+        uid, ip = ref.split(self.SEPARATOR)
+        client = self._get_client_by_ip(ip)
+        serialized = client.get(uid)
+        return cloudpickle.loads(serialized)
+
+    def _get_client_by_ip(self, ip_address: str):
+        if ip_address in self.client_cache:
+            return self.client_cache[ip_address]
+
+        base_client = redis.Redis(ip_address, 7777)
+
+        self.client_cache[ip_address] = base_client
+        return base_client
+
+    def _get_current_ip(self):
+        if self.current_ip is None:
+            self.current_ip = socket.gethostbyname(socket.gethostname())
+
+        return self.current_ip
+
+    def _create_ref(self, uid, ip):
+        return f"{uid}{self.SEPARATOR}{ip}"

deltacat/io/s3_object_store.py

@@ -0,0 +1,44 @@
+import logging
+from ray import cloudpickle
+import time
+from deltacat.io.object_store import IObjectStore
+from typing import Any, List
+from deltacat import logs
+import uuid
+from deltacat.aws import s3u as s3_utils
+
+logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
+
+
+class S3ObjectStore(IObjectStore):
+    """
+    An implementation of object store that uses S3.
+    """
+
+    def __init__(self, bucket_prefix: str) -> None:
+        self.bucket = bucket_prefix
+        super().__init__()
+
+    def put_many(self, objects: List[object], *args, **kwargs) -> List[Any]:
+        result = []
+        for obj in objects:
+            serialized = cloudpickle.dumps(obj)
+            ref = uuid.uuid4()
+
+            s3_utils.upload(f"s3://{self.bucket}/{ref}", serialized)
+            result.append(ref)
+
+        return result
+
+    def get_many(self, refs: List[Any], *args, **kwargs) -> List[object]:
+        result = []
+        start = time.monotonic()
+        for ref in refs:
+            cur = s3_utils.download(f"s3://{self.bucket}/{ref}")
+            serialized = cur["Body"].read()
+            loaded = cloudpickle.loads(serialized)
+            result.append(loaded)
+        end = time.monotonic()
+
+        logger.info(f"The total time taken to read all objects is: {end - start}")
+        return result

deltacat/storage/model/delta.py

@@ -256,7 +256,8 @@ class Delta(dict):
 class DeltaLocator(Locator, dict):
     @staticmethod
     def of(
-        partition_locator: Optional[PartitionLocator], stream_position: Optional[int]
+        partition_locator: Optional[PartitionLocator] = None,
+        stream_position: Optional[int] = None,
     ) -> DeltaLocator:
         """
         Creates a partition delta locator. Stream Position, if provided, should

deltacat/tests/compactor/test_compact_partition_params.py

@@ -0,0 +1,237 @@
+import json
+
+import unittest
+
+
+class TestCompactPartitionParams(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        from deltacat.types.media import ContentType
+
+        cls.VALID_COMPACT_PARTITION_PARAMS = {
+            "compaction_artifact_s3_bucket": "foobar",
+            "compacted_file_content_type": ContentType.PARQUET,
+            "deltacat_storage": "foobar",
+            "destination_partition_locator": {
+                "streamLocator": {
+                    "tableVersionLocator": {
+                        "tableLocator": {
+                            "namespaceLocator": {"namespace": "testNamespaceLocator"},
+                            "tableName": "TABLE_FOO",
+                        },
+                        "tableVersion": "1",
+                    },
+                    "streamId": "foobar",
+                    "storageType": "fooType",
+                },
+                "partitionValues": [],
+                "partitionId": None,
+            },
+            "hash_bucket_count": None,
+            "last_stream_position_to_compact": 168000000000,
+            "list_deltas_kwargs": {"equivalent_table_types": []},
+            "primary_keys": {"id"},
+            "properties": {
+                "parent_stream_position": "1688000000000",
+            },
+            "read_kwargs_provider": "foo",
+            "rebase_source_partition_high_watermark": 1688000000000,
+            "rebase_source_partition_locator": {
+                "stream_locator": {
+                    "table_version_locator": {
+                        "table_locator": {
+                            "namespace_locator": {"namespace": "testNamespaceLocator"},
+                            "tableName": "TABLE_FOO",
+                        },
+                        "table_version": "1",
+                    },
+                    "streamId": "foobar",
+                    "storageType": "fooType",
+                },
+                "partitionValues": [],
+                "partitionId": "79612ea39ac5493eae925abe60767d42",
+            },
+            "s3_table_writer_kwargs": {
+                "version": "1.0",
+                "flavor": "foobar",
+                "coerce_timestamps": "ms",
+            },
+            "source_partition_locator": {
+                "streamLocator": {
+                    "tableVersionLocator": {
+                        "tableLocator": {
+                            "namespaceLocator": {"namespace": "testNamespaceLocator"},
+                            "tableName": "TABLE_FOO",
+                        },
+                        "tableVersion": "2",
+                    },
+                    "streamId": "foobar",
+                    "storageType": "fooType",
+                },
+                "partitionValues": [],
+                "partitionId": "79612ea39ac5493eae925abe60767d42",
+            },
+        }
+
+        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"}
+        )
+        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 (
+            CompactPartitionParams,
+        )
+
+        params = CompactPartitionParams.of(
+            TestCompactPartitionParams.VALID_COMPACT_PARTITION_PARAMS
+        )
+        serialized_params = params.serialize()
+        assert isinstance(serialized_params, str)
+        assert (
+            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
+        )
+        assert (
+            json.loads(serialized_params)["hash_bucket_count"]
+            == params.hash_bucket_count
+        )
+        assert (
+            json.loads(serialized_params)["last_stream_position_to_compact"]
+            == params.last_stream_position_to_compact
+        )
+        assert (
+            json.loads(serialized_params)["list_deltas_kwargs"]
+            == 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
+        )
+        assert (
+            json.loads(serialized_params)["rebase_source_partition_locator"]
+            == params.rebase_source_partition_locator
+        )
+        assert (
+            json.loads(serialized_params)["source_partition_locator"]
+            == params.source_partition_locator
+        )
+        assert (
+            json.loads(serialized_params)["destination_partition_locator"]
+            == params.destination_partition_locator
+        )
+
+    def test_serialize_handles_sets(self):
+        from deltacat.compute.compactor.model.compact_partition_params import (
+            CompactPartitionParams,
+        )
+
+        params = CompactPartitionParams.of({"primary_keys": {"foo", "bar", "baz"}})
+        serialized_params = params.serialize()
+        self.assertCountEqual(
+            json.loads(serialized_params)["primary_keys"], ["foo", "bar", "baz"]
+        )
+
+    def test_serialize_handles_objects_with_toJSON_method(self):
+        from deltacat.compute.compactor.model.compact_partition_params import (
+            CompactPartitionParams,
+        )
+
+        class MyObject:
+            def toJSON(self) -> str:
+                return "my-json-object"
+
+        params = CompactPartitionParams.of({"compacted_file_content_type": MyObject()})
+        serialized_params = params.serialize()
+        assert (
+            json.loads(serialized_params)["compacted_file_content_type"]
+            == "my-json-object"
+        )
+
+    def test_json_handler_for_compact_partition_params_serializes_set_to_list(self):
+        from deltacat.compute.compactor.model.compact_partition_params import (
+            CompactPartitionParams,
+        )
+
+        my_set = {1, 2, 3}
+        json_string = json.dumps(
+            my_set,
+            default=CompactPartitionParams.json_handler_for_compact_partition_params,
+        )
+        assert json.loads(json_string) == [1, 2, 3]
+
+    def test_json_handler_for_compact_partition_params_serializes_object_with_toJSON_method_to_dict(
+        self,
+    ):
+        from dataclasses import dataclass
+
+        from deltacat.compute.compactor.model.compact_partition_params import (
+            CompactPartitionParams,
+        )
+
+        @dataclass
+        class DummyObject:
+            some_property: str = "foo"
+
+            def toJSON(self):
+                return self.__dict__
+
+        dummy_object = DummyObject()
+        json_string = json.dumps(
+            dummy_object,
+            default=CompactPartitionParams.json_handler_for_compact_partition_params,
+        )
+        assert json.loads(json_string) == {"some_property": "foo"}
+
+    def test_json_handler_for_compact_partition_params_returns_class_name_for_unknown_objects(
+        self,
+    ):
+        from deltacat.compute.compactor.model.compact_partition_params import (
+            CompactPartitionParams,
+        )
+
+        my_object = object()
+        json_string = json.dumps(
+            my_object,
+            default=CompactPartitionParams.json_handler_for_compact_partition_params,
+        )
+        assert json.loads(json_string) == "object"

deltacat/tests/compactor/utils/test_io.py

@@ -9,8 +9,18 @@ class TestFitInputDeltas(unittest.TestCase):
         cls.module_patcher = mock.patch.dict("sys.modules", {"ray": mock.MagicMock()})
         cls.module_patcher.start()
 
+        from deltacat.compute.compactor.model.compaction_session_audit_info import (
+            CompactionSessionAuditInfo,
+        )
+
+        cls.COMPACTION_AUDIT = CompactionSessionAuditInfo("1.0", "test")
+
         super().setUpClass()
 
+    @classmethod
+    def tearDownClass(cls) -> None:
+        cls.module_patcher.stop()
+
     def test_sanity(self):
         from deltacat.compute.compactor.utils import io
 
@@ -19,12 +29,18 @@ class TestFitInputDeltas(unittest.TestCase):
             hash_bucket_count,
             high_watermark,
             require_multiple_rounds,
-        ) = io.fit_input_deltas([TEST_DELTA], {"CPU": 1, "memory": 20000000}, None)
+        ) = io.fit_input_deltas(
+            [TEST_DELTA], {"CPU": 1, "memory": 20000000}, self.COMPACTION_AUDIT, None
+        )
 
         self.assertIsNotNone(hash_bucket_count)
         self.assertTrue(1, len(delta_list))
         self.assertIsNotNone(high_watermark)
         self.assertFalse(require_multiple_rounds)
+        self.assertIsNotNone(hash_bucket_count, self.COMPACTION_AUDIT.hash_bucket_count)
+        self.assertIsNotNone(self.COMPACTION_AUDIT.input_file_count)
+        self.assertIsNotNone(self.COMPACTION_AUDIT.input_size_bytes)
+        self.assertIsNotNone(self.COMPACTION_AUDIT.total_cluster_memory_bytes)
 
     def test_when_hash_bucket_count_overridden(self):
         from deltacat.compute.compactor.utils import io
@@ -34,7 +50,9 @@ class TestFitInputDeltas(unittest.TestCase):
             hash_bucket_count,
             high_watermark,
             require_multiple_rounds,
-        ) = io.fit_input_deltas([TEST_DELTA], {"CPU": 1, "memory": 20000000}, 20)
+        ) = io.fit_input_deltas(
+            [TEST_DELTA], {"CPU": 1, "memory": 20000000}, self.COMPACTION_AUDIT, 20
+        )
 
         self.assertEqual(20, hash_bucket_count)
         self.assertEqual(1, len(delta_list))
@@ -49,7 +67,9 @@ class TestFitInputDeltas(unittest.TestCase):
             hash_bucket_count,
             high_watermark,
             require_multiple_rounds,
-        ) = io.fit_input_deltas([TEST_DELTA], {"CPU": 2, "memory": 10}, 20)
+        ) = io.fit_input_deltas(
+            [TEST_DELTA], {"CPU": 2, "memory": 10}, self.COMPACTION_AUDIT, 20
+        )
 
         self.assertIsNotNone(hash_bucket_count)
         self.assertTrue(2, len(delta_list))
@@ -60,10 +80,12 @@ class TestFitInputDeltas(unittest.TestCase):
         from deltacat.compute.compactor.utils import io
 
         with self.assertRaises(AssertionError):
-            io.fit_input_deltas([], {"CPU": 100, "memory": 20000.0}, None)
+            io.fit_input_deltas(
+                [], {"CPU": 100, "memory": 20000.0}, self.COMPACTION_AUDIT, None
+            )
 
     def test_when_cpu_resources_is_not_passed(self):
         from deltacat.compute.compactor.utils import io
 
         with self.assertRaises(KeyError):
-            io.fit_input_deltas([], {}, None)
+            io.fit_input_deltas([], {}, self.COMPACTION_AUDIT, None)

deltacat/tests/io/test_file_object_store.py

@@ -0,0 +1,86 @@
+import unittest
+from unittest import mock
+
+
+class TestFileObjectStore(unittest.TestCase):
+
+    TEST_VALUE = "test-value"
+
+    @classmethod
+    def setUpClass(cls):
+        cls.ray_mock = mock.MagicMock()
+        cls.os_mock = mock.MagicMock()
+
+        cls.module_patcher = mock.patch.dict(
+            "sys.modules", {"ray": cls.ray_mock, "os": cls.os_mock}
+        )
+        cls.module_patcher.start()
+
+        super().setUpClass()
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        cls.module_patcher.stop()
+
+    @mock.patch(
+        "deltacat.io.file_object_store.open",
+        new_callable=mock.mock_open,
+        read_data="data",
+    )
+    def test_put_many_sanity(self, mock_file):
+        from deltacat.io.file_object_store import FileObjectStore
+
+        object_store = FileObjectStore(dir_path="")
+        self.ray_mock.cloudpickle.dumps.return_value = self.TEST_VALUE
+        result = object_store.put_many(["a", "b"])
+
+        self.assertEqual(2, len(result))
+        self.assertEqual(2, mock_file.call_count)
+
+    @mock.patch(
+        "deltacat.io.file_object_store.open",
+        new_callable=mock.mock_open,
+        read_data="data",
+    )
+    def test_get_many_sanity(self, mock_file):
+        from deltacat.io.file_object_store import FileObjectStore
+
+        object_store = FileObjectStore(dir_path="")
+        self.ray_mock.cloudpickle.loads.return_value = self.TEST_VALUE
+
+        result = object_store.get_many(["test", "test"])
+
+        self.assertEqual(2, len(result))
+        self.assertEqual(2, mock_file.call_count)
+
+    @mock.patch(
+        "deltacat.io.file_object_store.open",
+        new_callable=mock.mock_open,
+        read_data="data",
+    )
+    def test_get_sanity(self, mock_file):
+        from deltacat.io.file_object_store import FileObjectStore
+
+        object_store = FileObjectStore(dir_path="")
+        self.ray_mock.cloudpickle.loads.return_value = self.TEST_VALUE
+
+        result = object_store.get("test")
+
+        self.assertEqual(self.TEST_VALUE, result)
+        self.assertEqual(1, mock_file.call_count)
+
+    @mock.patch(
+        "deltacat.io.file_object_store.open",
+        new_callable=mock.mock_open,
+        read_data="data",
+    )
+    def test_put_sanity(self, mock_file):
+        from deltacat.io.file_object_store import FileObjectStore
+
+        object_store = FileObjectStore(dir_path="")
+        self.ray_mock.cloudpickle.dumps.return_value = self.TEST_VALUE
+
+        result = object_store.put("test")
+
+        self.assertIsNotNone(result)
+        self.assertEqual(1, mock_file.call_count)