deltacat 0.2.11__py3-none-any.whl → 1.0.1__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.11"
+__version__ = "1.0.1"
 
 
 __all__ = [

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}"
+        )