[data] Optimization to reduce ArrowBlock building time for blocks of size 1

[stephanie-wang]

Why are these changes needed?

Many Data ops depend on converting numpy batches to Arrow blocks. A single np array -> pyarrow is normally zero-copy, but blocks with multiple rows will need a copy to make the column of np arrays into one contiguous ndarray. This PR avoids this step for blocks of a single row by using np.expand_dims to reshape the array instead of copying it.

Related issue number

Needed for #37474.

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

[stephanie-wang]

For a single-row block [{"field": 100MB np.array}]:

Before:
In [11]: %timeit b.build()
385 ms ± 8.88 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

After:
In [13]: %timeit b.build()
116 µs ± 3.59 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

[ericl]

Nice!

…

On Thu, Aug 24, 2023, 10:43 AM Stephanie Wang ***@***.***> wrote: For a single-row block [{"field": 100MB np.array}]: Before: In [11]: %timeit b.build() 385 ms ± 8.88 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) After: In [13]: %timeit b.build() 116 µs ± 3.59 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each) — Reply to this email directly, view it on GitHub <#38833 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAADUSTVV4C5L4HF7YBA77LXW6HCXANCNFSM6AAAAAA35LQNKM> . You are receiving this because your review was requested.Message ID: ***@***.***>

[c21]

LGTM

[zhe-thoughts]

Please verify tests-ok and then I will review / merge. Thanks!

[stephanie-wang]

Hey @zhe-thoughts this is ready to merge.

[zhe-thoughts]

"buildkite/premerge" is a x

[stephanie-wang]

"buildkite/premerge" is a x

The failure is unrelated (it's flaky on master)

[c21]

I guess buildkite/premerge is mandatory to pass before merging. Let me retry it.

[zhe-thoughts]

Yes other checks we can use manual judgement, buildkite/premerge must pass. cc @aslonnie

[stephanie-wang]

Yes other checks we can use manual judgement, buildkite/premerge must pass. cc @aslonnie

It's passing now.