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

Issue:
[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.


Hints:
As a workaround you should be able to do `pd.Series(constant, index=df.index)`
Here is another issue with the function update; the [docstring says columns may be overridden](https://github.com/ray-project/ray/blob/8fe05532e7037aac1b9b3755f0053047dc0f1321/python/ray/data/dataset.py#L743C13-L744C39), but there is a [check raising an error if the column already exists](https://github.com/ray-project/ray/blob/8fe05532e7037aac1b9b3755f0053047dc0f1321/python/ray/data/dataset.py#L780C17-L781C55).

As a side note, after using `add_column` you can no longer get the `ds.columns()` without the `Dataset` being at least partially evaluated, but it should be possible to know the new column names as it is the same as before but with the additional column name provided as the first argument to `add_column`. My thought being it would be good to have a function that can add a column iff the column is not present without the cost of materializing the dataset like `if col not in ds.columns(): ds.add_column(col, func)`. However, this can also be accomplished with (pseudo-code lambda) `ds.map_batches(lambda df: (df[col] = func(df)) if col not in df else df)`.
@ivanthewebber good catch -- would you be willing to submit a PR?
If someone gets to it first I am finding other good first issues as well.

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)