feat: Reinterpret binary data to fixed size numerical array

crates/polars-compute/src/cast/binview_to.rs

@@ -1,12 +1,15 @@
+use std::ptr::copy_nonoverlapping;
+
 use arrow::array::*;
+use arrow::bitmap::MutableBitmap;
 #[cfg(feature = "dtype-decimal")]
 use arrow::compute::decimal::deserialize_decimal;
-use arrow::datatypes::{ArrowDataType, TimeUnit};
+use arrow::datatypes::{ArrowDataType, Field, TimeUnit};
 use arrow::offset::Offset;
 use arrow::types::NativeType;
 use chrono::Datelike;
 use num_traits::FromBytes;
-use polars_error::PolarsResult;
+use polars_error::{PolarsResult, polars_err};
 
 use super::CastOptionsImpl;
 use super::binary_to::Parse;
@@ -191,3 +194,135 @@ where
         is_little_endian,
     )))
 }
+
+/// Casts a [`BinaryViewArray`] to a [`FixedSizeListArray`], making any un-castable value a Null.
+///
+/// # Arguments
+///
+/// * `from`: The array to reinterpret.
+/// * `array_width`: The number of items in each `Array`.
+pub(super) fn try_binview_to_fixed_size_list<T, const IS_LITTLE_ENDIAN: bool>(
+    from: &BinaryViewArray,
+    array_width: usize,
+) -> PolarsResult<FixedSizeListArray>
+where
+    T: FromBytes + NativeType,
+    for<'a> &'a <T as FromBytes>::Bytes: TryFrom<&'a [u8]>,
+{
+    let element_size = std::mem::size_of::<T>();
+    // The maximum number of primitives in the result:
+    let primitive_length = from.len().checked_mul(array_width).ok_or_else(|| {
+        polars_err!(
+            InvalidOperation:
+            "array chunk length * number of items ({} * {}) is too large",
+            from.len(),
+            array_width
+        )
+    })?;
+    // The size of each array, in bytes:
+    let array_bytes_size = element_size.checked_mul(array_width).ok_or_else(|| {
+        polars_err!(
+            InvalidOperation:
+            "array size in bytes ({} * {}) is too large",
+            element_size,
+            array_width
+        )
+    })?;
+    let mut out: Vec<T> = Vec::with_capacity(primitive_length);
+    let mut validity = MutableBitmap::from_len_set(from.len());
+
+    for (index, value) in from.iter().enumerate() {
+        if let Some(value) = value
+            && value.len() == array_bytes_size
+        {
+            if cfg!(target_endian = "little") && IS_LITTLE_ENDIAN {
+                // Fast path, we can just copy the data with no need to
+                // reinterpret.
+                let write_index = array_width * index;
+                debug_assert!(write_index < primitive_length);
+                debug_assert!((write_index + (array_width - 1)) < primitive_length);
+                // # Safety
+                // - The target index is smaller than the vector's pre-allocated
+                //   capacity.
+                // - We made sure `value` has byte length
+                //   `array_width * element_size`.
+                unsafe {
+                    copy_nonoverlapping(
+                        value.as_ptr(),
+                        out.as_mut_ptr().add(write_index) as *mut u8,
+                        value.len(),
+                    );
+                }
+            } else {
+                // Slow path, reinterpret items one by one.
+                for j in 0..array_width {
+                    let jth_range = (j * element_size)..((j + 1) * element_size);
+                    debug_assert!(value.get(jth_range.clone()).is_some());
+                    // # Safety
+                    // We made sure the range is smaller than `value` length.
+                    let jth_bytes = unsafe { value.get_unchecked(jth_range) };
+                    // # Safety
+                    // We just made sure that the slice has length `element_size`
+                    let byte_array = unsafe { jth_bytes.try_into().unwrap_unchecked() };
+                    let jth_value = if IS_LITTLE_ENDIAN {
+                        <T as FromBytes>::from_le_bytes(byte_array)
+                    } else {
+                        <T as FromBytes>::from_be_bytes(byte_array)
+                    };
+
+                    let write_index = array_width * index + j;
+                    debug_assert!(write_index < primitive_length);
+                    // # Safety
+                    // - The target index is smaller than the vector's pre-allocated capacity.
+                    unsafe {
+                        std::ptr::write(out.as_mut_ptr().add(write_index), jth_value);
+                    }
+                }
+            }
+        } else {
+            validity.set(index, false);
+        };
+    }
+
+    // # Safety
+    // `out` was created with capacity primitive_length.
+    unsafe { out.set_len(primitive_length) };
+
+    FixedSizeListArray::try_new(
+        ArrowDataType::FixedSizeList(
+            Box::new(Field::new("".into(), T::PRIMITIVE.into(), true)),
+            array_width,
+        ),
+        from.len(),
+        Box::new(PrimitiveArray::<T>::from_vec(out)),
+        validity.into(),
+    )
+}
+
+/// Casts a `dyn` [`Array`] to a [`FixedSizeListArray`], making any un-castable value a Null.
+///
+/// # Arguments
+///
+/// * `from`: The array to reinterpret.
+/// * `array_width`: The number of items in each `Array`.
+///
+/// # Panics
+///    Panics if `from` is not `BinaryViewArray`.
+pub fn binview_to_fixed_size_list_dyn<T>(
+    from: &dyn Array,
+    array_width: usize,
+    is_little_endian: bool,
+) -> PolarsResult<Box<dyn Array>>
+where
+    T: FromBytes + NativeType,
+    for<'a> &'a <T as FromBytes>::Bytes: TryFrom<&'a [u8]>,
+{
+    let from = from.as_any().downcast_ref().unwrap();
+
+    let result = if is_little_endian {
+        try_binview_to_fixed_size_list::<T, true>(from, array_width)
+    } else {
+        try_binview_to_fixed_size_list::<T, false>(from, array_width)
+    }?;
+    Ok(Box::new(result))
+}

crates/polars-compute/src/cast/mod.rs

@@ -21,11 +21,11 @@ use arrow::array::*;
 use arrow::datatypes::*;
 use arrow::match_integer_type;
 use arrow::offset::{Offset, Offsets};
-pub use binview_to::binview_to_primitive_dyn;
 use binview_to::{
     binview_to_dictionary, utf8view_to_date32_dyn, utf8view_to_dictionary,
     utf8view_to_naive_timestamp_dyn, view_to_binary,
 };
+pub use binview_to::{binview_to_fixed_size_list_dyn, binview_to_primitive_dyn};
 use dictionary_to::*;
 use polars_error::{PolarsResult, polars_bail, polars_ensure, polars_err};
 use polars_utils::IdxSize;

crates/polars-ops/src/chunked_array/binary/namespace.rs

@@ -1,13 +1,14 @@
 #[cfg(feature = "binary_encoding")]
 use std::borrow::Cow;
 
-use arrow::with_match_primitive_type;
+#[cfg(feature = "binary_encoding")]
+use arrow::array::Array;
 #[cfg(feature = "binary_encoding")]
 use base64::Engine as _;
 #[cfg(feature = "binary_encoding")]
 use base64::engine::general_purpose;
 use memchr::memmem::find;
-use polars_compute::cast::binview_to_primitive_dyn;
+use polars_compute::cast::{binview_to_fixed_size_list_dyn, binview_to_primitive_dyn};
 use polars_compute::size::binary_size_bytes;
 use polars_core::prelude::arity::{broadcast_binary_elementwise_values, unary_elementwise_values};
 
@@ -156,29 +157,63 @@ pub trait BinaryNameSpaceImpl: AsBinary {
 
     #[cfg(feature = "binary_encoding")]
     fn reinterpret(&self, dtype: &DataType, is_little_endian: bool) -> PolarsResult<Series> {
+        unsafe {
+            Ok(Series::from_chunks_and_dtype_unchecked(
+                self.as_binary().name().clone(),
+                self._reinterpret_inner(dtype, is_little_endian)?,
+                dtype,
+            ))
+        }
+    }
+
+    #[cfg(feature = "binary_encoding")]
+    fn _reinterpret_inner(
+        &self,
+        dtype: &DataType,
+        is_little_endian: bool,
+    ) -> PolarsResult<Vec<Box<dyn Array>>> {
+        use polars_core::with_match_physical_numeric_polars_type;
+
         let ca = self.as_binary();
-        let arrow_type = dtype.to_arrow(CompatLevel::newest());
 
-        match arrow_type.to_physical_type() {
-            arrow::datatypes::PhysicalType::Primitive(ty) => {
-                with_match_primitive_type!(ty, |$T| {
+        match dtype {
+            dtype if dtype.is_primitive_numeric() || dtype.is_temporal() => {
+                let dtype = dtype.to_physical();
+                let arrow_data_type = dtype
+                    .to_arrow(CompatLevel::newest())
+                    .underlying_physical_type();
+                with_match_physical_numeric_polars_type!(dtype, |$T| {
                     unsafe {
-                        Ok(Series::from_chunks_and_dtype_unchecked(
-                            ca.name().clone(),
-                            ca.chunks().iter().map(|chunk| {
-                                binview_to_primitive_dyn::<$T>(
-                                    &**chunk,
-                                    &arrow_type,
-                                    is_little_endian,
-                                )
-                            }).collect::<PolarsResult<Vec<_>>>()?,
-                            dtype
-                        ))
+                        ca.chunks().iter().map(|chunk| {
+                            binview_to_primitive_dyn::<<$T as PolarsNumericType>::Native>(
+                                &**chunk,
+                                &arrow_data_type,
+                                is_little_endian,
+                            )
+                        }).collect()
                     }
                 })
             },
+            #[cfg(feature = "dtype-array")]
+            DataType::Array(inner_dtype, array_width)
+                if inner_dtype.is_primitive_numeric() || inner_dtype.is_temporal() =>
+            {
+                let inner_dtype = inner_dtype.to_physical();
+                let result: Vec<ArrayRef> = with_match_physical_numeric_polars_type!(inner_dtype, |$T| {
+                    unsafe {
+                        ca.chunks().iter().map(|chunk| {
+                            binview_to_fixed_size_list_dyn::<<$T as PolarsNumericType>::Native>(
+                                &**chunk,
+                                *array_width,
+                                is_little_endian
+                            )
+                        }).collect::<Result<Vec<ArrayRef>, _>>()
+                    }
+                })?;
+                Ok(result)
+            },
             _ => Err(
-                polars_err!(InvalidOperation:"unsupported data type in reinterpret. Only numerical types are allowed."),
+                polars_err!(InvalidOperation: "unsupported data type {:?} in reinterpret. Only numeric or temporal types, or Arrays of those, are allowed.", dtype),
             ),
         }
     }

crates/polars-plan/src/dsl/function_expr/binary.rs

@@ -20,6 +20,8 @@ pub enum BinaryFunction {
     Base64Encode,
     Size,
     #[cfg(feature = "binary_encoding")]
+    /// The parameters are destination type, and whether to use little endian
+    /// encoding.
     Reinterpret(DataTypeExpr, bool),
 }
 

crates/polars-plan/src/plans/conversion/dsl_to_ir/functions.rs

@@ -112,8 +112,19 @@ pub(super) fn convert_functions(
                 B::Base64Encode => IB::Base64Encode,
                 B::Size => IB::Size,
                 #[cfg(feature = "binary_encoding")]
-                B::Reinterpret(data_type, v) => {
-                    IB::Reinterpret(data_type.into_datatype(ctx.schema)?, v)
+                B::Reinterpret(dtype_expr, v) => {
+                    let dtype = dtype_expr.into_datatype(ctx.schema)?;
+                    let can_reinterpret_to =
+                        |dt: &DataType| dt.is_primitive_numeric() || dt.is_temporal();
+                    polars_ensure!(
+                        can_reinterpret_to(&dtype) || (
+                            dtype.is_array() && dtype.inner_dtype().map(can_reinterpret_to) == Some(true)
+                        ),
+                        InvalidOperation:
+                        "cannot reinterpret binary to dtype {:?}. Only numeric or temporal dtype, or Arrays of these, are supported. Hint: To reinterpret to a nested Array, first reinterpret to a linear Array, and then use reshape",
+                        dtype
+                    );
+                    IB::Reinterpret(dtype, v)
                 },
             })
         },

py-polars/polars/expr/binary.py

@@ -301,7 +301,10 @@ def reinterpret(
         self, *, dtype: PolarsDataType | DataTypeExpr, endianness: Endianness = "little"
     ) -> Expr:
         r"""
-        Interpret a buffer as a numerical Polars type.
+        Interpret bytes as another type.
+
+        Supported types are numerical or temporal dtypes, or an ``Array`` of
+        these dtypes.
 
         Parameters
         ----------
@@ -314,8 +317,9 @@ def reinterpret(
         -------
         Expr
             Expression of data type `dtype`.
-            Note that if binary array is too short value will be null.
-            If binary array is too long, remainder will be ignored.
+            Note that rows of the binary array where the length does not match
+            the size in bytes of the output array (number of items * byte size
+            of item) will become NULL.
 
         Examples
         --------

py-polars/polars/series/binary.py

@@ -220,7 +220,10 @@ def reinterpret(
         self, *, dtype: PolarsDataType, endianness: Endianness = "little"
     ) -> Series:
         r"""
-        Interpret a buffer as a numerical polars type.
+        Interpret bytes as another type.
+
+        Supported types are numerical or temporal dtypes, or an ``Array`` of
+        these dtypes.
 
         Parameters
         ----------
@@ -233,8 +236,9 @@ def reinterpret(
         -------
         Series
             Series of data type `dtype`.
-            Note that if binary array is too short value will be null.
-            If binary array is too long, remainder will be ignored.
+            Note that rows of the binary array where the length does not match
+            the size in bytes of the output array (number of items * byte size
+            of item) will become NULL.
 
         Examples
         --------