[Data] Ray Data heap OOMs when using fused pandas

[bveeramani]

What happened + What you expected to happen

Ray Data map tasks buffer outputs until the buffer contains at least 128 MiB of data. However, Ray Data doesn't correctly count the size of some pandas DataFrames, so if you use the pandas batch format, Ray Data might buffer too much data and your task will heap OOM.

See the related #44577 for more information.

Versions / Dependencies

2.32

Reproduction script

This doesn't work.

import numpy as np

import ray

ray.init(num_cpus=1)


def yield_blocks(batch):
    while True:
        yield {"data": np.zeros((128, 1024, 1024), dtype=np.uint8)}


def convert_to_dataframe(batch):
    return batch


ds = (
    ray.data.range(1, override_num_blocks=1)
    .map_batches(yield_blocks, batch_size=None)
    .map_batches(convert_to_dataframe, batch_size=None, batch_format="pandas")
)
for _ in ds.iter_batches(batch_size=None):  # This should eventually heap OOM
    pass

This does work. For some reason, pandas' memory estimate is correct in this case but not the previous.

import numpy as np
import pandas as pd

import ray

ray.init(num_cpus=1)


def yield_blocks(block):
    while True:
        yield pd.DataFrame({"data": 128 * [np.zeros((1024, 1024), dtype=np.uint8)]})


ds = ray.data.range(1, override_num_blocks=1).map_batches(yield_blocks, batch_size=None)
for _ in ds.iter_batches(batch_size=None):
    pass

Issue Severity

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