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

vortex_zstd/
zstd_buffers.rs
