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

fugue_spark/dataframe.py

@@ -14,9 +14,14 @@ from fugue.dataframe import (
     IterableDataFrame,
     LocalBoundedDataFrame,
 )
+from fugue.dataframe.utils import pa_table_as_array, pa_table_as_dicts
 from fugue.exceptions import FugueDataFrameOperationError
 from fugue.plugins import (
+    as_array,
+    as_array_iterable,
     as_arrow,
+    as_dict_iterable,
+    as_dicts,
     as_local_bounded,
     as_pandas,
     count,
@@ -152,23 +157,22 @@ class SparkDataFrame(DataFrame):
     def as_array(
         self, columns: Optional[List[str]] = None, type_safe: bool = False
     ) -> List[Any]:
-        sdf = self._select_columns(columns)
-        return sdf.as_local().as_array(type_safe=type_safe)
+        return _spark_as_array(self.native, columns=columns, type_safe=type_safe)
 
     def as_array_iterable(
         self, columns: Optional[List[str]] = None, type_safe: bool = False
     ) -> Iterable[Any]:
-        if is_spark_connect(self.native):  # pragma: no cover
-            yield from self.as_array(columns, type_safe=type_safe)
-            return
-        sdf = self._select_columns(columns)
-        if not type_safe:
-            for row in to_type_safe_input(sdf.native.rdd.toLocalIterator(), sdf.schema):
-                yield row
-        else:
-            df = IterableDataFrame(sdf.as_array_iterable(type_safe=False), sdf.schema)
-            for row in df.as_array_iterable(type_safe=True):
-                yield row
+        yield from _spark_as_array_iterable(
+            self.native, columns=columns, type_safe=type_safe
+        )
+
+    def as_dicts(self, columns: Optional[List[str]] = None) -> List[Dict[str, Any]]:
+        return _spark_as_dicts(self.native, columns=columns)
+
+    def as_dict_iterable(
+        self, columns: Optional[List[str]] = None
+    ) -> Iterable[Dict[str, Any]]:
+        yield from _spark_as_dict_iterable(self.native, columns=columns)
 
     def head(
         self, n: int, columns: Optional[List[str]] = None
@@ -288,6 +292,58 @@ def _spark_df_head(
     return SparkDataFrame(res).as_local() if as_fugue else to_pandas(res)
 
 
+@as_array.candidate(lambda df, *args, **kwargs: is_spark_dataframe(df))
+def _spark_as_array(
+    df: ps.DataFrame, columns: Optional[List[str]] = None, type_safe: bool = False
+) -> List[Any]:
+    assert_or_throw(columns is None or len(columns) > 0, ValueError("empty columns"))
+    _df = df if columns is None or len(columns) == 0 else df[columns]
+    return pa_table_as_array(to_arrow(_df), columns)
+
+
+@as_array_iterable.candidate(lambda df, *args, **kwargs: is_spark_dataframe(df))
+def _spark_as_array_iterable(
+    df: ps.DataFrame, columns: Optional[List[str]] = None, type_safe: bool = False
+) -> Iterable[Any]:
+    if is_spark_connect(df):  # pragma: no cover
+        yield from _spark_as_array(df, columns, type_safe=type_safe)
+    else:
+        assert_or_throw(
+            columns is None or len(columns) > 0, ValueError("empty columns")
+        )
+        _df = df if columns is None or len(columns) == 0 else df[columns]
+        if not type_safe:
+            for row in to_type_safe_input(
+                _df.rdd.toLocalIterator(), to_schema(_df.schema)
+            ):
+                yield list(row)
+        else:
+            tdf = IterableDataFrame(
+                _spark_as_array_iterable(_df, type_safe=False), to_schema(_df.schema)
+            )
+            yield from tdf.as_array_iterable(type_safe=True)
+
+
+@as_dicts.candidate(lambda df, *args, **kwargs: is_spark_dataframe(df))
+def _spark_as_dicts(
+    df: ps.DataFrame, columns: Optional[List[str]] = None, type_safe: bool = False
+) -> List[Dict[str, Any]]:
+    assert_or_throw(columns is None or len(columns) > 0, ValueError("empty columns"))
+    _df = df if columns is None or len(columns) == 0 else df[columns]
+    return pa_table_as_dicts(to_arrow(_df), columns)
+
+
+@as_dict_iterable.candidate(lambda df, *args, **kwargs: is_spark_dataframe(df))
+def _spark_as_dict_iterable(
+    df: ps.DataFrame, columns: Optional[List[str]] = None, type_safe: bool = False
+) -> Iterable[Dict[str, Any]]:
+    assert_or_throw(columns is None or len(columns) > 0, ValueError("empty columns"))
+    _df = df if columns is None or len(columns) == 0 else df[columns]
+    cols = list(_df.columns)
+    for row in _spark_as_array_iterable(_df, type_safe=type_safe):
+        yield dict(zip(cols, row))
+
+
 def _rename_spark_dataframe(df: ps.DataFrame, names: Dict[str, Any]) -> ps.DataFrame:
     cols: List[ps.Column] = []
     for f in df.schema:

fugue_spark/execution_engine.py

@@ -11,7 +11,7 @@ from pyspark.rdd import RDD
 from pyspark.sql import SparkSession
 from pyspark.sql.functions import broadcast, col, lit, row_number
 from pyspark.sql.window import Window
-from triad import FileSystem, IndexedOrderedDict, ParamDict, Schema, SerializableRLock
+from triad import IndexedOrderedDict, ParamDict, Schema, SerializableRLock
 from triad.utils.assertion import assert_arg_not_none, assert_or_throw
 from triad.utils.hash import to_uuid
 from triad.utils.iter import EmptyAwareIterable
@@ -360,13 +360,12 @@ class SparkExecutionEngine(ExecutionEngine):
         cf.update(ParamDict(conf))
         super().__init__(cf)
         self._lock = SerializableRLock()
-        self._fs = FileSystem()
         self._log = logging.getLogger()
         self._broadcast_func = RunOnce(
             self._broadcast, lambda *args, **kwargs: id(args[0])
         )
         self._persist_func = RunOnce(self._persist, lambda *args, **kwargs: id(args[0]))
-        self._io = SparkIO(self.spark_session, self.fs)
+        self._io = SparkIO(self.spark_session)
         self._registered_dfs: Dict[str, SparkDataFrame] = {}
 
     def __repr__(self) -> str:
@@ -395,10 +394,6 @@ class SparkExecutionEngine(ExecutionEngine):
     def log(self) -> logging.Logger:
         return self._log
 
-    @property
-    def fs(self) -> FileSystem:
-        return self._fs
-
     def create_default_sql_engine(self) -> SQLEngine:
         return SparkSQLEngine(self)
 

fugue_test/builtin_suite.py

@@ -12,11 +12,12 @@ import pickle
 from typing import Any, Callable, Dict, Iterable, Iterator, List, Optional
 from unittest import TestCase
 from uuid import uuid4
-
+from triad.utils.io import write_text, join
 import numpy as np
 import pandas as pd
 import pyarrow as pa
 import pytest
+from fsspec.implementations.local import LocalFileSystem
 from pytest import raises
 from triad import SerializableRLock
 
@@ -28,7 +29,6 @@ from fugue import (
     DataFrame,
     DataFrames,
     ExecutionEngine,
-    FileSystem,
     FugueWorkflow,
     LocalDataFrame,
     OutputCoTransformer,
@@ -65,6 +65,8 @@ from fugue.exceptions import (
     FugueWorkflowRuntimeValidationError,
 )
 
+_LOCAL_FS = LocalFileSystem(auto_mkdir=True)
+
 
 class BuiltInTests(object):
     """Workflow level general test suite. It is a more general end to end
@@ -633,9 +635,8 @@ class BuiltInTests(object):
             tmpdir = str(self.tmpdir)
 
             def incr():
-                fs = FileSystem(auto_close=False).makedirs(tmpdir, recreate=True)
-                fs.writetext(str(uuid4()) + ".txt", "")
-                return fs.glob("*.txt").count().files
+                write_text(join(tmpdir, str(uuid4()) + ".txt"), "")
+                return len(_LOCAL_FS.glob(join(tmpdir, "*.txt")))
 
             def t1(df: Iterable[Dict[str, Any]]) -> Iterable[Dict[str, Any]]:
                 for row in df:
@@ -717,9 +718,8 @@ class BuiltInTests(object):
             tmpdir = str(self.tmpdir)
 
             def incr():
-                fs = FileSystem(auto_close=False).makedirs(tmpdir, recreate=True)
-                fs.writetext(str(uuid4()) + ".txt", "")
-                return fs.glob("*.tx" "t").count().files
+                write_text(join(tmpdir, str(uuid4()) + ".txt"), "")
+                return len(_LOCAL_FS.glob(join(tmpdir, "*.txt")))
 
             def t1(
                 df: Iterable[Dict[str, Any]], df2: pd.DataFrame
@@ -1348,7 +1348,7 @@ class BuiltInTests(object):
                 b.partition(num=3).save(path, fmt="parquet", single=True)
                 b.save(path2, header=True)
             dag.run(self.engine)
-            assert FileSystem().isfile(path)
+            assert _LOCAL_FS.isfile(path)
             with FugueWorkflow() as dag:
                 a = dag.load(path, fmt="parquet", columns=["a", "c"])
                 a.assert_eq(dag.df([[1, 6], [7, 2]], "a:long,c:int"))
@@ -1359,9 +1359,9 @@ class BuiltInTests(object):
                 b = dag.df([[6, 1], [2, 7]], "c:int,a:long")
                 b.partition(by="c").save(path3, fmt="parquet", single=False)
             dag.run(self.engine)
-            assert FileSystem().isdir(path3)
-            assert FileSystem().isdir(os.path.join(path3, "c=6"))
-            assert FileSystem().isdir(os.path.join(path3, "c=2"))
+            assert _LOCAL_FS.isdir(path3)
+            assert _LOCAL_FS.isdir(os.path.join(path3, "c=6"))
+            assert _LOCAL_FS.isdir(os.path.join(path3, "c=2"))
             # TODO: in test below, once issue #288 is fixed, use dag.load
             #  instead of pd.read_parquet
             pdf = pd.read_parquet(path3).sort_values("a").reset_index(drop=True)

fugue_test/dataframe_suite.py

@@ -208,8 +208,22 @@ class DataFrameTests(object):
         def test_as_dict_iterable(self):
             df = self.df([[pd.NaT, 1]], "a:datetime,b:int")
             assert [dict(a=None, b=1)] == list(fi.as_dict_iterable(df))
+            df = self.df([[pd.NaT, 1]], "a:datetime,b:int")
+            assert [dict(b=1)] == list(fi.as_dict_iterable(df, ["b"]))
             df = self.df([[pd.Timestamp("2020-01-01"), 1]], "a:datetime,b:int")
             assert [dict(a=datetime(2020, 1, 1), b=1)] == list(fi.as_dict_iterable(df))
+            df = self.df([[pd.Timestamp("2020-01-01"), 1]], "a:datetime,b:int")
+            assert [dict(b=1)] == list(fi.as_dict_iterable(df, ["b"]))
+
+        def test_as_dicts(self):
+            df = self.df([[pd.NaT, 1]], "a:datetime,b:int")
+            assert [dict(a=None, b=1)] == fi.as_dicts(df)
+            df = self.df([[pd.NaT, 1]], "a:datetime,b:int")
+            assert [dict(b=1)] == fi.as_dicts(df, ["b"])
+            df = self.df([[pd.Timestamp("2020-01-01"), 1]], "a:datetime,b:int")
+            assert [dict(a=datetime(2020, 1, 1), b=1)] == fi.as_dicts(df)
+            df = self.df([[pd.Timestamp("2020-01-01"), 1]], "a:datetime,b:int")
+            assert [dict(b=1)] == fi.as_dicts(df, ["b"])
 
         def test_list_type(self):
             data = [[[30, 40]]]

fugue_test/execution_suite.py

@@ -15,8 +15,8 @@ from unittest import TestCase
 import pandas as pd
 import pytest
 from pytest import raises
-from triad.collections.fs import FileSystem
 from triad.exceptions import InvalidOperationError
+from triad.utils.io import isfile, makedirs, touch
 
 import fugue.api as fa
 import fugue.column.functions as ff
@@ -62,7 +62,6 @@ class ExecutionEngineTests(object):
         def test_init(self):
             print(self.engine)
             assert self.engine.log is not None
-            assert self.engine.fs is not None
             assert copy.copy(self.engine) is self.engine
             assert copy.deepcopy(self.engine) is self.engine
 
@@ -985,17 +984,16 @@ class ExecutionEngineTests(object):
             df_eq(res, [[1, "z1"]], "a:int,v:str", throw=True)
 
         @pytest.fixture(autouse=True)
-        def init_tmpdir(self, tmpdir):
+        def init_tmpdir(self, tmpdir, tmp_mem_dir):
             self.tmpdir = tmpdir
 
         def test_save_single_and_load_parquet(self):
-            e = self.engine
             b = ArrayDataFrame([[6, 1], [2, 7]], "c:int,a:long")
             path = os.path.join(self.tmpdir, "a", "b")
-            e.fs.makedirs(path, recreate=True)
+            makedirs(path, exist_ok=True)
             # over write folder with single file
             fa.save(b, path, format_hint="parquet", force_single=True)
-            assert e.fs.isfile(path)
+            assert isfile(path)
             c = fa.load(path, format_hint="parquet", columns=["a", "c"], as_fugue=True)
             df_eq(c, [[1, 6], [7, 2]], "a:long,c:int", throw=True)
 
@@ -1019,7 +1017,7 @@ class ExecutionEngineTests(object):
             path = os.path.join(self.tmpdir, "a", "b")
             fa.save(a, os.path.join(path, "a.parquet"), engine=native)
             fa.save(b, os.path.join(path, "b.parquet"), engine=native)
-            FileSystem().touch(os.path.join(path, "_SUCCESS"))
+            touch(os.path.join(path, "_SUCCESS"))
             c = fa.load(path, format_hint="parquet", columns=["a", "c"], as_fugue=True)
             df_eq(c, [[1, 6], [7, 2], [8, 4]], "a:long,c:int", throw=True)
 
@@ -1038,13 +1036,12 @@ class ExecutionEngineTests(object):
             df_eq(c, [[1, 6], [7, 2], [8, 4]], "a:long,c:int", throw=True)
 
         def test_save_single_and_load_csv(self):
-            e = self.engine
             b = ArrayDataFrame([[6.1, 1.1], [2.1, 7.1]], "c:double,a:double")
             path = os.path.join(self.tmpdir, "a", "b")
-            e.fs.makedirs(path, recreate=True)
+            makedirs(path, exist_ok=True)
             # over write folder with single file
             fa.save(b, path, format_hint="csv", header=True, force_single=True)
-            assert e.fs.isfile(path)
+            assert isfile(path)
             c = fa.load(
                 path, format_hint="csv", header=True, infer_schema=False, as_fugue=True
             )
@@ -1099,13 +1096,12 @@ class ExecutionEngineTests(object):
             df_eq(c, [["1.1", "60.1"], ["7.1", "20.1"]], "a:str,c:str", throw=True)
 
         def test_save_single_and_load_csv_no_header(self):
-            e = self.engine
             b = ArrayDataFrame([[6.1, 1.1], [2.1, 7.1]], "c:double,a:double")
             path = os.path.join(self.tmpdir, "a", "b")
-            e.fs.makedirs(path, recreate=True)
+            makedirs(path, exist_ok=True)
             # over write folder with single file
             fa.save(b, path, format_hint="csv", header=False, force_single=True)
-            assert e.fs.isfile(path)
+            assert isfile(path)
 
             with raises(ValueError):
                 c = fa.load(
@@ -1190,7 +1186,7 @@ class ExecutionEngineTests(object):
                 header=True,
                 engine=native,
             )
-            FileSystem().touch(os.path.join(path, "_SUCCESS"))
+            touch(os.path.join(path, "_SUCCESS"))
             c = fa.load(
                 path,
                 format_hint="csv",
@@ -1204,13 +1200,12 @@ class ExecutionEngineTests(object):
             )
 
         def test_save_single_and_load_json(self):
-            e = self.engine
             b = ArrayDataFrame([[6, 1], [2, 7]], "c:int,a:long")
             path = os.path.join(self.tmpdir, "a", "b")
-            e.fs.makedirs(path, recreate=True)
+            makedirs(path, exist_ok=True)
             # over write folder with single file
             fa.save(b, path, format_hint="json", force_single=True)
-            assert e.fs.isfile(path)
+            assert isfile(path)
             c = fa.load(path, format_hint="json", columns=["a", "c"], as_fugue=True)
             df_eq(c, [[1, 6], [7, 2]], "a:long,c:long", throw=True)
 
@@ -1241,7 +1236,7 @@ class ExecutionEngineTests(object):
             path = os.path.join(self.tmpdir, "a", "b")
             fa.save(a, os.path.join(path, "a.json"), format_hint="json", engine=native)
             fa.save(b, os.path.join(path, "b.json"), format_hint="json", engine=native)
-            FileSystem().touch(os.path.join(path, "_SUCCESS"))
+            touch(os.path.join(path, "_SUCCESS"))
             c = fa.load(path, format_hint="json", columns=["a", "c"], as_fugue=True)
             df_eq(c, [[1, 6], [7, 2], [8, 4], [4, 3]], "a:long,c:long", throw=True)
 

fugue_test/plugins/misc/__init__.py

@@ -0,0 +1,2 @@
+# flake8: noqa
+from .fixtures import tmp_mem_dir

fugue_test/plugins/misc/fixtures.py

@@ -0,0 +1,18 @@
+import uuid
+
+import pytest
+from triad.utils.io import makedirs, rm
+
+
+@pytest.fixture
+def tmp_mem_dir():
+    uuid_str = str(uuid.uuid4())[:5]
+    path = "memory://test_" + uuid_str
+    makedirs(path)
+    try:
+        yield path
+    finally:
+        try:
+            rm(path, recursive=True)
+        except Exception:  # pragma: no cover
+            pass