[Data] Groupby aggregation (min/max/count) hangs when the group column has NaN

[wingkitlee0]

What happened + What you expected to happen

Groupby min/max/count etc will hang (with increasing number of tasks) when the group column contains np.nan. For example,

ds = ray.data.from_items([1.0, 1.0, 2.0, np.nan])
ds.groupby("item").count().take_all()

However, map_groups would work:

ds.groupby("item").map_groups(lambda x: {"count": [len(x["item"])]}).take_all()

String column with None also work.

On Pandas, depending on the dropna=True/False, the np.nan will be treated as a separate group.

Versions / Dependencies

ray 2.40.0
numpy 2.1.3

Reproduction script

ds = ray.data.from_items([1.0, 1.0, 2.0, np.nan])
ds.groupby("item").count().take_all()

Issue Severity

Medium: It is a significant difficulty but I can work around it.

[richardliaw]

Great catch! Let me know if you can find the root cause. It's probably in our BlockAccessor or transform functions code.

[richardliaw]

Otherwise we may try to get to this in the next couple of weeks.