fugue 0.8.2.dev4__py3-none-any.whl → 0.8.4__py3-none-any.whl

fugue_spark/execution_engine.py

@@ -4,7 +4,6 @@ from uuid import uuid4
 
 import pandas as pd
 import pyarrow as pa
-import pyspark
 import pyspark.sql as ps
 from pyspark import StorageLevel
 from pyspark.rdd import RDD
@@ -25,7 +24,7 @@ from fugue.collections.partition import (
     PartitionSpec,
     parse_presort_exp,
 )
-from fugue.constants import KEYWORD_ROWCOUNT
+from fugue.constants import KEYWORD_PARALLELISM, KEYWORD_ROWCOUNT
 from fugue.dataframe import (
     ArrayDataFrame,
     ArrowDataFrame,
@@ -42,18 +41,13 @@ from fugue.dataframe.arrow_dataframe import _build_empty_arrow
 from fugue.dataframe.utils import get_join_schemas
 from fugue.exceptions import FugueDataFrameInitError
 from fugue.execution.execution_engine import ExecutionEngine, MapEngine, SQLEngine
-from fugue_spark._constants import (
-    FUGUE_SPARK_CONF_USE_PANDAS_UDF,
-    FUGUE_SPARK_DEFAULT_CONF,
-)
-from fugue_spark._utils.convert import to_schema, to_spark_schema, to_type_safe_input
-from fugue_spark._utils.io import SparkIO
-from fugue_spark._utils.partition import (
-    even_repartition,
-    hash_repartition,
-    rand_repartition,
-)
-from fugue_spark.dataframe import SparkDataFrame
+
+from ._constants import FUGUE_SPARK_CONF_USE_PANDAS_UDF, FUGUE_SPARK_DEFAULT_CONF
+from ._utils.convert import to_schema, to_spark_schema, to_type_safe_input
+from ._utils.io import SparkIO
+from ._utils.misc import is_spark_connect as _is_spark_connect, is_spark_dataframe
+from ._utils.partition import even_repartition, hash_repartition, rand_repartition
+from .dataframe import SparkDataFrame
 
 _TO_SPARK_JOIN_MAP: Dict[str, str] = {
     "inner": "inner",
@@ -103,12 +97,15 @@ class SparkMapEngine(MapEngine):
     def is_distributed(self) -> bool:
         return True
 
+    @property
+    def is_spark_connect(self) -> bool:
+        """Whether the spark session is created by spark connect"""
+        return self.execution_engine.is_spark_connect  # type:ignore
+
     def _should_use_pandas_udf(self, schema: Schema) -> bool:
+        if self.is_spark_connect:  # pragma: no cover
+            return True
         possible = hasattr(ps.DataFrame, "mapInPandas")  # must be new version of Spark
-        if pyspark.__version__ < "3":  # pragma: no cover
-            possible &= self.execution_engine.conf.get(
-                "spark.sql.execution.arrow.enabled", False
-            )
         # else:  # this condition seems to be unnecessary
         #    possible &= self.execution_engine.conf.get(
         #        "spark.sql.execution.arrow.pyspark.enabled", False
@@ -138,15 +135,25 @@ class SparkMapEngine(MapEngine):
         output_schema = Schema(output_schema)
         if self._should_use_pandas_udf(output_schema):
             # pandas udf can only be used for pyspark > 3
-            if len(partition_spec.partition_by) > 0 and partition_spec.algo != "even":
-                return self._group_map_by_pandas_udf(
-                    df,
-                    map_func=map_func,
-                    output_schema=output_schema,
-                    partition_spec=partition_spec,
-                    on_init=on_init,
-                    map_func_format_hint=map_func_format_hint,
-                )
+            if len(partition_spec.partition_by) > 0:
+                if partition_spec.algo == "coarse":
+                    return self._map_by_pandas_udf(
+                        df,
+                        map_func=map_func,
+                        output_schema=output_schema,
+                        partition_spec=partition_spec,
+                        on_init=on_init,
+                        map_func_format_hint=map_func_format_hint,
+                    )
+                elif partition_spec.algo != "even" or self.is_spark_connect:
+                    return self._group_map_by_pandas_udf(
+                        df,
+                        map_func=map_func,
+                        output_schema=output_schema,
+                        partition_spec=partition_spec,
+                        on_init=on_init,
+                        map_func_format_hint=map_func_format_hint,
+                    )
             elif len(partition_spec.partition_by) == 0:
                 return self._map_by_pandas_udf(
                     df,
@@ -187,7 +194,7 @@ class SparkMapEngine(MapEngine):
         def _udf_pandas(pdf: Any) -> pd.DataFrame:  # pragma: no cover
             if pdf.shape[0] == 0:
                 return PandasDataFrame([], output_schema).as_pandas()
-            if len(presort_keys) > 0:
+            if len(partition_spec.presort) > 0:
                 pdf = pdf.sort_values(presort_keys, ascending=presort_asc)
             input_df = PandasDataFrame(
                 pdf.reset_index(drop=True), input_schema, pandas_df_wrapper=True
@@ -239,6 +246,7 @@ class SparkMapEngine(MapEngine):
                         )
                         if not cursor_set:
                             cursor.set(lambda: pdf.peek_array(), 0, 0)
+                            cursor_set = True
                         yield pdf
 
             input_df = IterablePandasDataFrame(get_dfs(), input_schema)
@@ -273,6 +281,7 @@ class SparkMapEngine(MapEngine):
                         pdf = ArrowDataFrame(func(adf))
                         if not cursor_set:
                             cursor.set(lambda: pdf.peek_array(), 0, 0)
+                            cursor_set = True
                         yield pdf
 
             input_df = IterableArrowDataFrame(get_dfs(), input_schema)
@@ -316,7 +325,10 @@ class SparkExecutionEngine(ExecutionEngine):
             spark_session = SparkSession.builder.getOrCreate()
         self._spark_session = spark_session
         cf = dict(FUGUE_SPARK_DEFAULT_CONF)
-        cf.update({x[0]: x[1] for x in spark_session.sparkContext.getConf().getAll()})
+        if not self.is_spark_connect:
+            cf.update(
+                {x[0]: x[1] for x in spark_session.sparkContext.getConf().getAll()}
+            )
         cf.update(ParamDict(conf))
         super().__init__(cf)
         self._lock = SerializableRLock()
@@ -343,6 +355,10 @@ class SparkExecutionEngine(ExecutionEngine):
         )
         return self._spark_session
 
+    @property
+    def is_spark_connect(self) -> bool:
+        return _is_spark_connect(self.spark_session)
+
     @property
     def is_distributed(self) -> bool:
         return True
@@ -363,6 +379,11 @@ class SparkExecutionEngine(ExecutionEngine):
 
     def get_current_parallelism(self) -> int:
         spark = self.spark_session
+        if self.is_spark_connect:  # pragma: no cover
+            num = spark.conf.get("spark.default.parallelism", "")
+            if num != "":
+                return int(num)
+            return int(spark.conf.get("spark.sql.shuffle.partitions", "200"))
         e_cores = int(spark.conf.get("spark.executor.cores", "1"))
         tc = int(spark.conf.get("spark.task.cpus", "1"))
         sc = spark._jsc.sc()
@@ -403,10 +424,13 @@ class SparkExecutionEngine(ExecutionEngine):
             return df.count()
 
         df = self._to_spark_df(df)
-        num_funcs = {KEYWORD_ROWCOUNT: lambda: _persist_and_count(df)}
+        num_funcs = {
+            KEYWORD_ROWCOUNT: lambda: _persist_and_count(df),
+            KEYWORD_PARALLELISM: lambda: self.get_current_parallelism(),
+        }
         num = partition_spec.get_num_partitions(**num_funcs)
 
-        if partition_spec.algo == "hash":
+        if partition_spec.algo in ["hash", "coarse"]:
             sdf = hash_repartition(
                 self.spark_session, df.native, num, partition_spec.partition_by
             )
@@ -712,14 +736,16 @@ class SparkExecutionEngine(ExecutionEngine):
                 if isinstance(df, SparkDataFrame):
                     return df
                 if isinstance(df, ArrowDataFrame):
+                    raw_df: Any = df.as_pandas()
                     sdf = self.spark_session.createDataFrame(
-                        df.as_array(), to_spark_schema(df.schema)
+                        raw_df, to_spark_schema(df.schema)
                     )
                     return SparkDataFrame(sdf, df.schema)
                 if isinstance(df, (ArrayDataFrame, IterableDataFrame)):
                     adf = ArrowDataFrame(df.as_array(type_safe=False), df.schema)
+                    raw_df = adf.as_pandas()
                     sdf = self.spark_session.createDataFrame(
-                        adf.as_array(), to_spark_schema(df.schema)
+                        raw_df, to_spark_schema(df.schema)
                     )
                     return SparkDataFrame(sdf, df.schema)
                 if any(pa.types.is_struct(t) for t in df.schema.types):
@@ -731,7 +757,7 @@ class SparkExecutionEngine(ExecutionEngine):
                         df.as_pandas(), to_spark_schema(df.schema)
                     )
                 return SparkDataFrame(sdf, df.schema)
-            if isinstance(df, ps.DataFrame):
+            if is_spark_dataframe(df):
                 return SparkDataFrame(df, None if schema is None else to_schema(schema))
             if isinstance(df, RDD):
                 assert_arg_not_none(schema, "schema")
@@ -805,7 +831,7 @@ class _Mapper(object):  # pragma: no cover
             return
         if self.on_init is not None:
             self.on_init(no, df)
-        if self.partition_spec.empty:
+        if self.partition_spec.empty or self.partition_spec.algo == "coarse":
             partitions: Iterable[Tuple[int, int, EmptyAwareIterable]] = [
                 (0, 0, df.native)
             ]

fugue_spark/registry.py

@@ -18,18 +18,24 @@ from fugue.plugins import as_fugue_dataset, infer_execution_engine, parse_creato
 from fugue_spark.dataframe import SparkDataFrame
 from fugue_spark.execution_engine import SparkExecutionEngine
 
+from ._utils.misc import SparkConnectDataFrame, SparkConnectSession, is_spark_dataframe
+
 _is_sparksql = namespace_candidate("sparksql", lambda x: isinstance(x, str))
 
 
 @infer_execution_engine.candidate(
-    lambda objs: is_pandas_or(objs, (ps.DataFrame, SparkDataFrame))
+    lambda objs: (
+        is_pandas_or(objs, (ps.DataFrame, SparkConnectDataFrame, SparkDataFrame))
+        if SparkConnectDataFrame is not None
+        else is_pandas_or(objs, (ps.DataFrame, SparkDataFrame))
+    )
     or any(_is_sparksql(obj) for obj in objs)
 )
 def _infer_spark_client(obj: Any) -> Any:
     return SparkSession.builder.getOrCreate()
 
 
-@as_fugue_dataset.candidate(lambda df, **kwargs: isinstance(df, ps.DataFrame))
+@as_fugue_dataset.candidate(lambda df, **kwargs: is_spark_dataframe(df))
 def _spark_as_fugue_df(df: ps.DataFrame, **kwargs: Any) -> SparkDataFrame:
     return SparkDataFrame(df, **kwargs)
 
@@ -53,6 +59,12 @@ def _register_engines() -> None:
         lambda session, conf, **kwargs: SparkExecutionEngine(session, conf=conf),
         on_dup="ignore",
     )
+    if SparkConnectSession is not None:
+        register_execution_engine(
+            SparkConnectSession,
+            lambda session, conf, **kwargs: SparkExecutionEngine(session, conf=conf),
+            on_dup="ignore",
+        )
 
 
 @fugue_annotated_param(SparkExecutionEngine)
@@ -81,7 +93,7 @@ class _SparkDataFrameParam(DataFrameParam):
         return ctx.to_df(df).native
 
     def to_output_df(self, output: Any, schema: Any, ctx: Any) -> DataFrame:
-        assert isinstance(output, ps.DataFrame)
+        assert is_spark_dataframe(output)
         assert isinstance(ctx, SparkExecutionEngine)
         return ctx.to_df(output, schema=schema)
 

fugue_test/builtin_suite.py

@@ -57,7 +57,6 @@ from fugue.exceptions import (
     FugueWorkflowError,
     FugueWorkflowRuntimeValidationError,
 )
-from fugue_test._utils import _is_spark2
 
 
 class BuiltInTests(object):
@@ -98,7 +97,7 @@ class BuiltInTests(object):
             dag.run(self.engine)
 
         def test_create_df_equivalence(self):
-            ndf = self.engine.to_df(pd.DataFrame([[0]], columns=["a"]))
+            ndf = fa.as_fugue_engine_df(self.engine, pd.DataFrame([[0]], columns=["a"]))
             dag1 = FugueWorkflow()
             dag1.df(ndf).show()
             dag2 = FugueWorkflow()
@@ -1316,12 +1315,13 @@ class BuiltInTests(object):
             assert FileSystem().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)
+            pdf["c"] = pdf["c"].astype(int)
             pd.testing.assert_frame_equal(
-                pd.read_parquet(path3).sort_values("a").reset_index(drop=True),
-                pd.DataFrame({"c": pd.Categorical([6, 2]), "a": [1, 7]}).reset_index(
-                    drop=True
-                ),
+                pdf,
+                pd.DataFrame({"c": [6, 2], "a": [1, 7]}).reset_index(drop=True),
                 check_like=True,
+                check_dtype=False,
             )
 
         def test_save_and_use(self):
@@ -1675,9 +1675,7 @@ class BuiltInTests(object):
                 assert not isinstance(sdf4, DataFrame)
                 assert fa.is_local(sdf4)
 
-        @pytest.mark.skipif(
-            _is_spark2() or os.name == "nt", reason="Skip Spark<3 or Windows"
-        )
+        @pytest.mark.skipif(os.name == "nt", reason="Skip Windows")
         def test_any_column_name(self):
 
             f_parquet = os.path.join(str(self.tmpdir), "a.parquet")

fugue_test/dataframe_suite.py

@@ -415,7 +415,7 @@ class DataFrameTests(object):
 
             # str -> date
             df = self.df(
-                [["1", "2020-01-01"], ["2", "2020-01-02 01:02:03"], ["3", None]],
+                [["1", "2020-01-01"], ["2", "2020-01-02"], ["3", None]],
                 "a:str,b:str",
             )
             ndf = fi.alter_columns(df, "b:date,a:int", as_fugue=True)
@@ -428,12 +428,16 @@ class DataFrameTests(object):
 
             # str -> datetime
             df = self.df(
-                [["1", "2020-01-01"], ["2", "2020-01-02 01:02:03"], ["3", None]],
+                [
+                    ["1", "2020-01-01 01:02:03"],
+                    ["2", "2020-01-02 01:02:03"],
+                    ["3", None],
+                ],
                 "a:str,b:str",
             )
             ndf = fi.alter_columns(df, "b:datetime,a:int", as_fugue=True)
             assert [
-                [1, datetime(2020, 1, 1)],
+                [1, datetime(2020, 1, 1, 1, 2, 3)],
                 [2, datetime(2020, 1, 2, 1, 2, 3)],
                 [3, None],
             ] == fi.as_array(ndf, type_safe=True)