Skip to content

Commit 8859ff7

Browse files
committed
Adding try_append_value implementation.
1 parent 597c903 commit 8859ff7

File tree

1 file changed

+45
-6
lines changed

1 file changed

+45
-6
lines changed

arrow-array/src/builder/generic_bytes_view_builder.rs

Lines changed: 45 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -306,15 +306,30 @@ impl<T: ByteViewType + ?Sized> GenericByteViewBuilder<T> {
306306
/// - String length exceeds `u32::MAX`
307307
#[inline]
308308
pub fn append_value(&mut self, value: impl AsRef<T::Native>) {
309+
self.try_append_value(value).unwrap()
310+
}
311+
312+
/// Appends a value into the builder
313+
///
314+
/// # Errors
315+
///
316+
/// Returns an error if:
317+
/// - String buffer count exceeds `u32::MAX`
318+
/// - String length exceeds `u32::MAX`
319+
#[inline]
320+
pub fn try_append_value(&mut self, value: impl AsRef<T::Native>) -> Result<(), ArrowError> {
309321
let v: &[u8] = value.as_ref().as_ref();
310-
let length: u32 = v.len().try_into().unwrap();
322+
let length: u32 = v.len().try_into().map_err(|_| {
323+
ArrowError::InvalidArgumentError(format!("String length {} exceeds u32::MAX", v.len()))
324+
})?;
325+
311326
if length <= MAX_INLINE_VIEW_LEN {
312327
let mut view_buffer = [0; 16];
313328
view_buffer[0..4].copy_from_slice(&length.to_le_bytes());
314329
view_buffer[4..4 + v.len()].copy_from_slice(v);
315330
self.views_buffer.push(u128::from_le_bytes(view_buffer));
316331
self.null_buffer_builder.append_non_null();
317-
return;
332+
return Ok(());
318333
}
319334

320335
// Deduplication if:
@@ -339,7 +354,7 @@ impl<T: ByteViewType + ?Sized> GenericByteViewBuilder<T> {
339354
self.views_buffer.push(self.views_buffer[*idx]);
340355
self.null_buffer_builder.append_non_null();
341356
self.string_tracker = Some((ht, hasher));
342-
return;
357+
return Ok(());
343358
}
344359
Entry::Vacant(vacant) => {
345360
// o.w. we insert the (string hash -> view index)
@@ -356,17 +371,41 @@ impl<T: ByteViewType + ?Sized> GenericByteViewBuilder<T> {
356371
let to_reserve = v.len().max(self.block_size.next_size() as usize);
357372
self.in_progress.reserve(to_reserve);
358373
};
359-
let offset = self.in_progress.len() as u32;
374+
let offset: u32 = self.in_progress.len().try_into().map_err(|_| {
375+
ArrowError::InvalidArgumentError(format!(
376+
"In-progress buffer length {} exceeds u32::MAX",
377+
self.in_progress.len()
378+
))
379+
})?;
360380
self.in_progress.extend_from_slice(v);
361381

382+
let prefix = v
383+
.get(0..4)
384+
.and_then(|slice| slice.try_into().ok())
385+
.map(u32::from_le_bytes)
386+
.ok_or_else(|| {
387+
ArrowError::InvalidArgumentError(
388+
"String must be at least 4 bytes for non-inline view".to_string(),
389+
)
390+
})?;
391+
392+
let buffer_index: u32 = self.completed.len().try_into().map_err(|_| {
393+
ArrowError::InvalidArgumentError(format!(
394+
"Buffer count {} exceeds u32::MAX",
395+
self.completed.len()
396+
))
397+
})?;
398+
362399
let view = ByteView {
363400
length,
364-
prefix: u32::from_le_bytes(v[0..4].try_into().unwrap()),
365-
buffer_index: self.completed.len() as u32,
401+
prefix,
402+
buffer_index,
366403
offset,
367404
};
368405
self.views_buffer.push(view.into());
369406
self.null_buffer_builder.append_non_null();
407+
408+
Ok(())
370409
}
371410

372411
/// Append an `Option` value into the builder

0 commit comments

Comments
 (0)