1use 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
32pub 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
112pub 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 Ok(substring_by_char_impl(array, length, |val| {
151 ascii_bounds(val, start, length)
152 }))
153 } else {
154 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
162fn 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#[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#[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 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 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 let check_char_boundary = {
264 |offset: T::Offset| {
265 if !matches!(T::DATA_TYPE, DataType::Utf8 | DataType::LargeUtf8) {
266 return Ok(offset);
267 }
268 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 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 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 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 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 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 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 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 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 (0, None, input.clone()),
460 (0, Some(0), vec![Some(&[]), None, Some(&[])]),
462 (1000, Some(0), vec![Some(&[]), None, Some(&[])]),
464 (-1000, None, input.clone()),
466 (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 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 (0, None, input.clone()),
498 (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 (-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 (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 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 let array = GenericBinaryArray::<O>::new(offsets, values, nulls).slice(1, 2);
556 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 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 (0, None, input.clone()),
587 (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 (-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 (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 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 (0, None, input.clone()),
626 (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 (-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 (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 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 let array = FixedSizeBinaryArray::new(5, Buffer::from(values), nulls).slice(1, 2);
666 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 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 (0, None, input.clone()),
688 (0, Some(0), vec![Some(""), None, Some("")]),
690 (1000, Some(0), vec![Some(""), None, Some("")]),
692 (-1000, None, input.clone()),
694 (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 let base_case = gen_test_cases!(vec!["", "", ""], (0, None, vec!["", "", ""]));
719 let cases = gen_test_cases!(
720 input,
721 (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 (-1, None, vec!["o", "", "d"]),
729 (-2, None, vec!["lo", "", "rd"]),
730 (-3, None, vec!["llo", "", "ord"]),
731 (-10, None, input.clone()),
732 (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 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 let array = GenericStringArray::<O>::new(offsets, values, nulls).slice(1, 2);
780 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 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 (0, None, input.clone()),
808 (0, Some(0), vec![Some(""), None, Some("")]),
810 (1000, Some(0), vec![Some(""), None, Some("")]),
812 (-1000, None, input.clone()),
814 (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 let base_case = gen_test_cases!(vec!["", "", ""], (0, None, vec!["", "", ""]));
839 let cases = gen_test_cases!(
840 input,
841 (0, None, input.clone()),
843 (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 (-1, None, vec!["o", "", "T"]),
850 (-2, None, vec!["lo", "", ":T"]),
851 (-4, None, vec!["ello", "", "⊢x:T"]),
852 (-10, None, input.clone()),
853 (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 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 (0, None, input.clone()),
888 (1, None, vec![Some("ello"), None, Some(""), Some("ust")]),
890 (4, None, vec![Some("o"), None, Some(""), Some("")]),
891 (1000, None, vec![Some(""), None, Some(""), Some("")]),
893 (-1, None, vec![Some("o"), None, Some(""), Some("t")]),
895 (-4, None, vec![Some("ello"), None, Some(""), Some("rust")]),
896 (-1000, None, input.clone()),
898 (0, Some(0), vec![Some(""), None, Some(""), Some("")]),
900 (1, Some(2), vec![Some("el"), None, Some(""), Some("us")]),
902 (-4, Some(2), vec![Some("el"), None, Some(""), Some("ru")]),
903 (0, Some(1000), input.clone()),
905 (
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 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 let array = GenericStringArray::<O>::new(offsets, values, nulls).slice(1, 2);
946 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 #[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 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}