1#![allow(unsafe_op_in_unsafe_fn)]
2mod builder;
5pub use builder::*;
6mod ffi;
7pub(super) mod fmt;
8mod iterator;
9mod mutable;
10#[cfg(feature = "proptest")]
11pub mod proptest;
12mod view;
13
14use std::any::Any;
15use std::fmt::Debug;
16use std::marker::PhantomData;
17
18use polars_buffer::Buffer;
19use polars_error::*;
20use polars_utils::relaxed_cell::RelaxedCell;
21
22use crate::array::Array;
23use crate::bitmap::Bitmap;
24use crate::datatypes::ArrowDataType;
25
26mod private {
27 pub trait Sealed: Send + Sync {}
28
29 impl Sealed for str {}
30 impl Sealed for [u8] {}
31}
32pub use iterator::BinaryViewValueIter;
33pub use mutable::MutableBinaryViewArray;
34use polars_utils::aliases::{InitHashMaps, PlHashMap};
35use private::Sealed;
36
37use crate::array::binview::view::{validate_binary_views, validate_views_utf8_only};
38use crate::array::iterator::NonNullValuesIter;
39use crate::bitmap::utils::{BitmapIter, ZipValidity};
40pub type BinaryViewArray = BinaryViewArrayGeneric<[u8]>;
41pub type Utf8ViewArray = BinaryViewArrayGeneric<str>;
42pub type BinaryViewArrayBuilder = BinaryViewArrayGenericBuilder<[u8]>;
43pub type Utf8ViewArrayBuilder = BinaryViewArrayGenericBuilder<str>;
44pub use view::{View, validate_utf8_views};
45
46use super::Splitable;
47
48pub type MutablePlString = MutableBinaryViewArray<str>;
49pub type MutablePlBinary = MutableBinaryViewArray<[u8]>;
50
51static BIN_VIEW_TYPE: ArrowDataType = ArrowDataType::BinaryView;
52static UTF8_VIEW_TYPE: ArrowDataType = ArrowDataType::Utf8View;
53
54pub const BINVIEW_ARROW_BUFFER_LEN_LIMIT: usize = i32::MAX as usize;
56pub const BINVIEW_MAX_ROW_BYTE_LEN: usize = (u32::MAX - 1) as usize;
57const DEFAULT_BLOCK_SIZE: usize = 8 * 1024;
58const MAX_EXP_BLOCK_SIZE: usize = 16 * 1024 * 1024;
59
60pub trait ViewType: Sealed + 'static + PartialEq + AsRef<Self> {
61 const IS_UTF8: bool;
62 const DATA_TYPE: ArrowDataType;
63 type Owned: Debug + Clone + Sync + Send + AsRef<Self>;
64
65 unsafe fn from_bytes_unchecked(slice: &[u8]) -> &Self;
68 fn from_bytes(slice: &[u8]) -> Option<&Self>;
69
70 fn to_bytes(&self) -> &[u8];
71
72 #[allow(clippy::wrong_self_convention)]
73 fn into_owned(&self) -> Self::Owned;
74
75 fn dtype() -> &'static ArrowDataType;
76}
77
78impl ViewType for str {
79 const IS_UTF8: bool = true;
80 const DATA_TYPE: ArrowDataType = ArrowDataType::Utf8View;
81 type Owned = String;
82
83 #[inline(always)]
84 unsafe fn from_bytes_unchecked(slice: &[u8]) -> &Self {
85 std::str::from_utf8_unchecked(slice)
86 }
87 #[inline(always)]
88 fn from_bytes(slice: &[u8]) -> Option<&Self> {
89 std::str::from_utf8(slice).ok()
90 }
91
92 #[inline(always)]
93 fn to_bytes(&self) -> &[u8] {
94 self.as_bytes()
95 }
96
97 fn into_owned(&self) -> Self::Owned {
98 self.to_string()
99 }
100 fn dtype() -> &'static ArrowDataType {
101 &UTF8_VIEW_TYPE
102 }
103}
104
105impl ViewType for [u8] {
106 const IS_UTF8: bool = false;
107 const DATA_TYPE: ArrowDataType = ArrowDataType::BinaryView;
108 type Owned = Vec<u8>;
109
110 #[inline(always)]
111 unsafe fn from_bytes_unchecked(slice: &[u8]) -> &Self {
112 slice
113 }
114 #[inline(always)]
115 fn from_bytes(slice: &[u8]) -> Option<&Self> {
116 Some(slice)
117 }
118
119 #[inline(always)]
120 fn to_bytes(&self) -> &[u8] {
121 self
122 }
123
124 fn into_owned(&self) -> Self::Owned {
125 self.to_vec()
126 }
127
128 fn dtype() -> &'static ArrowDataType {
129 &BIN_VIEW_TYPE
130 }
131}
132
133pub struct BinaryViewArrayGeneric<T: ViewType + ?Sized> {
134 dtype: ArrowDataType,
135 views: Buffer<View>,
136 buffers: Buffer<Buffer<u8>>,
137 validity: Option<Bitmap>,
138 phantom: PhantomData<T>,
139 total_bytes_len: RelaxedCell<u64>,
141 total_buffer_len: usize,
143}
144
145impl<T: ViewType + ?Sized> PartialEq for BinaryViewArrayGeneric<T> {
146 fn eq(&self, other: &Self) -> bool {
147 self.len() == other.len() && self.into_iter().zip(other).all(|(l, r)| l == r)
148 }
149}
150
151impl<T: ViewType + ?Sized> Clone for BinaryViewArrayGeneric<T> {
152 fn clone(&self) -> Self {
153 Self {
154 dtype: self.dtype.clone(),
155 views: self.views.clone(),
156 buffers: self.buffers.clone(),
157 validity: self.validity.clone(),
158 phantom: Default::default(),
159 total_bytes_len: self.total_bytes_len.clone(),
160 total_buffer_len: self.total_buffer_len,
161 }
162 }
163}
164
165unsafe impl<T: ViewType + ?Sized> Send for BinaryViewArrayGeneric<T> {}
166unsafe impl<T: ViewType + ?Sized> Sync for BinaryViewArrayGeneric<T> {}
167
168const UNKNOWN_LEN: u64 = u64::MAX;
169
170impl<T: ViewType + ?Sized> BinaryViewArrayGeneric<T> {
171 pub unsafe fn new_unchecked(
176 dtype: ArrowDataType,
177 views: Buffer<View>,
178 buffers: Buffer<Buffer<u8>>,
179 validity: Option<Bitmap>,
180 total_bytes_len: Option<usize>,
181 total_buffer_len: usize,
182 ) -> Self {
183 #[cfg(debug_assertions)]
185 {
186 if let Some(validity) = validity.as_ref() {
187 assert_eq!(validity.len(), views.len());
188 }
189
190 for (i, view) in views.iter().enumerate() {
199 let is_valid = validity.as_ref().is_none_or(|v| v.get_bit(i));
200
201 if !is_valid {
202 continue;
203 }
204
205 if view.length > View::MAX_INLINE_SIZE {
207 assert!((view.buffer_idx as usize) < (buffers.len()));
208 assert!(
209 view.offset as usize + view.length as usize
210 <= buffers[view.buffer_idx as usize].len()
211 );
212 }
213 }
214
215 }
220
221 Self {
222 dtype,
223 views,
224 buffers,
225 validity,
226 phantom: Default::default(),
227 total_bytes_len: RelaxedCell::from(
228 total_bytes_len.map(|l| l as u64).unwrap_or(UNKNOWN_LEN),
229 ),
230 total_buffer_len,
231 }
232 }
233
234 pub unsafe fn new_unchecked_unknown_md(
239 dtype: ArrowDataType,
240 views: Buffer<View>,
241 buffers: Buffer<Buffer<u8>>,
242 validity: Option<Bitmap>,
243 total_buffer_len: Option<usize>,
244 ) -> Self {
245 let total_bytes_len = None;
246 let total_buffer_len =
247 total_buffer_len.unwrap_or_else(|| buffers.iter().map(|b| b.len()).sum());
248 Self::new_unchecked(
249 dtype,
250 views,
251 buffers,
252 validity,
253 total_bytes_len,
254 total_buffer_len,
255 )
256 }
257
258 pub fn data_buffers(&self) -> &Buffer<Buffer<u8>> {
259 &self.buffers
260 }
261
262 pub fn data_buffers_mut(&mut self) -> &mut Buffer<Buffer<u8>> {
263 &mut self.buffers
264 }
265
266 pub fn variadic_buffer_lengths(&self) -> Vec<i64> {
267 self.buffers.iter().map(|buf| buf.len() as i64).collect()
268 }
269
270 pub fn views(&self) -> &Buffer<View> {
271 &self.views
272 }
273
274 pub fn into_views(self) -> Vec<View> {
275 self.views.to_vec()
276 }
277
278 pub fn into_inner(
279 self,
280 ) -> (
281 Buffer<View>,
282 Buffer<Buffer<u8>>,
283 Option<Bitmap>,
284 Option<usize>,
285 usize,
286 ) {
287 let total_bytes_len = self.try_total_bytes_len();
288 let views = self.views;
289 let buffers = self.buffers;
290 let validity = self.validity;
291
292 (
293 views,
294 buffers,
295 validity,
296 total_bytes_len,
297 self.total_buffer_len,
298 )
299 }
300
301 pub unsafe fn apply_views<F: FnMut(View, &T) -> View>(&self, mut update_view: F) -> Self {
306 let arr = self.clone();
307 let (views, buffers, validity, _total_bytes_len, total_buffer_len) = arr.into_inner();
308
309 let mut total_bytes_len = 0;
310 let mut views = views.to_vec();
311 for v in views.iter_mut() {
312 let str_slice = T::from_bytes_unchecked(v.get_slice_unchecked(&buffers));
313 *v = update_view(*v, str_slice);
314 total_bytes_len += v.length as usize;
315 }
316
317 let len_valid = validity.is_none();
318 Self::new_unchecked(
319 self.dtype.clone(),
320 views.into(),
321 buffers,
322 validity,
323 len_valid.then_some(total_bytes_len),
324 total_buffer_len,
325 )
326 }
327
328 pub unsafe fn with_views_mut<F: FnOnce(&mut [View])>(&mut self, f: F) {
333 self.total_bytes_len.store(UNKNOWN_LEN);
334 if let Some(views) = self.views.get_mut_slice() {
335 f(views)
336 } else {
337 let mut views = self.views.as_slice().to_vec();
338 f(&mut views);
339 self.views = Buffer::from(views);
340 }
341 }
342
343 pub fn try_new(
344 dtype: ArrowDataType,
345 views: Buffer<View>,
346 buffers: Buffer<Buffer<u8>>,
347 validity: Option<Bitmap>,
348 ) -> PolarsResult<Self> {
349 if T::IS_UTF8 {
350 validate_utf8_views(views.as_ref(), buffers.as_ref())?;
351 } else {
352 validate_binary_views(views.as_ref(), buffers.as_ref())?;
353 }
354
355 if let Some(validity) = &validity {
356 polars_ensure!(validity.len()== views.len(), ComputeError: "validity mask length must match the number of values" )
357 }
358
359 unsafe {
360 Ok(Self::new_unchecked_unknown_md(
361 dtype, views, buffers, validity, None,
362 ))
363 }
364 }
365
366 #[inline]
368 pub fn new_empty(dtype: ArrowDataType) -> Self {
369 unsafe { Self::new_unchecked(dtype, Buffer::new(), Buffer::new(), None, Some(0), 0) }
370 }
371
372 #[inline]
374 pub fn new_null(dtype: ArrowDataType, length: usize) -> Self {
375 let validity = Some(Bitmap::new_zeroed(length));
376 unsafe {
377 Self::new_unchecked(
378 dtype,
379 Buffer::zeroed(length),
380 Buffer::new(),
381 validity,
382 Some(0),
383 0,
384 )
385 }
386 }
387
388 #[inline]
392 pub fn value(&self, i: usize) -> &T {
393 assert!(i < self.len());
394 unsafe { self.value_unchecked(i) }
395 }
396
397 #[inline]
402 pub unsafe fn value_unchecked(&self, i: usize) -> &T {
403 let v = self.views.get_unchecked(i);
404 T::from_bytes_unchecked(v.get_slice_unchecked(&self.buffers))
405 }
406
407 #[inline]
411 pub fn get(&self, i: usize) -> Option<&T> {
412 assert!(i < self.len());
413 unsafe { self.get_unchecked(i) }
414 }
415
416 #[inline]
421 pub unsafe fn get_unchecked(&self, i: usize) -> Option<&T> {
422 if self
423 .validity
424 .as_ref()
425 .is_none_or(|v| v.get_bit_unchecked(i))
426 {
427 let v = self.views.get_unchecked(i);
428 Some(T::from_bytes_unchecked(
429 v.get_slice_unchecked(&self.buffers),
430 ))
431 } else {
432 None
433 }
434 }
435
436 pub fn iter(&self) -> ZipValidity<&T, BinaryViewValueIter<'_, T>, BitmapIter<'_>> {
438 ZipValidity::new_with_validity(self.values_iter(), self.validity.as_ref())
439 }
440
441 pub fn values_iter(&self) -> BinaryViewValueIter<'_, T> {
443 BinaryViewValueIter::new(self)
444 }
445
446 pub fn len_iter(&self) -> impl Iterator<Item = u32> + '_ {
447 self.views.iter().map(|v| v.length)
448 }
449
450 pub fn non_null_values_iter(&self) -> NonNullValuesIter<'_, BinaryViewArrayGeneric<T>> {
452 NonNullValuesIter::new(self, self.validity())
453 }
454
455 pub fn non_null_views_iter(&self) -> NonNullValuesIter<'_, Buffer<View>> {
457 NonNullValuesIter::new(self.views(), self.validity())
458 }
459
460 impl_sliced!();
461 impl_into_array!();
462
463 #[must_use]
467 #[inline]
468 pub fn with_validity(mut self, validity: Option<Bitmap>) -> Self {
469 self.set_validity(validity);
470 self
471 }
472
473 #[inline]
477 pub fn set_validity(&mut self, validity: Option<Bitmap>) {
478 if matches!(&validity, Some(bitmap) if bitmap.len() != self.len()) {
479 panic!("validity must be equal to the array's length")
480 }
481 self.total_bytes_len.store(UNKNOWN_LEN);
482 self.validity = validity;
483 }
484
485 #[inline]
487 pub fn take_validity(&mut self) -> Option<Bitmap> {
488 self.total_bytes_len.store(UNKNOWN_LEN);
489 self.validity.take()
490 }
491
492 pub fn from_slice<S: AsRef<T>, P: AsRef<[Option<S>]>>(slice: P) -> Self {
493 let mutable = MutableBinaryViewArray::from_iterator(
494 slice.as_ref().iter().map(|opt_v| opt_v.as_ref()),
495 );
496 mutable.into()
497 }
498
499 pub fn from_slice_values<S: AsRef<T>, P: AsRef<[S]>>(slice: P) -> Self {
500 let mutable =
501 MutableBinaryViewArray::from_values_iter(slice.as_ref().iter().map(|v| v.as_ref()));
502 mutable.into()
503 }
504
505 pub fn total_bytes_len(&self) -> usize {
507 let total = self.total_bytes_len.load();
508 if total == UNKNOWN_LEN {
509 let total = ZipValidity::new_with_validity(self.len_iter(), self.validity.as_ref())
510 .map(|v| v.unwrap_or(0) as usize)
511 .sum::<usize>();
512 self.total_bytes_len.store(total as u64);
513 total
514 } else {
515 total as usize
516 }
517 }
518
519 pub fn try_total_bytes_len(&self) -> Option<usize> {
521 let b = self.total_bytes_len.load();
522 (b != UNKNOWN_LEN).then_some(b as usize)
523 }
524
525 pub fn total_buffer_len(&self) -> usize {
527 self.total_buffer_len
528 }
529
530 fn total_unshared_buffer_len(&self) -> usize {
531 self.buffers
535 .iter()
536 .map(|buf| {
537 if buf.storage_refcount() > 1 {
538 0
539 } else {
540 buf.len()
541 }
542 })
543 .sum()
544 }
545
546 #[inline(always)]
547 pub fn len(&self) -> usize {
548 self.views.len()
549 }
550
551 pub fn gc(self) -> Self {
553 if self.buffers.is_empty() {
554 return self;
555 }
556 let mut mutable = MutableBinaryViewArray::with_capacity(self.len());
557 let buffers = self.buffers.as_ref();
558
559 for view in self.views.as_ref() {
560 unsafe { mutable.push_view_unchecked(*view, buffers) }
561 }
562 mutable.freeze().with_validity(self.validity)
563 }
564
565 pub fn deshare(&self) -> Self {
566 if self.buffers.storage_refcount() == 1
567 && self.buffers.iter().all(|b| b.storage_refcount() == 1)
568 {
569 return self.clone();
570 }
571 self.clone().gc()
572 }
573
574 pub fn is_sliced(&self) -> bool {
575 !std::ptr::eq(self.views.as_ptr(), self.views.storage_ptr())
576 }
577
578 pub fn maybe_gc(self) -> Self {
579 const GC_MINIMUM_SAVINGS: usize = 16 * 1024; if self.total_buffer_len <= GC_MINIMUM_SAVINGS {
582 return self;
583 }
584
585 if self.buffers.storage_refcount() != 1 {
586 return self;
590 }
591
592 let total_bytes_len = self.total_bytes_len();
595 let buffer_req_lower_bound = total_bytes_len.saturating_sub(self.len() * 12);
596
597 let lower_bound_mem_usage_post_gc = self.len() * 16 + buffer_req_lower_bound;
598 let cur_mem_usage = self.len() * 16 + self.total_unshared_buffer_len();
600 let savings_upper_bound = cur_mem_usage.saturating_sub(lower_bound_mem_usage_post_gc);
601
602 if savings_upper_bound >= GC_MINIMUM_SAVINGS
603 && cur_mem_usage >= 4 * lower_bound_mem_usage_post_gc
604 {
605 self.gc()
606 } else {
607 self
608 }
609 }
610
611 pub fn make_mut(self) -> MutableBinaryViewArray<T> {
612 let views = self.views.to_vec();
613 let completed_buffers = self.buffers.to_vec();
614 let validity = self.validity.map(|bitmap| bitmap.make_mut());
615
616 let mut total_bytes_len = self.total_bytes_len.load();
618 if total_bytes_len == UNKNOWN_LEN {
619 total_bytes_len = views.iter().map(|view| view.length as u64).sum();
620 }
621 let total_bytes_len = total_bytes_len as usize;
622
623 MutableBinaryViewArray {
624 views,
625 completed_buffers,
626 in_progress_buffer: vec![],
627 validity,
628 phantom: Default::default(),
629 total_bytes_len,
630 total_buffer_len: self.total_buffer_len,
631 stolen_buffers: PlHashMap::new(),
632 }
633 }
634}
635
636impl BinaryViewArray {
637 pub fn validate_utf8(&self) -> PolarsResult<()> {
639 unsafe { validate_views_utf8_only(&self.views, &self.buffers, 0) }
641 }
642
643 pub fn to_utf8view(&self) -> PolarsResult<Utf8ViewArray> {
645 self.validate_utf8()?;
646 unsafe { Ok(self.to_utf8view_unchecked()) }
647 }
648
649 pub unsafe fn to_utf8view_unchecked(&self) -> Utf8ViewArray {
654 Utf8ViewArray::new_unchecked(
655 ArrowDataType::Utf8View,
656 self.views.clone(),
657 self.buffers.clone(),
658 self.validity.clone(),
659 self.try_total_bytes_len(),
660 self.total_buffer_len,
661 )
662 }
663}
664
665impl Utf8ViewArray {
666 pub fn to_binview(&self) -> BinaryViewArray {
667 unsafe {
669 BinaryViewArray::new_unchecked(
670 ArrowDataType::BinaryView,
671 self.views.clone(),
672 self.buffers.clone(),
673 self.validity.clone(),
674 self.try_total_bytes_len(),
675 self.total_buffer_len,
676 )
677 }
678 }
679}
680
681impl<T: ViewType + ?Sized> Array for BinaryViewArrayGeneric<T> {
682 fn as_any(&self) -> &dyn Any {
683 self
684 }
685
686 fn as_any_mut(&mut self) -> &mut dyn Any {
687 self
688 }
689
690 #[inline(always)]
691 fn len(&self) -> usize {
692 BinaryViewArrayGeneric::len(self)
693 }
694
695 #[inline(always)]
696 fn dtype(&self) -> &ArrowDataType {
697 &self.dtype
698 }
699
700 #[inline(always)]
701 fn dtype_mut(&mut self) -> &mut ArrowDataType {
702 &mut self.dtype
703 }
704
705 fn validity(&self) -> Option<&Bitmap> {
706 self.validity.as_ref()
707 }
708
709 fn split_at_boxed(&self, offset: usize) -> (Box<dyn Array>, Box<dyn Array>) {
710 let (lhs, rhs) = Splitable::split_at(self, offset);
711 (Box::new(lhs), Box::new(rhs))
712 }
713
714 unsafe fn split_at_boxed_unchecked(&self, offset: usize) -> (Box<dyn Array>, Box<dyn Array>) {
715 let (lhs, rhs) = unsafe { Splitable::split_at_unchecked(self, offset) };
716 (Box::new(lhs), Box::new(rhs))
717 }
718
719 fn slice(&mut self, offset: usize, length: usize) {
720 assert!(
721 offset + length <= self.len(),
722 "the offset of the new Buffer cannot exceed the existing length"
723 );
724 unsafe { self.slice_unchecked(offset, length) }
725 }
726
727 unsafe fn slice_unchecked(&mut self, offset: usize, length: usize) {
728 debug_assert!(offset + length <= self.len());
729 self.validity = self
730 .validity
731 .take()
732 .map(|bitmap| bitmap.sliced_unchecked(offset, length))
733 .filter(|bitmap| bitmap.unset_bits() > 0);
734 self.views.slice_in_place_unchecked(offset..offset + length);
735 self.total_bytes_len.store(UNKNOWN_LEN)
736 }
737
738 fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
739 debug_assert!(
740 validity.as_ref().is_none_or(|v| v.len() == self.len()),
741 "{} != {}",
742 validity.as_ref().unwrap().len(),
743 self.len()
744 );
745
746 let mut new = self.clone();
747 new.validity = validity;
748 Box::new(new)
749 }
750
751 fn to_boxed(&self) -> Box<dyn Array> {
752 Box::new(self.clone())
753 }
754}
755
756impl<T: ViewType + ?Sized> Splitable for BinaryViewArrayGeneric<T> {
757 fn check_bound(&self, offset: usize) -> bool {
758 offset <= self.len()
759 }
760
761 unsafe fn _split_at_unchecked(&self, offset: usize) -> (Self, Self) {
762 let (lhs_views, rhs_views) = unsafe { self.views.split_at_unchecked(offset) };
763 let (lhs_validity, rhs_validity) = unsafe { self.validity.split_at_unchecked(offset) };
764
765 unsafe {
766 (
767 Self::new_unchecked(
768 self.dtype.clone(),
769 lhs_views,
770 self.buffers.clone(),
771 lhs_validity,
772 (offset == 0).then_some(0),
773 self.total_buffer_len(),
774 ),
775 Self::new_unchecked(
776 self.dtype.clone(),
777 rhs_views,
778 self.buffers.clone(),
779 rhs_validity,
780 (offset == self.len()).then_some(0),
781 self.total_buffer_len(),
782 ),
783 )
784 }
785 }
786}