diff --git a/arrow-arith/src/boolean.rs b/arrow-arith/src/boolean.rs index d94df49de256..6bf438e64618 100644 --- a/arrow-arith/src/boolean.rs +++ b/arrow-arith/src/boolean.rs @@ -23,7 +23,7 @@ //! [here](https://doc.rust-lang.org/stable/core/arch/) for more information. use arrow_array::*; -use arrow_buffer::buffer::{bitwise_bin_op_helper, bitwise_quaternary_op_helper}; +use arrow_buffer::buffer::bitwise_quaternary_op_helper; use arrow_buffer::{BooleanBuffer, NullBuffer, buffer_bin_and_not}; use arrow_schema::ArrowError; @@ -74,7 +74,7 @@ pub fn and_kleene(left: &BooleanArray, right: &BooleanArray) -> Result Result { // Same as above - Some(bitwise_bin_op_helper( + Some(BooleanBuffer::from_bitwise_binary_op( right_null_buffer.buffer(), right_null_buffer.offset(), left_values.inner(), @@ -100,7 +100,7 @@ pub fn and_kleene(left: &BooleanArray, right: &BooleanArray) -> Result Result Result Result { // Same as above - Some(bitwise_bin_op_helper( + Some(BooleanBuffer::from_bitwise_binary_op( right_nulls.buffer(), right_nulls.offset(), left_values.inner(), @@ -195,7 +196,7 @@ pub fn or_kleene(left: &BooleanArray, right: &BooleanArray) -> Result Result( + left: impl AsRef<[u8]>, + left_offset_in_bits: usize, + right: impl AsRef<[u8]>, + right_offset_in_bits: usize, + len_in_bits: usize, + mut op: F, + ) -> Self + where + F: FnMut(u64, u64) -> u64, + { + let left = left.as_ref(); + let right = right.as_ref(); + // try fast path for aligned input + // If the underlying buffers are aligned to u64 we can apply the operation directly on the u64 slices + // to improve performance. + if left_offset_in_bits & 0x7 == 0 && right_offset_in_bits & 0x7 == 0 { + // align to byte boundary + let left = &left[left_offset_in_bits / 8..]; + let right = &right[right_offset_in_bits / 8..]; + + unsafe { + let (left_prefix, left_u64s, left_suffix) = left.align_to::(); + let (right_prefix, right_u64s, right_suffix) = right.align_to::(); + // if there is no prefix or suffix, both buffers are aligned and we can do the operation directly + // on u64s + // TODO also handle non empty suffixes by processing them separately + if left_prefix.is_empty() + && right_prefix.is_empty() + && left_suffix.is_empty() + && right_suffix.is_empty() + { + let result_u64s = left_u64s + .iter() + .zip(right_u64s.iter()) + .map(|(l, r)| op(*l, *r)) + .collect::>(); + return BooleanBuffer { + buffer: Buffer::from(result_u64s), + bit_offset: 0, + bit_len: len_in_bits, + }; + } + } + } + let left_chunks = BitChunks::new(left, left_offset_in_bits, len_in_bits); + let right_chunks = BitChunks::new(right, right_offset_in_bits, len_in_bits); + + let chunks = left_chunks + .iter() + .zip(right_chunks.iter()) + .map(|(left, right)| op(left, right)); + // Soundness: `BitChunks` is a `BitChunks` iterator which + // correctly reports its upper bound + let mut buffer = unsafe { MutableBuffer::from_trusted_len_iter(chunks) }; + + let remainder_bytes = bit_util::ceil(left_chunks.remainder_len(), 8); + let rem = op(left_chunks.remainder_bits(), right_chunks.remainder_bits()); + // we are counting its starting from the least significant bit, to to_le_bytes should be correct + let rem = &rem.to_le_bytes()[0..remainder_bytes]; + buffer.extend_from_slice(rem); + + BooleanBuffer { + buffer: Buffer::from(buffer), + bit_offset: 0, + bit_len: len_in_bits, + } + } + /// Returns the number of set bits in this buffer pub fn count_set_bits(&self) -> usize { self.buffer @@ -655,4 +761,42 @@ mod tests { assert_eq!(result, expected); } } + + #[test] + fn test_from_bitwise_binary_op() { + // pick random boolean inputs + let input_bools_left = (0..1024) + .map(|_| rand::random::()) + .collect::>(); + let input_bools_right = (0..1024) + .map(|_| rand::random::()) + .collect::>(); + let input_buffer_left = BooleanBuffer::from(&input_bools_left[..]); + let input_buffer_right = BooleanBuffer::from(&input_bools_right[..]); + + for left_offset in 0..200 { + for right_offset in [0, 4, 5, 17, 33, 24, 45, 64, 65, 100, 200] { + for len_offset in [0, 1, 44, 100, 256, 300, 512] { + let len = 1024 - len_offset - left_offset.max(right_offset); // ensure we don't go out of bounds + // compute with AND + let result = BooleanBuffer::from_bitwise_binary_op( + input_buffer_left.values(), + left_offset, + input_buffer_right.values(), + right_offset, + len, + |a, b| a & b, + ); + // compute directly from bools + let expected = input_bools_left[left_offset..] + .iter() + .zip(&input_bools_right[right_offset..]) + .take(len) + .map(|(a, b)| *a & *b) + .collect::(); + assert_eq!(result, expected); + } + } + } + } } diff --git a/arrow-buffer/src/buffer/ops.rs b/arrow-buffer/src/buffer/ops.rs index 05593504b1cf..c9e7971709e7 100644 --- a/arrow-buffer/src/buffer/ops.rs +++ b/arrow-buffer/src/buffer/ops.rs @@ -61,35 +61,30 @@ where /// Apply a bitwise operation `op` to two inputs and return the result as a Buffer. /// The inputs are treated as bitmaps, meaning that offsets and length are specified in number of bits. +#[deprecated( + since = "57.2.0", + note = "use BooleanBuffer::from_bitwise_binary_op instead" +)] pub fn bitwise_bin_op_helper( left: &Buffer, left_offset_in_bits: usize, right: &Buffer, right_offset_in_bits: usize, len_in_bits: usize, - mut op: F, + op: F, ) -> Buffer where F: FnMut(u64, u64) -> u64, { - let left_chunks = left.bit_chunks(left_offset_in_bits, len_in_bits); - let right_chunks = right.bit_chunks(right_offset_in_bits, len_in_bits); - - let chunks = left_chunks - .iter() - .zip(right_chunks.iter()) - .map(|(left, right)| op(left, right)); - // Soundness: `BitChunks` is a `BitChunks` iterator which - // correctly reports its upper bound - let mut buffer = unsafe { MutableBuffer::from_trusted_len_iter(chunks) }; - - let remainder_bytes = ceil(left_chunks.remainder_len(), 8); - let rem = op(left_chunks.remainder_bits(), right_chunks.remainder_bits()); - // we are counting its starting from the least significant bit, to to_le_bytes should be correct - let rem = &rem.to_le_bytes()[0..remainder_bytes]; - buffer.extend_from_slice(rem); - - buffer.into() + BooleanBuffer::from_bitwise_binary_op( + left, + left_offset_in_bits, + right, + right_offset_in_bits, + len_in_bits, + op, + ) + .into_inner() } /// Apply a bitwise operation `op` to one input and return the result as a Buffer. @@ -119,7 +114,7 @@ pub fn buffer_bin_and( right_offset_in_bits: usize, len_in_bits: usize, ) -> Buffer { - bitwise_bin_op_helper( + BooleanBuffer::from_bitwise_binary_op( left, left_offset_in_bits, right, @@ -127,6 +122,7 @@ pub fn buffer_bin_and( len_in_bits, |a, b| a & b, ) + .into_inner() } /// Apply a bitwise or to two inputs and return the result as a Buffer. @@ -138,7 +134,7 @@ pub fn buffer_bin_or( right_offset_in_bits: usize, len_in_bits: usize, ) -> Buffer { - bitwise_bin_op_helper( + BooleanBuffer::from_bitwise_binary_op( left, left_offset_in_bits, right, @@ -146,6 +142,7 @@ pub fn buffer_bin_or( len_in_bits, |a, b| a | b, ) + .into_inner() } /// Apply a bitwise xor to two inputs and return the result as a Buffer. @@ -157,7 +154,7 @@ pub fn buffer_bin_xor( right_offset_in_bits: usize, len_in_bits: usize, ) -> Buffer { - bitwise_bin_op_helper( + BooleanBuffer::from_bitwise_binary_op( left, left_offset_in_bits, right, @@ -165,6 +162,7 @@ pub fn buffer_bin_xor( len_in_bits, |a, b| a ^ b, ) + .into_inner() } /// Apply a bitwise and_not to two inputs and return the result as a Buffer. @@ -176,7 +174,7 @@ pub fn buffer_bin_and_not( right_offset_in_bits: usize, len_in_bits: usize, ) -> Buffer { - bitwise_bin_op_helper( + BooleanBuffer::from_bitwise_binary_op( left, left_offset_in_bits, right, @@ -184,11 +182,11 @@ pub fn buffer_bin_and_not( len_in_bits, |a, b| a & !b, ) + .into_inner() } /// Apply a bitwise not to one input and return the result as a Buffer. /// The input is treated as a bitmap, meaning that offset and length are specified in number of bits. pub fn buffer_unary_not(left: &Buffer, offset_in_bits: usize, len_in_bits: usize) -> Buffer { - // TODO: should we deprecate this function in favor of the Buffer ! impl ? BooleanBuffer::from_bitwise_unary_op(left, offset_in_bits, len_in_bits, |a| !a).into_inner() } diff --git a/arrow-select/src/nullif.rs b/arrow-select/src/nullif.rs index 211cabf7afc0..5edaf4dec464 100644 --- a/arrow-select/src/nullif.rs +++ b/arrow-select/src/nullif.rs @@ -18,12 +18,11 @@ //! Implements the `nullif` function for Arrow arrays. use arrow_array::{Array, ArrayRef, BooleanArray, make_array}; -use arrow_buffer::buffer::bitwise_bin_op_helper; use arrow_buffer::{BooleanBuffer, NullBuffer}; use arrow_schema::{ArrowError, DataType}; /// Returns a new array with the same values and the validity bit to false where -/// the corresponding element of`right` is true. +/// the corresponding element of `right` is true. /// /// This can be used to implement SQL `NULLIF` /// @@ -72,40 +71,20 @@ pub fn nullif(left: &dyn Array, right: &BooleanArray) -> Result { - let mut valid_count = 0; - let b = bitwise_bin_op_helper( - left.buffer(), - left.offset(), - right.inner(), - right.offset(), - len, - |l, r| { - let t = l & !r; - valid_count += t.count_ones() as usize; - t - }, - ); - (b, len - valid_count) - } - None => { - let mut null_count = 0; - let buffer = - BooleanBuffer::from_bitwise_unary_op(right.inner(), right.offset(), len, |b| { - let t = !b; - null_count += t.count_zeros() as usize; - t - }) - .into_inner(); - (buffer, null_count) - } + let combined = if let Some(left) = left_data.nulls() && left.null_count() > 0 { + BooleanBuffer::from_bitwise_binary_op( + left.buffer(), + left.offset(), + right.inner(), + right.offset(), + len, + |l, r| l & !r, + ) + } else { + BooleanBuffer::from_bitwise_unary_op(right.inner(), right.offset(), len, |b| !b) }; - let combined = BooleanBuffer::new(combined, 0, len); - // Safety: - // Counted nulls whilst computing - let nulls = unsafe { NullBuffer::new_unchecked(combined, null_count) }; + let nulls = NullBuffer::new(combined); let data = left_data.into_builder().nulls(Some(nulls)); // SAFETY: @@ -494,7 +473,18 @@ mod tests { let r_data = r.to_data(); r_data.validate().unwrap(); - assert_eq!(r.as_ref(), &expected); + assert_eq!( + r.as_ref(), + &expected, + "expected nulls: {:#?}\n\n\ + result nulls: {:#?}\n\n\\ + expected values: {:#?}\n\n\ + result values: {:#?}", + expected.nulls(), + r.nulls(), + expected.values(), + r.as_primitive::().values() + ); } #[test] diff --git a/arrow/Cargo.toml b/arrow/Cargo.toml index 4200cd7a6c78..c07d897ed423 100644 --- a/arrow/Cargo.toml +++ b/arrow/Cargo.toml @@ -306,6 +306,12 @@ name = "bitwise_kernel" harness = false required-features = ["test_utils"] +[[bench]] +name = "nullif_kernel" +harness = false +required-features = ["test_utils"] + + [[bench]] name = "lexsort" harness = false diff --git a/arrow/benches/nullif_kernel.rs b/arrow/benches/nullif_kernel.rs new file mode 100644 index 000000000000..edde39bd7450 --- /dev/null +++ b/arrow/benches/nullif_kernel.rs @@ -0,0 +1,56 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#[macro_use] +extern crate criterion; +use criterion::Criterion; + +use arrow::util::bench_util::{create_boolean_array, create_primitive_array}; + +use arrow::array::*; +use std::hint; +use arrow_array::types::{Int64Type}; +use arrow_select::nullif::nullif; + +fn bench_nullif(left: &dyn Array, right: &BooleanArray) { + hint::black_box(nullif(left, right).unwrap()); +} + + +fn add_benchmark(c: &mut Criterion) { + let size = 8192usize; + + // create input before benchmark to ensure allocations are consistent + let int64_no_nulls = create_primitive_array::(size, 0.0); + let int64_nulls = create_primitive_array::(size, 0.1); + + let mask_10 = create_boolean_array(size, 0.0, 0.1); + let mask_10_sliced = create_boolean_array(size+7, 0.0, 0.1).slice(7, size); + let mask_1 = create_boolean_array(size, 0.0, 0.01); + + c.bench_function("nullif no-nulls mask(10%)", |b| b.iter(|| bench_nullif(&int64_no_nulls, &mask_10))); + c.bench_function("nullif no-nulls mask(10%, sliced)", |b| b.iter(|| bench_nullif(&int64_no_nulls, &mask_10_sliced))); + c.bench_function("nullif no-nulls mask(1%)", |b| b.iter(|| bench_nullif(&int64_no_nulls, &mask_1))); + + c.bench_function("nullif nulls mask(10%)", |b| b.iter(|| bench_nullif(&int64_nulls, &mask_10))); + c.bench_function("nullif nulls mask(10%, sliced)", |b| b.iter(|| bench_nullif(&int64_nulls, &mask_10_sliced))); + c.bench_function("nullif nulls mask(1%)", |b| b.iter(|| bench_nullif(&int64_nulls, &mask_1))); + +} + +criterion_group!(benches, add_benchmark); +criterion_main!(benches); diff --git a/parquet/src/arrow/array_reader/struct_array.rs b/parquet/src/arrow/array_reader/struct_array.rs index b4a6a375334f..8df6a25c9102 100644 --- a/parquet/src/arrow/array_reader/struct_array.rs +++ b/parquet/src/arrow/array_reader/struct_array.rs @@ -129,8 +129,8 @@ impl ArrayReader for StructArrayReader { .len(children_array_len) .child_data( children_array - .into_iter() - .map(|x| x.into_data()) + .iter() + .map(|x| x.to_data()) .collect::>(), );