fix: fix ray lance sink error

[Jay-ju]

#3229
ray-project/ray#49214

[westonpace]

Ah, does this work on older versions of Ray as well? It's probably not a deal breaker if it doesn't (we can require a certain version of Ray) but I'm curious to know how we should document this.

[Jay-ju]

Yes, the changes here are compatible with older versions of Ray. So there is no need for a separate documentation. However, before the code I submitted to the Ray community is merged, only versions prior to Ray 2.38.0 can be used.

[chenkovsky]

Hi @Jay-ju @westonpace ray committer has created a PR ray-project/ray#49251 . maybe fix should be based on that PR.

[westonpace]

I've tried this locally. When I run the tests in python/python/tests/test_ray.py I get the following error:

========================================================================================== FAILURES ==========================================================================================
_______________________________________________________________________________________ test_ray_sink ________________________________________________________________________________________

tmp_path = PosixPath('/tmp/pytest-of-pace/pytest-4/test_ray_sink0')

    def test_ray_sink(tmp_path: Path):
        schema = pa.schema([pa.field("id", pa.int64()), pa.field("str", pa.string())])
    
        sink = LanceDatasink(tmp_path)
        ray.data.range(10).map(
            lambda x: {"id": x["id"], "str": f"str-{x['id']}"}
>       ).write_datasink(sink)

python/tests/test_ray.py:32: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
../../../miniconda3/envs/lance/lib/python3.10/site-packages/ray/data/dataset.py:3857: in write_datasink
    datasink.on_write_complete(raw_write_results)
python/lance/ray/sink.py:143: in on_write_complete
    write_results = [
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

.0 = <list_iterator object at 0x78db2e2ff8b0>

    write_results = [
>       result["write_result"].iloc[0]["result"] for result in write_results
    ]
E   TypeError: 'WriteResult' object is not subscriptable

python/lance/ray/sink.py:144: TypeError

I did some inspecting. result is indeed a pandas dataframe. It has 1 column, 1 row, 1 cell. The value in that cell is an object of type ray.data.datasource.datasink.WriteResult. This matches https://docs.ray.io/en/latest/data/api/doc/ray.data.Datasink.on_write_complete.html

A WriteResult object is not subscriptable (e.g. can't call ["result"] on it). From what I see in https://github.com/ray-project/ray/blob/4f1a30852f2d62b5d7dfae50bddcefbee9dd406b/python/ray/data/datasource/datasink.py#L15 there are only two properties on WriteResult which are num_rows and size_bytes and this is not enough information to commit the write.

I agree with @chenkovsky that we will need to wait for ray-project/ray#49251 to be released.

[westonpace]

(I tested with both ray==2.38 and ray==2.40)

[Jay-ju]

@westonpace I have made changes based on the latest code of ray. Please review it again.

[westonpace]

Thanks for continuing to work on this! I did confirm that all tests pass with this change using the nightly release from Ray and the tests continue to pass with older version (2.37) so i think this is good.

I have a few minor warning message / comment suggestions and then we can merge this. Please re-request review when you've been able to address them.

[Jay-ju]

@westonpace please help me review this PR again? tks