tetration 0.1.9

Chunked tensor .tet format, tet CLI, and optional C ABI (tetration-ffi)
Documentation
//! Dense embedder materialize API (`materialize_query_selection`, `materialize_query_transform_ram`).

use super::fixture::write_multichunk_2x3_tiles;
use crate::layout::mmap_file_read;
use crate::query::{
    DenseBuffer, materialize_query_selection, materialize_query_transform_ram, parse_query_json,
    plan_query_with_tet_mmap, plan_read_for_document, validate_query,
};

fn pop_std_1_to_6() -> f64 {
    let mean = 3.5;
    let var = (1..=6)
        .map(|n| {
            let d = f64::from(n) - mean;
            d * d
        })
        .sum::<f64>()
        / 6.0;
    var.sqrt()
}

#[test]
fn materialize_query_selection_f32_full_tensor() {
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("dense.tet");
    write_multichunk_2x3_tiles(&path, "a");
    let doc = parse_query_json(r#"{"dataset":"a"}"#).unwrap();
    validate_query(&doc).unwrap();
    let mmap = mmap_file_read(&path).unwrap();
    let out = materialize_query_selection(&doc, &mmap).unwrap();
    assert_eq!(out.shape, vec![2, 3]);
    match out.buffer {
        DenseBuffer::F32(vals) => {
            assert_eq!(vals.len(), 6);
            let want = [1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0];
            for (got, expected) in vals.iter().zip(want) {
                assert!((got - expected).abs() < 1e-5, "got {got}, want {expected}");
            }
        }
        other => panic!("expected f32 buffer, got {other:?}"),
    }
}

#[test]
fn materialize_query_selection_honors_slice() {
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("slice.tet");
    write_multichunk_2x3_tiles(&path, "a");
    let doc = parse_query_json(
        r#"{"dataset":"a","selection":[{"start":0,"stop":1},{"start":0,"stop":3}]}"#,
    )
    .unwrap();
    validate_query(&doc).unwrap();
    let mmap = mmap_file_read(&path).unwrap();
    let out = materialize_query_selection(&doc, &mmap).unwrap();
    assert_eq!(out.shape, vec![1, 3]);
    match out.buffer {
        DenseBuffer::F32(vals) => assert_eq!(vals, vec![1.0, 2.0, 3.0]),
        other => panic!("expected f32 buffer, got {other:?}"),
    }
}

#[test]
fn materialize_query_selection_rejects_operation() {
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("op.tet");
    write_multichunk_2x3_tiles(&path, "a");
    let doc = parse_query_json(r#"{"dataset":"a","mean":[]}"#).unwrap();
    validate_query(&doc).unwrap();
    let mmap = mmap_file_read(&path).unwrap();
    let err = materialize_query_selection(&doc, &mmap).unwrap_err();
    assert!(err.to_string().contains("selection-only"), "{err}");
}

#[test]
fn plan_read_for_document_matches_query_read_plan() {
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("plan.tet");
    write_multichunk_2x3_tiles(&path, "a");
    let doc = parse_query_json(
        r#"{"dataset":"a","selection":[{"start":0,"stop":1},{"start":0,"stop":3}]}"#,
    )
    .unwrap();
    validate_query(&doc).unwrap();
    let mmap = mmap_file_read(&path).unwrap();
    let planned = plan_read_for_document(&doc, &mmap).unwrap();
    let response = plan_query_with_tet_mmap(&doc, None, &mmap, None).unwrap();
    let rp = response.read_plan.as_ref().unwrap();
    assert_eq!(
        planned.read_plan.logical_selection_shape,
        rp.logical_selection_shape
    );
    assert_eq!(planned.read_plan.chunk_count, rp.chunk_count);
    assert_eq!(planned.read_plan.chunks.len(), rp.chunks.len());
}

#[test]
fn materialize_query_transform_ram_zscore_full_buffer() {
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("zscore.tet");
    write_multichunk_2x3_tiles(&path, "a");
    let doc = parse_query_json(r#"{"dataset":"a","transform":{"method":"zscore"},"write":"ram"}"#)
        .unwrap();
    validate_query(&doc).unwrap();
    let mmap = mmap_file_read(&path).unwrap();
    let out = materialize_query_transform_ram(&doc, &mmap, &path, None).unwrap();
    assert_eq!(out.shape, vec![2, 3]);
    let std = pop_std_1_to_6();
    match out.buffer {
        DenseBuffer::F32(vals) => {
            assert_eq!(vals.len(), 6);
            for (i, &v) in vals.iter().enumerate() {
                let raw = f64::from((i + 1) as f32);
                let expected = ((raw - 3.5) / std) as f32;
                assert!(
                    (v - expected).abs() < 1e-4,
                    "index {i}: got {v}, expected {expected}"
                );
            }
        }
        other => panic!("expected f32 buffer, got {other:?}"),
    }
}

#[test]
fn materialize_query_transform_ram_rejects_sidecar_write() {
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("sidecar.tet");
    write_multichunk_2x3_tiles(&path, "a");
    let doc = parse_query_json(
        r#"{"dataset":"a","transform":{"method":"zscore"},"write":{"target":"sidecar","timestamp":false}}"#,
    )
    .unwrap();
    validate_query(&doc).unwrap();
    let mmap = mmap_file_read(&path).unwrap();
    let err = materialize_query_transform_ram(&doc, &mmap, &path, None).unwrap_err();
    assert!(err.to_string().contains("ram"), "{err}");
}

#[test]
fn materialize_query_transform_ram_requires_transform_operation() {
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("plain.tet");
    write_multichunk_2x3_tiles(&path, "a");
    let doc = parse_query_json(r#"{"dataset":"a"}"#).unwrap();
    validate_query(&doc).unwrap();
    let mmap = mmap_file_read(&path).unwrap();
    let err = materialize_query_transform_ram(&doc, &mmap, &path, None).unwrap_err();
    assert!(err.to_string().contains("transform"), "{err}");
}