1use arrow_array::builder::{BinaryViewBuilder, BufferBuilder, StringViewBuilder};
24use arrow_array::cast::AsArray;
25use arrow_array::types::*;
26use arrow_array::*;
27use arrow_buffer::{ArrowNativeType, MutableBuffer, NullBuffer, OffsetBuffer};
28use arrow_schema::{ArrowError, DataType};
29use num_traits::Zero;
30use std::cmp::Ordering;
31use std::sync::Arc;
32
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 => {
85 byte_substring(array.as_binary::<i64>(), start, length.map(|e| e as i64))
86 }
87 DataType::Binary => byte_substring(
88 array.as_binary::<i32>(),
89 start as i32,
90 length.map(|e| e as i32),
91 ),
92 DataType::FixedSizeBinary(old_len) => {
93 let old_len: usize = (*old_len)
94 .try_into()
95 .expect("negative FixedSizeBinary value length");
96 fixed_size_binary_substring(array.as_fixed_size_binary(), old_len, start, length)
97 }
98 DataType::LargeUtf8 => {
99 byte_substring(array.as_string::<i64>(), start, length.map(|e| e as i64))
100 }
101 DataType::Utf8 => byte_substring(
102 array.as_string::<i32>(),
103 start as i32,
104 length.map(|e| e as i32),
105 ),
106 DataType::Utf8View => string_view_substring(array.as_string_view(), start, length),
107 DataType::BinaryView => binary_view_substring(array.as_binary_view(), start, length),
108 _ => Err(ArrowError::ComputeError(format!(
109 "substring does not support type {:?}",
110 array.data_type()
111 ))),
112 }
113}
114
115pub fn substring_by_char<OffsetSize: OffsetSizeTrait>(
145 array: &GenericStringArray<OffsetSize>,
146 start: i64,
147 length: Option<u64>,
148) -> Result<GenericStringArray<OffsetSize>, ArrowError> {
149 let length = length.map(|len| usize::try_from(len).unwrap_or(usize::MAX));
150
151 if array.is_ascii() {
152 Ok(substring_by_char_impl(array, length, |val| {
154 ascii_bounds(val, start, length)
155 }))
156 } else {
157 let max_element_len = length.map(|len| len.saturating_mul(4));
159 Ok(substring_by_char_impl(array, max_element_len, |val| {
160 utf8_bounds(val, start, length)
161 }))
162 }
163}
164
165fn substring_by_char_impl<OffsetSize: OffsetSizeTrait, F: Fn(&str) -> (usize, usize)>(
171 array: &GenericStringArray<OffsetSize>,
172 max_element_len: Option<usize>,
173 bounds: F,
174) -> GenericStringArray<OffsetSize> {
175 let mut vals = BufferBuilder::<u8>::new({
176 let offsets = array.value_offsets();
177 let input_len = (offsets[array.len()] - offsets[0]).to_usize().unwrap();
178 match max_element_len {
179 Some(len) => input_len.min(array.len().saturating_mul(len)),
180 None => input_len,
181 }
182 });
183 let mut new_offsets = BufferBuilder::<OffsetSize>::new(array.len() + 1);
184 new_offsets.append(OffsetSize::zero());
185
186 array.iter().for_each(|val| {
187 if let Some(val) = val {
188 let (start_offset, end_offset) = bounds(val);
189 vals.append_slice(&val.as_bytes()[start_offset..end_offset]);
190 }
191 new_offsets.append(OffsetSize::from_usize(vals.len()).unwrap());
192 });
193
194 let offsets = OffsetBuffer::new(new_offsets.finish().into());
195 let values = vals.finish();
196 let nulls = array
197 .nulls()
198 .map(|n| n.inner().sliced())
199 .and_then(|b| NullBuffer::from_unsliced_buffer(b, array.len()));
200 GenericStringArray::<OffsetSize>::new(offsets, values, nulls)
201}
202
203#[inline]
209fn ascii_bounds(val: &str, start: i64, length: Option<usize>) -> (usize, usize) {
210 let len = val.len();
211 let start_offset = if start >= 0 {
212 usize::try_from(start).unwrap_or(usize::MAX).min(len)
213 } else {
214 len.saturating_sub(usize::try_from(start.unsigned_abs()).unwrap_or(usize::MAX))
215 };
216 let end_offset = length.map_or(len, |length| start_offset.saturating_add(length).min(len));
217 (start_offset, end_offset)
218}
219
220#[inline]
225fn utf8_bounds(val: &str, start: i64, length: Option<usize>) -> (usize, usize) {
226 let len = val.len();
227 let start_offset = if start >= 0 {
228 val.char_indices()
229 .nth(usize::try_from(start).unwrap_or(usize::MAX))
230 .map_or(len, |(offset, _)| offset)
231 } else {
232 let back = usize::try_from(start.unsigned_abs()).unwrap_or(usize::MAX);
235 val.char_indices()
236 .nth_back(back - 1)
237 .map_or(0, |(offset, _)| offset)
238 };
239 let end_offset = length.map_or(len, |length| {
240 if length >= len - start_offset {
243 return len;
244 }
245 val[start_offset..]
246 .char_indices()
247 .nth(length)
248 .map_or(len, |(offset, _)| start_offset + offset)
249 });
250 (start_offset, end_offset)
251}
252
253fn view_substring_range(
255 original_length: usize,
256 start: i64,
257 substring_length: Option<u64>,
258) -> (usize, usize) {
259 let original_length = original_length as i64;
260 let new_start = match start.cmp(&0) {
261 Ordering::Greater => start.min(original_length),
262 Ordering::Equal => 0,
263 Ordering::Less => (original_length + start).max(0),
264 };
265 let new_end = match substring_length {
266 Some(length) => new_start.saturating_add(length as i64).min(original_length),
267 None => original_length,
268 };
269 (new_start as usize, new_end as usize)
270}
271
272fn string_view_substring(
273 array: &StringViewArray,
274 start: i64,
275 length: Option<u64>,
276) -> Result<ArrayRef, ArrowError> {
277 let mut builder = StringViewBuilder::with_capacity(array.len());
278
279 for idx in 0..array.len() {
280 if array.is_null(idx) {
281 builder.append_null();
282 continue;
283 }
284 let value = array.value(idx);
285 let (new_start, new_end) = view_substring_range(value.len(), start, length);
286 for offset in [new_start, new_end] {
287 if !value.is_char_boundary(offset) {
288 return Err(ArrowError::ComputeError(format!(
289 "The offset {offset} is at an invalid utf-8 boundary."
290 )));
291 }
292 }
293 builder.append_value(&value[new_start..new_end]);
294 }
295
296 Ok(Arc::new(builder.finish()))
297}
298
299fn binary_view_substring(
300 array: &BinaryViewArray,
301 start: i64,
302 length: Option<u64>,
303) -> Result<ArrayRef, ArrowError> {
304 let mut builder = BinaryViewBuilder::with_capacity(array.len());
305
306 for idx in 0..array.len() {
307 if array.is_null(idx) {
308 builder.append_null();
309 continue;
310 }
311 let value = array.value(idx);
312 let (new_start, new_end) = view_substring_range(value.len(), start, length);
313 builder.append_value(&value[new_start..new_end]);
314 }
315
316 Ok(Arc::new(builder.finish()))
317}
318
319fn byte_substring<T: ByteArrayType>(
320 array: &GenericByteArray<T>,
321 start: T::Offset,
322 length: Option<T::Offset>,
323) -> Result<ArrayRef, ArrowError>
324where
325 <T as ByteArrayType>::Native: PartialEq,
326{
327 let offsets = array.value_offsets();
328 let data = array.value_data();
329 let zero = <T::Offset as Zero>::zero();
330
331 let check_char_boundary = {
333 |offset: T::Offset| {
334 if !matches!(T::DATA_TYPE, DataType::Utf8 | DataType::LargeUtf8) {
335 return Ok(offset);
336 }
337 let data_str = unsafe { std::str::from_utf8_unchecked(data) };
339 let offset_usize = offset.as_usize();
340 if data_str.is_char_boundary(offset_usize) {
341 Ok(offset)
342 } else {
343 Err(ArrowError::ComputeError(format!(
344 "The offset {offset_usize} is at an invalid utf-8 boundary."
345 )))
346 }
347 }
348 };
349
350 let mut new_starts_ends: Vec<(T::Offset, T::Offset)> = Vec::with_capacity(array.len());
352 let mut new_offsets: Vec<T::Offset> = Vec::with_capacity(array.len() + 1);
353 let mut len_so_far = zero;
354 new_offsets.push(zero);
355
356 offsets
357 .windows(2)
358 .try_for_each(|pair| -> Result<(), ArrowError> {
359 let new_start = match start.cmp(&zero) {
360 Ordering::Greater => check_char_boundary((pair[0] + start).min(pair[1]))?,
361 Ordering::Equal => pair[0],
362 Ordering::Less => check_char_boundary((pair[1] + start).max(pair[0]))?,
363 };
364 let new_end = match length {
365 Some(length) => check_char_boundary((length + new_start).min(pair[1]))?,
366 None => pair[1],
367 };
368 len_so_far += new_end - new_start;
369 new_starts_ends.push((new_start, new_end));
370 new_offsets.push(len_so_far);
371 Ok(())
372 })?;
373
374 let mut new_values = MutableBuffer::new(new_offsets.last().unwrap().as_usize());
376
377 new_starts_ends
378 .iter()
379 .map(|(start, end)| {
380 let start = start.as_usize();
381 let end = end.as_usize();
382 &data[start..end]
383 })
384 .try_for_each(|slice| {
385 new_values
386 .try_extend_from_slice(slice)
387 .map_err(|e| ArrowError::MemoryError(e.to_string()))
388 })?;
389
390 let offsets = OffsetBuffer::new(new_offsets.into());
391 let values = new_values.into();
392 let nulls = array
393 .nulls()
394 .map(|n| n.inner().sliced())
395 .and_then(|b| NullBuffer::from_unsliced_buffer(b, array.len()));
396 Ok(Arc::new(GenericByteArray::<T>::new(offsets, values, nulls)))
397}
398
399fn fixed_size_binary_substring(
400 array: &FixedSizeBinaryArray,
401 old_len: usize,
402 start: i64,
403 length: Option<u64>,
404) -> Result<ArrayRef, ArrowError> {
405 let new_start = match start.cmp(&0) {
406 Ordering::Greater => usize::try_from(start).unwrap_or(usize::MAX).min(old_len),
407 Ordering::Equal => 0,
408 Ordering::Less => {
409 let offset = usize::try_from(start.unsigned_abs()).unwrap_or(usize::MAX);
410 old_len.saturating_sub(offset)
411 }
412 };
413
414 let new_len = match length {
415 Some(len) => usize::try_from(len)
416 .unwrap_or(usize::MAX)
417 .min(old_len - new_start),
418 None => old_len - new_start,
419 };
420
421 let num_of_elements = array.len();
423 let data = array.value_data();
424 let capacity = num_of_elements
425 .checked_mul(new_len)
426 .expect("capacity overflow");
427 let mut new_values = MutableBuffer::new(capacity);
428 (0..num_of_elements)
429 .map(|idx| {
430 let offset = idx * array.value_size();
431 (offset + new_start, offset + new_start + new_len)
432 })
433 .try_for_each(|(start, end)| {
434 new_values
435 .try_extend_from_slice(&data[start..end])
436 .map_err(|e| ArrowError::MemoryError(e.to_string()))
437 })?;
438
439 let mut nulls = array
440 .nulls()
441 .map(|n| n.inner().sliced())
442 .and_then(|b| NullBuffer::from_unsliced_buffer(b, num_of_elements));
443
444 if new_len == 0 && nulls.is_none() {
445 nulls = Some(NullBuffer::new_valid(num_of_elements));
450 }
451
452 let new_len: i32 = new_len.try_into().expect("new_len overflow");
453
454 Ok(Arc::new(FixedSizeBinaryArray::new(
455 new_len,
456 new_values.into(),
457 nulls,
458 )))
459}
460
461#[cfg(test)]
462mod tests {
463 use super::*;
464 use arrow_buffer::BooleanBuffer;
465 use arrow_buffer::Buffer;
466
467 macro_rules! gen_test_cases {
478 ($input:expr, $(($start:expr, $len:expr, $result:expr)), *) => {
479 [
480 $(
481 ($input.clone(), $start, $len, $result),
482 )*
483 ]
484 };
485 }
486
487 macro_rules! do_test {
494 ($cases:expr, $array_ty:ty, $substring_fn:ident) => {
495 $cases
496 .into_iter()
497 .for_each(|(array, start, length, expected)| {
498 let array = <$array_ty>::from(array);
499 let result = $substring_fn(&array, start, length).unwrap();
500 let result = result.as_any().downcast_ref::<$array_ty>().unwrap();
501 let expected = <$array_ty>::from(expected);
502 assert_eq!(&expected, result);
503 })
504 };
505 }
506
507 macro_rules! do_test_tryfrom {
509 ($cases:expr, $array_ty:ty, $substring_fn:ident) => {
510 $cases
511 .into_iter()
512 .for_each(|(array, start, length, expected)| {
513 let array = <$array_ty>::try_from(array).unwrap();
514 let result = $substring_fn(&array, start, length).unwrap();
515 let result = result.as_any().downcast_ref::<$array_ty>().unwrap();
516 let expected = <$array_ty>::try_from(expected).unwrap();
517 assert_eq!(&expected, result);
518 })
519 };
520 }
521
522 fn with_nulls_generic_binary<O: OffsetSizeTrait>() {
523 let input = vec![
524 Some("hello".as_bytes()),
525 None,
526 Some(&[0xf8, 0xf9, 0xff, 0xfa]),
527 ];
528 let base_case = gen_test_cases!(
530 vec![None, None, None],
531 (-1, Some(1), vec![None, None, None])
532 );
533 let cases = gen_test_cases!(
534 input,
535 (0, None, input.clone()),
537 (0, Some(0), vec![Some(&[]), None, Some(&[])]),
539 (1000, Some(0), vec![Some(&[]), None, Some(&[])]),
541 (-1000, None, input.clone()),
543 (0, Some(1000), input.clone())
545 );
546
547 do_test!(
548 [&base_case[..], &cases[..]].concat(),
549 GenericBinaryArray<O>,
550 substring
551 );
552 }
553
554 #[test]
555 fn with_nulls_binary() {
556 with_nulls_generic_binary::<i32>()
557 }
558
559 #[test]
560 fn with_nulls_large_binary() {
561 with_nulls_generic_binary::<i64>()
562 }
563
564 fn without_nulls_generic_binary<O: OffsetSizeTrait>() {
565 let input = vec!["hello".as_bytes(), b"", &[0xf8, 0xf9, 0xff, 0xfa]];
566 let base_case = gen_test_cases!(
568 vec!["".as_bytes(), b"", b""],
569 (2, Some(1), vec!["".as_bytes(), b"", b""])
570 );
571 let cases = gen_test_cases!(
572 input,
573 (0, None, input.clone()),
575 (1, None, vec![b"ello", b"", &[0xf9, 0xff, 0xfa]]),
577 (2, None, vec![b"llo", b"", &[0xff, 0xfa]]),
578 (3, None, vec![b"lo", b"", &[0xfa]]),
579 (10, None, vec![b"", b"", b""]),
580 (-1, None, vec![b"o", b"", &[0xfa]]),
582 (-2, None, vec![b"lo", b"", &[0xff, 0xfa]]),
583 (-3, None, vec![b"llo", b"", &[0xf9, 0xff, 0xfa]]),
584 (-10, None, input.clone()),
585 (1, Some(1), vec![b"e", b"", &[0xf9]]),
587 (1, Some(2), vec![b"el", b"", &[0xf9, 0xff]]),
588 (1, Some(3), vec![b"ell", b"", &[0xf9, 0xff, 0xfa]]),
589 (1, Some(4), vec![b"ello", b"", &[0xf9, 0xff, 0xfa]]),
590 (-3, Some(1), vec![b"l", b"", &[0xf9]]),
591 (-3, Some(2), vec![b"ll", b"", &[0xf9, 0xff]]),
592 (-3, Some(3), vec![b"llo", b"", &[0xf9, 0xff, 0xfa]]),
593 (-3, Some(4), vec![b"llo", b"", &[0xf9, 0xff, 0xfa]])
594 );
595
596 do_test!(
597 [&base_case[..], &cases[..]].concat(),
598 GenericBinaryArray<O>,
599 substring
600 );
601 }
602
603 #[test]
604 fn without_nulls_binary() {
605 without_nulls_generic_binary::<i32>()
606 }
607
608 #[test]
609 fn without_nulls_large_binary() {
610 without_nulls_generic_binary::<i64>()
611 }
612
613 fn generic_binary_with_non_zero_offset<O: OffsetSizeTrait>() {
614 let values = 0_u8..15;
615 let offsets = &[
616 O::zero(),
617 O::from_usize(5).unwrap(),
618 O::from_usize(10).unwrap(),
619 O::from_usize(15).unwrap(),
620 ];
621 let bitmap = [0b101_u8];
623
624 let offsets = OffsetBuffer::new(Buffer::from_slice_ref(offsets).into());
625 let values = Buffer::from_iter(values);
626 let nulls = Some(NullBuffer::new(BooleanBuffer::new(
627 Buffer::from(bitmap),
628 0,
629 3,
630 )));
631 let array = GenericBinaryArray::<O>::new(offsets, values, nulls).slice(1, 2);
633 let result = substring(&array, 1, None).unwrap();
635 let result = result
636 .as_any()
637 .downcast_ref::<GenericBinaryArray<O>>()
638 .unwrap();
639 let expected =
640 GenericBinaryArray::<O>::from_opt_vec(vec![None, Some(&[11_u8, 12, 13, 14])]);
641 assert_eq!(result, &expected);
642 }
643
644 #[test]
645 fn binary_with_non_zero_offset() {
646 generic_binary_with_non_zero_offset::<i32>()
647 }
648
649 #[test]
650 fn large_binary_with_non_zero_offset() {
651 generic_binary_with_non_zero_offset::<i64>()
652 }
653
654 #[test]
655 fn with_nulls_fixed_size_binary() {
656 let input = vec![Some("cat".as_bytes()), None, Some(&[0xf8, 0xf9, 0xff])];
657 let base_case =
659 gen_test_cases!(vec![None, None, None], (3, Some(2), vec![None, None, None]));
660 let cases = gen_test_cases!(
661 input,
662 (0, None, input.clone()),
664 (1, None, vec![Some(b"at"), None, Some(&[0xf9, 0xff])]),
666 (2, None, vec![Some(b"t"), None, Some(&[0xff])]),
667 (3, None, vec![Some(b""), None, Some(b"")]),
668 (10, None, vec![Some(b""), None, Some(b"")]),
669 (-1, None, vec![Some(b"t"), None, Some(&[0xff])]),
671 (-2, None, vec![Some(b"at"), None, Some(&[0xf9, 0xff])]),
672 (-3, None, input.clone()),
673 (-10, None, input.clone()),
674 (1, Some(1), vec![Some(b"a"), None, Some(&[0xf9])]),
676 (1, Some(2), vec![Some(b"at"), None, Some(&[0xf9, 0xff])]),
677 (1, Some(3), vec![Some(b"at"), None, Some(&[0xf9, 0xff])]),
678 (-3, Some(1), vec![Some(b"c"), None, Some(&[0xf8])]),
679 (-3, Some(2), vec![Some(b"ca"), None, Some(&[0xf8, 0xf9])]),
680 (-3, Some(3), input.clone()),
681 (-3, Some(4), input.clone())
682 );
683
684 do_test_tryfrom!(
685 [&base_case[..], &cases[..]].concat(),
686 FixedSizeBinaryArray,
687 substring
688 );
689 }
690
691 #[test]
692 fn without_nulls_fixed_size_binary() {
693 let input = vec!["cat".as_bytes(), b"dog", &[0xf8, 0xf9, 0xff]];
694 let base_case = gen_test_cases!(
696 vec!["".as_bytes(), &[], &[]],
697 (1, Some(2), vec!["".as_bytes(), &[], &[]])
698 );
699 let cases = gen_test_cases!(
700 input,
701 (0, None, input.clone()),
703 (1, None, vec![b"at", b"og", &[0xf9, 0xff]]),
705 (2, None, vec![b"t", b"g", &[0xff]]),
706 (3, None, vec![&[], &[], &[]]),
707 (10, None, vec![&[], &[], &[]]),
708 (-1, None, vec![b"t", b"g", &[0xff]]),
710 (-2, None, vec![b"at", b"og", &[0xf9, 0xff]]),
711 (-3, None, input.clone()),
712 (-10, None, input.clone()),
713 (1, Some(1), vec![b"a", b"o", &[0xf9]]),
715 (1, Some(2), vec![b"at", b"og", &[0xf9, 0xff]]),
716 (1, Some(3), vec![b"at", b"og", &[0xf9, 0xff]]),
717 (-3, Some(1), vec![b"c", b"d", &[0xf8]]),
718 (-3, Some(2), vec![b"ca", b"do", &[0xf8, 0xf9]]),
719 (-3, Some(3), input.clone()),
720 (-3, Some(4), input.clone())
721 );
722
723 do_test_tryfrom!(
724 [&base_case[..], &cases[..]].concat(),
725 FixedSizeBinaryArray,
726 substring
727 );
728 }
729
730 #[test]
731 fn fixed_size_binary_with_non_zero_offset() {
732 let values = b"hellotherearrow";
733 let bits_v = [0b101_u8];
735
736 let nulls = Some(NullBuffer::new(BooleanBuffer::new(
737 Buffer::from(bits_v),
738 0,
739 3,
740 )));
741 let array = FixedSizeBinaryArray::new(5, Buffer::from(values), nulls).slice(1, 2);
743 let result = substring(&array, 1, None).unwrap();
745 let result = result
746 .as_any()
747 .downcast_ref::<FixedSizeBinaryArray>()
748 .unwrap();
749 let expected = FixedSizeBinaryArray::try_from_sparse_iter_with_size(
750 vec![None, Some(b"rrow")].into_iter(),
751 4,
752 )
753 .unwrap();
754 assert_eq!(result, &expected);
755 }
756
757 fn with_nulls_generic_string<O: OffsetSizeTrait>() {
758 let input = vec![Some("hello"), None, Some("word")];
759 let base_case = gen_test_cases!(vec![None, None, None], (0, None, vec![None, None, None]));
761 let cases = gen_test_cases!(
762 input,
763 (0, None, input.clone()),
765 (0, Some(0), vec![Some(""), None, Some("")]),
767 (1000, Some(0), vec![Some(""), None, Some("")]),
769 (-1000, None, input.clone()),
771 (0, Some(1000), input.clone())
773 );
774
775 do_test!(
776 [&base_case[..], &cases[..]].concat(),
777 GenericStringArray<O>,
778 substring
779 );
780 }
781
782 #[test]
783 fn with_nulls_string() {
784 with_nulls_generic_string::<i32>()
785 }
786
787 #[test]
788 fn with_nulls_large_string() {
789 with_nulls_generic_string::<i64>()
790 }
791
792 fn without_nulls_generic_string<O: OffsetSizeTrait>() {
793 let input = vec!["hello", "", "word"];
794 let base_case = gen_test_cases!(vec!["", "", ""], (0, None, vec!["", "", ""]));
796 let cases = gen_test_cases!(
797 input,
798 (0, None, input.clone()),
800 (1, None, vec!["ello", "", "ord"]),
801 (2, None, vec!["llo", "", "rd"]),
802 (3, None, vec!["lo", "", "d"]),
803 (10, None, vec!["", "", ""]),
804 (-1, None, vec!["o", "", "d"]),
806 (-2, None, vec!["lo", "", "rd"]),
807 (-3, None, vec!["llo", "", "ord"]),
808 (-10, None, input.clone()),
809 (1, Some(1), vec!["e", "", "o"]),
811 (1, Some(2), vec!["el", "", "or"]),
812 (1, Some(3), vec!["ell", "", "ord"]),
813 (1, Some(4), vec!["ello", "", "ord"]),
814 (-3, Some(1), vec!["l", "", "o"]),
815 (-3, Some(2), vec!["ll", "", "or"]),
816 (-3, Some(3), vec!["llo", "", "ord"]),
817 (-3, Some(4), vec!["llo", "", "ord"])
818 );
819
820 do_test!(
821 [&base_case[..], &cases[..]].concat(),
822 GenericStringArray<O>,
823 substring
824 );
825 }
826
827 #[test]
828 fn without_nulls_string() {
829 without_nulls_generic_string::<i32>()
830 }
831
832 #[test]
833 fn without_nulls_large_string() {
834 without_nulls_generic_string::<i64>()
835 }
836
837 fn generic_string_with_non_zero_offset<O: OffsetSizeTrait>() {
838 let values = b"hellotherearrow";
839 let offsets = &[
840 O::zero(),
841 O::from_usize(5).unwrap(),
842 O::from_usize(10).unwrap(),
843 O::from_usize(15).unwrap(),
844 ];
845 let bitmap = [0b101_u8];
847
848 let offsets = OffsetBuffer::new(Buffer::from_slice_ref(offsets).into());
849 let values = Buffer::from(values);
850 let nulls = Some(NullBuffer::new(BooleanBuffer::new(
851 Buffer::from(bitmap),
852 0,
853 3,
854 )));
855 let array = GenericStringArray::<O>::new(offsets, values, nulls).slice(1, 2);
857 let result = substring(&array, 1, None).unwrap();
859 let result = result
860 .as_any()
861 .downcast_ref::<GenericStringArray<O>>()
862 .unwrap();
863 let expected = GenericStringArray::<O>::from(vec![None, Some("rrow")]);
864 assert_eq!(result, &expected);
865 }
866
867 #[test]
868 fn string_with_non_zero_offset() {
869 generic_string_with_non_zero_offset::<i32>()
870 }
871
872 #[test]
873 fn large_string_with_non_zero_offset() {
874 generic_string_with_non_zero_offset::<i64>()
875 }
876
877 fn with_nulls_generic_string_by_char<O: OffsetSizeTrait>() {
878 let input = vec![Some("hello"), None, Some("Γ ⊢x:T")];
879 let base_case = gen_test_cases!(vec![None, None, None], (0, None, vec![None, None, None]));
881 let cases = gen_test_cases!(
882 input,
883 (0, None, input.clone()),
885 (0, Some(0), vec![Some(""), None, Some("")]),
887 (1000, Some(0), vec![Some(""), None, Some("")]),
889 (-1000, None, input.clone()),
891 (0, Some(1000), input.clone())
893 );
894
895 do_test!(
896 [&base_case[..], &cases[..]].concat(),
897 GenericStringArray<O>,
898 substring_by_char
899 );
900 }
901
902 #[test]
903 fn with_nulls_string_by_char() {
904 with_nulls_generic_string_by_char::<i32>()
905 }
906
907 #[test]
908 fn with_nulls_large_string_by_char() {
909 with_nulls_generic_string_by_char::<i64>()
910 }
911
912 fn without_nulls_generic_string_by_char<O: OffsetSizeTrait>() {
913 let input = vec!["hello", "", "Γ ⊢x:T"];
914 let base_case = gen_test_cases!(vec!["", "", ""], (0, None, vec!["", "", ""]));
916 let cases = gen_test_cases!(
917 input,
918 (0, None, input.clone()),
920 (1, None, vec!["ello", "", " ⊢x:T"]),
922 (2, None, vec!["llo", "", "⊢x:T"]),
923 (3, None, vec!["lo", "", "x:T"]),
924 (10, None, vec!["", "", ""]),
925 (-1, None, vec!["o", "", "T"]),
927 (-2, None, vec!["lo", "", ":T"]),
928 (-4, None, vec!["ello", "", "⊢x:T"]),
929 (-10, None, input.clone()),
930 (1, Some(1), vec!["e", "", " "]),
932 (1, Some(2), vec!["el", "", " ⊢"]),
933 (1, Some(3), vec!["ell", "", " ⊢x"]),
934 (1, Some(6), vec!["ello", "", " ⊢x:T"]),
935 (-4, Some(1), vec!["e", "", "⊢"]),
936 (-4, Some(2), vec!["el", "", "⊢x"]),
937 (-4, Some(3), vec!["ell", "", "⊢x:"]),
938 (-4, Some(4), vec!["ello", "", "⊢x:T"])
939 );
940
941 do_test!(
942 [&base_case[..], &cases[..]].concat(),
943 GenericStringArray<O>,
944 substring_by_char
945 );
946 }
947
948 #[test]
949 fn without_nulls_string_by_char() {
950 without_nulls_generic_string_by_char::<i32>()
951 }
952
953 #[test]
954 fn without_nulls_large_string_by_char() {
955 without_nulls_generic_string_by_char::<i64>()
956 }
957
958 fn ascii_generic_string_by_char<O: OffsetSizeTrait>() {
960 let input = vec![Some("hello"), None, Some(""), Some("rust")];
961 let cases = gen_test_cases!(
962 input,
963 (0, None, input.clone()),
965 (1, None, vec![Some("ello"), None, Some(""), Some("ust")]),
967 (4, None, vec![Some("o"), None, Some(""), Some("")]),
968 (1000, None, vec![Some(""), None, Some(""), Some("")]),
970 (-1, None, vec![Some("o"), None, Some(""), Some("t")]),
972 (-4, None, vec![Some("ello"), None, Some(""), Some("rust")]),
973 (-1000, None, input.clone()),
975 (0, Some(0), vec![Some(""), None, Some(""), Some("")]),
977 (1, Some(2), vec![Some("el"), None, Some(""), Some("us")]),
979 (-4, Some(2), vec![Some("el"), None, Some(""), Some("ru")]),
980 (0, Some(1000), input.clone()),
982 (
984 1,
985 Some(u64::MAX),
986 vec![Some("ello"), None, Some(""), Some("ust")]
987 )
988 );
989
990 do_test!(cases, GenericStringArray<O>, substring_by_char);
991 }
992
993 #[test]
994 fn ascii_string_by_char() {
995 ascii_generic_string_by_char::<i32>()
996 }
997
998 #[test]
999 fn ascii_large_string_by_char() {
1000 ascii_generic_string_by_char::<i64>()
1001 }
1002
1003 fn generic_string_by_char_with_non_zero_offset<O: OffsetSizeTrait>() {
1004 let values = "S→T = Πx:S.T";
1005 let offsets = &[
1006 O::zero(),
1007 O::from_usize(values.char_indices().nth(3).map(|(pos, _)| pos).unwrap()).unwrap(),
1008 O::from_usize(values.char_indices().nth(6).map(|(pos, _)| pos).unwrap()).unwrap(),
1009 O::from_usize(values.len()).unwrap(),
1010 ];
1011 let bitmap = [0b101_u8];
1013
1014 let offsets = OffsetBuffer::new(Buffer::from_slice_ref(offsets).into());
1015 let values = Buffer::from(values.as_bytes());
1016 let nulls = Some(NullBuffer::new(BooleanBuffer::new(
1017 Buffer::from(bitmap),
1018 0,
1019 3,
1020 )));
1021 let array = GenericStringArray::<O>::new(offsets, values, nulls).slice(1, 2);
1023 let result = substring_by_char(&array, 1, None).unwrap();
1025 let expected = GenericStringArray::<O>::from(vec![None, Some("x:S.T")]);
1026 assert_eq!(result, expected);
1027 }
1028
1029 #[test]
1030 fn string_with_non_zero_offset_by_char() {
1031 generic_string_by_char_with_non_zero_offset::<i32>()
1032 }
1033
1034 #[test]
1035 fn large_string_with_non_zero_offset_by_char() {
1036 generic_string_by_char_with_non_zero_offset::<i64>()
1037 }
1038
1039 #[test]
1040 fn dictionary() {
1041 _dictionary::<Int8Type>();
1042 _dictionary::<Int16Type>();
1043 _dictionary::<Int32Type>();
1044 _dictionary::<Int64Type>();
1045 _dictionary::<UInt8Type>();
1046 _dictionary::<UInt16Type>();
1047 _dictionary::<UInt32Type>();
1048 _dictionary::<UInt64Type>();
1049 }
1050
1051 fn _dictionary<K: ArrowDictionaryKeyType>() {
1052 const TOTAL: i32 = 100;
1053
1054 let v = ["aaa", "bbb", "ccc", "ddd", "eee"];
1055 let data: Vec<Option<&str>> = (0..TOTAL)
1056 .map(|n| {
1057 let i = n % 5;
1058 if i == 3 { None } else { Some(v[i as usize]) }
1059 })
1060 .collect();
1061
1062 let dict_array: DictionaryArray<K> = data.clone().into_iter().collect();
1063
1064 let expected: Vec<Option<&str>> = data.iter().map(|opt| opt.map(|s| &s[1..3])).collect();
1065
1066 let res = substring(&dict_array, 1, Some(2)).unwrap();
1067 let actual = res.as_any().downcast_ref::<DictionaryArray<K>>().unwrap();
1068 let actual: Vec<Option<&str>> = actual
1069 .values()
1070 .as_any()
1071 .downcast_ref::<GenericStringArray<i32>>()
1072 .unwrap()
1073 .take_iter(actual.keys_iter())
1074 .collect();
1075
1076 for i in 0..TOTAL as usize {
1077 assert_eq!(expected[i], actual[i],);
1078 }
1079 }
1080
1081 #[test]
1082 fn check_invalid_array_type() {
1083 let array = Int32Array::from(vec![Some(1), Some(2), Some(3)]);
1084 let err = substring(&array, 0, None).unwrap_err().to_string();
1085 assert!(err.contains("substring does not support type"));
1086 }
1087
1088 #[test]
1090 fn check_start_index() {
1091 let array = StringArray::from(vec![Some("E=mc²"), Some("ascii")]);
1092 let err = substring(&array, -1, None).unwrap_err().to_string();
1093 assert!(err.contains("invalid utf-8 boundary"));
1094 }
1095
1096 #[test]
1097 fn check_length() {
1098 let array = StringArray::from(vec![Some("E=mc²"), Some("ascii")]);
1099 let err = substring(&array, 0, Some(5)).unwrap_err().to_string();
1100 assert!(err.contains("invalid utf-8 boundary"));
1101 }
1102
1103 #[test]
1104 fn non_utf8_bytes() {
1105 let bytes: &[u8] = &[0xE4, 0xBD, 0xA0, 0xE5, 0xA5, 0xBD, 0xE8, 0xAF, 0xAD];
1107 let array = BinaryArray::from(vec![Some(bytes)]);
1108 let arr = substring(&array, 0, Some(5)).unwrap();
1109 let actual = arr.as_any().downcast_ref::<BinaryArray>().unwrap();
1110
1111 let expected_bytes: &[u8] = &[0xE4, 0xBD, 0xA0, 0xE5, 0xA5];
1112 let expected = BinaryArray::from(vec![Some(expected_bytes)]);
1113 assert_eq!(expected, *actual);
1114 }
1115
1116 #[test]
1117 fn string_view_matches_utf8() {
1118 let values = vec![
1119 Some("hello world"),
1120 Some(""),
1121 None,
1122 Some("a"),
1123 Some("this one is definitely longer than twelve bytes"),
1124 ];
1125 let utf8 = StringArray::from(values.clone());
1126 let view = StringViewArray::from(values);
1127
1128 for (start, length) in [
1129 (0, None),
1130 (0, Some(0)),
1131 (0, Some(5)),
1132 (0, Some(1000)),
1133 (1, Some(3)),
1134 (5, None),
1135 (100, Some(2)),
1136 (100, None),
1137 (-3, None),
1138 (-3, Some(2)),
1139 (-100, Some(4)),
1140 (-100, None),
1141 ] {
1142 let expected = substring(&utf8, start, length).unwrap();
1143 let expected = expected.as_string::<i32>();
1144 let actual = substring(&view, start, length).unwrap();
1145 let actual = actual.as_string_view();
1146 assert_eq!(
1147 expected.iter().collect::<Vec<_>>(),
1148 actual.iter().collect::<Vec<_>>(),
1149 "start={start} length={length:?}"
1150 );
1151 }
1152 }
1153
1154 #[test]
1155 fn binary_view_matches_binary() {
1156 let values: Vec<Option<&[u8]>> = vec![
1157 Some(b"hello world"),
1158 Some(b""),
1159 None,
1160 Some(b"abc"),
1161 Some(b"this one is definitely longer than twelve bytes"),
1162 ];
1163 let binary = BinaryArray::from(values.clone());
1164 let view = BinaryViewArray::from(values);
1165
1166 for (start, length) in [
1167 (0, None),
1168 (0, Some(5)),
1169 (2, Some(3)),
1170 (-3, None),
1171 (100, Some(2)),
1172 ] {
1173 let expected = substring(&binary, start, length).unwrap();
1174 let expected = expected.as_binary::<i32>();
1175 let actual = substring(&view, start, length).unwrap();
1176 let actual = actual.as_binary_view();
1177 assert_eq!(
1178 expected.iter().collect::<Vec<_>>(),
1179 actual.iter().collect::<Vec<_>>(),
1180 "start={start} length={length:?}"
1181 );
1182 }
1183 }
1184
1185 #[test]
1186 fn string_view_slices_on_a_multi_byte_boundary() {
1187 let values = vec![Some("héllo wörld"), Some("日本語"), None];
1190 let utf8 = StringArray::from(values.clone());
1191 let view = StringViewArray::from(values);
1192
1193 for (start, length) in [
1194 (0, Some(3)),
1195 (3, Some(3)),
1196 (0, Some(6)),
1197 (3, None),
1198 (-3, None),
1199 (-6, None),
1200 ] {
1201 let expected = substring(&utf8, start, length).unwrap();
1202 let expected = expected.as_string::<i32>();
1203 let actual = substring(&view, start, length).unwrap();
1204 let actual = actual.as_string_view();
1205 assert_eq!(
1206 expected.iter().collect::<Vec<_>>(),
1207 actual.iter().collect::<Vec<_>>(),
1208 "start={start} length={length:?}"
1209 );
1210 }
1211
1212 let actual = substring(&view, 0, Some(3)).unwrap();
1213 let actual = actual.as_string_view();
1214 assert_eq!(actual.value(0), "hé");
1215 assert_eq!(actual.value(1), "日");
1216 }
1217
1218 #[test]
1219 fn string_view_rejects_an_invalid_char_boundary() {
1220 let view = StringViewArray::from(vec![Some("E=mc²")]);
1221 let err = substring(&view, 0, Some(5)).unwrap_err().to_string();
1222 assert!(err.contains("invalid utf-8 boundary"), "{err}");
1223 }
1224
1225 #[test]
1226 fn dictionary_of_string_view() {
1227 let view = StringViewArray::from(vec![Some("hello world"), Some("bye")]);
1228 let dict = DictionaryArray::new(Int32Array::from(vec![0, 1, 0]), Arc::new(view));
1229
1230 let actual = substring(&dict, 0, Some(3)).unwrap();
1231 let actual = actual.as_any_dictionary();
1232
1233 let values = actual.values().as_string_view();
1234 assert_eq!(
1235 values.iter().collect::<Vec<_>>(),
1236 vec![Some("hel"), Some("bye")]
1237 );
1238 }
1239}