Map_groups chaining bug?

Ray Libraries (Data, Train, Tune, Serve)
1

I am not sure whether this is a bug or lack of my understanding but it is reproducible. I created this example to show case the problem and I would like to do something similar in my code.

Ray 2.3.0

import ray
import pandas as pd

ds = ray.data.from_items([
    {"A": x % 3, "B": x} for x in range(100)]).groupby(
    "A")

ds.count().show()

## The goal is to call ds.map_groups(fnA).map_groups(fnB).write_parquet(s3_loc)
## 
def doub(df: pd.DataFrame) -> pd.DataFrame:
    df = df.copy()
    df["X"] = df["B"] * 2
    return df


def half(df: pd.DataFrame) -> pd.DataFrame:
    df = df.copy()
    df["X"] = df["B"] * 0.5
    return df


def doNothing(df: pd.DataFrame) -> pd.DataFrame:
    df = df.copy()
    df["X"] = df["B"] * 0.5 * 2
    return df

#
unpack_df = ds.map_groups(lambda x: doub(x), batch_format="pandas")
print(f"CHARU unpack_df type - {type(unpack_df)}")
print(f"CHARU unpack_df count - {unpack_df.show()}")
print(f"CHARU unpack_df print - {unpack_df}")

unpack_df2 = unpack_df
half_df = unpack_df2.groupby("A").map_groups(lambda x: half(x), batch_format="pandas")

print(f"CHARU half unpack_df type - {type(half_df)}")
print(f"CHARU half unpack_df count - {half_df.show()}")
print(f"CHARU half unpack_df print - {half_df}")

# It works when all the code in one function proving the method chaining may be an issue.
do_nothing_df = ds.map_groups(lambda x: doNothing(x), batch_format="pandas")

print(f"CHARU half unpack_df type - {type(do_nothing_df)}")
print(f"CHARU half unpack_df count - {do_nothing_df.show()}")
print(f"CHARU half unpack_df print - {do_nothing_df}")