feat(rust, python): add binary dtype

polars/Cargo.toml

@@ -163,6 +163,7 @@ dtype-full = [
   "dtype-u16",
   "dtype-categorical",
   "dtype-struct",
+  "dtype-binary",
 ]
 
 # sensible minimal set of opt-in datatypes
@@ -210,6 +211,12 @@ dtype-struct = [
   "polars-ops/dtype-struct",
   "polars-io/dtype-struct",
 ]
+dtype-binary = [
+  "polars-core/dtype-binary",
+  "polars-lazy/dtype-binary",
+  "polars-ops/dtype-binary",
+  "polars-io/dtype-binary",
+]
 
 docs-selection = [
   "csv-file",

polars/polars-arrow/src/array/default_arrays.rs

@@ -1,4 +1,4 @@
-use arrow::array::{BooleanArray, PrimitiveArray, Utf8Array};
+use arrow::array::{BinaryArray, BooleanArray, PrimitiveArray, Utf8Array};
 use arrow::bitmap::Bitmap;
 use arrow::buffer::Buffer;
 use arrow::datatypes::DataType;
@@ -40,3 +40,23 @@ impl FromDataUtf8 for Utf8Array<i64> {
         Utf8Array::from_data_unchecked(DataType::LargeUtf8, offsets, values, validity)
     }
 }
+
+pub trait FromDataBinary {
+    /// # Safety
+    /// `values` buffer must contain valid utf8 between every `offset`
+    unsafe fn from_data_unchecked_default(
+        offsets: Buffer<i64>,
+        values: Buffer<u8>,
+        validity: Option<Bitmap>,
+    ) -> Self;
+}
+
+impl FromDataBinary for BinaryArray<i64> {
+    unsafe fn from_data_unchecked_default(
+        offsets: Buffer<i64>,
+        values: Buffer<u8>,
+        validity: Option<Bitmap>,
+    ) -> Self {
+        BinaryArray::from_data_unchecked(DataType::LargeBinary, offsets, values, validity)
+    }
+}

polars/polars-arrow/src/array/get.rs

@@ -1,4 +1,4 @@
-use arrow::array::{Array, BooleanArray, ListArray, PrimitiveArray, Utf8Array};
+use arrow::array::{Array, BinaryArray, BooleanArray, ListArray, PrimitiveArray, Utf8Array};
 use arrow::types::NativeType;
 
 use crate::is_valid::IsValid;
@@ -79,6 +79,28 @@ impl<'a> ArrowGetItem for &'a Utf8Array<i64> {
     }
 }
 
+impl<'a> ArrowGetItem for &'a BinaryArray<i64> {
+    type Item = &'a [u8];
+
+    #[inline]
+    fn get(&self, item: usize) -> Option<Self::Item> {
+        if item >= self.len() {
+            None
+        } else {
+            unsafe { self.get_unchecked(item) }
+        }
+    }
+
+    #[inline]
+    unsafe fn get_unchecked(&self, item: usize) -> Option<Self::Item> {
+        if self.is_null_unchecked(item) {
+            None
+        } else {
+            Some(self.value_unchecked(item))
+        }
+    }
+}
+
 impl ArrowGetItem for ListArray<i64> {
     type Item = Box<dyn Array>;
 

polars/polars-arrow/src/array/mod.rs

@@ -1,4 +1,4 @@
-use arrow::array::{Array, BooleanArray, ListArray, PrimitiveArray, Utf8Array};
+use arrow::array::{Array, BinaryArray, BooleanArray, ListArray, PrimitiveArray, Utf8Array};
 use arrow::bitmap::MutableBitmap;
 use arrow::datatypes::DataType;
 use arrow::types::NativeType;
@@ -30,6 +30,12 @@ impl ValueSize for Utf8Array<i64> {
     }
 }
 
+impl ValueSize for BinaryArray<i64> {
+    fn get_values_size(&self) -> usize {
+        self.values().len()
+    }
+}
+
 impl ValueSize for ArrayRef {
     fn get_values_size(&self) -> usize {
         match self.data_type() {
@@ -179,6 +185,52 @@ pub trait ListFromIter {
             Some(validity.into()),
         )
     }
+
+    /// Create a list-array from an iterator.
+    /// Used in groupby agg-list
+    ///
+    /// # Safety
+    /// Will produce incorrect arrays if size hint is incorrect.
+    unsafe fn from_iter_binary_trusted_len<I, P, Ref>(iter: I, n_elements: usize) -> ListArray<i64>
+    where
+        I: IntoIterator<Item = Option<P>>,
+        P: IntoIterator<Item = Option<Ref>>,
+        Ref: AsRef<[u8]>,
+    {
+        let iterator = iter.into_iter();
+        let (lower, _) = iterator.size_hint();
+
+        let mut validity = MutableBitmap::with_capacity(lower);
+        let mut offsets = Vec::<i64>::with_capacity(lower + 1);
+        let mut length_so_far = 0i64;
+        offsets.push(length_so_far);
+        let values: BinaryArray<i64> = iterator
+            .filter_map(|opt_iter| match opt_iter {
+                Some(x) => {
+                    let it = x.into_iter();
+                    length_so_far += it.size_hint().0 as i64;
+                    validity.push(true);
+                    offsets.push(length_so_far);
+                    Some(it)
+                }
+                None => {
+                    validity.push(false);
+                    None
+                }
+            })
+            .flatten()
+            .trust_my_length(n_elements)
+            .collect();
+
+        // Safety:
+        // offsets are monotonically increasing
+        ListArray::new_unchecked(
+            ListArray::<i64>::default_datatype(DataType::LargeBinary),
+            offsets.into(),
+            Box::new(values),
+            Some(validity.into()),
+        )
+    }
 }
 impl ListFromIter for ListArray<i64> {}
 

polars/polars-arrow/src/compute/take/mod.rs

@@ -327,6 +327,21 @@ pub unsafe fn take_no_null_utf8_iter_unchecked<I: IntoIterator<Item = usize>>(
     Box::new(MutableUtf8Array::<i64>::from_trusted_len_values_iter_unchecked(iter).into())
 }
 
+/// # Safety
+/// - no bounds checks
+/// - iterator must be TrustedLen
+#[inline]
+pub unsafe fn take_no_null_binary_iter_unchecked<I: IntoIterator<Item = usize>>(
+    arr: &LargeBinaryArray,
+    indices: I,
+) -> Box<LargeBinaryArray> {
+    let iter = indices.into_iter().map(|idx| {
+        debug_assert!(idx < arr.len());
+        arr.value_unchecked(idx)
+    });
+    Box::new(MutableBinaryArray::<i64>::from_trusted_len_values_iter_unchecked(iter).into())
+}
+
 /// # Safety
 /// - no bounds checks
 /// - iterator must be TrustedLen
@@ -348,6 +363,27 @@ pub unsafe fn take_utf8_iter_unchecked<I: IntoIterator<Item = usize>>(
     Box::new(LargeStringArray::from_trusted_len_iter_unchecked(iter))
 }
 
+/// # Safety
+/// - no bounds checks
+/// - iterator must be TrustedLen
+#[inline]
+pub unsafe fn take_binary_iter_unchecked<I: IntoIterator<Item = usize>>(
+    arr: &LargeBinaryArray,
+    indices: I,
+) -> Box<LargeBinaryArray> {
+    let validity = arr.validity().expect("should have nulls");
+    let iter = indices.into_iter().map(|idx| {
+        debug_assert!(idx < arr.len());
+        if validity.get_bit_unchecked(idx) {
+            Some(arr.value_unchecked(idx))
+        } else {
+            None
+        }
+    });
+
+    Box::new(LargeBinaryArray::from_trusted_len_iter_unchecked(iter))
+}
+
 /// # Safety
 /// - no bounds checks
 /// - iterator must be TrustedLen
@@ -363,6 +399,21 @@ pub unsafe fn take_no_null_utf8_opt_iter_unchecked<I: IntoIterator<Item = Option
     Box::new(LargeStringArray::from_trusted_len_iter_unchecked(iter))
 }
 
+/// # Safety
+/// - no bounds checks
+/// - iterator must be TrustedLen
+#[inline]
+pub unsafe fn take_no_null_binary_opt_iter_unchecked<I: IntoIterator<Item = Option<usize>>>(
+    arr: &LargeBinaryArray,
+    indices: I,
+) -> Box<LargeBinaryArray> {
+    let iter = indices
+        .into_iter()
+        .map(|opt_idx| opt_idx.map(|idx| arr.value_unchecked(idx)));
+
+    Box::new(LargeBinaryArray::from_trusted_len_iter_unchecked(iter))
+}
+
 /// # Safety
 /// - no bounds checks
 /// - iterator must be TrustedLen
@@ -384,6 +435,27 @@ pub unsafe fn take_utf8_opt_iter_unchecked<I: IntoIterator<Item = Option<usize>>
     Box::new(LargeStringArray::from_trusted_len_iter_unchecked(iter))
 }
 
+/// # Safety
+/// - no bounds checks
+/// - iterator must be TrustedLen
+#[inline]
+pub unsafe fn take_binary_opt_iter_unchecked<I: IntoIterator<Item = Option<usize>>>(
+    arr: &LargeBinaryArray,
+    indices: I,
+) -> Box<LargeBinaryArray> {
+    let validity = arr.validity().expect("should have nulls");
+    let iter = indices.into_iter().map(|opt_idx| {
+        opt_idx.and_then(|idx| {
+            if validity.get_bit_unchecked(idx) {
+                Some(arr.value_unchecked(idx))
+            } else {
+                None
+            }
+        })
+    });
+    Box::new(LargeBinaryArray::from_trusted_len_iter_unchecked(iter))
+}
+
 /// # Safety
 /// caller must ensure indices are in bounds
 pub unsafe fn take_utf8_unchecked(
@@ -497,6 +569,119 @@ pub unsafe fn take_utf8_unchecked(
     ))
 }
 
+/// # Safety
+/// caller must ensure indices are in bounds
+pub unsafe fn take_binary_unchecked(
+    arr: &LargeBinaryArray,
+    indices: &IdxArr,
+) -> Box<LargeBinaryArray> {
+    let data_len = indices.len();
+
+    let mut offset_buf = vec![0; data_len + 1];
+    let offset_typed = offset_buf.as_mut_slice();
+
+    let mut length_so_far = 0;
+    offset_typed[0] = length_so_far;
+
+    let validity;
+
+    // The required size is yet unknown
+    // Allocate 2.0 times the expected size.
+    // where expected size is the length of bytes multiplied by the factor (take_len / current_len)
+    let mut values_capacity = if arr.len() > 0 {
+        ((arr.len() as f32 * 2.0) as usize) / arr.len() * indices.len() as usize
+    } else {
+        0
+    };
+
+    // 16 bytes per string as default alloc
+    let mut values_buf = Vec::<u8>::with_capacity(values_capacity);
+
+    // both 0 nulls
+    if !arr.has_validity() && !indices.has_validity() {
+        offset_typed
+            .iter_mut()
+            .skip(1)
+            .enumerate()
+            .for_each(|(idx, offset)| {
+                let index = indices.value_unchecked(idx) as usize;
+                let s = arr.value_unchecked(index);
+                length_so_far += s.len() as i64;
+                *offset = length_so_far;
+
+                if length_so_far as usize >= values_capacity {
+                    values_buf.reserve(values_capacity);
+                    values_capacity *= 2;
+                }
+
+                values_buf.extend_from_slice(s)
+            });
+        validity = None;
+    } else if !arr.has_validity() {
+        offset_typed
+            .iter_mut()
+            .skip(1)
+            .enumerate()
+            .for_each(|(idx, offset)| {
+                if indices.is_valid(idx) {
+                    let index = indices.value_unchecked(idx) as usize;
+                    let s = arr.value_unchecked(index);
+                    length_so_far += s.len() as i64;
+
+                    if length_so_far as usize >= values_capacity {
+                        values_buf.reserve(values_capacity);
+                        values_capacity *= 2;
+                    }
+
+                    values_buf.extend_from_slice(s)
+                }
+                *offset = length_so_far;
+            });
+        validity = indices.validity().cloned();
+    } else {
+        let mut builder = MutableBinaryArray::with_capacities(data_len, length_so_far as usize);
+        let validity_arr = arr.validity().expect("should have nulls");
+
+        if !indices.has_validity() {
+            (0..data_len).for_each(|idx| {
+                let index = indices.value_unchecked(idx) as usize;
+                builder.push(if validity_arr.get_bit_unchecked(index) {
+                    let s = arr.value_unchecked(index);
+                    Some(s)
+                } else {
+                    None
+                });
+            });
+        } else {
+            let validity_indices = indices.validity().expect("should have nulls");
+            (0..data_len).for_each(|idx| {
+                if validity_indices.get_bit_unchecked(idx) {
+                    let index = indices.value_unchecked(idx) as usize;
+
+                    if validity_arr.get_bit_unchecked(index) {
+                        let s = arr.value_unchecked(index);
+                        builder.push(Some(s));
+                    } else {
+                        builder.push_null();
+                    }
+                } else {
+                    builder.push_null();
+                }
+            });
+        }
+
+        let array: BinaryArray<i64> = builder.into();
+        return Box::new(array);
+    }
+
+    // Safety: all "values" are &str, and thus valid utf8
+    Box::new(BinaryArray::<i64>::from_data_unchecked_default(
+        offset_buf.into(),
+        values_buf.into(),
+        validity,
+    ))
+}
+
 /// Forked and adapted from arrow-rs
 /// This is faster because it does no bounds checks and allocates directly into aligned memory
 ///

polars/polars-arrow/src/data_types.rs

@@ -25,6 +25,7 @@ unsafe impl IsFloat for u16 {}
 unsafe impl IsFloat for u32 {}
 unsafe impl IsFloat for u64 {}
 unsafe impl IsFloat for &str {}
+unsafe impl IsFloat for &[u8] {}
 unsafe impl IsFloat for bool {}
 unsafe impl<T: IsFloat> IsFloat for Option<T> {}
 
@@ -41,6 +42,7 @@ mod private {
     impl Sealed for f32 {}
     impl Sealed for f64 {}
     impl Sealed for &str {}
+    impl Sealed for &[u8] {}
     impl Sealed for bool {}
     impl<T: Sealed> Sealed for Option<T> {}
 }