[data] Monotonically increasing id

[rishic3]

Description

Implements monotonically increasing ID expression. This closely follows the Spark implementation https://github.com/apache/spark/blob/9bbdc0743034b40a904ca87a08da4e0bf2b1386c/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/MonotonicallyIncreasingID.scala.

Example usage:

from ray.data.expressions import monotonically_increasing_id

ds = ray.data.range(100)
ds = ds.with_column("uid", monotonically_increasing_id())
train, test = ds_with_id.streaming_train_test_split(test_size=0.25, split_type="hash", hash_column="uid")

Related issues

Closes #57806

[gemini-code-assist]

Code Review

This pull request introduces add_unique_id to add a monotonically increasing ID column to a dataset, which is a valuable addition. The implementation is clear and follows the described logic from Spark. However, I've identified a significant correctness issue in the handling of pyarrow.Table objects, which results in creating a nested list column instead of a flat integer column. I've also provided suggestions to improve efficiency, remove unreachable code, and enhance the clarity of the documentation. The tests are well-structured, but they may not be catching the pyarrow bug I've pointed out.

[richardliaw]

awesome! thanks a bunch for this contribution. will let @gvspraveen find someone to shepherd this.

[gvspraveen]

Thanks for the contribution. @bveeramani to shepherd this.

[richardliaw]

hey @rishic3, thanks a bunch for the contribution! if you're interested in chatting more with the contributors, feel free to join our community sync -- https://docs.google.com/forms/d/e/1FAIpQLSeYWjNExnr6gbhO5rpM0i6wm4TBTdsm3y5S0LR8Syzk_2gelQ/viewform

[rishic3]

Previous commit implemented as a dataset method (see e0fc987). Pushed an update that instead implements as an expression as per original issue intended.