arrow_buffer/buffer/
ops.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use super::{Buffer, MutableBuffer};
19use crate::BooleanBuffer;
20use crate::util::bit_util::ceil;
21
22/// Apply a bitwise operation `op` to four inputs and return the result as a Buffer.
23///
24/// The inputs are treated as bitmaps, meaning that offsets and length are
25/// specified in number of bits.
26///
27/// NOTE: The operation `op` is applied to chunks of 64 bits (u64) and any bits
28/// outside the offsets and len are set to zero out before calling `op`.
29pub fn bitwise_quaternary_op_helper<F>(
30    buffers: [&Buffer; 4],
31    offsets: [usize; 4],
32    len_in_bits: usize,
33    op: F,
34) -> Buffer
35where
36    F: Fn(u64, u64, u64, u64) -> u64,
37{
38    let first_chunks = buffers[0].bit_chunks(offsets[0], len_in_bits);
39    let second_chunks = buffers[1].bit_chunks(offsets[1], len_in_bits);
40    let third_chunks = buffers[2].bit_chunks(offsets[2], len_in_bits);
41    let fourth_chunks = buffers[3].bit_chunks(offsets[3], len_in_bits);
42
43    let chunks = first_chunks
44        .iter()
45        .zip(second_chunks.iter())
46        .zip(third_chunks.iter())
47        .zip(fourth_chunks.iter())
48        .map(|(((first, second), third), fourth)| op(first, second, third, fourth));
49    // Soundness: `BitChunks` is a `BitChunks` iterator which
50    // correctly reports its upper bound
51    let mut buffer = unsafe { MutableBuffer::from_trusted_len_iter(chunks) };
52
53    let remainder_bytes = ceil(first_chunks.remainder_len(), 8);
54    let rem = op(
55        first_chunks.remainder_bits(),
56        second_chunks.remainder_bits(),
57        third_chunks.remainder_bits(),
58        fourth_chunks.remainder_bits(),
59    );
60    // we are counting its starting from the least significant bit, to to_le_bytes should be correct
61    let rem = &rem.to_le_bytes()[0..remainder_bytes];
62    buffer.extend_from_slice(rem);
63
64    buffer.into()
65}
66
67/// Apply a bitwise operation `op` to two inputs and return the result as a Buffer.
68///
69/// The inputs are treated as bitmaps, meaning that offsets and length are
70/// specified in number of bits.
71///
72/// NOTE: The operation `op` is applied to chunks of 64 bits (u64) and any bits
73/// outside the offsets and len are set to zero out before calling `op`.
74pub fn bitwise_bin_op_helper<F>(
75    left: &Buffer,
76    left_offset_in_bits: usize,
77    right: &Buffer,
78    right_offset_in_bits: usize,
79    len_in_bits: usize,
80    mut op: F,
81) -> Buffer
82where
83    F: FnMut(u64, u64) -> u64,
84{
85    let left_chunks = left.bit_chunks(left_offset_in_bits, len_in_bits);
86    let right_chunks = right.bit_chunks(right_offset_in_bits, len_in_bits);
87
88    let chunks = left_chunks
89        .iter()
90        .zip(right_chunks.iter())
91        .map(|(left, right)| op(left, right));
92    // Soundness: `BitChunks` is a `BitChunks` iterator which
93    // correctly reports its upper bound
94    let mut buffer = unsafe { MutableBuffer::from_trusted_len_iter(chunks) };
95
96    let remainder_bytes = ceil(left_chunks.remainder_len(), 8);
97    let rem = op(left_chunks.remainder_bits(), right_chunks.remainder_bits());
98    // we are counting its starting from the least significant bit, to to_le_bytes should be correct
99    let rem = &rem.to_le_bytes()[0..remainder_bytes];
100    buffer.extend_from_slice(rem);
101
102    buffer.into()
103}
104
105/// Apply a bitwise operation `op` to one input and return the result as a Buffer.
106///
107/// The input is treated as a bitmap, meaning that offset and length are
108/// specified in number of bits.
109///
110/// NOTE: The operation `op` is applied to chunks of 64 bits (u64) and any bits
111/// outside the offsets and len are set to zero out before calling `op`.
112pub fn bitwise_unary_op_helper<F>(
113    left: &Buffer,
114    offset_in_bits: usize,
115    len_in_bits: usize,
116    mut op: F,
117) -> Buffer
118where
119    F: FnMut(u64) -> u64,
120{
121    // reserve capacity and set length so we can get a typed view of u64 chunks
122    let mut result =
123        MutableBuffer::new(ceil(len_in_bits, 8)).with_bitset(len_in_bits / 64 * 8, false);
124
125    let left_chunks = left.bit_chunks(offset_in_bits, len_in_bits);
126
127    let result_chunks = result.typed_data_mut::<u64>().iter_mut();
128
129    result_chunks
130        .zip(left_chunks.iter())
131        .for_each(|(res, left)| {
132            *res = op(left);
133        });
134
135    let remainder_bytes = ceil(left_chunks.remainder_len(), 8);
136    let rem = op(left_chunks.remainder_bits());
137    // we are counting its starting from the least significant bit, to to_le_bytes should be correct
138    let rem = &rem.to_le_bytes()[0..remainder_bytes];
139    result.extend_from_slice(rem);
140
141    result.into()
142}
143
144/// Apply a bitwise and to two inputs and return the result as a Buffer.
145/// The inputs are treated as bitmaps, meaning that offsets and length are specified in number of bits.
146pub fn buffer_bin_and(
147    left: &Buffer,
148    left_offset_in_bits: usize,
149    right: &Buffer,
150    right_offset_in_bits: usize,
151    len_in_bits: usize,
152) -> Buffer {
153    BooleanBuffer::from_bitwise_binary_op(
154        left,
155        left_offset_in_bits,
156        right,
157        right_offset_in_bits,
158        len_in_bits,
159        |a, b| a & b,
160    )
161    .into_inner()
162}
163
164/// Apply a bitwise or to two inputs and return the result as a Buffer.
165/// The inputs are treated as bitmaps, meaning that offsets and length are specified in number of bits.
166pub fn buffer_bin_or(
167    left: &Buffer,
168    left_offset_in_bits: usize,
169    right: &Buffer,
170    right_offset_in_bits: usize,
171    len_in_bits: usize,
172) -> Buffer {
173    BooleanBuffer::from_bitwise_binary_op(
174        left,
175        left_offset_in_bits,
176        right,
177        right_offset_in_bits,
178        len_in_bits,
179        |a, b| a | b,
180    )
181    .into_inner()
182}
183
184/// Apply a bitwise xor to two inputs and return the result as a Buffer.
185/// The inputs are treated as bitmaps, meaning that offsets and length are specified in number of bits.
186pub fn buffer_bin_xor(
187    left: &Buffer,
188    left_offset_in_bits: usize,
189    right: &Buffer,
190    right_offset_in_bits: usize,
191    len_in_bits: usize,
192) -> Buffer {
193    BooleanBuffer::from_bitwise_binary_op(
194        left,
195        left_offset_in_bits,
196        right,
197        right_offset_in_bits,
198        len_in_bits,
199        |a, b| a ^ b,
200    )
201    .into_inner()
202}
203
204/// Apply a bitwise and_not to two inputs and return the result as a Buffer.
205/// The inputs are treated as bitmaps, meaning that offsets and length are specified in number of bits.
206pub fn buffer_bin_and_not(
207    left: &Buffer,
208    left_offset_in_bits: usize,
209    right: &Buffer,
210    right_offset_in_bits: usize,
211    len_in_bits: usize,
212) -> Buffer {
213    BooleanBuffer::from_bitwise_binary_op(
214        left,
215        left_offset_in_bits,
216        right,
217        right_offset_in_bits,
218        len_in_bits,
219        |a, b| a & !b,
220    )
221    .into_inner()
222}
223
224/// Apply a bitwise not to one input and return the result as a Buffer.
225/// The input is treated as a bitmap, meaning that offset and length are specified in number of bits.
226pub fn buffer_unary_not(left: &Buffer, offset_in_bits: usize, len_in_bits: usize) -> Buffer {
227    BooleanBuffer::from_bitwise_unary_op(left, offset_in_bits, len_in_bits, |a| !a).into_inner()
228}
arrow_buffer/buffer/ops.rs

arrow_buffer/buffer/
ops.rs
