[Data] Found Bug add_column

[ivanthewebber]

What happened + What you expected to happen

See reproduction which broke when updating from 2.40.0 to 2.44.0. It seems ray.data.Dataset.add_column is not setting for all elements properly.

Versions / Dependencies

Ray 2.44.0

Reproduction script

import unittest
import ray
import ray.data
import pandas as pd

class MyTest(unittest.TestCase):
    def test_add_column(self):
        def set_x(df: pd.DataFrame) -> pd.Series:
            return pd.Series([1] * len(df))
                
        df = pd.DataFrame({"a": [1, 2, 3]})

        # error does not present with 1 for override_num_blocks
        # second two are nan with 3 for override_num_blocks
        ds = ray.data.from_pandas(df, 2)

        for row in ds.add_column("x", set_x, batch_format="pandas").take(3):
            self.assertEqual(row["x"], 1) # last one is nan and fails

if __name__ == "__main__":
    unittest.main()

Issue Severity

High: It blocks me from completing my task.

[Bye-legumes]

I can reproduce on my side

FAIL: test_add_column (__main__.MyTest.test_add_column)
----------------------------------------------------------------------
Traceback (most recent call last):
 File "/home/zhilong/test_issue.py", line 18, in test_add_column
   self.assertEqual(row["x"], 1) # last one is nan and fails
   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError: nan != 1

----------------------------------------------------------------------
Ran 1 test in 3.777s

FAILED (failures=1)```