Skip to content
Closed
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
74 changes: 51 additions & 23 deletions arrow-array/src/array/fixed_size_binary_array.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ use arrow_buffer::buffer::NullBuffer;
use arrow_buffer::{ArrowNativeType, Buffer, MutableBuffer, bit_util};
use arrow_data::{ArrayData, ArrayDataBuilder};
use arrow_schema::{ArrowError, DataType};
use num_traits::{ToPrimitive};
use std::any::Any;
use std::sync::Arc;

Expand Down Expand Up @@ -55,7 +56,7 @@ pub struct FixedSizeBinaryArray {
value_data: Buffer,
nulls: Option<NullBuffer>,
len: usize,
value_length: i32,
value_length: i32, // NB: validated conversion to usize by `try_new`
}

impl FixedSizeBinaryArray {
Expand All @@ -71,7 +72,10 @@ impl FixedSizeBinaryArray {
/// Create a new [`Scalar`] from `value`
pub fn new_scalar(value: impl AsRef<[u8]>) -> Scalar<Self> {
let v = value.as_ref();
Scalar::new(Self::new(v.len() as _, Buffer::from(v), None))
let Some(value_len) = v.len().to_i32() else {
panic!("FixedSizeBinaryArray value length overflows i32");
};
Scalar::new(Self::new(value_len, Buffer::from(v), None))
}

/// Create a new [`FixedSizeBinaryArray`] from the provided parts, returning an error on failure
Expand All @@ -81,17 +85,18 @@ impl FixedSizeBinaryArray {
///
/// # Errors
///
/// * `size < 0`
/// * `value_length` can not be converted to `usize` (is negative)
/// * `value_length < 0`
/// * `values.len() / size != nulls.len()`
/// * `size == 0 && values.len() != 0`
pub fn try_new(
size: i32,
value_length: i32,
values: Buffer,
nulls: Option<NullBuffer>,
) -> Result<Self, ArrowError> {
let data_type = DataType::FixedSizeBinary(size);
let s = size.to_usize().ok_or_else(|| {
ArrowError::InvalidArgumentError(format!("Size cannot be negative, got {size}"))
let data_type = DataType::FixedSizeBinary(value_length);
let s = value_length.to_usize().ok_or_else(|| {
ArrowError::InvalidArgumentError(format!("Size cannot be negative, got {value_length}"))
})?;

let len = match values.len().checked_div(s) {
Expand Down Expand Up @@ -123,7 +128,7 @@ impl FixedSizeBinaryArray {
Ok(Self {
data_type,
value_data: values,
value_length: size,
value_length,
nulls,
len,
})
Expand Down Expand Up @@ -160,7 +165,8 @@ impl FixedSizeBinaryArray {
/// (but still well-defined) if [`is_null`](Self::is_null) returns true for the index.
///
/// # Panics
/// Panics if index `i` is out of bounds.
/// * if index `i` is out of bounds.
/// * if the resulting byte offset for index `i` would overflow `usize`
pub fn value(&self, i: usize) -> &[u8] {
assert!(
i < self.len(),
Expand All @@ -169,12 +175,11 @@ impl FixedSizeBinaryArray {
self.len()
);
let offset = i + self.offset();
let value_size = self.value_size();
let value_offset = offset.checked_mul(value_size).expect("offset overflow");
// SAFETY: value offset/length computed correctly using checked arithmetic
unsafe {
let pos = self.value_offset_at(offset);
std::slice::from_raw_parts(
self.value_data.as_ptr().offset(pos as isize),
(self.value_offset_at(offset + 1) - pos) as usize,
)
std::slice::from_raw_parts(self.value_data.as_ptr().add(value_offset), value_size)
}
}

Expand All @@ -186,7 +191,7 @@ impl FixedSizeBinaryArray {
/// # Safety
///
/// Caller is responsible for ensuring that the index is within the bounds
/// of the array
/// of the array and the resulting byte offset fits in `i32`
pub unsafe fn value_unchecked(&self, i: usize) -> &[u8] {
let offset = i + self.offset();
let pos = self.value_offset_at(offset);
Expand All @@ -208,12 +213,22 @@ impl FixedSizeBinaryArray {

/// Returns the length for an element.
///
/// All elements have the same length as the array is a fixed size.
/// Note this value always fits in a `usize`. See [`Self::value_size`]
#[inline]
pub fn value_length(&self) -> i32 {
self.value_length
}

/// Returns the length for each element as a [`usize`]
///
/// All elements have the same length as the array is a fixed size.
#[inline]
pub fn value_size(&self) -> usize {
// Note: `Self::try_new` verifies that `value_length` can be safely
// converted to `usize` without truncation
self.value_length as usize
}

/// Returns the values of this array.
///
/// Unlike [`Self::value_data`] this returns the [`Buffer`]
Expand All @@ -235,13 +250,17 @@ impl FixedSizeBinaryArray {
"the length + offset of the sliced FixedSizeBinaryArray cannot exceed the existing length"
);

let size = self.value_length as usize;
let scaled_offset = offset
.checked_mul(self.value_size())
.expect("Offset overflowed");

let scaled_len = len.checked_mul(self.value_size()).expect("Len overflowed");

Self {
data_type: self.data_type.clone(),
nulls: self.nulls.as_ref().map(|n| n.slice(offset, len)),
value_length: self.value_length,
value_data: self.value_data.slice_with_length(offset * size, len * size),
value_data: self.value_data.slice_with_length(scaled_offset, scaled_len),
len,
}
}
Expand Down Expand Up @@ -334,12 +353,17 @@ impl FixedSizeBinaryArray {

let nulls = NullBuffer::from_unsliced_buffer(null_buf, len);

let size = size.unwrap_or(0) as i32;
let Some(value_length) = size.unwrap_or(0).to_i32() else {
return Err(ArrowError::InvalidArgumentError(
"FixedSizeBinaryArray value_size overflows i32".to_owned(),
));
};

Ok(Self {
data_type: DataType::FixedSizeBinary(size),
data_type: DataType::FixedSizeBinary(value_length),
value_data: buffer.into(),
nulls,
value_length: size,
value_length,
len,
})
}
Expand Down Expand Up @@ -486,6 +510,8 @@ impl FixedSizeBinaryArray {
})
}

/// Returns the byte offset for the element at index `i` without
/// checking for overflow.
#[inline]
fn value_offset_at(&self, i: usize) -> i32 {
self.value_length * i as i32
Expand All @@ -511,8 +537,10 @@ impl From<ArrayData> for FixedSizeBinaryArray {
_ => panic!("Expected data type to be FixedSizeBinary"),
};

let size = value_length as usize;
let value_data = buffers[0].slice_with_length(offset * size, len * size);
let Some(value_size) = value_length.to_usize() else {
panic!("Invalid SizeBinaryArray value length{value_length}");
};
let value_data = buffers[0].slice_with_length(offset * value_size, len * value_size);

Self {
data_type,
Expand Down
Loading