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

deltacat/io/memcached_object_store.py

@@ -3,13 +3,14 @@ from ray import cloudpickle
 from collections import defaultdict
 import time
 from deltacat.io.object_store import IObjectStore
-from typing import Any, List
+from typing import Any, List, Optional
 from deltacat import logs
 import uuid
 import socket
 from pymemcache.client.base import Client
 from pymemcache.client.retrying import RetryingClient
 from pymemcache.exceptions import MemcacheUnexpectedCloseError
+from pymemcache.client.rendezvous import RendezvousHash
 
 logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
 
@@ -19,36 +20,46 @@ class MemcachedObjectStore(IObjectStore):
     An implementation of object store that uses Memcached.
     """
 
-    def __init__(self, port=11212) -> None:
+    def __init__(
+        self, storage_node_ips: Optional[List[str]] = None, port: Optional[int] = 11212
+    ) -> None:
         self.client_cache = {}
         self.current_ip = None
         self.SEPARATOR = "_"
         self.port = port
+        self.storage_node_ips = storage_node_ips
+        self.hasher = None
         super().__init__()
 
+    def initialize_hasher(self):
+        if not self.hasher and self.storage_node_ips:
+            self.hasher = RendezvousHash()
+            for n in self.storage_node_ips:
+                self.hasher.add_node(n)
+
     def put_many(self, objects: List[object], *args, **kwargs) -> List[Any]:
-        input = {}
+        input = defaultdict(dict)
         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
+            create_ref_ip = self._get_create_ref_ip(uid.__str__())
+            ref = self._create_ref(uid, create_ref_ip)
+            input[create_ref_ip][uid.__str__()] = serialized
             result.append(ref)
-
-        client = self._get_client_by_ip(current_ip)
-        if client.set_many(input, noreply=False):
-            raise RuntimeError("Unable to write few keys to cache")
+        for create_ref_ip, uid_to_object in input.items():
+            client = self._get_client_by_ip(create_ref_ip)
+            if client.set_many(uid_to_object, noreply=False):
+                raise RuntimeError("Unable to write few keys to cache")
 
         return result
 
     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)
+        create_ref_ip = self._get_create_ref_ip(uid.__str__())
+        ref = self._create_ref(uid, create_ref_ip)
+        client = self._get_client_by_ip(create_ref_ip)
 
         if client.set(uid.__str__(), serialized):
             return ref
@@ -99,6 +110,18 @@ class MemcachedObjectStore(IObjectStore):
     def _create_ref(self, uid, ip) -> str:
         return f"{uid}{self.SEPARATOR}{ip}"
 
+    def _get_storage_node_ip(self, key: str):
+        self.initialize_hasher()
+        storage_node_ip = self.hasher.get_node(key)
+        return storage_node_ip
+
+    def _get_create_ref_ip(self, uid: str):
+        if self.storage_node_ips:
+            create_ref_ip = self._get_storage_node_ip(uid)
+        else:
+            create_ref_ip = self._get_current_ip()
+        return create_ref_ip
+
     def _get_client_by_ip(self, ip_address: str):
         if ip_address in self.client_cache:
             return self.client_cache[ip_address]
@@ -108,7 +131,7 @@ class MemcachedObjectStore(IObjectStore):
             base_client,
             attempts=3,
             retry_delay=0.01,
-            retry_for=[MemcacheUnexpectedCloseError],
+            retry_for=[MemcacheUnexpectedCloseError, ConnectionResetError],
         )
 
         self.client_cache[ip_address] = client

deltacat/logs.py

@@ -143,9 +143,10 @@ def _configure_logger(
             log_dir, log_base_file_name, primary_log_level
         )
         _add_logger_handler(logger, handler)
-    ray_runtime_ctx = ray.get_runtime_context()
-    if ray_runtime_ctx.worker.connected:
-        logger = RayRuntimeContextLoggerAdapter(logger, ray_runtime_ctx)
+    if ray.is_initialized():
+        ray_runtime_ctx = ray.get_runtime_context()
+        if ray_runtime_ctx.worker.connected:
+            logger = RayRuntimeContextLoggerAdapter(logger, ray_runtime_ctx)
 
     return logger
 

deltacat/storage/interface.py

@@ -21,6 +21,7 @@ from deltacat.storage import (
     Table,
     TableVersion,
     SortKey,
+    PartitionLocator,
 )
 from deltacat.types.media import ContentType, StorageType, TableType
 from deltacat.utils.common import ReadKwargsProvider
@@ -105,7 +106,13 @@ def list_deltas(
 
 
 def list_partition_deltas(
-    partition: Partition, include_manifest: bool = False, *args, **kwargs
+    partition_like: Union[Partition, PartitionLocator],
+    first_stream_position: Optional[int] = None,
+    last_stream_position: Optional[int] = None,
+    ascending_order: bool = False,
+    include_manifest: bool = False,
+    *args,
+    **kwargs
 ) -> ListResult[Delta]:
     """
     Lists a page of deltas committed to the given partition.

deltacat/storage/model/types.py

@@ -1,6 +1,7 @@
 from enum import Enum
 from typing import List, Union, Any
 
+from pyarrow.parquet import ParquetFile
 import numpy as np
 import pandas as pd
 import pyarrow as pa
@@ -8,7 +9,7 @@ import pkg_resources
 from ray.data._internal.arrow_block import ArrowRow
 from ray.data.dataset import Dataset
 
-LocalTable = Union[pa.Table, pd.DataFrame, np.ndarray]
+LocalTable = Union[pa.Table, pd.DataFrame, np.ndarray, ParquetFile]
 LocalDataset = List[LocalTable]
 # Starting Ray 2.5.0, Dataset follows a strict mode (https://docs.ray.io/en/latest/data/faq.html#migrating-to-strict-mode),
 # and generic annotation is removed. So add a version checker to determine whether to use the old or new definition.

deltacat/tests/aws/test_clients.py

@@ -1,7 +1,7 @@
 from unittest.mock import patch
 import unittest
-import json
 from http import HTTPStatus
+import requests
 
 HAPPY_RESPONSE = {
     "AccessKeyId": "ASIA123456789",
@@ -20,7 +20,7 @@ class MockResponse:
     """
 
     def __init__(self, status_code: int, text: str, reason: str = "") -> None:
-        self.status_code = status_code
+        self.status_code: requests.Response.status_code = status_code
         self.text = text
         self.reason = reason
 
@@ -55,7 +55,7 @@ class TestBlockUntilInstanceMetadataServiceReturnsSuccess(unittest.TestCase):
         )
 
         requests_mock.get.side_effect = [
-            MockResponse(HTTPStatus.OK, json.dumps(HAPPY_RESPONSE)),
+            MockResponse(HTTPStatus.OK, "foo"),
             MockResponse(HTTPStatus.TOO_MANY_REQUESTS, "foo"),
             MockResponse(HTTPStatus.INTERNAL_SERVER_ERROR, "foo"),
             MockResponse(HTTPStatus.NOT_IMPLEMENTED, "bar"),
@@ -65,3 +65,16 @@ class TestBlockUntilInstanceMetadataServiceReturnsSuccess(unittest.TestCase):
         self.assertEqual(
             block_until_instance_metadata_service_returns_success().status_code, 200
         )
+
+    @patch("deltacat.aws.clients.requests")
+    def test_retrying_status_on_shortlist_returns_early(self, requests_mock):
+        from deltacat.aws.clients import (
+            block_until_instance_metadata_service_returns_success,
+        )
+
+        requests_mock.get.side_effect = [
+            MockResponse(HTTPStatus.FORBIDDEN, "foo"),
+        ]
+        self.assertEqual(
+            block_until_instance_metadata_service_returns_success().status_code, 403
+        )

deltacat/tests/compute/common.py

@@ -0,0 +1,96 @@
+# Allow classes to use self-referencing Type hints in Python 3.7.
+from __future__ import annotations
+from enum import Enum
+from typing import List
+import datetime as dt
+from datetime import timezone
+
+TEST_S3_RCF_BUCKET_NAME = "test-compaction-artifacts-bucket"
+# REBASE  src = spark compacted table to create an initial version of ray compacted table
+BASE_TEST_SOURCE_NAMESPACE = "source_test_namespace"
+BASE_TEST_SOURCE_TABLE_NAME = "test_table"
+BASE_TEST_SOURCE_TABLE_VERSION = "1"
+
+BASE_TEST_DESTINATION_NAMESPACE = "destination_test_namespace"
+BASE_TEST_DESTINATION_TABLE_NAME = "destination_test_table_RAY"
+BASE_TEST_DESTINATION_TABLE_VERSION = "1"
+
+HASH_BUCKET_COUNT: int = 1
+
+MAX_RECORDS_PER_FILE: int = 1
+
+UTC_ISO_8601_FORMAT_WITHOUT_MILLIS = "%Y-%m-%dT%H:%M:%SZ"  # '2018-09-05T14:09:03Z'
+
+
+class PartitionKeyType(str, Enum):
+    INT = "int"
+    STRING = "string"
+    TIMESTAMP = "timestamp"
+
+
+class PartitionKey(dict):
+    @staticmethod
+    def of(key_name: str, key_type: PartitionKeyType) -> PartitionKey:
+        return PartitionKey({"keyName": key_name, "keyType": key_type.value})
+
+    @property
+    def key_name(self) -> str:
+        return self["keyName"]
+
+    @property
+    def key_type(self) -> PartitionKeyType:
+        key_type = self["keyType"]
+        return None if key_type is None else PartitionKeyType(key_type)
+
+
+def setup_sort_and_partition_keys(sort_keys_param, partition_keys_param):
+    from deltacat.storage.model.sort_key import SortKey
+
+    sort_keys, partition_keys = None, None
+    if sort_keys_param is not None:
+        sort_keys = [SortKey.of(sort_key["key_name"]) for sort_key in sort_keys_param]
+    if partition_keys_param is not None:
+        partition_keys = [
+            PartitionKey.of(
+                partition_key["key_name"], PartitionKeyType(partition_key["key_type"])
+            )
+            for partition_key in partition_keys_param
+        ]
+    return sort_keys, partition_keys
+
+
+def offer_iso8601_timestamp_list(
+    periods: int,
+    unit_of_time: str,
+    end_time=dt.datetime(2023, 5, 3, 10, 0, 0, 0, tzinfo=timezone.utc),
+) -> List[str]:
+    """
+    Returns a list of ISO 8601 timestamps, each periods units of time before the start time.
+
+    Args:
+    periods: The number of timestamps to return.
+    unit_of_time: The unit of time to use for the timestamps. Must be one of "seconds", "minutes", "hours", "days", or "weeks".
+    end_time: The end time for the timestamps. Defaults to 2023-05-03T10:00:00Z.
+
+    Returns:
+    A list of ISO 8601 timestamps, each periods units of time before the start time.
+
+    Raises:
+    ValueError: If the unit_of_time argument is not one of "seconds", "minutes", "hours", "days", or "weeks".
+    """
+    import datetime as dt
+
+    acceptable_units_of_time = ["seconds", "minutes", "hours", "days", "weeks"]
+    if unit_of_time not in acceptable_units_of_time:
+        raise ValueError(
+            f"unit_of_time {unit_of_time} is not supported. Please use one of these time units: {acceptable_units_of_time}"
+        )
+    res = []
+    for i in range(periods):
+        kwarg = {unit_of_time: i}
+        res.append(
+            (end_time - dt.timedelta(**kwarg)).strftime(
+                UTC_ISO_8601_FORMAT_WITHOUT_MILLIS
+            )
+        )
+    return res

deltacat/tests/{test_repartition.py → compute/compactor/steps/test_repartition.py}

@@ -51,6 +51,7 @@ class TestRepartitionRange(unittest.TestCase):
         self.max_records_per_output_file = 2
         self.repartitioned_file_content_type = ContentType.PARQUET
         self.deltacat_storage = MagicMock()
+        self.deltacat_storage_kwargs = MagicMock()
 
     def test_repartition_range(self):
         result = repartition_range(
@@ -60,6 +61,7 @@ class TestRepartitionRange(unittest.TestCase):
             self.max_records_per_output_file,
             self.repartitioned_file_content_type,
             self.deltacat_storage,
+            self.deltacat_storage_kwargs,
         )
         # Assert that a RepartitionResult object is returned
         self.assertIsInstance(result, RepartitionResult)
@@ -85,6 +87,7 @@ class TestRepartitionRange(unittest.TestCase):
                 self.max_records_per_output_file,
                 self.repartitioned_file_content_type,
                 self.deltacat_storage,
+                self.deltacat_storage_kwargs,
             )
 
     def test_empty_ranges(self):
@@ -97,6 +100,7 @@ class TestRepartitionRange(unittest.TestCase):
                 self.max_records_per_output_file,
                 self.repartitioned_file_content_type,
                 self.deltacat_storage,
+                self.deltacat_storage_kwargs,
             )
 
     def test_one_value_in_ranges(self):
@@ -108,6 +112,7 @@ class TestRepartitionRange(unittest.TestCase):
             self.max_records_per_output_file,
             self.repartitioned_file_content_type,
             self.deltacat_storage,
+            self.deltacat_storage_kwargs,
         )
         self.assertEqual(len(result.range_deltas), 2)
 
@@ -120,6 +125,7 @@ class TestRepartitionRange(unittest.TestCase):
             self.max_records_per_output_file,
             self.repartitioned_file_content_type,
             self.deltacat_storage,
+            self.deltacat_storage_kwargs,
         )
         self.assertEqual(len(result.range_deltas), 3)
 
@@ -133,6 +139,7 @@ class TestRepartitionRange(unittest.TestCase):
             self.max_records_per_output_file,
             self.repartitioned_file_content_type,
             self.deltacat_storage,
+            self.deltacat_storage_kwargs,
         )
         self.assertLess(len(result.range_deltas), 2)
 
@@ -146,6 +153,7 @@ class TestRepartitionRange(unittest.TestCase):
                 self.max_records_per_output_file,
                 self.repartitioned_file_content_type,
                 self.deltacat_storage,
+                self.deltacat_storage_kwargs,
             )
 
     def test_unsorted_ranges(self):
@@ -161,6 +169,7 @@ class TestRepartitionRange(unittest.TestCase):
             self.max_records_per_output_file,
             self.repartitioned_file_content_type,
             self.deltacat_storage,
+            self.deltacat_storage_kwargs,
         )
         self.assertEqual(len(result.range_deltas), 4)
 
@@ -173,20 +182,24 @@ class TestRepartitionRange(unittest.TestCase):
             self.max_records_per_output_file,
             self.repartitioned_file_content_type,
             self.deltacat_storage,
+            self.deltacat_storage_kwargs,
         )
         self.assertEqual(len(result.range_deltas), 2)
 
     def test_ranges_with_inf(self):
         self.repartition_args["ranges"] = [1678665487112747, float("inf")]
-        result = repartition_range(
-            self.tables,
-            self.destination_partition,
-            self.repartition_args,
-            self.max_records_per_output_file,
-            self.repartitioned_file_content_type,
-            self.deltacat_storage,
+
+        self.assertRaises(
+            pa.lib.ArrowInvalid,
+            lambda: repartition_range(
+                self.tables,
+                self.destination_partition,
+                self.repartition_args,
+                self.max_records_per_output_file,
+                self.repartitioned_file_content_type,
+                self.deltacat_storage,
+            ),
         )
-        self.assertEqual(len(result.range_deltas), 2)
 
     def test_null_rows_are_not_dropped(self):
         # Add null value to the first table
@@ -211,6 +224,7 @@ class TestRepartitionRange(unittest.TestCase):
             self.max_records_per_output_file,
             self.repartitioned_file_content_type,
             self.deltacat_storage,
+            self.deltacat_storage_kwargs,
         )
 
         # Assuming range_deltas is a list of DataFrames,

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,
+            )