decorrelate shuffle size from target_batch_size

[Rich-T-kid]

Currently DistributedConfig.shuffle_batch_size sets the target_record_batch_size for all operators in a datafusion plan. This helps for buffering record batches before being sent over the network but it may cause users to unintentionally change the target_record_batch. target_record_batch and shuffle_batch_size should be independent config values. (see other issue with this in #541)

This PR re-introduces coalesce batching to provide the same buffering of arrow-flight data before being sent over the wire to minimize overhead. the target size for batches is set to shuffle_batch_size.

As @gabotechs mentions in #386, coalescing has a cost. It copies rows from source arrays into a new destination buffer. Looking at push_batch, in_progress_arrays briefly holds a pointer to each source array via set_source(), and copy_rows() is called against it on every push. Despite the name, copy_rows() does mostly zero-copy work: it takes a zero-copy slice() of the source array and pushes it into a Vec. It's cheap metadata bookkeeping, not a buffer copy.

The actual copy is deferred to finish_buffered_batch(), which calls concat over all the buffered slices to write them into a single pre-allocated, contiguous buffer. Because this is one bulk copy into pre-sized memory, we only pay for the writes not re-sizing buffers. Modern CPUs do memcpy-style bulk copies extremely fast, especially when the data stays in cache.
The tuning insight is that as shuffle_batch_size outpaces target_batch_size, the number of separate coalesce-and-send cycles shrinks.
Consider 12 incoming record batches of 8k rows each. With no coalescing, that's 12 Arrow Flight writes over the wire. If we set shuffle_batch_size to 32kb, we instead do 3 Flight writes, at the cost of coalescing 4 arrays into one buffer each time. That's a good trade: a handful of (possibly in-cache) CPU copies is cheaper than the per-message network overhead of sending many small batches over the wire.

a tradeoff of local CPU/copy cost vs. the overhead of many small network packets.

At worst, this performs on par with what main currently does, with the added benefit of being clearer to users. A follow up to this PR is discussed in #542

context:
#541

[Rich-T-kid]

this adds opt out behavior the users that want target_batch_size == shuffle_batch_size

[Rich-T-kid]

I think this is a sane default, happy to tune it

[gabotechs]

Could you share how does this perform over the remote benchmarks?

[Rich-T-kid]

Could you share how does this perform over the remote benchmarks?

@gabotechs

=== Comparing tpch_sf10 results from engine 'datafusion-distributed-main' [prev] with 'datafusion-distributed-decorrelate-shuffle-size' [new] ===
      q1: prev= 421 ms, new= 438 ms, diff=1.04 slower ✖
      q2: prev=1445 ms, new=1393 ms, diff=1.04 faster ✔
      q3: prev=1014 ms, new= 874 ms, diff=1.16 faster ✔
      q4: prev= 438 ms, new= 334 ms, diff=1.31 faster ✅
      q5: prev=1228 ms, new=1183 ms, diff=1.04 faster ✔
      q6: prev= 286 ms, new= 223 ms, diff=1.28 faster ✅
      q7: prev=1144 ms, new=1089 ms, diff=1.05 faster ✔
      q8: prev=1545 ms, new=1632 ms, diff=1.06 slower ✖
      q9: prev=1911 ms, new=1531 ms, diff=1.25 faster ✅
     q10: prev=1219 ms, new=1218 ms, diff=1.00 faster ✔
     q11: prev= 578 ms, new= 678 ms, diff=1.17 slower ✖
     q12: prev= 539 ms, new= 558 ms, diff=1.04 slower ✖
     q13: prev= 782 ms, new= 867 ms, diff=1.11 slower ✖
     q14: prev= 560 ms, new= 558 ms, diff=1.00 faster ✔
     q15: prev= 765 ms, new= 777 ms, diff=1.02 slower ✖
     q16: prev= 443 ms, new= 403 ms, diff=1.10 faster ✔
     q17: prev= 853 ms, new= 864 ms, diff=1.01 slower ✖
     q18: prev=1248 ms, new=1263 ms, diff=1.01 slower ✖
     q19: prev= 442 ms, new= 574 ms, diff=1.30 slower ❌
     q20: prev= 948 ms, new= 923 ms, diff=1.03 faster ✔
     q21: prev=1242 ms, new=1235 ms, diff=1.01 faster ✔
     q22: prev= 365 ms, new= 386 ms, diff=1.06 slower ✖
   TOTAL: prev=19416 ms, new=19001 ms, diff=1.02 faster ✔

this doesn't introduce any regressions & the differences seem to be mostly noise.
I also experimented with (8,16,32,64) * 8192 for shuffle_size and they are all within 1.00-1.08x so pretty minor difference.

[gabotechs]

If the difference is just within noise, I'd recommend to not merge this PR.

The problem with this approach is that it implies double-coalescing batches: one time for adhering to standard's datafusion.execution.batch_size and another time for adhering to distributed.shuffle_batch_size. Note that coalescing is an O(N) operation over the data at worst, as it copies it for coalescing.

Citing the description:

target_record_batch and shuffle_batch_size should be independent config values

We deliberately made the decision of not making these two values independent, mainly for avoiding double-coalescing and simplifying the setup. We prefer to keep to full pipeline under a single batch_size.

If you see that the benchmarks demonstrate that this decision was wrong, this is good to go, but otherwise I'd probably try to look for other perf opportunities.

[Rich-T-kid]

ah that make sense.

The problem with this approach is that it implies double-coalescing batches: one time for adhering to standard's datafusion.execution.batch_size and another time for adhering to distributed.shuffle_batch_size

This was actually the main issue I wanted to address. I spoke with @jayshrivastava offline about this. The goal of this PR was mostly as a building block for #542 to make that PR smaller.

If you see that the benchmarks demonstrate that this decision was wrong

With this in mind ill implement some of the optimizations I have in mind for #542 and then open a PR if I see good results.

thank you for the comments @gabotechs ill close this PR now

[Rich-T-kid]

We deliberately made the decision of not making these two values independent, mainly for avoiding double-coalescing and simplifying the setup. We prefer to keep to full pipeline under a single batch_size.

if we can avoid the double-coalescing do you think it would be worth it to make shuffle_size and target_batch_size independent? Im of the opinion that these should be distinct since they are tuned for different priorities, i'm curious as to what factors went into the decision.