deltacat 0.1.10.dev0__py3-none-any.whl → 0.1.12__py3-none-any.whl

deltacat/utils/numpy.py

@@ -1,21 +1,17 @@
+from typing import List, Optional
+
 import numpy as np
 import pyarrow as pa
-
 from fsspec import AbstractFileSystem
-
 from ray.data.datasource import BlockWritePathProvider
 
 from deltacat.types.media import ContentType
-from deltacat.utils import pyarrow as pa_utils
 from deltacat.utils import pandas as pd_utils
+from deltacat.utils import pyarrow as pa_utils
 from deltacat.utils.common import ReadKwargsProvider
 
-from typing import List, Optional
-
 
-def slice_ndarray(
-        np_array: np.ndarray,
-        max_len: Optional[int]) -> List[np.ndarray]:
+def slice_ndarray(np_array: np.ndarray, max_len: Optional[int]) -> List[np.ndarray]:
     """
     Iteratively creates max_len slices from the first dimension of an ndarray.
     """
@@ -23,17 +19,18 @@ def slice_ndarray(
         return [np_array]
 
     # Slice along the first dimension of the ndarray.
-    return [np_array[i:i + max_len] for i in range(0, len(np_array), max_len)]
+    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,
-        content_type: str,
-        content_encoding: str,
-        column_names: Optional[List[str]] = None,
-        include_columns: Optional[List[str]] = None,
-        pd_read_func_kwargs_provider: Optional[ReadKwargsProvider] = None,
-        **s3_client_kwargs) -> np.ndarray:
+    s3_url: str,
+    content_type: str,
+    content_encoding: str,
+    column_names: Optional[List[str]] = None,
+    include_columns: Optional[List[str]] = None,
+    pd_read_func_kwargs_provider: Optional[ReadKwargsProvider] = None,
+    **s3_client_kwargs
+) -> np.ndarray:
     # TODO: Compare perf to s3 -> pyarrow -> pandas [Series/DataFrame] -> numpy
     dataframe = pd_utils.s3_file_to_dataframe(
         s3_url,
@@ -52,12 +49,13 @@ def ndarray_size(np_array: np.ndarray) -> int:
 
 
 def ndarray_to_file(
-        np_array: np.ndarray,
-        path: str,
-        file_system: AbstractFileSystem,
-        block_path_provider: BlockWritePathProvider,
-        content_type: str = ContentType.PARQUET.value,
-        **kwargs) -> None:
+    np_array: np.ndarray,
+    path: str,
+    file_system: AbstractFileSystem,
+    block_path_provider: BlockWritePathProvider,
+    content_type: str = ContentType.PARQUET.value,
+    **kwargs
+) -> None:
     """
     Writes the given Numpy ndarray to a file.
     """

deltacat/utils/pandas.py

@@ -1,24 +1,25 @@
-import pandas as pd
 import csv
-import math
 import io
 import logging
-import pyarrow as pa
+import math
+from typing import Any, Callable, Dict, Iterable, List, Optional
 
+import pandas as pd
+import pyarrow as pa
 from fsspec import AbstractFileSystem
-
 from ray.data.datasource import BlockWritePathProvider
 
-from deltacat.types.media import ContentType, ContentEncoding, \
-    EXPLICIT_COMPRESSION_CONTENT_TYPES
-from deltacat.types.media import DELIMITED_TEXT_CONTENT_TYPES, \
-    TABULAR_CONTENT_TYPES
 from deltacat import logs
-from deltacat.utils.common import ReadKwargsProvider, ContentTypeKwargsProvider
+from deltacat.types.media import (
+    DELIMITED_TEXT_CONTENT_TYPES,
+    EXPLICIT_COMPRESSION_CONTENT_TYPES,
+    TABULAR_CONTENT_TYPES,
+    ContentEncoding,
+    ContentType,
+)
+from deltacat.utils.common import ContentTypeKwargsProvider, ReadKwargsProvider
 from deltacat.utils.performance import timed_invocation
 
-from typing import Any, Callable, Dict, Iterable, List, Optional
-
 logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
 
 
@@ -30,7 +31,7 @@ CONTENT_TYPE_TO_PD_READ_FUNC: Dict[str, Callable] = {
     ContentType.PARQUET.value: pd.read_parquet,
     ContentType.FEATHER.value: pd.read_feather,
     ContentType.ORC.value: pd.read_orc,
-    ContentType.JSON.value: pd.read_json
+    ContentType.JSON.value: pd.read_json,
 }
 
 
@@ -39,24 +40,21 @@ class ReadKwargsProviderPandasCsvPureUtf8(ContentTypeKwargsProvider):
     as UTF-8 strings (i.e. disables type inference). Useful for ensuring
     lossless reads of UTF-8 delimited text datasets and improving read
     performance in cases where type casting is not required."""
+
     def __init__(self, include_columns: Optional[Iterable[str]] = None):
         self.include_columns = include_columns
 
-    def _get_kwargs(
-            self,
-            content_type: str,
-            kwargs: Dict[str, Any]) -> Dict[str, Any]:
+    def _get_kwargs(self, content_type: str, kwargs: Dict[str, Any]) -> Dict[str, Any]:
         if content_type in DELIMITED_TEXT_CONTENT_TYPES:
-            include_columns = self.include_columns \
-                if self.include_columns else kwargs.get("usecols")
+            include_columns = (
+                self.include_columns if self.include_columns else kwargs.get("usecols")
+            )
             if not include_columns:
                 # read all columns as strings
                 kwargs["dtype"] = str
             else:
                 # read only the included columns as strings
-                kwargs["dtype"] = {
-                    column_name: str for column_name in include_columns
-                }
+                kwargs["dtype"] = {column_name: str for column_name in include_columns}
             # use the fastest available engine for pure utf-8 reads
             kwargs["engine"] = "pyarrow"
         return kwargs
@@ -71,24 +69,17 @@ def content_type_to_reader_kwargs(content_type: str) -> Dict[str, Any]:
             "keep_default_na": False,
         }
     if content_type == ContentType.TSV.value:
-        return {
-            "sep": "\t",
-            "header": None
-        }
+        return {"sep": "\t", "header": None}
     if content_type == ContentType.CSV.value:
-        return {
-            "sep": ",",
-            "header": None
-        }
+        return {"sep": ",", "header": None}
     if content_type == ContentType.PSV.value:
-        return {
-            "sep": "|",
-            "header": None
-        }
-    if content_type in {ContentType.PARQUET.value,
-                        ContentType.FEATHER.value,
-                        ContentType.ORC.value,
-                        ContentType.JSON.value}:
+        return {"sep": "|", "header": None}
+    if content_type in {
+        ContentType.PARQUET.value,
+        ContentType.FEATHER.value,
+        ContentType.ORC.value,
+        ContentType.JSON.value,
+    }:
         return {}
     raise ValueError(f"Unsupported content type: {content_type}")
 
@@ -96,13 +87,13 @@ def content_type_to_reader_kwargs(content_type: str) -> Dict[str, Any]:
 ENCODING_TO_PD_COMPRESSION: Dict[str, str] = {
     ContentEncoding.GZIP.value: "gzip",
     ContentEncoding.BZIP2.value: "bz2",
-    ContentEncoding.IDENTITY.value: "none"
+    ContentEncoding.IDENTITY.value: "none",
 }
 
 
 def slice_dataframe(
-        dataframe: pd.DataFrame,
-        max_len: Optional[int]) -> List[pd.DataFrame]:
+    dataframe: pd.DataFrame, max_len: Optional[int]
+) -> List[pd.DataFrame]:
     """
     Iteratively create dataframe slices.
     """
@@ -111,12 +102,11 @@ def slice_dataframe(
     dataframes = []
     num_slices = math.ceil(len(dataframe) / max_len)
     for i in range(num_slices):
-        dataframes.append(dataframe[i * max_len: (i + 1) * max_len])
+        dataframes.append(dataframe[i * max_len : (i + 1) * max_len])
     return dataframes
 
 
-def concat_dataframes(dataframes: List[pd.DataFrame]) \
-        -> Optional[pd.DataFrame]:
+def concat_dataframes(dataframes: List[pd.DataFrame]) -> Optional[pd.DataFrame]:
     if dataframes is None or not len(dataframes):
         return None
     if len(dataframes) == 1:
@@ -125,10 +115,11 @@ def concat_dataframes(dataframes: List[pd.DataFrame]) \
 
 
 def _add_column_kwargs(
-        content_type: str,
-        column_names: Optional[List[str]],
-        include_columns: Optional[List[str]],
-        kwargs: Dict[str, Any]):
+    content_type: str,
+    column_names: Optional[List[str]],
+    include_columns: Optional[List[str]],
+    kwargs: Dict[str, Any],
+):
 
     if content_type in DELIMITED_TEXT_CONTENT_TYPES:
         kwargs["names"] = column_names
@@ -140,25 +131,27 @@ def _add_column_kwargs(
             if include_columns:
                 logger.warning(
                     f"Ignoring request to include columns {include_columns} "
-                    f"for non-tabular content type {content_type}")
+                    f"for non-tabular content type {content_type}"
+                )
 
 
 def s3_file_to_dataframe(
-        s3_url: str,
-        content_type: str,
-        content_encoding: str,
-        column_names: Optional[List[str]] = None,
-        include_columns: Optional[List[str]] = None,
-        pd_read_func_kwargs_provider: Optional[ReadKwargsProvider] = None,
-        **s3_client_kwargs) -> pd.DataFrame:
+    s3_url: str,
+    content_type: str,
+    content_encoding: str,
+    column_names: Optional[List[str]] = None,
+    include_columns: Optional[List[str]] = None,
+    pd_read_func_kwargs_provider: Optional[ReadKwargsProvider] = None,
+    **s3_client_kwargs,
+) -> pd.DataFrame:
 
     from deltacat.aws import s3u as s3_utils
-    logger.debug(f"Reading {s3_url} to Pandas. Content type: {content_type}. "
-                 f"Encoding: {content_encoding}")
-    s3_obj = s3_utils.get_object_at_url(
-        s3_url,
-        **s3_client_kwargs
+
+    logger.debug(
+        f"Reading {s3_url} to Pandas. Content type: {content_type}. "
+        f"Encoding: {content_encoding}"
     )
+    s3_obj = s3_utils.get_object_at_url(s3_url, **s3_client_kwargs)
     logger.debug(f"Read S3 object from {s3_url}: {s3_obj}")
     pd_read_func = CONTENT_TYPE_TO_PD_READ_FUNC[content_type]
     args = [io.BytesIO(s3_obj["Body"].read())]
@@ -167,17 +160,12 @@ def s3_file_to_dataframe(
 
     if content_type in EXPLICIT_COMPRESSION_CONTENT_TYPES:
         kwargs["compression"] = ENCODING_TO_PD_COMPRESSION.get(
-            content_encoding,
-            "infer"
+            content_encoding, "infer"
         )
     if pd_read_func_kwargs_provider:
         kwargs = pd_read_func_kwargs_provider(content_type, kwargs)
     logger.debug(f"Reading {s3_url} via {pd_read_func} with kwargs: {kwargs}")
-    dataframe, latency = timed_invocation(
-        pd_read_func,
-        *args,
-        **kwargs
-    )
+    dataframe, latency = timed_invocation(pd_read_func, *args, **kwargs)
     logger.debug(f"Time to read {s3_url} into Pandas Dataframe: {latency}s")
     return dataframe
 
@@ -188,11 +176,8 @@ def dataframe_size(dataframe: pd.DataFrame) -> int:
 
 
 def write_csv(
-        dataframe: pd.DataFrame,
-        path: str,
-        *,
-        filesystem: AbstractFileSystem,
-        **kwargs) -> None:
+    dataframe: pd.DataFrame, path: str, *, filesystem: AbstractFileSystem, **kwargs
+) -> None:
     with filesystem.open(path, "wb") as f:
         # TODO (pdames): Add support for client-specified compression types.
         with pa.CompressedOutputStream(f, ContentEncoding.GZIP.value) as out:
@@ -200,31 +185,22 @@ def write_csv(
 
 
 def write_parquet(
-        dataframe: pd.DataFrame,
-        path: str,
-        *,
-        filesystem: AbstractFileSystem,
-        **kwargs) -> None:
+    dataframe: pd.DataFrame, path: str, *, filesystem: AbstractFileSystem, **kwargs
+) -> None:
     with filesystem.open(path, "wb") as f:
         dataframe.to_parquet(f, **kwargs)
 
 
 def write_feather(
-        dataframe: pd.DataFrame,
-        path: str,
-        *,
-        filesystem: AbstractFileSystem,
-        **kwargs) -> None:
+    dataframe: pd.DataFrame, path: str, *, filesystem: AbstractFileSystem, **kwargs
+) -> None:
     with filesystem.open(path, "wb") as f:
         dataframe.to_feather(f, **kwargs)
 
 
 def write_json(
-        dataframe: pd.DataFrame,
-        path: str,
-        *,
-        filesystem: AbstractFileSystem,
-        **kwargs) -> None:
+    dataframe: pd.DataFrame, path: str, *, filesystem: AbstractFileSystem, **kwargs
+) -> None:
     with filesystem.open(path, "wb") as f:
         # TODO (pdames): Add support for client-specified compression types.
         with pa.CompressedOutputStream(f, ContentEncoding.GZIP.value) as out:
@@ -283,12 +259,13 @@ def content_type_to_writer_kwargs(content_type: str) -> Dict[str, Any]:
 
 
 def dataframe_to_file(
-        dataframe: pd.DataFrame,
-        base_path: str,
-        file_system: AbstractFileSystem,
-        block_path_provider: BlockWritePathProvider,
-        content_type: str = ContentType.PARQUET.value,
-        **kwargs) -> None:
+    dataframe: pd.DataFrame,
+    base_path: str,
+    file_system: AbstractFileSystem,
+    block_path_provider: BlockWritePathProvider,
+    content_type: str = ContentType.PARQUET.value,
+    **kwargs,
+) -> None:
     """
     Writes the given Pandas Dataframe to a file.
     """
@@ -299,11 +276,7 @@ def dataframe_to_file(
         raise NotImplementedError(
             f"Pandas writer for content type '{content_type}' not "
             f"implemented. Known content types: "
-            f"{CONTENT_TYPE_TO_PD_WRITE_FUNC.keys}")
+            f"{CONTENT_TYPE_TO_PD_WRITE_FUNC.keys}"
+        )
     path = block_path_provider(base_path)
-    writer(
-        dataframe,
-        path,
-        filesystem=file_system,
-        **writer_kwargs
-    )
+    writer(dataframe, path, filesystem=file_system, **writer_kwargs)

deltacat/utils/performance.py

@@ -4,11 +4,8 @@ from typing import Any, Callable, Tuple
 
 
 def invoke_with_perf_counter(
-        counter: Counter,
-        counter_key: Any,
-        func: Callable,
-        *args,
-        **kwargs) -> Tuple[Any, float]:
+    counter: Counter, counter_key: Any, func: Callable, *args, **kwargs
+) -> Tuple[Any, float]:
 
     start = time.perf_counter()
     result = func(*args, **kwargs)
@@ -18,10 +15,7 @@ def invoke_with_perf_counter(
     return result, latency
 
 
-def timed_invocation(
-        func: Callable,
-        *args,
-        **kwargs) -> Tuple[Any, float]:
+def timed_invocation(func: Callable, *args, **kwargs) -> Tuple[Any, float]:
 
     start = time.perf_counter()
     result = func(*args, **kwargs)