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