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

deltacat/compute/compactor_v2/model/merge_input.py

@@ -47,6 +47,7 @@ class MergeInput(Dict):
         deltacat_storage=unimplemented_deltacat_storage,
         deltacat_storage_kwargs: Optional[Dict[str, Any]] = None,
         memory_logs_enabled: Optional[bool] = None,
+        disable_copy_by_reference: Optional[bool] = None,
     ) -> MergeInput:
 
         result = MergeInput()
@@ -69,6 +70,7 @@ class MergeInput(Dict):
         result["deltacat_storage"] = deltacat_storage
         result["deltacat_storage_kwargs"] = deltacat_storage_kwargs or {}
         result["memory_logs_enabled"] = memory_logs_enabled
+        result["disable_copy_by_reference"] = disable_copy_by_reference
         return result
 
     @property
@@ -148,3 +150,7 @@ class MergeInput(Dict):
     @property
     def delete_strategy(self) -> Optional[DeleteStrategy]:
         return self.get("delete_strategy")
+
+    @property
+    def disable_copy_by_reference(self) -> bool:
+        return self["disable_copy_by_reference"]

deltacat/compute/compactor_v2/steps/hash_bucket.py

@@ -29,14 +29,15 @@ from deltacat.utils.metrics import emit_timer_metrics, failure_metric, success_m
 from deltacat.utils.resources import (
     get_current_process_peak_memory_usage_in_bytes,
     ProcessUtilizationOverTimeRange,
-    timeout,
 )
 from deltacat.constants import BYTES_PER_GIBIBYTE
 from deltacat.compute.compactor_v2.constants import (
     HASH_BUCKET_TIME_IN_SECONDS,
     HASH_BUCKET_FAILURE_COUNT,
     HASH_BUCKET_SUCCESS_COUNT,
-    HASH_BUCKET_TASK_TIMEOUT_IN_SECONDS,
+)
+from deltacat.exceptions import (
+    categorize_errors,
 )
 
 if importlib.util.find_spec("memray"):
@@ -79,7 +80,7 @@ def _group_file_records_by_pk_hash_bucket(
         logger.info("Grouping by pk hash bucket")
         group_start = time.monotonic()
         hash_bucket_to_table = group_by_pk_hash_bucket(
-            dfe.table, num_hash_buckets, primary_keys
+            table=dfe.table, num_buckets=num_hash_buckets, primary_keys=primary_keys
         )
         group_end = time.monotonic()
         logger.info(f"Grouping took: {group_end - group_start}")
@@ -98,12 +99,9 @@ def _group_file_records_by_pk_hash_bucket(
     return hb_to_delta_file_envelopes, total_record_count, total_size_bytes
 
 
-# TODO: use timeout parameter in ray.remote
-# https://github.com/ray-project/ray/issues/18916
-# Note: order of decorators is important
 @success_metric(name=HASH_BUCKET_SUCCESS_COUNT)
 @failure_metric(name=HASH_BUCKET_FAILURE_COUNT)
-@timeout(HASH_BUCKET_TASK_TIMEOUT_IN_SECONDS)
+@categorize_errors
 def _timed_hash_bucket(input: HashBucketInput):
     task_id = get_current_ray_task_id()
     worker_id = get_current_ray_worker_id()

deltacat/compute/compactor_v2/steps/merge.py

@@ -28,7 +28,6 @@ from deltacat.utils.metrics import emit_timer_metrics, failure_metric, success_m
 from deltacat.utils.resources import (
     get_current_process_peak_memory_usage_in_bytes,
     ProcessUtilizationOverTimeRange,
-    timeout,
 )
 from deltacat.compute.compactor_v2.utils.primary_key_index import (
     generate_pk_hash_column,
@@ -47,9 +46,10 @@ from deltacat.compute.compactor_v2.constants import (
     MERGE_TIME_IN_SECONDS,
     MERGE_SUCCESS_COUNT,
     MERGE_FAILURE_COUNT,
-    MERGE_TASK_TIMEOUT_IN_SECONDS,
 )
-
+from deltacat.exceptions import (
+    categorize_errors,
+)
 
 if importlib.util.find_spec("memray"):
     import memray
@@ -284,16 +284,19 @@ def _can_copy_by_reference(
     Can copy by reference only if there are no deletes to merge in
     and previous compacted stream id matches that of new stream
     """
-    return (
+    copy_by_ref = (
         not has_delete
         and not merge_file_group.dfe_groups
         and input.round_completion_info is not None
-        and (
-            input.write_to_partition.stream_id
-            == input.round_completion_info.compacted_delta_locator.stream_id
-        )
     )
 
+    if input.disable_copy_by_reference:
+        copy_by_ref = False
+
+    logger.info(f"Copy by reference is {copy_by_ref} for {merge_file_group.hb_index}")
+
+    return copy_by_ref
+
 
 def _flatten_dfe_list(
     df_envelopes_list: List[List[DeltaFileEnvelope]],
@@ -486,12 +489,9 @@ def _copy_manifests_from_hash_bucketing(
     return materialized_results
 
 
-# TODO: use timeout parameter in ray.remote
-# https://github.com/ray-project/ray/issues/18916
-# Note: order of decorators is important
 @success_metric(name=MERGE_SUCCESS_COUNT)
 @failure_metric(name=MERGE_FAILURE_COUNT)
-@timeout(MERGE_TASK_TIMEOUT_IN_SECONDS)
+@categorize_errors
 def _timed_merge(input: MergeInput) -> MergeResult:
     task_id = get_current_ray_task_id()
     worker_id = get_current_ray_worker_id()

deltacat/compute/compactor_v2/utils/merge.py

@@ -133,4 +133,5 @@ def generate_local_merge_input(
         deltacat_storage_kwargs=params.deltacat_storage_kwargs,
         delete_strategy=delete_strategy,
         delete_file_envelopes=delete_file_envelopes,
+        disable_copy_by_reference=params.disable_copy_by_reference,
     )

deltacat/compute/compactor_v2/utils/primary_key_index.py

@@ -157,7 +157,12 @@ def _optimized_group_record_batches_by_hash_bucket(
 def group_by_pk_hash_bucket(
     table: pa.Table, num_buckets: int, primary_keys: List[str]
 ) -> np.ndarray:
-    table = generate_pk_hash_column([table], primary_keys, requires_sha1=True)[0]
+    new_tables = generate_pk_hash_column([table], primary_keys, requires_hash=True)
+    assert (
+        len(new_tables) == 1
+    ), f"Expected only 1 table in the result but found {len(new_tables)}"
+
+    table = generate_pk_hash_column([table], primary_keys, requires_hash=True)[0]
 
     # group hash bucket record indices
     result = group_record_indices_by_hash_bucket(
@@ -171,7 +176,7 @@ def group_by_pk_hash_bucket(
 def generate_pk_hash_column(
     tables: List[pa.Table],
     primary_keys: Optional[List[str]] = None,
-    requires_sha1: bool = False,
+    requires_hash: bool = False,
 ) -> List[pa.Table]:
     """
     Returns a new table list after generating the primary key hash if desired.
@@ -203,12 +208,12 @@ def generate_pk_hash_column(
     if primary_keys:
         hash_column_list = [_generate_pk_hash(table) for table in tables]
 
-        can_sha1 = requires_sha1 or _is_sha1_desired(hash_column_list)
+        can_sha1 = requires_hash or _is_sha1_desired(hash_column_list)
     else:
         hash_column_list = [_generate_uuid(table) for table in tables]
 
     logger.info(
-        f"can_generate_sha1={can_sha1} for the table and requires_sha1={requires_sha1}"
+        f"can_generate_sha1={can_sha1} for the table and requires_sha1={requires_hash}"
     )
 
     result = []

deltacat/compute/compactor_v2/utils/task_options.py

@@ -1,6 +1,4 @@
-import botocore
 import logging
-import tenacity
 from typing import Dict, Optional, List, Tuple, Any
 from deltacat import logs
 from deltacat.compute.compactor_v2.model.merge_file_group import (
@@ -21,8 +19,7 @@ from deltacat.compute.compactor_v2.utils.primary_key_index import (
 from deltacat.compute.compactor_v2.constants import (
     PARQUET_TO_PYARROW_INFLATION,
 )
-from daft.exceptions import DaftTransientError
-
+from deltacat.exceptions import RetryableError
 
 logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
 
@@ -79,14 +76,7 @@ def get_task_options(
 
     # List of possible botocore exceptions are available at
     # https://github.com/boto/botocore/blob/develop/botocore/exceptions.py
-    task_opts["retry_exceptions"] = [
-        botocore.exceptions.ConnectionError,
-        botocore.exceptions.HTTPClientError,
-        ConnectionError,
-        TimeoutError,
-        DaftTransientError,
-        tenacity.RetryError,
-    ]
+    task_opts["retry_exceptions"] = [RetryableError]
 
     return task_opts
 

deltacat/exceptions.py

@@ -1,14 +1,349 @@
-class RetryableError(Exception):
-    pass
+from __future__ import annotations
+from enum import Enum
+import botocore
+import ray
+import logging
+import tenacity
+from deltacat import logs
+from ray.exceptions import (
+    RayError,
+    RayTaskError,
+    RuntimeEnvSetupError,
+    WorkerCrashedError,
+    NodeDiedError,
+    OutOfMemoryError,
+)
+from deltacat.storage import interface as DeltaCatStorage
+from pyarrow.lib import ArrowException, ArrowInvalid, ArrowCapacityError
+from botocore.exceptions import BotoCoreError
+from typing import Callable
+from deltacat.utils.ray_utils.runtime import (
+    get_current_ray_task_id,
+)
+from daft.exceptions import DaftTransientError, DaftCoreException
 
+logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
 
-class NonRetryableError(Exception):
-    pass
+DELTACAT_STORAGE_PARAM = "deltacat_storage"
+DELTACAT_STORAGE_KWARGS_PARAM = "deltacat_storage_kwargs"
 
 
-class ConcurrentModificationError(Exception):
-    pass
+class DeltaCatErrorNames(str, Enum):
+
+    DEPENDENCY_RAY_ERROR = "DependencyRayError"
+    DEPENDENCY_RAY_WORKER_DIED_ERROR = "DependencyRayWorkerDiedError"
+    DEPENDENCY_RAY_OUT_OF_MEMORY_ERROR = "DependencyRayOOMError"
+    DEPENDENCY_RAY_RUNTIME_SETUP_ERROR = "DependencyRayRuntimeSetupError"
+    DEPENDENCY_BOTOCORE_ERROR = "DependencyBotocoreError"
+    DEPENDENCY_BOTOCORE_CONNECTION_ERROR = "DependencyBotocoreConnectionError"
+    DEPENDENCY_BOTOCORE_CREDENTIAL_ERROR = "DependencyBotocoreCredentialError"
+    DEPENDENCY_BOTOCORE_TIMEOUT_ERROR = "DependencyBotocoreTimeoutError"
+    NON_RETRYABLE_DOWNLOAD_TABLE_ERROR = "NonRetryableDownloadTableError"
+    NON_RETRYABLE_DOWNLOAD_FILE_ERROR = "NonRetryableDownloadFileError"
+    NON_RETRYABLE_UPLOAD_TABLE_ERROR = "NonRetryableUploadTableError"
+    NON_RETRYABLE_UPLOAD_FILE_ERROR = "NonRetryableUploadFileError"
+    DEPENDENCY_PYARROW_ERROR = "DependencyPyarrowError"
+    DEPENDENCY_PYARROW_INVALID_ERROR = "DependencyPyarrowInvalidError"
+    DEPENDENCY_PYARROW_CAPACITY_ERROR = "DependencyPyarrowCapacityError"
+    PYMEMCACHED_PUT_OBJECT_ERROR = "PymemcachedPutObjectError"
+    DEPENDENCY_DAFT_ERROR = "DependencyDaftError"
+
+    GENERAL_THROTTLING_ERROR = "GeneralThrottlingError"
+    RETRYABLE_UPLOAD_TABLE_ERROR = "RetryableUploadTableError"
+    RETRYABLE_UPLOAD_FILE_ERROR = "RetryableUploadFileError"
+    RETRYABLE_DOWNLOAD_FILE_ERROR = "RetryableDownloadFileError"
+    RETRYABLE_DOWNLOAD_TABLE_ERROR = "RetryableDownloadTableError"
+    RETRYABLE_TIMEOUT_ERROR = "RetryableTimeoutError"
+    DEPENDENCY_DAFT_TRANSIENT_ERROR = "DependencyDaftTransientError"
+
+    VALIDATION_ERROR = "ValidationError"
+    CONTENT_TYPE_VALIDATION_ERROR = "ContentTypeValidationError"
+
+    DELTACAT_SYSTEM_ERROR = "DeltaCatSystemError"
+    DELTACAT_TRANSIENT_ERROR = "DeltaCatTransientError"
+    UNCLASSIFIED_DELTACAT_ERROR = "UnclassifiedDeltaCatError"
+    UNRECOGNIZED_RAY_TASK_ERROR = "UnrecognizedRayTaskError"
+
+
+class DeltaCatError(Exception):
+    def __init__(self, *args, **kwargs):
+        task_id, node_ip = self._get_ray_task_id_and_node_ip()
+        self.task_id = task_id
+        self.node_ip = node_ip
+        super().__init__(*args, **kwargs)
+
+    def _get_ray_task_id_and_node_ip(self):
+        task_id = get_current_ray_task_id()
+        node_ip = ray.util.get_node_ip_address()
+        return task_id, node_ip
+
+
+class NonRetryableError(DeltaCatError):
+    is_retryable = False
+
+
+class RetryableError(DeltaCatError):
+    is_retryable = True
 
 
 class ValidationError(NonRetryableError):
-    pass
+    error_name = DeltaCatErrorNames.VALIDATION_ERROR.value
+
+
+class UnclassifiedDeltaCatError(NonRetryableError):
+    error_name = DeltaCatErrorNames.UNCLASSIFIED_DELTACAT_ERROR.value
+
+
+class DependencyRayError(NonRetryableError):
+    error_name = DeltaCatErrorNames.DEPENDENCY_RAY_ERROR.value
+
+
+class DeltaCatTransientError(RetryableError):
+    error_name = DeltaCatErrorNames.DELTACAT_TRANSIENT_ERROR.value
+
+
+class DependencyDaftError(NonRetryableError):
+    error_name = DeltaCatErrorNames.DEPENDENCY_DAFT_ERROR.value
+
+
+class DependencyRayWorkerDiedError(RetryableError):
+    error_name = DeltaCatErrorNames.DEPENDENCY_RAY_WORKER_DIED_ERROR.value
+
+
+class DependencyRayOutOfMemoryError(RetryableError):
+    error_name = DeltaCatErrorNames.DEPENDENCY_RAY_OUT_OF_MEMORY_ERROR.value
+
+
+class DependencyRayRuntimeSetupError(RetryableError):
+    error_name = DeltaCatErrorNames.DEPENDENCY_RAY_RUNTIME_SETUP_ERROR.value
+
+
+class DependencyPyarrowError(NonRetryableError):
+    error_name = DeltaCatErrorNames.DEPENDENCY_PYARROW_ERROR.value
+
+
+class DependencyPyarrowInvalidError(NonRetryableError):
+    error_name = DeltaCatErrorNames.DEPENDENCY_PYARROW_INVALID_ERROR.value
+
+
+class DependencyPyarrowCapacityError(NonRetryableError):
+    error_name = DeltaCatErrorNames.DEPENDENCY_PYARROW_CAPACITY_ERROR.value
+
+
+class PymemcachedPutObjectError(RetryableError):
+    error_name = DeltaCatErrorNames.PYMEMCACHED_PUT_OBJECT_ERROR.value
+
+
+class ContentTypeValidationError(NonRetryableError):
+    error_name = DeltaCatErrorNames.CONTENT_TYPE_VALIDATION_ERROR.value
+
+
+class DependencyBotocoreError(NonRetryableError):
+    error_name = DeltaCatErrorNames.DEPENDENCY_BOTOCORE_ERROR.value
+
+
+class DependencyBotocoreConnectionError(DeltaCatTransientError):
+    error_name = DeltaCatErrorNames.DEPENDENCY_BOTOCORE_CONNECTION_ERROR.value
+
+
+class DependencyBotocoreCredentialError(DeltaCatTransientError):
+    error_name = DeltaCatErrorNames.DEPENDENCY_BOTOCORE_CREDENTIAL_ERROR.value
+
+
+class DependencyBotocoreTimeoutError(DeltaCatTransientError):
+    error_name = DeltaCatErrorNames.DEPENDENCY_BOTOCORE_TIMEOUT_ERROR.value
+
+
+class NonRetryableDownloadFileError(NonRetryableError):
+    error_name = DeltaCatErrorNames.NON_RETRYABLE_DOWNLOAD_FILE_ERROR.value
+
+
+class NonRetryableDownloadTableError(NonRetryableDownloadFileError):
+    error_name = DeltaCatErrorNames.NON_RETRYABLE_DOWNLOAD_TABLE_ERROR.value
+
+
+class NonRetryableUploadFileError(NonRetryableError):
+    error_name = DeltaCatErrorNames.NON_RETRYABLE_UPLOAD_FILE_ERROR.value
+
+
+class NonRetryableUploadTableError(NonRetryableUploadFileError):
+    error_name = DeltaCatErrorNames.NON_RETRYABLE_UPLOAD_TABLE_ERROR.value
+
+
+class GeneralThrottlingError(RetryableError):
+    error_name = DeltaCatErrorNames.GENERAL_THROTTLING_ERROR.value
+
+
+class RetryableUploadFileError(RetryableError):
+    error_name = DeltaCatErrorNames.RETRYABLE_UPLOAD_FILE_ERROR.value
+
+
+class RetryableUploadTableError(RetryableUploadFileError):
+    error_name = DeltaCatErrorNames.RETRYABLE_UPLOAD_TABLE_ERROR.value
+
+
+class RetryableDownloadFileError(RetryableError):
+    error_name = DeltaCatErrorNames.RETRYABLE_DOWNLOAD_FILE_ERROR.value
+
+
+class RetryableDownloadTableError(RetryableDownloadFileError):
+    error_name = DeltaCatErrorNames.RETRYABLE_DOWNLOAD_TABLE_ERROR.value
+
+
+class RetryableTimeoutError(RetryableError):
+    error_name = DeltaCatErrorNames.RETRYABLE_TIMEOUT_ERROR.value
+
+
+class DependencyDaftTransientError(RetryableError):
+    error_name = DeltaCatErrorNames.DEPENDENCY_DAFT_TRANSIENT_ERROR.value
+
+
+class DeltaCatSystemError(NonRetryableError):
+    error_name = DeltaCatErrorNames.DELTACAT_SYSTEM_ERROR.value
+
+
+class UnrecognizedRayTaskError(NonRetryableError):
+    error_name = DeltaCatErrorNames.UNRECOGNIZED_RAY_TASK_ERROR.value
+
+
+def categorize_errors(func: Callable):
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except BaseException as e:
+            deltacat_storage = None
+            deltacat_storage_kwargs = {}
+            if kwargs:
+                deltacat_storage = kwargs.get(DELTACAT_STORAGE_PARAM)
+                deltacat_storage_kwargs = kwargs.get(DELTACAT_STORAGE_KWARGS_PARAM, {})
+            if not deltacat_storage and args:
+                for arg in args:
+                    if (
+                        isinstance(arg, dict)
+                        and arg.get(DELTACAT_STORAGE_PARAM) is not None
+                    ):
+                        deltacat_storage = arg.get(DELTACAT_STORAGE_PARAM)
+                        deltacat_storage_kwargs = arg.get(
+                            DELTACAT_STORAGE_KWARGS_PARAM, {}
+                        )
+                        break
+
+            categorize_deltacat_exception(e, deltacat_storage, deltacat_storage_kwargs)
+
+    return wrapper
+
+
+def categorize_deltacat_exception(
+    e: BaseException,
+    deltacat_storage: DeltaCatStorage = None,
+    deltacat_storage_kwargs: dict = None,
+):
+    if deltacat_storage_kwargs is None:
+        deltacat_storage_kwargs = {}
+
+    if isinstance(e, DeltaCatError):
+        raise e
+    elif deltacat_storage and deltacat_storage.can_categorize(
+        e, **deltacat_storage_kwargs
+    ):
+        deltacat_storage.raise_categorized_error(e, **deltacat_storage_kwargs)
+    elif isinstance(e, RayError):
+        _categorize_ray_error(e)
+    elif isinstance(e, tenacity.RetryError):
+        _categorize_tenacity_error(e)
+    elif isinstance(e, ArrowException):
+        _categorize_dependency_pyarrow_error(e)
+    elif isinstance(e, AssertionError):
+        _categorize_assertion_error(e)
+    elif isinstance(e, DaftCoreException):
+        _categorize_daft_error(e)
+    elif isinstance(e, BotoCoreError):
+        _categorize_botocore_error(e)
+    else:
+        _categorize_all_remaining_errors(e)
+
+    logger.error(f"Error categorization failed for {e}.", exc_info=True)
+    raise UnclassifiedDeltaCatError(
+        "Error could not categorized into DeltaCat error"
+    ) from e
+
+
+def _categorize_ray_error(e: RayError):
+    if isinstance(e, RuntimeEnvSetupError):
+        raise DependencyRayRuntimeSetupError("Ray failed to setup runtime env.") from e
+    elif isinstance(e, WorkerCrashedError) or isinstance(e, NodeDiedError):
+        raise DependencyRayWorkerDiedError("Ray worker died unexpectedly.") from e
+    elif isinstance(e, OutOfMemoryError):
+        raise DependencyRayOutOfMemoryError("Ray worker Out Of Memory.") from e
+    elif isinstance(e, RayTaskError):
+        if e.cause is not None and isinstance(e.cause, Exception):
+            categorize_deltacat_exception(e.cause)
+        else:
+            raise UnrecognizedRayTaskError(
+                "Unrecognized underlying error detected in a Ray task."
+            ) from e
+    else:
+        raise DependencyRayError("Dependency Ray error occurred.") from e
+
+
+def _categorize_tenacity_error(e: tenacity.RetryError):
+    if e.__cause__ is not None and isinstance(e.__cause__, Exception):
+        categorize_deltacat_exception(e.__cause__)
+    else:
+        raise RetryableError("Unrecognized retryable error occurred.") from e
+
+
+def _categorize_dependency_pyarrow_error(e: ArrowException):
+    if isinstance(e, ArrowInvalid):
+        raise DependencyPyarrowInvalidError(
+            f"Pyarrow Invalid error occurred. Reason: {e}"
+        ) from e
+    elif isinstance(e, ArrowCapacityError):
+        raise DependencyPyarrowCapacityError("Pyarrow Capacity error occurred.") from e
+    else:
+        raise DependencyPyarrowError("Pyarrow error occurred.") from e
+
+
+def _categorize_assertion_error(e: BaseException):
+    raise ValidationError(
+        f"One of the assertions in DeltaCAT has failed. Reason: {e}"
+    ) from e
+
+
+def _categorize_daft_error(e: DaftCoreException):
+    if isinstance(e, DaftTransientError):
+        raise DependencyDaftTransientError("Daft Transient error occurred.") from e
+    elif isinstance(e, DaftCoreException):
+        raise DependencyDaftError("Daft error occurred.") from e
+
+
+def _categorize_botocore_error(e: BotoCoreError):
+    if isinstance(e, botocore.exceptions.ConnectionError) or isinstance(
+        e, botocore.exceptions.HTTPClientError
+    ):
+        raise DependencyBotocoreConnectionError(
+            "Botocore connection error occurred."
+        ) from e
+    elif isinstance(e, botocore.exceptions.CredentialRetrievalError) or isinstance(
+        e, botocore.exceptions.NoCredentialsError
+    ):
+        raise DependencyBotocoreCredentialError(
+            "Botocore credential retrieval failed"
+        ) from e
+    elif isinstance(e, botocore.exceptions.ReadTimeoutError) or isinstance(
+        e, botocore.exceptions.ConnectTimeoutError
+    ):
+        raise DependencyBotocoreTimeoutError("Botocore connection timed out.") from e
+    else:
+        raise DependencyBotocoreError("Botocore error occurred.") from e
+
+
+def _categorize_all_remaining_errors(e: BaseException):
+    if isinstance(e, ConnectionError):
+        raise DeltaCatTransientError("Connection error has occurred.") from e
+    elif isinstance(e, TimeoutError):
+        raise DeltaCatTransientError("Timeout error has occurred.") from e
+    elif isinstance(e, OSError):
+        raise DeltaCatTransientError("OSError occurred.") from e
+    elif isinstance(e, SystemExit):
+        raise DeltaCatSystemError("Unexpected System error occurred.") from e

deltacat/io/dataset.py

@@ -6,9 +6,6 @@ from typing import Any, Callable, Dict, Optional, TypeVar, Union, cast
 import pyarrow as pa
 import s3fs
 from ray.data import Dataset
-from ray.data.datasource import BlockWritePathProvider, DefaultBlockWritePathProvider
-
-from deltacat.io.aws.redshift.redshift_datasource import RedshiftDatasource
 
 T = TypeVar("T")
 
@@ -27,7 +24,6 @@ class DeltacatDataset(Dataset[T]):
         filesystem: Optional[Union[pa.fs.FileSystem, s3fs.S3FileSystem]] = None,
         try_create_dir: bool = True,
         arrow_open_stream_args: Optional[Dict[str, Any]] = None,
-        block_path_provider: BlockWritePathProvider = DefaultBlockWritePathProvider(),
         arrow_parquet_args_fn: Callable[[], Dict[str, Any]] = lambda: {},
         **arrow_parquet_args,
     ) -> None:
@@ -59,9 +55,8 @@ class DeltacatDataset(Dataset[T]):
                 if True. Does nothing if all directories already exist.
             arrow_open_stream_args: kwargs passed to
                 pyarrow.fs.FileSystem.open_output_stream
-            block_path_provider: BlockWritePathProvider implementation
-                to write each dataset block to a custom output path. Uses
-                DefaultBlockWritePathProvider if None.
+            filename_provider: FilenameProvider implementation
+                to write each dataset block to a custom output path.
             arrow_parquet_args_fn: Callable that returns a dictionary of write
                 arguments to use when writing each block to a file. Overrides
                 any duplicate keys from arrow_parquet_args. This should be used
@@ -72,14 +67,7 @@ class DeltacatDataset(Dataset[T]):
                 pyarrow.parquet.write_table(), which is used to write out each
                 block to a file.
         """
-        self.write_datasource(
-            RedshiftDatasource(),
-            path=path,
-            dataset_uuid=self._uuid,
-            filesystem=filesystem,
-            try_create_dir=try_create_dir,
-            open_stream_args=arrow_open_stream_args,
-            block_path_provider=block_path_provider,
-            write_args_fn=arrow_parquet_args_fn,
-            **arrow_parquet_args,
+        raise NotImplementedError(
+            "Writing to Redshift is not yet supported. "
+            "Please use DeltacatDataset.write_parquet() instead."
         )

deltacat/io/memcached_object_store.py

@@ -12,6 +12,9 @@ from pymemcache.client.retrying import RetryingClient
 from pymemcache.exceptions import MemcacheUnexpectedCloseError
 from pymemcache.client.rendezvous import RendezvousHash
 from deltacat.utils.cloudpickle import dump_into_chunks
+from deltacat.exceptions import (
+    PymemcachedPutObjectError,
+)
 
 logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
 
@@ -72,7 +75,7 @@ class MemcachedObjectStore(IObjectStore):
         for create_ref_ip, ref_to_object in input.items():
             client = self._get_client_by_ip(create_ref_ip)
             if client.set_many(ref_to_object, noreply=self.noreply):
-                raise RuntimeError("Unable to write few keys to cache")
+                raise PymemcachedPutObjectError("Unable to write a few keys to cache")
 
         return result
 
@@ -87,10 +90,10 @@ class MemcachedObjectStore(IObjectStore):
 
             try:
                 if not client.set(ref, chunk, noreply=self.noreply):
-                    raise RuntimeError(f"Unable to write {ref} to cache")
+                    raise PymemcachedPutObjectError(f"Unable to write {ref} to cache")
             except BaseException as e:
-                raise RuntimeError(
-                    f"Received {e} while writing ref={ref} and obj size={len(chunk)}"
+                raise PymemcachedPutObjectError(
+                    f"Received {e} while writing ref={ref} and obj size={len(chunk)}",
                 )
 
         return self._create_ref(uid, create_ref_ip, len(serialized_list))

deltacat/storage/__init__.py

@@ -14,6 +14,20 @@ from deltacat.storage.model.stream import Stream, StreamLocator
 from deltacat.storage.model.table import Table, TableLocator
 from deltacat.storage.model.table_version import TableVersion, TableVersionLocator
 from deltacat.storage.model.delete_parameters import DeleteParameters
+from deltacat.storage.model.partition_spec import (
+    PartitionFilter,
+    PartitionValues,
+    DeltaPartitionSpec,
+    StreamPartitionSpec,
+)
+from deltacat.storage.model.transform import (
+    Transform,
+    TransformName,
+    TransformParameters,
+    BucketingStrategy,
+    BucketTransformParameters,
+    IdentityTransformParameters,
+)
 
 from deltacat.storage.model.types import (
     CommitState,
@@ -56,4 +70,14 @@ __all__ = [
     "TableVersionLocator",
     "SortKey",
     "SortOrder",
+    "PartitionFilter",
+    "PartitionValues",
+    "DeltaPartitionSpec",
+    "StreamPartitionSpec",
+    "Transform",
+    "TransformName",
+    "TransformParameters",
+    "BucketingStrategy",
+    "BucketTransformParameters",
+    "IdentityTransformParameters",
 ]