1use std::fmt::Display;
5use std::fmt::Formatter;
6
7use num_traits::AsPrimitive;
8use vortex_array::arrays::PrimitiveArray;
9use vortex_buffer::BufferAllocatorRef;
10use vortex_buffer::ByteBuffer;
11use vortex_error::VortexExpect;
12use vortex_error::VortexResult;
13use vortex_error::vortex_bail;
14use vortex_error::vortex_ensure;
15use vortex_error::vortex_ensure_eq;
16use vortex_error::vortex_err;
17
18use crate::ArrayRef;
19use crate::ArraySlots;
20use crate::ExecutionCtx;
21use crate::aggregate_fn::NumericalAggregateOpts;
22use crate::aggregate_fn::fns::min_max::min_max;
23use crate::array::Array;
24use crate::array::ArrayParts;
25use crate::array::TypedArrayRef;
26use crate::array::child_to_validity;
27use crate::array::validity_to_child;
28use crate::array_slots;
29use crate::arrays::VarBin;
30use crate::arrays::varbin::builder::VarBinBuilder;
31use crate::buffer::BufferHandle;
32use crate::dtype::DType;
33use crate::dtype::Nullability;
34use crate::dtype::OffsetBuilderPType;
35use crate::match_each_integer_ptype;
36use crate::validity::Validity;
37
38fn is_char_boundary_at(bytes: &[u8], index: usize) -> bool {
40 index == bytes.len() || bytes[index] & 0b1100_0000 != 0b1000_0000
42}
43
44pub(crate) fn offsets_tile_utf8<O: AsPrimitive<usize> + PartialOrd>(
51 offsets: &[O],
52 bytes: &[u8],
53) -> bool {
54 let (Some(first), Some(last)) = (offsets.first(), offsets.last()) else {
55 return false;
56 };
57 let (first, last): (usize, usize) = (first.as_(), last.as_());
58
59 first <= last
60 && last <= bytes.len()
61 && offsets.windows(2).all(|pair| pair[0] <= pair[1])
62 && simdutf8::basic::from_utf8(&bytes[first..last]).is_ok()
63 && offsets
64 .iter()
65 .all(|&offset| is_char_boundary_at(bytes, offset.as_()))
66}
67
68#[array_slots(VarBin)]
69pub struct VarBinSlots {
70 #[slot(0)]
72 pub offsets: ArrayRef,
73 #[slot(1)]
75 pub validity: Option<ArrayRef>,
76}
77
78#[derive(Clone, Debug)]
79pub struct VarBinData {
80 pub(super) bytes: BufferHandle,
81}
82
83impl Display for VarBinData {
84 fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result {
85 Ok(())
86 }
87}
88
89pub struct VarBinDataParts {
90 pub dtype: DType,
91 pub bytes: BufferHandle,
92 pub offsets: ArrayRef,
93 pub validity: Validity,
94}
95
96impl VarBinData {
97 pub fn build(
104 offsets: ArrayRef,
105 bytes: ByteBuffer,
106 dtype: DType,
107 validity: Validity,
108 ctx: &mut ExecutionCtx,
109 ) -> Self {
110 Self::try_build(offsets, bytes, dtype, validity, ctx).vortex_expect("VarBinArray new")
111 }
112
113 pub fn build_from_handle(
120 offset: ArrayRef,
121 bytes: BufferHandle,
122 dtype: DType,
123 validity: Validity,
124 ctx: &mut ExecutionCtx,
125 ) -> Self {
126 Self::try_build_from_handle(offset, bytes, dtype, validity, ctx)
127 .vortex_expect("VarBinArray new")
128 }
129
130 pub(crate) fn make_slots(offsets: ArrayRef, validity: &Validity, len: usize) -> ArraySlots {
131 VarBinSlots {
132 offsets,
133 validity: validity_to_child(validity, len),
134 }
135 .into_slots()
136 }
137
138 pub fn try_build(
147 offsets: ArrayRef,
148 bytes: ByteBuffer,
149 dtype: DType,
150 validity: Validity,
151 ctx: &mut ExecutionCtx,
152 ) -> VortexResult<Self> {
153 let bytes = BufferHandle::new_host(bytes);
154 Self::validate(&offsets, &bytes, &dtype, &validity, ctx)?;
155
156 Ok(unsafe { Self::new_unchecked_from_handle(bytes) })
158 }
159
160 pub fn try_build_from_handle(
170 offsets: ArrayRef,
171 bytes: BufferHandle,
172 dtype: DType,
173 validity: Validity,
174 ctx: &mut ExecutionCtx,
175 ) -> VortexResult<Self> {
176 Self::validate(&offsets, &bytes, &dtype, &validity, ctx)?;
177
178 Ok(unsafe { Self::new_unchecked_from_handle(bytes) })
180 }
181
182 pub unsafe fn new_unchecked(bytes: ByteBuffer) -> Self {
211 unsafe { Self::new_unchecked_from_handle(BufferHandle::new_host(bytes)) }
214 }
215
216 pub unsafe fn new_unchecked_from_handle(bytes: BufferHandle) -> Self {
223 Self { bytes }
224 }
225
226 pub fn validate(
230 offsets: &ArrayRef,
231 bytes: &BufferHandle,
232 dtype: &DType,
233 validity: &Validity,
234 ctx: &mut ExecutionCtx,
235 ) -> VortexResult<()> {
236 vortex_ensure!(
238 offsets.dtype().is_int() && !offsets.dtype().is_nullable(),
239 MismatchedTypes: "non nullable int", offsets.dtype()
240 );
241
242 vortex_ensure!(
244 matches!(dtype, DType::Binary(_) | DType::Utf8(_)),
245 MismatchedTypes: "utf8 or binary", dtype
246 );
247
248 vortex_ensure!(
250 dtype.is_nullable() != matches!(validity, Validity::NonNullable),
251 InvalidArgument: "incorrect validity {:?} for dtype {}",
252 validity,
253 dtype
254 );
255
256 vortex_ensure!(
258 !offsets.is_empty(),
259 InvalidArgument: "Offsets must have at least one element"
260 );
261
262 if let Some(validity_len) = validity.maybe_len() {
264 vortex_ensure_eq!(
265 validity_len,
266 offsets.len() - 1,
267 "Validity length doesn't match array length",
268 );
269 }
270
271 if offsets.is_host()
274 && bytes.is_on_host()
275 && matches!(dtype, DType::Utf8(_))
276 && let Some(bytes) = bytes.as_host_opt()
277 {
278 Self::validate_utf8(offsets, bytes.as_ref(), validity, ctx)?;
279 } else if offsets.is_host() {
280 Self::validate_offsets(offsets, bytes.len(), ctx)?;
281 }
282
283 Ok(())
284 }
285
286 fn validate_offsets(
288 offsets: &ArrayRef,
289 bytes_len: usize,
290 ctx: &mut ExecutionCtx,
291 ) -> VortexResult<()> {
292 let Some(is_sorted) = offsets.statistics().compute_is_sorted(ctx) else {
294 vortex_bail!(InvalidArgument: "offsets must report is_sorted statistic");
295 };
296 vortex_ensure!(is_sorted, InvalidArgument: "offsets must be sorted");
297
298 let Some(min_max) = min_max(offsets, ctx, NumericalAggregateOpts::default())? else {
301 vortex_bail!(
302 InvalidArgument: "offsets array with encoding {} must support min_max compute function",
303 offsets.encoding_id()
304 );
305 };
306
307 match_each_integer_ptype!(offsets.dtype().as_ptype(), |P| {
308 #[allow(clippy::absurd_extreme_comparisons, unused_comparisons)]
309 {
310 let max = min_max
311 .max
312 .as_primitive()
313 .as_::<P>()
314 .vortex_expect("offsets type must fit offsets values");
315 let min = min_max
316 .min
317 .as_primitive()
318 .as_::<P>()
319 .vortex_expect("offsets type must fit offsets values");
320
321 vortex_ensure!(
322 min >= 0,
323 InvalidArgument: "offsets minimum {min} outside valid range [0, {max}]"
324 );
325
326 vortex_ensure!(
328 P::try_from(bytes_len).ok().is_none_or(|len| max <= len),
329 InvalidArgument: "Max offset {max} is beyond the length of the bytes buffer {bytes_len}"
330 );
331 }
332 });
333
334 Ok(())
335 }
336
337 fn validate_utf8(
339 offsets: &ArrayRef,
340 bytes: &[u8],
341 validity: &Validity,
342 ctx: &mut ExecutionCtx,
343 ) -> VortexResult<()> {
344 let validate_at = |i: usize, start: usize, end: usize| -> VortexResult<()> {
345 let string_bytes = bytes.get(start..end).ok_or_else(|| {
346 vortex_err!(
347 InvalidArgument: "offsets {start}..{end} at index {i} are out of order or out of bounds for bytes of length {}",
348 bytes.len()
349 )
350 })?;
351 simdutf8::basic::from_utf8(string_bytes).map_err(|_| {
352 #[expect(clippy::unwrap_used)]
353 let err = simdutf8::compat::from_utf8(string_bytes).unwrap_err();
355 vortex_err!("invalid utf-8: {err} at index {i}")
356 })?;
357 Ok(())
358 };
359
360 let primitive_offsets = offsets.clone().execute::<PrimitiveArray>(ctx)?;
362
363 let mask = match validity {
367 Validity::Array(_) => {
368 Some(validity.execute_mask(primitive_offsets.len().saturating_sub(1), ctx)?)
369 }
370 _ => None,
371 };
372 let all_invalid = validity.definitely_all_null();
373
374 match_each_integer_ptype!(primitive_offsets.dtype().as_ptype(), |O| {
375 let offsets_slice = primitive_offsets.as_slice::<O>();
376
377 let last_offset: usize = offsets_slice[offsets_slice.len() - 1].as_();
378 vortex_ensure!(
379 last_offset <= bytes.len(),
380 InvalidArgument: "Last offset {} exceeds bytes length {}",
381 last_offset,
382 bytes.len()
383 );
384
385 if offsets_tile_utf8(offsets_slice, bytes) {
386 return Ok(());
387 }
388
389 Self::validate_offsets(offsets, bytes.len(), ctx)?;
392
393 for (i, (start, end)) in offsets_slice
395 .windows(2)
396 .map(|o| (o[0].as_(), o[1].as_()))
397 .enumerate()
398 {
399 let valid = mask.as_ref().map_or(!all_invalid, |mask| mask.value(i));
400 if valid {
401 validate_at(i, start, end)?;
402 }
403 }
404 });
405 Ok(())
406 }
407
408 #[inline]
416 pub fn bytes(&self) -> &ByteBuffer {
417 self.bytes.as_host()
418 }
419
420 #[inline]
422 pub fn bytes_handle(&self) -> &BufferHandle {
423 &self.bytes
424 }
425}
426
427pub trait VarBinArrayExt: VarBinArraySlotsExt {
428 fn dtype_parts(&self) -> (bool, Nullability) {
429 match self.dtype() {
430 DType::Utf8(nullability) => (true, *nullability),
431 DType::Binary(nullability) => (false, *nullability),
432 _ => unreachable!("VarBinArrayExt requires a utf8 or binary dtype"),
433 }
434 }
435
436 fn is_utf8(&self) -> bool {
437 self.dtype_parts().0
438 }
439
440 fn nullability(&self) -> Nullability {
441 self.dtype_parts().1
442 }
443
444 fn varbin_validity(&self) -> Validity {
445 child_to_validity(
446 self.slots()[VarBinSlots::VALIDITY].as_ref(),
447 self.nullability(),
448 )
449 }
450
451 fn offset_at(&self, index: usize, ctx: &mut ExecutionCtx) -> usize {
452 assert!(
453 index <= self.len(),
454 "Index {index} out of bounds 0..={}",
455 self.len()
456 );
457
458 (&self
459 .offsets()
460 .execute_scalar(index, ctx)
461 .vortex_expect("offsets must support execute_scalar"))
462 .try_into()
463 .vortex_expect("Failed to convert offset to usize")
464 }
465
466 fn bytes_at(&self, index: usize, ctx: &mut ExecutionCtx) -> ByteBuffer {
467 let start = self.offset_at(index, ctx);
468 let end = self.offset_at(index + 1, ctx);
469 self.bytes().slice(start..end)
470 }
471
472 fn sliced_bytes(&self, ctx: &mut ExecutionCtx) -> ByteBuffer {
473 let first_offset: usize = self.offset_at(0, ctx);
474 let last_offset = self.offset_at(self.len(), ctx);
475 self.bytes().slice(first_offset..last_offset)
476 }
477}
478impl<T: TypedArrayRef<VarBin>> VarBinArrayExt for T {}
479
480impl Array<VarBin> {
482 pub fn from_vec<T: AsRef<[u8]>>(vec: Vec<T>, dtype: DType) -> Self {
483 let size: usize = vec.iter().map(|v| v.as_ref().len()).sum();
484 if size < u32::MAX as usize {
485 Self::from_vec_sized::<u32, T>(vec, dtype)
486 } else {
487 Self::from_vec_sized::<u64, T>(vec, dtype)
488 }
489 }
490
491 #[expect(
492 clippy::same_name_method,
493 reason = "intentionally named from_iter like Iterator::from_iter"
494 )]
495 pub fn from_iter<T: AsRef<[u8]>, I: IntoIterator<Item = Option<T>>>(
496 iter: I,
497 dtype: DType,
498 ) -> Self {
499 let iter = iter.into_iter();
500 let mut builder = VarBinBuilder::<u32>::with_capacity_in(
501 dtype,
502 iter.size_hint().0,
503 BufferAllocatorRef::static_ref(),
504 );
505 for v in iter {
506 builder.append(v.as_ref().map(|o| o.as_ref()));
507 }
508 builder.finish_into_varbin()
509 }
510
511 pub fn from_iter_nonnull<T: AsRef<[u8]>, I: IntoIterator<Item = T>>(
512 iter: I,
513 dtype: DType,
514 ) -> Self {
515 let iter = iter.into_iter();
516 let mut builder = VarBinBuilder::<u32>::with_capacity_in(
517 dtype,
518 iter.size_hint().0,
519 BufferAllocatorRef::static_ref(),
520 );
521 for v in iter {
522 builder.append_value(v);
523 }
524 builder.finish_into_varbin()
525 }
526
527 fn from_vec_sized<O, T>(vec: Vec<T>, dtype: DType) -> Self
528 where
529 O: OffsetBuilderPType,
530 T: AsRef<[u8]>,
531 {
532 let mut builder = VarBinBuilder::<O>::with_capacity_in(
533 dtype,
534 vec.len(),
535 BufferAllocatorRef::static_ref(),
536 );
537 for v in vec {
538 builder.append_value(v.as_ref());
539 }
540 builder.finish_into_varbin()
541 }
542
543 pub fn from_strs(value: Vec<&str>) -> Self {
545 Self::from_vec(value, DType::Utf8(Nullability::NonNullable))
546 }
547
548 pub fn from_nullable_strs(value: Vec<Option<&str>>) -> Self {
550 Self::from_iter(value, DType::Utf8(Nullability::Nullable))
551 }
552
553 pub fn from_bytes(value: Vec<&[u8]>) -> Self {
555 Self::from_vec(value, DType::Binary(Nullability::NonNullable))
556 }
557
558 pub fn from_nullable_bytes(value: Vec<Option<&[u8]>>) -> Self {
560 Self::from_iter(value, DType::Binary(Nullability::Nullable))
561 }
562
563 pub fn into_data_parts(self) -> VarBinDataParts {
564 let dtype = self.dtype().clone();
565 let validity = self.varbin_validity();
566 let offsets = self.offsets().clone();
567 let data = self.into_data();
568 VarBinDataParts {
569 dtype,
570 bytes: data.bytes,
571 offsets,
572 validity,
573 }
574 }
575}
576
577impl Array<VarBin> {
578 pub fn new(
580 offsets: ArrayRef,
581 bytes: ByteBuffer,
582 dtype: DType,
583 validity: Validity,
584 ctx: &mut ExecutionCtx,
585 ) -> Self {
586 let len = offsets.len().saturating_sub(1);
587 let slots = VarBinData::make_slots(offsets, &validity, len);
588 let data = VarBinData::build(
589 slots[VarBinSlots::OFFSETS]
590 .as_ref()
591 .vortex_expect("VarBinArray offsets slot")
592 .clone(),
593 bytes,
594 dtype.clone(),
595 validity,
596 ctx,
597 );
598 unsafe { Array::from_parts_unchecked(ArrayParts::new(VarBin, dtype, len, data, slots)) }
599 }
600
601 pub unsafe fn new_unchecked(
607 offsets: ArrayRef,
608 bytes: ByteBuffer,
609 dtype: DType,
610 validity: Validity,
611 ) -> Self {
612 let len = offsets.len().saturating_sub(1);
613 let slots = VarBinData::make_slots(offsets, &validity, len);
614 let data = unsafe { VarBinData::new_unchecked(bytes) };
615 unsafe { Array::from_parts_unchecked(ArrayParts::new(VarBin, dtype, len, data, slots)) }
616 }
617
618 pub unsafe fn new_unchecked_from_handle(
624 offsets: ArrayRef,
625 bytes: BufferHandle,
626 dtype: DType,
627 validity: Validity,
628 ) -> Self {
629 let len = offsets.len().saturating_sub(1);
630 let slots = VarBinData::make_slots(offsets, &validity, len);
631 let data = unsafe { VarBinData::new_unchecked_from_handle(bytes) };
632 unsafe { Array::from_parts_unchecked(ArrayParts::new(VarBin, dtype, len, data, slots)) }
633 }
634
635 pub fn try_new(
638 offsets: ArrayRef,
639 bytes: ByteBuffer,
640 dtype: DType,
641 validity: Validity,
642 ctx: &mut ExecutionCtx,
643 ) -> VortexResult<Self> {
644 let len = offsets.len() - 1;
645 let bytes = BufferHandle::new_host(bytes);
646 VarBinData::validate(&offsets, &bytes, &dtype, &validity, ctx)?;
647 let slots = VarBinData::make_slots(offsets, &validity, len);
648 let data = unsafe { VarBinData::new_unchecked_from_handle(bytes) };
650 Ok(
651 unsafe {
652 Array::from_parts_unchecked(ArrayParts::new(VarBin, dtype, len, data, slots))
653 },
654 )
655 }
656}
657
658impl From<Vec<&[u8]>> for Array<VarBin> {
659 fn from(value: Vec<&[u8]>) -> Self {
660 Self::from_vec(value, DType::Binary(Nullability::NonNullable))
661 }
662}
663
664impl From<Vec<Vec<u8>>> for Array<VarBin> {
665 fn from(value: Vec<Vec<u8>>) -> Self {
666 Self::from_vec(value, DType::Binary(Nullability::NonNullable))
667 }
668}
669
670impl From<Vec<String>> for Array<VarBin> {
671 fn from(value: Vec<String>) -> Self {
672 Self::from_vec(value, DType::Utf8(Nullability::NonNullable))
673 }
674}
675
676impl From<Vec<&str>> for Array<VarBin> {
677 fn from(value: Vec<&str>) -> Self {
678 Self::from_vec(value, DType::Utf8(Nullability::NonNullable))
679 }
680}
681
682impl From<Vec<Option<&[u8]>>> for Array<VarBin> {
683 fn from(value: Vec<Option<&[u8]>>) -> Self {
684 Self::from_iter(value, DType::Binary(Nullability::Nullable))
685 }
686}
687
688impl From<Vec<Option<Vec<u8>>>> for Array<VarBin> {
689 fn from(value: Vec<Option<Vec<u8>>>) -> Self {
690 Self::from_iter(value, DType::Binary(Nullability::Nullable))
691 }
692}
693
694impl From<Vec<Option<String>>> for Array<VarBin> {
695 fn from(value: Vec<Option<String>>) -> Self {
696 Self::from_iter(value, DType::Utf8(Nullability::Nullable))
697 }
698}
699
700impl From<Vec<Option<&str>>> for Array<VarBin> {
701 fn from(value: Vec<Option<&str>>) -> Self {
702 Self::from_iter(value, DType::Utf8(Nullability::Nullable))
703 }
704}
705
706impl<'a> FromIterator<Option<&'a [u8]>> for Array<VarBin> {
707 fn from_iter<T: IntoIterator<Item = Option<&'a [u8]>>>(iter: T) -> Self {
708 Self::from_iter(iter, DType::Binary(Nullability::Nullable))
709 }
710}
711
712impl FromIterator<Option<Vec<u8>>> for Array<VarBin> {
713 fn from_iter<T: IntoIterator<Item = Option<Vec<u8>>>>(iter: T) -> Self {
714 Self::from_iter(iter, DType::Binary(Nullability::Nullable))
715 }
716}
717
718impl FromIterator<Option<String>> for Array<VarBin> {
719 fn from_iter<T: IntoIterator<Item = Option<String>>>(iter: T) -> Self {
720 Self::from_iter(iter, DType::Utf8(Nullability::Nullable))
721 }
722}
723
724impl<'a> FromIterator<Option<&'a str>> for Array<VarBin> {
725 fn from_iter<T: IntoIterator<Item = Option<&'a str>>>(iter: T) -> Self {
726 Self::from_iter(iter, DType::Utf8(Nullability::Nullable))
727 }
728}