vortex_array/arrays/piecewise_sequence/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Index encoding for concatenated sequential ranges.
5//!
6//! A `PiecewiseSequenceArray` represents the expanded index sequence
7//! `starts[i] + j * multipliers[i]` for `j` in `0..lengths[i]` for each piece `i`. It is
8//! intended for take operations that can gather regular runs without materializing one index per
9//! element.
10
11use itertools::Itertools;
12use num_traits::AsPrimitive;
13use vortex_buffer::BufferMut;
14use vortex_error::VortexExpect;
15use vortex_error::VortexResult;
16use vortex_error::vortex_bail;
17use vortex_error::vortex_ensure;
18use vortex_error::vortex_ensure_eq;
19use vortex_error::vortex_err;
20
21use crate::ArrayRef;
22use crate::Columnar;
23use crate::array::ArrayView;
24use crate::arrays::ConstantArray;
25use crate::arrays::PrimitiveArray;
26use crate::arrays::piecewise_sequence::array::PiecewiseSequenceArraySlotsExt;
27use crate::dtype::UnsignedPType;
28use crate::executor::ExecutionCtx;
29use crate::scalar::PValue;
30
31pub mod array;
32mod vtable;
33
34#[cfg(test)]
35mod tests;
36
37pub use array::PiecewiseSequenceArrayExt;
38pub use vtable::*;
39
40pub(crate) fn check_index_arrays(
41    starts: &ArrayRef,
42    lengths: &ArrayRef,
43    multipliers: &ArrayRef,
44) -> VortexResult<()> {
45    check_index_array("starts", starts)?;
46    check_index_array("lengths", lengths)?;
47    check_index_array("multipliers", multipliers)?;
48    vortex_ensure_eq!(
49        starts.len(),
50        lengths.len(),
51        "PiecewiseSequenceArray starts length does not match lengths length",
52    );
53    vortex_ensure_eq!(
54        starts.len(),
55        multipliers.len(),
56        "PiecewiseSequenceArray starts length does not match multipliers length",
57    );
58    Ok(())
59}
60
61pub(crate) fn execute_index_arrays(
62    array: ArrayView<'_, PiecewiseSequence>,
63    ctx: &mut ExecutionCtx,
64) -> VortexResult<(PrimitiveArray, PrimitiveArray, PrimitiveArray)> {
65    let starts = array.starts().clone().execute::<PrimitiveArray>(ctx)?;
66    let lengths = array.lengths().clone().execute::<PrimitiveArray>(ctx)?;
67    let multipliers = array.multipliers().clone().execute::<PrimitiveArray>(ctx)?;
68    Ok((starts, lengths, multipliers))
69}
70
71pub(crate) fn maybe_contiguous_slices(
72    array: ArrayView<'_, PiecewiseSequence>,
73    ctx: &mut ExecutionCtx,
74) -> VortexResult<Option<(PrimitiveArray, Columnar)>> {
75    if !is_constant_one(array.multipliers()) {
76        return Ok(None);
77    }
78
79    let starts = array.starts().clone().execute::<PrimitiveArray>(ctx)?;
80    let lengths = array.lengths().clone().execute::<Columnar>(ctx)?;
81    Ok(Some((starts, lengths)))
82}
83
84pub(crate) fn is_constant_one(multipliers: &ArrayRef) -> bool {
85    let Some(scalar) = multipliers.as_constant() else {
86        return false;
87    };
88    matches!(
89        scalar.as_primitive_opt().and_then(|scalar| scalar.pvalue()),
90        Some(PValue::U8(1) | PValue::U16(1) | PValue::U32(1) | PValue::U64(1))
91    )
92}
93
94pub(crate) fn constant_unsigned_usize(array: &ConstantArray) -> usize {
95    let pvalue = array
96        .scalar()
97        .as_primitive_opt()
98        .and_then(|scalar| scalar.pvalue())
99        .vortex_expect("validated PiecewiseSequence length constants are primitive");
100
101    match pvalue {
102        PValue::U8(value) => value as usize,
103        PValue::U16(value) => value as usize,
104        PValue::U32(value) => value as usize,
105        PValue::U64(value) => value.as_(),
106        _ => unreachable!("validated PiecewiseSequence length constants are unsigned"),
107    }
108}
109
110fn check_index_array(name: &str, array: &ArrayRef) -> VortexResult<()> {
111    vortex_ensure!(
112        array.dtype().is_unsigned_int(),
113        "PiecewiseSequenceArray {name} must have unsigned integer dtype, got {}",
114        array.dtype()
115    );
116    vortex_ensure!(
117        !array.dtype().is_nullable(),
118        "PiecewiseSequenceArray {name} must be non-nullable, got {}",
119        array.dtype()
120    );
121    Ok(())
122}
123
124pub(crate) fn materialize_ranges<S, L, M>(
125    starts: &PrimitiveArray,
126    lengths: &PrimitiveArray,
127    multipliers: &PrimitiveArray,
128    output_len: usize,
129) -> VortexResult<BufferMut<u64>>
130where
131    S: UnsignedPType,
132    L: UnsignedPType,
133    M: UnsignedPType,
134{
135    let starts = starts.as_slice::<S>();
136    let lengths = lengths.as_slice::<L>();
137    let multipliers = multipliers.as_slice::<M>();
138    let mut values = BufferMut::with_capacity(output_len);
139    let mut computed_len = 0usize;
140
141    for ((&start, &length), &multiplier) in starts.iter().zip_eq(lengths).zip_eq(multipliers) {
142        let start: usize = start.as_();
143        let length: usize = length.as_();
144        let multiplier: usize = multiplier.as_();
145        if length != 0 {
146            let last_offset = length - 1;
147            let last_delta = last_offset
148                .checked_mul(multiplier)
149                .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?;
150            start
151                .checked_add(last_delta)
152                .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?;
153        }
154        computed_len = computed_len
155            .checked_add(length)
156            .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?;
157
158        values.extend((0..length).map(|offset| (start + offset * multiplier) as u64));
159    }
160
161    if computed_len != output_len {
162        vortex_bail!(
163            "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}"
164        );
165    }
166    Ok(values)
167}
vortex_array/arrays/piecewise_sequence/mod.rs

vortex_array/arrays/piecewise_sequence/
mod.rs
