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

vortex_array/arrays/piecewise_sequence/
mod.rs
