Validate PiecewiseSequence buffer lengths after copy

Signed-off-by: Daniel King <dan@spiraldb.com>

Author: danking

vortex-array/src/arrays/bool/compute/take.rs

@@ -175,19 +175,16 @@ where
     L: UnsignedPType,
 {
     let mut values = BitBufferMut::with_capacity(output_len);
-    let mut computed_len = 0usize;
     for (&start, &length) in starts.iter().zip_eq(lengths) {
         let start = start.as_();
         let length = length.as_();
-        computed_len = computed_len
-            .checked_add(length)
-            .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?;
         values.append_buffer(&source.slice(start..).slice(..length));
     }
 
     vortex_ensure!(
-        computed_len == output_len,
-        "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}"
+        values.len() == output_len,
+        "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}",
+        values.len()
     );
     Ok(values.freeze())
 }

vortex-array/src/arrays/decimal/compute/take.rs

@@ -216,19 +216,16 @@ where
     T: NativeDecimalType,
 {
     let mut result = BufferMut::<T>::with_capacity(output_len);
-    let mut computed_len = 0usize;
     for (&start, &length) in starts.iter().zip_eq(lengths) {
         let start = start.as_();
         let length = length.as_();
-        computed_len = computed_len
-            .checked_add(length)
-            .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?;
         result.extend_from_slice(&values[start..][..length]);
     }
 
     vortex_ensure!(
-        computed_len == output_len,
-        "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}"
+        result.len() == output_len,
+        "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}",
+        result.len()
     );
     Ok(result.freeze())
 }

vortex-array/src/arrays/primitive/compute/take/mod.rs

@@ -242,19 +242,16 @@ where
     L: UnsignedPType,
 {
     let mut values = BufferMut::<T>::with_capacity(output_len);
-    let mut computed_len = 0usize;
     for (&start, &length) in starts.iter().zip_eq(lengths) {
         let start = start.as_();
         let length = length.as_();
-        computed_len = computed_len
-            .checked_add(length)
-            .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?;
         values.extend_from_slice(&source[start..][..length]);
     }
 
     vortex_ensure!(
-        computed_len == output_len,
-        "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}"
+        values.len() == output_len,
+        "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}",
+        values.len()
     );
     Ok(values.freeze())
 }

vortex-array/src/arrays/varbin/compute/take.rs

@@ -505,14 +505,10 @@ where
     let mut new_offsets = BufferMut::<NewOffset>::with_capacity(output_len + 1);
     new_offsets.push(NewOffset::zero());
     let mut output_bytes = 0usize;
-    let mut computed_len = 0usize;
 
     for (&start, &length) in starts.iter().zip_eq(lengths) {
         let start = start.as_();
         let length = length.as_();
-        computed_len = computed_len
-            .checked_add(length)
-            .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?;
         if length == 0 {
             continue;
         }
@@ -542,8 +538,9 @@ where
             .ok_or_else(|| vortex_err!("PiecewiseSequence VarBin output byte length overflow"))?;
     }
     vortex_ensure!(
-        computed_len == output_len,
-        "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}"
+        new_offsets.len() == output_len + 1,
+        "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}",
+        new_offsets.len() - 1
     );
 
     let mut new_data = ByteBufferMut::with_capacity(output_bytes);

vortex-array/src/arrays/varbinview/compute/take.rs

@@ -190,19 +190,16 @@ where
     L: UnsignedPType,
 {
     let mut views = BufferMut::<BinaryView>::with_capacity(output_len);
-    let mut computed_len = 0usize;
     for (&start, &length) in starts.iter().zip_eq(lengths) {
         let start = start.as_();
         let length = length.as_();
-        computed_len = computed_len
-            .checked_add(length)
-            .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?;
         views.extend_from_slice(&source[start..][..length]);
     }
 
     vortex_ensure!(
-        computed_len == output_len,
-        "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}"
+        views.len() == output_len,
+        "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}",
+        views.len()
     );
     Ok(views.freeze())
 }