Example 4: Appendix and embeddings

// Embed binary data (an appendix) in a pack: fixed-size embedding
// rows and an f32 tensor. Then read them back and bind each
// embedding to its entity.
use plxi_sdk::{
    write, AppendixSpec, BindProfile, OpenOptions, Pack, Record,
    COMPACT_RECORD_SIZE,
};

fn row(seed: u8) -> Vec<u8> {
    // Stand-in bytes. A real row is one Compact embedding record.
    (0..COMPACT_RECORD_SIZE)
        .map(|i| (i as u8).wrapping_add(seed))
        .collect()
}

fn main() -> Result<(), plxi_sdk::Error> {
    let lines = [
        r#"{"k":"ent","id":"e1","t":"Person"}"#,
        r#"{"k":"ent","id":"e2","t":"Person"}"#,
        // An `embref` points an entity at row N of appendix section M.
        concat!(
            r#"{"k":"embref","entity_id":"e1","target":"#,
            r#"{"local":{"section_index":0,"record_index":0}}}"#
        ),
        concat!(
            r#"{"k":"embref","entity_id":"e2","target":"#,
            r#"{"local":{"section_index":0,"record_index":1}}}"#
        ),
    ];
    let records = lines
        .iter()
        .map(|l| Record::from_json_line(l))
        .collect::<Result<Vec<_>, _>>()?;

    // The constructors return None on a caller mistake (wrong row
    // length, shape that does not match the data length).
    let appendix = AppendixSpec::compact(&[row(1), row(2)])
        .expect("every row is COMPACT_RECORD_SIZE bytes")
        .with_tensor_f32("probe.f32", &[2, 2], &[0.5, -1.25, 2.0, 0.0])
        .expect("2 x 2 = 4 values");
    assert!(AppendixSpec::compact(&[vec![0u8; 7]]).is_none());

    let bytes = write(&records, Some(&appendix))?;
    let pack = Pack::from_bytes(bytes, OpenOptions::default())?;
    assert!(pack.has_appendix());

    let app = pack.appendix()?.expect("this pack has an appendix");
    let sections = app.sections();
    assert_eq!(sections[0].section_type, 0x0020); // Compact
    assert_eq!(sections[0].record_count, 2);
    assert_eq!(sections[0].record_stride as usize, COMPACT_RECORD_SIZE);
    let data = app.section_bytes(0)?.expect("section 0 exists");
    assert_eq!(data.len(), 2 * COMPACT_RECORD_SIZE);
    assert_eq!(app.tensor_names(), ["probe.f32"]);
    let tensor = app.tensor_f32("probe.f32")?;
    assert_eq!(tensor, Some(vec![0.5, -1.25, 2.0, 0.0]));
    assert_eq!(app.tensor_f32("absent")?, None);

    // One bound embedding per `embref`, in record order.
    let bound = pack.bind_embrefs(BindProfile::Compact)?;
    assert_eq!(bound.len(), 2);
    assert_eq!(bound[0].entity_id, "e1");
    assert_eq!(bound[0].row, row(1));
    assert_eq!(bound[1].entity_id, "e2");
    assert_eq!(bound[1].row, row(2));

    let report = pack.verify()?;
    assert!(report.is_ok());
    let embrefs = &report.embrefs;
    assert_eq!((embrefs.total, embrefs.resolved), (2, 2));
    println!("row size = {COMPACT_RECORD_SIZE} bytes");
    Ok(())
}