Skip to main content

arrow_string/
substring.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
18//! Defines kernel to extract a substring of an Array
19//! Supported array types:
20//! [GenericStringArray], [GenericBinaryArray], [GenericByteViewArray],
21//! [FixedSizeBinaryArray], [DictionaryArray]
22
23use arrow_array::builder::{BinaryViewBuilder, BufferBuilder, StringViewBuilder};
24use arrow_array::cast::AsArray;
25use arrow_array::types::*;
26use arrow_array::*;
27use arrow_buffer::{ArrowNativeType, MutableBuffer, NullBuffer, OffsetBuffer};
28use arrow_schema::{ArrowError, DataType};
29use num_traits::Zero;
30use std::cmp::Ordering;
31use std::sync::Arc;
32
33/// Returns an [`ArrayRef`] with substrings of all the elements in `array`.
34///
35/// # Arguments
36///
37/// * `start` - The start index of all substrings.
38///   If `start >= 0`, then count from the start of the string,
39///   otherwise count from the end of the string.
40///
41/// * `length`(option) - The length of all substrings.
42///   If `length` is [None], then the substring is from `start` to the end of the string.
43///
44/// Attention: Both `start` and `length` are counted by byte, not by char.
45///
46/// # Basic usage
47/// ```
48/// # use arrow_array::StringArray;
49/// # use arrow_string::substring::substring;
50/// let array = StringArray::from(vec![Some("arrow"), None, Some("rust")]);
51/// let result = substring(&array, 1, Some(4)).unwrap();
52/// let result = result.as_any().downcast_ref::<StringArray>().unwrap();
53/// assert_eq!(result, &StringArray::from(vec![Some("rrow"), None, Some("ust")]));
54/// ```
55///
56/// # Error
57/// - The function errors when the passed array is not a [`GenericStringArray`],
58///   [`GenericBinaryArray`], [`FixedSizeBinaryArray`] or [`DictionaryArray`]
59///   with supported array type as its value type.
60/// - The function errors if the offset of a substring in the input array is
61///   at invalid char boundary (only for \[Large\]String array).
62///   It is recommended to use [`substring_by_char`] if the input array may
63///   contain non-ASCII chars.
64///
65/// ## Example of trying to get an invalid utf-8 format substring
66/// ```
67/// # use arrow_array::StringArray;
68/// # use arrow_string::substring::substring;
69/// let array = StringArray::from(vec![Some("E=mc²")]);
70/// let error = substring(&array, 0, Some(5)).unwrap_err().to_string();
71/// assert!(error.contains("invalid utf-8 boundary"));
72/// ```
73pub fn substring(
74    array: &dyn Array,
75    start: i64,
76    length: Option<u64>,
77) -> Result<ArrayRef, ArrowError> {
78    match array.data_type() {
79        DataType::Dictionary(_, _) => {
80            let dictionary = array.as_any_dictionary();
81            let values = substring(dictionary.values(), start, length)?;
82            Ok(Arc::new(dictionary.with_values(values)))
83        }
84        DataType::LargeBinary => {
85            byte_substring(array.as_binary::<i64>(), start, length.map(|e| e as i64))
86        }
87        DataType::Binary => byte_substring(
88            array.as_binary::<i32>(),
89            start as i32,
90            length.map(|e| e as i32),
91        ),
92        DataType::FixedSizeBinary(old_len) => {
93            let old_len: usize = (*old_len)
94                .try_into()
95                .expect("negative FixedSizeBinary value length");
96            fixed_size_binary_substring(array.as_fixed_size_binary(), old_len, start, length)
97        }
98        DataType::LargeUtf8 => {
99            byte_substring(array.as_string::<i64>(), start, length.map(|e| e as i64))
100        }
101        DataType::Utf8 => byte_substring(
102            array.as_string::<i32>(),
103            start as i32,
104            length.map(|e| e as i32),
105        ),
106        DataType::Utf8View => string_view_substring(array.as_string_view(), start, length),
107        DataType::BinaryView => binary_view_substring(array.as_binary_view(), start, length),
108        _ => Err(ArrowError::ComputeError(format!(
109            "substring does not support type {:?}",
110            array.data_type()
111        ))),
112    }
113}
114
115/// Substrings based on character index
116///
117/// # Arguments
118/// * `array` - The input string array
119///
120/// * `start` - The start index of all substrings.
121///   If `start >= 0`, then count from the start of the string,
122///   otherwise count from the end of the string.
123///
124/// * `length`(option) - The length of all substrings.
125///   If `length` is `None`, then the substring is from `start` to the end of the string.
126///
127/// Attention: Both `start` and `length` are counted by char.
128///
129/// # Performance
130///
131/// This function is slower than [substring]. Theoretically, the time complexity
132/// is `O(n)` where `n` is the length of the value buffer. If the array only
133/// contains ASCII chars, a fast path avoids decoding UTF-8 altogether and the
134/// performance is comparable to [substring].
135///
136/// # Basic usage
137/// ```
138/// # use arrow_array::StringArray;
139/// # use arrow_string::substring::substring_by_char;
140/// let array = StringArray::from(vec![Some("arrow"), None, Some("Γ ⊢x:T")]);
141/// let result = substring_by_char(&array, 1, Some(4)).unwrap();
142/// assert_eq!(result, StringArray::from(vec![Some("rrow"), None, Some(" ⊢x:")]));
143/// ```
144pub fn substring_by_char<OffsetSize: OffsetSizeTrait>(
145    array: &GenericStringArray<OffsetSize>,
146    start: i64,
147    length: Option<u64>,
148) -> Result<GenericStringArray<OffsetSize>, ArrowError> {
149    let length = length.map(|len| usize::try_from(len).unwrap_or(usize::MAX));
150
151    if array.is_ascii() {
152        // One char is one byte, so the byte offsets can be computed arithmetically.
153        Ok(substring_by_char_impl(array, length, |val| {
154            ascii_bounds(val, start, length)
155        }))
156    } else {
157        // A char occupies at most 4 bytes.
158        let max_element_len = length.map(|len| len.saturating_mul(4));
159        Ok(substring_by_char_impl(array, max_element_len, |val| {
160            utf8_bounds(val, start, length)
161        }))
162    }
163}
164
165/// Builds the output of [`substring_by_char`], delegating to `bounds` to locate the
166/// substring of each value.
167///
168/// * `max_element_len` - an upper bound, in bytes, on the size of a single output
169///   element, used to avoid over-allocating the value buffer. [None] if unbounded.
170fn substring_by_char_impl<OffsetSize: OffsetSizeTrait, F: Fn(&str) -> (usize, usize)>(
171    array: &GenericStringArray<OffsetSize>,
172    max_element_len: Option<usize>,
173    bounds: F,
174) -> GenericStringArray<OffsetSize> {
175    let mut vals = BufferBuilder::<u8>::new({
176        let offsets = array.value_offsets();
177        let input_len = (offsets[array.len()] - offsets[0]).to_usize().unwrap();
178        match max_element_len {
179            Some(len) => input_len.min(array.len().saturating_mul(len)),
180            None => input_len,
181        }
182    });
183    let mut new_offsets = BufferBuilder::<OffsetSize>::new(array.len() + 1);
184    new_offsets.append(OffsetSize::zero());
185
186    array.iter().for_each(|val| {
187        if let Some(val) = val {
188            let (start_offset, end_offset) = bounds(val);
189            vals.append_slice(&val.as_bytes()[start_offset..end_offset]);
190        }
191        new_offsets.append(OffsetSize::from_usize(vals.len()).unwrap());
192    });
193
194    let offsets = OffsetBuffer::new(new_offsets.finish().into());
195    let values = vals.finish();
196    let nulls = array
197        .nulls()
198        .map(|n| n.inner().sliced())
199        .and_then(|b| NullBuffer::from_unsliced_buffer(b, array.len()));
200    GenericStringArray::<OffsetSize>::new(offsets, values, nulls)
201}
202
203/// Returns the `start` and `end` byte offset of the substring of an ASCII `val`, where
204/// one char is one byte.
205///
206/// * `start` - the start char index of the substring, counted from the end if negative
207/// * `length` - the char length of the substring, or [None] to take the rest of `val`
208#[inline]
209fn ascii_bounds(val: &str, start: i64, length: Option<usize>) -> (usize, usize) {
210    let len = val.len();
211    let start_offset = if start >= 0 {
212        usize::try_from(start).unwrap_or(usize::MAX).min(len)
213    } else {
214        len.saturating_sub(usize::try_from(start.unsigned_abs()).unwrap_or(usize::MAX))
215    };
216    let end_offset = length.map_or(len, |length| start_offset.saturating_add(length).min(len));
217    (start_offset, end_offset)
218}
219
220/// Returns the `start` and `end` byte offset of the substring of an arbitrary UTF-8 `val`.
221///
222/// * `start` - the start char index of the substring, counted from the end if negative
223/// * `length` - the char length of the substring, or [None] to take the rest of `val`
224#[inline]
225fn utf8_bounds(val: &str, start: i64, length: Option<usize>) -> (usize, usize) {
226    let len = val.len();
227    let start_offset = if start >= 0 {
228        val.char_indices()
229            .nth(usize::try_from(start).unwrap_or(usize::MAX))
230            .map_or(len, |(offset, _)| offset)
231    } else {
232        // `start` is negative, so `nth_back` counts back from the last char. Strings with
233        // fewer than `-start` chars start at 0.
234        let back = usize::try_from(start.unsigned_abs()).unwrap_or(usize::MAX);
235        val.char_indices()
236            .nth_back(back - 1)
237            .map_or(0, |(offset, _)| offset)
238    };
239    let end_offset = length.map_or(len, |length| {
240        // Every char is at least one byte, so the rest of `val` is shorter than `length`
241        // chars and there is nothing to scan for.
242        if length >= len - start_offset {
243            return len;
244        }
245        val[start_offset..]
246            .char_indices()
247            .nth(length)
248            .map_or(len, |(offset, _)| start_offset + offset)
249    });
250    (start_offset, end_offset)
251}
252
253/// Byte range of one element, following the same rules as [`byte_substring`].
254fn view_substring_range(
255    original_length: usize,
256    start: i64,
257    substring_length: Option<u64>,
258) -> (usize, usize) {
259    let original_length = original_length as i64;
260    let new_start = match start.cmp(&0) {
261        Ordering::Greater => start.min(original_length),
262        Ordering::Equal => 0,
263        Ordering::Less => (original_length + start).max(0),
264    };
265    let new_end = match substring_length {
266        Some(length) => new_start.saturating_add(length as i64).min(original_length),
267        None => original_length,
268    };
269    (new_start as usize, new_end as usize)
270}
271
272fn string_view_substring(
273    array: &StringViewArray,
274    start: i64,
275    length: Option<u64>,
276) -> Result<ArrayRef, ArrowError> {
277    let mut builder = StringViewBuilder::with_capacity(array.len());
278
279    for idx in 0..array.len() {
280        if array.is_null(idx) {
281            builder.append_null();
282            continue;
283        }
284        let value = array.value(idx);
285        let (new_start, new_end) = view_substring_range(value.len(), start, length);
286        for offset in [new_start, new_end] {
287            if !value.is_char_boundary(offset) {
288                return Err(ArrowError::ComputeError(format!(
289                    "The offset {offset} is at an invalid utf-8 boundary."
290                )));
291            }
292        }
293        builder.append_value(&value[new_start..new_end]);
294    }
295
296    Ok(Arc::new(builder.finish()))
297}
298
299fn binary_view_substring(
300    array: &BinaryViewArray,
301    start: i64,
302    length: Option<u64>,
303) -> Result<ArrayRef, ArrowError> {
304    let mut builder = BinaryViewBuilder::with_capacity(array.len());
305
306    for idx in 0..array.len() {
307        if array.is_null(idx) {
308            builder.append_null();
309            continue;
310        }
311        let value = array.value(idx);
312        let (new_start, new_end) = view_substring_range(value.len(), start, length);
313        builder.append_value(&value[new_start..new_end]);
314    }
315
316    Ok(Arc::new(builder.finish()))
317}
318
319fn byte_substring<T: ByteArrayType>(
320    array: &GenericByteArray<T>,
321    start: T::Offset,
322    length: Option<T::Offset>,
323) -> Result<ArrayRef, ArrowError>
324where
325    <T as ByteArrayType>::Native: PartialEq,
326{
327    let offsets = array.value_offsets();
328    let data = array.value_data();
329    let zero = <T::Offset as Zero>::zero();
330
331    // When array is [Large]StringArray, we will check whether `offset` is at a valid char boundary.
332    let check_char_boundary = {
333        |offset: T::Offset| {
334            if !matches!(T::DATA_TYPE, DataType::Utf8 | DataType::LargeUtf8) {
335                return Ok(offset);
336            }
337            // Safety: a StringArray must contain valid UTF8 data
338            let data_str = unsafe { std::str::from_utf8_unchecked(data) };
339            let offset_usize = offset.as_usize();
340            if data_str.is_char_boundary(offset_usize) {
341                Ok(offset)
342            } else {
343                Err(ArrowError::ComputeError(format!(
344                    "The offset {offset_usize} is at an invalid utf-8 boundary."
345                )))
346            }
347        }
348    };
349
350    // start and end offsets of all substrings
351    let mut new_starts_ends: Vec<(T::Offset, T::Offset)> = Vec::with_capacity(array.len());
352    let mut new_offsets: Vec<T::Offset> = Vec::with_capacity(array.len() + 1);
353    let mut len_so_far = zero;
354    new_offsets.push(zero);
355
356    offsets
357        .windows(2)
358        .try_for_each(|pair| -> Result<(), ArrowError> {
359            let new_start = match start.cmp(&zero) {
360                Ordering::Greater => check_char_boundary((pair[0] + start).min(pair[1]))?,
361                Ordering::Equal => pair[0],
362                Ordering::Less => check_char_boundary((pair[1] + start).max(pair[0]))?,
363            };
364            let new_end = match length {
365                Some(length) => check_char_boundary((length + new_start).min(pair[1]))?,
366                None => pair[1],
367            };
368            len_so_far += new_end - new_start;
369            new_starts_ends.push((new_start, new_end));
370            new_offsets.push(len_so_far);
371            Ok(())
372        })?;
373
374    // concatenate substrings into a buffer
375    let mut new_values = MutableBuffer::new(new_offsets.last().unwrap().as_usize());
376
377    new_starts_ends
378        .iter()
379        .map(|(start, end)| {
380            let start = start.as_usize();
381            let end = end.as_usize();
382            &data[start..end]
383        })
384        .try_for_each(|slice| {
385            new_values
386                .try_extend_from_slice(slice)
387                .map_err(|e| ArrowError::MemoryError(e.to_string()))
388        })?;
389
390    let offsets = OffsetBuffer::new(new_offsets.into());
391    let values = new_values.into();
392    let nulls = array
393        .nulls()
394        .map(|n| n.inner().sliced())
395        .and_then(|b| NullBuffer::from_unsliced_buffer(b, array.len()));
396    Ok(Arc::new(GenericByteArray::<T>::new(offsets, values, nulls)))
397}
398
399fn fixed_size_binary_substring(
400    array: &FixedSizeBinaryArray,
401    old_len: usize,
402    start: i64,
403    length: Option<u64>,
404) -> Result<ArrayRef, ArrowError> {
405    let new_start = match start.cmp(&0) {
406        Ordering::Greater => usize::try_from(start).unwrap_or(usize::MAX).min(old_len),
407        Ordering::Equal => 0,
408        Ordering::Less => {
409            let offset = usize::try_from(start.unsigned_abs()).unwrap_or(usize::MAX);
410            old_len.saturating_sub(offset)
411        }
412    };
413
414    let new_len = match length {
415        Some(len) => usize::try_from(len)
416            .unwrap_or(usize::MAX)
417            .min(old_len - new_start),
418        None => old_len - new_start,
419    };
420
421    // build value buffer
422    let num_of_elements = array.len();
423    let data = array.value_data();
424    let capacity = num_of_elements
425        .checked_mul(new_len)
426        .expect("capacity overflow");
427    let mut new_values = MutableBuffer::new(capacity);
428    (0..num_of_elements)
429        .map(|idx| {
430            let offset = idx * array.value_size();
431            (offset + new_start, offset + new_start + new_len)
432        })
433        .try_for_each(|(start, end)| {
434            new_values
435                .try_extend_from_slice(&data[start..end])
436                .map_err(|e| ArrowError::MemoryError(e.to_string()))
437        })?;
438
439    let mut nulls = array
440        .nulls()
441        .map(|n| n.inner().sliced())
442        .and_then(|b| NullBuffer::from_unsliced_buffer(b, num_of_elements));
443
444    if new_len == 0 && nulls.is_none() {
445        // FixedSizeBinaryArray::new takes length from the values buffer, except when size == 0.
446        // In that case it uses the null buffer length, so preserve the original length here.
447        // Example: ["", "", ""] -> substring(..., 1, Some(2)) should keep len=3;
448        // otherwise it collapses to an empty array (len=0).
449        nulls = Some(NullBuffer::new_valid(num_of_elements));
450    }
451
452    let new_len: i32 = new_len.try_into().expect("new_len overflow");
453
454    Ok(Arc::new(FixedSizeBinaryArray::new(
455        new_len,
456        new_values.into(),
457        nulls,
458    )))
459}
460
461#[cfg(test)]
462mod tests {
463    use super::*;
464    use arrow_buffer::BooleanBuffer;
465    use arrow_buffer::Buffer;
466
467    /// A helper macro to generate test cases.
468    /// # Arguments
469    /// * `input` - A vector which array can be built from.
470    /// * `start` - The start index of the substring.
471    /// * `len` - The length of the substring.
472    /// * `result` - The expected result of substring, which is a vector that array can be built from.
473    /// # Return
474    /// A vector of `(input, start, len, result)`.
475    ///
476    /// Users can provide any number of `(start, len, result)` to generate test cases for one `input`.
477    macro_rules! gen_test_cases {
478        ($input:expr, $(($start:expr, $len:expr, $result:expr)), *) => {
479            [
480                $(
481                    ($input.clone(), $start, $len, $result),
482                )*
483            ]
484        };
485    }
486
487    /// A helper macro to test the substring functions.
488    /// # Arguments
489    /// * `cases` - The test cases which is a vector of `(input, start, len, result)`.
490    ///   Please look at [`gen_test_cases`] to find how to generate it.
491    /// * `array_ty` - The array type.
492    /// * `substring_fn` - Either [`substring`] or [`substring_by_char`].
493    macro_rules! do_test {
494        ($cases:expr, $array_ty:ty, $substring_fn:ident) => {
495            $cases
496                .into_iter()
497                .for_each(|(array, start, length, expected)| {
498                    let array = <$array_ty>::from(array);
499                    let result = $substring_fn(&array, start, length).unwrap();
500                    let result = result.as_any().downcast_ref::<$array_ty>().unwrap();
501                    let expected = <$array_ty>::from(expected);
502                    assert_eq!(&expected, result);
503                })
504        };
505    }
506
507    /// A helper macro to test the substring functions for array types only implementing TryFrom.
508    macro_rules! do_test_tryfrom {
509        ($cases:expr, $array_ty:ty, $substring_fn:ident) => {
510            $cases
511                .into_iter()
512                .for_each(|(array, start, length, expected)| {
513                    let array = <$array_ty>::try_from(array).unwrap();
514                    let result = $substring_fn(&array, start, length).unwrap();
515                    let result = result.as_any().downcast_ref::<$array_ty>().unwrap();
516                    let expected = <$array_ty>::try_from(expected).unwrap();
517                    assert_eq!(&expected, result);
518                })
519        };
520    }
521
522    fn with_nulls_generic_binary<O: OffsetSizeTrait>() {
523        let input = vec![
524            Some("hello".as_bytes()),
525            None,
526            Some(&[0xf8, 0xf9, 0xff, 0xfa]),
527        ];
528        // all-nulls array is always identical
529        let base_case = gen_test_cases!(
530            vec![None, None, None],
531            (-1, Some(1), vec![None, None, None])
532        );
533        let cases = gen_test_cases!(
534            input,
535            // identity
536            (0, None, input.clone()),
537            // 0 length -> Nothing
538            (0, Some(0), vec![Some(&[]), None, Some(&[])]),
539            // high start -> Nothing
540            (1000, Some(0), vec![Some(&[]), None, Some(&[])]),
541            // high negative start -> identity
542            (-1000, None, input.clone()),
543            // high length -> identity
544            (0, Some(1000), input.clone())
545        );
546
547        do_test!(
548            [&base_case[..], &cases[..]].concat(),
549            GenericBinaryArray<O>,
550            substring
551        );
552    }
553
554    #[test]
555    fn with_nulls_binary() {
556        with_nulls_generic_binary::<i32>()
557    }
558
559    #[test]
560    fn with_nulls_large_binary() {
561        with_nulls_generic_binary::<i64>()
562    }
563
564    fn without_nulls_generic_binary<O: OffsetSizeTrait>() {
565        let input = vec!["hello".as_bytes(), b"", &[0xf8, 0xf9, 0xff, 0xfa]];
566        // empty array is always identical
567        let base_case = gen_test_cases!(
568            vec!["".as_bytes(), b"", b""],
569            (2, Some(1), vec!["".as_bytes(), b"", b""])
570        );
571        let cases = gen_test_cases!(
572            input,
573            // identity
574            (0, None, input.clone()),
575            // increase start
576            (1, None, vec![b"ello", b"", &[0xf9, 0xff, 0xfa]]),
577            (2, None, vec![b"llo", b"", &[0xff, 0xfa]]),
578            (3, None, vec![b"lo", b"", &[0xfa]]),
579            (10, None, vec![b"", b"", b""]),
580            // increase start negatively
581            (-1, None, vec![b"o", b"", &[0xfa]]),
582            (-2, None, vec![b"lo", b"", &[0xff, 0xfa]]),
583            (-3, None, vec![b"llo", b"", &[0xf9, 0xff, 0xfa]]),
584            (-10, None, input.clone()),
585            // increase length
586            (1, Some(1), vec![b"e", b"", &[0xf9]]),
587            (1, Some(2), vec![b"el", b"", &[0xf9, 0xff]]),
588            (1, Some(3), vec![b"ell", b"", &[0xf9, 0xff, 0xfa]]),
589            (1, Some(4), vec![b"ello", b"", &[0xf9, 0xff, 0xfa]]),
590            (-3, Some(1), vec![b"l", b"", &[0xf9]]),
591            (-3, Some(2), vec![b"ll", b"", &[0xf9, 0xff]]),
592            (-3, Some(3), vec![b"llo", b"", &[0xf9, 0xff, 0xfa]]),
593            (-3, Some(4), vec![b"llo", b"", &[0xf9, 0xff, 0xfa]])
594        );
595
596        do_test!(
597            [&base_case[..], &cases[..]].concat(),
598            GenericBinaryArray<O>,
599            substring
600        );
601    }
602
603    #[test]
604    fn without_nulls_binary() {
605        without_nulls_generic_binary::<i32>()
606    }
607
608    #[test]
609    fn without_nulls_large_binary() {
610        without_nulls_generic_binary::<i64>()
611    }
612
613    fn generic_binary_with_non_zero_offset<O: OffsetSizeTrait>() {
614        let values = 0_u8..15;
615        let offsets = &[
616            O::zero(),
617            O::from_usize(5).unwrap(),
618            O::from_usize(10).unwrap(),
619            O::from_usize(15).unwrap(),
620        ];
621        // set the first and third element to be valid
622        let bitmap = [0b101_u8];
623
624        let offsets = OffsetBuffer::new(Buffer::from_slice_ref(offsets).into());
625        let values = Buffer::from_iter(values);
626        let nulls = Some(NullBuffer::new(BooleanBuffer::new(
627            Buffer::from(bitmap),
628            0,
629            3,
630        )));
631        // array is `[null, [10, 11, 12, 13, 14]]`
632        let array = GenericBinaryArray::<O>::new(offsets, values, nulls).slice(1, 2);
633        // result is `[null, [11, 12, 13, 14]]`
634        let result = substring(&array, 1, None).unwrap();
635        let result = result
636            .as_any()
637            .downcast_ref::<GenericBinaryArray<O>>()
638            .unwrap();
639        let expected =
640            GenericBinaryArray::<O>::from_opt_vec(vec![None, Some(&[11_u8, 12, 13, 14])]);
641        assert_eq!(result, &expected);
642    }
643
644    #[test]
645    fn binary_with_non_zero_offset() {
646        generic_binary_with_non_zero_offset::<i32>()
647    }
648
649    #[test]
650    fn large_binary_with_non_zero_offset() {
651        generic_binary_with_non_zero_offset::<i64>()
652    }
653
654    #[test]
655    fn with_nulls_fixed_size_binary() {
656        let input = vec![Some("cat".as_bytes()), None, Some(&[0xf8, 0xf9, 0xff])];
657        // all-nulls array is always identical
658        let base_case =
659            gen_test_cases!(vec![None, None, None], (3, Some(2), vec![None, None, None]));
660        let cases = gen_test_cases!(
661            input,
662            // identity
663            (0, None, input.clone()),
664            // increase start
665            (1, None, vec![Some(b"at"), None, Some(&[0xf9, 0xff])]),
666            (2, None, vec![Some(b"t"), None, Some(&[0xff])]),
667            (3, None, vec![Some(b""), None, Some(b"")]),
668            (10, None, vec![Some(b""), None, Some(b"")]),
669            // increase start negatively
670            (-1, None, vec![Some(b"t"), None, Some(&[0xff])]),
671            (-2, None, vec![Some(b"at"), None, Some(&[0xf9, 0xff])]),
672            (-3, None, input.clone()),
673            (-10, None, input.clone()),
674            // increase length
675            (1, Some(1), vec![Some(b"a"), None, Some(&[0xf9])]),
676            (1, Some(2), vec![Some(b"at"), None, Some(&[0xf9, 0xff])]),
677            (1, Some(3), vec![Some(b"at"), None, Some(&[0xf9, 0xff])]),
678            (-3, Some(1), vec![Some(b"c"), None, Some(&[0xf8])]),
679            (-3, Some(2), vec![Some(b"ca"), None, Some(&[0xf8, 0xf9])]),
680            (-3, Some(3), input.clone()),
681            (-3, Some(4), input.clone())
682        );
683
684        do_test_tryfrom!(
685            [&base_case[..], &cases[..]].concat(),
686            FixedSizeBinaryArray,
687            substring
688        );
689    }
690
691    #[test]
692    fn without_nulls_fixed_size_binary() {
693        let input = vec!["cat".as_bytes(), b"dog", &[0xf8, 0xf9, 0xff]];
694        // empty array is always identical
695        let base_case = gen_test_cases!(
696            vec!["".as_bytes(), &[], &[]],
697            (1, Some(2), vec!["".as_bytes(), &[], &[]])
698        );
699        let cases = gen_test_cases!(
700            input,
701            // identity
702            (0, None, input.clone()),
703            // increase start
704            (1, None, vec![b"at", b"og", &[0xf9, 0xff]]),
705            (2, None, vec![b"t", b"g", &[0xff]]),
706            (3, None, vec![&[], &[], &[]]),
707            (10, None, vec![&[], &[], &[]]),
708            // increase start negatively
709            (-1, None, vec![b"t", b"g", &[0xff]]),
710            (-2, None, vec![b"at", b"og", &[0xf9, 0xff]]),
711            (-3, None, input.clone()),
712            (-10, None, input.clone()),
713            // increase length
714            (1, Some(1), vec![b"a", b"o", &[0xf9]]),
715            (1, Some(2), vec![b"at", b"og", &[0xf9, 0xff]]),
716            (1, Some(3), vec![b"at", b"og", &[0xf9, 0xff]]),
717            (-3, Some(1), vec![b"c", b"d", &[0xf8]]),
718            (-3, Some(2), vec![b"ca", b"do", &[0xf8, 0xf9]]),
719            (-3, Some(3), input.clone()),
720            (-3, Some(4), input.clone())
721        );
722
723        do_test_tryfrom!(
724            [&base_case[..], &cases[..]].concat(),
725            FixedSizeBinaryArray,
726            substring
727        );
728    }
729
730    #[test]
731    fn fixed_size_binary_with_non_zero_offset() {
732        let values = b"hellotherearrow";
733        // set the first and third element to be valid
734        let bits_v = [0b101_u8];
735
736        let nulls = Some(NullBuffer::new(BooleanBuffer::new(
737            Buffer::from(bits_v),
738            0,
739            3,
740        )));
741        // array is `[null, "arrow"]`
742        let array = FixedSizeBinaryArray::new(5, Buffer::from(values), nulls).slice(1, 2);
743        // result is `[null, "rrow"]`
744        let result = substring(&array, 1, None).unwrap();
745        let result = result
746            .as_any()
747            .downcast_ref::<FixedSizeBinaryArray>()
748            .unwrap();
749        let expected = FixedSizeBinaryArray::try_from_sparse_iter_with_size(
750            vec![None, Some(b"rrow")].into_iter(),
751            4,
752        )
753        .unwrap();
754        assert_eq!(result, &expected);
755    }
756
757    fn with_nulls_generic_string<O: OffsetSizeTrait>() {
758        let input = vec![Some("hello"), None, Some("word")];
759        // all-nulls array is always identical
760        let base_case = gen_test_cases!(vec![None, None, None], (0, None, vec![None, None, None]));
761        let cases = gen_test_cases!(
762            input,
763            // identity
764            (0, None, input.clone()),
765            // 0 length -> Nothing
766            (0, Some(0), vec![Some(""), None, Some("")]),
767            // high start -> Nothing
768            (1000, Some(0), vec![Some(""), None, Some("")]),
769            // high negative start -> identity
770            (-1000, None, input.clone()),
771            // high length -> identity
772            (0, Some(1000), input.clone())
773        );
774
775        do_test!(
776            [&base_case[..], &cases[..]].concat(),
777            GenericStringArray<O>,
778            substring
779        );
780    }
781
782    #[test]
783    fn with_nulls_string() {
784        with_nulls_generic_string::<i32>()
785    }
786
787    #[test]
788    fn with_nulls_large_string() {
789        with_nulls_generic_string::<i64>()
790    }
791
792    fn without_nulls_generic_string<O: OffsetSizeTrait>() {
793        let input = vec!["hello", "", "word"];
794        // empty array is always identical
795        let base_case = gen_test_cases!(vec!["", "", ""], (0, None, vec!["", "", ""]));
796        let cases = gen_test_cases!(
797            input,
798            // identity
799            (0, None, input.clone()),
800            (1, None, vec!["ello", "", "ord"]),
801            (2, None, vec!["llo", "", "rd"]),
802            (3, None, vec!["lo", "", "d"]),
803            (10, None, vec!["", "", ""]),
804            // increase start negatively
805            (-1, None, vec!["o", "", "d"]),
806            (-2, None, vec!["lo", "", "rd"]),
807            (-3, None, vec!["llo", "", "ord"]),
808            (-10, None, input.clone()),
809            // increase length
810            (1, Some(1), vec!["e", "", "o"]),
811            (1, Some(2), vec!["el", "", "or"]),
812            (1, Some(3), vec!["ell", "", "ord"]),
813            (1, Some(4), vec!["ello", "", "ord"]),
814            (-3, Some(1), vec!["l", "", "o"]),
815            (-3, Some(2), vec!["ll", "", "or"]),
816            (-3, Some(3), vec!["llo", "", "ord"]),
817            (-3, Some(4), vec!["llo", "", "ord"])
818        );
819
820        do_test!(
821            [&base_case[..], &cases[..]].concat(),
822            GenericStringArray<O>,
823            substring
824        );
825    }
826
827    #[test]
828    fn without_nulls_string() {
829        without_nulls_generic_string::<i32>()
830    }
831
832    #[test]
833    fn without_nulls_large_string() {
834        without_nulls_generic_string::<i64>()
835    }
836
837    fn generic_string_with_non_zero_offset<O: OffsetSizeTrait>() {
838        let values = b"hellotherearrow";
839        let offsets = &[
840            O::zero(),
841            O::from_usize(5).unwrap(),
842            O::from_usize(10).unwrap(),
843            O::from_usize(15).unwrap(),
844        ];
845        // set the first and third element to be valid
846        let bitmap = [0b101_u8];
847
848        let offsets = OffsetBuffer::new(Buffer::from_slice_ref(offsets).into());
849        let values = Buffer::from(values);
850        let nulls = Some(NullBuffer::new(BooleanBuffer::new(
851            Buffer::from(bitmap),
852            0,
853            3,
854        )));
855        // array is `[null, "arrow"]`
856        let array = GenericStringArray::<O>::new(offsets, values, nulls).slice(1, 2);
857        // result is `[null, "rrow"]`
858        let result = substring(&array, 1, None).unwrap();
859        let result = result
860            .as_any()
861            .downcast_ref::<GenericStringArray<O>>()
862            .unwrap();
863        let expected = GenericStringArray::<O>::from(vec![None, Some("rrow")]);
864        assert_eq!(result, &expected);
865    }
866
867    #[test]
868    fn string_with_non_zero_offset() {
869        generic_string_with_non_zero_offset::<i32>()
870    }
871
872    #[test]
873    fn large_string_with_non_zero_offset() {
874        generic_string_with_non_zero_offset::<i64>()
875    }
876
877    fn with_nulls_generic_string_by_char<O: OffsetSizeTrait>() {
878        let input = vec![Some("hello"), None, Some("Γ ⊢x:T")];
879        // all-nulls array is always identical
880        let base_case = gen_test_cases!(vec![None, None, None], (0, None, vec![None, None, None]));
881        let cases = gen_test_cases!(
882            input,
883            // identity
884            (0, None, input.clone()),
885            // 0 length -> Nothing
886            (0, Some(0), vec![Some(""), None, Some("")]),
887            // high start -> Nothing
888            (1000, Some(0), vec![Some(""), None, Some("")]),
889            // high negative start -> identity
890            (-1000, None, input.clone()),
891            // high length -> identity
892            (0, Some(1000), input.clone())
893        );
894
895        do_test!(
896            [&base_case[..], &cases[..]].concat(),
897            GenericStringArray<O>,
898            substring_by_char
899        );
900    }
901
902    #[test]
903    fn with_nulls_string_by_char() {
904        with_nulls_generic_string_by_char::<i32>()
905    }
906
907    #[test]
908    fn with_nulls_large_string_by_char() {
909        with_nulls_generic_string_by_char::<i64>()
910    }
911
912    fn without_nulls_generic_string_by_char<O: OffsetSizeTrait>() {
913        let input = vec!["hello", "", "Γ ⊢x:T"];
914        // empty array is always identical
915        let base_case = gen_test_cases!(vec!["", "", ""], (0, None, vec!["", "", ""]));
916        let cases = gen_test_cases!(
917            input,
918            //identity
919            (0, None, input.clone()),
920            // increase start
921            (1, None, vec!["ello", "", " ⊢x:T"]),
922            (2, None, vec!["llo", "", "⊢x:T"]),
923            (3, None, vec!["lo", "", "x:T"]),
924            (10, None, vec!["", "", ""]),
925            // increase start negatively
926            (-1, None, vec!["o", "", "T"]),
927            (-2, None, vec!["lo", "", ":T"]),
928            (-4, None, vec!["ello", "", "⊢x:T"]),
929            (-10, None, input.clone()),
930            // increase length
931            (1, Some(1), vec!["e", "", " "]),
932            (1, Some(2), vec!["el", "", " ⊢"]),
933            (1, Some(3), vec!["ell", "", " ⊢x"]),
934            (1, Some(6), vec!["ello", "", " ⊢x:T"]),
935            (-4, Some(1), vec!["e", "", "⊢"]),
936            (-4, Some(2), vec!["el", "", "⊢x"]),
937            (-4, Some(3), vec!["ell", "", "⊢x:"]),
938            (-4, Some(4), vec!["ello", "", "⊢x:T"])
939        );
940
941        do_test!(
942            [&base_case[..], &cases[..]].concat(),
943            GenericStringArray<O>,
944            substring_by_char
945        );
946    }
947
948    #[test]
949    fn without_nulls_string_by_char() {
950        without_nulls_generic_string_by_char::<i32>()
951    }
952
953    #[test]
954    fn without_nulls_large_string_by_char() {
955        without_nulls_generic_string_by_char::<i64>()
956    }
957
958    /// The array is all-ASCII, which takes the fast path of [`substring_by_char`].
959    fn ascii_generic_string_by_char<O: OffsetSizeTrait>() {
960        let input = vec![Some("hello"), None, Some(""), Some("rust")];
961        let cases = gen_test_cases!(
962            input,
963            // identity
964            (0, None, input.clone()),
965            // increase start
966            (1, None, vec![Some("ello"), None, Some(""), Some("ust")]),
967            (4, None, vec![Some("o"), None, Some(""), Some("")]),
968            // high start -> nothing
969            (1000, None, vec![Some(""), None, Some(""), Some("")]),
970            // increase start negatively
971            (-1, None, vec![Some("o"), None, Some(""), Some("t")]),
972            (-4, None, vec![Some("ello"), None, Some(""), Some("rust")]),
973            // high negative start -> identity
974            (-1000, None, input.clone()),
975            // 0 length -> nothing
976            (0, Some(0), vec![Some(""), None, Some(""), Some("")]),
977            // increase length
978            (1, Some(2), vec![Some("el"), None, Some(""), Some("us")]),
979            (-4, Some(2), vec![Some("el"), None, Some(""), Some("ru")]),
980            // high length -> identity
981            (0, Some(1000), input.clone()),
982            // length past the end is clamped, including when it would overflow
983            (
984                1,
985                Some(u64::MAX),
986                vec![Some("ello"), None, Some(""), Some("ust")]
987            )
988        );
989
990        do_test!(cases, GenericStringArray<O>, substring_by_char);
991    }
992
993    #[test]
994    fn ascii_string_by_char() {
995        ascii_generic_string_by_char::<i32>()
996    }
997
998    #[test]
999    fn ascii_large_string_by_char() {
1000        ascii_generic_string_by_char::<i64>()
1001    }
1002
1003    fn generic_string_by_char_with_non_zero_offset<O: OffsetSizeTrait>() {
1004        let values = "S→T = Πx:S.T";
1005        let offsets = &[
1006            O::zero(),
1007            O::from_usize(values.char_indices().nth(3).map(|(pos, _)| pos).unwrap()).unwrap(),
1008            O::from_usize(values.char_indices().nth(6).map(|(pos, _)| pos).unwrap()).unwrap(),
1009            O::from_usize(values.len()).unwrap(),
1010        ];
1011        // set the first and third element to be valid
1012        let bitmap = [0b101_u8];
1013
1014        let offsets = OffsetBuffer::new(Buffer::from_slice_ref(offsets).into());
1015        let values = Buffer::from(values.as_bytes());
1016        let nulls = Some(NullBuffer::new(BooleanBuffer::new(
1017            Buffer::from(bitmap),
1018            0,
1019            3,
1020        )));
1021        // array is `[null, "Πx:S.T"]`
1022        let array = GenericStringArray::<O>::new(offsets, values, nulls).slice(1, 2);
1023        // result is `[null, "x:S.T"]`
1024        let result = substring_by_char(&array, 1, None).unwrap();
1025        let expected = GenericStringArray::<O>::from(vec![None, Some("x:S.T")]);
1026        assert_eq!(result, expected);
1027    }
1028
1029    #[test]
1030    fn string_with_non_zero_offset_by_char() {
1031        generic_string_by_char_with_non_zero_offset::<i32>()
1032    }
1033
1034    #[test]
1035    fn large_string_with_non_zero_offset_by_char() {
1036        generic_string_by_char_with_non_zero_offset::<i64>()
1037    }
1038
1039    #[test]
1040    fn dictionary() {
1041        _dictionary::<Int8Type>();
1042        _dictionary::<Int16Type>();
1043        _dictionary::<Int32Type>();
1044        _dictionary::<Int64Type>();
1045        _dictionary::<UInt8Type>();
1046        _dictionary::<UInt16Type>();
1047        _dictionary::<UInt32Type>();
1048        _dictionary::<UInt64Type>();
1049    }
1050
1051    fn _dictionary<K: ArrowDictionaryKeyType>() {
1052        const TOTAL: i32 = 100;
1053
1054        let v = ["aaa", "bbb", "ccc", "ddd", "eee"];
1055        let data: Vec<Option<&str>> = (0..TOTAL)
1056            .map(|n| {
1057                let i = n % 5;
1058                if i == 3 { None } else { Some(v[i as usize]) }
1059            })
1060            .collect();
1061
1062        let dict_array: DictionaryArray<K> = data.clone().into_iter().collect();
1063
1064        let expected: Vec<Option<&str>> = data.iter().map(|opt| opt.map(|s| &s[1..3])).collect();
1065
1066        let res = substring(&dict_array, 1, Some(2)).unwrap();
1067        let actual = res.as_any().downcast_ref::<DictionaryArray<K>>().unwrap();
1068        let actual: Vec<Option<&str>> = actual
1069            .values()
1070            .as_any()
1071            .downcast_ref::<GenericStringArray<i32>>()
1072            .unwrap()
1073            .take_iter(actual.keys_iter())
1074            .collect();
1075
1076        for i in 0..TOTAL as usize {
1077            assert_eq!(expected[i], actual[i],);
1078        }
1079    }
1080
1081    #[test]
1082    fn check_invalid_array_type() {
1083        let array = Int32Array::from(vec![Some(1), Some(2), Some(3)]);
1084        let err = substring(&array, 0, None).unwrap_err().to_string();
1085        assert!(err.contains("substring does not support type"));
1086    }
1087
1088    // tests for the utf-8 validation checking
1089    #[test]
1090    fn check_start_index() {
1091        let array = StringArray::from(vec![Some("E=mc²"), Some("ascii")]);
1092        let err = substring(&array, -1, None).unwrap_err().to_string();
1093        assert!(err.contains("invalid utf-8 boundary"));
1094    }
1095
1096    #[test]
1097    fn check_length() {
1098        let array = StringArray::from(vec![Some("E=mc²"), Some("ascii")]);
1099        let err = substring(&array, 0, Some(5)).unwrap_err().to_string();
1100        assert!(err.contains("invalid utf-8 boundary"));
1101    }
1102
1103    #[test]
1104    fn non_utf8_bytes() {
1105        // non-utf8 bytes
1106        let bytes: &[u8] = &[0xE4, 0xBD, 0xA0, 0xE5, 0xA5, 0xBD, 0xE8, 0xAF, 0xAD];
1107        let array = BinaryArray::from(vec![Some(bytes)]);
1108        let arr = substring(&array, 0, Some(5)).unwrap();
1109        let actual = arr.as_any().downcast_ref::<BinaryArray>().unwrap();
1110
1111        let expected_bytes: &[u8] = &[0xE4, 0xBD, 0xA0, 0xE5, 0xA5];
1112        let expected = BinaryArray::from(vec![Some(expected_bytes)]);
1113        assert_eq!(expected, *actual);
1114    }
1115
1116    #[test]
1117    fn string_view_matches_utf8() {
1118        let values = vec![
1119            Some("hello world"),
1120            Some(""),
1121            None,
1122            Some("a"),
1123            Some("this one is definitely longer than twelve bytes"),
1124        ];
1125        let utf8 = StringArray::from(values.clone());
1126        let view = StringViewArray::from(values);
1127
1128        for (start, length) in [
1129            (0, None),
1130            (0, Some(0)),
1131            (0, Some(5)),
1132            (0, Some(1000)),
1133            (1, Some(3)),
1134            (5, None),
1135            (100, Some(2)),
1136            (100, None),
1137            (-3, None),
1138            (-3, Some(2)),
1139            (-100, Some(4)),
1140            (-100, None),
1141        ] {
1142            let expected = substring(&utf8, start, length).unwrap();
1143            let expected = expected.as_string::<i32>();
1144            let actual = substring(&view, start, length).unwrap();
1145            let actual = actual.as_string_view();
1146            assert_eq!(
1147                expected.iter().collect::<Vec<_>>(),
1148                actual.iter().collect::<Vec<_>>(),
1149                "start={start} length={length:?}"
1150            );
1151        }
1152    }
1153
1154    #[test]
1155    fn binary_view_matches_binary() {
1156        let values: Vec<Option<&[u8]>> = vec![
1157            Some(b"hello world"),
1158            Some(b""),
1159            None,
1160            Some(b"abc"),
1161            Some(b"this one is definitely longer than twelve bytes"),
1162        ];
1163        let binary = BinaryArray::from(values.clone());
1164        let view = BinaryViewArray::from(values);
1165
1166        for (start, length) in [
1167            (0, None),
1168            (0, Some(5)),
1169            (2, Some(3)),
1170            (-3, None),
1171            (100, Some(2)),
1172        ] {
1173            let expected = substring(&binary, start, length).unwrap();
1174            let expected = expected.as_binary::<i32>();
1175            let actual = substring(&view, start, length).unwrap();
1176            let actual = actual.as_binary_view();
1177            assert_eq!(
1178                expected.iter().collect::<Vec<_>>(),
1179                actual.iter().collect::<Vec<_>>(),
1180                "start={start} length={length:?}"
1181            );
1182        }
1183    }
1184
1185    #[test]
1186    fn string_view_slices_on_a_multi_byte_boundary() {
1187        // "é" and "ö" are two bytes, the CJK chars three, so 0/3/6 are the offsets that land
1188        // on a char boundary in both values.
1189        let values = vec![Some("héllo wörld"), Some("日本語"), None];
1190        let utf8 = StringArray::from(values.clone());
1191        let view = StringViewArray::from(values);
1192
1193        for (start, length) in [
1194            (0, Some(3)),
1195            (3, Some(3)),
1196            (0, Some(6)),
1197            (3, None),
1198            (-3, None),
1199            (-6, None),
1200        ] {
1201            let expected = substring(&utf8, start, length).unwrap();
1202            let expected = expected.as_string::<i32>();
1203            let actual = substring(&view, start, length).unwrap();
1204            let actual = actual.as_string_view();
1205            assert_eq!(
1206                expected.iter().collect::<Vec<_>>(),
1207                actual.iter().collect::<Vec<_>>(),
1208                "start={start} length={length:?}"
1209            );
1210        }
1211
1212        let actual = substring(&view, 0, Some(3)).unwrap();
1213        let actual = actual.as_string_view();
1214        assert_eq!(actual.value(0), "hé");
1215        assert_eq!(actual.value(1), "日");
1216    }
1217
1218    #[test]
1219    fn string_view_rejects_an_invalid_char_boundary() {
1220        let view = StringViewArray::from(vec![Some("E=mc²")]);
1221        let err = substring(&view, 0, Some(5)).unwrap_err().to_string();
1222        assert!(err.contains("invalid utf-8 boundary"), "{err}");
1223    }
1224
1225    #[test]
1226    fn dictionary_of_string_view() {
1227        let view = StringViewArray::from(vec![Some("hello world"), Some("bye")]);
1228        let dict = DictionaryArray::new(Int32Array::from(vec![0, 1, 0]), Arc::new(view));
1229
1230        let actual = substring(&dict, 0, Some(3)).unwrap();
1231        let actual = actual.as_any_dictionary();
1232
1233        let values = actual.values().as_string_view();
1234        assert_eq!(
1235            values.iter().collect::<Vec<_>>(),
1236            vec![Some("hel"), Some("bye")]
1237        );
1238    }
1239}