Repository: ray-project/ray
Issue category: Bug Report

[Data] Ray 2.40 `add_column` Erroneous Assertion
### What happened + What you expected to happen

Ray 2.40 introduces an assertion that a value is of type `pd.Series`, but this is erroneous. A dataframe may have a constant assigned as in the following example or be a geopandas series (which should subclass).

```python
>>> import pandas as pd
>>>
>>> # Create a sample DataFrame
...     "A": [1, 2, 3],
...     "B": [4, 5, 6]
... })
>>>
>>> # Assign a constant value to a new column "C"
>>> df["C"] = 10
>>>
>>> print(df)
   A  B   C
0  1  4  10
1  2  5  10
2  3  6  10
>>>
```

As a result of the change valid expressions like:
```python
q_ds = q_ds.add_column("_q", lambda _: True)
s_ds = s_ds.add_column("_q", lambda _: False)
```
no longer succeed.

### Versions / Dependencies

Ray 2.40

### Reproduction script
```python
import pandas as pd
import ray

df= pd.DataFrame({"geometry": [1, 2, 3]})

# Create Ray datasets
ds = ray.data.from_pandas(df)

ds = ds.add_column("works", lambda _: True)
ds = ds.add_column("works", lambda _: True)

# evaluates and fails
ds.materialize()
ds.show()
```
### Issue Severity

Low: It annoys or frustrates me.


Functions to edit:
python/ray/data/dataset.py:Dataset.add_column

Patch:
diff --git a/python/ray/data/dataset.py b/python/ray/data/dataset.py
index d2e24d39cfa8..75f98de9f2e4 100644
--- a/python/ray/data/dataset.py
+++ b/python/ray/data/dataset.py
@@ -5,7 +5,6 @@
 import logging
 import time
 import warnings
-from collections.abc import Sequence
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -765,20 +764,9 @@ def add_column(
                 f"got: {batch_format}"
             )
 
-        def _raise_duplicate_column_error(col: str):
-            raise ValueError(f"Trying to add an existing column with name {col!r}")
-
         def add_column(batch: DataBatch) -> DataBatch:
             column = fn(batch)
             if batch_format == "pandas":
-                import pandas as pd
-
-                assert isinstance(column, (pd.Series, Sequence)), (
-                    f"For pandas batch format, the function must return a pandas "
-                    f"Series or sequence, got: {type(column)}"
-                )
-                if col in batch:
-                    _raise_duplicate_column_error(col)
                 batch.loc[:, col] = column
                 return batch
             elif batch_format == "pyarrow":
@@ -797,10 +785,9 @@ def add_column(batch: DataBatch) -> DataBatch:
                 # which case we'll want to append it
                 column_idx = batch.schema.get_field_index(col)
                 if column_idx == -1:
-                    # Append the column to the table
                     return batch.append_column(col, column)
                 else:
-                    _raise_duplicate_column_error(col)
+                    return batch.set_column(column_idx, col, column)
 
             else:
                 # batch format is assumed to be numpy since we checked at the
@@ -809,8 +796,6 @@ def add_column(batch: DataBatch) -> DataBatch:
                     f"For numpy batch format, the function must return a "
                     f"numpy.ndarray, got: {type(column)}"
                 )
-                if col in batch:
-                    _raise_duplicate_column_error(col)
                 batch[col] = column
                 return batch
 

Test patch: diff --git a/python/ray/data/tests/test_map.py b/python/ray/data/tests/test_map.py
index 0db0681d8d61..a29e5925815f 100644
--- a/python/ray/data/tests/test_map.py
+++ b/python/ray/data/tests/test_map.py
@@ -350,15 +350,11 @@ def test_add_column(ray_start_regular_shared):
     )
     assert ds.take(1) == [{"id": 0, "foo": 1}]
 
-    # Adding a column that is already there should result in an error
-    with pytest.raises(
-        ray.exceptions.UserCodeException,
-        match="Trying to add an existing column with name 'id'",
-    ):
-        ds = ray.data.range(5).add_column(
-            "id", lambda x: pc.add(x["id"], 1), batch_format="pyarrow"
-        )
-        assert ds.take(2) == [{"id": 1}, {"id": 2}]
+    # Adding a column that is already there should not result in an error
+    ds = ray.data.range(5).add_column(
+        "id", lambda x: pc.add(x["id"], 1), batch_format="pyarrow"
+    )
+    assert ds.take(2) == [{"id": 1}, {"id": 2}]
 
     # Adding a column in the wrong format should result in an error
     with pytest.raises(
@@ -378,15 +374,11 @@ def test_add_column(ray_start_regular_shared):
     )
     assert ds.take(1) == [{"id": 0, "foo": 1}]
 
-    # Adding a column that is already there should result in an error
-    with pytest.raises(
-        ray.exceptions.UserCodeException,
-        match="Trying to add an existing column with name 'id'",
-    ):
-        ds = ray.data.range(5).add_column(
-            "id", lambda x: np.add(x["id"], 1), batch_format="numpy"
-        )
-        assert ds.take(2) == [{"id": 1}, {"id": 2}]
+    # Adding a column that is already there should not result in an error
+    ds = ray.data.range(5).add_column(
+        "id", lambda x: np.add(x["id"], 1), batch_format="numpy"
+    )
+    assert ds.take(2) == [{"id": 1}, {"id": 2}]
 
     # Adding a column in the wrong format should result in an error
     with pytest.raises(
@@ -402,23 +394,20 @@ def test_add_column(ray_start_regular_shared):
     ds = ray.data.range(5).add_column("foo", lambda x: x["id"] + 1)
     assert ds.take(1) == [{"id": 0, "foo": 1}]
 
-    # Adding a column that is already there should result in an error
-    with pytest.raises(
-        ray.exceptions.UserCodeException,
-        match="Trying to add an existing column with name 'id'",
-    ):
-        ds = ray.data.range(5).add_column("id", lambda x: x["id"] + 1)
-        assert ds.take(2) == [{"id": 1}, {"id": 2}]
+    # Adding a column that is already there should not result in an error
+    ds = ray.data.range(5).add_column("id", lambda x: x["id"] + 1)
+    assert ds.take(2) == [{"id": 1}, {"id": 2}]
 
-    # Adding a column in the wrong format should result in an error
-    with pytest.raises(
-        ray.exceptions.UserCodeException, match="For pandas batch format"
-    ):
+    # Adding a column in the wrong format may result in an error
+    with pytest.raises(ray.exceptions.UserCodeException):
         ds = ray.data.range(5).add_column(
-            "id", lambda x: np.array([1]), batch_format="pandas"
+            "id", lambda x: range(7), batch_format="pandas"
         )
         assert ds.take(2) == [{"id": 1}, {"id": 2}]
 
+    ds = ray.data.range(5).add_column("const", lambda _: 3, batch_format="pandas")
+    assert ds.take(2) == [{"id": 0, "const": 3}, {"id": 1, "const": 3}]
+
     with pytest.raises(ValueError):
         ds = ray.data.range(5).add_column("id", 0)