Skip to content

Commit

Permalink
Use same codebase for boolean kernels (apache#2508)
Browse files Browse the repository at this point in the history
  • Loading branch information
viirya authored and amrltqt committed Aug 20, 2022
1 parent e50f969 commit b238bad
Show file tree
Hide file tree
Showing 2 changed files with 199 additions and 189 deletions.
47 changes: 47 additions & 0 deletions arrow/src/buffer/ops.rs
Expand Up @@ -18,6 +18,53 @@
use super::{Buffer, MutableBuffer};
use crate::util::bit_util::ceil;

/// Apply a bitwise operation `op` to four 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.
#[allow(clippy::too_many_arguments)]
pub(crate) fn bitwise_quaternary_op_helper<F>(
first: &Buffer,
first_offset_in_bits: usize,
second: &Buffer,
second_offset_in_bits: usize,
third: &Buffer,
third_offset_in_bits: usize,
fourth: &Buffer,
fourth_offset_in_bits: usize,
len_in_bits: usize,
op: F,
) -> Buffer
where
F: Fn(u64, u64, u64, u64) -> u64,
{
let first_chunks = first.bit_chunks(first_offset_in_bits, len_in_bits);
let second_chunks = second.bit_chunks(second_offset_in_bits, len_in_bits);
let third_chunks = third.bit_chunks(third_offset_in_bits, len_in_bits);
let fourth_chunks = fourth.bit_chunks(fourth_offset_in_bits, len_in_bits);

let chunks = first_chunks
.iter()
.zip(second_chunks.iter())
.zip(third_chunks.iter())
.zip(fourth_chunks.iter())
.map(|(((first, second), third), fourth)| op(first, second, third, fourth));
// 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(first_chunks.remainder_len(), 8);
let rem = op(
first_chunks.remainder_bits(),
second_chunks.remainder_bits(),
third_chunks.remainder_bits(),
fourth_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()
}

/// 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.
pub fn bitwise_bin_op_helper<F>(
Expand Down

0 comments on commit b238bad

Please sign in to comment.