IntCastingNaNError after outer join with int64 column

[NickSchouten]

Describe the issue:
An outer join with an int64 column, will keep the column as int64 even though it may have introduced NaNs.

Minimal Complete Verifiable Example:

from dask.distributed import Client
from dask.distributed import LocalCluster
import dask.dataframe as dd

cluster = LocalCluster()
client = Client(cluster)
client.cluster.scale(1)


dask_df = dd.from_dict(
    {
        "a": [1, 2],
    },
    npartitions=1,
)
dask_df2 = dd.from_dict(
    {
        "b": [1],
    },
    npartitions=1,
)
print(dask_df2.b.dtype) # int64
df = dask_df.join(dask_df2, how="left")
print(df.b.dtype) # int64, this would previously be a float
df.shuffle(on="a").compute() # causes IntCastingNaNError

Anything else we need to know?:

In a previous version (dask 2022.5.2 don't know which version the change first occurred) this same code would output a float column and therefore not throw an error.

Is this expected behaviour?

The problem may of course be prevented by first casting b to "Int64" (pandas nullable) or "float" manually.

Environment:

• Dask version:
  • dask 2023.9.1
  • dask-bigquery 2023.5.1
  • dask-glm 0.2.0
  • dask-kubernetes 2023.9.0
  • dask_labextension 7.0.0
  • dask-ml 2023.3.24
• Python version: Python 3.10.11
• Operating System: Linux
• Install method: conda

[jrbourbeau]

Thanks for the issue and nice minimal example @NickSchouten, I'm able to reproduce. For others, here's the full traceback

Details:

---------------------------------------------------------------------------
IntCastingNaNError                        Traceback (most recent call last)
File ~/projects/dask/distributed/distributed/shuffle/_core.py:342, in handle_transfer_errors()
    341 try:
--> 342     yield
    343 except ShuffleClosedError:

File ~/projects/dask/distributed/distributed/shuffle/_shuffle.py:63, in shuffle_transfer()
     62 with handle_transfer_errors(id):
---> 63     return get_worker_plugin().add_partition(
     64         input,
     65         input_partition,
     66         spec=DataFrameShuffleSpec(
     67             id=id,
     68             npartitions=npartitions,
     69             column=column,
     70             meta=meta,
     71             parts_out=parts_out,
     72         ),
     73     )

File ~/projects/dask/distributed/distributed/shuffle/_worker_plugin.py:140, in add_partition()
    139 shuffle = self.get_or_create_shuffle(spec)
--> 140 return sync(
    141     self.worker.loop,
    142     shuffle.add_partition,
    143     data=data,
    144     partition_id=partition_id,
    145     **kwargs,
    146 )

File ~/projects/dask/distributed/distributed/utils.py:426, in sync()
    425     typ, exc, tb = error
--> 426     raise exc.with_traceback(tb)
    427 else:

File ~/projects/dask/distributed/distributed/utils.py:399, in f()
    398     future = asyncio.ensure_future(future)
--> 399     result = yield future
    400 except Exception:

File ~/mambaforge/envs/dask-py39/lib/python3.9/site-packages/tornado/gen.py:767, in run()
    766 try:
--> 767     value = future.result()
    768 except Exception as e:
    769     # Save the exception for later. It's important that
    770     # gen.throw() not be called inside this try/except block
    771     # because that makes sys.exc_info behave unexpectedly.

File ~/projects/dask/distributed/distributed/shuffle/_core.py:217, in add_partition()
    216     raise RuntimeError(f"Cannot add more partitions to {self}")
--> 217 return await self._add_partition(data, partition_id)

File ~/projects/dask/distributed/distributed/shuffle/_shuffle.py:468, in _add_partition()
    466     return out
--> 468 out = await self.offload(_)
    469 await self._write_to_comm(out)

File ~/projects/dask/distributed/distributed/shuffle/_core.py:117, in offload()
    116 with self.time("cpu"):
--> 117     return await asyncio.get_running_loop().run_in_executor(
    118         self.executor,
    119         func,
    120         *args,
    121     )

File ~/mambaforge/envs/dask-py39/lib/python3.9/concurrent/futures/thread.py:58, in run()
     57 try:
---> 58     result = self.fn(*self.args, **self.kwargs)
     59 except BaseException as exc:

File ~/projects/dask/distributed/distributed/shuffle/_shuffle.py:459, in _()
    458 def _() -> dict[str, tuple[int, bytes]]:
--> 459     out = split_by_worker(
    460         data,
    461         self.column,
    462         self.meta,
    463         self.worker_for,
    464     )
    465     out = {k: (partition_id, serialize_table(t)) for k, t in out.items()}

File ~/projects/dask/distributed/distributed/shuffle/_shuffle.py:282, in split_by_worker()
    280 from dask.dataframe.dispatch import to_pyarrow_table_dispatch
--> 282 df = df.astype(meta.dtypes, copy=False)
    284 # (cudf support) Avoid pd.Series

File ~/mambaforge/envs/dask-py39/lib/python3.9/site-packages/pandas/core/generic.py:6305, in astype()
   6304 try:
-> 6305     res_col = col.astype(dtype=cdt, copy=copy, errors=errors)
   6306 except ValueError as ex:

File ~/mambaforge/envs/dask-py39/lib/python3.9/site-packages/pandas/core/generic.py:6324, in astype()
   6322 else:
   6323     # else, only a single dtype is given
-> 6324     new_data = self._mgr.astype(dtype=dtype, copy=copy, errors=errors)
   6325     return self._constructor(new_data).__finalize__(self, method="astype")

File ~/mambaforge/envs/dask-py39/lib/python3.9/site-packages/pandas/core/internals/managers.py:451, in astype()
    449     copy = False
--> 451 return self.apply(
    452     "astype",
    453     dtype=dtype,
    454     copy=copy,
    455     errors=errors,
    456     using_cow=using_copy_on_write(),
    457 )

File ~/mambaforge/envs/dask-py39/lib/python3.9/site-packages/pandas/core/internals/managers.py:352, in apply()
    351 else:
--> 352     applied = getattr(b, f)(**kwargs)
    353 result_blocks = extend_blocks(applied, result_blocks)

File ~/mambaforge/envs/dask-py39/lib/python3.9/site-packages/pandas/core/internals/blocks.py:511, in astype()
    509 values = self.values
--> 511 new_values = astype_array_safe(values, dtype, copy=copy, errors=errors)
    513 new_values = maybe_coerce_values(new_values)

File ~/mambaforge/envs/dask-py39/lib/python3.9/site-packages/pandas/core/dtypes/astype.py:242, in astype_array_safe()
    241 try:
--> 242     new_values = astype_array(values, dtype, copy=copy)
    243 except (ValueError, TypeError):
    244     # e.g. _astype_nansafe can fail on object-dtype of strings
    245     #  trying to convert to float

File ~/mambaforge/envs/dask-py39/lib/python3.9/site-packages/pandas/core/dtypes/astype.py:187, in astype_array()
    186 else:
--> 187     values = _astype_nansafe(values, dtype, copy=copy)
    189 # in pandas we don't store numpy str dtypes, so convert to object

File ~/mambaforge/envs/dask-py39/lib/python3.9/site-packages/pandas/core/dtypes/astype.py:105, in _astype_nansafe()
    104 elif np.issubdtype(arr.dtype, np.floating) and is_integer_dtype(dtype):
--> 105     return _astype_float_to_int_nansafe(arr, dtype, copy)
    107 elif is_object_dtype(arr.dtype):
    108     # if we have a datetime/timedelta array of objects
    109     # then coerce to datetime64[ns] and use DatetimeArray.astype

File ~/mambaforge/envs/dask-py39/lib/python3.9/site-packages/pandas/core/dtypes/astype.py:150, in _astype_float_to_int_nansafe()
    149 if not np.isfinite(values).all():
--> 150     raise IntCastingNaNError(
    151         "Cannot convert non-finite values (NA or inf) to integer"
    152     )
    153 if dtype.kind == "u":
    154     # GH#45151

IntCastingNaNError: Cannot convert non-finite values (NA or inf) to integer: Error while type casting for column 'b'

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

RuntimeError                              Traceback (most recent call last)
Cell In[1], line 25
     23 df = dask_df.join(dask_df2, how="left")
     24 print(df.b.dtype) # int64, this would previously be a float
---> 25 df.shuffle(on="a").compute() # causes IntCastingNaNError

File ~/projects/dask/dask/dask/base.py:342, in DaskMethodsMixin.compute(self, **kwargs)
    318 def compute(self, **kwargs):
    319     """Compute this dask collection
    320
    321     This turns a lazy Dask collection into its in-memory equivalent.
   (...)
    340     dask.compute
    341     """
--> 342     (result,) = compute(self, traverse=False, **kwargs)
    343     return result

File ~/projects/dask/dask/dask/base.py:628, in compute(traverse, optimize_graph, scheduler, get, *args, **kwargs)
    625     postcomputes.append(x.__dask_postcompute__())
    627 with shorten_traceback():
--> 628     results = schedule(dsk, keys, **kwargs)
    630 return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])

File ~/projects/dask/distributed/distributed/shuffle/_shuffle.py:63, in shuffle_transfer()
     53 def shuffle_transfer(
     54     input: pd.DataFrame,
     55     id: ShuffleId,
   (...)
     60     parts_out: set[int],
     61 ) -> int:
     62     with handle_transfer_errors(id):
---> 63         return get_worker_plugin().add_partition(
     64             input,
     65             input_partition,
     66             spec=DataFrameShuffleSpec(
     67                 id=id,
     68                 npartitions=npartitions,
     69                 column=column,
     70                 meta=meta,
     71                 parts_out=parts_out,
     72             ),
     73         )

File ~/mambaforge/envs/dask-py39/lib/python3.9/contextlib.py:137, in __exit__()
    135     value = typ()
    136 try:
--> 137     self.gen.throw(typ, value, traceback)
    138 except StopIteration as exc:
    139     # Suppress StopIteration *unless* it's the same exception that
    140     # was passed to throw().  This prevents a StopIteration
    141     # raised inside the "with" statement from being suppressed.
    142     return exc is not value

File ~/projects/dask/distributed/distributed/shuffle/_core.py:346, in handle_transfer_errors()
    344     raise Reschedule()
    345 except Exception as e:
--> 346     raise RuntimeError(f"P2P shuffling {id} failed during transfer phase") from e

RuntimeError: P2P shuffling 8afd65d270df8043cff2838636f5304a failed during transfer phase

Since the 2022.5.2 release, we've updated our default shuffling method to a new p2p algorithm. It looks like

If I specify shuffle="tasks" (the previous default) in your example, then things run as expected

In [8]: df.shuffle(on="a", shuffle="tasks").compute()
Out[8]:
   a    b
0  1  1.0
1  2  NaN

cc @hendrikmakait @fjetter for visibility

@NickSchouten, until this is fixed, I'd recommend either casting (as you mentioned) or using shuffle="tasks"