[Ray Data] repartition causes an error to be thrown for non-pyarrow compatible column types

[marwan116]

What happened + What you expected to happen

I have a simple pandas dataframe with a string column that contains numpy nans.

When I attempt to build a dataset from it and apply transformations it works as expected. However when I include a repartition call, it causes a TypeError to be thrown.

Versions / Dependencies

Version info:
3.10.11 (main, Dec 12 2023, 16:25:48) [Clang 15.0.0 (clang-1500.0.40.1)]
Ray version:
2.12.0

Reproduction script

See this code snippet to reproduce

import ray.data
import pandas as pd
import numpy as np

ds = ray.data.from_pandas(pd.DataFrame({"a": [1, 2, 3], "b": ["a", "b", np.nan]}))


def my_fn(batch):
    return batch


ds.repartition(2).map_batches(my_fn).take_all()
print("The above will pass!")

ds.repartition(2).map_batches(my_fn).to_pandas()
print("The above will crash!")

Issue Severity

Low: It annoys or frustrates me.

[982945902]

This is not a bug, conservatively speaking， if you change ds = ray.data.from_pandas(pd.DataFrame({"a": [1, 2, 3], "b": ["a", "b", np.nan]})) to
ds = ray.data.from_pandas(pd.DataFrame({"a": [1, 2, 3], "b": ["a", "b", "c"]}))
is will pass.

why?
look that:

a return block will try cover to arrow first if failed will be pandas.
["a", "b", np.nan] cannot cover to arrow, because arrow require same type (type("a") != type(np.nan))

so in you case , first block cover to arrow , second cover to pandas, i will cause block join check raise exception.

you can do like this

ds.repartition(2).map_batches(my_fn,batch_format="pandas").to_pandas()
pass

[marwan116]

Well if you try without the .repartition, the code will pass without having to explicitly set batch_format="pandas"

i.e. see the below code:

import ray.data
import pandas as pd
import numpy as np

ds = ray.data.from_pandas(pd.DataFrame({"a": [1, 2, 3], "b": ["a", "b", np.nan]}))


def my_fn(batch):
    return batch

ds.map_batches(my_fn).to_pandas()
print("The above will pass!")

I like this behavior more because it expects less work from the user - but regardless of my preferences, I would expect the behavior to be consistent whether a repartition operation is applied or not.

[bveeramani]

Yeah, this is a bug.

After the repartition, the dataset contains two pandas blocks. However, map_batches converts the first block to an Arrow table and the keeps the second block as a pandas DataFrame.

import numpy as np
import pandas as pd

import ray

ds = ray.data.from_pandas(pd.DataFrame({"a": [1, 2, 3], "b": ["a", "b", np.nan]}))


def my_fn(batch):
    return batch


refs = ds.repartition(2).map_batches(my_fn).get_internal_block_refs()
print([type(ray.get(ref)) for ref in refs])
# <class 'pyarrow.lib.Table'>, <class 'pandas.core.frame.DataFrame'>

A possible fix is to convert unexpected types in DelegatingBlockBuilder rather than raise an error.