deltacat 1.1.38__py3-none-any.whl → 2.0.0__py3-none-any.whl

deltacat/api.py

@@ -0,0 +1,578 @@
+import time
+from dataclasses import dataclass
+from typing import Any, Union, List, Optional, Dict, Callable, Tuple
+import logging
+
+import ray
+import deltacat as dc
+import pyarrow.fs as pafs
+
+from pyarrow.fs import FileType
+from ray.exceptions import OutOfMemoryError
+
+from deltacat.constants import BYTES_PER_GIBIBYTE
+from deltacat.io import (
+    read_deltacat,
+    DeltacatReadType,
+)
+from deltacat.storage import (
+    Namespace,
+    Table,
+    TableVersion,
+    Stream,
+    Partition,
+    Delta,
+    Dataset,
+    DistributedDataset,
+    ListResult,
+    LocalTable,
+    Metafile,
+)
+from deltacat.types.media import DatasetType
+from deltacat.utils.url import (
+    DeltaCatUrl,
+    DeltaCatUrlReader,
+    DeltaCatUrlWriter,
+)
+from deltacat.utils.common import ReadKwargsProvider
+from deltacat.types.tables import (
+    get_table_size,
+    get_table_length,
+)
+from deltacat.utils.filesystem import (
+    resolve_path_and_filesystem,
+    get_file_info,
+)
+from deltacat.utils.performance import timed_invocation
+from deltacat.utils.ray_utils.runtime import (
+    current_node_resources,
+    live_cpu_waiter,
+    live_node_resource_keys,
+    other_live_node_resource_keys,
+    find_max_single_node_resource_type,
+)
+from deltacat import logs
+
+logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
+
+"""
+    # CLI Example of Copying from Source to Dest without file conversion
+    # (i.e., register only - shallow copy):
+    $ dcat cp json+s3://my_bucket/log_manager/ dc://my_deltacat_catalog/log_manager/json_table
+    $ dcat cp json+s3://my_bucket/log_manager/ dc://my_deltacat_catalog/log_manager/json_table
+
+    # CLI Example of Copying from Source to Dest without file conversion
+    # (i.e., register only - deep copy):
+    $ dcat cp json+s3://my_bucket/log_manager/ dc://my_deltacat_catalog/log_manager/json_table -r
+    # The above command will make a deep copy of all JSON files found in the source
+    # to the catalog data file directory in the destination.
+
+    # CLI Example of Copying from Source to Dest with file conversion
+    # (i.e., deep copy with file content type transformation):
+    $ dcat convert json+s3://my_bucket/log_manager/ dc://my_deltacat_catalog/log_manager/ --type FEATHER
+    # The above command will read JSON files found in the source, transform them to
+    # Arrow Feather files, and register them in the destination.
+
+    # Python Example of Copying from Source to Dest with file conversion
+    # (i.e., deep copy with file content type transformation):
+    >>> ds = dc.get("json+s3://my_bucket/log_manager/")
+    >>> dc.put("dc://my_deltacat_catalog/log_manager/", dataset=ds, type=ContentType.FEATHER)
+    # Or, equivalently, we can do the write directly from the dataset:
+    >>> ds.write_deltacat("dc://my_deltacat_catalog/log_manager/", type=ContentType.FEATHER)
+"""
+
+
+def copy(
+    src: DeltaCatUrl,
+    dst: DeltaCatUrl,
+    *,
+    transforms: List[Callable[[Dataset, DeltaCatUrl], Dataset]] = [],
+    extension_to_memory_multiplier: Dict[str, float] = {
+        "pq": 5,
+        "parquet": 5,
+        "feather": 1.5,
+        "arrow": 1.5,
+        "csv": 1.5,
+        "tsv": 1.5,
+        "psv": 1.5,
+        "txt": 1.5,
+        "json": 1.5,
+        "jsonl": 1.5,
+        "gz": 35,
+        "bz2": 35,
+        "zip": 35,
+        "zst": 35,
+        "7z": 35,
+        "*": 2.5,
+    },
+    minimum_worker_cpus: int = 0,
+    reader_args: Dict[str, Any] = {},
+    writer_args: Dict[str, Any] = {},
+    filesystem: Optional[pafs.FileSystem] = None,
+) -> Union[Metafile, str]:
+    """
+    Copies data from the source datastore to the destination datastore. By
+    default, this method launches one parallel Ray process to read/transform
+    each input file found in the source followed by one parallel Ray process
+    to write each output file to the destination. To ensure that adequate
+    resources are available to complete the operation, you may optionally
+    specify minimum cluster and/or worker CPUs to wait for before starting
+    parallel processing.
+
+    Args:
+        src: DeltaCAT URL of the source datastore to read.
+        dst: DeltaCAT URL of the destination datastore to write.
+        transforms: List of transforms to apply to the source dataset prior
+            to write it to the destination datastore. Transforms take the in-memory
+            dataset type read (e.g., Polars DataFrame) and source DeltaCAT URL as
+            input and return the same dataset type as output. Transforms are
+            applied to the dataset in the order given.
+        extension_to_memory_multiplier: Dictionary of file extensions to
+            in-memory inflation estimates for that extension (i.e., the amount
+            of memory required to read a source file, apply transforms, and write
+            it back to a destination file).
+        minimum_worker_cpus: The minimum number of Ray worker CPUs
+            to wait for before starting distributed execution. Useful for cases
+            where the operation is known to suffer from resource starvation (e.g.,
+            out-of-memory errors) if started before the cluster has launched a
+            minimum number of required worker nodes.
+        reader_args: Additional keyword arguments to forward to the reader
+            associated with the in-memory dataset and datastore type to read
+            (e.g., polars.read_csv(**kwargs)).
+        writer_args: Additional keyword arguments to forward to the writer
+            associated with the in-memory dataset type read and datastore type to
+            write (e.g., polars.DataFrame.write_parquet(**kwargs)).
+        filesystem: Optional PyArrow filesystem to use for file IO. Will be
+            automatically resolved from the input path if not specified, and
+            will attempt to automatically resolve storage read/write
+            credentials for the associated source/dest file cloud provider(s).
+            Try providing your own filesystem with credentials, retry strategy,
+            etc. pre-configured if you encounter latency issues or errors
+            reading/writing files.
+
+    Returns:
+        None
+    """
+    if src.is_deltacat_catalog_url() or dst.is_deltacat_catalog_url():
+        return _copy_dc(src, dst, recursive=src.url.endswith("/**"))
+    else:
+        return _copy_external_ray(
+            src,
+            dst,
+            transforms=transforms,
+            extension_to_memory_multiplier=extension_to_memory_multiplier,
+            minimum_worker_cpus=minimum_worker_cpus,
+            reader_args=reader_args,
+            writer_args=writer_args,
+            filesystem=filesystem,
+        )
+
+
+def _copy_objects_in_order(
+    src_objects: List[Metafile],
+    destination: DeltaCatUrl,
+) -> Union[Metafile, List[Metafile]]:
+    dc_dest_url = DeltaCatUrl(destination.url)
+    catalog_name = dc_dest_url.catalog_name
+
+    copied_results = []
+
+    # Group objects by type for hierarchical copying
+    # Copy objects in strict hierarchical order
+    # Namespace -> Table -> TableVersion -> Stream -> Partition -> Delta
+    ordered_objects_by_type = {
+        Namespace: [],
+        Table: [],
+        TableVersion: [],
+        Stream: [],
+        Partition: [],
+        Delta: [],
+    }
+
+    for obj in src_objects:
+        obj_class = Metafile.get_class(obj.to_serializable())
+        ordered_objects_by_type[obj_class].append(obj)
+
+    # TODO(pdames): Support copying uncommitted streams/partitions.
+    # TODO(pdames): Support parallel/distributed copies.
+    for obj_class, objects in ordered_objects_by_type.items():
+        if objects:
+            logger.info(f"Copying {len(objects)} {obj_class} objects...")
+        if obj_class == TableVersion:
+            # sort table versions by ascending table version
+            objects.sort(key=lambda x: x.current_version_number())
+        if obj_class == Delta:
+            # sort deltas by ascending stream position
+            objects.sort(key=lambda x: x.stream_position)
+        for i, obj in enumerate(objects):
+            logger.info(f"Copying object {i+1}/{len(objects)}: {obj.url}")
+            dest_url = DeltaCatUrl(obj.url(catalog_name=catalog_name))
+            logger.info(f"Destination URL for object {i+1}/{len(objects)}: {dest_url}")
+            result = put(dest_url, metafile=obj)
+            copied_results.append(result)
+            logger.info(f"Successfully copied object {i+1}/{len(objects)}")
+    return copied_results[0] if len(copied_results) == 1 else copied_results
+
+
+def _copy_dc(
+    source: DeltaCatUrl,
+    destination: DeltaCatUrl,
+    recursive: bool = False,
+) -> Union[Metafile, List[Metafile]]:
+    dc.raise_if_not_initialized()
+    if len(source.url.split("/")) != len(destination.url.split("/")):
+        # TODO(pdames): Better error message.
+        raise ValueError(
+            f"Cannot copy {source} to {destination}. "
+            f"Source and destination must share the same type."
+        )
+    if recursive:
+        src_objects = list(DeltaCatUrl(source.url.rstrip("/**")), recursive=True)
+    elif source.url.endswith("/*"):
+        src_objects = list(DeltaCatUrl(source.url.rstrip("/*")))
+    else:
+        src_objects = [get(source)]
+    return _copy_objects_in_order(src_objects, destination)
+
+
+def concat(source, destination):
+    raise NotImplementedError
+
+
+def delete(source):
+    raise NotImplementedError
+
+
+def move(source, destination):
+    raise NotImplementedError
+
+
+def _list_all_metafiles(
+    url: DeltaCatUrl,
+    recursive: bool = False,
+    **kwargs,
+) -> List[Metafile]:
+    reader = DeltaCatUrlReader(url)
+    list_results: List[ListResult[Metafile]] = []
+    lister = reader.listers.pop(0)[0]
+    # the top-level lister doesn't have any missing keyword args
+    metafiles: ListResult[Metafile] = lister(**kwargs)
+    list_results.append(metafiles)
+    if recursive:
+        # Process each level of the hierarchy
+        current_level_metafiles = [mf for mf in metafiles.all_items()]
+
+        for lister, kwarg_name, kwarg_val_resolver_fn in reader.listers:
+            next_level_metafiles = []
+            # each subsequent lister needs to inject missing keyword args from the parent metafile
+            for metafile in current_level_metafiles:
+                kwargs_update = (
+                    {kwarg_name: kwarg_val_resolver_fn(metafile)}
+                    if kwarg_name and kwarg_val_resolver_fn
+                    else {}
+                )
+                lister_kwargs = {
+                    **kwargs,
+                    **kwargs_update,
+                }
+                child_metafiles = lister(**lister_kwargs)
+                list_results.append(child_metafiles)
+                next_level_metafiles.extend(child_metafiles.all_items())
+            # Move to the next level for the next iteration
+            current_level_metafiles = next_level_metafiles
+    return [
+        metafile for list_result in list_results for metafile in list_result.all_items()
+    ]
+
+
+class CustomReadKwargsProvider(ReadKwargsProvider):
+    def __init__(
+        self,
+        datasource_type: str,
+        kwargs: Dict[str, Any],
+    ):
+        self._datasource_type = datasource_type
+        self._kwargs = kwargs
+
+    def _get_kwargs(
+        self,
+        datasource_type: str,
+        kwargs: Dict[str, Any],
+    ) -> Dict[str, Any]:
+        if datasource_type == self._datasource_type:
+            kwargs.update(self._kwargs)
+        return kwargs
+
+
+def list(
+    url: DeltaCatUrl,
+    *,
+    recursive: bool = False,
+    dataset_type: Optional[DatasetType] = None,
+    **kwargs,
+) -> Union[List[Metafile], LocalTable, DistributedDataset]:
+    if not url.is_deltacat_catalog_url():
+        raise NotImplementedError("List only supports DeltaCAT Catalog URLs.")
+    if dataset_type in DatasetType.distributed():
+        if dataset_type == DatasetType.RAY_DATASET:
+            read_type = (
+                DeltacatReadType.METADATA_LIST
+                if not recursive
+                else DeltacatReadType.METADATA_LIST_RECURSIVE
+            )
+            return read_deltacat(
+                [url],
+                deltacat_read_type=read_type,
+                timestamp_as_of=None,
+                merge_on_read=False,
+                read_kwargs_provider=CustomReadKwargsProvider(
+                    datasource_type=url.datastore_type,
+                    kwargs=kwargs,
+                ),
+            )
+        else:
+            raise NotImplementedError(
+                f"Unsupported dataset type: {dataset_type.name}. "
+                f"Supported Dataset Types: {DatasetType.RAY_DATASET.name}",
+            )
+    else:
+        # return a local list of metafiles
+        # TODO(pdames): Cast the list to the appropriate local dataset type.
+        return _list_all_metafiles(
+            url=url,
+            recursive=recursive,
+            **kwargs,
+        )
+
+
+def get(
+    url,
+    *args,
+    **kwargs,
+) -> Union[Metafile, Dataset]:
+    reader = DeltaCatUrlReader(url)
+    return reader.read(*args, **kwargs)
+
+
+def put(
+    url: DeltaCatUrl,
+    metafile: Optional[Metafile] = None,
+    *args,
+    **kwargs,
+) -> Union[Metafile, str]:
+    writer = DeltaCatUrlWriter(url, metafile=metafile)
+    return writer.write(*args, **kwargs)
+
+
+def touch(path):
+    raise NotImplementedError
+
+
+def exists(path):
+    raise NotImplementedError
+
+
+def query(expression):
+    raise NotImplementedError
+
+
+def tail(path):
+    raise NotImplementedError
+
+
+def head(path):
+    raise NotImplementedError
+
+
+def _copy_external_ray(
+    src: DeltaCatUrl,
+    dst: DeltaCatUrl,
+    *,
+    transforms: List[Callable[[Dataset, DeltaCatUrl], Dataset]] = [],
+    extension_to_memory_multiplier: Dict[str, float] = {
+        "pq": 5,
+        "parquet": 5,
+        "feather": 1.5,
+        "arrow": 1.5,
+        "csv": 1.5,
+        "tsv": 1.5,
+        "psv": 1.5,
+        "txt": 1.5,
+        "json": 1.5,
+        "jsonl": 1.5,
+        "gz": 35,
+        "bz2": 35,
+        "zip": 35,
+        "zst": 35,
+        "7z": 35,
+        "*": 2.5,
+    },
+    minimum_worker_cpus: int = 0,
+    reader_args: Dict[str, Any] = {},
+    writer_args: Dict[str, Any] = {},
+    filesystem: pafs.FileSystem = None,
+) -> str:
+    logger.info(f"DeltaCAT Copy Invocation Received at: {time.time_ns()}")
+
+    if not isinstance(src, DeltaCatUrl):
+        raise ValueError(f"Expected `src` to be a `DeltaCatUrl` but got `{src}`.")
+
+    #  wait for required resources
+    head_cpu_count = int(current_node_resources()["CPU"])
+    if minimum_worker_cpus > 0:
+        logger.info(f"Waiting for {minimum_worker_cpus} worker CPUs...")
+        live_cpu_waiter(
+            min_live_cpus=minimum_worker_cpus + head_cpu_count,
+        )
+        logger.info(f"{minimum_worker_cpus} worker CPUs found!")
+    # start job execution
+    cluster_resources = ray.cluster_resources()
+    logger.info(f"Cluster Resources: {cluster_resources}")
+    logger.info(f"Available Cluster Resources: {ray.available_resources()}")
+    cluster_cpus = int(cluster_resources["CPU"])
+    logger.info(f"Cluster CPUs: {cluster_cpus}")
+    all_node_resource_keys = live_node_resource_keys()
+    logger.info(
+        f"Found {len(all_node_resource_keys)} live nodes: {all_node_resource_keys}"
+    )
+    worker_node_resource_keys = other_live_node_resource_keys()
+    logger.info(
+        f"Found {len(worker_node_resource_keys)} live worker nodes: {worker_node_resource_keys}"
+    )
+    worker_cpu_count = cluster_cpus - head_cpu_count
+    logger.info(f"Total worker CPUs: {worker_cpu_count}")
+
+    # estimate memory requirements based on file extension
+    estimated_memory_bytes = 0
+    if extension_to_memory_multiplier:
+        logger.info(f"Resolving stats collection filesystem for: {src.url_path}.")
+        path, filesystem = resolve_path_and_filesystem(src.url_path, filesystem)
+        if isinstance(filesystem, pafs.GcsFileSystem):
+            from datetime import timedelta
+
+            # Configure a retry time limit for GcsFileSystem so that it
+            # doesn't hang forever trying to get file info (e.g., when
+            # trying to get a public file w/o anonymous=True).
+            filesystem = pafs.GcsFileSystem(
+                anonymous=True,
+                retry_time_limit=timedelta(seconds=10),
+            )
+        logger.info(f"Using filesystem {type(filesystem)} to get file size of: {path}")
+        file_info = get_file_info(path, filesystem)
+        if file_info.type != FileType.File:
+            raise ValueError(
+                f"Expected `src` to be a file but got `{file_info.type}` at "
+                f"`{src.url_path}`."
+            )
+        inflation_multiplier = extension_to_memory_multiplier.get(file_info.extension)
+        if inflation_multiplier is None:
+            inflation_multiplier = extension_to_memory_multiplier.get("*")
+        estimated_memory_bytes = inflation_multiplier * file_info.size
+        logger.info(
+            f"Estimated Memory Required for Copy: "
+            f"{estimated_memory_bytes/BYTES_PER_GIBIBYTE} GiB"
+        )
+    logger.info(f"Starting DeltaCAT Copy at: {time.time_ns()}")
+
+    index_result = None
+    num_cpus = 1
+    # TODO(pdames): remove hard-coding - issues encountered when going greater
+    #  than 2 include verifying that the scope of schedulable nodes doesn't
+    #  result in all large files lining up for the one large node in the cluster
+    #  that can actually handle them (which is worse if it's also the head node)
+    max_allowed_cpus = 2
+    while not index_result:
+        copy_task_pending, latency = timed_invocation(
+            copy_task.options(num_cpus=num_cpus, memory=estimated_memory_bytes).remote,
+            src=src,
+            dest=dst,
+            dataset_type=DatasetType.POLARS,
+            transforms=transforms,
+            reader_args=reader_args,
+            writer_args=writer_args,
+        )
+        logger.info(f"Time to Launch Copy Task: {latency} seconds")
+        try:
+            index_result, latency = timed_invocation(
+                ray.get,
+                copy_task_pending,
+            )
+        except OutOfMemoryError as e:
+            logger.warning(f"Copy Task Ran Out of Memory: {e}")
+            max_single_node_cpus = min(
+                max_allowed_cpus, find_max_single_node_resource_type("CPU")
+            )
+            num_cpus += 1
+            if num_cpus > max_single_node_cpus:
+                raise e
+            logger.info(f"Retrying Failed Copy Task with {num_cpus} dedicated CPUs")
+
+    logger.info(f"Time to Launch Copy Task: {latency} seconds")
+    logger.info(f"Time to Complete Copy Task: {latency} seconds")
+
+    total_gib_indexed = index_result.table_size / BYTES_PER_GIBIBYTE
+
+    logger.info(f"Records Copied: {index_result.table_length}")
+    logger.info(f"Bytes Copied: {total_gib_indexed} GiB")
+    logger.info(f"Conversion Rate: {total_gib_indexed/latency} GiB/s")
+    logger.info(f"Finished Copy at: {time.time_ns()}")
+
+    return dst.url
+
+
+@ray.remote(scheduling_strategy="SPREAD")
+def copy_task(
+    src: DeltaCatUrl,
+    dest: DeltaCatUrl,
+    dataset_type: DatasetType,
+    transforms: List[Callable[[Dataset, DeltaCatUrl], Dataset]] = [],
+    reader_args: Dict[str, Any] = {},
+    writer_args: Dict[str, Any] = {},
+) -> Tuple[Optional[int], int]:
+    """
+    Indexes a DeltaCAT source URL into a DeltaCAT destination URL.
+    """
+    table, latency = timed_invocation(
+        read_table,
+        src=src,
+        dataset_type=dataset_type,
+        transforms=transforms,
+        reader_args=reader_args,
+    )
+    logger.debug(f"Time to read {src.url_path}: {latency} seconds")
+
+    table_size = get_table_size(table)
+    logger.debug(f"Table Size: {table_size/BYTES_PER_GIBIBYTE} GiB")
+
+    table_length = get_table_length(table)
+    logger.debug(f"Table Records: {table_length}")
+
+    writer = DeltaCatUrlWriter(dest, dataset_type)
+    written_file_path, latency = timed_invocation(
+        writer.write,
+        "",
+        table,
+        **writer_args,
+    )
+    logger.debug(f"Time to write {written_file_path}: {latency}")
+
+    return CopyResult(table_size, table_length)
+
+
+def read_table(
+    src: DeltaCatUrl,
+    dataset_type: DatasetType,
+    transforms: List[Callable[[Dataset, DeltaCatUrl], Dataset]] = [],
+    reader_args: Dict[str, Any] = {},
+) -> LocalTable:
+    reader = DeltaCatUrlReader(src, dataset_type)
+    table: LocalTable = reader.read(**reader_args)
+    for transform in transforms:
+        table = transform(table, src)
+    return table
+
+
+@dataclass(frozen=True)
+class CopyResult:
+    table_size: int
+    table_length: int

deltacat/aws/constants.py

@@ -1,32 +1,9 @@
-import botocore
 from typing import Set
-from daft.exceptions import DaftTransientError
 from deltacat.utils.common import env_integer, env_string
 
 
 DAFT_MAX_S3_CONNECTIONS_PER_FILE = env_integer("DAFT_MAX_S3_CONNECTIONS_PER_FILE", 8)
-DEFAULT_FILE_READ_TIMEOUT_MS = env_integer(
-    "DEFAULT_FILE_READ_TIMEOUT_MS", 300_000
-)  # 5 mins
 BOTO_MAX_RETRIES = env_integer("BOTO_MAX_RETRIES", 5)
 BOTO_TIMEOUT_ERROR_CODES: Set[str] = {"ReadTimeoutError", "ConnectTimeoutError"}
 BOTO_THROTTLING_ERROR_CODES: Set[str] = {"Throttling", "SlowDown"}
-RETRYABLE_TRANSIENT_ERRORS = (
-    OSError,
-    botocore.exceptions.ConnectionError,
-    botocore.exceptions.HTTPClientError,
-    botocore.exceptions.NoCredentialsError,
-    botocore.exceptions.ConnectTimeoutError,
-    botocore.exceptions.ReadTimeoutError,
-    DaftTransientError,
-)
 AWS_REGION = env_string("AWS_REGION", "us-east-1")
-UPLOAD_DOWNLOAD_RETRY_STOP_AFTER_DELAY = env_integer(
-    "UPLOAD_DOWNLOAD_RETRY_STOP_AFTER_DELAY", 10 * 60
-)
-UPLOAD_SLICED_TABLE_RETRY_STOP_AFTER_DELAY = env_integer(
-    "UPLOAD_SLICED_TABLE_RETRY_STOP_AFTER_DELAY", 30 * 60
-)
-DOWNLOAD_MANIFEST_ENTRY_RETRY_STOP_AFTER_DELAY = env_integer(
-    "DOWNLOAD_MANIFEST_ENTRY_RETRY_STOP_AFTER_DELAY", 30 * 60
-)