vortex_zstd/
zstd_buffers.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Debug;
5use std::fmt::Display;
6use std::fmt::Formatter;
7use std::hash::Hash;
8use std::hash::Hasher;
9use std::sync::Arc;
10
11use prost::Message as _;
12use vortex_array::Array;
13use vortex_array::ArrayEq;
14use vortex_array::ArrayHash;
15use vortex_array::ArrayId;
16use vortex_array::ArrayParts;
17use vortex_array::ArrayRef;
18use vortex_array::ArraySlots;
19use vortex_array::ArrayView;
20use vortex_array::EqMode;
21use vortex_array::ExecutionCtx;
22use vortex_array::ExecutionResult;
23use vortex_array::buffer::BufferHandle;
24use vortex_array::dtype::DType;
25use vortex_array::scalar::Scalar;
26use vortex_array::serde::ArrayChildren;
27use vortex_array::session::ArraySessionExt;
28use vortex_array::vtable::OperationsVTable;
29use vortex_array::vtable::VTable;
30use vortex_array::vtable::ValidityVTable;
31use vortex_buffer::Alignment;
32use vortex_buffer::ByteBuffer;
33use vortex_buffer::ByteBufferMut;
34use vortex_error::VortexResult;
35use vortex_error::vortex_ensure_eq;
36use vortex_error::vortex_err;
37use vortex_session::VortexSession;
38use vortex_session::registry::CachedId;
39
40use crate::ZstdBuffersMetadata;
41
42/// A [`ZstdBuffers`]-encoded Vortex array.
43pub type ZstdBuffersArray = Array<ZstdBuffers>;
44
45#[derive(Clone, Debug)]
46/// Encoding marker for buffer-level zstd compression.
47pub struct ZstdBuffers;
48
49impl ZstdBuffers {
50    /// Construct a [`ZstdBuffersArray`] from compressed buffer data.
51    pub fn try_new(
52        dtype: DType,
53        len: usize,
54        data: ZstdBuffersData,
55    ) -> VortexResult<ZstdBuffersArray> {
56        Array::try_from_parts(ArrayParts::new(ZstdBuffers, dtype, len, data))
57    }
58
59    /// Compress every top-level buffer of `array` independently with zstd.
60    ///
61    /// Children are preserved as slots and the wrapped array's serialized metadata is stored so the
62    /// original array can be rebuilt after decompression.
63    pub fn compress(
64        array: &ArrayRef,
65        level: i32,
66        session: &VortexSession,
67    ) -> VortexResult<ZstdBuffersArray> {
68        let encoding_id = array.encoding_id();
69        let metadata = session
70            .array_serialize(array)?
71            .ok_or_else(|| vortex_err!("[ZstdBuffers]: Array does not support serialization"))?;
72        let buffer_handles = array.buffer_handles();
73        let children = array.children();
74
75        let mut compressed_buffers = Vec::with_capacity(buffer_handles.len());
76        let mut uncompressed_sizes = Vec::with_capacity(buffer_handles.len());
77        let mut buffer_alignments = Vec::with_capacity(buffer_handles.len());
78
79        let mut compressor = zstd::bulk::Compressor::new(level)?;
80        // Compression is currently CPU-only, so we gather all buffers on the host.
81        for handle in &buffer_handles {
82            buffer_alignments.push(u32::from(handle.alignment()));
83            let host_buf = handle.clone().try_to_host_sync()?;
84            uncompressed_sizes.push(host_buf.len() as u64);
85            let compressed = compressor.compress(&host_buf)?;
86            compressed_buffers.push(BufferHandle::new_host(ByteBuffer::from(compressed)));
87        }
88
89        let data = ZstdBuffersData {
90            inner_encoding_id: encoding_id,
91            inner_metadata: metadata,
92            compressed_buffers,
93            uncompressed_sizes,
94            buffer_alignments,
95        };
96        let slots: ArraySlots = children.into_iter().map(Some).collect();
97        let compressed = Array::try_from_parts(
98            ArrayParts::new(ZstdBuffers, array.dtype().clone(), array.len(), data)
99                .with_slots(slots),
100        )?;
101        compressed.statistics().inherit_from(array.statistics());
102        Ok(compressed)
103    }
104
105    /// Rebuild the wrapped array from decompressed buffer handles.
106    pub fn build_inner(
107        array: &ZstdBuffersArray,
108        buffer_handles: &[BufferHandle],
109        session: &VortexSession,
110    ) -> VortexResult<ArrayRef> {
111        let registry = session.arrays().registry().clone();
112        let inner_vtable = registry
113            .find(&array.data().inner_encoding_id)
114            .ok_or_else(|| {
115                vortex_err!("Unknown inner encoding: {}", array.data().inner_encoding_id)
116            })?;
117
118        let children: Vec<ArrayRef> = array.slots().iter().flatten().cloned().collect();
119        inner_vtable.deserialize(
120            array.dtype(),
121            array.len(),
122            &array.data().inner_metadata,
123            buffer_handles,
124            &children.as_slice(),
125            session,
126        )
127    }
128
129    fn decompress_and_build_inner(
130        array: &ZstdBuffersArray,
131        session: &VortexSession,
132    ) -> VortexResult<ArrayRef> {
133        let decompressed_buffers = array.data().decompress_buffers()?;
134        Self::build_inner(array, &decompressed_buffers, session)
135    }
136}
137
138/// An encoding that ZSTD-compresses the buffers of any wrapped array.
139///
140/// Unlike [`ZstdArray`](crate::ZstdArray), which interleaves string lengths with content bytes,
141/// `ZstdBuffersArray` compresses each buffer independently. This enables zero-conversion
142/// GPU decompression since the original buffer layout is preserved after decompression.
143#[derive(Clone, Debug)]
144pub struct ZstdBuffersData {
145    inner_encoding_id: ArrayId,
146    inner_metadata: Vec<u8>,
147    compressed_buffers: Vec<BufferHandle>,
148    uncompressed_sizes: Vec<u64>,
149    buffer_alignments: Vec<u32>,
150}
151
152impl Display for ZstdBuffersData {
153    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
154        write!(f, "inner_encoding: {}", self.inner_encoding_id)
155    }
156}
157
158#[derive(Clone, Debug)]
159/// Decode plan for buffer-level zstd decompression.
160pub struct ZstdBuffersDecodePlan {
161    compressed_buffers: Vec<BufferHandle>,
162    frame_sizes: Arc<[usize]>,
163    output_sizes: Arc<[usize]>,
164    output_offsets: Vec<usize>,
165    output_alignments: Vec<Alignment>,
166    output_size_total: usize,
167    output_size_max: usize,
168}
169
170impl ZstdBuffersDecodePlan {
171    /// Compressed buffers to decode.
172    pub fn compressed_buffers(&self) -> &[BufferHandle] {
173        &self.compressed_buffers
174    }
175
176    /// Compressed frame sizes in bytes.
177    pub fn frame_sizes(&self) -> Arc<[usize]> {
178        Arc::clone(&self.frame_sizes)
179    }
180
181    /// Decompressed output size for each buffer.
182    pub fn output_sizes(&self) -> Arc<[usize]> {
183        Arc::clone(&self.output_sizes)
184    }
185
186    /// Byte offsets of each decompressed buffer in one contiguous output allocation.
187    pub fn output_offsets(&self) -> &[usize] {
188        &self.output_offsets
189    }
190
191    /// Total byte size of the planned contiguous output allocation.
192    pub fn output_size_total(&self) -> usize {
193        self.output_size_total
194    }
195
196    /// Largest single decompressed buffer size.
197    pub fn output_size_max(&self) -> usize {
198        self.output_size_max
199    }
200
201    /// Number of compressed frames/buffers in the plan.
202    pub fn num_frames(&self) -> usize {
203        self.compressed_buffers.len()
204    }
205
206    /// Split a contiguous decompressed output buffer into per-buffer handles using planned
207    /// offsets/sizes and enforce each buffer's required alignment.
208    pub fn split_output_handle(
209        &self,
210        output_handle: &BufferHandle,
211    ) -> VortexResult<Vec<BufferHandle>> {
212        self.output_offsets
213            .iter()
214            .zip(self.output_sizes.iter())
215            .zip(self.output_alignments.iter())
216            .map(|((&offset, &size), &alignment)| {
217                output_handle
218                    .slice(offset..offset + size)
219                    .ensure_aligned(alignment)
220            })
221            .collect::<VortexResult<Vec<_>>>()
222    }
223}
224
225impl ZstdBuffersData {
226    fn validate(&self) -> VortexResult<()> {
227        vortex_ensure_eq!(
228            self.compressed_buffers.len(),
229            self.uncompressed_sizes.len(),
230            "zstd_buffers metadata mismatch: {} compressed buffers vs {} sizes",
231            self.compressed_buffers.len(),
232            self.uncompressed_sizes.len()
233        );
234        vortex_ensure_eq!(
235            self.compressed_buffers.len(),
236            self.buffer_alignments.len(),
237            "zstd_buffers metadata mismatch: {} compressed buffers vs {} alignments",
238            self.compressed_buffers.len(),
239            self.buffer_alignments.len()
240        );
241        Ok(())
242    }
243
244    fn decompress_buffers(&self) -> VortexResult<Vec<BufferHandle>> {
245        // CPU decode path: zstd::bulk works on host bytes, so compressed buffers are
246        // materialized on the host via `try_to_host_sync`.
247        let mut decompressor = zstd::bulk::Decompressor::new()?;
248        let mut result = Vec::with_capacity(self.compressed_buffers.len());
249        for (i, (buf, &uncompressed_size)) in self
250            .compressed_buffers
251            .iter()
252            .zip(&self.uncompressed_sizes)
253            .enumerate()
254        {
255            let size = usize::try_from(uncompressed_size)?;
256            let alignment = self.buffer_alignments.get(i).copied().unwrap_or(1);
257
258            let aligned = Alignment::try_from(alignment)?;
259            let mut output = ByteBufferMut::with_capacity_aligned(size, aligned);
260            let spare = output.spare_capacity_mut();
261
262            // This is currently guaranteed, but still good to check because
263            // of the unsafe calls below.
264            if spare.len() < size {
265                return Err(vortex_err!(
266                    "Insufficient output capacity: expected at least {}, got {}",
267                    size,
268                    spare.len()
269                ));
270            }
271            // SAFETY: we only expose the first `size` bytes and mark them initialized via
272            // `set_len(size)` after zstd reports how many bytes were written.
273            let dst =
274                unsafe { std::slice::from_raw_parts_mut(spare.as_mut_ptr().cast::<u8>(), size) };
275            let compressed = buf.clone().try_to_host_sync()?;
276            let written = decompressor.decompress_to_buffer(compressed.as_slice(), dst)?;
277            if written != size {
278                return Err(vortex_err!(
279                    "Decompressed size mismatch: expected {}, got {}",
280                    size,
281                    written
282                ));
283            }
284            // SAFETY: zstd wrote exactly `size` initialized bytes into `dst`.
285            unsafe { output.set_len(size) };
286            result.push(BufferHandle::new_host(output.freeze()));
287        }
288        Ok(result)
289    }
290
291    /// Build a decode plan for external or device decompression.
292    pub fn decode_plan(&self) -> VortexResult<ZstdBuffersDecodePlan> {
293        // If invariants are somehow broken, device decompression could have UB, so ensure
294        // they still hold.
295        self.validate()?;
296
297        let output_sizes = self
298            .uncompressed_sizes
299            .iter()
300            .map(|&size| usize::try_from(size))
301            .collect::<Result<Vec<_>, _>>()?;
302        let output_size_max = output_sizes.iter().copied().max().unwrap_or(0);
303
304        let output_alignments = self
305            .buffer_alignments
306            .iter()
307            .map(|&alignment| Alignment::try_from(alignment))
308            .collect::<VortexResult<Vec<_>>>()?;
309
310        let (output_offsets, output_size_total) =
311            compute_output_layout(&output_sizes, &output_alignments);
312
313        let compressed_buffers = self.compressed_buffers.clone();
314        let frame_sizes: Arc<[usize]> = compressed_buffers
315            .iter()
316            .map(BufferHandle::len)
317            .collect::<Vec<_>>()
318            .into();
319        let output_sizes: Arc<[usize]> = output_sizes.into();
320
321        Ok(ZstdBuffersDecodePlan {
322            compressed_buffers,
323            frame_sizes,
324            output_sizes,
325            output_offsets,
326            output_alignments,
327            output_size_total,
328            output_size_max,
329        })
330    }
331}
332
333fn compute_output_layout(
334    output_sizes: &[usize],
335    output_alignments: &[Alignment],
336) -> (Vec<usize>, usize) {
337    // Compute aligned offsets for each decompressed buffer in one contiguous output allocation.
338    // Each buffer starts at the next multiple of its required alignment.
339    let mut offsets = Vec::with_capacity(output_sizes.len());
340    let mut total_size = 0usize;
341
342    for (&size, &alignment) in output_sizes.iter().zip(output_alignments.iter()) {
343        total_size = total_size.next_multiple_of(*alignment);
344        offsets.push(total_size);
345        total_size += size;
346    }
347
348    (offsets, total_size)
349}
350
351#[expect(clippy::disallowed_methods, reason = "interning a dynamic id")]
352fn array_id_from_string(s: &str) -> ArrayId {
353    ArrayId::new(s)
354}
355
356impl ArrayHash for ZstdBuffersData {
357    fn array_hash<H: Hasher>(&self, state: &mut H, accuracy: EqMode) {
358        self.inner_encoding_id.hash(state);
359        self.inner_metadata.hash(state);
360        for buf in &self.compressed_buffers {
361            buf.array_hash(state, accuracy);
362        }
363        self.uncompressed_sizes.hash(state);
364        self.buffer_alignments.hash(state);
365    }
366}
367
368impl ArrayEq for ZstdBuffersData {
369    fn array_eq(&self, other: &Self, accuracy: EqMode) -> bool {
370        self.inner_encoding_id == other.inner_encoding_id
371            && self.inner_metadata == other.inner_metadata
372            && self.compressed_buffers.len() == other.compressed_buffers.len()
373            && self
374                .compressed_buffers
375                .iter()
376                .zip(&other.compressed_buffers)
377                .all(|(a, b)| a.array_eq(b, accuracy))
378            && self.uncompressed_sizes == other.uncompressed_sizes
379            && self.buffer_alignments == other.buffer_alignments
380    }
381}
382
383impl VTable for ZstdBuffers {
384    type TypedArrayData = ZstdBuffersData;
385    type OperationsVTable = Self;
386    type ValidityVTable = Self;
387
388    fn id(&self) -> ArrayId {
389        static ID: CachedId = CachedId::new("vortex.zstd_buffers");
390        *ID
391    }
392
393    fn validate(
394        &self,
395        data: &Self::TypedArrayData,
396        _dtype: &DType,
397        _len: usize,
398        _slots: &[Option<ArrayRef>],
399    ) -> VortexResult<()> {
400        data.validate()
401    }
402
403    fn nbuffers(array: ArrayView<'_, Self>) -> usize {
404        array.compressed_buffers.len()
405    }
406
407    fn buffer(array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
408        array.compressed_buffers[idx].clone()
409    }
410
411    fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option<String> {
412        Some(format!("compressed_{idx}"))
413    }
414
415    fn with_buffers(
416        &self,
417        array: ArrayView<'_, Self>,
418        buffers: &[BufferHandle],
419    ) -> VortexResult<ArrayParts<Self>> {
420        let mut data = array.data().clone();
421        data.compressed_buffers = buffers.to_vec();
422        Ok(
423            ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data)
424                .with_slots(array.slots().iter().cloned().collect()),
425        )
426    }
427
428    fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
429        format!("child_{idx}")
430    }
431
432    fn serialize(
433        array: ArrayView<'_, Self>,
434        _session: &VortexSession,
435    ) -> VortexResult<Option<Vec<u8>>> {
436        let children: Vec<&ArrayRef> = array.slots().iter().flatten().collect();
437        let child_dtypes = children
438            .iter()
439            .map(|child| child.dtype().try_into())
440            .collect::<VortexResult<Vec<_>>>()?;
441        let child_lens = children.iter().map(|child| child.len() as u64).collect();
442
443        Ok(Some(
444            ZstdBuffersMetadata {
445                inner_encoding_id: array.inner_encoding_id.to_string(),
446                inner_metadata: array.inner_metadata.clone(),
447                uncompressed_sizes: array.uncompressed_sizes.clone(),
448                buffer_alignments: array.buffer_alignments.clone(),
449                child_dtypes,
450                child_lens,
451            }
452            .encode_to_vec(),
453        ))
454    }
455
456    fn deserialize(
457        &self,
458        dtype: &DType,
459        len: usize,
460        metadata: &[u8],
461        buffers: &[BufferHandle],
462        children: &dyn ArrayChildren,
463        session: &VortexSession,
464    ) -> VortexResult<ArrayParts<Self>> {
465        let metadata = ZstdBuffersMetadata::decode(metadata)?;
466        let compressed_buffers: Vec<BufferHandle> = buffers.to_vec();
467
468        // Children belong to inner encodings, and serialization doesn't
469        // preserve their dtypes and values. Check dtypes are recovered from
470        // metadata.
471        vortex_ensure_eq!(metadata.child_dtypes.len(), children.len());
472        vortex_ensure_eq!(metadata.child_lens.len(), children.len());
473
474        let slots: ArraySlots = (0..children.len())
475            .map(|i| {
476                let child_dtype = DType::from_proto(&metadata.child_dtypes[i], session)?;
477                let child_len = usize::try_from(metadata.child_lens[i])?;
478                children.get(i, &child_dtype, child_len).map(Some)
479            })
480            .collect::<VortexResult<Vec<_>>>()?
481            .into();
482
483        let data = ZstdBuffersData {
484            inner_encoding_id: array_id_from_string(&metadata.inner_encoding_id),
485            inner_metadata: metadata.inner_metadata.clone(),
486            compressed_buffers,
487            uncompressed_sizes: metadata.uncompressed_sizes.clone(),
488            buffer_alignments: metadata.buffer_alignments.clone(),
489        };
490
491        data.validate()?;
492        Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots))
493    }
494
495    // with_slots handles child replacement via the slots mechanism
496
497    fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
498        let session = ctx.session();
499        let inner_array = ZstdBuffers::decompress_and_build_inner(&array, session)?;
500        inner_array
501            .execute::<ArrayRef>(ctx)
502            .map(ExecutionResult::done)
503    }
504}
505
506impl OperationsVTable<ZstdBuffers> for ZstdBuffers {
507    fn scalar_at(
508        array: ArrayView<'_, ZstdBuffers>,
509        index: usize,
510        ctx: &mut ExecutionCtx,
511    ) -> VortexResult<Scalar> {
512        // TODO(os): maybe we should not support scalar_at, it is really slow, and adding a cache
513        // layer here is weird. Valid use of zstd buffers array would be by executing it first into
514        // canonical
515        let inner_array = ZstdBuffers::decompress_and_build_inner(
516            &array.into_owned(),
517            &vortex_array::LEGACY_SESSION,
518        )?;
519        inner_array.execute_scalar(index, ctx)
520    }
521}
522
523impl ValidityVTable<ZstdBuffers> for ZstdBuffers {
524    fn validity(
525        array: ArrayView<'_, ZstdBuffers>,
526    ) -> VortexResult<vortex_array::validity::Validity> {
527        if !array.dtype().is_nullable() {
528            return Ok(vortex_array::validity::Validity::NonNullable);
529        }
530
531        let inner_array = ZstdBuffers::decompress_and_build_inner(
532            &array.into_owned(),
533            &vortex_array::LEGACY_SESSION,
534        )?;
535        inner_array.validity()
536    }
537}
538
539#[cfg(test)]
540mod tests {
541    use rstest::rstest;
542    use vortex_array::ArrayContext;
543    use vortex_array::ArrayRef;
544    use vortex_array::IntoArray;
545    use vortex_array::VortexSessionExecute;
546    use vortex_array::array_session;
547    use vortex_array::arrays::PrimitiveArray;
548    use vortex_array::arrays::VarBinViewArray;
549    use vortex_array::assert_arrays_eq;
550    use vortex_array::expr::stats::Precision;
551    use vortex_array::expr::stats::Stat;
552    use vortex_array::expr::stats::StatsProvider;
553    use vortex_array::serde::SerializeOptions;
554    use vortex_array::serde::SerializedArray;
555    use vortex_array::session::ArraySessionExt;
556    use vortex_buffer::ByteBufferMut;
557    use vortex_error::VortexResult;
558    use vortex_session::registry::ReadContext;
559
560    use super::*;
561
562    fn make_primitive_array() -> ArrayRef {
563        PrimitiveArray::from_iter(0i32..100).into_array()
564    }
565
566    fn make_varbinview_array() -> ArrayRef {
567        VarBinViewArray::from_iter_str(["hello", "world", "foo", "bar", "a longer string here"])
568            .into_array()
569    }
570
571    fn make_nullable_primitive_array() -> ArrayRef {
572        PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, Some(5)]).into_array()
573    }
574
575    fn make_nullable_varbinview_array() -> ArrayRef {
576        VarBinViewArray::from_iter_nullable_str([
577            Some("hello"),
578            None,
579            Some("world"),
580            None,
581            Some("a moderately long string for testing"),
582        ])
583        .into_array()
584    }
585
586    fn make_empty_primitive_array() -> ArrayRef {
587        PrimitiveArray::from_iter(Vec::<i32>::new()).into_array()
588    }
589
590    fn make_inlined_varbinview_array() -> ArrayRef {
591        VarBinViewArray::from_iter_str(["hi", "ok", "yes", "no"]).into_array()
592    }
593
594    #[rstest]
595    #[case::primitive(make_primitive_array())]
596    #[case::varbinview(make_varbinview_array())]
597    #[case::nullable_primitive(make_nullable_primitive_array())]
598    #[case::nullable_varbinview(make_nullable_varbinview_array())]
599    #[case::empty_primitive(make_empty_primitive_array())]
600    #[case::inlined_varbinview(make_inlined_varbinview_array())]
601    fn test_roundtrip(#[case] input: ArrayRef) -> VortexResult<()> {
602        let compressed = ZstdBuffers::compress(&input, 3, &array_session())?;
603
604        assert_eq!(compressed.len(), input.len());
605        assert_eq!(compressed.dtype(), input.dtype());
606
607        let mut ctx = array_session().create_execution_ctx();
608        let decompressed = compressed.into_array().execute::<ArrayRef>(&mut ctx)?;
609
610        assert_arrays_eq!(input, decompressed, &mut ctx);
611        Ok(())
612    }
613
614    #[rstest]
615    #[case::primitive(make_primitive_array())]
616    #[case::varbinview(make_varbinview_array())]
617    #[case::nullable_primitive(make_nullable_primitive_array())]
618    #[case::nullable_varbinview(make_nullable_varbinview_array())]
619    #[case::empty_primitive(make_empty_primitive_array())]
620    #[case::inlined_varbinview(make_inlined_varbinview_array())]
621    fn test_serde_roundtrip(#[case] input: ArrayRef) -> VortexResult<()> {
622        let session = array_session();
623        session.arrays().register(ZstdBuffers);
624
625        let compressed = ZstdBuffers::compress(&input, 3, &session)?.into_array();
626        let dtype = compressed.dtype().clone();
627        let len = compressed.len();
628
629        let array_ctx = ArrayContext::empty();
630        let serialized =
631            compressed.serialize(&array_ctx, &session, &SerializeOptions::default())?;
632
633        let mut concat = ByteBufferMut::empty();
634        for buf in serialized {
635            concat.extend_from_slice(buf.as_ref());
636        }
637        let parts = SerializedArray::try_from(concat.freeze())?;
638        let decoded = parts.decode(&dtype, len, &ReadContext::new(array_ctx.to_ids()), &session)?;
639
640        let mut ctx = session.create_execution_ctx();
641        let decoded = decoded.execute::<ArrayRef>(&mut ctx)?;
642        assert_arrays_eq!(input, decoded, &mut ctx);
643        Ok(())
644    }
645
646    #[test]
647    fn test_compress_inherits_stats() -> VortexResult<()> {
648        let input = make_primitive_array();
649        input.statistics().set(Stat::Min, Precision::exact(0i32));
650
651        let compressed = ZstdBuffers::compress(&input, 3, &array_session())?;
652
653        assert!(!compressed.statistics().get(Stat::Min).is_absent());
654        Ok(())
655    }
656
657    #[test]
658    fn test_validity_delegates_for_nullable_input() -> VortexResult<()> {
659        let input = make_nullable_primitive_array();
660        let compressed = ZstdBuffers::compress(&input, 3, &array_session())?.into_array();
661
662        let mut ctx = array_session().create_execution_ctx();
663        assert_eq!(compressed.all_valid(&mut ctx)?, input.all_valid(&mut ctx)?);
664        assert_eq!(
665            compressed.all_invalid(&mut ctx)?,
666            input.all_invalid(&mut ctx)?
667        );
668
669        for i in 0..input.len() {
670            assert_eq!(
671                compressed.is_valid(i, &mut ctx)?,
672                input.is_valid(i, &mut ctx)?
673            );
674        }
675
676        Ok(())
677    }
678}
vortex_zstd/zstd_buffers.rs

vortex_zstd/
zstd_buffers.rs
