deltacat 0.2.9__py3-none-any.whl → 1.0.0__py3-none-any.whl

deltacat/__init__.py

@@ -44,7 +44,7 @@ from deltacat.types.tables import TableWriteMode
 
 deltacat.logs.configure_deltacat_logger(logging.getLogger(__name__))
 
-__version__ = "0.2.9"
+__version__ = "1.0.0"
 
 
 __all__ = [

deltacat/aws/redshift/__init__.py

@@ -1,4 +1,6 @@
 from deltacat.aws.redshift.model.manifest import (
+    EntryFileParams,
+    EntryType,
     Manifest,
     ManifestAuthor,
     ManifestEntry,
@@ -7,6 +9,8 @@ from deltacat.aws.redshift.model.manifest import (
 )
 
 __all__ = [
+    "EntryFileParams",
+    "EntryType",
     "Manifest",
     "ManifestAuthor",
     "ManifestEntry",

deltacat/aws/redshift/model/manifest.py

@@ -5,12 +5,63 @@ import itertools
 import logging
 from typing import Any, Dict, List, Optional
 from uuid import uuid4
+from enum import Enum
 
 from deltacat import logs
 
 logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
 
 
+class EntryType(str, Enum):
+    """
+    Enum representing all possible content categories of an manifest entry file
+    """
+
+    DATA = "data"
+    POSITIONAL_DELETE = "positional_delete"
+    EQUALITY_DELETE = "equality_delete"
+
+    @classmethod
+    def get_default(cls):
+        return EntryType.DATA
+
+    @classmethod
+    def list(cls):
+        return [c.value for c in EntryType]
+
+
+class EntryFileParams(dict):
+    """
+    Represents parameters relevant to the underlying contents of manifest entry. Contains all parameters required to support DELETEs
+    equality_column_names: List of column names that would be used to determine row equality for equality deletes.  Relevant only to equality deletes
+    position: Ordinal position of a deleted row in the target data file identified by uri, starting at 0. Relevant only to positional deletes
+    """
+
+    @staticmethod
+    def of(
+        equality_column_names: Optional[List[str]] = None,
+        position: Optional[int] = None,
+    ) -> EntryFileParams:
+        entry_file_params = EntryFileParams()
+        if equality_column_names is not None:
+            entry_file_params["equality_column_names"] = equality_column_names
+        if position is not None:
+            entry_file_params["position"] = position
+        return entry_file_params
+
+    @property
+    def equality_column_names(self) -> Optional[List[str]]:
+        return self.get("equality_column_names")
+
+    @property
+    def url(self) -> Optional[str]:
+        return self.get("url")
+
+    @property
+    def position(self) -> Optional[int]:
+        return self.get("position")
+
+
 class Manifest(dict):
     @staticmethod
     def _build_manifest(
@@ -18,6 +69,7 @@ class Manifest(dict):
         entries: Optional[ManifestEntryList],
         author: Optional[ManifestAuthor] = None,
         uuid: str = None,
+        entry_type: Optional[EntryType] = None,
     ) -> Manifest:
         if not uuid:
             uuid = str(uuid4())
@@ -29,6 +81,8 @@ class Manifest(dict):
             manifest["entries"] = entries
         if author is not None:
             manifest["author"] = author
+        if entry_type is not None:
+            manifest["entry_type"] = entry_type.value
         return manifest
 
     @staticmethod
@@ -36,6 +90,7 @@ class Manifest(dict):
         entries: ManifestEntryList,
         author: Optional[ManifestAuthor] = None,
         uuid: str = None,
+        entry_type: Optional[EntryType] = None,
     ) -> Manifest:
         if not uuid:
             uuid = str(uuid4())
@@ -78,8 +133,9 @@ class Manifest(dict):
             content_type,
             content_encoding,
             total_source_content_length,
+            entry_type=entry_type,
         )
-        manifest = Manifest._build_manifest(meta, entries, author, uuid)
+        manifest = Manifest._build_manifest(meta, entries, author, uuid, entry_type)
         return manifest
 
     @staticmethod
@@ -128,6 +184,7 @@ class ManifestMeta(dict):
         source_content_length: Optional[int] = None,
         credentials: Optional[Dict[str, str]] = None,
         content_type_parameters: Optional[List[Dict[str, str]]] = None,
+        entry_type: Optional[EntryType] = None,
     ) -> ManifestMeta:
         manifest_meta = ManifestMeta()
         if record_count is not None:
@@ -144,6 +201,8 @@ class ManifestMeta(dict):
             manifest_meta["content_encoding"] = content_encoding
         if credentials is not None:
             manifest_meta["credentials"] = credentials
+        if entry_type is not None:
+            manifest_meta["entry_type"] = entry_type.value
         return manifest_meta
 
     @property
@@ -178,6 +237,13 @@ class ManifestMeta(dict):
     def credentials(self) -> Optional[Dict[str, str]]:
         return self.get("credentials")
 
+    @property
+    def entry_type(self) -> Optional[EntryType]:
+        val = self.get("entry_type")
+        if val is not None:
+            return EntryType(self["entry_type"])
+        return val
+
 
 class ManifestAuthor(dict):
     @staticmethod
@@ -206,6 +272,8 @@ class ManifestEntry(dict):
         mandatory: bool = True,
         uri: Optional[str] = None,
         uuid: Optional[str] = None,
+        entry_type: Optional[EntryType] = None,
+        entry_file_params: Optional[EntryFileParams] = None,
     ) -> ManifestEntry:
         manifest_entry = ManifestEntry()
         if not (uri or url):
@@ -222,6 +290,16 @@ class ManifestEntry(dict):
             manifest_entry["mandatory"] = mandatory
         if uuid is not None:
             manifest_entry["id"] = uuid
+        if entry_type is not None:
+            manifest_entry["entry_type"] = entry_type.value
+        if entry_file_params is not None:
+            if entry_file_params.get("url") != manifest_entry.get("url"):
+                msg = (
+                    f"Expected manifest entry url: {manifest_entry.url}"
+                    f" and entry_file_params: '{entry_file_params.url}' to match"
+                )
+                raise ValueError(msg)
+            manifest_entry["entry_file_params"] = entry_file_params
         return manifest_entry
 
     @staticmethod
@@ -268,6 +346,20 @@ class ManifestEntry(dict):
     def id(self) -> Optional[str]:
         return self.get("id")
 
+    @property
+    def entry_type(self) -> Optional[EntryType]:
+        val = self.get("entry_type")
+        if val is not None:
+            return EntryType(self["entry_type"])
+        return val
+
+    @property
+    def entry_file_params(self) -> Optional[EntryFileParams]:
+        val: Dict[str, Any] = self.get("entry_file_params")
+        if val is not None and not isinstance(val, EntryFileParams):
+            self["entry_file_params"] = val = EntryFileParams(val)
+        return val
+
 
 class ManifestEntryList(List[ManifestEntry]):
     @staticmethod

deltacat/aws/s3u.py

@@ -22,7 +22,7 @@ from tenacity import (
     stop_after_delay,
     wait_random_exponential,
 )
-
+from deltacat.utils.ray_utils.concurrency import invoke_parallel
 import deltacat.aws.clients as aws_utils
 from deltacat import logs
 from deltacat.aws.constants import TIMEOUT_ERROR_CODES
@@ -35,10 +35,17 @@ from deltacat.storage import (
     ManifestEntry,
     ManifestEntryList,
 )
-from deltacat.types.media import ContentEncoding, ContentType, TableType
+from deltacat.types.media import (
+    ContentEncoding,
+    ContentType,
+    TableType,
+    DistributedDatasetType,
+)
 from deltacat.types.tables import (
     TABLE_CLASS_TO_SIZE_FUNC,
     TABLE_TYPE_TO_READER_FUNC,
+    TABLE_TYPE_TO_DATASET_CREATE_FUNC_REFS,
+    DISTRIBUTED_DATASET_TYPE_TO_READER_FUNC,
     get_table_length,
 )
 from deltacat.types.partial_download import PartialFileDownloadParams
@@ -284,59 +291,6 @@ def upload_sliced_table(
     return manifest_entries
 
 
-@ray.remote
-def _block_metadata(block: Block) -> BlockMetadata:
-    return BlockAccessor.for_block(block).get_metadata(
-        input_files=None,
-        exec_stats=None,
-    )
-
-
-def _get_metadata(
-    table: Union[LocalTable, DistributedDataset],
-    write_paths: List[str],
-    block_refs: List[ObjectRef[Block]],
-) -> List[BlockMetadata]:
-    metadata: List[BlockMetadata] = []
-    if not block_refs:
-        # this must be a local table - ensure it was written to only 1 file
-        assert len(write_paths) == 1, (
-            f"Expected table of type '{type(table)}' to be written to 1 "
-            f"file, but found {len(write_paths)} files."
-        )
-        table_size = None
-        table_size_func = TABLE_CLASS_TO_SIZE_FUNC.get(type(table))
-        if table_size_func:
-            table_size = table_size_func(table)
-        else:
-            logger.warning(f"Unable to estimate '{type(table)}' table size.")
-        metadata.append(
-            BlockMetadata(
-                num_rows=get_table_length(table),
-                size_bytes=table_size,
-                schema=None,
-                input_files=None,
-                exec_stats=None,
-            )
-        )
-    else:
-        # TODO(pdames): Expose BlockList metadata getter from Ray Dataset?
-        # ray 1.10
-        # metadata = dataset._blocks.get_metadata()
-        # ray 2.0.0dev
-        metadata = table._plan.execute().get_metadata()
-        if (
-            not metadata
-            or metadata[0].size_bytes is None
-            or metadata[0].num_rows is None
-        ):
-            metadata_futures = [
-                _block_metadata.remote(block_ref) for block_ref in block_refs
-            ]
-            metadata = ray.get(metadata_futures)
-    return metadata
-
-
 def upload_table(
     table: Union[LocalTable, DistributedDataset],
     s3_base_url: str,
@@ -403,17 +357,7 @@ def download_manifest_entry(
     content_encoding: Optional[ContentEncoding] = None,
 ) -> LocalTable:
 
-    conf = Config(retries={"max_attempts": BOTO_MAX_RETRIES, "mode": "adaptive"})
-    s3_client_kwargs = (
-        {
-            "aws_access_key_id": token_holder["accessKeyId"],
-            "aws_secret_access_key": token_holder["secretAccessKey"],
-            "aws_session_token": token_holder["sessionToken"],
-            "config": conf,
-        }
-        if token_holder
-        else {"config": conf}
-    )
+    s3_client_kwargs = _get_s3_client_kwargs_from_token(token_holder=token_holder)
     if not content_type:
         content_type = manifest_entry.meta.content_type
         assert (
@@ -458,51 +402,9 @@ def download_manifest_entry(
     return table
 
 
-def _download_manifest_entries(
-    manifest: Manifest,
-    token_holder: Optional[Dict[str, Any]] = None,
-    table_type: TableType = TableType.PYARROW,
-    column_names: Optional[List[str]] = None,
-    include_columns: Optional[List[str]] = None,
-    file_reader_kwargs_provider: Optional[ReadKwargsProvider] = None,
-) -> LocalDataset:
-
-    return [
-        download_manifest_entry(
-            e,
-            token_holder,
-            table_type,
-            column_names,
-            include_columns,
-            file_reader_kwargs_provider,
-        )
-        for e in manifest.entries
-    ]
-
-
-def _download_manifest_entries_parallel(
-    manifest: Manifest,
-    token_holder: Optional[Dict[str, Any]] = None,
-    table_type: TableType = TableType.PYARROW,
-    max_parallelism: Optional[int] = None,
-    column_names: Optional[List[str]] = None,
-    include_columns: Optional[List[str]] = None,
-    file_reader_kwargs_provider: Optional[ReadKwargsProvider] = None,
-) -> LocalDataset:
-
-    tables = []
-    pool = multiprocessing.Pool(max_parallelism)
-    downloader = partial(
-        download_manifest_entry,
-        token_holder=token_holder,
-        table_type=table_type,
-        column_names=column_names,
-        include_columns=include_columns,
-        file_reader_kwargs_provider=file_reader_kwargs_provider,
-    )
-    for table in pool.map(downloader, [e for e in manifest.entries]):
-        tables.append(table)
-    return tables
+@ray.remote
+def download_manifest_entry_ray(*args, **kwargs) -> ObjectRef[LocalTable]:
+    return download_manifest_entry(*args, **kwargs)
 
 
 def download_manifest_entries(
@@ -536,6 +438,42 @@ def download_manifest_entries(
         )
 
 
+def download_manifest_entries_distributed(
+    manifest: Manifest,
+    token_holder: Optional[Dict[str, Any]] = None,
+    table_type: TableType = TableType.PYARROW,
+    max_parallelism: Optional[int] = 1000,
+    column_names: Optional[List[str]] = None,
+    include_columns: Optional[List[str]] = None,
+    file_reader_kwargs_provider: Optional[ReadKwargsProvider] = None,
+    ray_options_provider: Callable[[int, Any], Dict[str, Any]] = None,
+    distributed_dataset_type: Optional[
+        DistributedDatasetType
+    ] = DistributedDatasetType.RAY_DATASET,
+) -> DistributedDataset:
+
+    params = {
+        "manifest": manifest,
+        "token_holder": token_holder,
+        "table_type": table_type,
+        "max_parallelism": max_parallelism,
+        "column_names": column_names,
+        "include_columns": include_columns,
+        "file_reader_kwargs_provider": file_reader_kwargs_provider,
+        "ray_options_provider": ray_options_provider,
+        "distributed_dataset_type": distributed_dataset_type,
+    }
+
+    if distributed_dataset_type == DistributedDatasetType.RAY_DATASET:
+        return _download_manifest_entries_ray_data_distributed(**params)
+    elif distributed_dataset_type is not None:
+        return _download_manifest_entries_all_dataset_distributed(**params)
+    else:
+        raise ValueError(
+            f"Distributed dataset type {distributed_dataset_type} not supported."
+        )
+
+
 def upload(s3_url: str, body, **s3_client_kwargs) -> Dict[str, Any]:
 
     # TODO (pdames): add tenacity retrying
@@ -574,3 +512,204 @@ def download(
         else:
             logger.info(f"file not found: {s3_url}")
     return None
+
+
+def _download_manifest_entries_parallel(
+    manifest: Manifest,
+    token_holder: Optional[Dict[str, Any]] = None,
+    table_type: TableType = TableType.PYARROW,
+    max_parallelism: Optional[int] = None,
+    column_names: Optional[List[str]] = None,
+    include_columns: Optional[List[str]] = None,
+    file_reader_kwargs_provider: Optional[ReadKwargsProvider] = None,
+) -> LocalDataset:
+
+    tables = []
+    pool = multiprocessing.Pool(max_parallelism)
+    downloader = partial(
+        download_manifest_entry,
+        token_holder=token_holder,
+        table_type=table_type,
+        column_names=column_names,
+        include_columns=include_columns,
+        file_reader_kwargs_provider=file_reader_kwargs_provider,
+    )
+    for table in pool.map(downloader, [e for e in manifest.entries]):
+        tables.append(table)
+    return tables
+
+
+def _download_manifest_entries(
+    manifest: Manifest,
+    token_holder: Optional[Dict[str, Any]] = None,
+    table_type: TableType = TableType.PYARROW,
+    column_names: Optional[List[str]] = None,
+    include_columns: Optional[List[str]] = None,
+    file_reader_kwargs_provider: Optional[ReadKwargsProvider] = None,
+) -> LocalDataset:
+
+    return [
+        download_manifest_entry(
+            manifest_entry=e,
+            token_holder=token_holder,
+            table_type=table_type,
+            column_names=column_names,
+            include_columns=include_columns,
+            file_reader_kwargs_provider=file_reader_kwargs_provider,
+        )
+        for e in manifest.entries
+    ]
+
+
+@ray.remote
+def _block_metadata(block: Block) -> BlockMetadata:
+    return BlockAccessor.for_block(block).get_metadata(
+        input_files=None,
+        exec_stats=None,
+    )
+
+
+def _get_s3_client_kwargs_from_token(token_holder) -> Dict[Any, Any]:
+    conf = Config(retries={"max_attempts": BOTO_MAX_RETRIES, "mode": "adaptive"})
+    return (
+        {
+            "aws_access_key_id": token_holder["accessKeyId"],
+            "aws_secret_access_key": token_holder["secretAccessKey"],
+            "aws_session_token": token_holder["sessionToken"],
+            "config": conf,
+        }
+        if token_holder
+        else {"config": conf}
+    )
+
+
+def _get_metadata(
+    table: Union[LocalTable, DistributedDataset],
+    write_paths: List[str],
+    block_refs: List[ObjectRef[Block]],
+) -> List[BlockMetadata]:
+    metadata: List[BlockMetadata] = []
+    if not block_refs:
+        # this must be a local table - ensure it was written to only 1 file
+        assert len(write_paths) == 1, (
+            f"Expected table of type '{type(table)}' to be written to 1 "
+            f"file, but found {len(write_paths)} files."
+        )
+        table_size = None
+        table_size_func = TABLE_CLASS_TO_SIZE_FUNC.get(type(table))
+        if table_size_func:
+            table_size = table_size_func(table)
+        else:
+            logger.warning(f"Unable to estimate '{type(table)}' table size.")
+        metadata.append(
+            BlockMetadata(
+                num_rows=get_table_length(table),
+                size_bytes=table_size,
+                schema=None,
+                input_files=None,
+                exec_stats=None,
+            )
+        )
+    else:
+        # TODO(pdames): Expose BlockList metadata getter from Ray Dataset?
+        # ray 1.10
+        # metadata = dataset._blocks.get_metadata()
+        # ray 2.0.0dev
+        metadata = table._plan.execute().get_metadata()
+        if (
+            not metadata
+            or metadata[0].size_bytes is None
+            or metadata[0].num_rows is None
+        ):
+            metadata_futures = [
+                _block_metadata.remote(block_ref) for block_ref in block_refs
+            ]
+            metadata = ray.get(metadata_futures)
+    return metadata
+
+
+def _download_manifest_entries_ray_data_distributed(
+    manifest: Manifest,
+    token_holder: Optional[Dict[str, Any]] = None,
+    table_type: TableType = TableType.PYARROW,
+    max_parallelism: Optional[int] = 1000,
+    column_names: Optional[List[str]] = None,
+    include_columns: Optional[List[str]] = None,
+    file_reader_kwargs_provider: Optional[ReadKwargsProvider] = None,
+    ray_options_provider: Callable[[int, Any], Dict[str, Any]] = None,
+) -> DistributedDataset:
+
+    table_pending_ids = []
+    manifest_entries = manifest.entries
+    if manifest_entries:
+        table_pending_ids = invoke_parallel(
+            manifest_entries,
+            download_manifest_entry_ray,
+            token_holder,
+            table_type,
+            column_names,
+            include_columns,
+            file_reader_kwargs_provider,
+            max_parallelism=max_parallelism,
+            options_provider=ray_options_provider,
+        )
+    return TABLE_TYPE_TO_DATASET_CREATE_FUNC_REFS[table_type](table_pending_ids)
+
+
+def _download_manifest_entries_all_dataset_distributed(
+    manifest: Manifest,
+    token_holder: Optional[Dict[str, Any]] = None,
+    table_type: TableType = TableType.PYARROW,
+    max_parallelism: Optional[int] = 1000,
+    column_names: Optional[List[str]] = None,
+    include_columns: Optional[List[str]] = None,
+    file_reader_kwargs_provider: Optional[ReadKwargsProvider] = None,
+    ray_options_provider: Callable[[int, Any], Dict[str, Any]] = None,
+    distributed_dataset_type: Optional[
+        DistributedDatasetType
+    ] = DistributedDatasetType.RAY_DATASET,
+) -> DistributedDataset:
+
+    entry_content_type = None
+    entry_content_encoding = None
+    uris = []
+    for entry in manifest.entries or []:
+        if (
+            entry_content_type is not None
+            and entry_content_type != entry.meta.content_type
+        ):
+            raise ValueError(
+                f"Mixed content types of ({entry_content_type},"
+                f" {entry.meta.content_type}) is not supported."
+            )
+
+        if (
+            entry_content_encoding is not None
+            and entry_content_encoding != entry.meta.content_encoding
+        ):
+            raise ValueError(
+                f"Mixed content encoding of {entry_content_encoding},"
+                f" {entry.meta.content_encoding} is not supported."
+            )
+
+        entry_content_type = entry.meta.content_type
+        entry_content_encoding = entry.meta.content_encoding
+        uris.append(entry.uri)
+
+    s3_client_kwargs = _get_s3_client_kwargs_from_token(token_holder=token_holder)
+
+    if distributed_dataset_type in DISTRIBUTED_DATASET_TYPE_TO_READER_FUNC:
+        return DISTRIBUTED_DATASET_TYPE_TO_READER_FUNC[distributed_dataset_type.value](
+            uris=uris,
+            content_type=entry_content_type,
+            content_encoding=entry_content_encoding,
+            column_names=column_names,
+            include_columns=include_columns,
+            read_func_kwargs_provider=file_reader_kwargs_provider,
+            ray_options_provider=ray_options_provider,
+            s3_client_kwargs=s3_client_kwargs,
+        )
+    else:
+        raise ValueError(
+            f"Unsupported distributed dataset type={distributed_dataset_type}"
+        )