[data] `rename_columns` doesn't work with `from_pandas`

[richardliaw]

What happened + What you expected to happen

RayTaskError(UserCodeException): ray::Project() (pid=52375, ip=127.0.0.1)
  File "/Users/rliaw/miniconda3/envs/oss/lib/python3.10/site-packages/pandas/core/generic.py", line 6318, in __getattr__
    return object.__getattribute__(self, name)
AttributeError: 'DataFrame' object has no attribute 'rename_columns'

The above exception was the direct cause of the following exception:

ray::Project() (pid=52375, ip=127.0.0.1)
  File "/Users/rliaw/miniconda3/envs/oss/lib/python3.10/site-packages/ray/data/_internal/execution/operators/map_operator.py", line 564, in _map_task
    for b_out in map_transformer.apply_transform(iter(blocks), ctx):
  File "/Users/rliaw/miniconda3/envs/oss/lib/python3.10/site-packages/ray/data/_internal/execution/operators/map_transformer.py", line 102, in __call__
    yield from self._post_process(results)
  File "/Users/rliaw/miniconda3/envs/oss/lib/python3.10/site-packages/ray/data/_internal/execution/operators/map_transformer.py", line 84, in _shape_blocks
    for result in results:
  File "/Users/rliaw/miniconda3/envs/oss/lib/python3.10/site-packages/ray/data/_internal/execution/operators/map_transformer.py", line 412, in _apply_transform
    yield from self._block_fn(blocks, ctx)
  File "/Users/rliaw/miniconda3/envs/oss/lib/python3.10/site-packages/ray/data/_internal/planner/plan_udf_map_op.py", line 620, in transform_fn
    out_block = fn(block)
  File "/Users/rliaw/miniconda3/envs/oss/lib/python3.10/site-packages/ray/data/_internal/planner/plan_udf_map_op.py", line 151, in _project_block
    _try_wrap_udf_exception(e)
  File "/Users/rliaw/miniconda3/envs/oss/lib/python3.10/site-packages/ray/data/_internal/planner/plan_udf_map_op.py", line 439, in _try_wrap_udf_exception
    raise UserCodeException("UDF failed to process a data block.") from e
ray.exceptions.UserCodeException: UDF failed to process a data block.

Versions / Dependencies

master

Reproduction script

import ray
import pandas as pd
pd.DataFrame({"doc_id": range(100)})
df = pd.DataFrame({"doc_id": range(100)})
ray.data.from_pandas(df).rename_columns({"doc": "doc_id"}).materialize()

Issue Severity

None

[400Ping]

Hi, may I try this?

[richardliaw]

sure! I'll assign it to you; if you end up not finding time to do this, let's unassign and find another owner.

[400Ping]

ok, thanks.

[richardliaw]

hi @400Ping i'm going to unassign you since it's been a week -- @kyuds will be taking a look

[kyuds]

thank you @richardliaw !

[kyuds]

Happy to announce that this issue is already resolved! I did a detailed investigation, so I'll write it down.

First, the original renaming logic (at least as of when this issue occurred) was introduced in #49393. The original implementation had an inner function fn in plan_udf_map_op.py:: plan_project_op. The rename logic was as followed: the Block type can be either pyarrow.Table or pandas.DataFrame. When the Block was a pandas.DataFrame, it was actually converted to a pyarrow.Table (this was done for another purpose: for handling None type, refer to comments in the PR.)

Now, it is guaranteed that the Block will always internally be of type pyarrow.Table. The renaming logic then used pyarrow.Table's rename_columns api. This api belongs to pyarrow, and does not exist for pandas. However, this does not matter, because the Block will always be a pyarrow.Table.

Later on, in #52980, this automatic conversion to a pyarrow.Table is dropped, however, the rename logic is kept the same. Now, Block type can be both pyarrow.Table and pandas.DataFrame. Therefore, during the renaming logic, Ray will attempt to call rename_columns on pandas.DataFrame, which results in the error above.

However, around the time I started my initial investigation into this issue, @goutamvenkat-anyscale merged this PR: #57855, essentially revamping the entire way how we rename columns! Therefore, I'm happy to report back that renaming the columns from from_pandas works!!

btw, script provided in the initial issue is actually invalid:

import ray
import pandas as pd
pd.DataFrame({"doc_id": range(100)})
df = pd.DataFrame({"doc_id": range(100)})

# original (this errors because `doc` is not the original column name, the order of `original`: `rename` is switched)
ray.data.from_pandas(df).rename_columns({"doc": "doc_id"}).materialize()

# modified (this works by supplying the correct key, value order)
ray.data.from_pandas(df).rename_columns({"doc_id": "doc"}).materialize()

but by modifying the script, it runs successfully, and I confirmed that the rename also works as expected!

ps: while investigating, I realized that the rename_columns api for BlockAccessor and their inheritors is actually dead code. Should I submit a PR to remove it?

[richardliaw]

removing dead code seems reasonable! thanks for the investigation :)

[gvspraveen]

Thanks for detailed investigation @kyuds