Excessive Arc-clone in HashJoinStream with StringView on build-side

[ctsk]

Describe the bug

An unfortunate pattern in the hash join implementation leads to excessive Arc-cloning: Assume the build-side carries a string-view column as a payload. Let N be the number of batches seen on the build side

1. In the build phase, datafusion concatenates the batches on the build side. The string-view column now holds references to at least N data buffers in a vec;

2. When constructing the output batch, the take implementation for string-views clones the data buffer vector of the concatenated build-side column - thus incrementing the references on all N data buffers.

To Reproduce

I noticed this issue when executing and profiling tpch query 18 - roughly 3% of the runtime is spent cloning these Arcs.

Expected behavior

No response

Additional context

• The concat during build:
  

  datafusion/datafusion/physical-plan/src/joins/hash_join.rs

  Lines 1013 to 1015 in 7002a00

  	// Merge all batches into a single batch, so we can directly index into the arrays
  	let single_batch = concat_batches(&schema, batches_iter)?;

• The take call during batch construction:
  

  datafusion/datafusion/physical-plan/src/joins/utils.rs

  Line 918 in 7002a00

  compute::take(array.as_ref(), build_indices, None)?

• The relevant bit of arrow-rs
  https://github.com/apache/arrow-rs/blob/7e85b48dc8f929afa82f2878b17db7b2df240b8b/arrow-select/src/take.rs#L565-L567

[jonathanc-n]

take

[alamb]

I have noticed something similar -- that the pattern of "take / concat" results in potentially non trival overhead.

I have been recently working on some related code (the same basic pattern happens for filter) Related:

• Optimize take/filter/concat from multiple input arrays to a single large output array arrow-rs#6692
• POC: Sketch out parquet cached filter result API arrow-rs#7513

[Dandandan]

Yeah - one of the options to optimize join is by avoiding the take + concat (which happens when generating fewer rows per input batch in join, so I think that could use a similar strategy).

Ideally, we also would avoid concat_batches in the start (and rather maintain aVec<RecordBatch>).
This was changed in the early days (by me) as it used to be slow.

But now that the hashmap is basically only comparing hash values during probe, I think we can try and avoid the concat_batches (and store a u32 / u32 of batch_id + index instead of a u64 index).

[Dandandan]

This might help here as well I think (to avoid the excessive cloning)?

apache/arrow-rs#6427

[jonathanc-n]

I'll try to give that a try, thanks @Dandandan

[Dandandan]

So if we're to change RecordBatch to Vec<RecordBatch> in the join code, I think this is what's roughly needed:

• Change batch to Vec<RecordBatch> in JoinLeftData
• Initialize JoinLeftData with a Vec<(usize, usize>) which has the batch_id and row_id for each index
• Pass a Vec<(usize, usize>) to methods based on the indices from the build side and the Vec<(usize, usize>) structure
• Change code to use interleave rather than take
• Adapt other places where needed

[ctsk]

I've re-encountered this issue and it (obviously, duh) gets amplified with larger scale facctors (>100). For example, the top 3 items by cycles when running tpch query 18 @ sf 300 are:

  18.07%  tokio-runtime-w  tpch                                              [.] arrow_select::take::take_byte_view                                                                                                                                                                    ◆
  14.52%  tokio-runtime-w  tpch                                              [.] <datafusion_physical_plan::coalesce_batches::CoalesceBatchesStream as futures_core::stream::Stream>::poll_next                                                                                        ▒
  12.96%  tokio-runtime-w  tpch                                              [.] core::ptr::drop_in_place<alloc::vec::Vec<arrow_buffer::buffer::immutable::Buffer>>    

That is close to 50% of cycles spent due to this issue. Because the batches that get produced by the hash join have a lot of data buffers, the CoalesceBatchesExec gets more expensive too: Inside the operator, we iterate over the data buffers to determine whether to perform a garbage collection.

I believe a reasonable fix for this issue is to perform a garbage collection on the build side after the concatenation. The build side would then only hold a single data buffer per string-view / byte-view column.

[Dandandan]

Thanks for sharing the insight @ctsk

I think there is some overlap here to the work @alamb to improve (and maybe remove over time) CoalesceBatchesExec and gc of views.

[alamb]

Thanks for sharing the insight @ctsk

I think there is some overlap here to the work @alamb to improve (and maybe remove over time) CoalesceBatchesExec and gc of views.

The relevant arrow ticket is

• Optimize take/filter/concat from multiple input arrays to a single large output array arrow-rs#6692

@Dandandan also made a PR that improves some view code here (but I don't think it affects take_view) here:

• Speedup interleave_views (4-7x faster) arrow-rs#7695

[ctsk]

I think the issue on the take+concat pattern is only tangentially related to this issue. Ultimately, you would need a version of the take operation, that does not simply clone the vector of data buffers, but rather eagerly allocates a new data buffer to hold the string view data. This modification could be applied to the current compute::take kernel too.

The interleave kernel already does something like this: It only takes references to the data buffers that are actually referenced in the new array - that is why it would help somewhat alleviate this issue. If the build side is really large, this could still mean allocating a vector of 1000+ buffer references. The subsequent gc (in CoalesceBatchesExec) drops the vector and allocate a new data buffer.

Instead, I think there are 2 direct solutions for this:

1. Extend the take (or interleave) kernels to eagerly allocate new data buffers for string views (i.e. add a compute::TakeOption - unfortunately also a breaking change).
2. Switch from Vec<Buffer> to Arc<[Buffer]> - also a breaking change in the arrow crate. This would not require writing new kernel-style code - the take/interleave operations would remain unchanged.

My PR on executing a GC on the hash join build-side is unfortunately not a good solution: This GC operation is expensive, and whether it pays off depends on the number of probe side batches - which is not known in advance. The benchmark shows that it significantly slows down smaller joins. It would also not prevent the CoalesceBatchesExec after the HJ from issuing another GC.

For now, the best way to fix this issue is to work on getting apache/arrow-rs#6427 merged. As far as I can tell, it is already fully implemented and got stuck due to lack of time of the author to add the ViewBuffers abstraction.

Lastly, I believe that using interleave, or a good solution to the take+concat pattern would be very helpful in making the HJ more efficient in general. Most importantly, it could allow avoiding the concat on the build-side batches, and remove the CoalesceBatchesExec after the join.

[alamb]

Extend the take (or interleave) kernels to eagerly allocate new data buffers for string views (i.e. add a compute::TakeOption - unfortunately also a breaking change).

BTW this eagerly allocating buffers / copying string views when needed is what the (in progress) coalesce kernel does:

• https://github.com/apache/arrow-rs/blob/main/arrow-select/src/coalesce.rs

Here is the logic that does it for StringView

• https://github.com/apache/arrow-rs/blob/main/arrow-select/src/coalesce/byte_view.rs

I hope to migrate DataFusion to use the BatchCoalescer soon.

There is a BatchCoalescer::push_batch_with_filter (that does the equivalent of filter + concat) that I am actively working to improve

I imagine implementing something like BatchCoalescer::push_batch_with_indices that does the equivalent of take + concat