[Data] Fix write_results type

[bveeramani]

Why are these changes needed?

Dataset.write_datasource passes a list[list[WriteResult]] to Datasource.on_write_complete instead of a list[WriteResult]. This PR fixes the bug.

Related issue number

Fixes #33935

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[c21]

I think we need to fix https://github.com/ray-project/ray/blob/master/python/ray/data/_internal/planner/write.py#L16:

from:

return [[datasource.write(blocks, ctx, **write_args)]]

to:

return [datasource.write(blocks, ctx, **write_args)]

@jianoaix - the above double [[]] looks like a mistake, right?

[bveeramani]

I think we need to fix https://github.com/ray-project/ray/blob/master/python/ray/data/_internal/planner/write.py#L16:

from:

return [[datasource.write(blocks, ctx, **write_args)]]

to:

return [datasource.write(blocks, ctx, **write_args)]

@jianoaix - the above double [[]] looks like a mistake, right?

I think it's correct as is.

Without the double brackets, we return an Iterable[WriteResult] instead of an Iterable[Block]. In this case, blocks are List[WriteResult].

[c21]

@bveeramani - ah sorry you are right. A Python list is a valid Block.

[jianoaix]

I think we need to fix https://github.com/ray-project/ray/blob/master/python/ray/data/_internal/planner/write.py#L16:
from:

return [[datasource.write(blocks, ctx, **write_args)]]

to:

return [datasource.write(blocks, ctx, **write_args)]

@jianoaix - the above double [[]] looks like a mistake, right?

I think it's correct as is.

Without the double brackets, we return an Iterable[WriteResult] instead of an Iterable[Block]. In this case, blocks are List[WriteResult].

Yep, the output of write op is a list of blocks with block type being a List.

[c21]

@bveeramani - can you also add a comment in https://github.com/ray-project/ray/blob/master/python/ray/data/_internal/planner/write.py ? Something like

# NOTE: return a nested list to conform the Block type (a Python list) 
return [[datasource.write(blocks, ctx, **write_args)]]

[c21]

cc @zhe-thoughts to approve, this is to fix a bug in Dataset write path. Would be low-risky surgical fix.

[zhe-thoughts]

Approved for merging into master