Repository: ray-project/ray
Pull request: [Data][Fix] Adding a column to pandas dataset
<!-- Thank you for your contribution! Please review https://github.com/ray-project/ray/blob/master/CONTRIBUTING.rst before opening a pull request. -->

<!-- Please add a reviewer to the assignee section when you create a PR. If you don't have the access to it, we will shortly find a reviewer and assign them to your PR. -->

## Why are these changes needed?
The index of the column must be set to align with the index of the batch when adding a new column to a pandas dataset. Otherwise it will result in a `Nan`. 
<!-- Please give a short summary of the change and the problem this solves. -->

## Related issue number
Closes #51758
<!-- For example: "Closes #1234" -->

## Checks

- [x] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR.
- [x] I've run `scripts/format.sh` to lint the changes in this PR.
- [ ] I've included any doc changes needed for https://docs.ray.io/en/master/.
    - [ ] I've added any new APIs to the API Reference. For example, if I added a
           method in Tune, I've added it in `doc/source/tune/api/` under the
           corresponding `.rst` file.
- [x] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
- Testing Strategy
   - [x] Unit tests
   - [ ] Release tests
   - [ ] This PR is not tested :(


Problem statement:
[Data] Found Bug `add_column`
### 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

```python
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.

Hints from issue discussion:
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)```

Commit message: fix add column to pandas

Signed-off-by: zhaoch23 <c233zhao@uwaterloo.ca>

Patch to review:
diff --git a/python/ray/data/dataset.py b/python/ray/data/dataset.py
index d82f112c65d7..b3658135c220 100644
--- a/python/ray/data/dataset.py
+++ b/python/ray/data/dataset.py
@@ -789,6 +789,9 @@ def add_column(
         def add_column(batch: DataBatch) -> DataBatch:
             column = fn(batch)
             if batch_format == "pandas":
+                # The index of the column must be set
+                # to align with the index of the batch.
+                column.index = batch.index
                 batch.loc[:, col] = column
                 return batch
             elif batch_format == "pyarrow":
@@ -5129,9 +5132,9 @@ def to_arrow_refs(self) -> List[ObjectRef["pyarrow.Table"]]:
         import pyarrow as pa
 
         ref_bundles: Iterator[RefBundle] = self.iter_internal_ref_bundles()
-        block_refs: List[
-            ObjectRef["pyarrow.Table"]
-        ] = _ref_bundles_iterator_to_block_refs_list(ref_bundles)
+        block_refs: List[ObjectRef["pyarrow.Table"]] = (
+            _ref_bundles_iterator_to_block_refs_list(ref_bundles)
+        )
         # Schema is safe to call since we have already triggered execution with
         # iter_internal_ref_bundles.
         schema = self.schema(fetch_if_missing=True)
diff --git a/python/ray/data/tests/test_pandas.py b/python/ray/data/tests/test_pandas.py
index 81a6c6eea8d5..7d1b5244f7e0 100644
--- a/python/ray/data/tests/test_pandas.py
+++ b/python/ray/data/tests/test_pandas.py
@@ -215,6 +215,17 @@ def test_read_pandas_data_array_column(ray_start_regular_shared):
     assert all(row["array"] == [1, 1, 1])
 
 
+def test_add_column_to_pandas(ray_start_regular_shared):
+    # Refer to issue https://github.com/ray-project/ray/issues/51758
+    ds = ray.data.from_pandas(pd.DataFrame({"a": list(range(20))}))
+
+    ds = ds.add_column(
+        "foo", lambda df: pd.Series([1] * len(df)), batch_format="pandas"
+    )
+    for row in ds.iter_rows():
+        assert row["foo"] == 1
+
+
 if __name__ == "__main__":
     import sys
 


Code review comments:
diff_hunk: @@ -215,6 +215,17 @@ def test_read_pandas_data_array_column(ray_start_regular_shared):
     assert all(row["array"] == [1, 1, 1])
 
 
+def test_add_column_to_pandas(ray_start_regular_shared):
+    # Refer to issue https://github.com/ray-project/ray/issues/51758
+    ds = ray.data.from_pandas(pd.DataFrame({"a": list(range(20))}))
path: python/ray/data/tests/test_pandas.py
text: @user1:
The current test works on master, so it's not testing the failure case -- , can you set override_num_blocks=2