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

fugue_ray/dataframe.py

@@ -17,6 +17,7 @@ from fugue.plugins import (
     rename,
 )
 
+from ._constants import _ZERO_COPY
 from ._utils.dataframe import build_empty, get_dataset_format
 
 
@@ -140,7 +141,10 @@ class RayDataFrame(DataFrame):
         if cols == self.columns:
             return self
         rdf = self.native.map_batches(
-            lambda b: b.select(cols), batch_format="pyarrow", **self._remote_args()
+            lambda b: b.select(cols),
+            batch_format="pyarrow",
+            **_ZERO_COPY,
+            **self._remote_args(),
         )
         return RayDataFrame(rdf, self.schema.extract(cols), internal_schema=True)
 
@@ -174,6 +178,7 @@ class RayDataFrame(DataFrame):
         rdf = self.native.map_batches(
             lambda b: b.rename_columns(new_cols),
             batch_format="pyarrow",
+            **_ZERO_COPY,
             **self._remote_args(),
         )
         return RayDataFrame(rdf, schema=new_schema, internal_schema=True)
@@ -188,7 +193,7 @@ class RayDataFrame(DataFrame):
         if self.schema == new_schema:
             return self
         rdf = self.native.map_batches(
-            _alter, batch_format="pyarrow", **self._remote_args()
+            _alter, batch_format="pyarrow", **_ZERO_COPY, **self._remote_args()
         )
         return RayDataFrame(rdf, schema=new_schema, internal_schema=True)
 
@@ -231,7 +236,9 @@ class RayDataFrame(DataFrame):
             return ArrowDataFrame(table).alter_columns(schema).native  # type: ignore
 
         return (
-            rdf.map_batches(_alter, batch_format="pyarrow", **self._remote_args()),
+            rdf.map_batches(
+                _alter, batch_format="pyarrow", **_ZERO_COPY, **self._remote_args()
+            ),
             schema,
         )
 
@@ -273,7 +280,9 @@ def _rename_ray_dataframe(df: rd.Dataset, columns: Dict[str, Any]) -> rd.Dataset
     if len(missing) > 0:
         raise FugueDataFrameOperationError("found nonexistent columns: {missing}")
     new_cols = [columns.get(name, name) for name in cols]
-    return df.map_batches(lambda b: b.rename_columns(new_cols), batch_format="pyarrow")
+    return df.map_batches(
+        lambda b: b.rename_columns(new_cols), batch_format="pyarrow", **_ZERO_COPY
+    )
 
 
 def _get_arrow_tables(df: rd.Dataset) -> Iterable[pa.Table]:

fugue_ray/execution_engine.py

@@ -2,7 +2,7 @@ from typing import Any, Callable, Dict, List, Optional, Type, Union
 
 import pyarrow as pa
 import ray
-from duckdb import DuckDBPyConnection, DuckDBPyRelation
+from duckdb import DuckDBPyConnection
 from triad import Schema, assert_or_throw, to_uuid
 from triad.utils.threading import RunOnce
 
@@ -15,14 +15,14 @@ from fugue import (
     PartitionCursor,
     PartitionSpec,
 )
-from fugue.constants import KEYWORD_ROWCOUNT
+from fugue.constants import KEYWORD_PARALLELISM, KEYWORD_ROWCOUNT
 from fugue.dataframe.arrow_dataframe import _build_empty_arrow
 from fugue_duckdb.dataframe import DuckDataFrame
 from fugue_duckdb.execution_engine import DuckExecutionEngine
 
 from ._constants import FUGUE_RAY_DEFAULT_BATCH_SIZE, FUGUE_RAY_ZERO_COPY
 from ._utils.cluster import get_default_partitions, get_default_shuffle_partitions
-from ._utils.dataframe import add_partition_key
+from ._utils.dataframe import add_coarse_partition_key, add_partition_key
 from ._utils.io import RayIO
 from .dataframe import RayDataFrame
 
@@ -72,12 +72,14 @@ class RayMapEngine(MapEngine):
         partition_spec: PartitionSpec,
         on_init: Optional[Callable[[int, DataFrame], Any]] = None,
     ) -> DataFrame:
-        presort = partition_spec.presort
+        output_schema = Schema(output_schema)
+        input_schema = df.schema
+        presort = partition_spec.get_sorts(
+            input_schema, with_partition_keys=partition_spec.algo == "coarse"
+        )
         presort_tuples = [
             (k, "ascending" if v else "descending") for k, v in presort.items()
         ]
-        output_schema = Schema(output_schema)
-        input_schema = df.schema
         cursor = partition_spec.get_cursor(input_schema, 0)
         on_init_once: Any = (
             None
@@ -91,7 +93,7 @@ class RayMapEngine(MapEngine):
             if adf.shape[0] == 0:
                 return _build_empty_arrow(output_schema)
             adf = adf.remove_column(len(input_schema))  # remove partition key
-            if len(presort_tuples) > 0:
+            if len(partition_spec.presort) > 0:
                 if pa.__version__ < "7":  # pragma: no cover
                     idx = pa.compute.sort_indices(
                         adf, options=pa.compute.SortOptions(presort_tuples)
@@ -118,12 +120,20 @@ class RayMapEngine(MapEngine):
                 _df = self.execution_engine.repartition(  # type: ignore
                     _df, PartitionSpec(num=n)
                 )
-        rdf, _ = add_partition_key(
-            _df.native,
-            keys=partition_spec.partition_by,
-            input_schema=input_schema,
-            output_key=_RAY_PARTITION_KEY,
-        )
+        if partition_spec.algo != "coarse":
+            rdf, _ = add_partition_key(
+                _df.native,
+                keys=partition_spec.partition_by,
+                input_schema=input_schema,
+                output_key=_RAY_PARTITION_KEY,
+            )
+        else:
+            rdf = add_coarse_partition_key(
+                _df.native,
+                keys=partition_spec.partition_by,
+                output_key=_RAY_PARTITION_KEY,
+                bucket=_df.num_partitions,
+            )
 
         gdf = rdf.groupby(_RAY_PARTITION_KEY)
         sdf = gdf.map_groups(
@@ -205,6 +215,7 @@ class RayExecutionEngine(DuckExecutionEngine):
     ):
         if not ray.is_initialized():  # pragma: no cover
             ray.init()
+
         super().__init__(conf, connection)
         self._io = RayIO(self)
 
@@ -235,12 +246,15 @@ class RayExecutionEngine(DuckExecutionEngine):
 
         rdf = self._to_ray_df(df)
 
-        num_funcs = {KEYWORD_ROWCOUNT: lambda: _persist_and_count(rdf)}
+        num_funcs = {
+            KEYWORD_ROWCOUNT: lambda: _persist_and_count(rdf),
+            KEYWORD_PARALLELISM: lambda: self.get_current_parallelism(),
+        }
         num = partition_spec.get_num_partitions(**num_funcs)
         pdf = rdf.native
 
         if num > 0:
-            if partition_spec.algo in ["hash", "even"]:
+            if partition_spec.algo in ["hash", "even", "coarse"]:
                 pdf = pdf.repartition(num)
             elif partition_spec.algo == "rand":
                 pdf = pdf.repartition(num, shuffle=True)
@@ -267,6 +281,16 @@ class RayExecutionEngine(DuckExecutionEngine):
             return df if not as_local else df.as_local()
         return super().convert_yield_dataframe(df, as_local)
 
+    def union(self, df1: DataFrame, df2: DataFrame, distinct: bool = True) -> DataFrame:
+        if distinct:
+            return super().union(df1, df2, distinct)
+        assert_or_throw(
+            df1.schema == df2.schema, ValueError(f"{df1.schema} != {df2.schema}")
+        )
+        tdf1 = self._to_ray_df(df1)
+        tdf2 = self._to_ray_df(df2)
+        return RayDataFrame(tdf1.native.union(tdf2.native), df1.schema)
+
     def load_df(  # type:ignore
         self,
         path: Union[str, List[str]],
@@ -315,12 +339,6 @@ class RayExecutionEngine(DuckExecutionEngine):
                 ValueError("schema must be None when df is a DataFrame"),
             )
             return df
-        if isinstance(df, DuckDBPyRelation):
-            assert_or_throw(
-                schema is None,
-                ValueError("schema must be None when df is a DuckDBPyRelation"),
-            )
-            return DuckDataFrame(df)
         return RayDataFrame(df, schema)
 
     def _get_remote_args(self) -> Dict[str, Any]:

fugue_spark/_utils/convert.py

@@ -1,29 +1,24 @@
 from typing import Any, Iterable, List, Tuple
 
+import cloudpickle
+import pandas as pd
 import pyarrow as pa
 import pyspark.sql as ps
 import pyspark.sql.types as pt
-
-try:  # pyspark < 3
-    from pyspark.sql.types import from_arrow_type, to_arrow_type  # type: ignore
-
-    # https://issues.apache.org/jira/browse/SPARK-29041
-    pt._acceptable_types[pt.BinaryType] = (bytearray, bytes)  # type: ignore  # pragma: no cover  # noqa: E501  # pylint: disable=line-too-long
-except ImportError:  # pyspark >=3
-    from pyspark.sql.pandas.types import from_arrow_type, to_arrow_type
-
 from pyarrow.types import is_list, is_struct, is_timestamp
+from pyspark.sql.pandas.types import from_arrow_type, to_arrow_type
 from triad.collections import Schema
 from triad.utils.assertion import assert_arg_not_none, assert_or_throw
 from triad.utils.pyarrow import TRIAD_DEFAULT_TIMESTAMP
 from triad.utils.schema import quote_name
+from .misc import is_spark_dataframe
 
 
 def to_spark_schema(obj: Any) -> pt.StructType:
     assert_arg_not_none(obj, "schema")
     if isinstance(obj, pt.StructType):
         return obj
-    if isinstance(obj, ps.DataFrame):
+    if is_spark_dataframe(obj):
         return obj.schema
     return _from_arrow_schema(Schema(obj).pa_schema)
 
@@ -32,7 +27,7 @@ def to_schema(obj: Any) -> Schema:
     assert_arg_not_none(obj, "obj")
     if isinstance(obj, pt.StructType):
         return Schema(_to_arrow_schema(obj))
-    if isinstance(obj, ps.DataFrame):
+    if is_spark_dataframe(obj):
         return to_schema(obj.schema)
     return Schema(obj)
 
@@ -113,6 +108,22 @@ def to_type_safe_input(rows: Iterable[ps.Row], schema: Schema) -> Iterable[List[
             yield r
 
 
+def to_pandas(df: ps.DataFrame) -> pd.DataFrame:
+    if pd.__version__ < "2" or not any(
+        isinstance(x.dataType, (pt.TimestampType, pt.TimestampNTZType))
+        for x in df.schema.fields
+    ):
+        return df.toPandas()
+
+    def serialize(dfs):  # pragma: no cover
+        for df in dfs:
+            data = cloudpickle.dumps(df)
+            yield pd.DataFrame([[data]], columns=["data"])
+
+    sdf = df.mapInPandas(serialize, schema="data binary")
+    return pd.concat(cloudpickle.loads(x.data) for x in sdf.collect())
+
+
 # TODO: the following function always set nullable to true,
 # but should we use field.nullable?
 def _to_arrow_type(dt: pt.DataType) -> pa.DataType:

fugue_spark/_utils/io.py

@@ -21,7 +21,6 @@ class SparkIO(object):
             "csv": self._load_csv,
             "parquet": self._load_parquet,
             "json": self._load_json,
-            "avro": self._load_avro,
         }
 
     def load_df(
@@ -136,15 +135,3 @@ class SparkIO(object):
             return SparkDataFrame(reader.load(p))[columns]
         schema = Schema(columns)
         return SparkDataFrame(reader.load(p)[schema.names], schema)
-
-    def _load_avro(self, p: List[str], columns: Any = None, **kwargs: Any) -> DataFrame:
-        reader = self._session.read.format(
-            "avro"
-        )  # avro is an external data source that has built-in support since spark 2.4
-        reader.options(**kwargs)
-        if columns is None:
-            return SparkDataFrame(reader.load(p))
-        if isinstance(columns, list):  # column names
-            return SparkDataFrame(reader.load(p))[columns]
-        schema = Schema(columns)
-        return SparkDataFrame(reader.load(p)[schema.names], schema)

fugue_spark/_utils/misc.py

@@ -0,0 +1,27 @@
+from typing import Any
+
+try:
+    from pyspark.sql.connect.session import SparkSession as SparkConnectSession
+    from pyspark.sql.connect.dataframe import DataFrame as SparkConnectDataFrame
+except ImportError:  # pragma: no cover
+    SparkConnectSession = None
+    SparkConnectDataFrame = None
+import pyspark.sql as ps
+
+
+def is_spark_connect(session: Any) -> bool:
+    return SparkConnectSession is not None and isinstance(
+        session, (SparkConnectSession, SparkConnectDataFrame)
+    )
+
+
+def is_spark_dataframe(df: Any) -> bool:
+    return isinstance(df, ps.DataFrame) or (
+        SparkConnectDataFrame is not None and isinstance(df, SparkConnectDataFrame)
+    )
+
+
+def is_spark_session(session: Any) -> bool:
+    return isinstance(session, ps.SparkSession) or (
+        SparkConnectSession is not None and isinstance(session, SparkConnectSession)
+    )

fugue_spark/_utils/partition.py

@@ -1,11 +1,12 @@
-import random
 from typing import Any, Iterable, List
 
 import pyspark.sql as ps
-from fugue_spark._utils.convert import to_schema, to_spark_schema
+import pyspark.sql.functions as psf
 from pyspark import RDD
 from pyspark.sql import SparkSession
-from pyspark.sql.functions import lit
+import warnings
+from .convert import to_schema, to_spark_schema
+from .misc import is_spark_connect
 
 _PARTITION_DUMMY_KEY = "__partition_dummy_key__"
 
@@ -28,16 +29,10 @@ def rand_repartition(
     if len(cols) > 0 or num <= 1:
         return hash_repartition(session, df, num, cols)
 
-    def _rand(rows: Iterable[Any], n: int) -> Iterable[Any]:  # pragma: no cover
-        for row in rows:
-            yield (random.randint(0, n - 1), row)
-
-    rdd = (
-        df.rdd.mapPartitions(lambda r: _rand(r, num))
-        .partitionBy(num, lambda k: k)
-        .mapPartitions(_to_rows)
+    tdf = df.withColumn(
+        _PARTITION_DUMMY_KEY, (psf.rand(0) * psf.lit(2**15 - 1)).cast("long")
     )
-    return session.createDataFrame(rdd, df.schema)
+    return tdf.repartition(num, _PARTITION_DUMMY_KEY)[df.schema.names]
 
 
 def even_repartition(
@@ -45,6 +40,9 @@ def even_repartition(
 ) -> ps.DataFrame:
     if num == 1:
         return _single_repartition(df)
+    if is_spark_connect(session):  # pragma: no cover
+        warnings.warn("Even repartitioning is not supported by Spark Connect")
+        return hash_repartition(session, df, num, cols)
     if len(cols) == 0:
         if num == 0:
             return df
@@ -82,7 +80,7 @@ def even_repartition(
 
 def _single_repartition(df: ps.DataFrame) -> ps.DataFrame:
     return (
-        df.withColumn(_PARTITION_DUMMY_KEY, lit(0))
+        df.withColumn(_PARTITION_DUMMY_KEY, psf.lit(0))
         .repartition(_PARTITION_DUMMY_KEY)
         .drop(_PARTITION_DUMMY_KEY)
     )
@@ -93,11 +91,6 @@ def _to_rows(rdd: Iterable[Any]) -> Iterable[Any]:  # pragma: no cover
         yield item[1]
 
 
-def _to_rows_with_key(rdd: Iterable[Any]) -> Iterable[Any]:  # pragma: no cover
-    for item in rdd:
-        yield list(item[1]) + [item[0]]
-
-
 def _zipWithIndex(rdd: RDD, to_rows: bool = False) -> RDD:
     """
     Modified from

fugue_spark/dataframe.py

@@ -30,7 +30,9 @@ from fugue.plugins import (
     rename,
     select_columns,
 )
-from fugue_spark._utils.convert import to_cast_expression, to_schema, to_type_safe_input
+
+from ._utils.convert import to_cast_expression, to_pandas, to_schema, to_type_safe_input
+from ._utils.misc import is_spark_connect, is_spark_dataframe
 
 
 class SparkDataFrame(DataFrame):
@@ -51,12 +53,12 @@ class SparkDataFrame(DataFrame):
 
     def __init__(self, df: Any = None, schema: Any = None):  # noqa: C901
         self._lock = SerializableRLock()
-        if isinstance(df, ps.DataFrame):
+        if is_spark_dataframe(df):
             if schema is not None:
                 schema = to_schema(schema).assert_not_empty()
                 has_cast, expr = to_cast_expression(df, schema, True)
                 if has_cast:
-                    df = df.selectExpr(*expr)
+                    df = df.selectExpr(*expr)  # type: ignore
             else:
                 schema = to_schema(df).assert_not_empty()
             self._native = df
@@ -94,7 +96,7 @@ class SparkDataFrame(DataFrame):
             data = list(to_type_safe_input(self.native.collect(), self.schema))
             res: LocalBoundedDataFrame = ArrayDataFrame(data, self.schema)
         else:
-            res = PandasDataFrame(self.native.toPandas(), self.schema)
+            res = PandasDataFrame(self.as_pandas(), self.schema)
         if self.has_metadata:
             res.reset_metadata(self.metadata)
         return res
@@ -126,7 +128,7 @@ class SparkDataFrame(DataFrame):
         return SparkDataFrame(self.native[schema.names])
 
     def as_pandas(self) -> pd.DataFrame:
-        return self.native.toPandas()
+        return to_pandas(self.native)
 
     def rename(self, columns: Dict[str, str]) -> DataFrame:
         try:
@@ -150,6 +152,9 @@ class SparkDataFrame(DataFrame):
     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):
@@ -182,47 +187,47 @@ class SparkDataFrame(DataFrame):
         return SparkDataFrame(self.native.select(*columns))
 
 
-@is_df.candidate(lambda df: isinstance(df, ps.DataFrame))
+@is_df.candidate(lambda df: is_spark_dataframe(df))
 def _spark_is_df(df: ps.DataFrame) -> bool:
     return True
 
 
-@get_num_partitions.candidate(lambda df: isinstance(df, ps.DataFrame))
+@get_num_partitions.candidate(lambda df: is_spark_dataframe(df))
 def _spark_num_partitions(df: ps.DataFrame) -> int:
     return df.rdd.getNumPartitions()
 
 
-@count.candidate(lambda df: isinstance(df, ps.DataFrame))
+@count.candidate(lambda df: is_spark_dataframe(df))
 def _spark_df_count(df: ps.DataFrame) -> int:
     return df.count()
 
 
-@is_bounded.candidate(lambda df: isinstance(df, ps.DataFrame))
+@is_bounded.candidate(lambda df: is_spark_dataframe(df))
 def _spark_df_is_bounded(df: ps.DataFrame) -> bool:
     return True
 
 
-@is_empty.candidate(lambda df: isinstance(df, ps.DataFrame))
+@is_empty.candidate(lambda df: is_spark_dataframe(df))
 def _spark_df_is_empty(df: ps.DataFrame) -> bool:
     return df.first() is None
 
 
-@is_local.candidate(lambda df: isinstance(df, ps.DataFrame))
+@is_local.candidate(lambda df: is_spark_dataframe(df))
 def _spark_df_is_local(df: ps.DataFrame) -> bool:
     return False
 
 
-@as_local_bounded.candidate(lambda df: isinstance(df, ps.DataFrame))
+@as_local_bounded.candidate(lambda df: is_spark_dataframe(df))
 def _spark_df_as_local(df: ps.DataFrame) -> pd.DataFrame:
-    return df.toPandas()
+    return to_pandas(df)
 
 
-@get_column_names.candidate(lambda df: isinstance(df, ps.DataFrame))
+@get_column_names.candidate(lambda df: is_spark_dataframe(df))
 def _get_spark_df_columns(df: ps.DataFrame) -> List[Any]:
     return df.columns
 
 
-@rename.candidate(lambda df, *args, **kwargs: isinstance(df, ps.DataFrame))
+@rename.candidate(lambda df, *args, **kwargs: is_spark_dataframe(df))
 def _rename_spark_df(
     df: ps.DataFrame, columns: Dict[str, Any], as_fugue: bool = False
 ) -> ps.DataFrame:
@@ -232,7 +237,7 @@ def _rename_spark_df(
     return _adjust_df(_rename_spark_dataframe(df, columns), as_fugue=as_fugue)
 
 
-@drop_columns.candidate(lambda df, *args, **kwargs: isinstance(df, ps.DataFrame))
+@drop_columns.candidate(lambda df, *args, **kwargs: is_spark_dataframe(df))
 def _drop_spark_df_columns(
     df: ps.DataFrame, columns: List[str], as_fugue: bool = False
 ) -> Any:
@@ -244,7 +249,7 @@ def _drop_spark_df_columns(
     return _adjust_df(df[cols], as_fugue=as_fugue)
 
 
-@select_columns.candidate(lambda df, *args, **kwargs: isinstance(df, ps.DataFrame))
+@select_columns.candidate(lambda df, *args, **kwargs: is_spark_dataframe(df))
 def _select_spark_df_columns(
     df: ps.DataFrame, columns: List[Any], as_fugue: bool = False
 ) -> Any:
@@ -254,7 +259,7 @@ def _select_spark_df_columns(
     return _adjust_df(df[columns], as_fugue=as_fugue)
 
 
-@head.candidate(lambda df, *args, **kwargs: isinstance(df, ps.DataFrame))
+@head.candidate(lambda df, *args, **kwargs: is_spark_dataframe(df))
 def _spark_df_head(
     df: ps.DataFrame,
     n: int,
@@ -264,7 +269,7 @@ def _spark_df_head(
     if columns is not None:
         df = df[columns]
     res = df.limit(n)
-    return SparkDataFrame(res).as_local() if as_fugue else res.toPandas()
+    return SparkDataFrame(res).as_local() if as_fugue else to_pandas(res)
 
 
 def _rename_spark_dataframe(df: ps.DataFrame, names: Dict[str, Any]) -> ps.DataFrame: