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    let null_sentinel = null_sentinel(opts);
225    for (value_idx, is_valid) in nulls.iter().enumerate() {
226        let offset = &mut offsets[value_idx + 1];
227        let end_offset = *offset + T::ENCODED_LEN;
228        if is_valid {
229            let to_write = &mut data[*offset..end_offset];
230            to_write[0] = 1;
231            let mut encoded = values[value_idx].encode();
232            if opts.descending {
233                // Flip bits to reverse order
234                encoded.as_mut().iter_mut().for_each(|v| *v = !*v)
235            }
236            to_write[1..].copy_from_slice(encoded.as_ref())
237        } else {
238            data[*offset] = null_sentinel;
239        }
240        *offset = end_offset;
241    }
242}
243
244/// Encoding for non-nullable primitive arrays.
245/// Iterates directly over the `values`, and skips NULLs-checking.
246pub fn encode_not_null<T: FixedLengthEncoding>(
247    data: &mut [u8],
248    offsets: &mut [usize],
249    values: &[T],
250    opts: SortOptions,
251) {
252    for (value_idx, val) in values.iter().enumerate() {
253        let offset = &mut offsets[value_idx + 1];
254        let end_offset = *offset + T::ENCODED_LEN;
255
256        let to_write = &mut data[*offset..end_offset];
257        to_write[0] = 1;
258        let mut encoded = val.encode();
259        if opts.descending {
260            // Flip bits to reverse order
261            encoded.as_mut().iter_mut().for_each(|v| *v = !*v)
262        }
263        to_write[1..].copy_from_slice(encoded.as_ref());
264
265        *offset = end_offset;
266    }
267}
268
269/// Boolean values are encoded as
270///
271/// - 1 byte `0` if null or `1` if valid
272/// - bytes of [`FixedLengthEncoding`]
273pub fn encode_boolean(
274    data: &mut [u8],
275    offsets: &mut [usize],
276    values: &BooleanBuffer,
277    nulls: &NullBuffer,
278    opts: SortOptions,
279) {
280    let null_sentinel = null_sentinel(opts);
281    for (idx, is_valid) in nulls.iter().enumerate() {
282        let offset = &mut offsets[idx + 1];
283        let end_offset = *offset + bool::ENCODED_LEN;
284        if is_valid {
285            let to_write = &mut data[*offset..end_offset];
286            to_write[0] = 1;
287            let mut encoded = values.value(idx).encode();
288            if opts.descending {
289                // Flip bits to reverse order
290                encoded.as_mut().iter_mut().for_each(|v| *v = !*v)
291            }
292            to_write[1..].copy_from_slice(encoded.as_ref())
293        } else {
294            data[*offset] = null_sentinel;
295        }
296        *offset = end_offset;
297    }
298}
299
300/// Encoding for non-nullable boolean arrays.
301/// Iterates directly over `values`, and skips NULLs-checking.
302pub fn encode_boolean_not_null(
303    data: &mut [u8],
304    offsets: &mut [usize],
305    values: &BooleanBuffer,
306    opts: SortOptions,
307) {
308    for (value_idx, val) in values.iter().enumerate() {
309        let offset = &mut offsets[value_idx + 1];
310        let end_offset = *offset + bool::ENCODED_LEN;
311
312        let to_write = &mut data[*offset..end_offset];
313        to_write[0] = 1;
314        let mut encoded = val.encode();
315        if opts.descending {
316            // Flip bits to reverse order
317            encoded.as_mut().iter_mut().for_each(|v| *v = !*v)
318        }
319        to_write[1..].copy_from_slice(encoded.as_ref());
320
321        *offset = end_offset;
322    }
323}
324
325pub fn encode_fixed_size_binary(
326    data: &mut [u8],
327    offsets: &mut [usize],
328    array: &FixedSizeBinaryArray,
329    opts: SortOptions,
330) {
331    let len = array.value_length() as usize;
332    let null_sentinel = null_sentinel(opts);
333    for (offset, maybe_val) in offsets.iter_mut().skip(1).zip(array.iter()) {
334        let end_offset = *offset + len + 1;
335        if let Some(val) = maybe_val {
336            let to_write = &mut data[*offset..end_offset];
337            to_write[0] = 1;
338            to_write[1..].copy_from_slice(&val[..len]);
339            if opts.descending {
340                // Flip bits to reverse order
341                to_write[1..1 + len].iter_mut().for_each(|v| *v = !*v)
342            }
343        } else {
344            data[*offset] = null_sentinel;
345        }
346        *offset = end_offset;
347    }
348}
349
350/// Splits `len` bytes from `src`
351#[inline]
352fn split_off<'a>(src: &mut &'a [u8], len: usize) -> &'a [u8] {
353    let v = &src[..len];
354    *src = &src[len..];
355    v
356}
357
358/// Decodes a `BooleanArray` from rows
359pub fn decode_bool(rows: &mut [&[u8]], options: SortOptions) -> BooleanArray {
360    let true_val = match options.descending {
361        true => !1,
362        false => 1,
363    };
364
365    let len = rows.len();
366
367    let mut nulls = MutableBuffer::new(bit_util::ceil(len, 64) * 8);
368    let mut values = MutableBuffer::new(bit_util::ceil(len, 64) * 8);
369
370    let chunks = len / 64;
371    let remainder = len % 64;
372    for chunk in 0..chunks {
373        let mut null_packed = 0;
374        let mut values_packed = 0;
375
376        for bit_idx in 0..64 {
377            let i = split_off(&mut rows[bit_idx + chunk * 64], 2);
378            let (null, value) = (i[0] == 1, i[1] == true_val);
379            null_packed |= (null as u64) << bit_idx;
380            values_packed |= (value as u64) << bit_idx;
381        }
382
383        nulls.push(null_packed);
384        values.push(values_packed);
385    }
386
387    if remainder != 0 {
388        let mut null_packed = 0;
389        let mut values_packed = 0;
390
391        for bit_idx in 0..remainder {
392            let i = split_off(&mut rows[bit_idx + chunks * 64], 2);
393            let (null, value) = (i[0] == 1, i[1] == true_val);
394            null_packed |= (null as u64) << bit_idx;
395            values_packed |= (value as u64) << bit_idx;
396        }
397
398        nulls.push(null_packed);
399        values.push(values_packed);
400    }
401
402    let nulls = NullBuffer::new(BooleanBuffer::new(nulls.into(), 0, len));
403    let nulls = (nulls.null_count() > 0).then_some(nulls);
404
405    BooleanArray::new(BooleanBuffer::new(values.into(), 0, len), nulls)
406}
407
408/// Decodes a single byte from each row, interpreting `0x01` as a valid value
409/// and all other values as a null. Optionally returns `None` if there are no
410/// null values.
411pub fn decode_nulls(rows: &[&[u8]]) -> Option<NullBuffer> {
412    let nulls = BooleanBuffer::collect_bool(rows.len(), |idx| rows[idx][0] == 1);
413    let nulls = NullBuffer::new(nulls);
414    (nulls.null_count() > 0).then_some(nulls)
415}
416
417/// Decodes a `PrimitiveArray` from rows
418pub fn decode_primitive<T: ArrowPrimitiveType>(
419    rows: &mut [&[u8]],
420    data_type: DataType,
421    options: SortOptions,
422) -> PrimitiveArray<T>
423where
424    T::Native: FixedLengthEncoding,
425{
426    assert!(PrimitiveArray::<T>::is_compatible(&data_type));
427
428    let nulls = decode_nulls(rows);
429    let values = rows
430        .iter_mut()
431        .map(|row| {
432            let i = split_off(row, T::Native::ENCODED_LEN);
433            let value = <T::Native as FixedLengthEncoding>::Encoded::from_slice(
434                &i[1..],
435                options.descending,
436            );
437            T::Native::decode(value)
438        })
439        .collect::<Vec<_>>();
440
441    PrimitiveArray::new(values.into(), nulls).with_data_type(data_type)
442}
443
444/// Decodes a `FixedLengthBinary` from rows
445///
446/// # Panics:
447/// Panics if `size` is negative
448pub fn decode_fixed_size_binary(
449    rows: &mut [&[u8]],
450    size: i32,
451    options: SortOptions,
452) -> FixedSizeBinaryArray {
453    if size < 0 {
454        panic!("cannot decode FixedSizeBinary({size})");
455    }
456    let num_rows = rows.len();
457    let mut values = MutableBuffer::new(size as usize * num_rows);
458    let nulls = decode_nulls(rows);
459
460    let encoded_len = size as usize + 1;
461
462    for row in rows {
463        let i = split_off(row, encoded_len);
464        values.extend_from_slice(&i[1..]);
465    }
466
467    if options.descending {
468        for v in values.as_slice_mut() {
469            *v = !*v;
470        }
471    }
472
473    // Need to set the length since when size is 0 and no nulls the length could not be determined by FixedSizeBinaryArray
474    FixedSizeBinaryArray::try_new_with_len(size, values.into(), nulls, num_rows).unwrap()
475}
arrow_row/fixed.rs

arrow_row/
fixed.rs
