ARROW-11300: [Rust][DataFusion] Further performance improvements on hash aggregation with small groups

[Dandandan]

Based on #9234, this PR improves the situation described in https://issues.apache.org/jira/browse/ARROW-11300.

The current situation is that we call take on arrays, which is fine, but causes a lot of small Arrays to be created / allocated. when we have only a small number of rows in each group.

This improves the results on the group by queries on db-benchmark:

PR:

q1 took 32 ms
q2 took 422 ms
q3 took 3468 ms
q4 took 44 ms
q5 took 3166 ms
q7 took 3081 ms

#9234 (different results from that PR description as this has now partitioning enabled and a custom allocator)

q1 took 34 ms
q2 took 389 ms
q3 took 4590 ms
q4 took 47 ms
q5 took 5152 ms
q7 took 3941 ms

The PR changes the algorithm to:

• Create indices / offsets of all keys / indices new in the batch.
• take the arrays based on indices in one go (so it only requires one bigger allocation for each array)
• Use slice based on the offsets to take values from the arrays and pass it to the accumulators.

[github-actions]

https://issues.apache.org/jira/browse/ARROW-11300

[Dandandan]

@jorgecarleitao while the change makes hash aggregates a bit faster already, it seems Array.slice doesn't really do a zero-copy as promised in the doc (and accounts for about 11% of the total instructions of the program). Maybe you have an idea what could be going on here?

[jorgecarleitao]

Yes, slicing is suboptimal atm. Also, IMO it should not be the Array to implement that method, but each implementation individually. I haven't touch that part yet, though.

That analysis is really cool, though. How did you came up with that????? 😮 Any blog post or readme??? :P

[Dandandan]

Thanks @jorgecarleitao makes sense.
I'm planning a blog post soon again, might put some details there.
In short it is cargo profiler installed from the repo https://github.com/svenstaro/cargo-profiler with callgrind and kcachegrind.

[Dandandan]

This PR in itself is ready for review now. The change in approach will help us in the future, even more with Array.slice being improved (will create a ticket later), execution times should be further reduced by something like 15-20%.

[codecov-io]

Codecov Report

Merging #9271 (f578e4c) into master (3e47777) will increase coverage by 0.00%.
The diff coverage is 85.03%.

@@           Coverage Diff           @@
##           master    #9271   +/-   ##
=======================================
  Coverage   81.89%   81.89%           
=======================================
  Files         215      215           
  Lines       52988    53003   +15     
=======================================
+ Hits        43392    43407   +15     
  Misses       9596     9596           

Impacted Files	Coverage Δ
rust/arrow/src/array/array_struct.rs	88.43% <ø> (ø)
rust/arrow/src/array/equal/utils.rs	75.49% <0.00%> (ø)
rust/arrow/src/bytes.rs	53.12% <ø> (ø)
rust/arrow/src/compute/kernels/aggregate.rs	74.93% <ø> (ø)
rust/arrow/src/ipc/gen/SparseTensor.rs	0.00% <0.00%> (ø)
rust/arrow/src/ipc/gen/Tensor.rs	0.00% <0.00%> (ø)
rust/benchmarks/src/bin/nyctaxi.rs	0.00% <ø> (ø)
rust/benchmarks/src/bin/tpch.rs	6.97% <0.00%> (ø)
rust/datafusion/src/datasource/memory.rs	79.75% <0.00%> (ø)
rust/datafusion/src/logical_plan/operators.rs	75.00% <ø> (ø)
... and 134 more

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 9199801...f578e4c. Read the comment docs.

[Dandandan]

@jorgecarleitao I found the "offending" code is this function in array/data.rs which does a self.clone(). Any idea how we could write a non copying version?
I think it was copied/adapted based on a function in the array module, but I think it was also cloning in previous versions.

    /// Creates a zero-copy slice of itself. This creates a new [ArrayData]
    /// with a different offset, len and a shifted null bitmap.
    ///
    /// # Panics
    ///
    /// Panics if `offset + length > self.len()`.
    pub fn slice(&self, offset: usize, length: usize) -> ArrayData {
        assert!((offset + length) <= self.len());

        let mut new_data = self.clone();

        new_data.len = length;
        new_data.offset = offset + self.offset;

        new_data.null_count =
            count_nulls(new_data.null_buffer(), new_data.offset, new_data.len);

        new_data
    }
    

[jorgecarleitao]

Isn't the data contained on a buffer Arced? I.e. Vec<Buffer>::clone() should be cheap, no?

[Dandandan]

@jorgecarleitao
Maybe, reasonably, yes (as long the underlying vecs have a few items).

I think the clone on the ArrayData structure itself is expensive when slicing for a few elements as it does some allocations for the struct, the vecs, execute some code, etc.

[jhorstmann]

Cloning the vector of Buffers and child ArrayData has some overhead. Incrementing the reference counts should be relatively cheap unless there are concurrent threads accessing the same Arc.

I tried replacing the vectors inside ArrayData with SmallVec some time ago. That made the slice benchmarks faster, but several other benchmarks slowed down because of it. Might be worth to revisit that.

[nevi-me]

I found the "offending" code is this function in array/data.rs which does a self.clone().

This relates to the other discussion that we had on how slicing an array does a clone without propagating offset information to child_data and buffers.

I find it interesting though that self.clone() would cost us so much. If you were to clone an array, and inspect the original and the cloned, the buffers point to the same location in memory. I thought that's what zero-copying would give us.

Is that not the case?

[Dandandan]

@nevi-me

I don't think indeed it is very expensive on large Arrays compared to the size / operations on the array, but it turns out to be expensive on very small arrays. For this PR I am using slice to make the hash aggregate code in DataFusion more efficient for small output groups with a small amount of rows (only 1 row / Array.slice(i, 1)) in extreme cases), in which case the slicing function becomes a bottleneck, because of the cloning here + make_array function and because it will be called many times, for example (I believe) in total 20M times for a table of 10M rows (it is one example of the db-benchmark benchmark).
It still is faster than taking for each group individually though as the benchmark results show.

I am wondering if instead of trying to make a new array when doing .slice(), we could create a data-structure for slicing instead that implements the Array interface and is supported in kernels, so creating the slice would be cheap?

Something like this :

struct ArraySlice {
    offset: usize,
    length: usize,
    array: ArrayRef
}

[Dandandan]

This PR itself is ready for review.

I think the performance for slicing small slices would be something to look at later.

[jorgecarleitao]

Thanks a lot for your points. I am learning a lot! :)

Note that for small arrays, we are basically in the metadata problem on which the "payload size" of transmitting 1 element is driven by its metadata, not the data itself. This will always be a problem, as the arrow format was designed to be performant for large arrays.

For example, all our buffers are shared via an Arc. There is a tradeoff between this indirection and mem-copying the memory region. The tradeoff works in Arc's favor for large memory regions and vice-versa.

With that said, we could consider replacing Arc<ArrayData> by ArrayData on all our arrays, to avoid the extra Arc: cloning an ArrayData is actually cheap. I am not sure if that would work for FFI, but we could certainly try.

Another idea is to use buffer1: Buffer, buffer2: Buffer instead of buffers: Vec<Buffer> to avoid the Vec. This is possible because arrow arrays support at most 2 buffers (3 with the null). For types of a single buffer, we are already incurring the cost of the Vec and thus adding a Buffer instead should not be a big issue (memory-wise). The advantage of this is that we avoid cloning the Vec on every operation as well as the extra bound check. The disadvantage is that we have to be more verbose when we want to apply an operation to every buffer.

[jorgecarleitao]

LGTM. cool idea :)

[Dandandan]

@jorgecarleitao

Thanks, that it a great summary of the situation! I think removing the Vec sounds like it might improve the "metadata overhead", maybe there are some more opportunities like that as well. Maybe the second buffer could be Option<Buffer> as well to avoid creating a buffer (if that would add extra overhead)?
Looking at Arraydata, that might be reasonable cheap as well to clone, not sure if it would improve in all situations as it also involves some extra copying.

[Dandandan]

I rebased this PR @jorgecarleitao @alamb

[jorgecarleitao]

Thanks @Dandandan . fmt missing :)

[alamb]

This looks good to me @Dandandan -- thank you. I read it fairly closely and it makes sense.

[alamb]

I wonder if you could get any additional performance by using the knowledge of the size of batch_keys

Suggested change

	let mut batch_indices: UInt32Builder = UInt32Builder::new(0);
	let mut batch_indices: UInt32Builder = UInt32Builder::new(batch_keys.len());

[alamb]

Or maybe you have to scale it by the number of accumulators too

[alamb]

Suggested change

	let mut offsets = vec![0];
	let mut offsets = Vec::with_capacity(batch_keys.len());
	offsets.push(0);