feat: elide overflow checks on i64 (#11563)

crates/nano-arrow/src/array/binary/mutable_values.rs

@@ -348,7 +348,7 @@ impl<O: Offset, T: AsRef<[u8]>> TryPush<T> for MutableBinaryValuesArray<O> {
     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())
     }
 }
 

crates/nano-arrow/src/array/list/mutable.rs

@@ -162,7 +162,7 @@ impl<O: Offset, M: MutableArray> MutableListArray<O, M> {
         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)
         }

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)
 }

crates/nano-arrow/src/array/utf8/mutable_values.rs

@@ -395,7 +395,7 @@ impl<O: Offset, T: AsRef<str>> TryPush<T> for MutableUtf8ValuesArray<O> {
     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())
     }
 }
 

crates/nano-arrow/src/compute/cast/binary_to.rs

@@ -123,7 +123,7 @@ pub(super) fn binary_to_dictionary_dyn<O: Offset, K: DictionaryKey>(
 fn fixed_size_to_offsets<O: Offset>(values_len: usize, fixed_size: usize) -> Offsets<O> {
     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`

crates/nano-arrow/src/compute/cast/primitive_to.rs

@@ -33,7 +33,7 @@ pub fn primitive_to_binary<T: NativeType + lexical_core::ToLexical, O: Offset>(
             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<T: NativeType + lexical_core::ToLexical, O: Offset>(
             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();

crates/nano-arrow/src/io/avro/read/nested.rs

@@ -36,7 +36,7 @@ impl<O: Offset> DynMutableListArray<O> {
         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)
         }

crates/nano-arrow/src/io/parquet/read/deserialize/binary/utils.rs

@@ -19,7 +19,7 @@ impl<O: Offset> Pushable<usize> for Offsets<O> {
 
     #[inline]
     fn push(&mut self, value: usize) {
-        self.try_push_usize(value).unwrap()
+        self.try_push(value).unwrap()
     }
 
     #[inline]
@@ -53,7 +53,7 @@ impl<O: Offset> Binary<O> {
         }
 
         self.values.extend(v);
-        self.offsets.try_push_usize(v.len()).unwrap()
+        self.offsets.try_push(v.len()).unwrap()
     }
 
     #[inline]

crates/nano-arrow/src/offset.rs

@@ -76,7 +76,7 @@ impl<O: Offset> Offsets<O> {
         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,34 +103,28 @@ impl<O: Offset> Offsets<O> {
         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.
     /// # Implementation
     /// 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

crates/polars-json/src/json/deserialize.rs

@@ -85,7 +85,7 @@ fn deserialize_list<'a, A: Borrow<BorrowedValue<'a>>>(
             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<BorrowedValue<'a>>>(
         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);