vortex_fsst/compute/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4mod byte_length;
5mod cast;
6mod compare;
7mod filter;
8mod like;
9
10use num_traits::ToPrimitive;
11use vortex_array::ArrayRef;
12use vortex_array::ArrayView;
13use vortex_array::ExecutionCtx;
14use vortex_array::IntoArray;
15use vortex_array::arrays::Primitive;
16use vortex_array::arrays::PrimitiveArray;
17use vortex_array::arrays::dict::TakeExecute;
18use vortex_array::arrays::primitive::PrimitiveArrayExt;
19use vortex_array::arrays::varbin::take_varbin;
20use vortex_array::builtins::ArrayBuiltins;
21use vortex_array::dtype::IntegerPType;
22use vortex_array::expr::stats::Stat;
23use vortex_array::expr::stats::StatsProviderExt;
24use vortex_array::match_each_integer_ptype;
25use vortex_array::match_each_unsigned_integer_ptype;
26use vortex_array::scalar::Scalar;
27use vortex_error::VortexResult;
28use vortex_mask::AllOr;
29
30use crate::FSST;
31use crate::FSSTArrayExt;
32use crate::FSSTArraySlotsExt;
33
34/// Small arrays gather only when their compressed payload is large enough to amortize setup.
35const MIN_GATHER_LEN: usize = 1024;
36/// Minimum compressed payload for gathering from a small array, measured by `fsst_take`.
37const MIN_GATHER_COMPRESSED_BYTES: u64 = 8 * 1024;
38
39/// Only takes with fewer indices than `1 / SPARSE_TAKE_DENOMINATOR` of the rows gather compressed
40/// rows. A denser take decodes every row once and gathers the decoded strings instead: gathering
41/// decodes a row again for every index that repeats it, and costs more per row than decoding
42/// in bulk.
43const SPARSE_TAKE_DENOMINATOR: usize = 4;
44
45impl TakeExecute for FSST {
46    fn take(
47        array: ArrayView<'_, Self>,
48        indices: &ArrayRef,
49        ctx: &mut ExecutionCtx,
50    ) -> VortexResult<Option<ArrayRef>> {
51        if indices.len().saturating_mul(SPARSE_TAKE_DENOMINATOR) >= array.len() {
52            return Ok(None);
53        }
54        if array.len() < MIN_GATHER_LEN {
55            let Some(bytes) = compressed_byte_span(array.codes_offsets())
56                .filter(|&bytes| bytes >= MIN_GATHER_COMPRESSED_BYTES)
57            else {
58                return Ok(None);
59            };
60            // Leave at least half the codes untouched to amortize gathering and the cost check.
61            if !take_fits_byte_budget(array.codes_offsets(), indices, bytes / 2, ctx)? {
62                return Ok(None);
63            }
64        }
65
66        Ok(Some(
67            FSST::try_new_with_symbol_table(
68                array
69                    .dtype()
70                    .clone()
71                    .union_nullability(indices.dtype().nullability()),
72                array.symbol_table(),
73                take_varbin(array.codes().as_view(), indices, ctx)?,
74                array
75                    .uncompressed_lengths()
76                    .take(indices.clone())?
77                    .fill_null(Scalar::zero_value(
78                        &array.uncompressed_lengths_dtype().clone(),
79                    ))?,
80                ctx,
81            )?
82            .into_array(),
83        ))
84    }
85}
86
87fn take_fits_byte_budget(
88    offsets: &ArrayRef,
89    indices: &ArrayRef,
90    budget: u64,
91    ctx: &mut ExecutionCtx,
92) -> VortexResult<bool> {
93    let offsets = offsets.clone().execute::<PrimitiveArray>(ctx)?;
94    let indices = indices.clone().execute::<PrimitiveArray>(ctx)?;
95    let validity = indices.validity()?.execute_mask(indices.len(), ctx)?;
96    // Dispatch only on widths: offsets and valid indices are nonnegative.
97    let offsets = offsets.reinterpret_cast(offsets.ptype().to_unsigned());
98    let indices = indices.reinterpret_cast(indices.ptype().to_unsigned());
99    Ok(match_each_unsigned_integer_ptype!(offsets.ptype(), |O| {
100        let offsets = offsets.as_slice::<O>();
101        match_each_unsigned_integer_ptype!(indices.ptype(), |I| {
102            let indices = indices.as_slice::<I>();
103            match validity.indices() {
104                AllOr::All => selected_bytes_fit(offsets, indices.iter().copied(), budget),
105                AllOr::None => true,
106                AllOr::Some(valid) => {
107                    selected_bytes_fit(offsets, valid.iter().map(|&i| indices[i]), budget)
108                }
109            }
110        })
111    }))
112}
113
114// Count repeats before copying codes: gathering decodes each selected occurrence separately.
115fn selected_bytes_fit<O: IntegerPType, I: IntegerPType>(
116    offsets: &[O],
117    mut indices: impl Iterator<Item = I>,
118    mut budget: u64,
119) -> bool {
120    indices.all(|index| {
121        let Some(index) = index.to_usize().filter(|&i| i < offsets.len() - 1) else {
122            return false;
123        };
124        let remaining = offsets[index + 1]
125            .to_u64()
126            .zip(offsets[index].to_u64())
127            .and_then(|(end, start)| end.checked_sub(start))
128            .and_then(|len| budget.checked_sub(len));
129        let Some(remaining) = remaining else {
130            return false;
131        };
132        budget = remaining;
133        true
134    })
135}
136
137// Slices retain their backing bytes, so only the logical offsets describe their decode cost.
138// Inspect already-materialized offsets or cached exact bounds without executing either child.
139fn compressed_byte_span(offsets: &ArrayRef) -> Option<u64> {
140    if let Some(offsets) = offsets.as_opt::<Primitive>()
141        && offsets.buffer_handle().as_host_opt().is_some()
142    {
143        return match_each_integer_ptype!(offsets.ptype(), |P| {
144            let values = offsets.as_slice::<P>();
145            let first = values.first()?.to_u64()?;
146            let last = values.last()?.to_u64()?;
147            last.checked_sub(first)
148        });
149    }
150    let stats = offsets.statistics();
151    let first = stats.get_as::<u64>(Stat::Min).as_exact()?;
152    let last = stats.get_as::<u64>(Stat::Max).as_exact()?;
153    last.checked_sub(first)
154}
155
156#[cfg(test)]
157mod tests {
158    use fsst::Compressor;
159    use fsst::Symbol;
160    use rstest::rstest;
161    use vortex_array::ArrayRef;
162    use vortex_array::ExecutionCtx;
163    use vortex_array::IntoArray;
164    use vortex_array::VortexSessionExecute;
165    use vortex_array::array_session;
166    use vortex_array::arrays::DictArray;
167    use vortex_array::arrays::PrimitiveArray;
168    use vortex_array::arrays::VarBinArray;
169    use vortex_array::arrays::dict::TakeExecute;
170    use vortex_array::assert_arrays_eq;
171    use vortex_array::compute::conformance::consistency::test_array_consistency;
172    use vortex_array::compute::conformance::take::test_take_conformance;
173    use vortex_array::dtype::DType;
174    use vortex_array::dtype::Nullability;
175    use vortex_array::expr::stats::Precision;
176    use vortex_array::expr::stats::Stat;
177    use vortex_array::validity::Validity;
178    use vortex_buffer::buffer;
179    use vortex_error::VortexResult;
180
181    use super::compressed_byte_span;
182    use crate::FSST;
183    use crate::FSSTArray;
184    use crate::fsst_compress;
185    use crate::fsst_train_compressor;
186
187    #[test]
188    fn test_take_null() -> VortexResult<()> {
189        let session = array_session();
190        crate::initialize(&session);
191        let mut ctx = session.create_execution_ctx();
192        let arr =
193            VarBinArray::from_iter([Some("h")], DType::Utf8(Nullability::NonNullable)).into_array();
194        let compr = fsst_train_compressor(&arr, &mut ctx)?;
195        let fsst = fsst_compress(&arr, &compr, &mut ctx)?;
196
197        let idx1: PrimitiveArray = (0..1).collect();
198
199        assert_arrays_eq!(fsst.take(idx1.into_array())?, arr, &mut ctx);
200
201        let idx2: PrimitiveArray = PrimitiveArray::from_option_iter(vec![Some(0)]);
202
203        assert_arrays_eq!(
204            fsst.take(idx2.into_array())?,
205            VarBinArray::from_iter([Some("h")], DType::Utf8(Nullability::Nullable)),
206            &mut ctx
207        );
208        Ok(())
209    }
210
211    #[rstest]
212    #[case(VarBinArray::from_iter(
213        ["hello world", "testing fsst", "compression test", "data array", "vortex encoding"].map(Some),
214        DType::Utf8(Nullability::NonNullable),
215    ))]
216    #[case(VarBinArray::from_iter(
217        [Some("hello"), None, Some("world"), Some("test"), None],
218        DType::Utf8(Nullability::Nullable),
219    ))]
220    #[case(VarBinArray::from_iter(
221        ["single element"].map(Some),
222        DType::Utf8(Nullability::NonNullable),
223    ))]
224    fn test_take_fsst_conformance(#[case] varbin: VarBinArray) -> VortexResult<()> {
225        let session = array_session();
226        crate::initialize(&session);
227        let mut ctx = session.create_execution_ctx();
228        let varbin = varbin.into_array();
229        let compressor = fsst_train_compressor(&varbin, &mut ctx)?;
230        let array = fsst_compress(&varbin, &compressor, &mut ctx)?;
231        test_take_conformance(&array.into_array(), &mut ctx);
232        Ok(())
233    }
234
235    #[rstest]
236    #[case(2048, 1, false, true)]
237    #[case(2048, 3, true, true)]
238    #[case(2048, 4, false, false)]
239    #[case(2, 1, false, false)]
240    fn test_take_small_byte_gate(
241        #[case] repeats: usize,
242        #[case] selected: usize,
243        #[case] nullable: bool,
244        #[case] gathers: bool,
245    ) -> VortexResult<()> {
246        let mut ctx = array_session().create_execution_ctx();
247        let input = VarBinArray::from_iter(
248            (0..16).map(|row| {
249                (!(nullable && row == 0)).then(|| format!("{}{row}", "caption ".repeat(repeats)))
250            }),
251            DType::Utf8(Nullability::Nullable),
252        )
253        .into_array();
254        let compressor = Compressor::rebuild_from([Symbol::from_slice(b"caption ")], [8u8]);
255        let array = fsst_compress(&input, &compressor, &mut ctx)?;
256        let indices = if nullable {
257            PrimitiveArray::from_option_iter([Some(15u64), None, Some(0)]).into_array()
258        } else {
259            PrimitiveArray::from_iter((0..selected).map(|index| (15 - index) as u64)).into_array()
260        };
261        assert_eq!(
262            FSST::take(array.as_view(), &indices, &mut ctx)?.is_some(),
263            gathers
264        );
265        assert_arrays_eq!(
266            array.into_array().take(indices.clone())?,
267            input.take(indices)?,
268            &mut ctx
269        );
270        Ok(())
271    }
272
273    #[test]
274    fn test_take_small_slice_uses_logical_byte_span() -> VortexResult<()> {
275        let mut ctx = array_session().create_execution_ctx();
276        let input = VarBinArray::from_iter(
277            (0..16).map(|row| Some("caption ".repeat(if row < 4 { 4096 } else { 2 }))),
278            DType::Utf8(Nullability::NonNullable),
279        )
280        .into_array();
281        let compressor = Compressor::rebuild_from([Symbol::from_slice(b"caption ")], [8u8]);
282        let array = fsst_compress(&input, &compressor, &mut ctx)?.into_array();
283        let indices = buffer![0u64].into_array();
284        for (range, gathers) in [(0..8, true), (4..16, false)] {
285            let sliced = array.slice(range.clone())?;
286            assert_eq!(
287                FSST::take(sliced.as_::<FSST>(), &indices, &mut ctx)?.is_some(),
288                gathers
289            );
290            assert_arrays_eq!(
291                sliced.take(indices.clone())?,
292                input.slice(range)?.take(indices.clone())?,
293                &mut ctx
294            );
295        }
296        Ok(())
297    }
298
299    #[rstest]
300    #[case::repeated_long_row(1, PrimitiveArray::from_iter([0u64; 63]).into_array(), false)]
301    #[case::distinct_long_rows(16, PrimitiveArray::from_iter(0i32..16).into_array(), false)]
302    #[case::short_row_with_null_index(1, PrimitiveArray::new(
303        buffer![u64::MAX, 1], Validity::from_iter([false, true]),
304    ).into_array(), true)]
305    #[case::encoded_indices(1, DictArray::new(
306        buffer![0u8, 0].into_array(), buffer![1u32].into_array(),
307    ).into_array(), true)]
308    fn test_take_small_uneven_lengths(
309        #[case] long_rows: usize,
310        #[case] indices: ArrayRef,
311        #[case] gathers: bool,
312    ) -> VortexResult<()> {
313        let mut ctx = array_session().create_execution_ctx();
314        let input = VarBinArray::from_iter(
315            (0..256).map(|row| Some("caption ".repeat(if row < long_rows { 8192 } else { 1 }))),
316            DType::Utf8(Nullability::NonNullable),
317        )
318        .into_array();
319        let compressor = Compressor::rebuild_from([Symbol::from_slice(b"caption ")], [8u8]);
320        let array = fsst_compress(&input, &compressor, &mut ctx)?;
321        assert_eq!(
322            FSST::take(array.as_view(), &indices, &mut ctx)?.is_some(),
323            gathers
324        );
325        assert_arrays_eq!(
326            array.into_array().take(indices.clone())?,
327            input.take(indices)?,
328            &mut ctx
329        );
330        Ok(())
331    }
332
333    #[test]
334    fn test_take_compressed_span_requires_exact_cached_bounds() {
335        let offsets = DictArray::new(
336            buffer![0u32, 1].into_array(),
337            buffer![1000u64, 20000].into_array(),
338        )
339        .into_array();
340        assert_eq!(compressed_byte_span(&offsets), None);
341        offsets
342            .statistics()
343            .set(Stat::Min, Precision::exact(1000u64));
344        offsets
345            .statistics()
346            .set(Stat::Max, Precision::exact(20000u64));
347        assert_eq!(compressed_byte_span(&offsets), Some(19000));
348        offsets
349            .statistics()
350            .set(Stat::Min, Precision::Inexact(1000u64.into()));
351        assert_eq!(compressed_byte_span(&offsets), None);
352    }
353
354    type FsstBuilder = fn(&mut ExecutionCtx) -> FSSTArray;
355
356    #[rstest]
357    // Basic string arrays
358    #[case::fsst_simple(|ctx: &mut ExecutionCtx| {
359        let array = VarBinArray::from_iter(
360            ["hello world", "testing fsst", "compression test", "data array", "vortex encoding"].map(Some),
361            DType::Utf8(Nullability::NonNullable),
362        ).into_array();
363        let compressor = fsst_train_compressor(&array, ctx).unwrap();
364        fsst_compress(&array, &compressor, ctx).unwrap()
365    })]
366    // Nullable strings
367    #[case::fsst_nullable(|ctx: &mut ExecutionCtx| {
368        let array = VarBinArray::from_iter(
369            [Some("hello"), None, Some("world"), Some("test"), None],
370            DType::Utf8(Nullability::Nullable),
371        ).into_array();
372        let compressor = fsst_train_compressor(&array, ctx).unwrap();
373        fsst_compress(&array, &compressor, ctx).unwrap()
374    })]
375    // Repetitive patterns (good for FSST compression)
376    #[case::fsst_repetitive(|ctx: &mut ExecutionCtx| {
377        let array = VarBinArray::from_iter(
378            ["http://example.com", "http://test.com", "http://vortex.dev", "http://data.org"].map(Some),
379            DType::Utf8(Nullability::NonNullable),
380        ).into_array();
381        let compressor = fsst_train_compressor(&array, ctx).unwrap();
382        fsst_compress(&array, &compressor, ctx).unwrap()
383    })]
384    // Edge cases
385    #[case::fsst_single(|ctx: &mut ExecutionCtx| {
386        let array = VarBinArray::from_iter(
387            ["single element"].map(Some),
388            DType::Utf8(Nullability::NonNullable),
389        ).into_array();
390        let compressor = fsst_train_compressor(&array, ctx).unwrap();
391        fsst_compress(&array, &compressor, ctx).unwrap()
392    })]
393    #[case::fsst_empty_strings(|ctx: &mut ExecutionCtx| {
394        let array = VarBinArray::from_iter(
395            ["", "test", "", "hello", ""].map(Some),
396            DType::Utf8(Nullability::NonNullable),
397        ).into_array();
398        let compressor = fsst_train_compressor(&array, ctx).unwrap();
399        fsst_compress(&array, &compressor, ctx).unwrap()
400    })]
401    // Large arrays
402    #[case::fsst_large(|ctx: &mut ExecutionCtx| {
403        let data: Vec<Option<&str>> = (0..1500)
404            .map(|i| Some(match i % 10 {
405                0 => "https://www.example.com/page",
406                1 => "https://www.test.org/data",
407                2 => "https://www.vortex.dev/docs",
408                3 => "https://www.github.com/apache/arrow",
409                4 => "https://www.rust-lang.org/learn",
410                5 => "SELECT * FROM table WHERE id = ",
411                6 => "INSERT INTO users (name, email) VALUES",
412                7 => "UPDATE records SET status = 'active'",
413                8 => "DELETE FROM logs WHERE timestamp < ",
414                _ => "CREATE TABLE data (id INT, value TEXT)",
415            }))
416            .collect();
417        let array = VarBinArray::from_iter(data, DType::Utf8(Nullability::NonNullable)).into_array();
418        let compressor = fsst_train_compressor(&array, ctx).unwrap();
419        fsst_compress(&array, &compressor, ctx).unwrap()
420    })]
421
422    fn test_fsst_consistency(#[case] build: FsstBuilder) {
423        let session = array_session();
424        crate::initialize(&session);
425        let mut ctx = session.create_execution_ctx();
426        let array = build(&mut ctx);
427        test_array_consistency(&array.into_array(), &mut ctx);
428    }
429}
vortex_fsst/compute/mod.rs

vortex_fsst/compute/
mod.rs
