Skip to main content

vortex_fsst/
ops.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_array::ArrayView;
5use vortex_array::ExecutionCtx;
6use vortex_array::IntoArray;
7use vortex_array::ProbeState;
8use vortex_array::RepeatedArrayProbe;
9use vortex_array::arrays::varbin::varbin_scalar;
10use vortex_array::scalar::Scalar;
11use vortex_array::vtable::OperationsVTable;
12use vortex_buffer::ByteBuffer;
13use vortex_error::VortexExpect;
14use vortex_error::VortexResult;
15
16use crate::FSST;
17use crate::FSSTArrayExt;
18
19/// The codes array is rebuilt from the codes buffer and the offsets slot on every read, so a
20/// repeated probe keeps one probe over it and the offsets child keeps its preparation.
21#[derive(Default)]
22pub struct FsstProbeState {
23    codes: Option<RepeatedArrayProbe>,
24}
25
26impl OperationsVTable<FSST> for FSST {
27    type ProbeState = FsstProbeState;
28
29    fn probe_scalar(
30        state: &mut ProbeState<'_, FSST>,
31        index: usize,
32        ctx: &mut ExecutionCtx,
33    ) -> VortexResult<Scalar> {
34        if !state.is_valid(index, ctx)? {
35            return Ok(Scalar::null(state.array().dtype().clone()));
36        }
37        let array = state.array();
38        let compressed = match state.retained() {
39            None => array.codes().into_array().execute_scalar(index, ctx)?,
40            Some(retained) => retained
41                .codes
42                .get_or_insert_with(|| RepeatedArrayProbe::new(array.codes().into_array()))
43                .execute_scalar(index, ctx)?,
44        };
45        let binary_datum = compressed.as_binary().value().vortex_expect("non-null");
46
47        let decoded_buffer = ByteBuffer::from(array.decompressor().decompress(binary_datum));
48        Ok(varbin_scalar(decoded_buffer, array.dtype()))
49    }
50
51    fn scalar_at(
52        array: ArrayView<'_, FSST>,
53        index: usize,
54        ctx: &mut ExecutionCtx,
55    ) -> VortexResult<Scalar> {
56        Self::probe_scalar(&mut ProbeState::once(array), index, ctx)
57    }
58}