From fdf5beabc5d64555b7e6c4e0f3864ec25dd4f3df Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Mon, 12 Sep 2022 20:30:54 -0500 Subject: [PATCH] Fix assigning pandas Series to column when `ddf.columns.min()` raises --- dask/dataframe/core.py | 13 ++++++++----- dask/dataframe/tests/test_dataframe.py | 8 ++++++++ 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py index 184ef693802..00c5a167486 100644 --- a/dask/dataframe/core.py +++ b/dask/dataframe/core.py @@ -6017,15 +6017,18 @@ def is_broadcastable(dfs, s): """ This Series is broadcastable against another dataframe in the sequence """ + + def compare(s, df): + try: + return s.divisions == (df.columns.min(), df.columns.max()) + except TypeError: + return False + return ( isinstance(s, Series) and s.npartitions == 1 and s.known_divisions - and any( - s.divisions == (df.columns.min(), df.columns.max()) - for df in dfs - if isinstance(df, DataFrame) - ) + and any(compare(s, df) for df in dfs if isinstance(df, DataFrame)) ) diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py index 932abfa0bf1..a2e16cadf00 100644 --- a/dask/dataframe/tests/test_dataframe.py +++ b/dask/dataframe/tests/test_dataframe.py @@ -1623,6 +1623,14 @@ def test_assign_dtypes(): ) +def test_assign_pandas_series(): + # Make sure we handle when `df.columns.min()` raises TypeError + df = pd.DataFrame({"a": [1, 2], 1: [5, 6]}) + ddf = dd.from_pandas(df, npartitions=2) + ddf = ddf.assign(c=df["a"]) + assert_eq(ddf, df.assign(c=df["a"])) + + def test_map(): df = pd.DataFrame( {"a": range(9), "b": [4, 5, 6, 1, 2, 3, 0, 0, 0]},