perf: avoid re-allocation if buffer is not shared during BooleanArray::take_n_true

[Rich-T-kid]

Which issue does this PR close?

• Closes use Buffer.into_mutable to reuse the allocation if possible in take_n_true #10251.

Rationale for this change

If a boolean buffer is not shared we should re-use its allocation instead of building a new builder from scratch.

What changes are included in this PR?

match on into_mutable, in the error case the old path is taken. If the buffer can be re-used we use it directly.

Are these changes tested?

yes, existing test cover this behavior

Are there any user-facing changes?

no

[Rich-T-kid]

Im sure there are faster ways to do this. I looked up a few ways to do this and im not too familar with bit operations so I left is as a basic loop for now

[alamb]

See also a related PR

• Add BooleanBuffer::try_into_builder and BooleanBuffer::into_builder_or_clone #10397

[Rich-T-kid]

@alamb this is ready for review

[Jefffrey]

maybe we should do a microbenchmark to doublecheck this is a performance gain? we dont necessarily need to commit it, but just something to use as reference for reviewing this

[Rich-T-kid]

I'll make a seperate PR

maybe we should do a microbenchmark to doublecheck this is a performance gain? we dont necessarily need to commit it, but just something to use as reference for reviewing this

makes sense to me, created this PR so we can run the benchmarks on it https://github.com/apache/arrow-rs/pull/10703/changes

[Rich-T-kid]

ooh wait the two benchmarks share the same name. Thats my mistake, fixed it here #10705

[Jefffrey]

run benchmark boolean_array
env:
BENCH_FILTER: take_n_true

[adriangbot]

🤖 Arrow criterion benchmark completed (GKE) | trigger

Instance: c4a-highmem-16 (12 vCPU / 65 GiB)

Comparing rich-T-kid/optimize-take-n-boolBuff (2e2815d) to 5ce0ebe (merge-base) diff

Run configuration

run benchmark boolean_array
env:
  BENCH_FILTER: "take_n_true"

CPU Details (lscpu)

Architecture:                            aarch64
CPU op-mode(s):                          64-bit
Byte Order:                              Little Endian
CPU(s):                                  16
On-line CPU(s) list:                     0-15
Vendor ID:                               ARM
Model name:                              Neoverse-V2
Model:                                   1
Thread(s) per core:                      1
Core(s) per cluster:                     16
Socket(s):                               -
Cluster(s):                              1
Stepping:                                r0p1
BogoMIPS:                                2000.00
Flags:                                   fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache:                               1 MiB (16 instances)
L1i cache:                               1 MiB (16 instances)
L2 cache:                                32 MiB (16 instances)
L3 cache:                                80 MiB (1 instance)
NUMA node(s):                            1
NUMA node0 CPU(s):                       0-15
Vulnerability Gather data sampling:      Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit:             Not affected
Vulnerability L1tf:                      Not affected
Vulnerability Mds:                       Not affected
Vulnerability Meltdown:                  Not affected
Vulnerability Mmio stale data:           Not affected
Vulnerability Reg file data sampling:    Not affected
Vulnerability Retbleed:                  Not affected
Vulnerability Spec rstack overflow:      Not affected
Vulnerability Spec store bypass:         Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:                Mitigation; __user pointer sanitization
Vulnerability Spectre v2:                Mitigation; CSV2, BHB
Vulnerability Srbds:                     Not affected
Vulnerability Tsa:                       Not affected
Vulnerability Tsx async abort:           Not affected
Vulnerability Vmscape:                   Not affected

Details

group                        main                                   rich-T-kid_optimize-take-n-boolBuff
-----                        ----                                   -----------------------------------
take_n_true_shared(32)       1.02    105.5±5.81ns        ? ?/sec    1.00    103.0±1.75ns        ? ?/sec
take_n_true_shared(32768)    1.00     16.3±0.07µs        ? ?/sec    1.00     16.3±0.05µs        ? ?/sec
take_n_true_shared(512)      1.00    345.6±3.69ns        ? ?/sec    1.00    344.1±1.62ns        ? ?/sec
take_n_true_unique(32)       1.45    143.8±4.99ns        ? ?/sec    1.00     98.9±2.73ns        ? ?/sec
take_n_true_unique(32768)    1.00     16.4±0.29µs        ? ?/sec    1.62     26.7±0.04µs        ? ?/sec
take_n_true_unique(512)      1.00   384.0±10.26ns        ? ?/sec    1.26    484.2±2.18ns        ? ?/sec

Resource Usage

base (merge-base)

Metric	Value
Wall time	55.0s
Peak memory	321.9 MiB
Avg memory	83.0 MiB
CPU user	51.3s
CPU sys	1.6s
Peak spill	0 B

branch

Metric	Value
Wall time	60.0s
Peak memory	298.5 MiB
Avg memory	98.9 MiB
CPU user	53.4s
CPU sys	1.6s
Peak spill	0 B

---

File an issue against this benchmark runner

[Jefffrey]

run benchmark boolean_array
env:
BENCH_FILTER: take_n_true

[adriangbot]

🤖 Arrow criterion benchmark completed (GKE) | trigger

Instance: c4a-highmem-16 (12 vCPU / 65 GiB)

Comparing rich-T-kid/optimize-take-n-boolBuff (2e2815d) to 5ce0ebe (merge-base) diff

Run configuration

run benchmark boolean_array
env:
  BENCH_FILTER: "take_n_true"

CPU Details (lscpu)

Architecture:                            aarch64
CPU op-mode(s):                          64-bit
Byte Order:                              Little Endian
CPU(s):                                  16
On-line CPU(s) list:                     0-15
Vendor ID:                               ARM
Model name:                              Neoverse-V2
Model:                                   1
Thread(s) per core:                      1
Core(s) per cluster:                     16
Socket(s):                               -
Cluster(s):                              1
Stepping:                                r0p1
BogoMIPS:                                2000.00
Flags:                                   fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache:                               1 MiB (16 instances)
L1i cache:                               1 MiB (16 instances)
L2 cache:                                32 MiB (16 instances)
L3 cache:                                80 MiB (1 instance)
NUMA node(s):                            1
NUMA node0 CPU(s):                       0-15
Vulnerability Gather data sampling:      Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit:             Not affected
Vulnerability L1tf:                      Not affected
Vulnerability Mds:                       Not affected
Vulnerability Meltdown:                  Not affected
Vulnerability Mmio stale data:           Not affected
Vulnerability Reg file data sampling:    Not affected
Vulnerability Retbleed:                  Not affected
Vulnerability Spec rstack overflow:      Not affected
Vulnerability Spec store bypass:         Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:                Mitigation; __user pointer sanitization
Vulnerability Spectre v2:                Mitigation; CSV2, BHB
Vulnerability Srbds:                     Not affected
Vulnerability Tsa:                       Not affected
Vulnerability Tsx async abort:           Not affected
Vulnerability Vmscape:                   Not affected

Details

group                        main                                   rich-T-kid_optimize-take-n-boolBuff
-----                        ----                                   -----------------------------------
take_n_true_shared(32)       1.00    102.5±4.44ns        ? ?/sec    1.01    103.5±1.98ns        ? ?/sec
take_n_true_shared(32768)    1.00     16.3±0.05µs        ? ?/sec    1.00     16.3±0.04µs        ? ?/sec
take_n_true_shared(512)      1.00    344.1±2.98ns        ? ?/sec    1.00    344.5±1.75ns        ? ?/sec
take_n_true_unique(32)       1.45    143.5±4.93ns        ? ?/sec    1.00     98.9±2.92ns        ? ?/sec
take_n_true_unique(32768)    1.00     16.3±0.07µs        ? ?/sec    1.63     26.6±0.04µs        ? ?/sec
take_n_true_unique(512)      1.00    382.3±5.96ns        ? ?/sec    1.27    486.0±4.70ns        ? ?/sec

Resource Usage

base (merge-base)

Metric	Value
Wall time	55.0s
Peak memory	300.1 MiB
Avg memory	85.6 MiB
CPU user	51.2s
CPU sys	1.7s
Peak spill	0 B

branch

Metric	Value
Wall time	60.0s
Peak memory	319.4 MiB
Avg memory	97.5 MiB
CPU user	53.4s
CPU sys	1.7s
Peak spill	0 B

---

File an issue against this benchmark runner

[Jefffrey]

maybe we need to find a faster way to unset the bits?

[Rich-T-kid]

maybe we need to find a faster way to unset the bits?

@Jefffrey setting bits in a loop seems comparably slower than truncating and then appending false all at once. I changed the PR to use the same approach as @/alamb used in his draft pr #10397

[Jefffrey]

how about something like this

        let mut_buffer_result = self.values.into_inner().into_mutable();
        match mut_buffer_result {
            Ok(mut mutable_buffer) => {
                let raw_bytes = mutable_buffer.as_slice_mut();
                let byte_idx_of_end = end / 8;
                let bits_to_preserve = end % 8;
                // end on a byte boundary, so just easily rewrite at byte level
                if bits_to_preserve == 0 {
                    // TODO: this technically can modify bits beyond what the
                    //       boolean buffer actually points to, but given we
                    //       have unique ownership it should be fine? there could
                    //       be pathological case where if this buffer was sliced
                    //       there could be unused bytes that we still process,
                    //       if we wanna bother with that edge case
                    raw_bytes
                        .iter_mut()
                        .skip(byte_idx_of_end)
                        .for_each(|b| *b = 0);
                } else {
                    // if end in middle of a byte, need to unset only higher bits
                    raw_bytes[byte_idx_of_end] &= (1_u8 << bits_to_preserve) - 1;
                    raw_bytes
                        .iter_mut()
                        // +1 since we account for one byte above
                        .skip(byte_idx_of_end + 1)
                        .for_each(|b| *b = 0);
                }
                // TODO: this offset is wrong?
                let boolean_buf = BooleanBuffer::new(mutable_buffer.into(), 0, len);
                BooleanArray::new(boolean_buf, self.nulls)
            }
            Err(buf) => {
                let mut builder = BooleanBufferBuilder::new(len);
                builder.append_buffer(&BooleanBuffer::new(buf, 0, end));
                builder.append_n(len - end, false);
                BooleanArray::new(builder.finish(), self.nulls)
            }
        }

essentially rewrite at the byte level, except for if the end was inside a byte so we need to do some bit ops there

i did a single benchmark run and it seems promising, though i havent carefully checked for edge cases yet 🤔