ARROW-11045: [Rust] Fix performance issues of allocator

[jorgecarleitao]

This PR addresses a performance issue in how we allocate and reallocate the MutableBuffer.

Problem

See #9032

This PR

This PR changes MutableBuffer::reserve to call std::alloc::alloc instead of std::alloc::alloc_zeroed, which improves performance when building buffers with unknown sizes (such as strings and nested types).

This required changing some calls of MutableBuffer that assumed a zero initialized buffer even when reserve was used.
It also changed reserve's signature to reserve(additional) instead of reserve(new_len), which is the notation used throughout Rust's std library.

critcmp master-simd-e1b38cdaa4f2a7d35e2e576463e12b38875f29f3 alloc2-simd-88fc0ae819c24239ac9363fa462f9c6e1ddfd9fc -t 10

group                                 alloc2-simd-88fc0ae819c24239ac9363fa462f9c6e1ddfd9fc    master-simd-e1b38cdaa4f2a7d35e2e576463e12b38875f29f3
-----                                 ----------------------------------------------------    ----------------------------------------------------
add 512                               1.00   549.1±17.26ns        ? B/sec                     1.14  624.4±118.72ns        ? B/sec
buffer_bit_ops or                     1.00    369.3±7.51ns        ? B/sec                     1.16   427.1±20.25ns        ? B/sec
cast float32 to int32 512             1.00      3.3±0.09µs        ? B/sec                     1.21      4.0±0.09µs        ? B/sec
cast float64 to float32 512           1.00      3.0±0.06µs        ? B/sec                     1.24      3.7±0.11µs        ? B/sec
cast float64 to uint64 512            1.00      3.6±0.33µs        ? B/sec                     1.22      4.4±0.29µs        ? B/sec
cast int32 to float32 512             1.00      2.9±0.10µs        ? B/sec                     1.15      3.4±0.09µs        ? B/sec
cast int32 to float64 512             1.00      2.9±0.06µs        ? B/sec                     1.14      3.3±0.06µs        ? B/sec
cast int32 to uint32 512              1.00      4.0±0.09µs        ? B/sec                     1.23      4.9±0.12µs        ? B/sec
concat str 1024                       1.00      8.4±0.26µs        ? B/sec                     1.14      9.6±0.22µs        ? B/sec
equal_nulls_512                       1.00      3.3±0.07µs        ? B/sec                     1.22      4.0±0.10µs        ? B/sec
filter context u8 high selectivity    1.00      3.7±0.09µs        ? B/sec                     1.27      4.6±0.13µs        ? B/sec
filter u8 high selectivity            1.00     10.8±0.47µs        ? B/sec                     1.10     11.9±0.49µs        ? B/sec
like_utf8 scalar equals               1.00    149.9±6.36µs        ? B/sec                     1.12    167.3±3.45µs        ? B/sec
like_utf8 scalar starts with          1.00   338.2±12.26µs        ? B/sec                     1.15    388.4±7.62µs        ? B/sec
min string 512                        1.13      6.0±0.17µs        ? B/sec                     1.00      5.3±0.09µs        ? B/sec
nlike_utf8 scalar starts with         1.00   367.8±39.16µs        ? B/sec                     1.16    425.4±9.38µs        ? B/sec
subtract 512                          1.00   567.9±10.77ns        ? B/sec                     1.21  686.0±209.80ns        ? B/sec
sum 512                               1.32     67.8±0.59ns        ? B/sec                     1.00     51.3±0.70ns        ? B/sec
take str 1024                         1.19      6.1±0.10µs        ? B/sec                     1.00      5.1±0.03µs        ? B/sec
take str 512                          1.12      3.9±0.09µs        ? B/sec                     1.00      3.4±0.07µs        ? B/sec
take str null indices 1024            1.18      6.1±0.16µs        ? B/sec                     1.00      5.2±0.11µs        ? B/sec
take str null indices 512             1.12      3.9±0.09µs        ? B/sec                     1.00      3.5±0.03µs        ? B/sec
take str null values 1024             1.19      6.1±0.51µs        ? B/sec                     1.00      5.2±0.11µs        ? B/sec

The take is slower because it is using Buffer::from, which is always worse than using a MutableBuffer.

[jorgecarleitao]

@nevi-me , there is a failing test in parquet that I am unable to fix. I think that we may be creating bitmap buffers that are too large.

My current hypothesis is that the buffer old_bitmap at consume_bitmap_buffer is not being resized in the same way the other buffer in consume_rep_levels is, which is causing it to be too large. Do you agree with this (before I proceed to patch it)?

[codecov-io]

Codecov Report

Merging #9044 (ff3efa0) into master (eb17687) will decrease coverage by 0.00%.
The diff coverage is 93.51%.

@@            Coverage Diff             @@
##           master    #9044      +/-   ##
==========================================
- Coverage   82.57%   82.56%   -0.01%     
==========================================
  Files         203      203              
  Lines       50036    49875     -161     
==========================================
- Hits        41316    41180     -136     
+ Misses       8720     8695      -25     

Impacted Files	Coverage Δ
rust/arrow/src/array/array_list.rs	92.72% <ø> (-0.38%)	⬇️
rust/arrow/src/bytes.rs	53.12% <ø> (-5.21%)	⬇️
rust/arrow/src/compute/kernels/aggregate.rs	74.93% <ø> (-0.07%)	⬇️
rust/arrow/src/compute/kernels/comparison.rs	95.91% <ø> (ø)
rust/arrow/src/datatypes.rs	75.02% <ø> (ø)
rust/arrow/src/array/builder.rs	85.42% <86.84%> (+1.42%)	⬆️
rust/arrow/src/buffer.rs	97.38% <94.87%> (-0.84%)	⬇️
rust/arrow/src/array/array_primitive.rs	92.28% <100.00%> (-0.04%)	⬇️
rust/arrow/src/array/raw_pointer.rs	100.00% <100.00%> (ø)
rust/arrow/src/array/transform/list.rs	83.33% <100.00%> (-0.54%)	⬇️
... and 15 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 eb17687...ff3efa0. Read the comment docs.

[sunchao]

Is there any reason to change the API in this PR? don't see how it could be related to the performance issue.
Also better to split it to make PR easier to review. We also need to update the doc.

[jorgecarleitao]

I agree that it makes it more difficult. The signature is related to performance, though:

It was difficult to reason with reserve(new_len) when Vec::reserve and RawVec::reserve uses additional (and I was basing this PR on that). This PR fixes a bug in one Builder that was calling reserve(1) in append, when it should be calling reserve(self.len + 1). I also spent 4hs tracking a bug because of this difference (I was using reserve(additional) when its signature was reserve(new_len).

I can try to split that to another PR, though. IMO this PR will still need to be based on that PR.

[sunchao]

I think originally we referred C++ impl for the signature design to make it consistent. Since buffer module is public, we'll need to be careful when making any breaking change on API.

[jorgecarleitao]

Yap, and I think it was a good decision back then, when the goal was to port C++ code to Rust. I understand your concern.

However, IMO we should not keep using C++-motivated APIs because it was historically like that. additional is the standard in virtually every container in Rust (HashMap, Vec, tokio::bytes::Bytes). IMO this change will need to be made at some point, and delaying it will only cause more pain in the future.

In this particular instance, it was motivated by an already existing confusion on our own code base, which IMO evidences that rust developers expect additional. Most people that use this API are probably passing additional to it already.

Regardless, we are bumping the major version, so there is already no expectation of backward compatibility as far as cargo and semver is concerned. My goal here is to make MutableBuffer as close as possible to Vec so that people can use them without even thinking that they are using a special allocation.

[sunchao]

Thanks, if we're going to bump the major version then I'm fine with the compatibility issue. I don't have strong opinion on this so either is fine to me. IMO a good documentation should be suffice for developers to understand what the method does.

[sunchao]

we could have another method with value = 0 since this seems to be the dominate use case (it might enable future optimizations for the special case as well).

[jorgecarleitao]

I based this signature on Vec::resize. I.e. I am assuming here that we use MutableBuffer as a special Vec<u8> that uses our custom allocator, and offer users an API that they are used to (Vec).

The main use-case for value: 255u8 is when we want to allocate and grow a set bitmap to later unset bits. We did not offer that before, which is the reason MutableArrayData does not do that (even though is most likely more efficient for many use-cases).

Note that even though rust std::RawVec offers a with_capacity_zeroed, std::Vec does not expose that as a public API: resize and resize_with are the public APIs to resize a vec, and none of them uses std::alloc::alloc_zeroed, i.e. they also just call std::ptr::write_bytes.

I agree with you that resize(len, 0) can potentially be optimized when the mutable buffer has zero capacity by calling std::alloc::alloc_zeroed instead of std::ptr::write_bytes.

[Dandandan]

@jorgecarleitao I think this PR broke master:

   --> datafusion/src/physical_plan/parquet.rs:712:25
    |
712 |             data_buffer.resize(data_buffer.len() + data_size);
    |                         ^^^^^^ ----------------------------- supplied 1 argument
    |                         |
    |                         expected 2 arguments
    |
note: associated function defined here
   --> rust/arrow/src/buffer.rs:833:12
    |
833 |     pub fn resize(&mut self, new_len: usize, value: u8) {

[jorgecarleitao]

likely. It was kind of expected, as it did some backward incompatible changes. I was trying to merge it first to avoid breaking, but I guess I was not fast for the speed on which PRs are merged into master after the green light on the mailing list :P