fugue 0.8.7.dev7__py3-none-any.whl → 0.9.0__py3-none-any.whl

fugue_ray/_utils/io.py

@@ -4,13 +4,14 @@ from typing import Any, Callable, Dict, Iterable, List, Optional, Union
 
 import pyarrow as pa
 import ray.data as rd
+from packaging import version
 from pyarrow import csv as pacsv
 from pyarrow import json as pajson
 from ray.data.datasource import FileExtensionFilter
 from triad.collections import Schema
 from triad.collections.dict import ParamDict
 from triad.utils.assertion import assert_or_throw
-from triad.utils.io import exists, makedirs, rm
+from triad.utils.io import exists, makedirs, rm, isfile
 
 from fugue import ExecutionEngine
 from fugue._utils.io import FileParser, save_df
@@ -18,6 +19,8 @@ from fugue.collections.partition import PartitionSpec
 from fugue.dataframe import DataFrame
 from fugue_ray.dataframe import RayDataFrame
 
+from .._constants import RAY_VERSION
+
 
 class RayIO(object):
     def __init__(self, engine: ExecutionEngine):
@@ -148,6 +151,18 @@ class RayIO(object):
         if infer_schema and columns is not None and not isinstance(columns, list):
             raise ValueError("can't set columns as a schema when infer schema is true")
 
+        if RAY_VERSION >= version.parse("2.10"):
+            if len(p) == 1 and isfile(p[0]):  # TODO: very hacky
+                params: Dict[str, Any] = {}
+            else:
+                params = {"file_extensions": ["csv"]}
+        else:  # pragma: no cover
+            params = {
+                "partition_filter": _FileFiler(
+                    file_extensions=["csv"], exclude=["_SUCCESS"]
+                ),
+            }
+
         def _read_csv(to_str: bool) -> RayDataFrame:
             res = rd.read_csv(
                 p,
@@ -155,9 +170,7 @@ class RayIO(object):
                 read_options=pacsv.ReadOptions(**read_options),
                 parse_options=pacsv.ParseOptions(**parse_options),
                 convert_options=pacsv.ConvertOptions(**convert_options),
-                partition_filter=_FileFiler(
-                    file_extensions=["csv"], exclude=["_SUCCESS"]
-                ),
+                **params,
             )
             if to_str:
                 _schema = res.schema(fetch_if_missing=True)
@@ -195,16 +208,32 @@ class RayIO(object):
         read_options: Dict[str, Any] = {"use_threads": False}
         parse_options: Dict[str, Any] = {}
 
-        def _read_json() -> RayDataFrame:
+        def _read_json() -> RayDataFrame:  # pragma: no cover
+            if RAY_VERSION >= version.parse("2.10"):
+                if len(p) == 1 and isfile(p[0]):  # TODO: very hacky
+                    params: Dict[str, Any] = {"file_extensions": None}
+                else:
+                    params = {"file_extensions": ["json"]}
+            elif RAY_VERSION >= version.parse("2.9"):  # pragma: no cover
+                params = {
+                    "file_extensions": None,
+                    "partition_filter": _FileFiler(
+                        file_extensions=["json"], exclude=["_SUCCESS"]
+                    ),
+                }
+            else:  # pragma: no cover
+                params = {
+                    "partition_filter": _FileFiler(
+                        file_extensions=["json"], exclude=["_SUCCESS"]
+                    ),
+                }
             return RayDataFrame(
                 rd.read_json(
                     p,
                     ray_remote_args=self._remote_args(),
                     read_options=pajson.ReadOptions(**read_options),
                     parse_options=pajson.ParseOptions(**parse_options),
-                    partition_filter=_FileFiler(
-                        file_extensions=["json"], exclude=["_SUCCESS"]
-                    ),
+                    **params,
                 )
             )
 
@@ -221,7 +250,7 @@ class RayIO(object):
         return {"num_cpus": 1}
 
 
-class _FileFiler(FileExtensionFilter):
+class _FileFiler(FileExtensionFilter):  # pragma: no cover
     def __init__(self, file_extensions: Union[str, List[str]], exclude: Iterable[str]):
         super().__init__(file_extensions, allow_if_no_extension=True)
         self._exclude = set(exclude)

fugue_ray/execution_engine.py

@@ -191,8 +191,7 @@ class RayMapEngine(MapEngine):
             mb_args["batch_size"] = self.conf.get_or_throw(
                 FUGUE_RAY_DEFAULT_BATCH_SIZE, int
             )
-        if ray.__version__ >= "2.3":
-            mb_args["zero_copy_batch"] = self.conf.get(FUGUE_RAY_ZERO_COPY, True)
+        mb_args["zero_copy_batch"] = self.conf.get(FUGUE_RAY_ZERO_COPY, True)
         sdf = rdf.native.map_batches(
             _udf,
             batch_format="pyarrow",

fugue_ray/registry.py

@@ -14,6 +14,7 @@ from fugue.plugins import as_fugue_dataset, infer_execution_engine
 
 from .dataframe import RayDataFrame
 from .execution_engine import RayExecutionEngine
+from .tester import RayTestBackend  # noqa: F401  # pylint: disable-all
 
 
 @infer_execution_engine.candidate(

fugue_ray/tester.py

@@ -0,0 +1,22 @@
+from contextlib import contextmanager
+from typing import Any, Dict, Iterator
+
+import ray
+
+import fugue.test as ft
+
+
+@ft.fugue_test_backend
+class RayTestBackend(ft.FugueTestBackend):
+    name = "ray"
+    default_session_conf = {"num_cpus": 2}
+    default_fugue_conf = {
+        "fugue.ray.zero_copy": True,
+        "fugue.ray.default.batch_size": 10000,
+    }
+
+    @classmethod
+    @contextmanager
+    def session_context(cls, session_conf: Dict[str, Any]) -> Iterator[Any]:
+        with ray.init(**session_conf):
+            yield "ray"

fugue_spark/execution_engine.py

@@ -492,11 +492,6 @@ class SparkExecutionEngine(ExecutionEngine):
         res.reset_metadata(df.metadata)
         return res
 
-    def register(self, df: DataFrame, name: str) -> SparkDataFrame:
-        sdf = self._to_spark_df(df)
-        sdf.native.createOrReplaceTempView(name)
-        return sdf
-
     def join(
         self,
         df1: DataFrame,
@@ -679,6 +674,11 @@ class SparkExecutionEngine(ExecutionEngine):
 
         # If partition exists
         else:
+            if len(_presort.keys()) == 0 and n == 1:
+                return self._to_spark_df(
+                    d.dropDuplicates(subset=partition_spec.partition_by), df.schema
+                )
+
             w = Window.partitionBy([col(x) for x in partition_spec.partition_by])
 
             if len(_presort.keys()) > 0:

fugue_spark/registry.py

@@ -19,6 +19,13 @@ from fugue_spark.dataframe import SparkDataFrame
 from fugue_spark.execution_engine import SparkExecutionEngine
 
 from ._utils.misc import SparkConnectDataFrame, SparkConnectSession, is_spark_dataframe
+from .tester import SparkTestBackend  # noqa: F401  # pylint: disable-all
+
+try:
+    from .tester import SparkConnectTestBackend  # noqa: F401  # pylint: disable-all
+except ImportError:  # pragma: no cover
+    pass
+
 
 _is_sparksql = namespace_candidate("sparksql", lambda x: isinstance(x, str))
 
@@ -31,7 +38,12 @@ _is_sparksql = namespace_candidate("sparksql", lambda x: isinstance(x, str))
     )
     or any(_is_sparksql(obj) for obj in objs)
 )
-def _infer_spark_client(obj: Any) -> Any:
+def _infer_spark_client(objs: Any) -> Any:
+    obj = objs[0]
+    if isinstance(obj, SparkDataFrame):
+        obj = obj.native
+    if hasattr(obj, "sparkSession"):
+        return obj.sparkSession
     return SparkSession.builder.getOrCreate()
 
 

fugue_spark/tester.py

@@ -0,0 +1,78 @@
+from contextlib import contextmanager
+from typing import Any, Dict, Iterator
+
+from pyspark.sql import SparkSession
+
+import fugue.test as ft
+
+from ._utils.misc import SparkConnectSession
+
+
+@ft.fugue_test_backend
+class SparkTestBackend(ft.FugueTestBackend):
+    name = "spark"
+    default_session_conf = {
+        "spark.app.name": "fugue-test-spark",
+        "spark.master": "local[*]",
+        "spark.default.parallelism": 4,
+        "spark.dynamicAllocation.enabled": "false",
+        "spark.executor.cores": 4,
+        "spark.executor.instances": 1,
+        "spark.io.compression.codec": "lz4",
+        "spark.rdd.compress": "false",
+        "spark.sql.shuffle.partitions": 4,
+        "spark.shuffle.compress": "false",
+        "spark.sql.catalogImplementation": "in-memory",
+        "spark.sql.execution.arrow.pyspark.enabled": True,
+        "spark.sql.adaptive.enabled": False,
+    }
+
+    @classmethod
+    def transform_session_conf(cls, conf: Dict[str, Any]) -> Dict[str, Any]:
+        return ft.extract_conf(conf, "spark.", remove_prefix=False)
+
+    @classmethod
+    @contextmanager
+    def session_context(cls, session_conf: Dict[str, Any]) -> Iterator[Any]:
+        with _create_session(session_conf).getOrCreate() as spark:
+            yield spark
+
+
+if SparkConnectSession is not None:
+
+    @ft.fugue_test_backend
+    class SparkConnectTestBackend(SparkTestBackend):
+        name = "sparkconnect"
+        default_session_conf = {
+            "spark.default.parallelism": 4,
+            "spark.sql.shuffle.partitions": 4,
+            "spark.sql.execution.arrow.pyspark.enabled": True,
+            "spark.sql.adaptive.enabled": False,
+        }
+
+        @classmethod
+        def transform_session_conf(
+            cls, conf: Dict[str, Any]
+        ) -> Dict[str, Any]:  # pragma: no cover
+            # replace sparkconnect. with spark.
+            return {
+                "spark." + k: v
+                for k, v in ft.extract_conf(
+                    conf, cls.name + ".", remove_prefix=True
+                ).items()
+            }
+
+        @classmethod
+        @contextmanager
+        def session_context(
+            cls, session_conf: Dict[str, Any]
+        ) -> Iterator[Any]:  # pragma: no cover
+            spark = _create_session(session_conf).remote("sc://localhost").getOrCreate()
+            yield spark
+
+
+def _create_session(conf: Dict[str, Any]) -> Any:
+    sb = SparkSession.builder
+    for k, v in conf.items():
+        sb = sb.config(k, v)
+    return sb

fugue_test/__init__.py

@@ -0,0 +1,82 @@
+from typing import Any, Dict, Tuple
+
+import pyarrow as pa
+import pytest
+from triad.utils.pyarrow import to_pa_datatype
+
+_FUGUE_TEST_CONF_NAME = "fugue_test_conf"
+
+
+def pytest_addoption(parser: Any):  # pragma: no cover
+    parser.addini(
+        _FUGUE_TEST_CONF_NAME,
+        help="Configs for fugue testing execution engines",
+        type="linelist",
+    )
+
+
+def pytest_configure(config: Any):
+    from fugue.test.plugins import _set_global_conf
+
+    options = config.getini(_FUGUE_TEST_CONF_NAME)
+    conf: Dict[str, Any] = {}
+    if options:
+        for line in options:
+            line = line.strip()
+            if not line.startswith("#"):
+                k, v = _parse_line(line)
+                conf[k] = v
+    _set_global_conf(conf)
+
+
+def pytest_report_header(config, start_path):
+    from fugue.test.plugins import _get_all_ini_conf
+
+    header_lines = []
+    header_lines.append("Fugue tests will be initialized with options:")
+    for k, v in _get_all_ini_conf().items():
+        header_lines.append(f"\t{k} = {v}")
+    return "\n".join(header_lines)
+
+
+def _parse_line(line: str) -> Tuple[str, Any]:
+    try:
+        kv = line.split("=", 1)
+        if len(kv) == 1:
+            raise ValueError()
+        kt = kv[0].split(":", 1)
+        if len(kt) == 1:
+            tp = pa.string()
+        else:
+            tp = to_pa_datatype(kt[1].strip())
+        key = kt[0].strip()
+        if key == "":
+            raise ValueError()
+        value = pa.compute.cast([kv[1].strip()], tp).to_pylist()[0]
+        return key, value
+    except Exception:
+        raise ValueError(
+            f"Invalid config line: {line}, it must be in format: key[:type]=value"
+        )
+
+
+@pytest.fixture(scope="class")
+def backend_context(request: Any):
+    from fugue.test.plugins import _make_backend_context, _parse_backend
+
+    c, _ = _parse_backend(request.param)
+    session = request.getfixturevalue(c + "_session")
+    with _make_backend_context(request.param, session) as ctx:
+        yield ctx
+
+
+@pytest.fixture(scope="class")
+def _class_backend_context(request, backend_context):
+    from fugue.test.plugins import FugueTestContext
+
+    request.cls._test_context = FugueTestContext(
+        engine=backend_context.engine,
+        session=backend_context.session,
+        name=backend_context.name,
+    )
+    yield

fugue_test/builtin_suite.py

@@ -10,9 +10,8 @@ import datetime
 import os
 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
@@ -20,8 +19,10 @@ import pytest
 from fsspec.implementations.local import LocalFileSystem
 from pytest import raises
 from triad import SerializableRLock
+from triad.utils.io import join, write_text
 
 import fugue.api as fa
+import fugue.test as ft
 from fugue import (
     AnyDataFrame,
     ArrayDataFrame,
@@ -46,7 +47,6 @@ from fugue import (
     outputter,
     processor,
     register_creator,
-    register_default_sql_engine,
     register_output_transformer,
     register_outputter,
     register_processor,
@@ -56,7 +56,6 @@ from fugue import (
 from fugue.column import col
 from fugue.column import functions as ff
 from fugue.column import lit
-from fugue.dataframe.utils import _df_eq as df_eq
 from fugue.exceptions import (
     FugueInterfacelessError,
     FugueWorkflowCompileError,
@@ -78,26 +77,10 @@ class BuiltInTests(object):
     add correspondent tests here
     """
 
-    class Tests(TestCase):
-        @classmethod
-        def setUpClass(cls):
-            register_default_sql_engine(lambda engine: engine.sql_engine)
-            cls._engine = cls.make_engine(cls)
-
-        @property
-        def engine(self) -> ExecutionEngine:
-            return self._engine  # type: ignore
-
-        @classmethod
-        def tearDownClass(cls):
-            cls._engine.stop()
-
-        def make_engine(self) -> ExecutionEngine:  # pragma: no cover
-            raise NotImplementedError
-
+    class Tests(ft.FugueTestSuite):
         def test_workflows(self):
             a = FugueWorkflow().df([[0]], "a:int")
-            df_eq(a.compute(self.engine), [[0]], "a:int")
+            self.df_eq(a.compute(self.engine), [[0]], "a:int")
 
         def test_create_show(self):
             with FugueWorkflow() as dag:
@@ -1706,7 +1689,7 @@ class BuiltInTests(object):
                 """,
                     x=sdf3,
                 ).run()
-                df_eq(
+                self.df_eq(
                     res["res"],
                     [[3, 4, 13]],
                     schema="a:long,b:int,c:long",
@@ -1739,9 +1722,9 @@ class BuiltInTests(object):
                 df1 = pd.DataFrame([[0, 1], [2, 3]], columns=["a b", " "])
                 df2 = pd.DataFrame([[0, 10], [20, 3]], columns=["a b", "d"])
                 r = fa.inner_join(df1, df2, as_fugue=True)
-                df_eq(r, [[0, 1, 10]], "`a b`:long,` `:long,d:long", throw=True)
+                self.df_eq(r, [[0, 1, 10]], "`a b`:long,` `:long,d:long", throw=True)
                 r = fa.transform(r, tr)
-                df_eq(
+                self.df_eq(
                     r,
                     [[0, 1, 10, 2]],
                     "`a b`:long,` `:long,d:long,`c *`:long",
@@ -1755,7 +1738,7 @@ class BuiltInTests(object):
                     col("d"),
                     col("c *").cast(int),
                 )
-                df_eq(
+                self.df_eq(
                     r,
                     [[0, 1, 10, 2]],
                     "`a b `:long,`x y`:long,d:long,`c *`:long",
@@ -1764,13 +1747,13 @@ class BuiltInTests(object):
                 r = fa.rename(r, {"a b ": "a b"})
                 fa.save(r, f_csv, header=True, force_single=True)
                 fa.save(r, f_parquet)
-                df_eq(
+                self.df_eq(
                     fa.load(f_parquet, columns=["x y", "d", "c *"], as_fugue=True),
                     [[1, 10, 2]],
                     "`x y`:long,d:long,`c *`:long",
                     throw=True,
                 )
-                df_eq(
+                self.df_eq(
                     fa.load(
                         f_csv,
                         header=True,
@@ -1782,7 +1765,7 @@ class BuiltInTests(object):
                     "d:str,`c *`:str",
                     throw=True,
                 )
-                df_eq(
+                self.df_eq(
                     fa.load(
                         f_csv,
                         header=True,
@@ -1802,14 +1785,14 @@ class BuiltInTests(object):
                 """,
                     as_fugue=True,
                 )
-                df_eq(r, [[0, 1, 10]], "`a b`:long,` `:long,d:long", throw=True)
+                self.df_eq(r, [[0, 1, 10]], "`a b`:long,` `:long,d:long", throw=True)
                 r = fa.fugue_sql(
                     """
                 TRANSFORM r USING tr SCHEMA *,`c *`:long
                 """,
                     as_fugue=True,
                 )
-                df_eq(
+                self.df_eq(
                     r,
                     [[0, 1, 10, 2]],
                     "`a b`:long,` `:long,d:long,`c *`:long",
@@ -1821,7 +1804,7 @@ class BuiltInTests(object):
                 """,
                     as_fugue=True,
                 )
-                df_eq(
+                self.df_eq(
                     r,
                     [[0, 1, 10, 2]],
                     "`a b`:long,` `:long,d:long,`c *`:long",
@@ -1842,19 +1825,19 @@ class BuiltInTests(object):
                     f_parquet=f_parquet,
                     f_csv=f_csv,
                 ).run()
-                df_eq(
+                self.df_eq(
                     res["r1"],
                     [[1, 10, 2]],
                     "`x y`:long,d:long,`c *`:long",
                     throw=True,
                 )
-                df_eq(
+                self.df_eq(
                     res["r2"],
                     [["1", "10", "2"]],
                     "`x y`:str,d:str,`c *`:str",
                     throw=True,
                 )
-                df_eq(
+                self.df_eq(
                     res["r3"],
                     [[0, 1, 10, 2]],
                     "`a b`:long,`x y`:long,d:long,`c *`:long",
@@ -1875,13 +1858,13 @@ def mock_processor(df1: List[List[Any]], df2: List[List[Any]]) -> DataFrame:
 
 
 def mock_processor2(e: ExecutionEngine, dfs: DataFrames) -> DataFrame:
-    assert "test" in e.conf
+    assert "fugue.test" in e.conf
     return ArrayDataFrame([[sum(s.count() for s in dfs.values())]], "a:int")
 
 
 class MockProcessor3(Processor):
     def process(self, dfs):
-        assert "test" in self.workflow_conf
+        assert "fugue.test" in self.workflow_conf
         return ArrayDataFrame([[sum(s.count() for s in dfs.values())]], "a:int")
 
 
@@ -1915,11 +1898,11 @@ class MockOutputter4(Outputter):
 
 class MockTransform1(Transformer):
     def get_output_schema(self, df: DataFrame) -> Any:
-        assert "test" in self.workflow_conf
+        assert "fugue.test" in self.workflow_conf
         return [df.schema, "ct:int,p:int"]
 
     def on_init(self, df: DataFrame) -> None:
-        assert "test" in self.workflow_conf
+        assert "fugue.test" in self.workflow_conf
         self.pn = self.cursor.physical_partition_no
         self.ks = self.key_schema
         if "on_init_called" not in self.__dict__:
@@ -1929,7 +1912,7 @@ class MockTransform1(Transformer):
 
     def transform(self, df: LocalDataFrame) -> LocalDataFrame:
         assert 1 == self.on_init_called
-        assert "test" in self.workflow_conf
+        assert "fugue.test" in self.workflow_conf
         pdf = df.as_pandas()
         pdf["p"] = self.params.get("p", 1)
         pdf["ct"] = pdf.shape[0]
@@ -1971,7 +1954,7 @@ def mock_tf3(df: Iterable[Dict[str, Any]]) -> Iterable[Dict[str, Any]]:
 
 class MockCoTransform1(CoTransformer):
     def get_output_schema(self, dfs: DataFrames) -> Any:
-        assert "test" in self.workflow_conf
+        assert "fugue.test" in self.workflow_conf
         assert 2 == len(dfs)
         if self.params.get("named", False):
             assert dfs.has_key
@@ -1980,7 +1963,7 @@ class MockCoTransform1(CoTransformer):
         return [self.key_schema, "ct1:int,ct2:int,p:int"]
 
     def on_init(self, dfs: DataFrames) -> None:
-        assert "test" in self.workflow_conf
+        assert "fugue.test" in self.workflow_conf
         assert 2 == len(dfs)
         if self.params.get("named", False):
             assert dfs.has_key
@@ -1995,7 +1978,7 @@ class MockCoTransform1(CoTransformer):
 
     def transform(self, dfs: DataFrames) -> LocalDataFrame:
         assert 1 == self.on_init_called
-        assert "test" in self.workflow_conf
+        assert "fugue.test" in self.workflow_conf
         assert 2 == len(dfs)
         if self.params.get("named", False):
             assert dfs.has_key

fugue_test/dataframe_suite.py

@@ -2,15 +2,14 @@
 
 from datetime import date, datetime
 from typing import Any
-from unittest import TestCase
 
 import numpy as np
 import pandas as pd
 from pytest import raises
 
 import fugue.api as fi
+import fugue.test as ft
 from fugue.dataframe import ArrowDataFrame, DataFrame
-from fugue.dataframe.utils import _df_eq as df_eq
 from fugue.exceptions import FugueDataFrameOperationError, FugueDatasetEmptyError
 
 
@@ -19,15 +18,7 @@ class DataFrameTests(object):
     All new DataFrame types should pass this test suite.
     """
 
-    class Tests(TestCase):
-        @classmethod
-        def setUpClass(cls):
-            pass
-
-        @classmethod
-        def tearDownClass(cls):
-            pass
-
+    class Tests(ft.FugueTestSuite):
         def df(self, data: Any = None, schema: Any = None) -> Any:  # pragma: no cover
             raise NotImplementedError
 
@@ -129,7 +120,7 @@ class DataFrameTests(object):
             assert [[1]] == fi.as_array(df, type_safe=True)
 
             df = self.df([["a", 1, 2]], "a:str,b:int,c:int")
-            df_eq(
+            self.df_eq(
                 fi.as_fugue_df(fi.select_columns(df, ["c", "a"])),
                 [[2, "a"]],
                 "a:str,c:int",
@@ -140,13 +131,13 @@ class DataFrameTests(object):
                 df = self.df(data, "a:str,b:int")
                 df2 = fi.rename(df, columns=dict(a="aa"))
                 assert fi.get_schema(df) == "a:str,b:int"
-                df_eq(fi.as_fugue_df(df2), data, "aa:str,b:int", throw=True)
+                self.df_eq(fi.as_fugue_df(df2), data, "aa:str,b:int", throw=True)
 
             for data in [[["a", 1]], []]:
                 df = self.df(data, "a:str,b:int")
                 df3 = fi.rename(df, columns={})
                 assert fi.get_schema(df3) == "a:str,b:int"
-                df_eq(fi.as_fugue_df(df3), data, "a:str,b:int", throw=True)
+                self.df_eq(fi.as_fugue_df(df3), data, "a:str,b:int", throw=True)
 
         def test_rename_invalid(self):
             df = self.df([["a", 1]], "a:str,b:int")