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

deltacat/utils/export.py

@@ -5,7 +5,9 @@ import pyarrow.parquet
 import pyarrow.feather
 from typing import Callable, Dict
 
-from deltacat.storage.rivulet.reader.query_expression import QueryExpression
+from deltacat.experimental.storage.rivulet.reader.query_expression import (
+    QueryExpression,
+)
 from deltacat import logs
 
 logger = logs.configure_deltacat_logger(logging.getLogger(__name__))

deltacat/utils/filesystem.py

@@ -2,12 +2,13 @@ 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
 import pyarrow as pa
 from pyarrow.fs import (
     _resolve_filesystem_and_path,
@@ -17,15 +18,59 @@ from pyarrow.fs import (
     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: pyarrow.fs.FileSystem = None,
-) -> Tuple[List[str], pyarrow.fs.FileSystem]:
+    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
@@ -113,19 +158,26 @@ def resolve_paths_and_filesystem(
             else:
                 raise
         if filesystem is None:
-            filesystem = resolved_filesystem
+            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[pyarrow.fs.FileSystem] = None,
-) -> Tuple[str, pyarrow.fs.FileSystem]:
+    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.
@@ -148,7 +200,7 @@ def resolve_path_and_filesystem(
 
 def list_directory(
     path: str,
-    filesystem: pyarrow.fs.FileSystem,
+    filesystem: FileSystem,
     exclude_prefixes: Optional[List[str]] = None,
     ignore_missing_path: bool = False,
     recursive: bool = False,
@@ -199,7 +251,7 @@ def list_directory(
 
 def get_file_info(
     path: str,
-    filesystem: pyarrow.fs.FileSystem,
+    filesystem: FileSystem,
     ignore_missing_path: bool = False,
 ) -> FileInfo:
     """Get the file info for the provided path."""
@@ -213,6 +265,62 @@ def get_file_info(
     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]],
@@ -227,6 +335,9 @@ def _handle_read_os_error(
         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
@@ -243,9 +354,28 @@ def _handle_read_os_error(
                 "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
 

deltacat/utils/metafile_locator.py

@@ -1,6 +1,7 @@
 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
 
@@ -28,7 +29,7 @@ def _find_first_child_with_rev(
     )
     for child in children:
         if child.type == pyarrow.fs.FileType.Directory:
-            rev_path = posixpath.join(child.path, "rev")
+            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}")

deltacat/utils/numpy.py

@@ -1,14 +1,21 @@
-from typing import List, Optional, Callable, Union
+from typing import List, Optional, Callable, Union, Dict, Any
 
+import pandas as pd
 import numpy as np
-import pyarrow as pa
 from fsspec import AbstractFileSystem
+import pyarrow.fs as pafs
+import logging
 
 from ray.data.datasource import FilenameProvider
-from deltacat.types.media import ContentType
+from deltacat.types.media import ContentType, ContentEncoding
 from deltacat.utils import pandas as pd_utils
-from deltacat.utils import pyarrow as pa_utils
+
 from deltacat.utils.common import ReadKwargsProvider
+from deltacat import logs
+from deltacat.utils.performance import timed_invocation
+from deltacat.types.partial_download import PartialFileDownloadParams
+
+logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
 
 
 def slice_ndarray(np_array: np.ndarray, max_len: Optional[int]) -> List[np.ndarray]:
@@ -22,26 +29,61 @@ def slice_ndarray(np_array: np.ndarray, max_len: Optional[int]) -> List[np.ndarr
     return [np_array[i : i + max_len] for i in range(0, len(np_array), max_len)]
 
 
-def s3_file_to_ndarray(
-    s3_url: str,
+def file_to_ndarray(
+    path: str,
     content_type: str,
-    content_encoding: str,
+    content_encoding: str = ContentEncoding.IDENTITY.value,
+    filesystem: Optional[Union[AbstractFileSystem, pafs.FileSystem]] = None,
     column_names: Optional[List[str]] = None,
     include_columns: Optional[List[str]] = None,
     pd_read_func_kwargs_provider: Optional[ReadKwargsProvider] = None,
-    **s3_client_kwargs
+    partial_file_download_params: Optional[PartialFileDownloadParams] = None,
+    fs_open_kwargs: Dict[str, Any] = {},
+    **kwargs,
 ) -> np.ndarray:
-    # TODO: Compare perf to s3 -> pyarrow -> pandas [Series/DataFrame] -> numpy
-    dataframe = pd_utils.s3_file_to_dataframe(
-        s3_url,
-        content_type,
-        content_encoding,
-        column_names,
-        include_columns,
-        pd_read_func_kwargs_provider,
-        **s3_client_kwargs
+    """
+    Read a file into a NumPy ndarray using any filesystem.
+
+    This function delegates to the pandas file_to_dataframe function and converts
+    the resulting DataFrame to a NumPy ndarray.
+
+    Args:
+        path: The file path to read
+        content_type: The content type of the file (e.g., ContentType.CSV.value)
+        content_encoding: The content encoding (default: IDENTITY)
+        filesystem: The filesystem to use (if None, will be inferred from path)
+        column_names: Optional column names to assign
+        include_columns: Optional columns to include in the result
+        pd_read_func_kwargs_provider: Optional kwargs provider for customization
+        fs_open_kwargs: Optional kwargs for filesystem open operations
+        **kwargs: Additional kwargs passed to the reader function
+
+    Returns:
+        np.ndarray: The loaded data as a NumPy ndarray
+    """
+    logger.debug(
+        f"Reading {path} to NumPy ndarray. Content type: {content_type}. "
+        f"Encoding: {content_encoding}"
     )
-    return dataframe.to_numpy()
+
+    dataframe, latency = timed_invocation(
+        pd_utils.file_to_dataframe,
+        path=path,
+        content_type=content_type,
+        content_encoding=content_encoding,
+        filesystem=filesystem,
+        column_names=column_names,
+        include_columns=include_columns,
+        pd_read_func_kwargs_provider=pd_read_func_kwargs_provider,
+        partial_file_download_params=partial_file_download_params,
+        fs_open_kwargs=fs_open_kwargs,
+        **kwargs,
+    )
+
+    ndarray, conversion_latency = timed_invocation(dataframe.to_numpy)
+    total_latency = latency + conversion_latency
+    logger.debug(f"Time to read {path} into NumPy ndarray: {total_latency}s")
+    return ndarray
 
 
 def ndarray_size(np_array: np.ndarray) -> int:
@@ -51,22 +93,72 @@ def ndarray_size(np_array: np.ndarray) -> int:
 def ndarray_to_file(
     np_array: np.ndarray,
     path: str,
-    file_system: AbstractFileSystem,
+    filesystem: Optional[Union[AbstractFileSystem, pafs.FileSystem]],
     block_path_provider: Union[FilenameProvider, Callable],
     content_type: str = ContentType.PARQUET.value,
-    **kwargs
+    **kwargs,
 ) -> None:
     """
     Writes the given Numpy ndarray to a file.
     """
+    import pyarrow as pa
 
-    # PyArrow only supports 1D ndarrays, so convert to list of 1D arrays
-    np_arrays = [array for array in np_array]
-    pa_utils.table_to_file(
-        pa.table({"data": np_arrays}),
+    # Extract schema from kwargs if available
+    schema = kwargs.pop("schema", None)
+
+    # Convert to pandas DataFrame with proper column names if schema is available
+    if schema and isinstance(schema, pa.Schema):
+        if np_array.ndim == 1:
+            # 1D array: single column
+            column_names = [schema.names[0]] if schema.names else ["0"]
+            df = pd.DataFrame({column_names[0]: np_array})
+        elif np_array.ndim == 2:
+            # 2D array: multiple columns
+            column_names = (
+                schema.names
+                if len(schema.names) == np_array.shape[1]
+                else [f"{i}" for i in range(np_array.shape[1])]
+            )
+            df = pd.DataFrame(np_array, columns=column_names)
+        else:
+            raise ValueError(
+                f"NumPy arrays with {np_array.ndim} dimensions are not supported"
+            )
+    else:
+        # Fallback to generic column names
+        df = pd.DataFrame(np_array)
+
+    pd_utils.dataframe_to_file(
+        df,
         path,
-        file_system,
+        filesystem,
         block_path_provider,
         content_type,
-        **kwargs
+        **kwargs,
     )
+
+
+def concat_ndarrays(arrays: List[np.ndarray]) -> Optional[np.ndarray]:
+    """
+    Concatenate a list of NumPy ndarrays into a single ndarray.
+
+    Args:
+        arrays: List of NumPy ndarrays to concatenate
+
+    Returns:
+        Concatenated NumPy ndarray, or None if input is empty
+    """
+    if arrays is None or not len(arrays):
+        return None
+    if len(arrays) == 1:
+        return next(iter(arrays))
+    return np.concatenate(arrays, axis=0)
+
+
+def append_column_to_ndarray(
+    np_array: np.ndarray,
+    column_name: str,
+    column_value: Any,
+) -> np.ndarray:
+    # Add a new column with value repeating for each row of np_array
+    return np.concatenate((np_array, np.full((len(np_array), 1), column_value)), axis=1)