Specialize constant PiecewiseSequence runs

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

Author: danking

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

@@ -20,8 +20,10 @@ use crate::arrays::PiecewiseSequence;
 use crate::arrays::PrimitiveArray;
 use crate::arrays::bool::BoolArrayExt;
 use crate::arrays::dict::TakeExecute;
+use crate::arrays::piecewise_sequence::UnitMultiplierLengths;
 use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays;
 use crate::arrays::piecewise_sequence::validate_index_ranges;
+use crate::arrays::piecewise_sequence::validate_index_ranges_constant;
 use crate::builtins::ArrayBuiltins;
 use crate::dtype::UnsignedPType;
 use crate::executor::ExecutionCtx;
@@ -76,16 +78,32 @@ fn take_piecewise_sequence(
     let Some((starts, lengths)) = execute_unit_multiplier_index_arrays(indices, ctx)? else {
         return Ok(None);
     };
-    let buffer = match_each_unsigned_integer_ptype!(starts.ptype(), |S| {
-        match_each_unsigned_integer_ptype!(lengths.ptype(), |L| {
-            take_piecewise_bits(
-                &array.to_bit_buffer(),
-                starts.as_slice::<S>(),
-                lengths.as_slice::<L>(),
-                indices_ref.len(),
-            )?
-        })
-    });
+    let source = array.to_bit_buffer();
+    let output_len = indices_ref.len();
+    let buffer = match &lengths {
+        UnitMultiplierLengths::Constant(length) => {
+            match_each_unsigned_integer_ptype!(starts.ptype(), |S| {
+                take_piecewise_bits_constant_length(
+                    &source,
+                    starts.as_slice::<S>(),
+                    *length,
+                    output_len,
+                )?
+            })
+        }
+        UnitMultiplierLengths::Array(lengths) => {
+            match_each_unsigned_integer_ptype!(starts.ptype(), |S| {
+                match_each_unsigned_integer_ptype!(lengths.ptype(), |L| {
+                    take_piecewise_bits(
+                        &source,
+                        starts.as_slice::<S>(),
+                        lengths.as_slice::<L>(),
+                        output_len,
+                    )?
+                })
+            })
+        }
+    };
 
     Ok(Some(
         BoolArray::new(buffer, array.validity()?.take(indices_ref)?).into_array(),
@@ -119,6 +137,26 @@ fn take_bool_impl<I: AsPrimitive<usize>>(bools: BitBufferView<'_>, indices: &[I]
     })
 }
 
+fn take_piecewise_bits_constant_length<S>(
+    source: &BitBuffer,
+    starts: &[S],
+    length: usize,
+    output_len: usize,
+) -> VortexResult<BitBuffer>
+where
+    S: UnsignedPType,
+{
+    validate_index_ranges_constant(source.len(), starts, length, output_len)?;
+
+    let mut values = BitBufferMut::with_capacity(output_len);
+    for &start in starts {
+        let start = start.as_();
+        values.append_buffer(&source.slice(start..start + length));
+    }
+
+    Ok(values.freeze())
+}
+
 fn take_piecewise_bits<S, L>(
     source: &BitBuffer,
     starts: &[S],

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

@@ -14,8 +14,10 @@ use crate::arrays::DecimalArray;
 use crate::arrays::PiecewiseSequence;
 use crate::arrays::PrimitiveArray;
 use crate::arrays::dict::TakeExecute;
+use crate::arrays::piecewise_sequence::UnitMultiplierLengths;
 use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays;
 use crate::arrays::piecewise_sequence::validate_index_ranges;
+use crate::arrays::piecewise_sequence::validate_index_ranges_constant;
 use crate::dtype::IntegerPType;
 use crate::dtype::NativeDecimalType;
 use crate::dtype::UnsignedPType;
@@ -67,10 +69,57 @@ fn take_piecewise_sequence(
     };
     let validity = array.validity()?.take(indices_ref)?;
     let output_len = indices_ref.len();
-    let taken = take_piecewise_sequence_lengths(array, &starts, &lengths, validity, output_len)?;
+    let taken = match lengths {
+        UnitMultiplierLengths::Constant(length) => {
+            take_piecewise_sequence_constant_length(array, &starts, length, validity, output_len)?
+        }
+        UnitMultiplierLengths::Array(lengths) => {
+            take_piecewise_sequence_lengths(array, &starts, &lengths, validity, output_len)?
+        }
+    };
     Ok(Some(taken))
 }
 
+fn take_piecewise_sequence_constant_length(
+    array: ArrayView<'_, Decimal>,
+    starts: &PrimitiveArray,
+    length: usize,
+    validity: Validity,
+    output_len: usize,
+) -> VortexResult<ArrayRef> {
+    match_each_decimal_value_type!(array.values_type(), |D| {
+        take_piecewise_sequence_constant_length_typed::<D>(
+            array, starts, length, validity, output_len,
+        )
+    })
+}
+
+fn take_piecewise_sequence_constant_length_typed<D>(
+    array: ArrayView<'_, Decimal>,
+    starts: &PrimitiveArray,
+    length: usize,
+    validity: Validity,
+    output_len: usize,
+) -> VortexResult<ArrayRef>
+where
+    D: NativeDecimalType,
+{
+    match_each_unsigned_integer_ptype!(starts.ptype(), |S| {
+        let values = take_piecewise_constant_length_to_buffer::<S, D>(
+            starts.as_slice::<S>(),
+            length,
+            array.buffer::<D>().as_slice(),
+            output_len,
+        )?;
+
+        // SAFETY: contiguous gather preserves the decimal dtype and value representation.
+        Ok(
+            unsafe { DecimalArray::new_unchecked(values, array.decimal_dtype(), validity) }
+                .into_array(),
+        )
+    })
+}
+
 fn take_piecewise_sequence_lengths(
     array: ArrayView<'_, Decimal>,
     starts: &PrimitiveArray,
@@ -131,6 +180,27 @@ fn take_to_buffer<I: IntegerPType, T: NativeDecimalType>(indices: &[I], values:
     indices.iter().map(|idx| values[idx.as_()]).collect()
 }
 
+fn take_piecewise_constant_length_to_buffer<S, T>(
+    starts: &[S],
+    length: usize,
+    values: &[T],
+    output_len: usize,
+) -> VortexResult<Buffer<T>>
+where
+    S: UnsignedPType,
+    T: NativeDecimalType,
+{
+    validate_index_ranges_constant(values.len(), starts, length, output_len)?;
+
+    let mut result = BufferMut::<T>::with_capacity(output_len);
+    for &start in starts {
+        let start = start.as_();
+        result.extend_from_slice(&values[start..start + length]);
+    }
+
+    Ok(result.freeze())
+}
+
 fn take_piecewise_to_buffer<S, L, T>(
     starts: &[S],
     lengths: &[L],

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

@@ -84,7 +84,7 @@ fn take_empty_fsl(
         FixedSizeListArray::new_unchecked(new_elements, array.list_size(), new_validity, new_len)
     }
     .into_array()
-    .optimize()
+    .optimize_ctx(ctx.session())
 }
 
 fn take_non_empty_fsl(
@@ -147,7 +147,7 @@ fn take_non_empty_degenerate_fsl(
         )
     }
     .into_array()
-    .optimize()
+    .optimize_ctx(ctx.session())
 }
 
 fn take_non_empty_non_degenerate_fsl<I: IntegerPType>(
@@ -174,7 +174,7 @@ fn take_non_empty_non_degenerate_fsl<I: IntegerPType>(
         FixedSizeListArray::new_unchecked(new_elements, array.list_size(), new_validity, new_len)
     }
     .into_array()
-    .optimize()
+    .optimize_ctx(ctx.session())
 }
 
 fn take_non_empty_non_degenerate_elements<I: IntegerPType>(