quatzal-storage 0.1.0

// SPDX-License-Identifier: Apache-2.0
//! PAX (Partition Attributes Across) row codec.
//!
//! Every stored record separates the small, frequently-mutated scalar columns from the
//! large, rarely-mutated vector column. A scalar-only update encodes its `vector` field as
//! `VectorField::Unchanged` instead of re-serializing the vector bytes, so the WAL/MemTable
//! write path never pays to move vector-sized data for a metadata-only change. See
//! `UACE-FR-1.2` (encode/decode identity) and `UACE-FR-1.3` (scalar updates don't touch the
//! vector bytes) in `Documents/SRS.md`.

use std::collections::BTreeMap;

use quatzal_schema::{Row, UaceError, UaceResult, Value};

#[derive(Debug, Clone, PartialEq)]
pub enum VectorField {
    /// This row has never had a vector column.
    Unset,
    /// This record doesn't change the vector; the current value lives in an older record
    /// for the same key (resolved by walking MemTable -> L0 -> L1, newest to oldest).
    Unchanged,
    /// This record sets the vector to exactly these encoded bytes (bincode of `Vec<f32>`).
    Set(Vec<u8>),
}

#[derive(Debug, Clone, PartialEq)]
pub struct PaxRecord {
    pub key: Vec<u8>,
    pub seq: u64,
    pub tombstone: bool,
    /// bincode-encoded `BTreeMap<String, Value>` (empty map for pure tombstones).
    pub scalar_blob: Vec<u8>,
    pub vector: VectorField,
}

impl PaxRecord {
    /// Encode a full row (insert / full replace): the vector, if present, is always `Set`.
    pub fn encode_full(seq: u64, row: &Row) -> UaceResult<Self> {
        let scalar_blob = bincode::serialize(&row.scalars)
            .map_err(|e| UaceError::Codec(format!("scalar encode: {e}")))?;
        let vector = match &row.vector {
            Some(v) => VectorField::Set(
                bincode::serialize(v)
                    .map_err(|e| UaceError::Codec(format!("vector encode: {e}")))?,
            ),
            None => VectorField::Unset,
        };
        Ok(PaxRecord {
            key: row.key.clone(),
            seq,
            tombstone: false,
            scalar_blob,
            vector,
        })
    }

    /// Encode a scalar-only update: `UACE-FR-1.3` — the vector field is always `Unchanged`,
    /// never a re-serialization of whatever the vector currently is.
    pub fn encode_scalar_update(
        seq: u64,
        key: &[u8],
        scalars: &BTreeMap<String, Value>,
    ) -> UaceResult<Self> {
        let scalar_blob = bincode::serialize(scalars)
            .map_err(|e| UaceError::Codec(format!("scalar encode: {e}")))?;
        Ok(PaxRecord {
            key: key.to_vec(),
            seq,
            tombstone: false,
            scalar_blob,
            vector: VectorField::Unchanged,
        })
    }

    pub fn encode_tombstone(seq: u64, key: &[u8]) -> Self {
        PaxRecord {
            key: key.to_vec(),
            seq,
            tombstone: true,
            scalar_blob: Vec::new(),
            vector: VectorField::Unset,
        }
    }

    /// Decode this record back into a `Row`. Only valid when `vector` is `Set` or `Unset` —
    /// callers holding an `Unchanged` record must resolve the vector first (see
    /// `crate::engine::resolve`) before calling this.
    pub fn decode(&self) -> UaceResult<Row> {
        self.decode_with_vector(None)
    }

    pub fn decode_with_vector(&self, resolved_vector: Option<Vec<f32>>) -> UaceResult<Row> {
        let scalars: BTreeMap<String, Value> = bincode::deserialize(&self.scalar_blob)
            .map_err(|e| UaceError::Codec(format!("scalar decode: {e}")))?;
        let vector = match &self.vector {
            VectorField::Set(bytes) => Some(
                bincode::deserialize::<Vec<f32>>(bytes)
                    .map_err(|e| UaceError::Codec(format!("vector decode: {e}")))?,
            ),
            VectorField::Unset => None,
            VectorField::Unchanged => resolved_vector,
        };
        Ok(Row {
            key: self.key.clone(),
            scalars,
            vector,
        })
    }

    /// Frame this record for the WAL / SSTable file format:
    /// `[key_len:u32][key][seq:u64][flags:u8][scalar_len:u32][scalar][vector_len:u32][vector]?`
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut buf = Vec::with_capacity(4 + self.key.len() + 8 + 1 + 4 + self.scalar_blob.len());
        buf.extend_from_slice(&(self.key.len() as u32).to_le_bytes());
        buf.extend_from_slice(&self.key);
        buf.extend_from_slice(&self.seq.to_le_bytes());

        let vector_tag: u8 = match &self.vector {
            VectorField::Unset => 0,
            VectorField::Unchanged => 1,
            VectorField::Set(_) => 2,
        };
        let flags = (self.tombstone as u8) | (vector_tag << 1);
        buf.push(flags);

        buf.extend_from_slice(&(self.scalar_blob.len() as u32).to_le_bytes());
        buf.extend_from_slice(&self.scalar_blob);

        if let VectorField::Set(bytes) = &self.vector {
            buf.extend_from_slice(&(bytes.len() as u32).to_le_bytes());
            buf.extend_from_slice(bytes);
        }
        buf
    }

    /// Parse one record starting at `buf[0]`, returning the record and the number of bytes
    /// consumed. Used to walk a WAL/SSTable file record-by-record.
    pub fn from_bytes(buf: &[u8]) -> UaceResult<(Self, usize)> {
        let mut pos = 0usize;
        let key_len = read_u32(buf, &mut pos)? as usize;
        let key = read_slice(buf, &mut pos, key_len)?.to_vec();
        let seq = read_u64(buf, &mut pos)?;
        let flags = read_u8(buf, &mut pos)?;
        let tombstone = flags & 0x1 != 0;
        let vector_tag = (flags >> 1) & 0x3;

        let scalar_len = read_u32(buf, &mut pos)? as usize;
        let scalar_blob = read_slice(buf, &mut pos, scalar_len)?.to_vec();

        let vector = match vector_tag {
            0 => VectorField::Unset,
            1 => VectorField::Unchanged,
            2 => {
                let vlen = read_u32(buf, &mut pos)? as usize;
                let vbytes = read_slice(buf, &mut pos, vlen)?.to_vec();
                VectorField::Set(vbytes)
            }
            other => return Err(UaceError::Codec(format!("unknown vector tag {other}"))),
        };

        Ok((
            PaxRecord {
                key,
                seq,
                tombstone,
                scalar_blob,
                vector,
            },
            pos,
        ))
    }
}

fn read_u8(buf: &[u8], pos: &mut usize) -> UaceResult<u8> {
    let s = read_slice(buf, pos, 1)?;
    Ok(s[0])
}

fn read_u32(buf: &[u8], pos: &mut usize) -> UaceResult<u32> {
    let s = read_slice(buf, pos, 4)?;
    Ok(u32::from_le_bytes(s.try_into().unwrap()))
}

fn read_u64(buf: &[u8], pos: &mut usize) -> UaceResult<u64> {
    let s = read_slice(buf, pos, 8)?;
    Ok(u64::from_le_bytes(s.try_into().unwrap()))
}

fn read_slice<'a>(buf: &'a [u8], pos: &mut usize, len: usize) -> UaceResult<&'a [u8]> {
    if *pos + len > buf.len() {
        return Err(UaceError::Codec("truncated record".into()));
    }
    let s = &buf[*pos..*pos + len];
    *pos += len;
    Ok(s)
}