arrow_row/
fixed.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 crate::array::PrimitiveArray;
19use crate::null_sentinel;
20use arrow_array::{ArrowPrimitiveType, BooleanArray, FixedSizeBinaryArray};
21use arrow_buffer::{
22    BooleanBuffer, IntervalDayTime, IntervalMonthDayNano, MutableBuffer, NullBuffer, bit_util, i256,
23};
24use arrow_schema::{DataType, SortOptions};
25use half::f16;
26
27pub trait FromSlice {
28    fn from_slice(slice: &[u8], invert: bool) -> Self;
29}
30
31impl<const N: usize> FromSlice for [u8; N] {
32    #[inline]
33    fn from_slice(slice: &[u8], invert: bool) -> Self {
34        let mut t: Self = slice.try_into().unwrap();
35        if invert {
36            t.iter_mut().for_each(|o| *o = !*o);
37        }
38        t
39    }
40}
41
42/// Encodes a value of a particular fixed width type into bytes according to the rules
43/// described on [`super::RowConverter`]
44pub trait FixedLengthEncoding: Copy {
45    const ENCODED_LEN: usize = 1 + std::mem::size_of::<Self::Encoded>();
46
47    type Encoded: Sized + Copy + FromSlice + AsRef<[u8]> + AsMut<[u8]>;
48
49    fn encode(self) -> Self::Encoded;
50
51    fn decode(encoded: Self::Encoded) -> Self;
52}
53
54impl FixedLengthEncoding for bool {
55    type Encoded = [u8; 1];
56
57    fn encode(self) -> [u8; 1] {
58        [self as u8]
59    }
60
61    fn decode(encoded: Self::Encoded) -> Self {
62        encoded[0] != 0
63    }
64}
65
66macro_rules! encode_signed {
67    ($n:expr, $t:ty) => {
68        impl FixedLengthEncoding for $t {
69            type Encoded = [u8; $n];
70
71            fn encode(self) -> [u8; $n] {
72                let mut b = self.to_be_bytes();
73                // Toggle top "sign" bit to ensure consistent sort order
74                b[0] ^= 0x80;
75                b
76            }
77
78            fn decode(mut encoded: Self::Encoded) -> Self {
79                // Toggle top "sign" bit
80                encoded[0] ^= 0x80;
81                Self::from_be_bytes(encoded)
82            }
83        }
84    };
85}
86
87encode_signed!(1, i8);
88encode_signed!(2, i16);
89encode_signed!(4, i32);
90encode_signed!(8, i64);
91encode_signed!(16, i128);
92encode_signed!(32, i256);
93
94macro_rules! encode_unsigned {
95    ($n:expr, $t:ty) => {
96        impl FixedLengthEncoding for $t {
97            type Encoded = [u8; $n];
98
99            fn encode(self) -> [u8; $n] {
100                self.to_be_bytes()
101            }
102
103            fn decode(encoded: Self::Encoded) -> Self {
104                Self::from_be_bytes(encoded)
105            }
106        }
107    };
108}
109
110encode_unsigned!(1, u8);
111encode_unsigned!(2, u16);
112encode_unsigned!(4, u32);
113encode_unsigned!(8, u64);
114
115impl FixedLengthEncoding for f16 {
116    type Encoded = [u8; 2];
117
118    fn encode(self) -> [u8; 2] {
119        // https://github.com/rust-lang/rust/blob/9c20b2a8cc7588decb6de25ac6a7912dcef24d65/library/core/src/num/f32.rs#L1176-L1260
120        let s = self.to_bits() as i16;
121        let val = s ^ (((s >> 15) as u16) >> 1) as i16;
122        val.encode()
123    }
124
125    fn decode(encoded: Self::Encoded) -> Self {
126        let bits = i16::decode(encoded);
127        let val = bits ^ (((bits >> 15) as u16) >> 1) as i16;
128        Self::from_bits(val as u16)
129    }
130}
131
132impl FixedLengthEncoding for f32 {
133    type Encoded = [u8; 4];
134
135    fn encode(self) -> [u8; 4] {
136        // https://github.com/rust-lang/rust/blob/9c20b2a8cc7588decb6de25ac6a7912dcef24d65/library/core/src/num/f32.rs#L1176-L1260
137        let s = self.to_bits() as i32;
138        let val = s ^ (((s >> 31) as u32) >> 1) as i32;
139        val.encode()
140    }
141
142    fn decode(encoded: Self::Encoded) -> Self {
143        let bits = i32::decode(encoded);
144        let val = bits ^ (((bits >> 31) as u32) >> 1) as i32;
145        Self::from_bits(val as u32)
146    }
147}
148
149impl FixedLengthEncoding for f64 {
150    type Encoded = [u8; 8];
151
152    fn encode(self) -> [u8; 8] {
153        // https://github.com/rust-lang/rust/blob/9c20b2a8cc7588decb6de25ac6a7912dcef24d65/library/core/src/num/f32.rs#L1176-L1260
154        let s = self.to_bits() as i64;
155        let val = s ^ (((s >> 63) as u64) >> 1) as i64;
156        val.encode()
157    }
158
159    fn decode(encoded: Self::Encoded) -> Self {
160        let bits = i64::decode(encoded);
161        let val = bits ^ (((bits >> 63) as u64) >> 1) as i64;
162        Self::from_bits(val as u64)
163    }
164}
165
166impl FixedLengthEncoding for IntervalDayTime {
167    type Encoded = [u8; 8];
168
169    fn encode(self) -> Self::Encoded {
170        let mut out = [0_u8; 8];
171        out[..4].copy_from_slice(&self.days.encode());
172        out[4..].copy_from_slice(&self.milliseconds.encode());
173        out
174    }
175
176    fn decode(encoded: Self::Encoded) -> Self {
177        Self {
178            days: i32::decode(encoded[..4].try_into().unwrap()),
179            milliseconds: i32::decode(encoded[4..].try_into().unwrap()),
180        }
181    }
182}
183
184impl FixedLengthEncoding for IntervalMonthDayNano {
185    type Encoded = [u8; 16];
186
187    fn encode(self) -> Self::Encoded {
188        let mut out = [0_u8; 16];
189        out[..4].copy_from_slice(&self.months.encode());
190        out[4..8].copy_from_slice(&self.days.encode());
191        out[8..].copy_from_slice(&self.nanoseconds.encode());
192        out
193    }
194
195    fn decode(encoded: Self::Encoded) -> Self {
196        Self {
197            months: i32::decode(encoded[..4].try_into().unwrap()),
198            days: i32::decode(encoded[4..8].try_into().unwrap()),
199            nanoseconds: i64::decode(encoded[8..].try_into().unwrap()),
200        }
201    }
202}
203
204/// Returns the total encoded length (including null byte) for a value of type `T::Native`
205pub const fn encoded_len<T>(_col: &PrimitiveArray<T>) -> usize
206where
207    T: ArrowPrimitiveType,
208    T::Native: FixedLengthEncoding,
209{
210    T::Native::ENCODED_LEN
211}
212
213/// Fixed width types are encoded as
214///
215/// - 1 byte `0` if null or `1` if valid
216/// - bytes of [`FixedLengthEncoding`]
217pub fn encode<T: FixedLengthEncoding>(
218    data: &mut [u8],
219    offsets: &mut [usize],
220    values: &[T],
221    nulls: &NullBuffer,
222    opts: SortOptions,
223) {
224    for (value_idx, is_valid) in nulls.iter().enumerate() {
225        let offset = &mut offsets[value_idx + 1];
226        let end_offset = *offset + T::ENCODED_LEN;
227        if is_valid {
228            let to_write = &mut data[*offset..end_offset];
229            to_write[0] = 1;
230            let mut encoded = values[value_idx].encode();
231            if opts.descending {
232                // Flip bits to reverse order
233                encoded.as_mut().iter_mut().for_each(|v| *v = !*v)
234            }
235            to_write[1..].copy_from_slice(encoded.as_ref())
236        } else {
237            data[*offset] = null_sentinel(opts);
238        }
239        *offset = end_offset;
240    }
241}
242
243/// Encoding for non-nullable primitive arrays.
244/// Iterates directly over the `values`, and skips NULLs-checking.
245pub fn encode_not_null<T: FixedLengthEncoding>(
246    data: &mut [u8],
247    offsets: &mut [usize],
248    values: &[T],
249    opts: SortOptions,
250) {
251    for (value_idx, val) in values.iter().enumerate() {
252        let offset = &mut offsets[value_idx + 1];
253        let end_offset = *offset + T::ENCODED_LEN;
254
255        let to_write = &mut data[*offset..end_offset];
256        to_write[0] = 1;
257        let mut encoded = val.encode();
258        if opts.descending {
259            // Flip bits to reverse order
260            encoded.as_mut().iter_mut().for_each(|v| *v = !*v)
261        }
262        to_write[1..].copy_from_slice(encoded.as_ref());
263
264        *offset = end_offset;
265    }
266}
267
268/// Boolean values are encoded as
269///
270/// - 1 byte `0` if null or `1` if valid
271/// - bytes of [`FixedLengthEncoding`]
272pub fn encode_boolean(
273    data: &mut [u8],
274    offsets: &mut [usize],
275    values: &BooleanBuffer,
276    nulls: &NullBuffer,
277    opts: SortOptions,
278) {
279    for (idx, is_valid) in nulls.iter().enumerate() {
280        let offset = &mut offsets[idx + 1];
281        let end_offset = *offset + bool::ENCODED_LEN;
282        if is_valid {
283            let to_write = &mut data[*offset..end_offset];
284            to_write[0] = 1;
285            let mut encoded = values.value(idx).encode();
286            if opts.descending {
287                // Flip bits to reverse order
288                encoded.as_mut().iter_mut().for_each(|v| *v = !*v)
289            }
290            to_write[1..].copy_from_slice(encoded.as_ref())
291        } else {
292            data[*offset] = null_sentinel(opts);
293        }
294        *offset = end_offset;
295    }
296}
297
298/// Encoding for non-nullable boolean arrays.
299/// Iterates directly over `values`, and skips NULLs-checking.
300pub fn encode_boolean_not_null(
301    data: &mut [u8],
302    offsets: &mut [usize],
303    values: &BooleanBuffer,
304    opts: SortOptions,
305) {
306    for (value_idx, val) in values.iter().enumerate() {
307        let offset = &mut offsets[value_idx + 1];
308        let end_offset = *offset + bool::ENCODED_LEN;
309
310        let to_write = &mut data[*offset..end_offset];
311        to_write[0] = 1;
312        let mut encoded = val.encode();
313        if opts.descending {
314            // Flip bits to reverse order
315            encoded.as_mut().iter_mut().for_each(|v| *v = !*v)
316        }
317        to_write[1..].copy_from_slice(encoded.as_ref());
318
319        *offset = end_offset;
320    }
321}
322
323pub fn encode_fixed_size_binary(
324    data: &mut [u8],
325    offsets: &mut [usize],
326    array: &FixedSizeBinaryArray,
327    opts: SortOptions,
328) {
329    let len = array.value_length() as usize;
330    for (offset, maybe_val) in offsets.iter_mut().skip(1).zip(array.iter()) {
331        let end_offset = *offset + len + 1;
332        if let Some(val) = maybe_val {
333            let to_write = &mut data[*offset..end_offset];
334            to_write[0] = 1;
335            to_write[1..].copy_from_slice(&val[..len]);
336            if opts.descending {
337                // Flip bits to reverse order
338                to_write[1..1 + len].iter_mut().for_each(|v| *v = !*v)
339            }
340        } else {
341            data[*offset] = null_sentinel(opts);
342        }
343        *offset = end_offset;
344    }
345}
346
347/// Splits `len` bytes from `src`
348#[inline]
349fn split_off<'a>(src: &mut &'a [u8], len: usize) -> &'a [u8] {
350    let v = &src[..len];
351    *src = &src[len..];
352    v
353}
354
355/// Decodes a `BooleanArray` from rows
356pub fn decode_bool(rows: &mut [&[u8]], options: SortOptions) -> BooleanArray {
357    let true_val = match options.descending {
358        true => !1,
359        false => 1,
360    };
361
362    let len = rows.len();
363
364    let mut nulls = MutableBuffer::new(bit_util::ceil(len, 64) * 8);
365    let mut values = MutableBuffer::new(bit_util::ceil(len, 64) * 8);
366
367    let chunks = len / 64;
368    let remainder = len % 64;
369    for chunk in 0..chunks {
370        let mut null_packed = 0;
371        let mut values_packed = 0;
372
373        for bit_idx in 0..64 {
374            let i = split_off(&mut rows[bit_idx + chunk * 64], 2);
375            let (null, value) = (i[0] == 1, i[1] == true_val);
376            null_packed |= (null as u64) << bit_idx;
377            values_packed |= (value as u64) << bit_idx;
378        }
379
380        nulls.push(null_packed);
381        values.push(values_packed);
382    }
383
384    if remainder != 0 {
385        let mut null_packed = 0;
386        let mut values_packed = 0;
387
388        for bit_idx in 0..remainder {
389            let i = split_off(&mut rows[bit_idx + chunks * 64], 2);
390            let (null, value) = (i[0] == 1, i[1] == true_val);
391            null_packed |= (null as u64) << bit_idx;
392            values_packed |= (value as u64) << bit_idx;
393        }
394
395        nulls.push(null_packed);
396        values.push(values_packed);
397    }
398
399    let nulls = NullBuffer::new(BooleanBuffer::new(nulls.into(), 0, len));
400    let nulls = (nulls.null_count() > 0).then_some(nulls);
401
402    BooleanArray::new(BooleanBuffer::new(values.into(), 0, len), nulls)
403}
404
405/// Decodes a single byte from each row, interpreting `0x01` as a valid value
406/// and all other values as a null. Optionally returns `None` if there are no
407/// null values.
408pub fn decode_nulls(rows: &[&[u8]]) -> Option<NullBuffer> {
409    let nulls = BooleanBuffer::collect_bool(rows.len(), |idx| rows[idx][0] == 1);
410    let nulls = NullBuffer::new(nulls);
411    (nulls.null_count() > 0).then_some(nulls)
412}
413
414/// Decodes a `PrimitiveArray` from rows
415pub fn decode_primitive<T: ArrowPrimitiveType>(
416    rows: &mut [&[u8]],
417    data_type: DataType,
418    options: SortOptions,
419) -> PrimitiveArray<T>
420where
421    T::Native: FixedLengthEncoding,
422{
423    assert!(PrimitiveArray::<T>::is_compatible(&data_type));
424
425    let nulls = decode_nulls(rows);
426    let values = rows
427        .iter_mut()
428        .map(|row| {
429            let i = split_off(row, T::Native::ENCODED_LEN);
430            let value = <T::Native as FixedLengthEncoding>::Encoded::from_slice(
431                &i[1..],
432                options.descending,
433            );
434            T::Native::decode(value)
435        })
436        .collect::<Vec<_>>();
437
438    PrimitiveArray::new(values.into(), nulls).with_data_type(data_type)
439}
440
441/// Decodes a `FixedLengthBinary` from rows
442///
443/// # Panics:
444/// Panics if `size` is negative
445pub fn decode_fixed_size_binary(
446    rows: &mut [&[u8]],
447    size: i32,
448    options: SortOptions,
449) -> FixedSizeBinaryArray {
450    if size < 0 {
451        panic!("cannot decode FixedSizeBinary({size})");
452    }
453    let num_rows = rows.len();
454    let mut values = MutableBuffer::new(size as usize * num_rows);
455    let nulls = decode_nulls(rows);
456
457    let encoded_len = size as usize + 1;
458
459    for row in rows {
460        let i = split_off(row, encoded_len);
461        values.extend_from_slice(&i[1..]);
462    }
463
464    if options.descending {
465        for v in values.as_slice_mut() {
466            *v = !*v;
467        }
468    }
469
470    // Need to set the length since when size is 0 and no nulls the length could not be determined by FixedSizeBinaryArray
471    FixedSizeBinaryArray::try_new_with_len(size, values.into(), nulls, num_rows).unwrap()
472}
arrow_row/fixed.rs

arrow_row/
fixed.rs
