Repository: dask/dask
Pull request #11756 (2025-02-18 12:32:37)

Ensure `map_partitions` returns Series object if function returns scalar
Supersedes https://github.com/dask/dask/pull/10739
xref https://github.com/dask/dask/pull/11736

- [x] Closes https://github.com/dask-contrib/dask-expr/issues/625
- [x] Tests added / passed
- [x] Passes `pre-commit run --all-files`


github-actions[bot]: ## Unit Test Results
_See [test report](https://dask.github.io/dask/test_report.html) for an extended history of previous test failures. This is useful for diagnosing flaky tests._

      9 files  ± 0        9 suites  ±0   3h 26m 30s ⏱️ + 4m 14s
 17 718 tests + 2   16 506 ✅ +1   1 211 💤 ± 0  1 ❌ +1 
158 536 runs  +18  146 391 ✅ +4  12 144 💤 +13  1 ❌ +1 

For more details on these failures, see [this check](https://github.com/dask/dask/runs/37399936531).

Results for commit 82d1a3e6. ± Comparison against base commit 388bb0d7.

[test-results]:data:application/gzip;base64,H4sIAESItGcC/02Myw7CIBQFf6Vh7YJb4EL9GcMzIfZhKKyM/y40FrucOTnzJiHOfif3YboNZC8xd3Al6Ry3tSKMbKJV1TG3GaQEdfJjL9Y2iYLiXz7j63gCdBd0nJvrwqe0pWpaO5X1SAslGP64pzmyCbrsac5Pd0kffC3bbVlirkDU6EAzj5Jah4FKSx0NQXJptfcBjREIihlOPl+2QB3QGAEAAA==


Patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
index a0aa0d9fe..32d8bbcb6 100644
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -18,6 +18,7 @@ from dask.dataframe.utils import (
     has_known_categories,
     is_dataframe_like,
     is_index_like,
+    is_scalar,
     is_series_like,
     meta_frame_constructor,
     meta_series_constructor,
@@ -95,7 +96,13 @@ def apply_and_enforce(*args, **kwargs):
     func = kwargs.pop("_func")
     meta = kwargs.pop("_meta")
     df = func(*args, **kwargs)
-    if is_dataframe_like(df) or is_series_like(df) or is_index_like(df):
+
+    if any(
+        bool(is_dataframe_like(obj) or is_series_like(obj) or is_index_like(obj))
+        for obj in [df, meta]
+    ):
+        if is_scalar(df):
+            df = pd.Series(df)
         if not len(df):
             return meta
         if is_dataframe_like(df):
diff --git a/dask/dataframe/dask_expr/_collection.py b/dask/dataframe/dask_expr/_collection.py
index 70a18d8a3..8335aba4d 100644
--- a/dask/dataframe/dask_expr/_collection.py
+++ b/dask/dataframe/dask_expr/_collection.py
@@ -106,7 +106,6 @@ from dask.dataframe.dask_expr._util import (
     _raise_if_object_series,
     _validate_axis,
     get_specified_shuffle,
-    is_scalar,
 )
 from dask.dataframe.dask_expr.io import FromPandasDivisions, FromScalars
 from dask.dataframe.dispatch import (
@@ -121,6 +120,7 @@ from dask.dataframe.utils import (
     has_known_categories,
     index_summary,
     insert_meta_param_description,
+    is_scalar,
     meta_frame_constructor,
     meta_series_constructor,
     pyarrow_strings_enabled,
diff --git a/dask/dataframe/dask_expr/_expr.py b/dask/dataframe/dask_expr/_expr.py
index 032717ba5..3f9578967 100644
--- a/dask/dataframe/dask_expr/_expr.py
+++ b/dask/dataframe/dask_expr/_expr.py
@@ -37,7 +37,6 @@ from dask.dataframe.dask_expr._util import (
     _calc_maybe_new_divisions,
     _convert_to_list,
     _tokenize_partial,
-    is_scalar,
 )
 from dask.dataframe.dispatch import make_meta, meta_nonempty
 from dask.dataframe.rolling import CombinedOutput, _head_timedelta, overlap_chunk
@@ -45,6 +44,7 @@ from dask.dataframe.shuffle import drop_overlap, get_overlap
 from dask.dataframe.utils import (
     clear_known_categories,
     drop_by_shallow_copy,
+    is_scalar,
     raise_on_meta_error,
     valid_divisions,
 )
diff --git a/dask/dataframe/dask_expr/_groupby.py b/dask/dataframe/dask_expr/_groupby.py
index d56cacd39..8c555c0e8 100644
--- a/dask/dataframe/dask_expr/_groupby.py
+++ b/dask/dataframe/dask_expr/_groupby.py
@@ -45,7 +45,6 @@ from dask.dataframe.dask_expr._util import (
     PANDAS_GE_300,
     _convert_to_list,
     get_specified_shuffle,
-    is_scalar,
 )
 from dask.dataframe.dispatch import concat, make_meta, meta_nonempty
 from dask.dataframe.groupby import (
@@ -82,7 +81,7 @@ from dask.dataframe.groupby import (
     _var_agg,
     _var_chunk,
 )
-from dask.dataframe.utils import insert_meta_param_description
+from dask.dataframe.utils import insert_meta_param_description, is_scalar
 from dask.typing import Key
 from dask.utils import M, derived_from, is_index_like
 
diff --git a/dask/dataframe/dask_expr/_indexing.py b/dask/dataframe/dask_expr/_indexing.py
index 12bd43f04..7758c032d 100644
--- a/dask/dataframe/dask_expr/_indexing.py
+++ b/dask/dataframe/dask_expr/_indexing.py
@@ -20,7 +20,6 @@ from dask.dataframe.dask_expr._expr import (
     Projection,
     are_co_aligned,
 )
-from dask.dataframe.dask_expr._util import is_scalar
 from dask.dataframe.dispatch import meta_nonempty
 from dask.dataframe.indexing import (
     _coerce_loc_index,
@@ -28,6 +27,7 @@ from dask.dataframe.indexing import (
     _partition_of_index_value,
     _partitions_of_index_values,
 )
+from dask.dataframe.utils import is_scalar
 from dask.typing import Key
 from dask.utils import is_arraylike, is_series_like
 
diff --git a/dask/dataframe/dask_expr/_merge.py b/dask/dataframe/dask_expr/_merge.py
index 190cde1e2..d6fecac69 100644
--- a/dask/dataframe/dask_expr/_merge.py
+++ b/dask/dataframe/dask_expr/_merge.py
@@ -33,7 +33,7 @@ from dask.dataframe.dask_expr._shuffle import (
     _is_numeric_cast_type,
     _select_columns_or_index,
 )
-from dask.dataframe.dask_expr._util import _convert_to_list, is_scalar
+from dask.dataframe.dask_expr._util import _convert_to_list
 from dask.dataframe.dispatch import make_meta, meta_nonempty
 from dask.dataframe.multi import (
     _concat_wrapper,
@@ -42,6 +42,7 @@ from dask.dataframe.multi import (
     merge_chunk,
 )
 from dask.dataframe.shuffle import partitioning_index
+from dask.dataframe.utils import is_scalar
 from dask.tokenize import _tokenize_deterministic
 from dask.typing import Key
 from dask.utils import apply, get_default_shuffle_method
diff --git a/dask/dataframe/dask_expr/_reductions.py b/dask/dataframe/dask_expr/_reductions.py
index 5ba667896..da9fb6a2a 100644
--- a/dask/dataframe/dask_expr/_reductions.py
+++ b/dask/dataframe/dask_expr/_reductions.py
@@ -38,8 +38,8 @@ from dask.dataframe.dask_expr._expr import (
     determine_column_projection,
     plain_column_projection,
 )
-from dask.dataframe.dask_expr._util import is_scalar
 from dask.dataframe.dispatch import make_meta, meta_nonempty
+from dask.dataframe.utils import is_scalar
 from dask.tokenize import _tokenize_deterministic
 from dask.typing import no_default
 from dask.utils import M, apply, funcname
diff --git a/dask/dataframe/dask_expr/_util.py b/dask/dataframe/dask_expr/_util.py
index aef4d8e90..c9de28dab 100644
--- a/dask/dataframe/dask_expr/_util.py
+++ b/dask/dataframe/dask_expr/_util.py
@@ -2,10 +2,9 @@ from __future__ import annotations
 
 import functools
 from collections import OrderedDict, UserDict
-from collections.abc import Hashable, Iterable, Sequence
+from collections.abc import Hashable
 from typing import Any, Literal, TypeVar, cast
 
-import numpy as np
 import pandas as pd
 from packaging.version import Version
 
@@ -79,22 +78,6 @@ def _convert_to_list(column) -> list | None:
     return column
 
 
-def is_scalar(x):
-    # np.isscalar does not work for some pandas scalars, for example pd.NA
-    if isinstance(x, (Sequence, Iterable)) and not isinstance(x, str):
-        return False
-    elif hasattr(x, "dtype"):
-        return isinstance(x, np.ScalarType)
-    if isinstance(x, dict):
-        return False
-    if isinstance(x, (str, int)) or x is None:
-        return True
-
-    from dask.dataframe.dask_expr._expr import Expr
-
-    return not isinstance(x, Expr)
-
-
 def _tokenize_partial(expr, ignore: list | None = None) -> str:
     # Helper function to "tokenize" the operands
     # that are not in the `ignore` list
diff --git a/dask/dataframe/dask_expr/diagnostics/_analyze.py b/dask/dataframe/dask_expr/diagnostics/_analyze.py
index 29d2cadc4..0dfd822ca 100644
--- a/dask/dataframe/dask_expr/diagnostics/_analyze.py
+++ b/dask/dataframe/dask_expr/diagnostics/_analyze.py
@@ -7,12 +7,12 @@ import pandas as pd
 
 from dask.base import DaskMethodsMixin
 from dask.dataframe.dask_expr._expr import Blockwise, Expr
-from dask.dataframe.dask_expr._util import is_scalar
 from dask.dataframe.dask_expr.diagnostics._explain import (
     _add_graphviz_edges,
     _explain_info,
 )
 from dask.dataframe.dask_expr.io.io import FusedIO
+from dask.dataframe.utils import is_scalar
 from dask.sizeof import sizeof
 from dask.tokenize import _tokenize_deterministic
 from dask.utils import format_bytes, import_required
diff --git a/dask/dataframe/utils.py b/dask/dataframe/utils.py
index 1911e5de7..9e156c41c 100644
--- a/dask/dataframe/utils.py
+++ b/dask/dataframe/utils.py
@@ -5,7 +5,7 @@ import re
 import sys
 import textwrap
 import traceback
-from collections.abc import Callable, Mapping
+from collections.abc import Callable, Iterable, Mapping, Sequence
 from contextlib import contextmanager
 from numbers import Number
 from typing import TypeVar, overload
@@ -43,6 +43,22 @@ except ImportError:
     pass
 
 
+def is_scalar(x):
+    # np.isscalar does not work for some pandas scalars, for example pd.NA
+    if isinstance(x, (Sequence, Iterable)) and not isinstance(x, str):
+        return False
+    elif hasattr(x, "dtype"):
+        return isinstance(x, np.ScalarType)
+    if isinstance(x, dict):
+        return False
+    if isinstance(x, (str, int)) or x is None:
+        return True
+
+    from dask.dataframe.dask_expr._expr import Expr
+
+    return not isinstance(x, Expr)
+
+
 def is_integer_na_dtype(t):
     dtype = getattr(t, "dtype", t)
     types = (


Test patch:
diff --git a/dask/dataframe/dask_expr/tests/test_collection.py b/dask/dataframe/dask_expr/tests/test_collection.py
index a4d2d15c4..36ababee1 100644
--- a/dask/dataframe/dask_expr/tests/test_collection.py
+++ b/dask/dataframe/dask_expr/tests/test_collection.py
@@ -22,7 +22,6 @@ from dask.dataframe.dask_expr import (
     expr,
     from_dict,
     from_pandas,
-    is_scalar,
     isna,
     optimize,
     to_datetime,
@@ -40,7 +39,7 @@ from dask.dataframe.dask_expr._shuffle import Shuffle
 from dask.dataframe.dask_expr.datasets import timeseries
 from dask.dataframe.dask_expr.io import FromPandas
 from dask.dataframe.dask_expr.tests._util import _backend_library, assert_eq, xfail_gpu
-from dask.dataframe.utils import UNKNOWN_CATEGORIES, pyarrow_strings_enabled
+from dask.dataframe.utils import UNKNOWN_CATEGORIES, is_scalar, pyarrow_strings_enabled
 from dask.utils import M
 
 # Set DataFrame backend for this module
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
index e9716e420..52e8e8d15 100644
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -903,6 +903,7 @@ def test_map_partitions():
         d.map_partitions(lambda df: 1),
         pd.Series([1, 1, 1], dtype=np.int64),
         check_divisions=False,
+        check_index=False,
     )
 
 
@@ -5353,3 +5354,19 @@ def test_array_to_df_conversion():
     result = arr.map_blocks(foo, meta=pd.DataFrame(np.random.random((1, 10))))
     expected = pd.DataFrame(arr.compute())
     assert_eq(result, expected, check_index=False)
+
+
+@pytest.mark.parametrize("npartitions", [1, 5])
+def test_map_partitions_always_series(npartitions):
+    pdf = pd.DataFrame({"x": range(50)})
+    ddf = dd.from_pandas(pdf, npartitions)
+
+    def assert_series(x):
+        assert isinstance(x, pd.Series)
+        return x
+
+    res = ddf.x.map_partitions(M.min).map_partitions(assert_series).compute()
+    assert len(res) == ddf.npartitions
+    assert min(res) == min(pdf.x)
+
+    assert ddf.x.map_partitions(M.min).count().compute() == ddf.npartitions