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. It is recommended to
130/// use [substring] if the input array only contains ASCII chars.
131///
132/// # Basic usage
133/// ```
134/// # use arrow_array::StringArray;
135/// # use arrow_string::substring::substring_by_char;
136/// let array = StringArray::from(vec![Some("arrow"), None, Some("Γ ⊢x:T")]);
137/// let result = substring_by_char(&array, 1, Some(4)).unwrap();
138/// assert_eq!(result, StringArray::from(vec![Some("rrow"), None, Some(" ⊢x:")]));
139/// ```
140pub fn substring_by_char<OffsetSize: OffsetSizeTrait>(
141    array: &GenericStringArray<OffsetSize>,
142    start: i64,
143    length: Option<u64>,
144) -> Result<GenericStringArray<OffsetSize>, ArrowError> {
145    let mut vals = BufferBuilder::<u8>::new({
146        let offsets = array.value_offsets();
147        (offsets[array.len()] - offsets[0]).to_usize().unwrap()
148    });
149    let mut new_offsets = BufferBuilder::<OffsetSize>::new(array.len() + 1);
150    new_offsets.append(OffsetSize::zero());
151    let length = length.map(|len| len.to_usize().unwrap());
152
153    array.iter().for_each(|val| {
154        if let Some(val) = val {
155            let char_count = val.chars().count();
156            let start = if start >= 0 {
157                start.to_usize().unwrap()
158            } else {
159                char_count - (-start).to_usize().unwrap().min(char_count)
160            };
161            let (start_offset, end_offset) = get_start_end_offset(val, start, length);
162            vals.append_slice(&val.as_bytes()[start_offset..end_offset]);
163        }
164        new_offsets.append(OffsetSize::from_usize(vals.len()).unwrap());
165    });
166    let offsets = OffsetBuffer::new(new_offsets.finish().into());
167    let values = vals.finish();
168    let nulls = array
169        .nulls()
170        .map(|n| n.inner().sliced())
171        .and_then(|b| NullBuffer::from_unsliced_buffer(b, array.len()));
172    Ok(GenericStringArray::<OffsetSize>::new(
173        offsets, values, nulls,
174    ))
175}
176
177/// * `val` - string
178/// * `start` - the start char index of the substring
179/// * `length` - the char length of the substring
180///
181/// Return the `start` and `end` offset (by byte) of the substring
182fn get_start_end_offset(val: &str, start: usize, length: Option<usize>) -> (usize, usize) {
183    let len = val.len();
184    let mut offset_char_iter = val.char_indices();
185    let start_offset = offset_char_iter
186        .nth(start)
187        .map_or(len, |(offset, _)| offset);
188    let end_offset = length.map_or(len, |length| {
189        if length > 0 {
190            offset_char_iter
191                .nth(length - 1)
192                .map_or(len, |(offset, _)| offset)
193        } else {
194            start_offset
195        }
196    });
197    (start_offset, end_offset)
198}
199
200fn byte_substring<T: ByteArrayType>(
201    array: &GenericByteArray<T>,
202    start: T::Offset,
203    length: Option<T::Offset>,
204) -> Result<ArrayRef, ArrowError>
205where
206    <T as ByteArrayType>::Native: PartialEq,
207{
208    let offsets = array.value_offsets();
209    let data = array.value_data();
210    let zero = <T::Offset as Zero>::zero();
211
212    // When array is [Large]StringArray, we will check whether `offset` is at a valid char boundary.
213    let check_char_boundary = {
214        |offset: T::Offset| {
215            if !matches!(T::DATA_TYPE, DataType::Utf8 | DataType::LargeUtf8) {
216                return Ok(offset);
217            }
218            // Safety: a StringArray must contain valid UTF8 data
219            let data_str = unsafe { std::str::from_utf8_unchecked(data) };
220            let offset_usize = offset.as_usize();
221            if data_str.is_char_boundary(offset_usize) {
222                Ok(offset)
223            } else {
224                Err(ArrowError::ComputeError(format!(
225                    "The offset {offset_usize} is at an invalid utf-8 boundary."
226                )))
227            }
228        }
229    };
230
231    // start and end offsets of all substrings
232    let mut new_starts_ends: Vec<(T::Offset, T::Offset)> = Vec::with_capacity(array.len());
233    let mut new_offsets: Vec<T::Offset> = Vec::with_capacity(array.len() + 1);
234    let mut len_so_far = zero;
235    new_offsets.push(zero);
236
237    offsets
238        .windows(2)
239        .try_for_each(|pair| -> Result<(), ArrowError> {
240            let new_start = match start.cmp(&zero) {
241                Ordering::Greater => check_char_boundary((pair[0] + start).min(pair[1]))?,
242                Ordering::Equal => pair[0],
243                Ordering::Less => check_char_boundary((pair[1] + start).max(pair[0]))?,
244            };
245            let new_end = match length {
246                Some(length) => check_char_boundary((length + new_start).min(pair[1]))?,
247                None => pair[1],
248            };
249            len_so_far += new_end - new_start;
250            new_starts_ends.push((new_start, new_end));
251            new_offsets.push(len_so_far);
252            Ok(())
253        })?;
254
255    // concatenate substrings into a buffer
256    let mut new_values = MutableBuffer::new(new_offsets.last().unwrap().as_usize());
257
258    new_starts_ends
259        .iter()
260        .map(|(start, end)| {
261            let start = start.as_usize();
262            let end = end.as_usize();
263            &data[start..end]
264        })
265        .for_each(|slice| new_values.extend_from_slice(slice));
266
267    let offsets = OffsetBuffer::new(new_offsets.into());
268    let values = new_values.into();
269    let nulls = array
270        .nulls()
271        .map(|n| n.inner().sliced())
272        .and_then(|b| NullBuffer::from_unsliced_buffer(b, array.len()));
273    Ok(Arc::new(GenericByteArray::<T>::new(offsets, values, nulls)))
274}
275
276fn fixed_size_binary_substring(
277    array: &FixedSizeBinaryArray,
278    old_len: usize,
279    start: i64,
280    length: Option<u64>,
281) -> Result<ArrayRef, ArrowError> {
282    let new_start = match start.cmp(&0) {
283        Ordering::Greater => usize::try_from(start).unwrap_or(usize::MAX).min(old_len),
284        Ordering::Equal => 0,
285        Ordering::Less => {
286            let offset = usize::try_from(start.unsigned_abs()).unwrap_or(usize::MAX);
287            old_len.saturating_sub(offset)
288        }
289    };
290
291    let new_len = match length {
292        Some(len) => usize::try_from(len)
293            .unwrap_or(usize::MAX)
294            .min(old_len - new_start),
295        None => old_len - new_start,
296    };
297
298    // build value buffer
299    let num_of_elements = array.len();
300    let data = array.value_data();
301    let capacity = num_of_elements
302        .checked_mul(new_len)
303        .expect("capacity overflow");
304    let mut new_values = MutableBuffer::new(capacity);
305    (0..num_of_elements)
306        .map(|idx| {
307            let offset = idx * array.value_size();
308            (offset + new_start, offset + new_start + new_len)
309        })
310        .for_each(|(start, end)| new_values.extend_from_slice(&data[start..end]));
311
312    let mut nulls = array
313        .nulls()
314        .map(|n| n.inner().sliced())
315        .and_then(|b| NullBuffer::from_unsliced_buffer(b, num_of_elements));
316
317    if new_len == 0 && nulls.is_none() {
318        // FixedSizeBinaryArray::new takes length from the values buffer, except when size == 0.
319        // In that case it uses the null buffer length, so preserve the original length here.
320        // Example: ["", "", ""] -> substring(..., 1, Some(2)) should keep len=3;
321        // otherwise it collapses to an empty array (len=0).
322        nulls = Some(NullBuffer::new_valid(num_of_elements));
323    }
324
325    let new_len: i32 = new_len.try_into().expect("new_len overflow");
326
327    Ok(Arc::new(FixedSizeBinaryArray::new(
328        new_len,
329        new_values.into(),
330        nulls,
331    )))
332}
333
334#[cfg(test)]
335mod tests {
336    use super::*;
337    use arrow_buffer::BooleanBuffer;
338    use arrow_buffer::Buffer;
339
340    /// A helper macro to generate test cases.
341    /// # Arguments
342    /// * `input` - A vector which array can be built from.
343    /// * `start` - The start index of the substring.
344    /// * `len` - The length of the substring.
345    /// * `result` - The expected result of substring, which is a vector that array can be built from.
346    /// # Return
347    /// A vector of `(input, start, len, result)`.
348    ///
349    /// Users can provide any number of `(start, len, result)` to generate test cases for one `input`.
350    macro_rules! gen_test_cases {
351        ($input:expr, $(($start:expr, $len:expr, $result:expr)), *) => {
352            [
353                $(
354                    ($input.clone(), $start, $len, $result),
355                )*
356            ]
357        };
358    }
359
360    /// A helper macro to test the substring functions.
361    /// # Arguments
362    /// * `cases` - The test cases which is a vector of `(input, start, len, result)`.
363    ///   Please look at [`gen_test_cases`] to find how to generate it.
364    /// * `array_ty` - The array type.
365    /// * `substring_fn` - Either [`substring`] or [`substring_by_char`].
366    macro_rules! do_test {
367        ($cases:expr, $array_ty:ty, $substring_fn:ident) => {
368            $cases
369                .into_iter()
370                .for_each(|(array, start, length, expected)| {
371                    let array = <$array_ty>::from(array);
372                    let result = $substring_fn(&array, start, length).unwrap();
373                    let result = result.as_any().downcast_ref::<$array_ty>().unwrap();
374                    let expected = <$array_ty>::from(expected);
375                    assert_eq!(&expected, result);
376                })
377        };
378    }
379
380    /// A helper macro to test the substring functions for array types only implementing TryFrom.
381    macro_rules! do_test_tryfrom {
382        ($cases:expr, $array_ty:ty, $substring_fn:ident) => {
383            $cases
384                .into_iter()
385                .for_each(|(array, start, length, expected)| {
386                    let array = <$array_ty>::try_from(array).unwrap();
387                    let result = $substring_fn(&array, start, length).unwrap();
388                    let result = result.as_any().downcast_ref::<$array_ty>().unwrap();
389                    let expected = <$array_ty>::try_from(expected).unwrap();
390                    assert_eq!(&expected, result);
391                })
392        };
393    }
394
395    fn with_nulls_generic_binary<O: OffsetSizeTrait>() {
396        let input = vec![
397            Some("hello".as_bytes()),
398            None,
399            Some(&[0xf8, 0xf9, 0xff, 0xfa]),
400        ];
401        // all-nulls array is always identical
402        let base_case = gen_test_cases!(
403            vec![None, None, None],
404            (-1, Some(1), vec![None, None, None])
405        );
406        let cases = gen_test_cases!(
407            input,
408            // identity
409            (0, None, input.clone()),
410            // 0 length -> Nothing
411            (0, Some(0), vec![Some(&[]), None, Some(&[])]),
412            // high start -> Nothing
413            (1000, Some(0), vec![Some(&[]), None, Some(&[])]),
414            // high negative start -> identity
415            (-1000, None, input.clone()),
416            // high length -> identity
417            (0, Some(1000), input.clone())
418        );
419
420        do_test!(
421            [&base_case[..], &cases[..]].concat(),
422            GenericBinaryArray<O>,
423            substring
424        );
425    }
426
427    #[test]
428    fn with_nulls_binary() {
429        with_nulls_generic_binary::<i32>()
430    }
431
432    #[test]
433    fn with_nulls_large_binary() {
434        with_nulls_generic_binary::<i64>()
435    }
436
437    fn without_nulls_generic_binary<O: OffsetSizeTrait>() {
438        let input = vec!["hello".as_bytes(), b"", &[0xf8, 0xf9, 0xff, 0xfa]];
439        // empty array is always identical
440        let base_case = gen_test_cases!(
441            vec!["".as_bytes(), b"", b""],
442            (2, Some(1), vec!["".as_bytes(), b"", b""])
443        );
444        let cases = gen_test_cases!(
445            input,
446            // identity
447            (0, None, input.clone()),
448            // increase start
449            (1, None, vec![b"ello", b"", &[0xf9, 0xff, 0xfa]]),
450            (2, None, vec![b"llo", b"", &[0xff, 0xfa]]),
451            (3, None, vec![b"lo", b"", &[0xfa]]),
452            (10, None, vec![b"", b"", b""]),
453            // increase start negatively
454            (-1, None, vec![b"o", b"", &[0xfa]]),
455            (-2, None, vec![b"lo", b"", &[0xff, 0xfa]]),
456            (-3, None, vec![b"llo", b"", &[0xf9, 0xff, 0xfa]]),
457            (-10, None, input.clone()),
458            // increase length
459            (1, Some(1), vec![b"e", b"", &[0xf9]]),
460            (1, Some(2), vec![b"el", b"", &[0xf9, 0xff]]),
461            (1, Some(3), vec![b"ell", b"", &[0xf9, 0xff, 0xfa]]),
462            (1, Some(4), vec![b"ello", b"", &[0xf9, 0xff, 0xfa]]),
463            (-3, Some(1), vec![b"l", b"", &[0xf9]]),
464            (-3, Some(2), vec![b"ll", b"", &[0xf9, 0xff]]),
465            (-3, Some(3), vec![b"llo", b"", &[0xf9, 0xff, 0xfa]]),
466            (-3, Some(4), vec![b"llo", b"", &[0xf9, 0xff, 0xfa]])
467        );
468
469        do_test!(
470            [&base_case[..], &cases[..]].concat(),
471            GenericBinaryArray<O>,
472            substring
473        );
474    }
475
476    #[test]
477    fn without_nulls_binary() {
478        without_nulls_generic_binary::<i32>()
479    }
480
481    #[test]
482    fn without_nulls_large_binary() {
483        without_nulls_generic_binary::<i64>()
484    }
485
486    fn generic_binary_with_non_zero_offset<O: OffsetSizeTrait>() {
487        let values = 0_u8..15;
488        let offsets = &[
489            O::zero(),
490            O::from_usize(5).unwrap(),
491            O::from_usize(10).unwrap(),
492            O::from_usize(15).unwrap(),
493        ];
494        // set the first and third element to be valid
495        let bitmap = [0b101_u8];
496
497        let offsets = OffsetBuffer::new(Buffer::from_slice_ref(offsets).into());
498        let values = Buffer::from_iter(values);
499        let nulls = Some(NullBuffer::new(BooleanBuffer::new(
500            Buffer::from(bitmap),
501            0,
502            3,
503        )));
504        // array is `[null, [10, 11, 12, 13, 14]]`
505        let array = GenericBinaryArray::<O>::new(offsets, values, nulls).slice(1, 2);
506        // result is `[null, [11, 12, 13, 14]]`
507        let result = substring(&array, 1, None).unwrap();
508        let result = result
509            .as_any()
510            .downcast_ref::<GenericBinaryArray<O>>()
511            .unwrap();
512        let expected =
513            GenericBinaryArray::<O>::from_opt_vec(vec![None, Some(&[11_u8, 12, 13, 14])]);
514        assert_eq!(result, &expected);
515    }
516
517    #[test]
518    fn binary_with_non_zero_offset() {
519        generic_binary_with_non_zero_offset::<i32>()
520    }
521
522    #[test]
523    fn large_binary_with_non_zero_offset() {
524        generic_binary_with_non_zero_offset::<i64>()
525    }
526
527    #[test]
528    fn with_nulls_fixed_size_binary() {
529        let input = vec![Some("cat".as_bytes()), None, Some(&[0xf8, 0xf9, 0xff])];
530        // all-nulls array is always identical
531        let base_case =
532            gen_test_cases!(vec![None, None, None], (3, Some(2), vec![None, None, None]));
533        let cases = gen_test_cases!(
534            input,
535            // identity
536            (0, None, input.clone()),
537            // increase start
538            (1, None, vec![Some(b"at"), None, Some(&[0xf9, 0xff])]),
539            (2, None, vec![Some(b"t"), None, Some(&[0xff])]),
540            (3, None, vec![Some(b""), None, Some(b"")]),
541            (10, None, vec![Some(b""), None, Some(b"")]),
542            // increase start negatively
543            (-1, None, vec![Some(b"t"), None, Some(&[0xff])]),
544            (-2, None, vec![Some(b"at"), None, Some(&[0xf9, 0xff])]),
545            (-3, None, input.clone()),
546            (-10, None, input.clone()),
547            // increase length
548            (1, Some(1), vec![Some(b"a"), None, Some(&[0xf9])]),
549            (1, Some(2), vec![Some(b"at"), None, Some(&[0xf9, 0xff])]),
550            (1, Some(3), vec![Some(b"at"), None, Some(&[0xf9, 0xff])]),
551            (-3, Some(1), vec![Some(b"c"), None, Some(&[0xf8])]),
552            (-3, Some(2), vec![Some(b"ca"), None, Some(&[0xf8, 0xf9])]),
553            (-3, Some(3), input.clone()),
554            (-3, Some(4), input.clone())
555        );
556
557        do_test_tryfrom!(
558            [&base_case[..], &cases[..]].concat(),
559            FixedSizeBinaryArray,
560            substring
561        );
562    }
563
564    #[test]
565    fn without_nulls_fixed_size_binary() {
566        let input = vec!["cat".as_bytes(), b"dog", &[0xf8, 0xf9, 0xff]];
567        // empty array is always identical
568        let base_case = gen_test_cases!(
569            vec!["".as_bytes(), &[], &[]],
570            (1, Some(2), vec!["".as_bytes(), &[], &[]])
571        );
572        let cases = gen_test_cases!(
573            input,
574            // identity
575            (0, None, input.clone()),
576            // increase start
577            (1, None, vec![b"at", b"og", &[0xf9, 0xff]]),
578            (2, None, vec![b"t", b"g", &[0xff]]),
579            (3, None, vec![&[], &[], &[]]),
580            (10, None, vec![&[], &[], &[]]),
581            // increase start negatively
582            (-1, None, vec![b"t", b"g", &[0xff]]),
583            (-2, None, vec![b"at", b"og", &[0xf9, 0xff]]),
584            (-3, None, input.clone()),
585            (-10, None, input.clone()),
586            // increase length
587            (1, Some(1), vec![b"a", b"o", &[0xf9]]),
588            (1, Some(2), vec![b"at", b"og", &[0xf9, 0xff]]),
589            (1, Some(3), vec![b"at", b"og", &[0xf9, 0xff]]),
590            (-3, Some(1), vec![b"c", b"d", &[0xf8]]),
591            (-3, Some(2), vec![b"ca", b"do", &[0xf8, 0xf9]]),
592            (-3, Some(3), input.clone()),
593            (-3, Some(4), input.clone())
594        );
595
596        do_test_tryfrom!(
597            [&base_case[..], &cases[..]].concat(),
598            FixedSizeBinaryArray,
599            substring
600        );
601    }
602
603    #[test]
604    fn fixed_size_binary_with_non_zero_offset() {
605        let values = b"hellotherearrow";
606        // set the first and third element to be valid
607        let bits_v = [0b101_u8];
608
609        let nulls = Some(NullBuffer::new(BooleanBuffer::new(
610            Buffer::from(bits_v),
611            0,
612            3,
613        )));
614        // array is `[null, "arrow"]`
615        let array = FixedSizeBinaryArray::new(5, Buffer::from(values), nulls).slice(1, 2);
616        // result is `[null, "rrow"]`
617        let result = substring(&array, 1, None).unwrap();
618        let result = result
619            .as_any()
620            .downcast_ref::<FixedSizeBinaryArray>()
621            .unwrap();
622        let expected = FixedSizeBinaryArray::try_from_sparse_iter_with_size(
623            vec![None, Some(b"rrow")].into_iter(),
624            4,
625        )
626        .unwrap();
627        assert_eq!(result, &expected);
628    }
629
630    fn with_nulls_generic_string<O: OffsetSizeTrait>() {
631        let input = vec![Some("hello"), None, Some("word")];
632        // all-nulls array is always identical
633        let base_case = gen_test_cases!(vec![None, None, None], (0, None, vec![None, None, None]));
634        let cases = gen_test_cases!(
635            input,
636            // identity
637            (0, None, input.clone()),
638            // 0 length -> Nothing
639            (0, Some(0), vec![Some(""), None, Some("")]),
640            // high start -> Nothing
641            (1000, Some(0), vec![Some(""), None, Some("")]),
642            // high negative start -> identity
643            (-1000, None, input.clone()),
644            // high length -> identity
645            (0, Some(1000), input.clone())
646        );
647
648        do_test!(
649            [&base_case[..], &cases[..]].concat(),
650            GenericStringArray<O>,
651            substring
652        );
653    }
654
655    #[test]
656    fn with_nulls_string() {
657        with_nulls_generic_string::<i32>()
658    }
659
660    #[test]
661    fn with_nulls_large_string() {
662        with_nulls_generic_string::<i64>()
663    }
664
665    fn without_nulls_generic_string<O: OffsetSizeTrait>() {
666        let input = vec!["hello", "", "word"];
667        // empty array is always identical
668        let base_case = gen_test_cases!(vec!["", "", ""], (0, None, vec!["", "", ""]));
669        let cases = gen_test_cases!(
670            input,
671            // identity
672            (0, None, input.clone()),
673            (1, None, vec!["ello", "", "ord"]),
674            (2, None, vec!["llo", "", "rd"]),
675            (3, None, vec!["lo", "", "d"]),
676            (10, None, vec!["", "", ""]),
677            // increase start negatively
678            (-1, None, vec!["o", "", "d"]),
679            (-2, None, vec!["lo", "", "rd"]),
680            (-3, None, vec!["llo", "", "ord"]),
681            (-10, None, input.clone()),
682            // increase length
683            (1, Some(1), vec!["e", "", "o"]),
684            (1, Some(2), vec!["el", "", "or"]),
685            (1, Some(3), vec!["ell", "", "ord"]),
686            (1, Some(4), vec!["ello", "", "ord"]),
687            (-3, Some(1), vec!["l", "", "o"]),
688            (-3, Some(2), vec!["ll", "", "or"]),
689            (-3, Some(3), vec!["llo", "", "ord"]),
690            (-3, Some(4), vec!["llo", "", "ord"])
691        );
692
693        do_test!(
694            [&base_case[..], &cases[..]].concat(),
695            GenericStringArray<O>,
696            substring
697        );
698    }
699
700    #[test]
701    fn without_nulls_string() {
702        without_nulls_generic_string::<i32>()
703    }
704
705    #[test]
706    fn without_nulls_large_string() {
707        without_nulls_generic_string::<i64>()
708    }
709
710    fn generic_string_with_non_zero_offset<O: OffsetSizeTrait>() {
711        let values = b"hellotherearrow";
712        let offsets = &[
713            O::zero(),
714            O::from_usize(5).unwrap(),
715            O::from_usize(10).unwrap(),
716            O::from_usize(15).unwrap(),
717        ];
718        // set the first and third element to be valid
719        let bitmap = [0b101_u8];
720
721        let offsets = OffsetBuffer::new(Buffer::from_slice_ref(offsets).into());
722        let values = Buffer::from(values);
723        let nulls = Some(NullBuffer::new(BooleanBuffer::new(
724            Buffer::from(bitmap),
725            0,
726            3,
727        )));
728        // array is `[null, "arrow"]`
729        let array = GenericStringArray::<O>::new(offsets, values, nulls).slice(1, 2);
730        // result is `[null, "rrow"]`
731        let result = substring(&array, 1, None).unwrap();
732        let result = result
733            .as_any()
734            .downcast_ref::<GenericStringArray<O>>()
735            .unwrap();
736        let expected = GenericStringArray::<O>::from(vec![None, Some("rrow")]);
737        assert_eq!(result, &expected);
738    }
739
740    #[test]
741    fn string_with_non_zero_offset() {
742        generic_string_with_non_zero_offset::<i32>()
743    }
744
745    #[test]
746    fn large_string_with_non_zero_offset() {
747        generic_string_with_non_zero_offset::<i64>()
748    }
749
750    fn with_nulls_generic_string_by_char<O: OffsetSizeTrait>() {
751        let input = vec![Some("hello"), None, Some("Γ ⊢x:T")];
752        // all-nulls array is always identical
753        let base_case = gen_test_cases!(vec![None, None, None], (0, None, vec![None, None, None]));
754        let cases = gen_test_cases!(
755            input,
756            // identity
757            (0, None, input.clone()),
758            // 0 length -> Nothing
759            (0, Some(0), vec![Some(""), None, Some("")]),
760            // high start -> Nothing
761            (1000, Some(0), vec![Some(""), None, Some("")]),
762            // high negative start -> identity
763            (-1000, None, input.clone()),
764            // high length -> identity
765            (0, Some(1000), input.clone())
766        );
767
768        do_test!(
769            [&base_case[..], &cases[..]].concat(),
770            GenericStringArray<O>,
771            substring_by_char
772        );
773    }
774
775    #[test]
776    fn with_nulls_string_by_char() {
777        with_nulls_generic_string_by_char::<i32>()
778    }
779
780    #[test]
781    fn with_nulls_large_string_by_char() {
782        with_nulls_generic_string_by_char::<i64>()
783    }
784
785    fn without_nulls_generic_string_by_char<O: OffsetSizeTrait>() {
786        let input = vec!["hello", "", "Γ ⊢x:T"];
787        // empty array is always identical
788        let base_case = gen_test_cases!(vec!["", "", ""], (0, None, vec!["", "", ""]));
789        let cases = gen_test_cases!(
790            input,
791            //identity
792            (0, None, input.clone()),
793            // increase start
794            (1, None, vec!["ello", "", " ⊢x:T"]),
795            (2, None, vec!["llo", "", "⊢x:T"]),
796            (3, None, vec!["lo", "", "x:T"]),
797            (10, None, vec!["", "", ""]),
798            // increase start negatively
799            (-1, None, vec!["o", "", "T"]),
800            (-2, None, vec!["lo", "", ":T"]),
801            (-4, None, vec!["ello", "", "⊢x:T"]),
802            (-10, None, input.clone()),
803            // increase length
804            (1, Some(1), vec!["e", "", " "]),
805            (1, Some(2), vec!["el", "", " ⊢"]),
806            (1, Some(3), vec!["ell", "", " ⊢x"]),
807            (1, Some(6), vec!["ello", "", " ⊢x:T"]),
808            (-4, Some(1), vec!["e", "", "⊢"]),
809            (-4, Some(2), vec!["el", "", "⊢x"]),
810            (-4, Some(3), vec!["ell", "", "⊢x:"]),
811            (-4, Some(4), vec!["ello", "", "⊢x:T"])
812        );
813
814        do_test!(
815            [&base_case[..], &cases[..]].concat(),
816            GenericStringArray<O>,
817            substring_by_char
818        );
819    }
820
821    #[test]
822    fn without_nulls_string_by_char() {
823        without_nulls_generic_string_by_char::<i32>()
824    }
825
826    #[test]
827    fn without_nulls_large_string_by_char() {
828        without_nulls_generic_string_by_char::<i64>()
829    }
830
831    fn generic_string_by_char_with_non_zero_offset<O: OffsetSizeTrait>() {
832        let values = "S→T = Πx:S.T";
833        let offsets = &[
834            O::zero(),
835            O::from_usize(values.char_indices().nth(3).map(|(pos, _)| pos).unwrap()).unwrap(),
836            O::from_usize(values.char_indices().nth(6).map(|(pos, _)| pos).unwrap()).unwrap(),
837            O::from_usize(values.len()).unwrap(),
838        ];
839        // set the first and third element to be valid
840        let bitmap = [0b101_u8];
841
842        let offsets = OffsetBuffer::new(Buffer::from_slice_ref(offsets).into());
843        let values = Buffer::from(values.as_bytes());
844        let nulls = Some(NullBuffer::new(BooleanBuffer::new(
845            Buffer::from(bitmap),
846            0,
847            3,
848        )));
849        // array is `[null, "Πx:S.T"]`
850        let array = GenericStringArray::<O>::new(offsets, values, nulls).slice(1, 2);
851        // result is `[null, "x:S.T"]`
852        let result = substring_by_char(&array, 1, None).unwrap();
853        let expected = GenericStringArray::<O>::from(vec![None, Some("x:S.T")]);
854        assert_eq!(result, expected);
855    }
856
857    #[test]
858    fn string_with_non_zero_offset_by_char() {
859        generic_string_by_char_with_non_zero_offset::<i32>()
860    }
861
862    #[test]
863    fn large_string_with_non_zero_offset_by_char() {
864        generic_string_by_char_with_non_zero_offset::<i64>()
865    }
866
867    #[test]
868    fn dictionary() {
869        _dictionary::<Int8Type>();
870        _dictionary::<Int16Type>();
871        _dictionary::<Int32Type>();
872        _dictionary::<Int64Type>();
873        _dictionary::<UInt8Type>();
874        _dictionary::<UInt16Type>();
875        _dictionary::<UInt32Type>();
876        _dictionary::<UInt64Type>();
877    }
878
879    fn _dictionary<K: ArrowDictionaryKeyType>() {
880        const TOTAL: i32 = 100;
881
882        let v = ["aaa", "bbb", "ccc", "ddd", "eee"];
883        let data: Vec<Option<&str>> = (0..TOTAL)
884            .map(|n| {
885                let i = n % 5;
886                if i == 3 { None } else { Some(v[i as usize]) }
887            })
888            .collect();
889
890        let dict_array: DictionaryArray<K> = data.clone().into_iter().collect();
891
892        let expected: Vec<Option<&str>> = data.iter().map(|opt| opt.map(|s| &s[1..3])).collect();
893
894        let res = substring(&dict_array, 1, Some(2)).unwrap();
895        let actual = res.as_any().downcast_ref::<DictionaryArray<K>>().unwrap();
896        let actual: Vec<Option<&str>> = actual
897            .values()
898            .as_any()
899            .downcast_ref::<GenericStringArray<i32>>()
900            .unwrap()
901            .take_iter(actual.keys_iter())
902            .collect();
903
904        for i in 0..TOTAL as usize {
905            assert_eq!(expected[i], actual[i],);
906        }
907    }
908
909    #[test]
910    fn check_invalid_array_type() {
911        let array = Int32Array::from(vec![Some(1), Some(2), Some(3)]);
912        let err = substring(&array, 0, None).unwrap_err().to_string();
913        assert!(err.contains("substring does not support type"));
914    }
915
916    // tests for the utf-8 validation checking
917    #[test]
918    fn check_start_index() {
919        let array = StringArray::from(vec![Some("E=mc²"), Some("ascii")]);
920        let err = substring(&array, -1, None).unwrap_err().to_string();
921        assert!(err.contains("invalid utf-8 boundary"));
922    }
923
924    #[test]
925    fn check_length() {
926        let array = StringArray::from(vec![Some("E=mc²"), Some("ascii")]);
927        let err = substring(&array, 0, Some(5)).unwrap_err().to_string();
928        assert!(err.contains("invalid utf-8 boundary"));
929    }
930
931    #[test]
932    fn non_utf8_bytes() {
933        // non-utf8 bytes
934        let bytes: &[u8] = &[0xE4, 0xBD, 0xA0, 0xE5, 0xA5, 0xBD, 0xE8, 0xAF, 0xAD];
935        let array = BinaryArray::from(vec![Some(bytes)]);
936        let arr = substring(&array, 0, Some(5)).unwrap();
937        let actual = arr.as_any().downcast_ref::<BinaryArray>().unwrap();
938
939        let expected_bytes: &[u8] = &[0xE4, 0xBD, 0xA0, 0xE5, 0xA5];
940        let expected = BinaryArray::from(vec![Some(expected_bytes)]);
941        assert_eq!(expected, *actual);
942    }
943}
arrow_string/substring.rs

arrow_string/
substring.rs
