fugue 0.8.7.dev5__py3-none-any.whl → 0.8.7.dev7__py3-none-any.whl

fugue/__init__.py

@@ -1,6 +1,5 @@
 # flake8: noqa
 from triad.collections import Schema
-from triad.collections.fs import FileSystem
 
 from fugue.api import out_transform, transform
 from fugue.bag.array_bag import ArrayBag

fugue/_utils/io.py

@@ -1,14 +1,14 @@
 import os
 import pathlib
 from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union
-from urllib.parse import urlparse
 
-import fs as pfs
 import pandas as pd
+from fsspec import AbstractFileSystem
+from fsspec.implementations.local import LocalFileSystem
 from triad.collections.dict import ParamDict
-from triad.collections.fs import FileSystem
 from triad.collections.schema import Schema
 from triad.utils.assertion import assert_or_throw
+from triad.utils.io import join, url_to_fs
 from triad.utils.pandas_like import PD_UTILS
 
 from fugue.dataframe import LocalBoundedDataFrame, LocalDataFrame, PandasDataFrame
@@ -16,23 +16,14 @@ from fugue.dataframe import LocalBoundedDataFrame, LocalDataFrame, PandasDataFra
 
 class FileParser(object):
     def __init__(self, path: str, format_hint: Optional[str] = None):
-        last = len(path)
-        has_glob = False
         self._orig_format_hint = format_hint
-        for i in range(len(path)):
-            if path[i] in ["/", "\\"]:
-                last = i
-            if path[i] in ["*", "?"]:
-                has_glob = True
-                break
-        if not has_glob:
-            self._uri = urlparse(path)
-            self._glob_pattern = ""
-            self._path = self._uri.path
+        self._has_glob = "*" in path or "?" in path
+        self._raw_path = path
+        self._fs, self._fs_path = url_to_fs(path)
+        if not self.is_local:
+            self._path = self._fs.unstrip_protocol(self._fs_path)
         else:
-            self._uri = urlparse(path[:last])
-            self._glob_pattern = path[last + 1 :]
-            self._path = pfs.path.combine(self._uri.path, self._glob_pattern)
+            self._path = os.path.abspath(self._fs._strip_protocol(path))
 
         if format_hint is None or format_hint == "":
             for k, v in _FORMAT_MAP.items():
@@ -48,56 +39,64 @@ class FileParser(object):
             self._format = format_hint
 
     def assert_no_glob(self) -> "FileParser":
-        assert_or_throw(self.glob_pattern == "", f"{self.path} has glob pattern")
+        assert_or_throw(not self.has_glob, f"{self.raw_path} has glob pattern")
         return self
 
-    def with_glob(self, glob: str, format_hint: Optional[str] = None) -> "FileParser":
-        uri = self.uri
-        if glob != "":
-            uri = pfs.path.combine(uri, glob)
-        return FileParser(uri, format_hint or self._orig_format_hint)
-
     @property
-    def glob_pattern(self) -> str:
-        return self._glob_pattern
+    def has_glob(self):
+        return self._has_glob
 
     @property
-    def uri(self) -> str:
-        return self._uri.geturl()
+    def is_local(self):
+        return isinstance(self._fs, LocalFileSystem)
 
-    @property
-    def uri_with_glob(self) -> str:
-        if self.glob_pattern == "":
-            return self.uri
-        return pfs.path.combine(self.uri, self.glob_pattern)
+    def join(self, path: str, format_hint: Optional[str] = None) -> "FileParser":
+        if not self.has_glob:
+            _path = join(self.path, path)
+        else:
+            _path = join(self.parent, path)
+        return FileParser(_path, format_hint or self._orig_format_hint)
 
     @property
     def parent(self) -> str:
-        dn = os.path.dirname(self.uri)
-        return dn if dn != "" else "."
-
-    @property
-    def scheme(self) -> str:
-        return self._uri.scheme
+        return self._fs.unstrip_protocol(self._fs._parent(self._fs_path))
 
     @property
     def path(self) -> str:
         return self._path
 
+    @property
+    def raw_path(self) -> str:
+        return self._raw_path
+
     @property
     def suffix(self) -> str:
-        return "".join(pathlib.Path(self.path.lower()).suffixes)
+        return "".join(pathlib.Path(self.raw_path.lower()).suffixes)
 
     @property
     def file_format(self) -> str:
         return self._format
 
+    def make_parent_dirs(self) -> None:
+        self._fs.makedirs(self._fs._parent(self._fs_path), exist_ok=True)
+
+    def find_all(self) -> Iterable["FileParser"]:
+        if self.has_glob:
+            for x in self._fs.glob(self._fs_path):
+                yield FileParser(self._fs.unstrip_protocol(x))
+        else:
+            yield self
+
+    def open(self, *args: Any, **kwargs: Any) -> Any:
+        self.assert_no_glob()
+        return self._fs.open(self._fs_path, *args, **kwargs)
+
 
 def load_df(
     uri: Union[str, List[str]],
     format_hint: Optional[str] = None,
     columns: Any = None,
-    fs: Optional[FileSystem] = None,
+    fs: Optional[AbstractFileSystem] = None,
     **kwargs: Any,
 ) -> LocalBoundedDataFrame:
     if isinstance(uri, str):
@@ -117,7 +116,7 @@ def save_df(
     uri: str,
     format_hint: Optional[str] = None,
     mode: str = "overwrite",
-    fs: Optional[FileSystem] = None,
+    fs: Optional[AbstractFileSystem] = None,
     **kwargs: Any,
 ) -> None:
     assert_or_throw(
@@ -125,40 +124,28 @@ def save_df(
     )
     p = FileParser(uri, format_hint).assert_no_glob()
     if fs is None:
-        fs = FileSystem()
+        fs, _ = url_to_fs(uri)
     if fs.exists(uri):
         assert_or_throw(mode == "overwrite", FileExistsError(uri))
         try:
-            fs.remove(uri)
-        except Exception:
-            try:
-                fs.removetree(uri)
-            except Exception:  # pragma: no cover
-                pass
+            fs.rm(uri, recursive=True)
+        except Exception:  # pragma: no cover
+            pass
     _FORMAT_SAVE[p.file_format](df, p, **kwargs)
 
 
 def _get_single_files(
-    fp: Iterable[FileParser], fs: Optional[FileSystem]
+    fp: Iterable[FileParser], fs: Optional[AbstractFileSystem]
 ) -> Iterable[FileParser]:
-    if fs is None:
-        fs = FileSystem()
     for f in fp:
-        if f.glob_pattern != "":
-            files = [
-                FileParser(pfs.path.combine(f.uri, pfs.path.basename(x.path)))
-                for x in fs.opendir(f.uri).glob(f.glob_pattern)
-            ]
-            yield from _get_single_files(files, fs)
-        else:
-            yield f
+        yield from f.find_all()
 
 
 def _save_parquet(df: LocalDataFrame, p: FileParser, **kwargs: Any) -> None:
     PD_UTILS.to_parquet_friendly(
         df.as_pandas(), partition_cols=kwargs.get("partition_cols", [])
     ).to_parquet(
-        p.uri,
+        p.assert_no_glob().path,
         **{
             "engine": "pyarrow",
             "schema": df.schema.pa_schema,
@@ -171,34 +158,36 @@ def _load_parquet(
     p: FileParser, columns: Any = None, **kwargs: Any
 ) -> Tuple[pd.DataFrame, Any]:
     if columns is None:
-        pdf = pd.read_parquet(p.uri, **{"engine": "pyarrow", **kwargs})
+        pdf = pd.read_parquet(p.path, **{"engine": "pyarrow", **kwargs})
         return pdf, None
     if isinstance(columns, list):  # column names
-        pdf = pd.read_parquet(p.uri, columns=columns, **{"engine": "pyarrow", **kwargs})
+        pdf = pd.read_parquet(
+            p.path, columns=columns, **{"engine": "pyarrow", **kwargs}
+        )
         return pdf, None
     schema = Schema(columns)
     pdf = pd.read_parquet(
-        p.uri, columns=schema.names, **{"engine": "pyarrow", **kwargs}
+        p.path, columns=schema.names, **{"engine": "pyarrow", **kwargs}
     )
     return pdf, schema
 
 
 def _save_csv(df: LocalDataFrame, p: FileParser, **kwargs: Any) -> None:
-    df.as_pandas().to_csv(p.uri, **{"index": False, "header": False, **kwargs})
+    with p.open("w") as f:
+        df.as_pandas().to_csv(f, **{"index": False, "header": False, **kwargs})
 
 
-def _safe_load_csv(path: str, **kwargs: Any) -> pd.DataFrame:
+def _safe_load_csv(p: FileParser, **kwargs: Any) -> pd.DataFrame:
     def load_dir() -> pd.DataFrame:
-        fs = FileSystem()
-        return pd.concat(
-            [
-                pd.read_csv(pfs.path.combine(path, pfs.path.basename(x.path)), **kwargs)
-                for x in fs.opendir(path).glob("*.csv")
-            ]
-        )
+        dfs: List[pd.DataFrame] = []
+        for _p in p.join("*.csv").find_all():  # type: ignore
+            with _p.open("r") as f:
+                dfs.append(pd.read_csv(f, **kwargs))
+        return pd.concat(dfs)
 
     try:
-        return pd.read_csv(path, **kwargs)
+        with p.open("r") as f:
+            return pd.read_csv(f, **kwargs)
     except IsADirectoryError:
         return load_dir()
     except pd.errors.ParserError:  # pragma: no cover
@@ -224,7 +213,7 @@ def _load_csv(  # noqa: C901
         header = kw["header"]
         del kw["header"]
     if str(header) in ["True", "0"]:
-        pdf = _safe_load_csv(p.uri, **{"index_col": False, "header": 0, **kw})
+        pdf = _safe_load_csv(p, **{"index_col": False, "header": 0, **kw})
         if columns is None:
             return pdf, None
         if isinstance(columns, list):  # column names
@@ -236,12 +225,14 @@ def _load_csv(  # noqa: C901
             raise ValueError("columns must be set if without header")
         if isinstance(columns, list):  # column names
             pdf = _safe_load_csv(
-                p.uri, **{"index_col": False, "header": None, "names": columns, **kw}
+                p,
+                **{"index_col": False, "header": None, "names": columns, **kw},
             )
             return pdf, None
         schema = Schema(columns)
         pdf = _safe_load_csv(
-            p.uri, **{"index_col": False, "header": None, "names": schema.names, **kw}
+            p,
+            **{"index_col": False, "header": None, "names": schema.names, **kw},
         )
         return pdf, schema
     else:
@@ -249,27 +240,31 @@ def _load_csv(  # noqa: C901
 
 
 def _save_json(df: LocalDataFrame, p: FileParser, **kwargs: Any) -> None:
-    df.as_pandas().to_json(p.uri, **{"orient": "records", "lines": True, **kwargs})
+    with p.open("w") as f:
+        df.as_pandas().to_json(f, **{"orient": "records", "lines": True, **kwargs})
 
 
-def _safe_load_json(path: str, **kwargs: Any) -> pd.DataFrame:
+def _safe_load_json(p: FileParser, **kwargs: Any) -> pd.DataFrame:
     kw = {"orient": "records", "lines": True, **kwargs}
+
+    def load_dir() -> pd.DataFrame:
+        dfs: List[pd.DataFrame] = []
+        for _p in p.join("*.json").find_all():  # type: ignore
+            with _p.open("r") as f:
+                dfs.append(pd.read_json(f, **kw))
+        return pd.concat(dfs)
+
     try:
-        return pd.read_json(path, **kw)
+        with p.open("r") as f:
+            return pd.read_json(f, **kw)
     except (IsADirectoryError, PermissionError):
-        fs = FileSystem()
-        return pd.concat(
-            [
-                pd.read_json(pfs.path.combine(path, pfs.path.basename(x.path)), **kw)
-                for x in fs.opendir(path).glob("*.json")
-            ]
-        )
+        return load_dir()
 
 
 def _load_json(
     p: FileParser, columns: Any = None, **kwargs: Any
 ) -> Tuple[pd.DataFrame, Any]:
-    pdf = _safe_load_json(p.uri, **kwargs).reset_index(drop=True)
+    pdf = _safe_load_json(p, **kwargs).reset_index(drop=True)
     if columns is None:
         return pdf, None
     if isinstance(columns, list):  # column names

fugue/api.py

@@ -6,6 +6,7 @@ from .dataframe.api import (
     as_array_iterable,
     as_arrow,
     as_dict_iterable,
+    as_dicts,
     as_fugue_df,
     as_pandas,
     drop_columns,

fugue/dataframe/api.py

@@ -116,15 +116,32 @@ def as_array_iterable(
     return as_fugue_df(df).as_array_iterable(columns=columns, type_safe=type_safe)
 
 
+@fugue_plugin
+def as_dicts(
+    df: AnyDataFrame, columns: Optional[List[str]] = None
+) -> List[Dict[str, Any]]:
+    """Convert any dataframe to a list of python dicts
+
+    :param df: the object that can be recognized as a dataframe by Fugue
+    :param columns: columns to extract, defaults to None
+    :return: a list of python dicts
+
+    .. note::
+
+        The default implementation enforces ``type_safe`` True
+    """
+    return as_fugue_df(df).as_dicts(columns=columns)
+
+
 @fugue_plugin
 def as_dict_iterable(
     df: AnyDataFrame, columns: Optional[List[str]] = None
 ) -> Iterable[Dict[str, Any]]:
-    """Convert any dataframe to iterable of native python dicts
+    """Convert any dataframe to iterable of python dicts
 
     :param df: the object that can be recognized as a dataframe by Fugue
     :param columns: columns to extract, defaults to None
-    :return: iterable of native python dicts
+    :return: iterable of python dicts
 
     .. note::
 

fugue/dataframe/arrow_dataframe.py

@@ -21,6 +21,10 @@ from fugue.exceptions import FugueDataFrameOperationError
 
 from .api import (
     alter_columns,
+    as_array,
+    as_array_iterable,
+    as_dict_iterable,
+    as_dicts,
     as_pandas,
     drop_columns,
     get_column_names,
@@ -30,6 +34,12 @@ from .api import (
     select_columns,
 )
 from .dataframe import DataFrame, LocalBoundedDataFrame, _input_schema
+from .utils import (
+    pa_table_as_array,
+    pa_table_as_array_iterable,
+    pa_table_as_dict_iterable,
+    pa_table_as_dicts,
+)
 
 
 class ArrowDataFrame(LocalBoundedDataFrame):
@@ -174,21 +184,20 @@ class ArrowDataFrame(LocalBoundedDataFrame):
     def as_array(
         self, columns: Optional[List[str]] = None, type_safe: bool = False
     ) -> List[Any]:
-        return list(self.as_array_iterable(columns, type_safe=type_safe))
+        return pa_table_as_array(self.native, columns=columns)
+
+    def as_dicts(self, columns: Optional[List[str]] = None) -> List[Dict[str, Any]]:
+        return pa_table_as_dicts(self.native, columns=columns)
 
     def as_array_iterable(
         self, columns: Optional[List[str]] = None, type_safe: bool = False
     ) -> Iterable[Any]:
-        if self.empty:
-            return
-        if columns is not None:
-            for x in self[columns].as_array_iterable(type_safe=type_safe):
-                yield x
-        else:
-            d = self.native.to_pydict()
-            cols = [d[n] for n in self.columns]
-            for arr in zip(*cols):
-                yield list(arr)
+        yield from pa_table_as_array_iterable(self.native, columns=columns)
+
+    def as_dict_iterable(
+        self, columns: Optional[List[str]] = None
+    ) -> Iterable[Dict[str, Any]]:
+        yield from pa_table_as_dict_iterable(self.native, columns=columns)
 
 
 @as_local.candidate(lambda df: isinstance(df, pa.Table))
@@ -212,6 +221,34 @@ def _pa_table_as_pandas(df: pa.Table) -> pd.DataFrame:
     )
 
 
+@as_array.candidate(lambda df, *args, **kwargs: isinstance(df, pa.Table))
+def _pa_table_as_array(
+    df: pa.Table, columns: Optional[List[str]] = None, type_safe: bool = False
+) -> List[Any]:
+    return pa_table_as_array(df, columns=columns)
+
+
+@as_array_iterable.candidate(lambda df, *args, **kwargs: isinstance(df, pa.Table))
+def _pa_table_as_array_iterable(
+    df: pa.Table, columns: Optional[List[str]] = None, type_safe: bool = False
+) -> Iterable[Any]:
+    yield from pa_table_as_array_iterable(df, columns=columns)
+
+
+@as_dicts.candidate(lambda df, *args, **kwargs: isinstance(df, pa.Table))
+def _pa_table_as_dicts(
+    df: pa.Table, columns: Optional[List[str]] = None
+) -> List[Dict[str, Any]]:
+    return pa_table_as_dicts(df, columns=columns)
+
+
+@as_dict_iterable.candidate(lambda df, *args, **kwargs: isinstance(df, pa.Table))
+def _pa_table_as_dict_iterable(
+    df: pa.Table, columns: Optional[List[str]] = None
+) -> Iterable[Dict[str, Any]]:
+    yield from pa_table_as_dict_iterable(df, columns=columns)
+
+
 @alter_columns.candidate(lambda df, *args, **kwargs: isinstance(df, pa.Table))
 def _pa_table_alter_columns(
     df: pa.Table, columns: Any, as_fugue: bool = False

fugue/dataframe/dataframe.py

@@ -237,13 +237,31 @@ class DataFrame(Dataset):
         """
         raise NotImplementedError
 
+    def as_dicts(self, columns: Optional[List[str]] = None) -> List[Dict[str, Any]]:
+        """Convert to a list of python dicts
+
+        :param columns: columns to extract, defaults to None
+        :return: a list of python dicts
+
+        .. note::
+
+            The default implementation enforces ``type_safe`` True
+        """
+        if columns is None:
+            columns = self.columns
+        idx = range(len(columns))
+        return [
+            {columns[i]: x[i] for i in idx}
+            for x in self.as_array(columns, type_safe=True)
+        ]
+
     def as_dict_iterable(
         self, columns: Optional[List[str]] = None
     ) -> Iterable[Dict[str, Any]]:
-        """Convert to iterable of native python dicts
+        """Convert to iterable of python dicts
 
         :param columns: columns to extract, defaults to None
-        :return: iterable of native python dicts
+        :return: iterable of python dicts
 
         .. note::
 

fugue/dataframe/function_wrapper.py

@@ -269,7 +269,7 @@ class _EmptyAwareIterableListParam(_LocalNoSchemaDataFrameParam):
 class _ListDictParam(_LocalNoSchemaDataFrameParam):
     @no_type_check
     def to_input_data(self, df: DataFrame, ctx: Any) -> List[Dict[str, Any]]:
-        return list(df.as_local().as_dict_iterable())
+        return df.as_local().as_dicts()
 
     @no_type_check
     def to_output_df(

fugue/dataframe/iterable_dataframe.py

@@ -105,6 +105,9 @@ class IterableDataFrame(LocalUnboundedDataFrame):
     ) -> List[Any]:
         return list(self.as_array_iterable(columns, type_safe=type_safe))
 
+    def as_dicts(self, columns: Optional[List[str]] = None) -> List[Dict[str, Any]]:
+        return list(self.as_dict_iterable(columns))
+
     def as_array_iterable(
         self, columns: Optional[List[str]] = None, type_safe: bool = False
     ) -> Iterable[Any]:

fugue/dataframe/pandas_dataframe.py

@@ -1,8 +1,11 @@
 from typing import Any, Dict, Iterable, List, Optional, Tuple
 
 import pandas as pd
+import pyarrow as pa
+from triad import assert_or_throw
 from triad.collections.schema import Schema
 from triad.utils.pandas_like import PD_UTILS
+from triad.utils.pyarrow import pa_batch_to_dicts
 
 from fugue.dataset.api import (
     as_fugue_dataset,
@@ -17,6 +20,10 @@ from fugue.dataset.api import (
 from fugue.exceptions import FugueDataFrameOperationError
 
 from .api import (
+    as_array,
+    as_array_iterable,
+    as_dict_iterable,
+    as_dicts,
     drop_columns,
     get_column_names,
     get_schema,
@@ -134,6 +141,9 @@ class PandasDataFrame(LocalBoundedDataFrame):
             return self
         return PandasDataFrame(self.native, new_schema)
 
+    def as_arrow(self, type_safe: bool = False) -> pa.Table:
+        return PD_UTILS.as_arrow(self.native, schema=self.schema.pa_schema)
+
     def as_array(
         self, columns: Optional[List[str]] = None, type_safe: bool = False
     ) -> List[Any]:
@@ -150,6 +160,18 @@ class PandasDataFrame(LocalBoundedDataFrame):
         ):
             yield row
 
+    def as_dicts(self, columns: Optional[List[str]] = None) -> List[Dict[str, Any]]:
+        res: List[Dict[str, Any]] = []
+        for block in _to_dicts(self.native, columns, self.schema):
+            res += block
+        return res
+
+    def as_dict_iterable(
+        self, columns: Optional[List[str]] = None
+    ) -> Iterable[Dict[str, Any]]:
+        for block in _to_dicts(self.native, columns, self.schema):
+            yield from block
+
     def head(
         self, n: int, columns: Optional[List[str]] = None
     ) -> LocalBoundedDataFrame:
@@ -272,6 +294,43 @@ def _pd_head(
     return _adjust_df(df.head(n), as_fugue=as_fugue)
 
 
+@as_array.candidate(lambda df, *args, **kwargs: isinstance(df, pd.DataFrame))
+def _pd_as_array(
+    df: pd.DataFrame, columns: Optional[List[str]] = None, type_safe: bool = False
+) -> List[Any]:
+    return list(_pd_as_array_iterable(df, columns, type_safe=type_safe))
+
+
+@as_array_iterable.candidate(lambda df, *args, **kwargs: isinstance(df, pd.DataFrame))
+def _pd_as_array_iterable(
+    df: pd.DataFrame, columns: Optional[List[str]] = None, type_safe: bool = False
+) -> Iterable[Any]:
+    for row in PD_UTILS.as_array_iterable(
+        df,
+        columns=columns,
+        type_safe=type_safe,
+    ):
+        yield row
+
+
+@as_dicts.candidate(lambda df, *args, **kwargs: isinstance(df, pd.DataFrame))
+def _pd_as_dicts(
+    df: pd.DataFrame, columns: Optional[List[str]] = None
+) -> List[Dict[str, Any]]:
+    res: List[Dict[str, Any]] = []
+    for block in _to_dicts(df, columns):
+        res += block
+    return res
+
+
+@as_dict_iterable.candidate(lambda df, *args, **kwargs: isinstance(df, pd.DataFrame))
+def _pd_as_dict_iterable(
+    df: pa.Table, columns: Optional[List[str]] = None
+) -> Iterable[Dict[str, Any]]:
+    for block in _to_dicts(df, columns):
+        yield from block
+
+
 def _adjust_df(res: pd.DataFrame, as_fugue: bool):
     return res if not as_fugue else PandasDataFrame(res)
 
@@ -280,3 +339,17 @@ def _assert_no_missing(df: pd.DataFrame, columns: Iterable[Any]) -> None:
     missing = [x for x in columns if x not in df.columns]
     if len(missing) > 0:
         raise FugueDataFrameOperationError("found nonexistent columns: {missing}")
+
+
+def _to_dicts(
+    df: pd.DataFrame,
+    columns: Optional[List[str]] = None,
+    schema: Optional[Schema] = None,
+) -> Iterable[List[Dict[str, Any]]]:
+    cols = list(df.columns) if columns is None else columns
+    assert_or_throw(len(cols) > 0, ValueError("columns cannot be empty"))
+    pa_schema = schema.extract(cols).pa_schema if schema is not None else None
+    adf = PD_UTILS.as_arrow(df[cols], schema=pa_schema)
+    for batch in adf.to_batches():
+        if batch.num_rows > 0:
+            yield pa_batch_to_dicts(batch)