Skip to main content

urna_format/writer/
encoding_choice.rs

1//! Wire-encoding and dtype choices exposed by the builder. The two
2//! enums are 1:1 mapped to `SECTION_ENCODING_*` values; this layer
3//! keeps the public API enum-friendly and the layout module pure.
4
5use crate::layout::{
6    SECTION_ENCODING_FLOAT16, SECTION_ENCODING_INT4, SECTION_ENCODING_INT8, SECTION_ENCODING_RAW,
7    SECTION_ENCODING_ZSTD,
8};
9
10/// Timestamp written to the manifest when the builder is in reproducible
11/// mode. Chosen so that two builds with identical inputs produce
12/// identical bytes regardless of when they ran.
13pub const REPRODUCIBLE_CREATED: &str = "1970-01-01T00:00:00Z";
14
15/// Wire encoding choice for a non-embedding section. Embedding encoding
16/// is controlled by `EmbeddingDType` since dtype and encoding are
17/// intertwined.
18#[derive(Clone, Copy, Debug, PartialEq, Eq)]
19pub enum SectionEncoding {
20    Raw,
21    Zstd,
22}
23
24impl SectionEncoding {
25    pub fn id(self) -> u32 {
26        match self {
27            Self::Raw => SECTION_ENCODING_RAW,
28            Self::Zstd => SECTION_ENCODING_ZSTD,
29        }
30    }
31}
32
33/// Embedding dtype + on-disk encoding. The two are 1:1 in v1 - float32
34/// implies raw f32 LE, float16 implies raw f16 LE, int8 implies the
35/// quantized prefix layout (see `encoding::encode_int8_embeddings`), int4
36/// implies the block-64 per-group layout (see `encode_int4_embeddings`).
37#[derive(Clone, Copy, Debug, PartialEq, Eq)]
38pub enum EmbeddingDType {
39    Float32,
40    Float16,
41    Int8,
42    Int4,
43}
44
45impl EmbeddingDType {
46    pub fn manifest_str(self) -> &'static str {
47        match self {
48            Self::Float32 => "float32",
49            Self::Float16 => "float16",
50            Self::Int8 => "int8",
51            Self::Int4 => "int4",
52        }
53    }
54    pub fn encoding(self) -> u32 {
55        match self {
56            Self::Float32 => SECTION_ENCODING_RAW,
57            Self::Float16 => SECTION_ENCODING_FLOAT16,
58            Self::Int8 => SECTION_ENCODING_INT8,
59            Self::Int4 => SECTION_ENCODING_INT4,
60        }
61    }
62}