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