tidaldb/tidal/tests/m12p2_ann_retrieve.rs

#![allow(
    clippy::too_many_lines,
    clippy::unwrap_used,
    clippy::cast_possible_truncation
)]
//! m12p2 — ANN candidate generation + cached `SignalRanked` in RETRIEVE.
//!
//! These prove the G1 unblock end-to-end through the engine: `for_you` sources
//! its candidates by nearest-neighbour over the user's preference vector (not an
//! arbitrary low-id scan slice), and `trending` sources its candidates from the
//! cached per-signal-type top-K (reaching the actually-viewed items at any id).
//! The distinguishing trick: seed a corpus far larger than the scan cap and
//! place the relevant items at HIGH ids — a scan (capped at the lowest ~240 ids)
//! cannot reach them, so their presence in the feed proves the index path ran.

use std::collections::HashMap;

use tidaldb::{
    TidalDb,
    query::retrieve::{ProfileRef, RetrieveBuilder},
    schema::{DecaySpec, EntityId, EntityKind, SchemaBuilder, Timestamp, Window},
};

const DIM: usize = 8;
/// Corpus far larger than the executor's scan cap (~240) so high-id relevant
/// items are unreachable by a scan — their presence proves the index path.
const N: u64 = 2000;

/// A one-hot embedding on axis `i % DIM`: items sharing an axis are identical
/// (distance 0), so a preference vector on axis A makes every axis-A item — at
/// any id — a nearest neighbour.
fn one_hot(axis: usize) -> Vec<f32> {
    let mut v = vec![0.0_f32; DIM];
    v[axis] = 1.0;
    v
}

fn schema_with_embeddings() -> tidaldb::schema::Schema {
    let mut builder = SchemaBuilder::new();
    let _ = builder
        .signal(
            "view",
            EntityKind::Item,
            DecaySpec::Exponential {
                half_life: std::time::Duration::from_secs(7 * 24 * 3600),
            },
        )
        .windows(&[Window::OneHour, Window::TwentyFourHours])
        .velocity(true)
        .add();
    let _ = builder
        .signal(
            "like",
            EntityKind::Item,
            DecaySpec::Exponential {
                half_life: std::time::Duration::from_secs(30 * 24 * 3600),
            },
        )
        .windows(&[Window::TwentyFourHours])
        .velocity(false)
        .positive_engagement(true)
        .add();
    builder.embedding_slot("content", EntityKind::Item, DIM);
    builder.build().unwrap()
}

#[test]
fn for_you_ann_surfaces_nearest_to_preference_across_the_whole_corpus() {
    let db = TidalDb::builder()
        .ephemeral()
        .with_schema(schema_with_embeddings())
        .open()
        .unwrap();

    let ts = Timestamp::now();
    for id in 1..=N {
        db.write_item_with_metadata(EntityId::new(id), &HashMap::new())
            .unwrap();
        db.write_item_embedding(EntityId::new(id), &one_hot((id % DIM as u64) as usize))
            .unwrap();
    }

    // A HIGH-id axis-5 item, made engaging with views. Its id (1997) is far beyond
    // the executor's scan pool (limit × 10 = 500 lowest ids), so a scan can NEVER
    // surface it — only ANN candidate-gen over the axis-5 preference vector reaches
    // it. (1997 % 8 == 5.)
    let far = 1997u64;
    for _ in 0..50 {
        db.signal("view", EntityId::new(far), 1.0, ts).unwrap();
    }

    // User 7's learned taste is axis 5. Set the preference vector directly (the
    // signal-driven path that builds it is covered elsewhere; here we isolate the
    // ANN candidate-gen). Axis-5 items are ids 5, 13, 21, …, 1997.
    let user = 7u64;
    assert!(db.preference_vectors().set(user, one_hot(5)));

    let results = db
        .retrieve(
            &RetrieveBuilder::new(EntityKind::Item, ProfileRef::new("for_you"))
                .for_user(user)
                .limit(50)
                .build()
                .unwrap(),
        )
        .unwrap();
    let ids: Vec<u64> = results.items.iter().map(|r| r.entity_id.as_u64()).collect();
    assert_eq!(ids.len(), 50, "expected a full page");

    let axis5 = ids.iter().filter(|id| *id % DIM as u64 == 5).count();
    // ANN over the axis-5 preference vector ⇒ the page is dominated by axis-5
    // items (for_you injects ~10% exploration, so allow for that). A scan would
    // return ~1/8 axis-5 from the low-id slice — nowhere near this.
    assert!(
        axis5 >= 30,
        "expected the page dominated by axis-5 (ANN); got {axis5}/50 — looks like a scan"
    );
    // The engaging high-id item is reachable ONLY through the ANN candidate pool
    // (it sits beyond the scan cap), so its presence proves the ANN index — not the
    // universe scan — generated these candidates.
    assert!(
        ids.contains(&far),
        "ANN must reach the engaging high-id nearest neighbour {far} that the scan pool \
         (ids 1..=500) cannot; got {ids:?}"
    );
}

#[test]
fn for_you_without_preference_vector_falls_back_to_scan() {
    // An anonymous-ish read (a user with no preference vector) must still serve —
    // the ANN strategy degrades to a scan rather than erroring or returning empty.
    let db = TidalDb::builder()
        .ephemeral()
        .with_schema(schema_with_embeddings())
        .open()
        .unwrap();
    for id in 1..=300u64 {
        db.write_item_with_metadata(EntityId::new(id), &HashMap::new())
            .unwrap();
        db.write_item_embedding(EntityId::new(id), &one_hot((id % DIM as u64) as usize))
            .unwrap();
    }
    let results = db
        .retrieve(
            &RetrieveBuilder::new(EntityKind::Item, ProfileRef::new("for_you"))
                .for_user(999) // no preference vector set
                .limit(20)
                .build()
                .unwrap(),
        )
        .unwrap();
    assert!(
        !results.items.is_empty(),
        "for_you with no preference vector must fall back to scan, not return empty"
    );
}

#[test]
fn related_ann_surfaces_seed_neighbours_across_the_whole_corpus() {
    let db = TidalDb::builder()
        .ephemeral()
        .with_schema(schema_with_embeddings())
        .open()
        .unwrap();
    let ts = Timestamp::now();
    for id in 1..=N {
        db.write_item_with_metadata(EntityId::new(id), &HashMap::new())
            .unwrap();
        db.write_item_embedding(EntityId::new(id), &one_hot((id % DIM as u64) as usize))
            .unwrap();
    }
    // Seed item 3 is axis 3; an engaging high-id axis-3 item (1995 % 8 == 3) is
    // reachable only through the seed's ANN neighbour pool, never the scan slice.
    let far = 1995u64;
    for _ in 0..50 {
        db.signal("view", EntityId::new(far), 1.0, ts).unwrap();
    }
    let seed = EntityId::new(3);
    let results = db
        .retrieve(
            &RetrieveBuilder::new(EntityKind::Item, ProfileRef::new("related"))
                .similar_to(seed)
                .limit(50)
                .build()
                .unwrap(),
        )
        .unwrap();
    let ids: Vec<u64> = results.items.iter().map(|r| r.entity_id.as_u64()).collect();
    assert_eq!(ids.len(), 50, "expected a full page");
    // The decisive proof of ANN candidate generation: the seed's high-id
    // neighbour (1995, beyond the scan pool of ids 1..=500) is in the page. Only
    // the ANN search over the seed embedding can reach it; a scan never could.
    // (related ranks the ANN-similar pool by engagement/recency — surfacing
    // SIMILAR items by similarity score is a related-scoring follow-up; here we
    // assert candidate generation, the m12p2 work item.)
    assert!(
        ids.contains(&far),
        "related ANN must reach the seed's high-id neighbour {far} beyond the scan pool; got {ids:?}"
    );
    // Axis-3 (the seed's true neighbours) appear well above the ~1/8 a scan slice
    // would yield, confirming the pool is the seed's neighbourhood.
    let axis3 = ids.iter().filter(|id| *id % DIM as u64 == 3).count();
    assert!(
        axis3 >= 8,
        "related's candidate pool should be the seed's axis-3 neighbourhood; got {axis3}/50"
    );
}

#[test]
fn trending_signalranked_reaches_high_id_viewed_items() {
    let db = TidalDb::builder()
        .ephemeral()
        .with_schema(schema_with_embeddings())
        .open()
        .unwrap();
    let ts = Timestamp::now();
    for id in 1..=N {
        db.write_item_with_metadata(EntityId::new(id), &HashMap::new())
            .unwrap();
    }
    // Heavy views on three HIGH-id items (unreachable by a low-id scan), light
    // views on a few low ids as noise.
    let hot = [1500u64, 1700, 1900];
    for &id in &hot {
        for _ in 0..100 {
            db.signal("view", EntityId::new(id), 1.0, ts).unwrap();
        }
    }
    for id in 1..=20u64 {
        db.signal("view", EntityId::new(id), 1.0, ts).unwrap();
    }

    let results = db
        .retrieve(
            &RetrieveBuilder::new(EntityKind::Item, ProfileRef::new("trending"))
                .limit(20)
                .build()
                .unwrap(),
        )
        .unwrap();
    let ids: std::collections::HashSet<u64> =
        results.items.iter().map(|r| r.entity_id.as_u64()).collect();
    for &id in &hot {
        assert!(
            ids.contains(&id),
            "trending must surface heavily-viewed high-id item {id} via SignalRanked; \
             a scan (low ~240 ids) could never reach it. got {ids:?}"
        );
    }
}

#[test]
fn trending_cache_reflects_new_writes() {
    // The SignalRanked top-K cache must not serve stale results across writes on a
    // small ledger (always-fresh policy): a newly-hot item appears immediately.
    let db = TidalDb::builder()
        .ephemeral()
        .with_schema(schema_with_embeddings())
        .open()
        .unwrap();
    let ts = Timestamp::now();
    for id in 1..=100u64 {
        db.write_item_with_metadata(EntityId::new(id), &HashMap::new())
            .unwrap();
    }
    for _ in 0..50 {
        db.signal("view", EntityId::new(10), 1.0, ts).unwrap();
    }
    let q = || {
        db.retrieve(
            &RetrieveBuilder::new(EntityKind::Item, ProfileRef::new("trending"))
                .limit(10)
                .build()
                .unwrap(),
        )
        .unwrap()
    };
    let first: std::collections::HashSet<u64> =
        q().items.iter().map(|r| r.entity_id.as_u64()).collect();
    assert!(
        first.contains(&10),
        "item 10 should be trending after its views"
    );

    // New burst on a different item AFTER the first query (which warmed the cache).
    for _ in 0..200 {
        db.signal("view", EntityId::new(90), 1.0, ts).unwrap();
    }
    let second: std::collections::HashSet<u64> =
        q().items.iter().map(|r| r.entity_id.as_u64()).collect();
    assert!(
        second.contains(&90),
        "the cache must reflect the new burst on item 90 (always-fresh on a small ledger); got {second:?}"
    );
}