Make write an operator as part of the execution plan

[jianoaix]

Why are these changes needed?

Today the "write" is not captured as an operator in the execution plan: it just fully executes the plan and then launches Ray tasks to perform the writing.

By make it a write operator, we will have a cleaner/unified model to handle how "write" is supported, enable query optimization for write (e.g. read->map->write can be fused into one operator node) for better performance.

Note: supporting this write operator in the new query plan will be in a follow-up PR.

Benchmarked on a single node:

import ray
import time

def make_ds(size_gb: int, parallelism: int):
    # Dataset of 10KiB tensor records.
    total_size = 1024 * 1024 * 1024 * size_gb
    record_dim = 1280
    record_size = record_dim * 8
    num_records = int(total_size / record_size)
    dataset = ray.data.range_tensor(num_records, shape=(record_dim,), parallelism=parallelism)
    print("Created dataset", dataset, "of size", dataset.size_bytes())
    return dataset

sizes = [1, 5, 10]
parallelism = [1, 5, 10, 20, 100]

perf = {}

for sz in sizes:
    for p in parallelism:
        path = f"/tmp/write/test-{sz}-{p}"
        time_start = time.perf_counter()
        ds = make_ds(sz, p).map_batches(lambda x: x).write_numpy(path)
        perf[path] = time.perf_counter() - time_start

print(perf)

Before:

{'/tmp/write/test-1-1': 9.953020322000157, '/tmp/write/test-1-5': 3.428301415000533, '/tmp/write/test-1-10': 4.723238538999794, '/tmp/write/test-1-20': 5.163609835999523, '/tmp/write/test-1-100': 4.33488237399979, '/tmp/write/test-5-1': 47.80200573299953, '/tmp/write/test-5-5': 57.3792794459996, '/tmp/write/test-5-10': 21.089490562999345, '/tmp/write/test-5-20': 49.24927898700025, '/tmp/write/test-5-100': 47.518018443000074, '/tmp/write/test-10-1': 182.66500996800005, '/tmp/write/test-10-5': 139.72414283000035, '/tmp/write/test-10-10': 129.0919650550004, '/tmp/write/test-10-20': 109.1896361629997, '/tmp/write/test-10-100': 61.62153517600018}

real    14m35.661s
user    0m13.973s
sys     0m2.241s

After:

{'/tmp/write/test-1-1': 14.615472456999669, '/tmp/write/test-1-5': 3.4766328169998815, '/tmp/write/test-1-10': 2.675365926000268, '/tmp/write/test-1-20': 2.3434579160002613, '/tmp/write/test-1-100': 2.7332164729996293, '/tmp/write/test-5-1': 29.679511864000233, '/tmp/write/test-5-5': 18.722395105999567, '/tmp/write/test-5-10': 18.884924998000315, '/tmp/write/test-5-20': 20.06004047999977, '/tmp/write/test-5-100': 18.469305754000743, '/tmp/write/test-10-1': 160.16473960199983, '/tmp/write/test-10-5': 47.103051861000495, '/tmp/write/test-10-10': 39.20541486499951, '/tmp/write/test-10-20': 38.90085198700035, '/tmp/write/test-10-100': 40.142852582999694}

real    7m39.632s
user    0m9.006s
sys     0m2.634s

The write fusion generally makes it faster. And it took only half of the time with fusion to complete this entire benchmark.

Related issue number

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