fugue 0.9.9.dev0__tar.gz → 0.9.9.dev1__tar.gz

{fugue-0.9.9.dev0 → fugue-0.9.9.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fugue
-Version: 0.9.9.dev0
+Version: 0.9.9.dev1
 Summary: An abstraction layer for distributed computing
 Author-email: The Fugue Development Team <hello@fugue.ai>
 License: Apache-2.0

{fugue-0.9.9.dev0 → fugue-0.9.9.dev1}/fugue.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fugue
-Version: 0.9.9.dev0
+Version: 0.9.9.dev1
 Summary: An abstraction layer for distributed computing
 Author-email: The Fugue Development Team <hello@fugue.ai>
 License: Apache-2.0

{fugue-0.9.9.dev0 → fugue-0.9.9.dev1}/fugue_spark/_utils/convert.py

@@ -34,6 +34,27 @@ except ImportError:  # pragma: no cover
 _PYSPARK_ARROW_FRIENDLY = version.parse(pyspark.__version__) >= version.parse("3.3")
 
 
+def _combine_arrow_chunks(df: pd.DataFrame) -> pd.DataFrame:
+    # pyspark < 4.4 can't serialize multi-chunk Arrow columns
+    # (SPARK-58625 UDF output, SPARK-46776 createDataFrame).
+    # Delete once fugue requires pyspark >= 4.4.
+    cols = {}
+    for c, s in df.items():
+        d = s.dtype
+        if isinstance(d, pd.ArrowDtype) or (
+            isinstance(d, pd.StringDtype) and d.storage in ("pyarrow", "pyarrow_numpy")
+        ):
+            arr = s.array.__arrow_array__()
+            if isinstance(arr, pa.ChunkedArray) and arr.num_chunks > 1:
+                cols[c] = pd.array(arr.combine_chunks(), dtype=d)
+    if not cols:
+        return df
+    df = df.copy(deep=False)
+    for c, a in cols.items():
+        df[c] = a
+    return df
+
+
 def pandas_udf_can_accept(schema: Schema, is_input: bool) -> bool:
     try:
         # pyspark's own from_arrow_schema to_arrow_schema
@@ -166,7 +187,7 @@ def pd_to_spark_df(
     session: ps.SparkSession, df: pd.DataFrame, schema: Any = None
 ) -> ps.DataFrame:
     if _PYSPARK_ARROW_FRIENDLY:
-        return session.createDataFrame(df, schema=schema)
+        return session.createDataFrame(_combine_arrow_chunks(df), schema=schema)
     else:  # pragma: no cover
         # Not efficient for pyspark<3.3
         return session.createDataFrame(df.astype(object), schema=schema)

{fugue-0.9.9.dev0 → fugue-0.9.9.dev1}/fugue_spark/execution_engine.py

@@ -45,6 +45,7 @@ from fugue.execution.execution_engine import ExecutionEngine, MapEngine, SQLEngi
 from ._constants import FUGUE_SPARK_CONF_USE_PANDAS_UDF, FUGUE_SPARK_DEFAULT_CONF
 from ._utils.convert import (
     _PYSPARK_ARROW_FRIENDLY,
+    _combine_arrow_chunks,
     pandas_udf_can_accept,
     to_schema,
     to_spark_df,
@@ -864,4 +865,4 @@ class _Mapper(object):  # pragma: no cover
 
 
 def _to_safe_spark_worker_pandas(df: pd.DataFrame):  # pragma: no cover
-    return df if _PYSPARK_ARROW_FRIENDLY else df.astype(object)
+    return _combine_arrow_chunks(df) if _PYSPARK_ARROW_FRIENDLY else df.astype(object)

{fugue-0.9.9.dev0 → fugue-0.9.9.dev1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "fugue"
-version = "0.9.9.dev0"
+version = "0.9.9.dev1"
 description = "An abstraction layer for distributed computing"
 readme = "README.md"
 license = { text = "Apache-2.0" }