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

deltacat/utils/export.py

@@ -0,0 +1,61 @@
+import logging
+import json
+import pyarrow as pa
+import pyarrow.parquet
+import pyarrow.feather
+from typing import Callable, Dict
+
+from deltacat.experimental.storage.rivulet.reader.query_expression import (
+    QueryExpression,
+)
+from deltacat import logs
+
+logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
+
+
+def export_parquet(dataset, file_uri: str, query: QueryExpression = QueryExpression()):
+    records = dataset.scan(query).to_arrow()
+    table = pa.Table.from_batches(records)
+    pyarrow.parquet.write_table(table, file_uri)
+
+
+def export_feather(dataset, file_uri: str, query: QueryExpression = QueryExpression()):
+    records = dataset.scan(query).to_arrow()
+    table = pa.Table.from_batches(records)
+    pyarrow.feather.write_feather(table, file_uri)
+
+
+def export_json(dataset, file_uri: str, query: QueryExpression = QueryExpression()):
+    with open(file_uri, "w") as f:
+        for batch in dataset.scan(query).to_pydict():
+            json.dump(batch, f, indent=2)
+            f.write("\n")
+
+
+def export_dataset(dataset, file_uri: str, format: str = "parquet", query=None):
+    """
+    Export the dataset to a file.
+
+    TODO: Make this pluggable for custom formats.
+
+    Args:
+        dataset: The dataset to export.
+        file_uri: The URI to write the dataset to.
+        format: The format to write the dataset in. Options are [parquet, feather, json].
+        query: QueryExpression to filter the dataset before exporting.
+    """
+    # Supported format handlers
+    export_handlers: Dict[str, Callable] = {
+        "parquet": export_parquet,
+        "feather": export_feather,
+        "json": export_json,
+    }
+
+    if format not in export_handlers:
+        raise ValueError(
+            f"Unsupported format: {format}. Supported formats are {list(export_handlers.keys())}"
+        )
+
+    export_handlers[format](dataset, file_uri, query or QueryExpression())
+
+    logger.info(f"Dataset exported to {file_uri} in {format} format.")

deltacat/utils/filesystem.py

@@ -0,0 +1,450 @@
+from __future__ import annotations
+
+import re
+from typing import Optional, Tuple, Union, List
+from datetime import timedelta
+from enum import Enum
+
+import sys
+import urllib
+import pathlib
+
+import pyarrow as pa
+from pyarrow.fs import (
+    _resolve_filesystem_and_path,
+    FileSelector,
+    FileInfo,
+    FileType,
+    FileSystem,
+    FSSpecHandler,
+    PyFileSystem,
+    GcsFileSystem,
+    LocalFileSystem,
+    S3FileSystem,
+    AzureFileSystem,
+    HadoopFileSystem,
+)
+
+_LOCAL_SCHEME = "local"
+
+
+class FilesystemType(str, Enum):
+    LOCAL = "local"
+    S3 = "s3"
+    GCS = "gcs"
+    AZURE = "azure"
+    HADOOP = "hadoop"
+    UNKNOWN = "unknown"
+
+    @classmethod
+    def from_filesystem(cls, filesystem: FileSystem) -> FilesystemType:
+        if isinstance(filesystem, LocalFileSystem):
+            return cls.LOCAL
+        elif isinstance(filesystem, S3FileSystem):
+            return cls.S3
+        elif isinstance(filesystem, GcsFileSystem):
+            return cls.GCS
+        elif isinstance(filesystem, AzureFileSystem):
+            return cls.AZURE
+        elif isinstance(filesystem, HadoopFileSystem):
+            return cls.HADOOP
+        else:
+            return cls.UNKNOWN
+
+    @classmethod
+    def to_filesystem(cls, filesystem_type: FilesystemType) -> FileSystem:
+        if filesystem_type == cls.LOCAL:
+            return LocalFileSystem()
+        elif filesystem_type == cls.S3:
+            return S3FileSystem()
+        elif filesystem_type == cls.GCS:
+            return GcsFileSystem()
+        elif filesystem_type == cls.AZURE:
+            return AzureFileSystem()
+        elif filesystem_type == cls.HADOOP:
+            return HadoopFileSystem()
+        else:
+            raise ValueError(f"Unsupported filesystem type: {filesystem_type}")
+
+
+def resolve_paths_and_filesystem(
+    paths: Union[str, List[str]],
+    filesystem: FileSystem = None,
+) -> Tuple[List[str], FileSystem]:
+    """
+    Resolves and normalizes all provided paths, infers a filesystem from the
+    paths or validates the provided filesystem against the paths and ensures
+    that all paths use the same filesystem.
+
+    Args:
+        paths: A single file/directory path or a list of file/directory paths.
+            A list of paths can contain both files and directories.
+        filesystem: The filesystem implementation that should be used for
+            reading these files. If None, a filesystem will be inferred. If not
+            None, the provided filesystem will still be validated against all
+            filesystems inferred from the provided paths to ensure
+            compatibility.
+    """
+    if isinstance(paths, str):
+        paths = [paths]
+    if isinstance(paths, pathlib.Path):
+        paths = [str(paths)]
+    elif not isinstance(paths, list) or any(not isinstance(p, str) for p in paths):
+        raise ValueError(
+            "Expected `paths` to be a `str`, `pathlib.Path`, or `list[str]`, but got "
+            f"`{paths}`."
+        )
+    elif len(paths) == 0:
+        raise ValueError("Must provide at least one path.")
+
+    need_unwrap_path_protocol = True
+    if filesystem and not isinstance(filesystem, FileSystem):
+        err_msg = (
+            f"The filesystem passed must either conform to "
+            f"pyarrow.fs.FileSystem, or "
+            f"fsspec.spec.AbstractFileSystem. The provided "
+            f"filesystem was: {filesystem}"
+        )
+        try:
+            import fsspec
+            from fsspec.implementations.http import HTTPFileSystem
+        except ModuleNotFoundError:
+            # If filesystem is not a pyarrow filesystem and fsspec isn't
+            # installed, then filesystem is neither a pyarrow filesystem nor
+            # an fsspec filesystem, so we raise a TypeError.
+            raise TypeError(err_msg) from None
+        if not isinstance(filesystem, fsspec.spec.AbstractFileSystem):
+            raise TypeError(err_msg) from None
+        if isinstance(filesystem, HTTPFileSystem):
+            # If filesystem is fsspec HTTPFileSystem, the protocol/scheme of paths
+            # should not be unwrapped/removed, because HTTPFileSystem expects full file
+            # paths including protocol/scheme. This is different behavior compared to
+            # file systems implementation in pyarrow.fs.FileSystem.
+            need_unwrap_path_protocol = False
+
+        filesystem = PyFileSystem(FSSpecHandler(filesystem))
+
+    resolved_paths = []
+    for path in paths:
+        path = _resolve_custom_scheme(path)
+        try:
+            resolved_filesystem, resolved_path = _resolve_filesystem_and_path(
+                path, filesystem
+            )
+        except pa.lib.ArrowInvalid as e:
+            if "Cannot parse URI" in str(e):
+                resolved_filesystem, resolved_path = _resolve_filesystem_and_path(
+                    _encode_url(path), filesystem
+                )
+                resolved_path = _decode_url(resolved_path)
+            elif "Unrecognized filesystem type in URI" in str(e):
+                scheme = urllib.parse.urlparse(path, allow_fragments=False).scheme
+                if scheme in ["http", "https"]:
+                    # If scheme of path is HTTP and filesystem is not resolved,
+                    # try to use fsspec HTTPFileSystem. This expects fsspec is
+                    # installed.
+                    try:
+                        from fsspec.implementations.http import HTTPFileSystem
+                    except ModuleNotFoundError:
+                        raise ImportError(
+                            "Please install fsspec to read files from HTTP."
+                        ) from None
+
+                    resolved_filesystem = PyFileSystem(FSSpecHandler(HTTPFileSystem()))
+                    resolved_path = path
+                    need_unwrap_path_protocol = False
+                else:
+                    raise
+            else:
+                raise
+        if filesystem is None:
+            if isinstance(resolved_filesystem, GcsFileSystem):
+                # 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 = GcsFileSystem(
+                    retry_time_limit=timedelta(seconds=60),
+                )
+            else:
+                filesystem = resolved_filesystem
+        elif need_unwrap_path_protocol:
+            resolved_path = _unwrap_protocol(resolved_path)
+        resolved_path = filesystem.normalize_path(resolved_path)
+        resolved_paths.append(resolved_path)
+    return resolved_paths, filesystem
+
+
+def resolve_path_and_filesystem(
+    path: str,
+    filesystem: Optional[FileSystem] = None,
+) -> Tuple[str, FileSystem]:
+    """
+    Resolves and normalizes the provided path, infers a filesystem from the
+    path or validates the provided filesystem against the path.
+
+    Args:
+        path: A single file/directory path.
+        filesystem: The filesystem implementation that should be used for
+            reading these files. If None, a filesystem will be inferred. If not
+            None, the provided filesystem will still be validated against all
+            filesystems inferred from the provided paths to ensure
+            compatibility.
+    """
+    paths, filesystem = resolve_paths_and_filesystem(
+        paths=path,
+        filesystem=filesystem,
+    )
+    assert len(paths) == 1, len(paths)
+    return paths[0], filesystem
+
+
+def list_directory(
+    path: str,
+    filesystem: FileSystem,
+    exclude_prefixes: Optional[List[str]] = None,
+    ignore_missing_path: bool = False,
+    recursive: bool = False,
+) -> List[Tuple[str, int]]:
+    """
+    Expand the provided directory path to a list of file paths.
+
+    Args:
+        path: The directory path to expand.
+        filesystem: The filesystem implementation that should be used for
+            reading these files.
+        exclude_prefixes: The file relative path prefixes that should be
+            excluded from the returned file set. Default excluded prefixes are
+            "." and "_".
+        recursive: Whether to expand subdirectories or not.
+
+    Returns:
+        An iterator of (file_path, file_size) tuples.
+    """
+    if exclude_prefixes is None:
+        exclude_prefixes = [".", "_"]
+
+    selector = FileSelector(
+        base_dir=path,
+        recursive=recursive,
+        allow_not_found=ignore_missing_path,
+    )
+    try:
+        files = filesystem.get_file_info(selector)
+    except OSError as e:
+        if isinstance(e, FileNotFoundError):
+            files = []
+        else:
+            _handle_read_os_error(e, path)
+    base_path = selector.base_dir
+    out = []
+    for file_ in files:
+        file_path = file_.path
+        if not file_path.startswith(base_path):
+            continue
+        relative = file_path[len(base_path) :]
+        if any(relative.startswith(prefix) for prefix in exclude_prefixes):
+            continue
+        out.append((file_path, file_.size))
+    # We sort the paths to guarantee a stable order.
+    return sorted(out)
+
+
+def get_file_info(
+    path: str,
+    filesystem: FileSystem,
+    ignore_missing_path: bool = False,
+) -> FileInfo:
+    """Get the file info for the provided path."""
+    try:
+        file_info = filesystem.get_file_info(path)
+    except OSError as e:
+        _handle_read_os_error(e, path)
+    if file_info.type == FileType.NotFound and not ignore_missing_path:
+        raise FileNotFoundError(path)
+
+    return file_info
+
+
+def write_file(
+    path: str,
+    data: Union[str, bytes],
+    filesystem: Optional[FileSystem] = None,
+) -> None:
+    """
+    Write data to a file using any filesystem.
+
+    Args:
+        path: The file path to write to.
+        data: The data to write (string or bytes).
+        filesystem: The filesystem implementation to use. If None, will be inferred from the path.
+    """
+    resolved_path, resolved_filesystem = resolve_path_and_filesystem(
+        path=path,
+        filesystem=filesystem,
+    )
+
+    # Convert string to bytes if necessary
+    if isinstance(data, str):
+        data = data.encode("utf-8")
+
+    with resolved_filesystem.open_output_stream(resolved_path) as f:
+        f.write(data)
+
+
+def read_file(
+    path: str,
+    filesystem: Optional[FileSystem] = None,
+    fail_if_not_found: bool = True,
+) -> Optional[bytes]:
+    """
+    Read data from a file using any filesystem.
+
+    Args:
+        path: The file path to read from.
+        filesystem: The filesystem implementation to use. If None, will be inferred from the path.
+        fail_if_not_found: Whether to raise an error if the file is not found.
+
+    Returns:
+        The file data as bytes, or None if file not found and fail_if_not_found is False.
+    """
+    try:
+        resolved_path, resolved_filesystem = resolve_path_and_filesystem(
+            path=path,
+            filesystem=filesystem,
+        )
+
+        with resolved_filesystem.open_input_stream(resolved_path) as f:
+            return f.read()
+    except FileNotFoundError:
+        if fail_if_not_found:
+            raise
+        return None
+
+
+def _handle_read_os_error(
+    error: OSError,
+    paths: Union[str, List[str]],
+) -> str:
+    # NOTE: this is not comprehensive yet, and should be extended as more errors arise.
+    # NOTE: The latter patterns are raised in Arrow 10+, while the former is raised in
+    # Arrow < 10.
+    aws_error_pattern = (
+        r"^(?:(.*)AWS Error \[code \d+\]: No response body\.(.*))|"
+        r"(?:(.*)AWS Error UNKNOWN \(HTTP status 400\) during HeadObject operation: "
+        r"No response body\.(.*))|"
+        r"(?:(.*)AWS Error ACCESS_DENIED during HeadObject operation: No response "
+        r"body\.(.*))$"
+    )
+    gcp_error_pattern = (
+        r"^(?:(.*)google::cloud::Status\(UNAVAILABLE:(.*?)Couldn't resolve host name)"
+    )
+    if re.match(aws_error_pattern, str(error)):
+        # Specially handle AWS error when reading files, to give a clearer error
+        # message to avoid confusing users. The real issue is most likely that the AWS
+        # S3 file credentials have not been properly configured yet.
+        if isinstance(paths, str):
+            # Quote to highlight single file path in error message for better
+            # readability. List of file paths will be shown up as ['foo', 'boo'],
+            # so only quote single file path here.
+            paths = f'"{paths}"'
+        raise OSError(
+            (
+                f"Failing to read AWS S3 file(s): {paths}. "
+                "Please check that file exists and has properly configured access. "
+                "You can also run AWS CLI command to get more detailed error message "
+                "(e.g., aws s3 ls <file-name>). "
+                "See https://awscli.amazonaws.com/v2/documentation/api/latest/reference/s3/index.html "  # noqa
+                "and https://arrow.apache.org/docs/python/generated/pyarrow.fs.S3FileSystem.html "
+                "for more information."
+            )
+        )
+    elif re.match(gcp_error_pattern, str(error)):
+        # Special handling for GCP errors (e.g., handling the special case of
+        # requiring the filesystem to be instantiated with anonymous access to
+        # read public files).
+        if isinstance(paths, str):
+            paths = f'"{paths}"'
+        raise OSError(
+            (
+                f"Failing to read GCP GS file(s): {paths}. "
+                "Please check that file exists and has properly configured access. "
+                "If this is a public file, please instantiate a filesystem with "
+                "anonymous access via `pyarrow.fs.GcsFileSystem(anonymous=True)` "
+                "to read it. See https://google.aip.dev/auth/4110 and "
+                "https://arrow.apache.org/docs/python/generated/pyarrow.fs.GcsFileSystem.html"  # noqa
+                "for more information."
+            )
+        )
+
+    else:
+        raise error
+
+
+def _is_local_windows_path(path: str) -> bool:
+    """Determines if path is a Windows file-system location."""
+    if sys.platform != "win32":
+        return False
+
+    if len(path) >= 1 and path[0] == "\\":
+        return True
+    if (
+        len(path) >= 3
+        and path[1] == ":"
+        and (path[2] == "/" or path[2] == "\\")
+        and path[0].isalpha()
+    ):
+        return True
+    return False
+
+
+def _unwrap_protocol(path):
+    """
+    Slice off any protocol prefixes on path.
+    """
+    if sys.platform == "win32" and _is_local_windows_path(path):
+        # Represent as posix path such that downstream functions properly handle it.
+        # This is executed when 'file://' is NOT included in the path.
+        return pathlib.Path(path).as_posix()
+
+    parsed = urllib.parse.urlparse(path, allow_fragments=False)  # support '#' in path
+    query = "?" + parsed.query if parsed.query else ""  # support '?' in path
+    netloc = parsed.netloc
+    if parsed.scheme == "s3" and "@" in parsed.netloc:
+        # If the path contains an @, it is assumed to be an anonymous
+        # credentialed path, and we need to strip off the credentials.
+        netloc = parsed.netloc.split("@")[-1]
+
+    parsed_path = parsed.path
+    # urlparse prepends the path with a '/'. This does not work on Windows
+    # so if this is the case strip the leading slash.
+    if (
+        sys.platform == "win32"
+        and not netloc
+        and len(parsed_path) >= 3
+        and parsed_path[0] == "/"  # The problematic leading slash
+        and parsed_path[1].isalpha()  # Ensure it is a drive letter.
+        and parsed_path[2:4] in (":", ":/")
+    ):
+        parsed_path = parsed_path[1:]
+
+    return netloc + parsed_path + query
+
+
+def _encode_url(path):
+    return urllib.parse.quote(path, safe="/:")
+
+
+def _decode_url(path):
+    return urllib.parse.unquote(path)
+
+
+def _resolve_custom_scheme(path: str) -> str:
+    """Returns the resolved path if the given path follows a Ray-specific custom
+    scheme. Othewise, returns the path unchanged.
+
+    The supported custom schemes are: "local", "example".
+    """
+    parsed_uri = urllib.parse.urlparse(path)
+    if parsed_uri.scheme == _LOCAL_SCHEME:
+        path = parsed_uri.netloc + parsed_uri.path
+    return path

deltacat/utils/metafile_locator.py

@@ -0,0 +1,74 @@
+import posixpath
+import pyarrow.fs
+
+from deltacat.constants import REV_DIR_NAME
+from deltacat.storage.model.partition import PartitionLocator
+from deltacat.utils.filesystem import resolve_path_and_filesystem
+
+"""
+Helper functions to work with deltacat metadata paths.
+TODO: Replace with direct calls to Deltacat storage interface.
+"""
+
+
+def _find_first_child_with_rev(
+    parent_path: str, filesystem: pyarrow.fs.FileSystem
+) -> str:
+    """
+    Walks the filesystem to find the first child directory with a `rev/` folder.
+
+    This is a temporary solution to locate the first Namespace and Table directories.
+    The Deltacat Storage interface will provide a more robust way to locate these directories.
+
+    param: parent_path: The parent directory to search for a child with a `rev/` folder.
+    param: filesystem: The filesystem to search for the child directory.
+    returns: The name of the first child directory with a `rev/` folder.
+    """
+    children = filesystem.get_file_info(
+        pyarrow.fs.FileSelector(parent_path, allow_not_found=True)
+    )
+    for child in children:
+        if child.type == pyarrow.fs.FileType.Directory:
+            rev_path = posixpath.join(child.path, REV_DIR_NAME)
+            if filesystem.get_file_info(rev_path).type == pyarrow.fs.FileType.Directory:
+                return child.base_name
+    raise ValueError(f"No directory with 'rev/' found under {parent_path}")
+
+
+def _find_table_path(root_path: str, filesystem: pyarrow.fs.FileSystem):
+    """
+    Finds a path with structure: root/namespace_id/table_id
+    Uses _find_first_child_with_rev to determine the namespace and table ids.
+
+    param: root_path: The root directory to search for the namespace and table directories.
+    param: filesystem: The filesystem to search for the namespace and table directories.
+    returns: The path to the table directory.
+    raises: ValueError if the namespace or table directories are not found.
+    """
+    try:
+        # Find Namespace (first directory under root with rev/)
+        namespace_id = _find_first_child_with_rev(root_path, filesystem)
+        namespace_path = posixpath.join(root_path, namespace_id)
+
+        # Find Table (first directory under namespace with rev/)
+        table_id = _find_first_child_with_rev(namespace_path, filesystem)
+        return posixpath.join(namespace_path, table_id)
+
+    except ValueError as e:
+        raise ValueError(f"Failed to locate Namespace or Table: {e}") from e
+
+
+def _find_partition_path(root_path: str, locator: PartitionLocator) -> str:
+    """
+    Finds the path to the partition directory for the specified locator.
+
+    param: root_uri: The root URI of the dataset.
+    param: locator: The DeltaLocator for the delta.
+    returns: The path to the delta directory.
+    """
+    root_path, filesystem = resolve_path_and_filesystem(root_path)
+    return posixpath.join(
+        _find_table_path(root_path, filesystem),
+        locator.table_version,
+        locator.stream_id,
+    )