[Data] Add num_rows_per_file parameter to file-based writes

[bveeramani]

Why are these changes needed?

To avoid creating too many files, you might want write more rows to each file. But, there's no way to do so without using hacks. This PR fixes the issue by adding a num_rows_per_file parameter.

Related issue number

Fixes #41219

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

[raulchen]

Looks good on the high level. Please ping me for reviewing again when the PR is ready.

[raulchen]

The number of rows to write to each file.
Add a note that this number is not accurate? The actual number of rows will be slightly larger.

If None, Ray Data writes a system-chosen number of rows to each file.
We can say the size is based on the max target block size by default.

[bveeramani]

It's typically less, right? Unless each block contains more rows than the number specified?

We can say the size is based on the max target block size by default.

Not strongly opposed, but this is kinda leaking our implementation? Not sure if we want to expose this as part of the interface.

[raulchen]

It's typically less, right? Unless each block contains more rows than the number specified?

why is it less? IIRC, the bundler will accumulate input blocks until the total number of rows is more than the target.

Not strongly opposed, but this is kinda leaking our implementation? Not sure if we want to expose this as part of the interface.

agreed not to leak implementation. I felt "system-chosen" was too vague, and wanted to suggest something like "the file size will be based on the upstream outputs" in the beginning, which leaked more details. then I changed to "based on max target size". It's still leaking implementation to some extend though, but doesn't seem too bad because the block size is already a public api.
I don't have strong preference either. it's fine to keep it as it is.

[bveeramani]

why is it less? IIRC, the bundler will accumulate input blocks until the total number of rows is more than the target.

Not exactly. If an input block would cause the number of rows to exceed the target, it's excluded.

Here's a simple example: you add two two-row bundles to a bundler with a target of three rows. When you get next bundle, there are two rows instead of four.

import pyarrow as pa

import ray
from ray.data._internal.execution.interfaces import RefBundle
from ray.data._internal.execution.operators.map_operator import _BlockRefBundler
from ray.data.block import BlockAccessor

table = pa.Table.from_pylist([{"spam": 0}, {"spam": 1}])
metadata = BlockAccessor.for_block(table).get_metadata(None, None)
bundle = RefBundle(((ray.put(table), metadata),), owns_blocks=True)


bundler = _BlockRefBundler(3)
bundler.add_bundle(bundle)
bundler.add_bundle(bundle)
assert bundler.has_bundle()
print(bundler.get_next_bundle().num_rows())  # 2