1use arrow_array::builder::{BinaryViewBuilder, StringViewBuilder};
24use arrow_array::cast::AsArray;
25use arrow_array::types::*;
26use arrow_array::*;
27use arrow_buffer::{ArrowNativeType, MutableBuffer, NullBuffer, OffsetBuffer};
28use arrow_schema::{ArrowError, DataType};
29use num_traits::{CheckedAdd, Zero};
30use std::cmp::Ordering;
31use std::sync::Arc;
32
33pub fn substring(
74 array: &dyn Array,
75 start: i64,
76 length: Option<u64>,
77) -> Result<ArrayRef, ArrowError> {
78 match array.data_type() {
79 DataType::Dictionary(_, _) => {
80 let dictionary = array.as_any_dictionary();
81 let values = substring(dictionary.values(), start, length)?;
82 Ok(Arc::new(dictionary.with_values(values)))
83 }
84 DataType::LargeBinary => byte_substring(
85 array.as_binary::<i64>(),
86 start,
87 length.map(u64_to_i64_saturating),
89 ),
90 DataType::Binary => byte_substring(
91 array.as_binary::<i32>(),
92 i64_to_i32_saturating(start),
94 length.map(u64_to_i32_saturating),
95 ),
96 DataType::FixedSizeBinary(old_len) => {
97 let old_len: usize = (*old_len)
98 .try_into()
99 .expect("negative FixedSizeBinary value length");
100 fixed_size_binary_substring(array.as_fixed_size_binary(), old_len, start, length)
101 }
102 DataType::LargeUtf8 => byte_substring(
103 array.as_string::<i64>(),
104 start,
105 length.map(u64_to_i64_saturating),
107 ),
108 DataType::Utf8 => byte_substring(
109 array.as_string::<i32>(),
110 i64_to_i32_saturating(start),
112 length.map(u64_to_i32_saturating),
113 ),
114 DataType::Utf8View => string_view_substring(array.as_string_view(), start, length),
115 DataType::BinaryView => binary_view_substring(array.as_binary_view(), start, length),
116 _ => Err(ArrowError::ComputeError(format!(
117 "substring does not support type {:?}",
118 array.data_type()
119 ))),
120 }
121}
122
123pub fn substring_by_char<OffsetSize: OffsetSizeTrait>(
153 array: &GenericStringArray<OffsetSize>,
154 start: i64,
155 length: Option<u64>,
156) -> Result<GenericStringArray<OffsetSize>, ArrowError> {
157 let length = length.map(|len| usize::try_from(len).unwrap_or(usize::MAX));
158
159 if array.is_ascii() {
160 Ok(substring_by_char_impl(array, length, |val| {
162 ascii_bounds(val, start, length)
163 }))
164 } else {
165 let max_element_len = length.map(|len| len.saturating_mul(4));
167 Ok(substring_by_char_impl(array, max_element_len, |val| {
168 utf8_bounds(val, start, length)
169 }))
170 }
171}
172
173fn substring_by_char_impl<OffsetSize: OffsetSizeTrait, F: Fn(&str) -> (usize, usize)>(
179 array: &GenericStringArray<OffsetSize>,
180 max_element_len: Option<usize>,
181 bounds: F,
182) -> GenericStringArray<OffsetSize> {
183 let mut vals = Vec::with_capacity({
184 let offsets = array.value_offsets();
185 let input_len = (offsets[array.len()] - offsets[0]).to_usize().unwrap();
186 match max_element_len {
187 Some(len) => input_len.min(array.len().saturating_mul(len)),
188 None => input_len,
189 }
190 });
191 let mut new_offsets = Vec::with_capacity(array.len() + 1);
192 new_offsets.push(OffsetSize::zero());
193
194 array.iter().for_each(|val| {
195 if let Some(val) = val {
196 let (start_offset, end_offset) = bounds(val);
197 vals.extend_from_slice(&val.as_bytes()[start_offset..end_offset]);
198 }
199 new_offsets.push(OffsetSize::from_usize(vals.len()).unwrap());
200 });
201
202 let offsets = OffsetBuffer::new(new_offsets.into());
203 let values = vals.into();
204 let nulls = array
205 .nulls()
206 .map(|n| n.inner().sliced())
207 .and_then(|b| NullBuffer::from_unsliced_buffer(b, array.len()));
208 GenericStringArray::<OffsetSize>::new(offsets, values, nulls)
209}
210
211#[inline]
217fn ascii_bounds(val: &str, start: i64, length: Option<usize>) -> (usize, usize) {
218 let len = val.len();
219 let start_offset = if start >= 0 {
220 usize::try_from(start).unwrap_or(usize::MAX).min(len)
221 } else {
222 len.saturating_sub(usize::try_from(start.unsigned_abs()).unwrap_or(usize::MAX))
223 };
224 let end_offset = length.map_or(len, |length| start_offset.saturating_add(length).min(len));
225 (start_offset, end_offset)
226}
227
228#[inline]
233fn utf8_bounds(val: &str, start: i64, length: Option<usize>) -> (usize, usize) {
234 let len = val.len();
235 let start_offset = if start >= 0 {
236 val.char_indices()
237 .nth(usize::try_from(start).unwrap_or(usize::MAX))
238 .map_or(len, |(offset, _)| offset)
239 } else {
240 let back = usize::try_from(start.unsigned_abs()).unwrap_or(usize::MAX);
243 val.char_indices()
244 .nth_back(back - 1)
245 .map_or(0, |(offset, _)| offset)
246 };
247 let end_offset = length.map_or(len, |length| {
248 if length >= len - start_offset {
251 return len;
252 }
253 val[start_offset..]
254 .char_indices()
255 .nth(length)
256 .map_or(len, |(offset, _)| start_offset + offset)
257 });
258 (start_offset, end_offset)
259}
260
261fn view_substring_range(
263 original_length: usize,
264 start: i64,
265 substring_length: Option<u64>,
266) -> (usize, usize) {
267 let original_length = original_length as i64;
268 let new_start = match start.cmp(&0) {
269 Ordering::Greater => start.min(original_length),
270 Ordering::Equal => 0,
271 Ordering::Less => (original_length + start).max(0),
272 };
273 let new_end = match substring_length {
274 Some(length) => new_start.saturating_add(length as i64).min(original_length),
275 None => original_length,
276 };
277 (new_start as usize, new_end as usize)
278}
279
280fn string_view_substring(
281 array: &StringViewArray,
282 start: i64,
283 length: Option<u64>,
284) -> Result<ArrayRef, ArrowError> {
285 let mut builder = StringViewBuilder::with_capacity(array.len());
286
287 for idx in 0..array.len() {
288 if array.is_null(idx) {
289 builder.append_null();
290 continue;
291 }
292 let value = array.value(idx);
293 let (new_start, new_end) = view_substring_range(value.len(), start, length);
294 for offset in [new_start, new_end] {
295 if !value.is_char_boundary(offset) {
296 return Err(ArrowError::ComputeError(format!(
297 "The offset {offset} is at an invalid utf-8 boundary."
298 )));
299 }
300 }
301 builder.append_value(&value[new_start..new_end]);
302 }
303
304 Ok(Arc::new(builder.finish()))
305}
306
307fn binary_view_substring(
308 array: &BinaryViewArray,
309 start: i64,
310 length: Option<u64>,
311) -> Result<ArrayRef, ArrowError> {
312 let mut builder = BinaryViewBuilder::with_capacity(array.len());
313
314 for idx in 0..array.len() {
315 if array.is_null(idx) {
316 builder.append_null();
317 continue;
318 }
319 let value = array.value(idx);
320 let (new_start, new_end) = view_substring_range(value.len(), start, length);
321 builder.append_value(&value[new_start..new_end]);
322 }
323
324 Ok(Arc::new(builder.finish()))
325}
326
327fn i64_to_i32_saturating(value: i64) -> i32 {
328 value.clamp(i32::MIN as i64, i32::MAX as i64) as i32
329}
330
331fn u64_to_i32_saturating(value: u64) -> i32 {
332 value.min(i32::MAX as u64) as i32
333}
334
335fn u64_to_i64_saturating(value: u64) -> i64 {
336 value.min(i64::MAX as u64) as i64
337}
338
339fn byte_substring<T: ByteArrayType>(
340 array: &GenericByteArray<T>,
341 start: T::Offset,
342 length: Option<T::Offset>,
343) -> Result<ArrayRef, ArrowError>
344where
345 <T as ByteArrayType>::Native: PartialEq,
346{
347 let offsets = array.value_offsets();
348 let data = array.value_data();
349 let zero = <T::Offset as Zero>::zero();
350
351 let check_char_boundary = {
353 |offset: T::Offset| {
354 if !matches!(T::DATA_TYPE, DataType::Utf8 | DataType::LargeUtf8) {
355 return Ok(offset);
356 }
357 let data_str = unsafe { std::str::from_utf8_unchecked(data) };
359 let offset_usize = offset.as_usize();
360 if data_str.is_char_boundary(offset_usize) {
361 Ok(offset)
362 } else {
363 Err(ArrowError::ComputeError(format!(
364 "The offset {offset_usize} is at an invalid utf-8 boundary."
365 )))
366 }
367 }
368 };
369
370 let mut new_starts_ends: Vec<(T::Offset, T::Offset)> = Vec::with_capacity(array.len());
372 let mut new_offsets: Vec<T::Offset> = Vec::with_capacity(array.len() + 1);
373 let mut len_so_far = zero;
374 new_offsets.push(zero);
375
376 offsets
377 .windows(2)
378 .try_for_each(|pair| -> Result<(), ArrowError> {
379 let new_start = match start.cmp(&zero) {
380 Ordering::Greater => {
381 let shifted = pair[0].checked_add(&start).unwrap_or(pair[1]);
385 check_char_boundary(shifted.min(pair[1]))?
386 }
387 Ordering::Equal => pair[0],
388 Ordering::Less => check_char_boundary((pair[1] + start).max(pair[0]))?,
389 };
390 let new_end = match length {
391 Some(length) => {
392 let end = length.checked_add(&new_start).unwrap_or(pair[1]);
393 check_char_boundary(end.min(pair[1]))?
394 }
395 None => pair[1],
396 };
397 len_so_far += new_end - new_start;
398 new_starts_ends.push((new_start, new_end));
399 new_offsets.push(len_so_far);
400 Ok(())
401 })?;
402
403 let mut new_values = MutableBuffer::new(new_offsets.last().unwrap().as_usize());
405
406 new_starts_ends
407 .iter()
408 .map(|(start, end)| {
409 let start = start.as_usize();
410 let end = end.as_usize();
411 &data[start..end]
412 })
413 .try_for_each(|slice| {
414 new_values
415 .try_extend_from_slice(slice)
416 .map_err(|e| ArrowError::MemoryError(e.to_string()))
417 })?;
418
419 let offsets = OffsetBuffer::new(new_offsets.into());
420 let values = new_values.into();
421 let nulls = array
422 .nulls()
423 .map(|n| n.inner().sliced())
424 .and_then(|b| NullBuffer::from_unsliced_buffer(b, array.len()));
425 Ok(Arc::new(GenericByteArray::<T>::new(offsets, values, nulls)))
426}
427
428fn fixed_size_binary_substring(
429 array: &FixedSizeBinaryArray,
430 old_len: usize,
431 start: i64,
432 length: Option<u64>,
433) -> Result<ArrayRef, ArrowError> {
434 let new_start = match start.cmp(&0) {
435 Ordering::Greater => usize::try_from(start).unwrap_or(usize::MAX).min(old_len),
436 Ordering::Equal => 0,
437 Ordering::Less => {
438 let offset = usize::try_from(start.unsigned_abs()).unwrap_or(usize::MAX);
439 old_len.saturating_sub(offset)
440 }
441 };
442
443 let new_len = match length {
444 Some(len) => usize::try_from(len)
445 .unwrap_or(usize::MAX)
446 .min(old_len - new_start),
447 None => old_len - new_start,
448 };
449
450 let num_of_elements = array.len();
452 let data = array.value_data();
453 let capacity = num_of_elements
454 .checked_mul(new_len)
455 .expect("capacity overflow");
456 let mut new_values = MutableBuffer::new(capacity);
457 (0..num_of_elements)
458 .map(|idx| {
459 let offset = idx * array.value_size();
460 (offset + new_start, offset + new_start + new_len)
461 })
462 .try_for_each(|(start, end)| {
463 new_values
464 .try_extend_from_slice(&data[start..end])
465 .map_err(|e| ArrowError::MemoryError(e.to_string()))
466 })?;
467
468 let mut nulls = array
469 .nulls()
470 .map(|n| n.inner().sliced())
471 .and_then(|b| NullBuffer::from_unsliced_buffer(b, num_of_elements));
472
473 if new_len == 0 && nulls.is_none() {
474 nulls = Some(NullBuffer::new_valid(num_of_elements));
479 }
480
481 let new_len: i32 = new_len.try_into().expect("new_len overflow");
482
483 Ok(Arc::new(FixedSizeBinaryArray::new(
484 new_len,
485 new_values.into(),
486 nulls,
487 )))
488}
489
490#[cfg(test)]
491mod tests {
492 use super::*;
493 use arrow_buffer::BooleanBuffer;
494 use arrow_buffer::Buffer;
495
496 macro_rules! gen_test_cases {
507 ($input:expr, $(($start:expr, $len:expr, $result:expr)), *) => {
508 [
509 $(
510 ($input.clone(), $start, $len, $result),
511 )*
512 ]
513 };
514 }
515
516 macro_rules! do_test {
523 ($cases:expr, $array_ty:ty, $substring_fn:ident) => {
524 $cases
525 .into_iter()
526 .for_each(|(array, start, length, expected)| {
527 let array = <$array_ty>::from(array);
528 let result = $substring_fn(&array, start, length).unwrap();
529 let result = result.as_any().downcast_ref::<$array_ty>().unwrap();
530 let expected = <$array_ty>::from(expected);
531 assert_eq!(&expected, result);
532 })
533 };
534 }
535
536 macro_rules! do_test_tryfrom {
538 ($cases:expr, $array_ty:ty, $substring_fn:ident) => {
539 $cases
540 .into_iter()
541 .for_each(|(array, start, length, expected)| {
542 let array = <$array_ty>::try_from(array).unwrap();
543 let result = $substring_fn(&array, start, length).unwrap();
544 let result = result.as_any().downcast_ref::<$array_ty>().unwrap();
545 let expected = <$array_ty>::try_from(expected).unwrap();
546 assert_eq!(&expected, result);
547 })
548 };
549 }
550
551 fn with_nulls_generic_binary<O: OffsetSizeTrait>() {
552 let input = vec![
553 Some(b"hello".as_slice()),
554 None,
555 Some(&[0xf8, 0xf9, 0xff, 0xfa]),
556 ];
557 let base_case = gen_test_cases!(
559 vec![None, None, None],
560 (-1, Some(1), vec![None, None, None])
561 );
562 let cases = gen_test_cases!(
563 input,
564 (0, None, input.clone()),
566 (0, Some(0), vec![Some(&[]), None, Some(&[])]),
568 (1000, Some(0), vec![Some(&[]), None, Some(&[])]),
570 (-1000, None, input.clone()),
572 (0, Some(1000), input.clone())
574 );
575
576 do_test!(
577 [&base_case[..], &cases[..]].concat(),
578 GenericBinaryArray<O>,
579 substring
580 );
581 }
582
583 #[test]
584 fn with_nulls_binary() {
585 with_nulls_generic_binary::<i32>()
586 }
587
588 #[test]
589 fn with_nulls_large_binary() {
590 with_nulls_generic_binary::<i64>()
591 }
592
593 fn without_nulls_generic_binary<O: OffsetSizeTrait>() {
594 let input = vec![b"hello".as_slice(), b"", &[0xf8, 0xf9, 0xff, 0xfa]];
595 let base_case = gen_test_cases!(
597 vec![b"".as_slice(), b"", b""],
598 (2, Some(1), vec![b"".as_slice(), b"", b""])
599 );
600 let cases = gen_test_cases!(
601 input,
602 (0, None, input.clone()),
604 (1, None, vec![b"ello", b"", &[0xf9, 0xff, 0xfa]]),
606 (2, None, vec![b"llo", b"", &[0xff, 0xfa]]),
607 (3, None, vec![b"lo", b"", &[0xfa]]),
608 (10, None, vec![b"", b"", b""]),
609 (-1, None, vec![b"o", b"", &[0xfa]]),
611 (-2, None, vec![b"lo", b"", &[0xff, 0xfa]]),
612 (-3, None, vec![b"llo", b"", &[0xf9, 0xff, 0xfa]]),
613 (-10, None, input.clone()),
614 (1, Some(1), vec![b"e", b"", &[0xf9]]),
616 (1, Some(2), vec![b"el", b"", &[0xf9, 0xff]]),
617 (1, Some(3), vec![b"ell", b"", &[0xf9, 0xff, 0xfa]]),
618 (1, Some(4), vec![b"ello", b"", &[0xf9, 0xff, 0xfa]]),
619 (-3, Some(1), vec![b"l", b"", &[0xf9]]),
620 (-3, Some(2), vec![b"ll", b"", &[0xf9, 0xff]]),
621 (-3, Some(3), vec![b"llo", b"", &[0xf9, 0xff, 0xfa]]),
622 (-3, Some(4), vec![b"llo", b"", &[0xf9, 0xff, 0xfa]])
623 );
624
625 do_test!(
626 [&base_case[..], &cases[..]].concat(),
627 GenericBinaryArray<O>,
628 substring
629 );
630 }
631
632 #[test]
633 fn without_nulls_binary() {
634 without_nulls_generic_binary::<i32>()
635 }
636
637 #[test]
638 fn without_nulls_large_binary() {
639 without_nulls_generic_binary::<i64>()
640 }
641
642 fn generic_binary_with_non_zero_offset<O: OffsetSizeTrait>() {
643 let values = 0_u8..15;
644 let offsets = &[
645 O::zero(),
646 O::from_usize(5).unwrap(),
647 O::from_usize(10).unwrap(),
648 O::from_usize(15).unwrap(),
649 ];
650 let bitmap = [0b101_u8];
652
653 let offsets = OffsetBuffer::new(Buffer::from_slice_ref(offsets).into());
654 let values = Buffer::from_iter(values);
655 let nulls = Some(NullBuffer::new(BooleanBuffer::new(
656 Buffer::from(bitmap),
657 0,
658 3,
659 )));
660 let array = GenericBinaryArray::<O>::new(offsets, values, nulls).slice(1, 2);
662 let result = substring(&array, 1, None).unwrap();
664 let result = result
665 .as_any()
666 .downcast_ref::<GenericBinaryArray<O>>()
667 .unwrap();
668 let expected =
669 GenericBinaryArray::<O>::from_opt_vec(vec![None, Some(&[11_u8, 12, 13, 14])]);
670 assert_eq!(result, &expected);
671 }
672
673 #[test]
674 fn binary_with_non_zero_offset() {
675 generic_binary_with_non_zero_offset::<i32>()
676 }
677
678 #[test]
679 fn large_binary_with_non_zero_offset() {
680 generic_binary_with_non_zero_offset::<i64>()
681 }
682
683 #[test]
684 fn with_nulls_fixed_size_binary() {
685 let input = vec![Some(b"cat".as_slice()), None, Some(&[0xf8, 0xf9, 0xff])];
686 let base_case =
688 gen_test_cases!(vec![None, None, None], (3, Some(2), vec![None, None, None]));
689 let cases = gen_test_cases!(
690 input,
691 (0, None, input.clone()),
693 (1, None, vec![Some(b"at"), None, Some(&[0xf9, 0xff])]),
695 (2, None, vec![Some(b"t"), None, Some(&[0xff])]),
696 (3, None, vec![Some(b""), None, Some(b"")]),
697 (10, None, vec![Some(b""), None, Some(b"")]),
698 (-1, None, vec![Some(b"t"), None, Some(&[0xff])]),
700 (-2, None, vec![Some(b"at"), None, Some(&[0xf9, 0xff])]),
701 (-3, None, input.clone()),
702 (-10, None, input.clone()),
703 (1, Some(1), vec![Some(b"a"), None, Some(&[0xf9])]),
705 (1, Some(2), vec![Some(b"at"), None, Some(&[0xf9, 0xff])]),
706 (1, Some(3), vec![Some(b"at"), None, Some(&[0xf9, 0xff])]),
707 (-3, Some(1), vec![Some(b"c"), None, Some(&[0xf8])]),
708 (-3, Some(2), vec![Some(b"ca"), None, Some(&[0xf8, 0xf9])]),
709 (-3, Some(3), input.clone()),
710 (-3, Some(4), input.clone())
711 );
712
713 do_test_tryfrom!(
714 [&base_case[..], &cases[..]].concat(),
715 FixedSizeBinaryArray,
716 substring
717 );
718 }
719
720 #[test]
721 fn without_nulls_fixed_size_binary() {
722 let input = vec![b"cat".as_slice(), b"dog", &[0xf8, 0xf9, 0xff]];
723 let base_case = gen_test_cases!(
725 vec![b"".as_slice(), &[], &[]],
726 (1, Some(2), vec![b"".as_slice(), &[], &[]])
727 );
728 let cases = gen_test_cases!(
729 input,
730 (0, None, input.clone()),
732 (1, None, vec![b"at", b"og", &[0xf9, 0xff]]),
734 (2, None, vec![b"t", b"g", &[0xff]]),
735 (3, None, vec![&[], &[], &[]]),
736 (10, None, vec![&[], &[], &[]]),
737 (-1, None, vec![b"t", b"g", &[0xff]]),
739 (-2, None, vec![b"at", b"og", &[0xf9, 0xff]]),
740 (-3, None, input.clone()),
741 (-10, None, input.clone()),
742 (1, Some(1), vec![b"a", b"o", &[0xf9]]),
744 (1, Some(2), vec![b"at", b"og", &[0xf9, 0xff]]),
745 (1, Some(3), vec![b"at", b"og", &[0xf9, 0xff]]),
746 (-3, Some(1), vec![b"c", b"d", &[0xf8]]),
747 (-3, Some(2), vec![b"ca", b"do", &[0xf8, 0xf9]]),
748 (-3, Some(3), input.clone()),
749 (-3, Some(4), input.clone())
750 );
751
752 do_test_tryfrom!(
753 [&base_case[..], &cases[..]].concat(),
754 FixedSizeBinaryArray,
755 substring
756 );
757 }
758
759 #[test]
760 fn fixed_size_binary_with_non_zero_offset() {
761 let values = b"hellotherearrow";
762 let bits_v = [0b101_u8];
764
765 let nulls = Some(NullBuffer::new(BooleanBuffer::new(
766 Buffer::from(bits_v),
767 0,
768 3,
769 )));
770 let array = FixedSizeBinaryArray::new(5, Buffer::from(values), nulls).slice(1, 2);
772 let result = substring(&array, 1, None).unwrap();
774 let result = result
775 .as_any()
776 .downcast_ref::<FixedSizeBinaryArray>()
777 .unwrap();
778 let expected = FixedSizeBinaryArray::try_from_sparse_iter_with_size(
779 vec![None, Some(b"rrow")].into_iter(),
780 4,
781 )
782 .unwrap();
783 assert_eq!(result, &expected);
784 }
785
786 fn with_nulls_generic_string<O: OffsetSizeTrait>() {
787 let input = vec![Some("hello"), None, Some("word")];
788 let base_case = gen_test_cases!(vec![None, None, None], (0, None, vec![None, None, None]));
790 let cases = gen_test_cases!(
791 input,
792 (0, None, input.clone()),
794 (0, Some(0), vec![Some(""), None, Some("")]),
796 (1000, Some(0), vec![Some(""), None, Some("")]),
798 (-1000, None, input.clone()),
800 (0, Some(1000), input.clone())
802 );
803
804 do_test!(
805 [&base_case[..], &cases[..]].concat(),
806 GenericStringArray<O>,
807 substring
808 );
809 }
810
811 #[test]
812 fn with_nulls_string() {
813 with_nulls_generic_string::<i32>()
814 }
815
816 #[test]
817 fn with_nulls_large_string() {
818 with_nulls_generic_string::<i64>()
819 }
820
821 fn without_nulls_generic_string<O: OffsetSizeTrait>() {
822 let input = vec!["hello", "", "word"];
823 let base_case = gen_test_cases!(vec!["", "", ""], (0, None, vec!["", "", ""]));
825 let cases = gen_test_cases!(
826 input,
827 (0, None, input.clone()),
829 (1, None, vec!["ello", "", "ord"]),
830 (2, None, vec!["llo", "", "rd"]),
831 (3, None, vec!["lo", "", "d"]),
832 (10, None, vec!["", "", ""]),
833 (-1, None, vec!["o", "", "d"]),
835 (-2, None, vec!["lo", "", "rd"]),
836 (-3, None, vec!["llo", "", "ord"]),
837 (-10, None, input.clone()),
838 (1, Some(1), vec!["e", "", "o"]),
840 (1, Some(2), vec!["el", "", "or"]),
841 (1, Some(3), vec!["ell", "", "ord"]),
842 (1, Some(4), vec!["ello", "", "ord"]),
843 (-3, Some(1), vec!["l", "", "o"]),
844 (-3, Some(2), vec!["ll", "", "or"]),
845 (-3, Some(3), vec!["llo", "", "ord"]),
846 (-3, Some(4), vec!["llo", "", "ord"])
847 );
848
849 do_test!(
850 [&base_case[..], &cases[..]].concat(),
851 GenericStringArray<O>,
852 substring
853 );
854 }
855
856 #[test]
857 fn without_nulls_string() {
858 without_nulls_generic_string::<i32>()
859 }
860
861 #[test]
862 fn without_nulls_large_string() {
863 without_nulls_generic_string::<i64>()
864 }
865
866 fn generic_string_with_non_zero_offset<O: OffsetSizeTrait>() {
867 let values = b"hellotherearrow";
868 let offsets = &[
869 O::zero(),
870 O::from_usize(5).unwrap(),
871 O::from_usize(10).unwrap(),
872 O::from_usize(15).unwrap(),
873 ];
874 let bitmap = [0b101_u8];
876
877 let offsets = OffsetBuffer::new(Buffer::from_slice_ref(offsets).into());
878 let values = Buffer::from(values);
879 let nulls = Some(NullBuffer::new(BooleanBuffer::new(
880 Buffer::from(bitmap),
881 0,
882 3,
883 )));
884 let array = GenericStringArray::<O>::new(offsets, values, nulls).slice(1, 2);
886 let result = substring(&array, 1, None).unwrap();
888 let result = result
889 .as_any()
890 .downcast_ref::<GenericStringArray<O>>()
891 .unwrap();
892 let expected = GenericStringArray::<O>::from(vec![None, Some("rrow")]);
893 assert_eq!(result, &expected);
894 }
895
896 #[test]
897 fn string_with_non_zero_offset() {
898 generic_string_with_non_zero_offset::<i32>()
899 }
900
901 #[test]
902 fn large_string_with_non_zero_offset() {
903 generic_string_with_non_zero_offset::<i64>()
904 }
905
906 fn with_nulls_generic_string_by_char<O: OffsetSizeTrait>() {
907 let input = vec![Some("hello"), None, Some("Γ ⊢x:T")];
908 let base_case = gen_test_cases!(vec![None, None, None], (0, None, vec![None, None, None]));
910 let cases = gen_test_cases!(
911 input,
912 (0, None, input.clone()),
914 (0, Some(0), vec![Some(""), None, Some("")]),
916 (1000, Some(0), vec![Some(""), None, Some("")]),
918 (-1000, None, input.clone()),
920 (0, Some(1000), input.clone())
922 );
923
924 do_test!(
925 [&base_case[..], &cases[..]].concat(),
926 GenericStringArray<O>,
927 substring_by_char
928 );
929 }
930
931 #[test]
932 fn with_nulls_string_by_char() {
933 with_nulls_generic_string_by_char::<i32>()
934 }
935
936 #[test]
937 fn with_nulls_large_string_by_char() {
938 with_nulls_generic_string_by_char::<i64>()
939 }
940
941 fn without_nulls_generic_string_by_char<O: OffsetSizeTrait>() {
942 let input = vec!["hello", "", "Γ ⊢x:T"];
943 let base_case = gen_test_cases!(vec!["", "", ""], (0, None, vec!["", "", ""]));
945 let cases = gen_test_cases!(
946 input,
947 (0, None, input.clone()),
949 (1, None, vec!["ello", "", " ⊢x:T"]),
951 (2, None, vec!["llo", "", "⊢x:T"]),
952 (3, None, vec!["lo", "", "x:T"]),
953 (10, None, vec!["", "", ""]),
954 (-1, None, vec!["o", "", "T"]),
956 (-2, None, vec!["lo", "", ":T"]),
957 (-4, None, vec!["ello", "", "⊢x:T"]),
958 (-10, None, input.clone()),
959 (1, Some(1), vec!["e", "", " "]),
961 (1, Some(2), vec!["el", "", " ⊢"]),
962 (1, Some(3), vec!["ell", "", " ⊢x"]),
963 (1, Some(6), vec!["ello", "", " ⊢x:T"]),
964 (-4, Some(1), vec!["e", "", "⊢"]),
965 (-4, Some(2), vec!["el", "", "⊢x"]),
966 (-4, Some(3), vec!["ell", "", "⊢x:"]),
967 (-4, Some(4), vec!["ello", "", "⊢x:T"])
968 );
969
970 do_test!(
971 [&base_case[..], &cases[..]].concat(),
972 GenericStringArray<O>,
973 substring_by_char
974 );
975 }
976
977 #[test]
978 fn without_nulls_string_by_char() {
979 without_nulls_generic_string_by_char::<i32>()
980 }
981
982 #[test]
983 fn without_nulls_large_string_by_char() {
984 without_nulls_generic_string_by_char::<i64>()
985 }
986
987 fn ascii_generic_string_by_char<O: OffsetSizeTrait>() {
989 let input = vec![Some("hello"), None, Some(""), Some("rust")];
990 let cases = gen_test_cases!(
991 input,
992 (0, None, input.clone()),
994 (1, None, vec![Some("ello"), None, Some(""), Some("ust")]),
996 (4, None, vec![Some("o"), None, Some(""), Some("")]),
997 (1000, None, vec![Some(""), None, Some(""), Some("")]),
999 (-1, None, vec![Some("o"), None, Some(""), Some("t")]),
1001 (-4, None, vec![Some("ello"), None, Some(""), Some("rust")]),
1002 (-1000, None, input.clone()),
1004 (0, Some(0), vec![Some(""), None, Some(""), Some("")]),
1006 (1, Some(2), vec![Some("el"), None, Some(""), Some("us")]),
1008 (-4, Some(2), vec![Some("el"), None, Some(""), Some("ru")]),
1009 (0, Some(1000), input.clone()),
1011 (
1013 1,
1014 Some(u64::MAX),
1015 vec![Some("ello"), None, Some(""), Some("ust")]
1016 )
1017 );
1018
1019 do_test!(cases, GenericStringArray<O>, substring_by_char);
1020 }
1021
1022 #[test]
1023 fn ascii_string_by_char() {
1024 ascii_generic_string_by_char::<i32>()
1025 }
1026
1027 #[test]
1028 fn ascii_large_string_by_char() {
1029 ascii_generic_string_by_char::<i64>()
1030 }
1031
1032 fn generic_string_by_char_with_non_zero_offset<O: OffsetSizeTrait>() {
1033 let values = "S→T = Πx:S.T";
1034 let offsets = &[
1035 O::zero(),
1036 O::from_usize(values.char_indices().nth(3).map(|(pos, _)| pos).unwrap()).unwrap(),
1037 O::from_usize(values.char_indices().nth(6).map(|(pos, _)| pos).unwrap()).unwrap(),
1038 O::from_usize(values.len()).unwrap(),
1039 ];
1040 let bitmap = [0b101_u8];
1042
1043 let offsets = OffsetBuffer::new(Buffer::from_slice_ref(offsets).into());
1044 let values = Buffer::from(values.as_bytes());
1045 let nulls = Some(NullBuffer::new(BooleanBuffer::new(
1046 Buffer::from(bitmap),
1047 0,
1048 3,
1049 )));
1050 let array = GenericStringArray::<O>::new(offsets, values, nulls).slice(1, 2);
1052 let result = substring_by_char(&array, 1, None).unwrap();
1054 let expected = GenericStringArray::<O>::from(vec![None, Some("x:S.T")]);
1055 assert_eq!(result, expected);
1056 }
1057
1058 #[test]
1059 fn string_with_non_zero_offset_by_char() {
1060 generic_string_by_char_with_non_zero_offset::<i32>()
1061 }
1062
1063 #[test]
1064 fn large_string_with_non_zero_offset_by_char() {
1065 generic_string_by_char_with_non_zero_offset::<i64>()
1066 }
1067
1068 #[test]
1069 fn dictionary() {
1070 _dictionary::<Int8Type>();
1071 _dictionary::<Int16Type>();
1072 _dictionary::<Int32Type>();
1073 _dictionary::<Int64Type>();
1074 _dictionary::<UInt8Type>();
1075 _dictionary::<UInt16Type>();
1076 _dictionary::<UInt32Type>();
1077 _dictionary::<UInt64Type>();
1078 }
1079
1080 fn _dictionary<K: ArrowDictionaryKeyType>() {
1081 const TOTAL: i32 = 100;
1082
1083 let v = ["aaa", "bbb", "ccc", "ddd", "eee"];
1084 let data: Vec<Option<&str>> = (0..TOTAL)
1085 .map(|n| {
1086 let i = n % 5;
1087 if i == 3 { None } else { Some(v[i as usize]) }
1088 })
1089 .collect();
1090
1091 let dict_array: DictionaryArray<K> = data.clone().into_iter().collect();
1092
1093 let expected: Vec<Option<&str>> = data.iter().map(|opt| opt.map(|s| &s[1..3])).collect();
1094
1095 let res = substring(&dict_array, 1, Some(2)).unwrap();
1096 let actual = res.as_any().downcast_ref::<DictionaryArray<K>>().unwrap();
1097 let actual: Vec<Option<&str>> = actual
1098 .values()
1099 .as_any()
1100 .downcast_ref::<GenericStringArray<i32>>()
1101 .unwrap()
1102 .take_iter(actual.keys_iter())
1103 .collect();
1104
1105 for i in 0..TOTAL as usize {
1106 assert_eq!(expected[i], actual[i],);
1107 }
1108 }
1109
1110 #[test]
1111 fn check_invalid_array_type() {
1112 let array = Int32Array::from(vec![Some(1), Some(2), Some(3)]);
1113 let err = substring(&array, 0, None).unwrap_err().to_string();
1114 assert!(err.contains("substring does not support type"));
1115 }
1116
1117 #[test]
1119 fn check_start_index() {
1120 let array = StringArray::from(vec![Some("E=mc²"), Some("ascii")]);
1121 let err = substring(&array, -1, None).unwrap_err().to_string();
1122 assert!(err.contains("invalid utf-8 boundary"));
1123 }
1124
1125 #[test]
1126 fn check_length() {
1127 let array = StringArray::from(vec![Some("E=mc²"), Some("ascii")]);
1128 let err = substring(&array, 0, Some(5)).unwrap_err().to_string();
1129 assert!(err.contains("invalid utf-8 boundary"));
1130 }
1131
1132 #[test]
1133 fn non_utf8_bytes() {
1134 let bytes: &[u8] = &[0xE4, 0xBD, 0xA0, 0xE5, 0xA5, 0xBD, 0xE8, 0xAF, 0xAD];
1136 let array = BinaryArray::from(vec![Some(bytes)]);
1137 let arr = substring(&array, 0, Some(5)).unwrap();
1138 let actual = arr.as_any().downcast_ref::<BinaryArray>().unwrap();
1139
1140 let expected_bytes: &[u8] = &[0xE4, 0xBD, 0xA0, 0xE5, 0xA5];
1141 let expected = BinaryArray::from(vec![Some(expected_bytes)]);
1142 assert_eq!(expected, *actual);
1143 }
1144
1145 #[test]
1146 fn string_view_matches_utf8() {
1147 let values = vec![
1148 Some("hello world"),
1149 Some(""),
1150 None,
1151 Some("a"),
1152 Some("this one is definitely longer than twelve bytes"),
1153 ];
1154 let utf8 = StringArray::from(values.clone());
1155 let view = StringViewArray::from(values);
1156
1157 for (start, length) in [
1158 (0, None),
1159 (0, Some(0)),
1160 (0, Some(5)),
1161 (0, Some(1000)),
1162 (1, Some(3)),
1163 (5, None),
1164 (100, Some(2)),
1165 (100, None),
1166 (-3, None),
1167 (-3, Some(2)),
1168 (-100, Some(4)),
1169 (-100, None),
1170 ] {
1171 let expected = substring(&utf8, start, length).unwrap();
1172 let expected = expected.as_string::<i32>();
1173 let actual = substring(&view, start, length).unwrap();
1174 let actual = actual.as_string_view();
1175 assert_eq!(
1176 expected.iter().collect::<Vec<_>>(),
1177 actual.iter().collect::<Vec<_>>(),
1178 "start={start} length={length:?}"
1179 );
1180 }
1181 }
1182
1183 #[test]
1184 fn binary_view_matches_binary() {
1185 let values: Vec<Option<&[u8]>> = vec![
1186 Some(b"hello world"),
1187 Some(b""),
1188 None,
1189 Some(b"abc"),
1190 Some(b"this one is definitely longer than twelve bytes"),
1191 ];
1192 let binary = BinaryArray::from(values.clone());
1193 let view = BinaryViewArray::from(values);
1194
1195 for (start, length) in [
1196 (0, None),
1197 (0, Some(5)),
1198 (2, Some(3)),
1199 (-3, None),
1200 (100, Some(2)),
1201 ] {
1202 let expected = substring(&binary, start, length).unwrap();
1203 let expected = expected.as_binary::<i32>();
1204 let actual = substring(&view, start, length).unwrap();
1205 let actual = actual.as_binary_view();
1206 assert_eq!(
1207 expected.iter().collect::<Vec<_>>(),
1208 actual.iter().collect::<Vec<_>>(),
1209 "start={start} length={length:?}"
1210 );
1211 }
1212 }
1213
1214 #[test]
1215 fn string_view_slices_on_a_multi_byte_boundary() {
1216 let values = vec![Some("héllo wörld"), Some("日本語"), None];
1219 let utf8 = StringArray::from(values.clone());
1220 let view = StringViewArray::from(values);
1221
1222 for (start, length) in [
1223 (0, Some(3)),
1224 (3, Some(3)),
1225 (0, Some(6)),
1226 (3, None),
1227 (-3, None),
1228 (-6, None),
1229 ] {
1230 let expected = substring(&utf8, start, length).unwrap();
1231 let expected = expected.as_string::<i32>();
1232 let actual = substring(&view, start, length).unwrap();
1233 let actual = actual.as_string_view();
1234 assert_eq!(
1235 expected.iter().collect::<Vec<_>>(),
1236 actual.iter().collect::<Vec<_>>(),
1237 "start={start} length={length:?}"
1238 );
1239 }
1240
1241 let actual = substring(&view, 0, Some(3)).unwrap();
1242 let actual = actual.as_string_view();
1243 assert_eq!(actual.value(0), "hé");
1244 assert_eq!(actual.value(1), "日");
1245 }
1246
1247 #[test]
1248 fn string_view_rejects_an_invalid_char_boundary() {
1249 let view = StringViewArray::from(vec![Some("E=mc²")]);
1250 let err = substring(&view, 0, Some(5)).unwrap_err().to_string();
1251 assert!(err.contains("invalid utf-8 boundary"), "{err}");
1252 }
1253
1254 #[test]
1255 fn dictionary_of_string_view() {
1256 let view = StringViewArray::from(vec![Some("hello world"), Some("bye")]);
1257 let dict = DictionaryArray::new(Int32Array::from(vec![0, 1, 0]), Arc::new(view));
1258
1259 let actual = substring(&dict, 0, Some(3)).unwrap();
1260 let actual = actual.as_any_dictionary();
1261
1262 let values = actual.values().as_string_view();
1263 assert_eq!(
1264 values.iter().collect::<Vec<_>>(),
1265 vec![Some("hel"), Some("bye")]
1266 );
1267 }
1268
1269 #[test]
1270 fn out_of_range_start_and_length_match_the_64_bit_arms() {
1271 let values = vec![Some("hello"), Some("world"), None];
1274 let utf8 = StringArray::from(values.clone());
1275 let large = LargeStringArray::from(values.clone());
1276 let binary = BinaryArray::from_iter(values.iter().map(|v| v.map(|s| s.as_bytes())));
1277 let large_binary =
1278 LargeBinaryArray::from_iter(values.iter().map(|v| v.map(|s| s.as_bytes())));
1279
1280 let starts: [i64; 4] = [1 << 31, (1 << 31) + 5, 1 << 32, i64::MAX];
1282 let lengths: [Option<u64>; 3] = [None, Some(1 << 31), Some(u64::MAX)];
1283
1284 for start in starts {
1285 for length in lengths {
1286 let narrow = substring(&utf8, start, length).unwrap();
1287 let wide = substring(&large, start, length).unwrap();
1288 assert_eq!(
1289 narrow.as_string::<i32>().iter().collect::<Vec<_>>(),
1290 wide.as_string::<i64>().iter().collect::<Vec<_>>(),
1291 "Utf8 and LargeUtf8 disagree at start {start}, length {length:?}"
1292 );
1293
1294 let narrow = substring(&binary, start, length).unwrap();
1295 let wide = substring(&large_binary, start, length).unwrap();
1296 assert_eq!(
1297 narrow.as_binary::<i32>().iter().collect::<Vec<_>>(),
1298 wide.as_binary::<i64>().iter().collect::<Vec<_>>(),
1299 "Binary and LargeBinary disagree at start {start}, length {length:?}"
1300 );
1301 }
1302 }
1303
1304 let out = substring(&utf8, 1 << 31, None).unwrap();
1307 assert_eq!(
1308 out.as_string::<i32>().iter().collect::<Vec<_>>(),
1309 vec![Some(""), Some(""), None]
1310 );
1311
1312 let out = substring(&utf8, 3, None).unwrap();
1314 assert_eq!(
1315 out.as_string::<i32>().iter().collect::<Vec<_>>(),
1316 vec![Some("lo"), Some("ld"), None]
1317 );
1318 }
1319}