arrow_row/
variable.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::null_sentinel;
19use arrow_array::builder::BufferBuilder;
20use arrow_array::types::ByteArrayType;
21use arrow_array::*;
22use arrow_buffer::bit_util::ceil;
23use arrow_buffer::{
24    ArrowNativeType, BooleanBuffer, MutableBuffer, NullBuffer, OffsetBuffer, ScalarBuffer,
25};
26use arrow_data::MAX_INLINE_VIEW_LEN;
27use arrow_schema::SortOptions;
28use builder::make_view;
29
30/// The block size of the variable length encoding
31pub const BLOCK_SIZE: usize = 32;
32
33/// The first block is split into `MINI_BLOCK_COUNT` mini-blocks
34///
35/// This helps to reduce the space amplification for small strings
36pub const MINI_BLOCK_COUNT: usize = 4;
37
38/// The mini block size
39pub const MINI_BLOCK_SIZE: usize = BLOCK_SIZE / MINI_BLOCK_COUNT;
40
41/// The continuation token
42pub const BLOCK_CONTINUATION: u8 = 0xFF;
43
44/// Indicates an empty string
45pub const EMPTY_SENTINEL: u8 = 1;
46
47/// Indicates a non-empty string
48pub const NON_EMPTY_SENTINEL: u8 = 2;
49
50/// Indicates a Null value (for DataType::Null)
51pub const NULL_VALUE_SENTINEL: u8 = 3;
52
53/// Returns the padded length of the encoded length of the given length
54#[inline]
55pub fn padded_length(a: Option<usize>) -> usize {
56    match a {
57        Some(a) => non_null_padded_length(a),
58        None => 1,
59    }
60}
61
62/// Returns the padded length of the encoded length of the given length
63#[inline]
64pub(crate) fn non_null_padded_length(len: usize) -> usize {
65    if len <= BLOCK_SIZE {
66        1 + ceil(len, MINI_BLOCK_SIZE) * (MINI_BLOCK_SIZE + 1)
67    } else {
68        // Each miniblock ends with a 1 byte continuation, therefore add
69        // `(MINI_BLOCK_COUNT - 1)` additional bytes over non-miniblock size
70        MINI_BLOCK_COUNT + ceil(len, BLOCK_SIZE) * (BLOCK_SIZE + 1)
71    }
72}
73
74/// Decodes a single byte from each row, where a valid value is determined via
75/// a configurable sentinel. Optionally returns `None` if there are no null values.
76pub(crate) fn decode_nulls_sentinel(rows: &[&[u8]], options: SortOptions) -> Option<NullBuffer> {
77    let null_sentinel = null_sentinel(options);
78    let nulls = BooleanBuffer::collect_bool(rows.len(), |x| rows[x][0] != null_sentinel);
79    let nulls = NullBuffer::new(nulls);
80    (nulls.null_count() > 0).then_some(nulls)
81}
82
83/// Variable length values are encoded as
84///
85/// - single `0_u8` if null
86/// - single `1_u8` if empty array
87/// - `2_u8` if not empty, followed by one or more blocks
88///
89/// where a block is encoded as
90///
91/// - [`BLOCK_SIZE`] bytes of string data, padded with 0s
92/// - `0xFF_u8` if this is not the last block for this string
93/// - otherwise the length of the block as a `u8`
94pub fn encode<'a, I: Iterator<Item = Option<&'a [u8]>>>(
95    data: &mut [u8],
96    offsets: &mut [usize],
97    i: I,
98    opts: SortOptions,
99) {
100    for (offset, maybe_val) in offsets.iter_mut().skip(1).zip(i) {
101        *offset += encode_one(&mut data[*offset..], maybe_val, opts);
102    }
103}
104
105/// Calls [`encode`] with optimized iterator for generic byte arrays
106pub(crate) fn encode_generic_byte_array<T: ByteArrayType>(
107    data: &mut [u8],
108    offsets: &mut [usize],
109    input_array: &GenericByteArray<T>,
110    opts: SortOptions,
111) {
112    let input_offsets = input_array.value_offsets();
113    let bytes = input_array.values().as_slice();
114
115    if let Some(null_buffer) = input_array.nulls().filter(|x| x.null_count() > 0) {
116        let input_iter =
117            input_offsets
118                .windows(2)
119                .zip(null_buffer.iter())
120                .map(|(start_end, is_valid)| {
121                    if is_valid {
122                        let item_range = start_end[0].as_usize()..start_end[1].as_usize();
123                        // SAFETY: the offsets of the input are valid by construction
124                        // so it is ok to use unsafe here
125                        let item = unsafe { bytes.get_unchecked(item_range) };
126                        Some(item)
127                    } else {
128                        None
129                    }
130                });
131
132        encode(data, offsets, input_iter, opts);
133    } else {
134        // Skip null checks
135        let input_iter = input_offsets.windows(2).map(|start_end| {
136            let item_range = start_end[0].as_usize()..start_end[1].as_usize();
137            // SAFETY: the offsets of the input are valid by construction
138            // so it is ok to use unsafe here
139            let item = unsafe { bytes.get_unchecked(item_range) };
140            Some(item)
141        });
142
143        encode(data, offsets, input_iter, opts);
144    }
145}
146
147pub fn encode_null(out: &mut [u8], opts: SortOptions) -> usize {
148    out[0] = null_sentinel(opts);
149    1
150}
151
152pub fn encode_empty(out: &mut [u8], opts: SortOptions) -> usize {
153    out[0] = match opts.descending {
154        true => !EMPTY_SENTINEL,
155        false => EMPTY_SENTINEL,
156    };
157    1
158}
159
160/// Ensure `NullArray`s don't get encoded as empty lists which can lose their length
161pub fn encode_null_value(out: &mut [u8], opts: SortOptions) -> usize {
162    out[0] = match opts.descending {
163        true => !NON_EMPTY_SENTINEL,
164        false => NON_EMPTY_SENTINEL,
165    };
166    out[1] = match opts.descending {
167        true => !NULL_VALUE_SENTINEL,
168        false => NULL_VALUE_SENTINEL,
169    };
170    2
171}
172
173#[inline]
174pub fn encode_one(out: &mut [u8], val: Option<&[u8]>, opts: SortOptions) -> usize {
175    match val {
176        None => encode_null(out, opts),
177        Some([]) => encode_empty(out, opts),
178        Some(val) => {
179            // Write `2_u8` to demarcate as non-empty, non-null string
180            out[0] = NON_EMPTY_SENTINEL;
181
182            let len = if val.len() <= BLOCK_SIZE {
183                1 + encode_blocks::<MINI_BLOCK_SIZE>(&mut out[1..], val)
184            } else {
185                let (initial, rem) = val.split_at(BLOCK_SIZE);
186                let offset = encode_blocks::<MINI_BLOCK_SIZE>(&mut out[1..], initial);
187                out[offset] = BLOCK_CONTINUATION;
188                1 + offset + encode_blocks::<BLOCK_SIZE>(&mut out[1 + offset..], rem)
189            };
190
191            if opts.descending {
192                // Invert bits
193                out[..len].iter_mut().for_each(|v| *v = !*v)
194            }
195            len
196        }
197    }
198}
199
200/// Writes `val` in `SIZE` blocks with the appropriate continuation tokens
201#[inline]
202fn encode_blocks<const SIZE: usize>(out: &mut [u8], val: &[u8]) -> usize {
203    let block_count = ceil(val.len(), SIZE);
204    let end_offset = block_count * (SIZE + 1);
205    let to_write = &mut out[..end_offset];
206
207    let chunks = val.chunks_exact(SIZE);
208    let remainder = chunks.remainder();
209    for (input, output) in chunks.clone().zip(to_write.chunks_exact_mut(SIZE + 1)) {
210        let input: &[u8; SIZE] = input.try_into().unwrap();
211        let out_block: &mut [u8; SIZE] = (&mut output[..SIZE]).try_into().unwrap();
212
213        *out_block = *input;
214
215        // Indicate that there are further blocks to follow
216        output[SIZE] = BLOCK_CONTINUATION;
217    }
218
219    if !remainder.is_empty() {
220        let start_offset = (block_count - 1) * (SIZE + 1);
221        to_write[start_offset..start_offset + remainder.len()].copy_from_slice(remainder);
222        *to_write.last_mut().unwrap() = remainder.len() as u8;
223    } else {
224        // We must overwrite the continuation marker written by the loop above
225        *to_write.last_mut().unwrap() = SIZE as u8;
226    }
227    end_offset
228}
229
230/// Decodes a single block of data
231/// The `f` function accepts a slice of the decoded data, it may be called multiple times
232pub fn decode_blocks(row: &[u8], options: SortOptions, mut f: impl FnMut(&[u8])) -> usize {
233    let (non_empty_sentinel, continuation) = match options.descending {
234        true => (!NON_EMPTY_SENTINEL, !BLOCK_CONTINUATION),
235        false => (NON_EMPTY_SENTINEL, BLOCK_CONTINUATION),
236    };
237
238    if row[0] != non_empty_sentinel {
239        // Empty or null string
240        return 1;
241    }
242
243    // Extracts the block length from the sentinel
244    let block_len = |sentinel: u8| match options.descending {
245        true => !sentinel as usize,
246        false => sentinel as usize,
247    };
248
249    let mut idx = 1;
250    for _ in 0..MINI_BLOCK_COUNT {
251        let sentinel = row[idx + MINI_BLOCK_SIZE];
252        if sentinel != continuation {
253            f(&row[idx..idx + block_len(sentinel)]);
254            return idx + MINI_BLOCK_SIZE + 1;
255        }
256        f(&row[idx..idx + MINI_BLOCK_SIZE]);
257        idx += MINI_BLOCK_SIZE + 1;
258    }
259
260    loop {
261        let sentinel = row[idx + BLOCK_SIZE];
262        if sentinel != continuation {
263            f(&row[idx..idx + block_len(sentinel)]);
264            return idx + BLOCK_SIZE + 1;
265        }
266        f(&row[idx..idx + BLOCK_SIZE]);
267        idx += BLOCK_SIZE + 1;
268    }
269}
270
271/// Returns the number of bytes of encoded data
272fn decoded_len(row: &[u8], options: SortOptions) -> usize {
273    let mut len = 0;
274    decode_blocks(row, options, |block| len += block.len());
275    len
276}
277
278/// Decodes a binary array from `rows` with the provided `options`
279pub fn decode_binary<I: OffsetSizeTrait>(
280    rows: &mut [&[u8]],
281    options: SortOptions,
282) -> GenericBinaryArray<I> {
283    let len = rows.len();
284    let nulls = decode_nulls_sentinel(rows, options);
285
286    let values_capacity = rows.iter().map(|row| decoded_len(row, options)).sum();
287    let mut offsets = BufferBuilder::<I>::new(len + 1);
288    offsets.append(I::zero());
289    let mut values = MutableBuffer::new(values_capacity);
290
291    for row in rows {
292        let offset = decode_blocks(row, options, |b| values.extend_from_slice(b));
293        *row = &row[offset..];
294        offsets.append(I::from_usize(values.len()).expect("offset overflow"))
295    }
296
297    if options.descending {
298        values.as_slice_mut().iter_mut().for_each(|o| *o = !*o)
299    }
300
301    // SAFETY:
302    // Valid by construction above
303    unsafe {
304        GenericBinaryArray::new_unchecked(
305            OffsetBuffer::new(ScalarBuffer::from(offsets)),
306            values.into(),
307            nulls,
308        )
309    }
310}
311
312fn decode_binary_view_inner<const VALIDATE_UTF8: bool>(
313    rows: &mut [&[u8]],
314    options: SortOptions,
315) -> BinaryViewArray {
316    let len = rows.len();
317    let inline_str_max_len = MAX_INLINE_VIEW_LEN as usize;
318
319    let nulls = decode_nulls_sentinel(rows, options);
320
321    // Capacity for all long strings plus room for one short string
322    let mut values_capacity = inline_str_max_len;
323    let mut inline_capacity = 0;
324    for row in rows.iter() {
325        let len = decoded_len(row, options);
326        if len > inline_str_max_len {
327            values_capacity += len;
328        } else if VALIDATE_UTF8 {
329            inline_capacity += len;
330        }
331    }
332    let mut values = MutableBuffer::new(values_capacity);
333    let mut view_utf8_validation_buffer = if VALIDATE_UTF8 {
334        Vec::with_capacity(inline_capacity)
335    } else {
336        Vec::new()
337    };
338
339    let mut views = vec![0_u128; len];
340    for (i, row) in rows.iter_mut().enumerate() {
341        let start_offset = values.len();
342        let offset = decode_blocks(row, options, |b| values.extend_from_slice(b));
343        // Measure string length via change in values buffer. This way we can
344        // overwrite short strings in the values buffer as we inline those to
345        // views.
346        let decoded_len = values.len() - start_offset;
347        if row[0] == null_sentinel(options) {
348            debug_assert_eq!(offset, 1);
349            debug_assert_eq!(start_offset, values.len());
350        } else {
351            // Safety: we just appended the data to the end of the buffer
352            let val = unsafe { values.get_unchecked_mut(start_offset..) };
353
354            if options.descending {
355                val.iter_mut().for_each(|o| *o = !*o);
356            }
357
358            views[i] = make_view(val, 0, start_offset as u32);
359
360            if decoded_len <= inline_str_max_len {
361                if VALIDATE_UTF8 {
362                    view_utf8_validation_buffer.extend_from_slice(val);
363                }
364                values.truncate(start_offset);
365            }
366        }
367        *row = &row[offset..];
368    }
369
370    if VALIDATE_UTF8 {
371        std::str::from_utf8(&values).unwrap();
372        std::str::from_utf8(&view_utf8_validation_buffer).unwrap();
373    }
374
375    // SAFETY:
376    // Valid by construction above
377    unsafe { BinaryViewArray::new_unchecked(views.into(), [values.into()], nulls) }
378}
379
380/// Decodes a binary view array from `rows` with the provided `options`
381pub fn decode_binary_view(rows: &mut [&[u8]], options: SortOptions) -> BinaryViewArray {
382    decode_binary_view_inner::<false>(rows, options)
383}
384
385/// Decodes a string array from `rows` with the provided `options`
386///
387/// # Safety
388///
389/// The row must contain valid UTF-8 data
390pub unsafe fn decode_string<I: OffsetSizeTrait>(
391    rows: &mut [&[u8]],
392    options: SortOptions,
393    validate_utf8: bool,
394) -> GenericStringArray<I> {
395    let decoded = decode_binary::<I>(rows, options);
396
397    if validate_utf8 {
398        return GenericStringArray::from(decoded);
399    }
400
401    let (offsets, values, nulls) = decoded.into_parts();
402
403    // SAFETY:
404    // Row data must have come from a valid UTF-8 array
405    unsafe { GenericStringArray::new_unchecked(offsets, values, nulls) }
406}
407
408/// Decodes a string view array from `rows` with the provided `options`
409///
410/// # Safety
411///
412/// The row must contain valid UTF-8 data
413pub unsafe fn decode_string_view(
414    rows: &mut [&[u8]],
415    options: SortOptions,
416    validate_utf8: bool,
417) -> StringViewArray {
418    let view = if validate_utf8 {
419        decode_binary_view_inner::<true>(rows, options)
420    } else {
421        decode_binary_view_inner::<false>(rows, options)
422    };
423    unsafe { view.to_string_view_unchecked() }
424}
425
426pub fn decode_null_value(rows: &mut [&[u8]], options: SortOptions) {
427    for row in rows.iter_mut() {
428        let (sentinel1, sentinel2) = match options.descending {
429            true => (!NON_EMPTY_SENTINEL, !NULL_VALUE_SENTINEL),
430            false => (NON_EMPTY_SENTINEL, NULL_VALUE_SENTINEL),
431        };
432        debug_assert_eq!(row[0], sentinel1, "Expected NULL_VALUE_SENTINEL at byte 0");
433        debug_assert_eq!(row[1], sentinel2, "Expected NULL_VALUE_SENTINEL at byte 1");
434        *row = &row[2..];
435    }
436}
arrow_row/variable.rs

arrow_row/
variable.rs
