Rich-T-kid commented on issue #7001: URL: https://github.com/apache/datafusion/issues/7001#issuecomment-5801503499
Okay I think I have a few ideas for how we can speed up RepartionExec. first lets do a review so everyone is on the same page. <img width="2928" height="1362" alt="Image" src="https://github.com/user-attachments/assets/a0807852-4857-469d-8240-c39bc24c0b8e"; /> At a high level this is how repartionExec, specifically the HashRepartion path. 1. hash the key columns. `M`,`N`,`K` in `Hash(M,N,K)` 2. Call `take_arrays()` once per input batch to produce one reordered batch, then slice() it into per-partition sub-batches (https://github.com/apache/datafusion/pull/22159) 3. Each sub-batch is pushed into its output partition's `SharedCoalescer`, which buffers rows behind a Mutex until `target_batch_size` is reached 4. When `target_batch_size` is reached, `push_and_drain()` returns the completed batch, which is then sent over the partition's async channel to the consumer side The most CPU-intensive parts of this are `take_arrays()` and `push_batch()`, since every row passes through both. Key hashing only scales with the number of columns being hashed, and as @gabotechs found [here](https://github.com/apache/datafusion/pull/23720#issuecomment-5104141874), improvements there would only be marginal. I've been working on a [series of take optimizations](https://github.com/apache/arrow-rs/issues/8879) in arrow-rs that build on top of one another nicely — [optimize take bool & take utf8View](https://github.com/apache/arrow-rs/pull/10813), [optimize take list](https://github.com/apache/arrow-rs/pull/10812), [optimize fsl](https://github.com/apache/arrow-rs/pull/10441), etc. The most important piece is the final one, [#10945](https://github.com/apache/arrow-rs/pull/10945), which has shown up to a 25% speedup. Simply swapping the existing methods for the new ones introduced in that PR should yield significant improvements here, given how heavily RepartitionExec relies on take. That leaves SharedCoalescer, which is a thin wrapper around [BatchCoalescer](https://arrow.apache.org/rust/arrow/compute/struct.BatchCoalescer). I've made a few attempts to [optimize](https://github.com/apache/arrow-rs/issues/7761) this, with little to no success, but I think the best return on effort is [Optimize BatchCoalescer::push_batch_with_indices](https://github.com/apache/arrow-rs/issues/8957) (by @Dandandan). The core issue is that we're constrained by having to copy repeatedly. `take()` has to copy every byte from input to output, and BatchCoalescer (except for primitives, utf8View, and fixed-size binary) has to allocate 2x memory and then copy the bytes from source to destination. Right now this copy happens twice; simply reducing that to a single copy could be a significant win. <img width="2828" height="1396" alt="Image" src="https://github.com/user-attachments/assets/e5fa6c96-88bb-4afb-8e92-c3fca814cfff"; /> What Im proposing is to avoid the `take_arrays()` call in place for passing a zero-copy slice to `SharedCoalescer` as well as the indices that would have went to `take_arrays()`. then SharedCoacler would use `push_batch_with_indices()`. the issue with using this today is that this method is just a wrapper for `take_arrays()` https://github.com/apache/arrow-rs/blob/e491ce6c6f343ff9d7a235376d0353d12b365330/arrow-select/src/coalesce.rs#L291. The first step towards this is updating the backing trait `InprogressArrays` to take advantage of `push_batch_with_indices()` by writing the intermediate state directly to the output buffer. That is introduced in this PR https://github.com/apache/arrow-rs/pull/11174. Besides that I think the other largest wins are trying to speed up the concat kernel as this is the generic path (non-primitves/View types). -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
