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    bitwise_bin_op_helper(
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}
162
163/// Apply a bitwise or to two inputs and return the result as a Buffer.
164/// The inputs are treated as bitmaps, meaning that offsets and length are specified in number of bits.
165pub fn buffer_bin_or(
166    left: &Buffer,
167    left_offset_in_bits: usize,
168    right: &Buffer,
169    right_offset_in_bits: usize,
170    len_in_bits: usize,
171) -> Buffer {
172    bitwise_bin_op_helper(
173        left,
174        left_offset_in_bits,
175        right,
176        right_offset_in_bits,
177        len_in_bits,
178        |a, b| a | b,
179    )
180}
181
182/// Apply a bitwise xor to two inputs and return the result as a Buffer.
183/// The inputs are treated as bitmaps, meaning that offsets and length are specified in number of bits.
184pub fn buffer_bin_xor(
185    left: &Buffer,
186    left_offset_in_bits: usize,
187    right: &Buffer,
188    right_offset_in_bits: usize,
189    len_in_bits: usize,
190) -> Buffer {
191    bitwise_bin_op_helper(
192        left,
193        left_offset_in_bits,
194        right,
195        right_offset_in_bits,
196        len_in_bits,
197        |a, b| a ^ b,
198    )
199}
200
201/// Apply a bitwise and_not to two inputs and return the result as a Buffer.
202/// The inputs are treated as bitmaps, meaning that offsets and length are specified in number of bits.
203pub fn buffer_bin_and_not(
204    left: &Buffer,
205    left_offset_in_bits: usize,
206    right: &Buffer,
207    right_offset_in_bits: usize,
208    len_in_bits: usize,
209) -> Buffer {
210    bitwise_bin_op_helper(
211        left,
212        left_offset_in_bits,
213        right,
214        right_offset_in_bits,
215        len_in_bits,
216        |a, b| a & !b,
217    )
218}
219
220/// Apply a bitwise not to one input and return the result as a Buffer.
221/// The input is treated as a bitmap, meaning that offset and length are specified in number of bits.
222pub fn buffer_unary_not(left: &Buffer, offset_in_bits: usize, len_in_bits: usize) -> Buffer {
223    // TODO: should we deprecate this function in favor of the Buffer ! impl ?
224    BooleanBuffer::from_bitwise_unary_op(left, offset_in_bits, len_in_bits, |a| !a).into_inner()
225}
arrow_buffer/buffer/ops.rs

arrow_buffer/buffer/
ops.rs
