[Data] Add Datasource.on_write_start

[bveeramani]

Why are these changes needed?

Currently, we attempt to create a directory in every write task. This can cause rate limiting issues with S3. To address this problem, this PR adds a on_write_start method that is executed once per write job.

Related issue number

Fixes #36690

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 :(

[ericl]

Nice find!

…

On Wed, Aug 9, 2023, 10:35 PM Balaji Veeramani ***@***.***> wrote: @bveeramani <https://github.com/bveeramani> requested your review on: #38298 <#38298> [Data] Add Datasource.on_write_start as a code owner. — Reply to this email directly, view it on GitHub <#38298 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAADUSWH6RCKSGQNPYCLIKLXURXK3ANCNFSM6AAAAAA3K6RD4U> . You are receiving this because your review was requested.Message ID: ***@***.***>

[bveeramani]

@c21 Here's the currently failing test:

  | # Ensure that directory was created for only the non-empty dataset.
  | if fs is None:
  | assert not os.path.isdir(all_empty_path)
  | assert os.path.isdir(some_empty_path)
  | # Only files for the non-empty blocks should be created.
  | file_list = os.listdir(some_empty_path)
  | file_list.sort()
  | assert file_list == [
  | f"{some_empty_key}_00000{i}_000000.parquet" for i in range(2)
  | ]
  | else:
  | >           assert (
  | fs.get_file_info(_unwrap_protocol(all_empty_path)).type
  | == pa.fs.FileType.NotFound
  | )
  | E           assert <FileType.Directory: 3> == <FileType.NotFound: 0>
  | E             +<FileType.Directory: 3>
  | E             -<FileType.NotFound: 0>
  |  
  | python/ray/data/tests/test_parquet.py:861: AssertionError

[c21]

The change with on_write_start LGTM. As discussed offline, let's delete the staging directory if no file is written.