Repository: dask/dask
Pull request #11757 (2025-02-18 12:45:33)

Remove bytes as keys
This is hopefully only a minor change. I ran into this every now and then, particularly with mypy. The type `bytes` is extremely unnatural for keys and I very strongly suspect that this is still a remnant of python 2

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 23m 29s ⏱️ + 1m 13s
 17 716 tests ±0   16 504 ✅  -  1   1 211 💤 ± 0  1 ❌ +1 
158 518 runs  ±0  146 372 ✅  - 15  12 145 💤 +14  1 ❌ +1 

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

Results for commit c39b3fbc. ± Comparison against base commit 388bb0d7.

[test-results]:data:application/gzip;base64,H4sIAP+JtGcC/03MSw7CIBSF4a00jB1weVzAzZhCISH2YSiMjHsXGosdkf+74bxJiLPfyX0wt4HsJeYeU0ljjttaExijjeoxtzMoBXj2Yy/ONURJxR+f8XX8BOgWxjg36+BT2lIVWiWV9ZiWWoL+dZ8WyBXr2KeFPO0yffR12W3LEnMN4rixPFhnjRjrY5jgoLSnjiEa5ZGiQirsRD5fr0klIRgBAAA=


Patch:
diff --git a/dask/_task_spec.py b/dask/_task_spec.py
index 80df3cfbd..5506e5553 100644
--- a/dask/_task_spec.py
+++ b/dask/_task_spec.py
@@ -224,7 +224,7 @@ def convert_legacy_task(
             new_args.append(new)
         return Task(key, func, *new_args)
     try:
-        if isinstance(task, (bytes, int, float, str, tuple)):
+        if isinstance(task, (int, float, str, tuple)):
             if task in all_keys:
                 if key is None:
                     return Alias(task)
diff --git a/dask/core.py b/dask/core.py
index bf0c2dceb..dd61cc980 100644
--- a/dask/core.py
+++ b/dask/core.py
@@ -149,7 +149,7 @@ def keys_in_tasks(keys: Collection[Key], tasks: Iterable[Any], as_list: bool = F
 def iskey(key: object) -> bool:
     """Return True if the given object is a potential dask key; False otherwise.
 
-    The definition of a key in a Dask graph is any str, bytes, int, float, or tuple
+    The definition of a key in a Dask graph is any str, int, float, or tuple
     thereof.
 
     See Also
@@ -161,7 +161,7 @@ def iskey(key: object) -> bool:
     typ = type(key)
     if typ is tuple:
         return all(iskey(i) for i in cast(tuple, key))
-    return typ in {bytes, int, float, str}
+    return typ in {int, float, str}
 
 
 def validate_key(key: object) -> None:
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 = (
diff --git a/dask/typing.py b/dask/typing.py
index 6c67ff746..8a970f87e 100644
--- a/dask/typing.py
+++ b/dask/typing.py
@@ -25,7 +25,7 @@ CollType_co = TypeVar("CollType_co", bound="DaskCollection", covariant=True)
 PostComputeCallable = Callable
 
 
-Key: TypeAlias = Union[str, bytes, int, float, tuple["Key", ...]]
+Key: TypeAlias = Union[str, int, float, tuple["Key", ...]]
 # FIXME: This type is a little misleading. Low level graphs are often
 # MutableMappings but HLGs are not
 Graph: TypeAlias = Mapping[Key, Any]


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