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}