deltacat 1.1.8__py3-none-any.whl → 1.1.10__py3-none-any.whl

deltacat/tests/local_deltacat_storage/__init__.py

@@ -12,6 +12,7 @@ import io
 from deltacat.tests.test_utils.storage import create_empty_delta
 from deltacat.utils.common import current_time_ms
 
+
 from deltacat.storage import (
     Delta,
     DeltaLocator,
@@ -40,6 +41,12 @@ from deltacat.storage import (
     ManifestEntry,
     ManifestEntryList,
     DeleteParameters,
+    PartitionFilter,
+    PartitionValues,
+    DeltaPartitionSpec,
+    StreamPartitionSpec,
+    TransformName,
+    IdentityTransformParameters,
 )
 from deltacat.types.media import (
     ContentType,
@@ -49,6 +56,10 @@ from deltacat.types.media import (
     DistributedDatasetType,
 )
 from deltacat.utils.common import ReadKwargsProvider
+from deltacat.tests.local_deltacat_storage.exceptions import (
+    InvalidNamespaceError,
+    LocalStorageValidationError,
+)
 
 SQLITE_CUR_ARG = "sqlite3_cur"
 SQLITE_CON_ARG = "sqlite3_con"
@@ -97,6 +108,19 @@ def _get_manifest_entry_uri(manifest_entry_id: str) -> str:
     return f"cloudpickle://{manifest_entry_id}"
 
 
+def _merge_and_promote(
+    partition_deltas: List[Delta], previous_partition_deltas: List[Delta]
+):
+    previous_partition_deltas_spos_gt: List[Delta] = [
+        delta
+        for delta in previous_partition_deltas
+        if delta.stream_position > partition_deltas[0].stream_position
+    ]
+    # handle the case if the previous partition deltas have a greater stream position than the partition_delta
+    partition_deltas = previous_partition_deltas_spos_gt + partition_deltas
+    return partition_deltas
+
+
 def list_namespaces(*args, **kwargs) -> ListResult[Namespace]:
     cur, con = _get_sqlite3_cursor_con(kwargs)
     res = cur.execute("SELECT * FROM namespaces")
@@ -176,12 +200,13 @@ def list_stream_partitions(stream: Stream, *args, **kwargs) -> ListResult[Partit
 def list_deltas(
     namespace: str,
     table_name: str,
-    partition_values: Optional[List[Any]] = None,
+    partition_values: Optional[PartitionValues] = None,
     table_version: Optional[str] = None,
     first_stream_position: Optional[int] = None,
     last_stream_position: Optional[int] = None,
     ascending_order: Optional[bool] = None,
     include_manifest: bool = False,
+    partition_filter: Optional[PartitionFilter] = None,
     *args,
     **kwargs,
 ) -> ListResult[Delta]:
@@ -189,6 +214,13 @@ def list_deltas(
     if stream is None:
         return ListResult.of([], None, None)
 
+    if partition_values is not None and partition_filter is not None:
+        raise ValueError(
+            "Only one of partition_values or partition_filter must be provided"
+        )
+    if partition_filter is not None:
+        partition_values = partition_filter.partition_values
+
     partition = get_partition(stream.locator, partition_values, *args, **kwargs)
 
     all_deltas = list_partition_deltas(
@@ -279,15 +311,25 @@ def get_delta(
     namespace: str,
     table_name: str,
     stream_position: int,
-    partition_values: Optional[List[Any]] = None,
+    partition_values: Optional[PartitionValues] = None,
     table_version: Optional[str] = None,
     include_manifest: bool = False,
+    partition_filter: Optional[PartitionFilter] = None,
     *args,
     **kwargs,
 ) -> Optional[Delta]:
     cur, con = _get_sqlite3_cursor_con(kwargs)
 
     stream = get_stream(namespace, table_name, table_version, *args, **kwargs)
+
+    if partition_values is not None and partition_filter is not None:
+        raise ValueError(
+            "Only one of partition_values or partition_filter must be provided"
+        )
+
+    if partition_filter is not None:
+        partition_values = partition_filter.partition_values
+
     partition = get_partition(stream.locator, partition_values, *args, **kwargs)
     delta_locator = DeltaLocator.of(partition.locator, stream_position)
 
@@ -310,22 +352,24 @@ def get_delta(
 def get_latest_delta(
     namespace: str,
     table_name: str,
-    partition_values: Optional[List[Any]] = None,
+    partition_values: Optional[PartitionValues] = None,
     table_version: Optional[str] = None,
     include_manifest: bool = False,
+    partition_filter: Optional[PartitionFilter] = None,
     *args,
     **kwargs,
 ) -> Optional[Delta]:
 
     deltas = list_deltas(
-        namespace,
-        table_name,
-        partition_values,
-        table_version,
-        None,
-        None,
-        False,
-        include_manifest,
+        namespace=namespace,
+        table_name=table_name,
+        partition_values=partition_values,
+        table_version=table_version,
+        first_stream_position=None,
+        last_stream_position=None,
+        ascending_order=False,
+        include_manifest=include_manifest,
+        partition_filter=partition_filter,
         *args,
         **kwargs,
     ).all_items()
@@ -345,13 +389,24 @@ def download_delta(
     file_reader_kwargs_provider: Optional[ReadKwargsProvider] = None,
     ray_options_provider: Callable[[int, Any], Dict[str, Any]] = None,
     distributed_dataset_type: DistributedDatasetType = DistributedDatasetType.RAY_DATASET,
+    partition_filter: Optional[PartitionFilter] = None,
     *args,
     **kwargs,
 ) -> Union[LocalDataset, DistributedDataset]:  # type: ignore
     result = []
     manifest = get_delta_manifest(delta_like, *args, **kwargs)
 
+    partition_values: PartitionValues = None
+    if partition_filter is not None:
+        partition_values = partition_filter.partition_values
+
     for entry_index in range(len(manifest.entries)):
+        if (
+            partition_values is not None
+            and partition_values != manifest.entries[entry_index].meta.partition_values
+        ):
+            continue
+
         result.append(
             download_delta_manifest_entry(
                 delta_like=delta_like,
@@ -506,11 +561,29 @@ def create_table_version(
     table_description: Optional[str] = None,
     table_properties: Optional[Dict[str, str]] = None,
     supported_content_types: Optional[List[ContentType]] = None,
+    partition_spec: Optional[StreamPartitionSpec] = None,
     *args,
     **kwargs,
 ) -> Stream:
     cur, con = _get_sqlite3_cursor_con(kwargs)
 
+    if partition_keys is not None and partition_spec is not None:
+        raise ValueError(
+            "Only one of partition_keys or partition_spec must be provided"
+        )
+    if partition_spec is not None:
+        assert (
+            partition_spec.ordered_transforms is not None
+        ), "Ordered transforms must be specified when partition_spec is specified"
+        partition_keys = []
+        for transform in partition_spec.ordered_transforms:
+            assert transform.name == TransformName.IDENTITY, (
+                "Local DeltaCAT storage does not support creating table versions "
+                "with non identity transform partition spec"
+            )
+            transform_params: IdentityTransformParameters = transform.parameters
+            partition_keys.append(transform_params.column_name)
+
     latest_version = get_latest_table_version(namespace, table_name, *args, **kwargs)
     if (
         table_version is not None
@@ -758,7 +831,7 @@ def delete_stream(
 
 
 def stage_partition(
-    stream: Stream, partition_values: Optional[List[Any]] = None, *args, **kwargs
+    stream: Stream, partition_values: Optional[PartitionValues] = None, *args, **kwargs
 ) -> Partition:
     cur, con = _get_sqlite3_cursor_con(kwargs)
     partition_id = uuid.uuid4().__str__()
@@ -820,19 +893,19 @@ def commit_partition(
         ).all_items()
         or []
     )
+
     partition_deltas: Optional[List[Delta]] = (
         list_partition_deltas(
             partition, ascending_order=False, *args, **kwargs
         ).all_items()
         or []
     )
-    previous_partition_deltas_spos_gt: List[Delta] = [
-        delta
-        for delta in previous_partition_deltas
-        if delta.stream_position > partition_deltas[0].stream_position
-    ]
-    # handle the case if the previous partition deltas have a greater stream position than the partition_delta
-    partition_deltas = previous_partition_deltas_spos_gt + partition_deltas
+
+    # if previous_partition is passed in, table is in-place compacted and we need to run merge-and-promote
+    if previous_partition:
+        partition_deltas = _merge_and_promote(
+            partition_deltas, previous_partition_deltas
+        )
 
     stream_position = (
         partition_deltas[0].stream_position
@@ -840,13 +913,14 @@ def commit_partition(
         else partition.stream_position
     )
 
-    partition.state = CommitState.COMMITTED
     partition.stream_position = stream_position
+    if partition_deltas:
+        partition.locator = partition_deltas[0].partition_locator
+
+    partition.state = CommitState.COMMITTED
     partition.previous_stream_position = (
         pv_partition.stream_position if pv_partition else None
     )
-    if partition_deltas:
-        partition.locator = partition_deltas[0].partition_locator
     params = (json.dumps(partition), partition.locator.canonical_string())
     cur.execute("UPDATE partitions SET value = ? WHERE locator = ?", params)
     con.commit()
@@ -858,7 +932,7 @@ def delete_partition(
     namespace: str,
     table_name: str,
     table_version: Optional[str] = None,
-    partition_values: Optional[List[Any]] = None,
+    partition_values: Optional[PartitionValues] = None,
     *args,
     **kwargs,
 ) -> None:
@@ -875,7 +949,7 @@ def delete_partition(
 
 def get_partition(
     stream_locator: StreamLocator,
-    partition_values: Optional[List[Any]] = None,
+    partition_values: Optional[PartitionValues] = None,
     *args,
     **kwargs,
 ) -> Optional[Partition]:
@@ -916,12 +990,14 @@ def stage_delta(
     s3_table_writer_kwargs: Optional[Dict[str, Any]] = None,
     content_type: ContentType = ContentType.PARQUET,
     delete_parameters: Optional[DeleteParameters] = None,
+    partition_spec: Optional[DeltaPartitionSpec] = None,
+    partition_values: Optional[PartitionValues] = None,
     *args,
     **kwargs,
 ) -> Delta:
     cur, con = _get_sqlite3_cursor_con(kwargs)
-    manifest_entry_id = uuid.uuid4().__str__()
-    uri = _get_manifest_entry_uri(manifest_entry_id)
+    manifest_id = uuid.uuid4().__str__()
+    uri = _get_manifest_entry_uri(manifest_id)
 
     if data is None:
         delta = create_empty_delta(
@@ -929,7 +1005,7 @@ def stage_delta(
             delta_type,
             author,
             properties=properties,
-            manifest_entry_id=manifest_entry_id,
+            manifest_entry_id=manifest_id,
         )
         cur.execute("INSERT OR IGNORE INTO data VALUES (?, ?)", (uri, None))
         params = (delta.locator.canonical_string(), "staged_delta", json.dumps(delta))
@@ -937,6 +1013,12 @@ def stage_delta(
         con.commit()
         return delta
 
+    if partition_spec:
+        assert partition_values is not None, (
+            "partition_values must be provided as local "
+            "storage does not support computing it from input data"
+        )
+
     serialized_data = None
     if content_type == ContentType.PARQUET:
         buffer = io.BytesIO()
@@ -961,18 +1043,19 @@ def stage_delta(
         content_type=content_type,
         content_encoding=ContentEncoding.IDENTITY,
         source_content_length=data.nbytes,
+        partition_values=partition_values,
     )
 
     manifest = Manifest.of(
         entries=ManifestEntryList.of(
             [
                 ManifestEntry.of(
-                    uri=uri, url=uri, meta=meta, mandatory=True, uuid=manifest_entry_id
+                    uri=uri, url=uri, meta=meta, mandatory=True, uuid=manifest_id
                 )
             ]
         ),
         author=author,
-        uuid=manifest_entry_id,
+        uuid=manifest_id,
     )
 
     delta = Delta.of(
@@ -1162,3 +1245,15 @@ def get_table_version_column_names(
     **kwargs,
 ) -> Optional[List[str]]:
     raise NotImplementedError("Fetching column names is not supported")
+
+
+def can_categorize(e: BaseException, **kwargs) -> bool:
+    if isinstance(e, InvalidNamespaceError):
+        return True
+    else:
+        return False
+
+
+def raise_categorized_error(e: BaseException, **kwargs):
+    if isinstance(e, InvalidNamespaceError):
+        raise LocalStorageValidationError("Namespace provided is invalid!")

deltacat/tests/local_deltacat_storage/exceptions.py

@@ -0,0 +1,10 @@
+class InvalidNamespaceError(Exception):
+    error_name = "InvalidNamespaceError"
+
+
+class LocalStorageValidationError(Exception):
+    error_name = "LocalStorageValidationError"
+
+
+class LocalStorageError(Exception):
+    error_name = "LocalStorageError"

deltacat/tests/test_exceptions.py

@@ -0,0 +1,100 @@
+import unittest
+from deltacat.exceptions import categorize_errors
+import ray
+from deltacat.exceptions import (
+    DependencyPyarrowCapacityError,
+    NonRetryableDownloadTableError,
+    RetryableError,
+    NonRetryableError,
+    DeltaCatTransientError,
+    DependencyDaftTransientError,
+    UnclassifiedDeltaCatError,
+)
+from daft.exceptions import DaftTransientError
+from deltacat.tests.local_deltacat_storage.exceptions import (
+    InvalidNamespaceError,
+    LocalStorageValidationError,
+)
+from botocore.exceptions import NoCredentialsError
+from tenacity import retry, retry_if_exception_type, stop_after_attempt
+
+from pyarrow.lib import ArrowCapacityError
+import deltacat.tests.local_deltacat_storage as ds
+
+
+class MockUnknownException(Exception):
+    pass
+
+
+@categorize_errors
+def mock_raise_exception(exception_to_raise, deltacat_storage=ds):
+    raise exception_to_raise
+
+
+@retry(retry=retry_if_exception_type(NoCredentialsError), stop=stop_after_attempt(2))
+def mock_tenacity_wrapped_method(exception_to_raise):
+    mock_raise_exception(exception_to_raise)
+
+
+@ray.remote
+def mock_remote_task(exception_to_raise):
+    mock_raise_exception(exception_to_raise)
+
+
+class TestCategorizeErrors(unittest.TestCase):
+    def test_pyarrow_exception_categorizer(self):
+        self.assertRaises(
+            DependencyPyarrowCapacityError,
+            lambda: mock_raise_exception(ArrowCapacityError),
+        )
+
+    def test_storage_exception_categorizer(self):
+        self.assertRaises(
+            LocalStorageValidationError,
+            lambda: mock_raise_exception(InvalidNamespaceError, deltacat_storage=ds),
+        )
+
+    def test_non_retryable_error(self):
+        self.assertRaises(
+            NonRetryableError,
+            lambda: mock_raise_exception(NonRetryableDownloadTableError),
+        )
+
+    def test_retryable_error(self):
+        self.assertRaises(RetryableError, lambda: mock_raise_exception(ConnectionError))
+
+    def test_ray_task_returns_wrapped_exception(self):
+        self.assertRaises(
+            DeltaCatTransientError,
+            lambda: ray.get(mock_remote_task.remote(ConnectionError)),
+        )
+
+    def test_daft_transient_error(self):
+        self.assertRaises(
+            DependencyDaftTransientError,
+            lambda: ray.get(mock_remote_task.remote(DaftTransientError)),
+        )
+
+    def test_tenacity_underlying_error_returned(self):
+        self.assertRaises(
+            DeltaCatTransientError,
+            lambda: mock_tenacity_wrapped_method(NoCredentialsError),
+        )
+
+    def test_unclassified_error_when_error_cannot_be_categorized(self):
+        self.assertRaises(
+            UnclassifiedDeltaCatError,
+            lambda: ray.get(mock_remote_task.remote(MockUnknownException)),
+        )
+
+    def test_deltacat_exception_contains_attributes(self):
+
+        try:
+            mock_raise_exception(ConnectionError)
+        except DeltaCatTransientError as e:
+            self.assertTrue(hasattr(e, "is_retryable"))
+            self.assertTrue(hasattr(e, "error_name"))
+            assert e.error_name == "DeltaCatTransientError"
+            return
+
+        self.assertFalse(True)

deltacat/tests/test_logs.py

@@ -38,6 +38,7 @@ class TestJsonFormatter(unittest.TestCase):
         self.assertEqual({"message": "test_message"}, result)
 
     def test_format_sanity(self):
+        ray.shutdown()
         formatter = JsonFormatter({"message": "msg"})
 
         record = LogRecord(

deltacat/tests/test_utils/pyarrow.py

@@ -66,7 +66,10 @@ def download_delta(delta_like: Union[Delta, DeltaLocator], *args, **kwargs) -> D
 
 
 def commit_delta_to_partition(
-    partition: Partition, file_paths: List[str], *args, **kwargs
+    partition: Union[Partition, PartitionLocator],
+    file_paths: List[str],
+    *args,
+    **kwargs,
 ) -> Delta:
     tables = []
 

deltacat/tests/utils/ray_utils/test_dataset.py

@@ -0,0 +1,66 @@
+from ray.data import from_items
+from typing import Any
+import pytest
+import fsspec
+from fsspec import AbstractFileSystem
+from ray.data.datasource import FilenameProvider
+from deltacat.types.media import ContentType
+import ray
+
+
+class TestDatasetToFile:
+
+    BASE_PATH = "/tmp"
+    SUB_PATH = "abcd"
+
+    @pytest.fixture(autouse=True, scope="module")
+    def ensure_ray_down(self):
+        # ray.data fails when ray is instantiated in local mode
+        ray.shutdown()
+
+    @pytest.fixture(scope="module")
+    def mock_dataset(self):
+        return from_items([{"col1": i, "col2": i * 2} for i in range(1000)])
+
+    @pytest.fixture(scope="module")
+    def mock_filename_provider(self):
+        class MockFilenameProvider(FilenameProvider):
+            def get_filename_for_block(
+                self, block: Any, task_index: int, block_index: int
+            ) -> str:
+                return TestDatasetToFile.SUB_PATH
+
+        return MockFilenameProvider()
+
+    def test_parquet_sanity(self, mock_dataset, mock_filename_provider):
+        from deltacat.utils.ray_utils.dataset import dataset_to_file
+
+        fs: AbstractFileSystem = fsspec.filesystem("local")
+
+        dataset_to_file(
+            mock_dataset,
+            self.BASE_PATH,
+            file_system=fs,
+            block_path_provider=mock_filename_provider,
+        )
+
+        file_expected_at = f"{self.BASE_PATH}/{self.SUB_PATH}"
+        assert fs.exists(file_expected_at), "file was not written"
+        fs.delete(file_expected_at)
+
+    def test_csv_sanity(self, mock_dataset, mock_filename_provider):
+        from deltacat.utils.ray_utils.dataset import dataset_to_file
+
+        fs: AbstractFileSystem = fsspec.filesystem("local")
+
+        dataset_to_file(
+            mock_dataset,
+            self.BASE_PATH,
+            file_system=fs,
+            block_path_provider=mock_filename_provider,
+            content_type=ContentType.CSV.value,
+        )
+
+        file_expected_at = f"{self.BASE_PATH}/{self.SUB_PATH}"
+        assert fs.exists(file_expected_at), "file was not written"
+        fs.delete(file_expected_at)

deltacat/tests/utils/test_daft.py

@@ -1,7 +1,6 @@
 import unittest
 from deltacat.types.media import ContentEncoding, ContentType
 from deltacat.utils.daft import daft_s3_file_to_table, s3_files_to_dataframe
-
 from deltacat.utils.pyarrow import ReadKwargsProviderPyArrowSchemaOverride
 from deltacat.types.partial_download import PartialParquetParameters
 import pyarrow as pa

deltacat/tests/utils/test_resources.py

@@ -1,8 +1,6 @@
 import unittest
 from unittest import mock
 import time
-from multiprocessing import Pool
-import platform
 
 
 class TestGetCurrentClusterUtilization(unittest.TestCase):
@@ -72,29 +70,3 @@ class TestProcessUtilizationOverTimeRange(unittest.TestCase):
             nu.schedule_callback(test_callback, 1)
             time.sleep(3)
             self.assertTrue(nu.test_field_set)
-
-
-class TestTimeoutDecorator(unittest.TestCase):
-    from deltacat.utils.resources import timeout
-
-    @staticmethod
-    @timeout(2)
-    def something_that_runs_xs(x, *args, **kwargs):
-        time.sleep(x)
-
-    def test_timeout(self):
-        if platform.system() != "Windows":
-            self.assertRaises(
-                TimeoutError, lambda: self.something_that_runs_xs(3, test=10)
-            )
-
-    def test_sanity_in_multiprocess(self):
-        if platform.system() != "Windows":
-            # An alarm works per process
-            # https://pubs.opengroup.org/onlinepubs/9699919799/functions/alarm.html
-            with Pool(3) as p:
-                p.map(self.something_that_runs_xs, [1, 1.1, 1.2])
-
-    def test_sanity(self):
-        if platform.system() != "Windows":
-            self.something_that_runs_xs(1, test=10)

deltacat/utils/daft.py

@@ -16,6 +16,7 @@ from deltacat.aws.constants import (
     BOTO_MAX_RETRIES,
     DAFT_MAX_S3_CONNECTIONS_PER_FILE,
     AWS_REGION,
+    DEFAULT_FILE_READ_TIMEOUT_MS,
 )
 from deltacat.utils.performance import timed_invocation
 
@@ -112,6 +113,7 @@ def daft_s3_file_to_table(
     coerce_int96_timestamp_unit = TimeUnit.from_str(
         kwargs.get("coerce_int96_timestamp_unit", "ms")
     )
+    file_timeout_ms = kwargs.get("file_timeout_ms", DEFAULT_FILE_READ_TIMEOUT_MS)
 
     row_groups = None
     if (
@@ -132,6 +134,7 @@ def daft_s3_file_to_table(
         io_config=io_config,
         coerce_int96_timestamp_unit=coerce_int96_timestamp_unit,
         multithreaded_io=False,
+        file_timeout_ms=file_timeout_ms,
     )
 
     logger.debug(f"Time to read S3 object from {s3_url} into daft table: {latency}s")

deltacat/utils/numpy.py

@@ -1,10 +1,10 @@
-from typing import List, Optional
+from typing import List, Optional, Callable, Union
 
 import numpy as np
 import pyarrow as pa
 from fsspec import AbstractFileSystem
-from ray.data.datasource import BlockWritePathProvider
 
+from ray.data.datasource import FilenameProvider
 from deltacat.types.media import ContentType
 from deltacat.utils import pandas as pd_utils
 from deltacat.utils import pyarrow as pa_utils
@@ -52,7 +52,7 @@ def ndarray_to_file(
     np_array: np.ndarray,
     path: str,
     file_system: AbstractFileSystem,
-    block_path_provider: BlockWritePathProvider,
+    block_path_provider: Union[FilenameProvider, Callable],
     content_type: str = ContentType.PARQUET.value,
     **kwargs
 ) -> None:

deltacat/utils/pandas.py

@@ -2,12 +2,12 @@ import csv
 import io
 import logging
 import math
-from typing import Any, Callable, Dict, Iterable, List, Optional
+from typing import Any, Callable, Dict, Iterable, List, Optional, Union
 
 import pandas as pd
 import pyarrow as pa
 from fsspec import AbstractFileSystem
-from ray.data.datasource import BlockWritePathProvider
+from ray.data.datasource import FilenameProvider
 
 from deltacat import logs
 from deltacat.types.media import (
@@ -262,7 +262,7 @@ def dataframe_to_file(
     dataframe: pd.DataFrame,
     base_path: str,
     file_system: AbstractFileSystem,
-    block_path_provider: BlockWritePathProvider,
+    block_path_provider: Union[Callable, FilenameProvider],
     content_type: str = ContentType.PARQUET.value,
     **kwargs,
 ) -> None: