vortex_zstd/
lib.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Zstd-backed compression encodings for variable-width Vortex arrays.
5//!
6//! [`ZstdArray`] stores UTF-8 or binary values as one or more zstd frames, optionally sharing a
7//! trained dictionary across frames. Frame metadata lets slices decompress only the frames that can
8//! contribute values to the requested row range.
9//!
10//! [`ZstdBuffers`] stores the buffers of another encoding as independently compressed zstd
11//! buffers while preserving the inner encoding metadata.
12//!
13//! This crate exposes array encodings only. Compression scheme selection is wired through
14//! `vortex-btrblocks` and file writing. To deserialize arrays manually, register the encoding in the
15//! array session:
16//!
17//! ```rust
18//! use vortex_array::session::ArraySessionExt;
19//!
20//! let session = vortex_array::array_session();
21//! session.arrays().register(vortex_zstd::Zstd);
22//! ```
23
24pub use array::*;
25use vortex_array::dtype::proto::dtype as pb;
26use vortex_array::session::ArraySessionExt;
27use vortex_edition::EditionSessionExt;
28use vortex_error::VortexExpect;
29use vortex_error::VortexResult;
30use vortex_error::vortex_ensure;
31use vortex_error::vortex_err;
32use vortex_session::VortexSession;
33pub use zstd_buffers::*;
34
35mod array;
36mod compute;
37pub mod editions;
38mod rules;
39mod slice;
40mod zstd_buffers;
41
42#[cfg(test)]
43mod test;
44
45/// Register the Zstd encodings and their optional edition declaration with a Vortex session.
46pub fn initialize(session: &VortexSession) {
47    session.arrays().register(Zstd);
48    session.arrays().register(ZstdBuffers);
49    if session.editions().find(&editions::ZSTD_2026_02).is_none() {
50        session
51            .editions()
52            .declare_family(&editions::FAMILY)
53            .map_err(|error| vortex_err!("{error}"))
54            .vortex_expect("Zstd edition family is valid");
55        session
56            .register_edition(&editions::DECLARATION)
57            .map_err(|error| vortex_err!("{error}"))
58            .vortex_expect("Zstd edition declaration is valid");
59    }
60}
61
62/// Ensure Vortex metadata agrees with the content size declared by a zstd frame.
63pub(crate) fn validate_frame_content_size(
64    frame: &[u8],
65    metadata_size: u64,
66    index: usize,
67) -> VortexResult<()> {
68    let frame_content_size = zstd::zstd_safe::get_frame_content_size(frame)
69        .map_err(|error| vortex_err!("Invalid zstd frame {index}: {error}"))?
70        .ok_or_else(|| vortex_err!("Zstd frame {index} does not declare a content size"))?;
71    vortex_ensure!(
72        metadata_size == frame_content_size,
73        "Zstd frame {index} metadata declares {metadata_size} uncompressed bytes, but its header declares {frame_content_size}"
74    );
75    Ok(())
76}
77
78#[derive(Clone, prost::Message)]
79/// Metadata for one zstd frame.
80pub struct ZstdFrameMetadata {
81    /// Uncompressed byte size of this frame.
82    #[prost(uint64, tag = "1")]
83    pub uncompressed_size: u64,
84    /// Number of valid values stored in this frame.
85    #[prost(uint64, tag = "2")]
86    pub n_values: u64,
87}
88
89#[derive(Clone, prost::Message)]
90/// Serialized metadata for a [`ZstdArray`].
91pub struct ZstdMetadata {
92    // optional, will be 0 if there's no dictionary
93    /// Dictionary size in bytes, or `0` when no dictionary is present.
94    #[prost(uint32, tag = "1")]
95    pub dictionary_size: u32,
96    /// Metadata for each compressed frame.
97    #[prost(message, repeated, tag = "2")]
98    pub frames: Vec<ZstdFrameMetadata>,
99}
100
101#[derive(Clone, prost::Message)]
102/// Serialized metadata for the unstable `ZstdBuffers` encoding.
103pub struct ZstdBuffersMetadata {
104    /// Encoding id of the inner array whose buffers were compressed.
105    #[prost(string, tag = "1")]
106    pub inner_encoding_id: String,
107    /// Serialized metadata of the inner array.
108    #[prost(bytes = "vec", tag = "2")]
109    pub inner_metadata: Vec<u8>,
110    /// Uncompressed byte size of each compressed buffer.
111    #[prost(uint64, repeated, tag = "3")]
112    pub uncompressed_sizes: Vec<u64>,
113    /// Alignment of each buffer in bytes (must be a power of two).
114    #[prost(uint32, repeated, tag = "4")]
115    pub buffer_alignments: Vec<u32>,
116    /// DType of child arrays. Children belong to inner encodings, and their
117    /// dtypes don't persist after serialization, so we need to retrieve them
118    /// from metadata.
119    #[prost(message, repeated, tag = "5")]
120    pub child_dtypes: Vec<pb::DType>,
121    /// Length of each child array, ordered as "child_dtypes"
122    #[prost(uint64, repeated, tag = "6")]
123    pub child_lens: Vec<u64>,
124}
vortex_zstd/lib.rs

vortex_zstd/
lib.rs
