pg_infer Design Document
========================

PURPOSE
-------

pg_infer makes transformer model internals queryable through SQL.
Instead of treating models as opaque prediction machines, pg_infer
exposes the learned features, gate activations, and semantic
relationships that models encode in their weights.

The key insight: transformer models organize knowledge in
interpretable gate vectors.  Each "feature" in a layer encodes a
concept (a token or cluster of tokens the model has learned to
detect).  By querying these gate vectors with an input embedding,
you can ask "what does the model know about X?" without running
a forward pass.


CORE CONCEPTS
-------------

Vindex (Vector Index)
    An extracted, query-optimized representation of a model's
    internal knowledge.  Contains:
    - Gate vectors: per-layer feature activation patterns (f16)
    - Token embeddings: the model's learned word representations
    - Feature metadata: labels mapping features to concepts
    - Tokenizer: the model's subword tokenizer

    A vindex is built offline from model weights by the infer
    extraction pipeline and is read-only at query time.

Gate KNN
    The core query primitive.  Given a text query:
    1. Tokenize and compute an embedding (average of token vectors)
    2. For each layer, find the K features whose gate vectors
       have the highest dot product with the query embedding
    3. Return feature indices and activation scores

    This is analogous to a nearest-neighbor search, but over the
    model's internal feature space rather than an external vector
    database.

Model Handle
    A loaded, ready-to-query model.  Contains the mmap'd vindex,
    tokenizer, embeddings, and configuration.  Cached per-backend
    in a process-local HashMap wrapped in Arc for safe concurrent
    access.


ARCHITECTURE
------------

PostgreSQL Extension Layer (pgrx 0.17)
    lib.rs            Extension entry, schema bootstrap, _PG_init
    gucs.rs           GUC parameters (7 settings)
    error.rs          Unified error type with PG error mapping

Query Functions
    fn_walk.rs        Raw gate activations per layer
    fn_describe.rs    Labeled knowledge edges (deduped, filtered)
    fn_similar.rs     Semantic similarity via shared gate activations
    fn_implies.rs     Directional relationship test via describe
    fn_infer.rs       Forward-pass predictions (feature-gated)

Model Management
    registry.rs       Per-backend model cache (Arc<ModelHandle>)
    model_mgmt.rs     SQL functions for model registration
    options.rs        WITH option parsing and validation

Index Access Method
    am.rs             IndexAmRoutine callbacks
    build.rs          ambuild: reads vindex, writes PG pages
    pages.rs          Page format structs (meta, gate, embed, etc.)
    page_reader.rs    Read-back from PG pages to ModelHandle
    scan.rs           Brute-force heap scan for ORDER BY <~>


DATA FLOW
---------

Registration:
    infer_create_model('name', '/path/to/model.vindex')
        |
        v
    resolve_source() -- validates path is under data_directory
        |
        v
    load_from_path() -- validates vindex can be loaded
        |
        v
    INSERT INTO infer.models -- legacy registry entry
        |
        v
    CREATE INDEX USING infer -- copies vindex data into PG pages

Query (describe/walk/similar_to):
    SQL call
        |
        v
    resolve_model_name() -- explicit param or GUC default
        |
        v
    with_model() -- acquire Arc from cache (mutex held only for lookup)
        |
        v
    embed_text() -- tokenize + average token embeddings
        |
        v
    gate_knn() per layer -- brute-force dot product + top-K
        |
        v
    format results

Index Scan (ORDER BY col <~> 'query'):
    amgettuple()
        |
        v
    compute_distances() -- first call only
        |
        +-- get_heap_attnum() -- direct from rd_index
        +-- use heapRelation from scan descriptor
        +-- embed_text(query) -- pre-compute once
        |
        +-- for each heap tuple:
        |       embed_text(row_value)
        |       similar_to_with_embedding(row, query_embed)
        |       score_to_distance()
        |
        +-- sort by distance ascending
        |
        v
    return TIDs in distance order


STORAGE FORMAT
--------------

Vindex data is stored in standard 8KB PostgreSQL pages, WAL-logged
via GenericXLog.  This provides crash safety and replication.

Page Layout:
    Block 0:     Metapage (model name, dimensions, block ranges)
    Block 1:     Layer directory (start/end blocks per layer)
    Blocks N..M: Gate vectors (f16, packed into pages)
    Blocks M..P: Token embeddings (f32)
    Blocks P..Q: Feature metadata (labels, compressed)
    Blocks Q..R: Tokenizer (JSON blob, chunked across pages)

Each page carries a 16-byte opaque area at the end:
    magic:     u32  (0x494E4652 = "INFR")
    page_type: u8   (Meta/LayerDir/Gate/Embed/DownMeta/Blob)
    flags:     u8
    padding:   [u8; 2]
    next_blk:  u32  (chain pointer for multi-page data)
    reserved:  u32


DISTANCE METRIC
---------------

The <~> operator computes semantic distance between two texts:

    score = similar_to(a, b)
    distance = 1 / (1 + score)

This maps similarity scores to [0, 1) range:
    - score = 0    -> distance = 1.0  (unrelated)
    - score = inf  -> distance = 0.0  (identical)
    - Monotonically decreasing, no discontinuities

The similarity score itself is the maximum shared gate activation:
    1. Embed both texts (average token embeddings)
    2. For each layer, find top-50 activated features for each text
    3. Find features that appear in both sets
    4. Return the minimum activation score of the best shared feature
    5. If no features overlap, fall back to cosine similarity * 10


COST MODEL
----------

The planner cost estimate reflects the actual O(N) scan cost:

    startup_cost = 100.0     (query embedding computation)
    total_cost   = 100 + N * 50  (per-row: tokenize + embed + KNN)
    selectivity  = 1.0       (scans all rows; LIMIT prunes)
    correlation  = 0.0       (output in distance order, not physical)
    index_pages  = heap_pages

The 50x multiplier on cpu_tuple_cost makes the planner prefer
sequential scan for large tables unless ORDER BY + LIMIT makes
the index path worthwhile.


SECURITY
--------

Path restriction: infer_create_model() requires local paths to be
under infer.data_directory (resolved relative to PGDATA).  This
prevents arbitrary filesystem reads via the source parameter.

Reloption validation: amoptions rejects unrecognized WITH options
rather than silently accepting them.


CONCURRENCY
-----------

Each PostgreSQL backend maintains its own model cache in a
process-local HashMap.  The cache is protected by a Mutex, but
the lock is only held during HashMap lookup/insert (microseconds).
The actual model data is wrapped in Arc, so the closure that
performs queries holds an Arc reference without blocking other
cache operations.

The underlying mmap pages are shared by the OS kernel across all
backends that have the same vindex open, so physical memory is
shared even though each backend has its own logical handle.


PERFORMANCE PROFILE (Qwen 0.5B, 24 layers, 4864 features)
----------------------------------------------------------

    Operation              Time        Category
    -------------------    --------    ---------
    tokenization           3-27 us     negligible
    embedding average      0.1-4 us    negligible
    gate KNN (1 layer)     6.4 ms      core cost
    gate KNN (24 layers)   154 ms      bottleneck
    describe()             176 ms      1x gate KNN
    similar_to()           354 ms      2x gate KNN
    scan per row           138 ms      dominates scans
    model loading          2327 ms     one-time

The gate KNN brute-force scan is the dominant cost.  Future
optimization targets:
    - Batched gate KNN (pre-compute query features, 50x potential)
    - Parallel layers via rayon (6-8x potential)
    - SIMD dot products (2-4x potential)
    - Approximate KNN (10-100x potential)


TESTING
-------

13 pgrx integration tests covering:
    - Extension loading and schema creation
    - GUC parameter registration and default values
    - Function signature presence in pg_proc
    - Operator and operator class existence
    - Model management (create/drop/list)
    - New Phase 8 GUCs (describe_top_k, walk_embed_mode)

Note: amoptions reloption rejection cannot be tested in pgrx's
test framework because pgrx::error!() from bare C callbacks
(without pg_guard) causes SIGABRT in the test backend.  This is
verified manually.


FUTURE WORK
-----------

- amcanparallel: parallel-aware index scans
- amvalidate: full operator class signature validation
- Batched gate KNN for 50x scan speedup
- LRU eviction based on infer.max_memory GUC
- HuggingFace download support (hf:// URIs)
- Model weight extraction pipeline integration
