Gate PiecewiseSequence take paths on unit multipliers

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

Author: danking

vortex-array/benches/take_fsl.rs

@@ -26,7 +26,7 @@ use vortex_array::VortexSessionExecute;
 use vortex_array::array_session;
 use vortex_array::arrays::ConstantArray;
 use vortex_array::arrays::FixedSizeListArray;
-use vortex_array::arrays::PiecewiseSequentialArray;
+use vortex_array::arrays::PiecewiseSequenceArray;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt;
 use vortex_array::dtype::IntegerPType;
@@ -143,7 +143,7 @@ fn take_fsl_f16_force_per_index<const LIST_SIZE: usize>(bencher: Bencher, num_in
 }
 
 #[divan::bench(args = NUM_INDICES, consts = F16_STRATEGY_LIST_SIZES)]
-fn take_fsl_f16_force_piecewise_sequential<const LIST_SIZE: usize>(
+fn take_fsl_f16_force_piecewise_sequence<const LIST_SIZE: usize>(
     bencher: Bencher,
     num_indices: usize,
 ) {
@@ -154,7 +154,7 @@ fn take_fsl_f16_force_piecewise_sequential<const LIST_SIZE: usize>(
         .counter(BytesCount::of_many::<f16>(num_indices * LIST_SIZE))
         .with_inputs(|| (&fsl, &indices, SESSION.create_execution_ctx()))
         .bench_refs(|(array, indices, execution_ctx)| {
-            take_fsl_f16_piecewise_sequential_strategy::<LIST_SIZE>(array, indices)
+            take_fsl_f16_piecewise_sequence_strategy::<LIST_SIZE>(array, indices)
                 .into_array()
                 .execute::<RecursiveCanonical>(execution_ctx)
                 .unwrap()
@@ -211,7 +211,7 @@ fn take_fsl_f16_per_index_strategy<const LIST_SIZE: usize, E: IntegerPType>(
     }
 }
 
-fn take_fsl_f16_piecewise_sequential_strategy<const LIST_SIZE: usize>(
+fn take_fsl_f16_piecewise_sequence_strategy<const LIST_SIZE: usize>(
     array: &FixedSizeListArray,
     indices: &Buffer<u64>,
 ) -> FixedSizeListArray {
@@ -223,11 +223,17 @@ fn take_fsl_f16_piecewise_sequential_strategy<const LIST_SIZE: usize>(
     let run_count = starts.len();
     let starts = PrimitiveArray::from_iter(starts).into_array();
     let lengths = ConstantArray::new(LIST_SIZE as u64, run_count).into_array();
+    let multipliers = ConstantArray::new(1u64, run_count).into_array();
 
-    // SAFETY: benchmark indices are generated in-bounds, lengths is a non-nullable unsigned
-    // constant, and output length is exactly `indices.len() * LIST_SIZE`.
+    // SAFETY: benchmark indices are generated in-bounds; lengths and multiplier 1 are
+    // non-nullable unsigned constants; output length is exactly `indices.len() * LIST_SIZE`.
     let element_indices = unsafe {
-        PiecewiseSequentialArray::new_unchecked(starts, lengths, indices.len() * LIST_SIZE)
+        PiecewiseSequenceArray::new_unchecked(
+            starts,
+            lengths,
+            multipliers,
+            indices.len() * LIST_SIZE,
+        )
     }
     .into_array();
     let elements = array.elements().take(element_indices).unwrap();

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

@@ -75,12 +75,10 @@ fn take_piecewise_sequence(
                 let validity = array.validity()?.take(indices_ref)?;
 
                 // SAFETY: contiguous gather preserves the decimal dtype and value representation.
-                Ok(
-                    Some(unsafe {
-                        DecimalArray::new_unchecked(values, array.decimal_dtype(), validity)
-                    }
-                    .into_array()),
-                )
+                Ok(Some(
+                    unsafe { DecimalArray::new_unchecked(values, array.decimal_dtype(), validity) }
+                        .into_array(),
+                ))
             })
         })
     })

vortex-array/src/arrays/piecewise_sequence/mod.rs

@@ -163,16 +163,16 @@ where
     for (&start, &length) in starts.iter().zip_eq(lengths) {
         let start: usize = start.as_();
         let length: usize = length.as_();
-        let end = start.checked_add(length).ok_or_else(|| {
-            vortex_err!("PiecewiseSequenceArray range overflows usize")
-        })?;
+        let end = start
+            .checked_add(length)
+            .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?;
         vortex_ensure!(
             end <= source_len,
             "PiecewiseSequenceArray range {start}..{end} exceeds source length {source_len}"
         );
-        computed_len = computed_len.checked_add(length).ok_or_else(|| {
-            vortex_err!("PiecewiseSequenceArray output length overflows usize")
-        })?;
+        computed_len = computed_len
+            .checked_add(length)
+            .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?;
     }
 
     vortex_ensure!(

vortex-array/src/arrays/piecewise_sequence/tests.rs

@@ -187,6 +187,26 @@ fn primitive_take_consumes_piecewise_indices() -> VortexResult<()> {
     Ok(())
 }
 
+#[test]
+fn primitive_take_handles_non_unit_multiplier() -> VortexResult<()> {
+    let values = PrimitiveArray::from_iter(0i32..20).into_array();
+    let indices = PiecewiseSequenceArray::try_new(
+        buffer![3u64].into_array(),
+        buffer![3u64].into_array(),
+        buffer![2u64].into_array(),
+        3,
+    )?
+    .into_array();
+    let taken = values.take(indices)?;
+
+    assert_arrays_eq!(
+        taken,
+        PrimitiveArray::from_iter([3i32, 5, 7]).into_array(),
+        &mut array_session().create_execution_ctx()
+    );
+    Ok(())
+}
+
 #[test]
 fn bool_take_consumes_piecewise_indices() -> VortexResult<()> {
     let values = BoolArray::from_iter([true, false, true, true, false, false]).into_array();

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

@@ -93,13 +93,15 @@ fn take_piecewise_sequence(
     // SAFETY: ranges were validated against the source views, and copied views still reference the
     // same backing data buffers.
     unsafe {
-        Ok(Some(VarBinViewArray::new_handle_unchecked(
-            BufferHandle::new_host(views.into_byte_buffer()),
-            Arc::clone(array.data_buffers()),
-            array.dtype().clone(),
-            validity,
-        )
-        .into_array()))
+        Ok(Some(
+            VarBinViewArray::new_handle_unchecked(
+                BufferHandle::new_host(views.into_byte_buffer()),
+                Arc::clone(array.data_buffers()),
+                array.dtype().clone(),
+                validity,
+            )
+            .into_array(),
+        ))
     }
 }