1use std::fmt::Display;
5use std::fmt::Formatter;
6use std::mem::size_of;
7use std::sync::Arc;
8
9use smallvec::smallvec;
10use vortex_buffer::Alignment;
11use vortex_buffer::Buffer;
12use vortex_buffer::ByteBuffer;
13use vortex_error::VortexExpect;
14use vortex_error::VortexResult;
15use vortex_error::vortex_bail;
16use vortex_error::vortex_ensure;
17use vortex_error::vortex_err;
18use vortex_error::vortex_panic;
19use vortex_mask::AllOr;
20use vortex_mask::Mask;
21
22use crate::ArrayRef;
23use crate::ArraySlots;
24use crate::ExecutionCtx;
25#[cfg(debug_assertions)]
26use crate::VortexSessionExecute;
27use crate::array::Array;
28use crate::array::ArrayParts;
29use crate::array::TypedArrayRef;
30use crate::array::child_to_validity;
31use crate::array::validity_to_child;
32use crate::array_slots;
33use crate::arrays::VarBinView;
34use crate::arrays::varbinview::BinaryView;
35use crate::buffer::BufferHandle;
36use crate::builders::ArrayBuilder;
37use crate::builders::VarBinViewBuilder;
38use crate::dtype::DType;
39use crate::dtype::Nullability;
40#[cfg(debug_assertions)]
41use crate::legacy_session;
42use crate::validity::Validity;
43
44#[array_slots(VarBinView)]
45pub struct VarBinViewSlots {
46 #[slot(0)]
48 pub validity: Option<ArrayRef>,
49}
50
51#[derive(Clone, Debug)]
111pub struct VarBinViewData {
112 pub(super) buffers: Arc<[BufferHandle]>,
113 pub(super) views: BufferHandle,
114}
115
116impl Display for VarBinViewData {
117 fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result {
118 Ok(())
119 }
120}
121
122pub struct VarBinViewDataParts {
123 pub dtype: DType,
124 pub buffers: Arc<[BufferHandle]>,
125 pub views: BufferHandle,
126 pub validity: Validity,
127}
128
129fn for_each_invalid_range(mask: &Mask, len: usize, mut f: impl FnMut(usize, usize)) {
131 match mask.bit_buffer() {
132 AllOr::All => {}
133 AllOr::None => f(0, len),
134 AllOr::Some(buffer) => {
135 let mut prev = 0;
136 for (start, end) in buffer.set_slices() {
137 if start > prev {
138 f(prev, start);
139 }
140 prev = end;
141 }
142 if prev < len {
143 f(prev, len);
144 }
145 }
146 }
147}
148
149impl VarBinViewData {
150 fn dtype_parts(dtype: &DType) -> VortexResult<(bool, Nullability)> {
151 match dtype {
152 DType::Utf8(nullability) => Ok((true, *nullability)),
153 DType::Binary(nullability) => Ok((false, *nullability)),
154 _ => vortex_bail!(InvalidArgument: "invalid DType {dtype} for `VarBinViewArray`"),
155 }
156 }
157
158 pub(super) fn make_slots(validity: &Validity, len: usize) -> ArraySlots {
160 smallvec![validity_to_child(validity, len)]
161 }
162
163 pub fn new(
170 views: Buffer<BinaryView>,
171 buffers: Arc<[ByteBuffer]>,
172 dtype: DType,
173 validity: Validity,
174 ctx: &mut ExecutionCtx,
175 ) -> Self {
176 Self::try_new(views, buffers, dtype, validity, ctx)
177 .vortex_expect("VarBinViewArray construction failed")
178 }
179
180 pub fn new_handle(
187 views: BufferHandle,
188 buffers: Arc<[BufferHandle]>,
189 dtype: DType,
190 validity: Validity,
191 ) -> Self {
192 Self::try_new_handle(views, buffers, dtype, validity)
193 .vortex_expect("VarbinViewArray construction failed")
194 }
195
196 pub fn try_new(
205 views: Buffer<BinaryView>,
206 buffers: Arc<[ByteBuffer]>,
207 dtype: DType,
208 validity: Validity,
209 ctx: &mut ExecutionCtx,
210 ) -> VortexResult<Self> {
211 let views = Self::validate_and_fix(views, &buffers, &dtype, &validity, ctx)?;
212
213 Ok(unsafe { Self::new_unchecked(views, buffers, dtype, validity) })
215 }
216
217 pub fn try_new_handle(
226 views: BufferHandle,
227 buffers: Arc<[BufferHandle]>,
228 dtype: DType,
229 validity: Validity,
230 ) -> VortexResult<Self> {
231 let views_nbytes = views.len();
232 vortex_ensure!(
233 views_nbytes.is_multiple_of(size_of::<BinaryView>()),
234 "Expected views buffer length ({views_nbytes}) to be a multiple of {}",
235 size_of::<BinaryView>()
236 );
237
238 if let Some(host) = views.as_host_opt() {
240 vortex_ensure!(
241 host.is_aligned(Alignment::of::<BinaryView>()),
242 "Views on host must be 16 byte aligned"
243 );
244 }
245
246 Ok(unsafe { Self::new_handle_unchecked(views, buffers, dtype, validity) })
248 }
249
250 pub unsafe fn new_unchecked(
280 views: Buffer<BinaryView>,
281 buffers: Arc<[ByteBuffer]>,
282 dtype: DType,
283 validity: Validity,
284 ) -> Self {
285 #[cfg(debug_assertions)]
286 #[expect(clippy::disallowed_methods)]
287 Self::validate(
288 &views,
289 &buffers,
290 &dtype,
291 &validity,
292 &mut legacy_session().create_execution_ctx(),
293 )
294 .vortex_expect("[Debug Assertion]: Invalid `VarBinViewArray` parameters");
295
296 let handles: Vec<BufferHandle> = buffers
297 .iter()
298 .cloned()
299 .map(BufferHandle::new_host)
300 .collect();
301
302 let handles = Arc::from(handles);
303 let view_handle = BufferHandle::new_host(views.into_byte_buffer());
304 unsafe { Self::new_handle_unchecked(view_handle, handles, dtype, validity) }
305 }
306
307 pub unsafe fn new_handle_unchecked(
313 views: BufferHandle,
314 buffers: Arc<[BufferHandle]>,
315 dtype: DType,
316 _validity: Validity,
317 ) -> Self {
318 let _ =
319 Self::dtype_parts(&dtype).vortex_expect("VarBinViewArray dtype must be utf8 or binary");
320 Self { buffers, views }
321 }
322
323 pub fn validate(
328 views: &Buffer<BinaryView>,
329 buffers: &Arc<[ByteBuffer]>,
330 dtype: &DType,
331 validity: &Validity,
332 ctx: &mut ExecutionCtx,
333 ) -> VortexResult<()> {
334 let check_utf8 = Self::check_nullability(dtype, validity)?;
335 match validity {
336 Validity::Array(_) => {
340 let mask = validity.execute_mask(views.len(), ctx)?;
341 for ((idx, view), valid) in views.iter().enumerate().zip(mask.iter()) {
342 if valid {
343 Self::validate_view(idx, view, buffers, check_utf8)?;
344 }
345 }
346 }
347 Validity::AllInvalid => {}
349 Validity::NonNullable | Validity::AllValid => {
351 for (idx, view) in views.iter().enumerate() {
352 Self::validate_view(idx, view, buffers, check_utf8)?;
353 }
354 }
355 }
356 Ok(())
357 }
358
359 pub(crate) fn validate_and_fix(
361 views: Buffer<BinaryView>,
362 buffers: &Arc<[ByteBuffer]>,
363 dtype: &DType,
364 validity: &Validity,
365 ctx: &mut ExecutionCtx,
366 ) -> VortexResult<Buffer<BinaryView>> {
367 let check_utf8 = Self::check_nullability(dtype, validity)?;
368 let empty = BinaryView::empty_view();
369 let len = views.len();
370
371 match validity {
372 Validity::Array(_) => {
373 let mask = validity.execute_mask(len, ctx)?;
374 match views.try_into_mut() {
375 Ok(mut views) => {
376 let slice = views.as_mut_slice();
377 for (idx, valid) in mask.iter().enumerate() {
378 if valid {
379 Self::validate_view(idx, &slice[idx], buffers, check_utf8)?;
380 } else {
381 slice[idx] = empty;
382 }
383 }
384 Ok(views.freeze())
385 }
386 Err(views) => {
387 let mut needs_replace = false;
388 for ((idx, view), valid) in views.iter().enumerate().zip(mask.iter()) {
389 if valid {
390 Self::validate_view(idx, view, buffers, check_utf8)?;
391 } else if *view != empty {
392 needs_replace = true;
393 }
394 }
395 if !needs_replace {
396 return Ok(views);
397 }
398 let mut views = views.into_mut();
399 let slice = views.as_mut_slice();
400 for_each_invalid_range(&mask, len, |start, end| {
401 slice[start..end].fill(empty)
402 });
403 Ok(views.freeze())
404 }
405 }
406 }
407 Validity::AllInvalid => match views.try_into_mut() {
409 Ok(mut views) => {
410 views.as_mut_slice().fill(empty);
411 Ok(views.freeze())
412 }
413 Err(views) if views.iter().all(|view| *view == empty) => Ok(views),
414 Err(views) => {
415 let mut views = views.into_mut();
416 views.as_mut_slice().fill(empty);
417 Ok(views.freeze())
418 }
419 },
420 Validity::NonNullable | Validity::AllValid => {
422 for (idx, view) in views.iter().enumerate() {
423 Self::validate_view(idx, view, buffers, check_utf8)?;
424 }
425 Ok(views)
426 }
427 }
428 }
429
430 fn check_nullability(dtype: &DType, validity: &Validity) -> VortexResult<bool> {
431 let (is_utf8, nullability) = Self::dtype_parts(dtype)?;
432 vortex_ensure!(
433 validity.nullability() == nullability,
434 InvalidArgument: "validity {:?} incompatible with nullability {:?}",
435 validity,
436 nullability
437 );
438 Ok(is_utf8)
439 }
440
441 fn validate_view(
442 idx: usize,
443 view: &BinaryView,
444 buffers: &Arc<[ByteBuffer]>,
445 check_utf8: bool,
446 ) -> VortexResult<()> {
447 let valid_utf8 = |bytes: &[u8]| !check_utf8 || simdutf8::basic::from_utf8(bytes).is_ok();
448 if view.is_inlined() {
449 let bytes = &view.as_inlined().data[..view.len() as usize];
451 vortex_ensure!(
452 valid_utf8(bytes),
453 InvalidArgument: "view at index {idx}: inlined bytes failed utf-8 validation"
454 );
455 } else {
456 let view = view.as_view();
458 let buf_index = view.buffer_index as usize;
459 let start_offset = view.offset as usize;
460 let end_offset = start_offset.saturating_add(view.size as usize);
461
462 let buf = buffers.get(buf_index).ok_or_else(||
463 vortex_err!(InvalidArgument: "view at index {idx} references invalid buffer: {buf_index} out of bounds for VarBinViewData with {} buffers",
464 buffers.len()))?;
465
466 vortex_ensure!(
467 start_offset < buf.len(),
468 InvalidArgument: "start offset {start_offset} out of bounds for buffer {buf_index} with size {}",
469 buf.len(),
470 );
471
472 vortex_ensure!(
473 end_offset <= buf.len(),
474 InvalidArgument: "end offset {end_offset} out of bounds for buffer {buf_index} with size {}",
475 buf.len(),
476 );
477
478 let bytes = &buf[start_offset..end_offset];
480 vortex_ensure!(
481 view.prefix == bytes[..4],
482 InvalidArgument: "VarBinView prefix does not match full string"
483 );
484
485 vortex_ensure!(
487 valid_utf8(bytes),
488 InvalidArgument: "view at index {idx}: outlined bytes fails utf-8 validation"
489 );
490 }
491 Ok(())
492 }
493
494 pub fn len(&self) -> usize {
496 self.views.len() / size_of::<BinaryView>()
497 }
498
499 pub fn is_empty(&self) -> bool {
501 self.len() == 0
502 }
503
504 #[inline]
510 pub fn views(&self) -> &[BinaryView] {
511 let host_views = self.views.as_host();
512 let len = host_views.len() / size_of::<BinaryView>();
513
514 unsafe { std::slice::from_raw_parts(host_views.as_ptr().cast(), len) }
516 }
517
518 pub fn views_handle(&self) -> &BufferHandle {
520 &self.views
521 }
522
523 #[inline]
527 pub fn bytes_at(&self, index: usize) -> ByteBuffer {
528 let views = self.views();
529 let view = &views[index];
530 if !view.is_inlined() {
532 let view_ref = view.as_view();
533 self.buffer(view_ref.buffer_index as usize)
534 .slice(view_ref.as_range())
535 } else {
536 self.views_handle()
538 .as_host()
539 .clone()
540 .into_byte_buffer()
541 .slice_ref(view.as_inlined().value())
542 }
543 }
544
545 #[inline]
552 pub fn buffer(&self, idx: usize) -> &ByteBuffer {
553 if idx >= self.data_buffers().len() {
554 vortex_panic!(
555 "{idx} buffer index out of bounds, there are {} buffers",
556 self.data_buffers().len()
557 );
558 }
559 self.buffers[idx].as_host()
560 }
561
562 #[inline]
564 pub fn data_buffers(&self) -> &Arc<[BufferHandle]> {
565 &self.buffers
566 }
567
568 #[expect(
570 clippy::same_name_method,
571 reason = "intentionally named from_iter like Iterator::from_iter"
572 )]
573 pub fn from_iter<T: AsRef<[u8]>, I: IntoIterator<Item = Option<T>>>(
574 iter: I,
575 dtype: DType,
576 ) -> Self {
577 let iter = iter.into_iter();
578 let mut builder = VarBinViewBuilder::with_capacity(dtype, iter.size_hint().0);
579
580 for item in iter {
581 match item {
582 None => builder.append_null(),
583 Some(v) => builder.append_value(v),
584 }
585 }
586
587 builder.finish_into_varbinview().into_data()
588 }
589
590 pub fn from_iter_str<T: AsRef<str>, I: IntoIterator<Item = T>>(iter: I) -> Self {
591 let iter = iter.into_iter();
592 let mut builder = VarBinViewBuilder::with_capacity(
593 DType::Utf8(Nullability::NonNullable),
594 iter.size_hint().0,
595 );
596
597 for item in iter {
598 builder.append_value(item.as_ref());
599 }
600
601 builder.finish_into_varbinview().into_data()
602 }
603
604 pub fn from_iter_nullable_str<T: AsRef<str>, I: IntoIterator<Item = Option<T>>>(
605 iter: I,
606 ) -> Self {
607 let iter = iter.into_iter();
608 let mut builder = VarBinViewBuilder::with_capacity(
609 DType::Utf8(Nullability::Nullable),
610 iter.size_hint().0,
611 );
612
613 for item in iter {
614 match item {
615 None => builder.append_null(),
616 Some(v) => builder.append_value(v.as_ref()),
617 }
618 }
619
620 builder.finish_into_varbinview().into_data()
621 }
622
623 pub fn from_iter_bin<T: AsRef<[u8]>, I: IntoIterator<Item = T>>(iter: I) -> Self {
624 let iter = iter.into_iter();
625 let mut builder = VarBinViewBuilder::with_capacity(
626 DType::Binary(Nullability::NonNullable),
627 iter.size_hint().0,
628 );
629
630 for item in iter {
631 builder.append_value(item.as_ref());
632 }
633
634 builder.finish_into_varbinview().into_data()
635 }
636
637 pub fn from_iter_nullable_bin<T: AsRef<[u8]>, I: IntoIterator<Item = Option<T>>>(
638 iter: I,
639 ) -> Self {
640 let iter = iter.into_iter();
641 let mut builder = VarBinViewBuilder::with_capacity(
642 DType::Binary(Nullability::Nullable),
643 iter.size_hint().0,
644 );
645
646 for item in iter {
647 match item {
648 None => builder.append_null(),
649 Some(v) => builder.append_value(v.as_ref()),
650 }
651 }
652
653 builder.finish_into_varbinview().into_data()
654 }
655}
656
657pub trait VarBinViewArrayExt: TypedArrayRef<VarBinView> {
658 fn dtype_parts(&self) -> (bool, Nullability) {
659 match self.as_ref().dtype() {
660 DType::Utf8(nullability) => (true, *nullability),
661 DType::Binary(nullability) => (false, *nullability),
662 _ => unreachable!("VarBinViewArrayExt requires a utf8 or binary dtype"),
663 }
664 }
665
666 fn varbinview_validity(&self) -> Validity {
667 child_to_validity(
668 self.as_ref().slots()[VarBinViewSlots::VALIDITY].as_ref(),
669 self.dtype_parts().1,
670 )
671 }
672}
673impl<T: TypedArrayRef<VarBinView>> VarBinViewArrayExt for T {}
674
675impl Array<VarBinView> {
676 #[inline]
677 fn from_prevalidated_data(dtype: DType, data: VarBinViewData, slots: ArraySlots) -> Self {
678 let len = data.len();
679 unsafe {
680 Array::from_parts_unchecked(
681 ArrayParts::new(VarBinView, dtype, len, data).with_slots(slots),
682 )
683 }
684 }
685
686 #[expect(
688 clippy::same_name_method,
689 reason = "intentionally named from_iter like Iterator::from_iter"
690 )]
691 pub fn from_iter<T: AsRef<[u8]>, I: IntoIterator<Item = Option<T>>>(
692 iter: I,
693 dtype: DType,
694 ) -> Self {
695 let iter = iter.into_iter();
696 let mut builder = VarBinViewBuilder::with_capacity(dtype, iter.size_hint().0);
697 for value in iter {
698 match value {
699 Some(value) => builder.append_value(value),
700 None => builder.append_null(),
701 }
702 }
703 builder.finish_into_varbinview()
704 }
705
706 pub fn from_iter_str<T: AsRef<str>, I: IntoIterator<Item = T>>(iter: I) -> Self {
707 let iter = iter.into_iter();
708 let mut builder = VarBinViewBuilder::with_capacity(
709 DType::Utf8(Nullability::NonNullable),
710 iter.size_hint().0,
711 );
712 for value in iter {
713 builder.append_value(value.as_ref());
714 }
715 builder.finish_into_varbinview()
716 }
717
718 pub fn from_iter_nullable_str<T: AsRef<str>, I: IntoIterator<Item = Option<T>>>(
719 iter: I,
720 ) -> Self {
721 let iter = iter.into_iter();
722 let mut builder = VarBinViewBuilder::with_capacity(
723 DType::Utf8(Nullability::Nullable),
724 iter.size_hint().0,
725 );
726 for value in iter {
727 match value {
728 Some(value) => builder.append_value(value.as_ref()),
729 None => builder.append_null(),
730 }
731 }
732 builder.finish_into_varbinview()
733 }
734
735 pub fn from_iter_bin<T: AsRef<[u8]>, I: IntoIterator<Item = T>>(iter: I) -> Self {
736 let iter = iter.into_iter();
737 let mut builder = VarBinViewBuilder::with_capacity(
738 DType::Binary(Nullability::NonNullable),
739 iter.size_hint().0,
740 );
741 for value in iter {
742 builder.append_value(value.as_ref());
743 }
744 builder.finish_into_varbinview()
745 }
746
747 pub fn from_iter_nullable_bin<T: AsRef<[u8]>, I: IntoIterator<Item = Option<T>>>(
748 iter: I,
749 ) -> Self {
750 let iter = iter.into_iter();
751 let mut builder = VarBinViewBuilder::with_capacity(
752 DType::Binary(Nullability::Nullable),
753 iter.size_hint().0,
754 );
755 for value in iter {
756 match value {
757 Some(value) => builder.append_value(value.as_ref()),
758 None => builder.append_null(),
759 }
760 }
761 builder.finish_into_varbinview()
762 }
763
764 pub fn try_new(
766 views: Buffer<BinaryView>,
767 buffers: Arc<[ByteBuffer]>,
768 dtype: DType,
769 validity: Validity,
770 ctx: &mut ExecutionCtx,
771 ) -> VortexResult<Self> {
772 let data = VarBinViewData::try_new(views, buffers, dtype.clone(), validity.clone(), ctx)?;
773 let slots = VarBinViewData::make_slots(&validity, data.len());
774 Ok(Self::from_prevalidated_data(dtype, data, slots))
775 }
776
777 pub unsafe fn new_unchecked(
783 views: Buffer<BinaryView>,
784 buffers: Arc<[ByteBuffer]>,
785 dtype: DType,
786 validity: Validity,
787 ) -> Self {
788 let data = unsafe {
789 VarBinViewData::new_unchecked(views, buffers, dtype.clone(), validity.clone())
790 };
791 let slots = VarBinViewData::make_slots(&validity, data.len());
792 Self::from_prevalidated_data(dtype, data, slots)
793 }
794
795 pub fn new_handle(
797 views: BufferHandle,
798 buffers: Arc<[BufferHandle]>,
799 dtype: DType,
800 validity: Validity,
801 ) -> Self {
802 let data = VarBinViewData::new_handle(views, buffers, dtype.clone(), validity.clone());
803 let slots = VarBinViewData::make_slots(&validity, data.len());
804 Self::from_prevalidated_data(dtype, data, slots)
805 }
806
807 pub unsafe fn new_handle_unchecked(
813 views: BufferHandle,
814 buffers: Arc<[BufferHandle]>,
815 dtype: DType,
816 validity: Validity,
817 ) -> Self {
818 let data = unsafe {
819 VarBinViewData::new_handle_unchecked(views, buffers, dtype.clone(), validity.clone())
820 };
821 let slots = VarBinViewData::make_slots(&validity, data.len());
822 Self::from_prevalidated_data(dtype, data, slots)
823 }
824
825 pub fn into_data_parts(self) -> VarBinViewDataParts {
826 let dtype = self.dtype().clone();
827 let validity = self.varbinview_validity();
828 let data = self.into_data();
829 VarBinViewDataParts {
830 dtype,
831 buffers: data.buffers,
832 views: data.views,
833 validity,
834 }
835 }
836}
837
838impl<'a> FromIterator<Option<&'a [u8]>> for VarBinViewData {
839 fn from_iter<T: IntoIterator<Item = Option<&'a [u8]>>>(iter: T) -> Self {
840 Self::from_iter_nullable_bin(iter)
841 }
842}
843
844impl FromIterator<Option<Vec<u8>>> for VarBinViewData {
845 fn from_iter<T: IntoIterator<Item = Option<Vec<u8>>>>(iter: T) -> Self {
846 Self::from_iter_nullable_bin(iter)
847 }
848}
849
850impl FromIterator<Option<String>> for VarBinViewData {
851 fn from_iter<T: IntoIterator<Item = Option<String>>>(iter: T) -> Self {
852 Self::from_iter_nullable_str(iter)
853 }
854}
855
856impl<'a> FromIterator<Option<&'a str>> for VarBinViewData {
857 fn from_iter<T: IntoIterator<Item = Option<&'a str>>>(iter: T) -> Self {
858 Self::from_iter_nullable_str(iter)
859 }
860}
861
862impl<'a> FromIterator<Option<&'a [u8]>> for Array<VarBinView> {
865 fn from_iter<T: IntoIterator<Item = Option<&'a [u8]>>>(iter: T) -> Self {
866 Self::from_iter(iter, DType::Binary(Nullability::Nullable))
867 }
868}
869
870impl FromIterator<Option<Vec<u8>>> for Array<VarBinView> {
871 fn from_iter<T: IntoIterator<Item = Option<Vec<u8>>>>(iter: T) -> Self {
872 Self::from_iter(iter, DType::Binary(Nullability::Nullable))
873 }
874}
875
876impl FromIterator<Option<String>> for Array<VarBinView> {
877 fn from_iter<T: IntoIterator<Item = Option<String>>>(iter: T) -> Self {
878 Self::from_iter_nullable_str(iter)
879 }
880}
881
882impl<'a> FromIterator<Option<&'a str>> for Array<VarBinView> {
883 fn from_iter<T: IntoIterator<Item = Option<&'a str>>>(iter: T) -> Self {
884 Self::from_iter_nullable_str(iter)
885 }
886}