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

arrow_row/
variable.rs
