Repository: dask/dask (python)

Issue:
KeyError when indexing into Series after calling `to_series` on Scalar
<!-- Please include a self-contained copy-pastable example that generates the issue if possible.

Please be concise with code posted. See guidelines below on how to provide a good bug report:

- Craft Minimal Bug Reports http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports
- Minimal Complete Verifiable Examples https://stackoverflow.com/help/mcve

Bug reports that follow these guidelines are easier to diagnose, and so are often handled much more quickly.
-->

**Describe the issue**:

**Minimal Complete Verifiable Example**:

```python
In [1]: import dask.dataframe as dd

In [2]: import pandas as pd

In [3]: data = {"a": [1, 3, 2]}

In [4]: df = dd.from_pandas(pd.DataFrame(data), npartitions=2)

In [5]: df['a'].sum().to_series().fillna(0)[0].compute()
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
File ~/polars-api-compat-dev/.venv/lib/python3.12/site-packages/pandas/core/indexes/base.py:3805, in Index.get_loc(self, key)
   3804 try:
-> 3805     return self._engine.get_loc(casted_key)
   3806 except KeyError as err:

File index.pyx:167, in pandas._libs.index.IndexEngine.get_loc()

File index.pyx:196, in pandas._libs.index.IndexEngine.get_loc()

File pandas/_libs/hashtable_class_helper.pxi:2606, in pandas._libs.hashtable.Int64HashTable.get_item()

File pandas/_libs/hashtable_class_helper.pxi:2630, in pandas._libs.hashtable.Int64HashTable.get_item()

KeyError: 0

The above exception was the direct cause of the following exception:

KeyError                                  Traceback (most recent call last)
Cell In[5], line 1
----> 1 df['a'].sum().to_series().fillna(0)[0].compute()

File ~/polars-api-compat-dev/.venv/lib/python3.12/site-packages/dask/dataframe/dask_expr/_collection.py:489, in FrameBase.compute(self, fuse, concatenate, **kwargs)
    487     out = out.repartition(npartitions=1)
    488 out = out.optimize(fuse=fuse)
--> 489 return DaskMethodsMixin.compute(out, **kwargs)

File ~/polars-api-compat-dev/.venv/lib/python3.12/site-packages/dask/base.py:374, in DaskMethodsMixin.compute(self, **kwargs)
    350 def compute(self, **kwargs):
    351     """Compute this dask collection
    352 
    353     This turns a lazy Dask collection into its in-memory equivalent.
   (...)
    372     dask.compute
    373     """
--> 374     (result,) = compute(self, traverse=False, **kwargs)
    375     return result

File ~/polars-api-compat-dev/.venv/lib/python3.12/site-packages/dask/base.py:662, in compute(traverse, optimize_graph, scheduler, get, *args, **kwargs)
    659     postcomputes.append(x.__dask_postcompute__())
    661 with shorten_traceback():
--> 662     results = schedule(dsk, keys, **kwargs)
    664 return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])

File ~/polars-api-compat-dev/.venv/lib/python3.12/site-packages/pandas/core/indexes/base.py:3812, in Index.get_loc(self, key)
   3807     if isinstance(casted_key, slice) or (
   3808         isinstance(casted_key, abc.Iterable)
   3809         and any(isinstance(x, slice) for x in casted_key)
   3810     ):
   3811         raise InvalidIndexError(key)
-> 3812     raise KeyError(key) from err
   3813 except TypeError:
   3814     # If we have a listlike key, _check_indexing_error will raise
   3815     #  InvalidIndexError. Otherwise we fall through and re-raise
   3816     #  the TypeError.
   3817     self._check_indexing_error(key)

KeyError: 0
```
If I `compute` before the `[0]`, then I get:
```
In [6]: df['a'].sum().to_series().fillna(0).compute()
Out[6]: 
0    6
dtype: int64
```
so I'd have expected the `[0]` to work

**Anything else we need to know?**:
spotted in Narwhals

**Environment**:

- Dask version: 2025.1.0
- Python version: 3.12
- Operating System: linux
- Install method (conda, pip, source): pip


Pull request:
Fix Projection logic for Series objects
- [ ] Closes #11746
- [ ] Tests added / passed
- [ ] Passes `pre-commit run --all-files`


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..7093e7235 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,
 )
@@ -3867,7 +3867,9 @@ def plain_column_projection(expr, parent, dependents, additional_columns=None):
         # we are accesing the index
         column_union = []
 
-    if column_union == expr.frame.columns:
+    if column_union == expr.frame.columns or not column_union and expr.ndim < 2:
+        # this projection is for the index, but the elements are unknown, so
+        # don't project
         return
     result = type(expr)(expr.frame[column_union], *expr.operands[1:])
     if column_union == parent.operand("columns"):
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/dask_expr/io/io.py b/dask/dataframe/dask_expr/io/io.py
index e0edf672d..ab416156f 100644
--- a/dask/dataframe/dask_expr/io/io.py
+++ b/dask/dataframe/dask_expr/io/io.py
@@ -445,6 +445,8 @@ class FromPandas(PartitionsFiltered, BlockwiseIO):
             try:
                 return list(self.frame.columns)
             except AttributeError:
+                if self.ndim == 1:
+                    return [self.name]
                 return []
         else:
             return _convert_to_list(columns_operand)
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/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 958c6e13e..f086d4405 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -125,7 +125,8 @@ jobs:
         # ensure this runs even if pytest fails
         if: >
           always() &&
-          (steps.run_tests.outcome == 'success' || steps.run_tests.outcome == 'failure')
+          (steps.run_tests.outcome == 'success' || steps.run_tests.outcome == 'failure') &&
+          matrix.array-expr == 'false'
         uses: actions/upload-artifact@v4
         with:
           name: ${{ matrix.os }}-${{ matrix.environment }}-${{ matrix.extra }}
diff --git a/dask/dataframe/dask_expr/tests/test_collection.py b/dask/dataframe/dask_expr/tests/test_collection.py
index a4d2d15c4..af5bba927 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
@@ -2785,3 +2784,11 @@ def test_to_backend_simplify():
         assert str(df2.expr) != str(df[["y"]].expr)
         df3 = df2.simplify()
         assert str(df3.expr) == str(df[["y"]].expr)
+
+
+def test_projection_on_series():
+    pdf = pd.DataFrame(data={"a": [1, 3, 2]}).a
+    df = from_pandas(pdf, npartitions=1)
+    result = df.replace(1, 5)
+    assert_eq(result, pdf.replace(1, 5))
+    assert_eq(result.fillna(0), pdf.replace(1, 5).fillna(0))
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