diff --git a/crates/nano-arrow/src/array/binary/mutable_values.rs b/crates/nano-arrow/src/array/binary/mutable_values.rs index e73f0223ec44..08be36d6f38d 100644 --- a/crates/nano-arrow/src/array/binary/mutable_values.rs +++ b/crates/nano-arrow/src/array/binary/mutable_values.rs @@ -348,7 +348,7 @@ impl> TryPush for MutableBinaryValuesArray { fn try_push(&mut self, value: T) -> Result<()> { let bytes = value.as_ref(); self.values.extend_from_slice(bytes); - self.offsets.try_push_usize(bytes.len()) + self.offsets.try_push(bytes.len()) } } diff --git a/crates/nano-arrow/src/array/list/mutable.rs b/crates/nano-arrow/src/array/list/mutable.rs index 39dc22da3cb0..91c36ff42d21 100644 --- a/crates/nano-arrow/src/array/list/mutable.rs +++ b/crates/nano-arrow/src/array/list/mutable.rs @@ -162,7 +162,7 @@ impl MutableListArray { let offset = self.offsets.last().to_usize(); let length = total_length.checked_sub(offset).ok_or(Error::Overflow)?; - self.offsets.try_push_usize(length)?; + self.offsets.try_push(length)?; if let Some(validity) = &mut self.validity { validity.push(true) } diff --git a/crates/nano-arrow/src/array/physical_binary.rs b/crates/nano-arrow/src/array/physical_binary.rs index 694e61a7ea63..36e4ecf52d35 100644 --- a/crates/nano-arrow/src/array/physical_binary.rs +++ b/crates/nano-arrow/src/array/physical_binary.rs @@ -28,7 +28,7 @@ where if let Some(item) = item? { null.push_unchecked(true); let s = item.as_ref(); - length += O::from_usize(s.len()).unwrap(); + length += O::from_as_usize(s.len()); values.extend_from_slice(s); } else { null.push_unchecked(false); @@ -147,7 +147,7 @@ where for item in iterator { let bytes = item.as_ref(); values.extend_from_slice(bytes); - offsets.try_push_usize(bytes.len()).unwrap(); + offsets.try_push(bytes.len()).unwrap(); } offsets.len_proxy() - start_index } @@ -205,7 +205,7 @@ where for item in iterator { let s = item.as_ref(); values.extend_from_slice(s); - offsets.try_push_usize(s.len()).unwrap(); + offsets.try_push(s.len()).unwrap(); } (offsets, values) } diff --git a/crates/nano-arrow/src/array/utf8/mutable_values.rs b/crates/nano-arrow/src/array/utf8/mutable_values.rs index 8810d30febb5..f500bb79877f 100644 --- a/crates/nano-arrow/src/array/utf8/mutable_values.rs +++ b/crates/nano-arrow/src/array/utf8/mutable_values.rs @@ -395,7 +395,7 @@ impl> TryPush for MutableUtf8ValuesArray { fn try_push(&mut self, value: T) -> Result<()> { let bytes = value.as_ref().as_bytes(); self.values.extend_from_slice(bytes); - self.offsets.try_push_usize(bytes.len()) + self.offsets.try_push(bytes.len()) } } diff --git a/crates/nano-arrow/src/compute/cast/binary_to.rs b/crates/nano-arrow/src/compute/cast/binary_to.rs index 52038f9caefa..4f5a1fb2b610 100644 --- a/crates/nano-arrow/src/compute/cast/binary_to.rs +++ b/crates/nano-arrow/src/compute/cast/binary_to.rs @@ -123,7 +123,7 @@ pub(super) fn binary_to_dictionary_dyn( fn fixed_size_to_offsets(values_len: usize, fixed_size: usize) -> Offsets { let offsets = (0..(values_len + 1)) .step_by(fixed_size) - .map(|v| O::from_usize(v).unwrap()) + .map(|v| O::from_as_usize(v)) .collect(); // Safety // * every element is `>= 0` diff --git a/crates/nano-arrow/src/compute/cast/primitive_to.rs b/crates/nano-arrow/src/compute/cast/primitive_to.rs index a83569ee165c..661e40bc6343 100644 --- a/crates/nano-arrow/src/compute/cast/primitive_to.rs +++ b/crates/nano-arrow/src/compute/cast/primitive_to.rs @@ -33,7 +33,7 @@ pub fn primitive_to_binary( let len = lexical_core::write_unchecked(*x, bytes).len(); offset += len; - offsets.push(O::from_usize(offset).unwrap()); + offsets.push(O::from_as_usize(offset)); } values.set_len(offset); values.shrink_to_fit(); @@ -101,7 +101,7 @@ pub fn primitive_to_utf8( let len = lexical_core::write_unchecked(*x, bytes).len(); offset += len; - offsets.push(O::from_usize(offset).unwrap()); + offsets.push(O::from_as_usize(offset)); } values.set_len(offset); values.shrink_to_fit(); diff --git a/crates/nano-arrow/src/io/avro/read/nested.rs b/crates/nano-arrow/src/io/avro/read/nested.rs index 056d9a8f836e..fd5bb6b7dbbd 100644 --- a/crates/nano-arrow/src/io/avro/read/nested.rs +++ b/crates/nano-arrow/src/io/avro/read/nested.rs @@ -36,7 +36,7 @@ impl DynMutableListArray { let offset = self.offsets.last().to_usize(); let length = total_length.checked_sub(offset).ok_or(Error::Overflow)?; - self.offsets.try_push_usize(length)?; + self.offsets.try_push(length)?; if let Some(validity) = &mut self.validity { validity.push(true) } diff --git a/crates/nano-arrow/src/io/parquet/read/deserialize/binary/utils.rs b/crates/nano-arrow/src/io/parquet/read/deserialize/binary/utils.rs index 0a2a0f3466f8..961268db2beb 100644 --- a/crates/nano-arrow/src/io/parquet/read/deserialize/binary/utils.rs +++ b/crates/nano-arrow/src/io/parquet/read/deserialize/binary/utils.rs @@ -19,7 +19,7 @@ impl Pushable for Offsets { #[inline] fn push(&mut self, value: usize) { - self.try_push_usize(value).unwrap() + self.try_push(value).unwrap() } #[inline] @@ -53,7 +53,7 @@ impl Binary { } self.values.extend(v); - self.offsets.try_push_usize(v.len()).unwrap() + self.offsets.try_push(v.len()).unwrap() } #[inline] diff --git a/crates/nano-arrow/src/offset.rs b/crates/nano-arrow/src/offset.rs index 409e695ba66a..5bd06aeb7e57 100644 --- a/crates/nano-arrow/src/offset.rs +++ b/crates/nano-arrow/src/offset.rs @@ -76,7 +76,7 @@ impl Offsets { let (lower, _) = iterator.size_hint(); let mut offsets = Self::with_capacity(lower); for item in iterator { - offsets.try_push_usize(item)? + offsets.try_push(item)? } Ok(offsets) } @@ -103,20 +103,6 @@ impl Offsets { self.0.shrink_to_fit(); } - /// Pushes a new element with a given length. - /// # Error - /// This function errors iff the new last item is larger than what `O` supports. - /// # Panic - /// This function asserts that `length > 0`. - #[inline] - pub fn try_push(&mut self, length: O) -> Result<(), Error> { - let old_length = self.last(); - assert!(length >= O::zero()); - let new_length = old_length.checked_add(&length).ok_or(Error::Overflow)?; - self.0.push(new_length); - Ok(()) - } - /// Pushes a new element with a given length. /// # Error /// This function errors iff the new last item is larger than what `O` supports. @@ -124,13 +110,21 @@ impl Offsets { /// This function: /// * checks that this length does not overflow #[inline] - pub fn try_push_usize(&mut self, length: usize) -> Result<(), Error> { - let length = O::from_usize(length).ok_or(Error::Overflow)?; + pub fn try_push(&mut self, length: usize) -> Result<(), Error> { + if O::IS_LARGE { + let length = O::from_as_usize(length); + let old_length = self.last(); + let new_length = *old_length + length; + self.0.push(new_length); + Ok(()) + } else { + let length = O::from_usize(length).ok_or(Error::Overflow)?; - let old_length = self.last(); - let new_length = old_length.checked_add(&length).ok_or(Error::Overflow)?; - self.0.push(new_length); - Ok(()) + let old_length = self.last(); + let new_length = old_length.checked_add(&length).ok_or(Error::Overflow)?; + self.0.push(new_length); + Ok(()) + } } /// Returns [`Offsets`] assuming that `offsets` fulfills its invariants diff --git a/crates/polars-json/src/json/deserialize.rs b/crates/polars-json/src/json/deserialize.rs index a27a14184b8e..0ce3831a46eb 100644 --- a/crates/polars-json/src/json/deserialize.rs +++ b/crates/polars-json/src/json/deserialize.rs @@ -85,7 +85,7 @@ fn deserialize_list<'a, A: Borrow>>( inner.extend(value.iter()); validity.push(true); offsets - .try_push_usize(value.len()) + .try_push(value.len()) .expect("List offset is too large :/"); }, BorrowedValue::Static(StaticNode::Null) => { @@ -95,9 +95,7 @@ fn deserialize_list<'a, A: Borrow>>( value @ (BorrowedValue::Static(_) | BorrowedValue::String(_)) => { inner.push(value); validity.push(true); - offsets - .try_push_usize(1) - .expect("List offset is too large :/"); + offsets.try_push(1).expect("List offset is too large :/"); }, _ => { validity.push(false);