perf(retain): optimize embedding_to_pgvector serialization via zero-copy orjson

[Sanderhoff-alt]

Summary & Key Results

This PR optimizes embedding_to_pgvector—the foundational vector literal serialization operator called on critical Retain, link generation, and import paths when binding embeddings to PostgreSQL vector columns.

By declaring numpy as an explicit core dependency and replacing the Python generator and element-wise repr(float()) formatting with a zero-copy buffer view (np.frombuffer on PackedEmbedding) and Rust Ryu SIMD float serialization (orjson.OPT_SERIALIZE_NUMPY), we achieve:

• ~9.3x – 9.8x CPU speedup on typical Retain batches (e.g. 200 facts @ 1536d drops from 92.6ms $\rightarrow$ 9.45ms; 500 facts from 221.8ms $\rightarrow$ 23.8ms);
• Throughput scales from 3.32 Mfloat/s up to 32.50 Mfloat/s (+879%) on typical 200-fact batches, and up to 18.5 Mfloat/s even on single-vector calls;
• ~44% peak heap memory reduction on 500-fact batches (15.0 MB $\rightarrow$ 8.4 MB);
• Zero intermediate Python PyFloat or PyUnicode object allocations;
• Bit-identical float32 roundtrip fidelity and 100% backward compatibility with all existing types (array('f'), list[float], tuple, np.ndarray, str).

================================================================================
 WORKLOAD SPEEDUP & LATENCY IMPROVEMENT (Apple Silicon / best of 5 runs)
================================================================================
 single_bge_384 (1x 384d)           [3.4x]  ■■■▌            0.136ms -> 0.040ms
 single_openai_1536 (1x 1536d)      [5.8x]  ■■■■■▊          0.489ms -> 0.085ms
 batch_20_gemini_768 (20x 768d)     [8.9x]  ■■■■■■■■▉       4.580ms -> 0.514ms
 batch_200_openai_1536 (200x 1536d) [9.8x]  ■■■■■■■■■▉      92.64ms -> 9.45ms
 batch_500_large_doc (500x 1536d)   [9.3x]  ■■■■■■■■■▍     221.78ms -> 23.84ms
 batch_200_raw_list (200x 1536d)    [7.6x]  ■■■■■■■▋        87.70ms -> 11.48ms
================================================================================

---

1. Problem Analysis & Bottleneck Trace

During Retain (insert_facts_batch, compute_semantic_links_ann), fact embeddings must be formatted as pgvector text literals ("[0.1,0.2,...]") for asyncpg to bind to the PostgreSQL vector column.

Previous Implementation:

def embedding_to_pgvector(embedding: EmbeddingLike) -> str:
    if isinstance(embedding, str):
        return embedding
    return "[" + ",".join(repr(float(value)) for value in embedding) + "]"

Allocation & Latency Bottleneck:

flowchart TD
    subgraph Baseline [Baseline Python Execution: ~220 ms / 500 facts]
        A["500 Facts @ 1536d (768,000 floats)"]
        --> B["768,000 float() unpackings -> 768,000 PyFloat objects"]
        --> C["768,000 repr() calls -> 768,000 PyUnicode string objects"]
        --> D["join() collects 768,000 strings into list -> ~15 MB heap"]
        --> E["Final string concatenation -> Heavy GC churn"]
    end

Loading

For a typical 500-fact document batch at 1536d, this allocated over 1.53 million transient Python heap objects and consumed ~220ms of pure single-core CPU time.

---

2. Technical Architecture & Solution

flowchart LR
    subgraph Input [Polymorphic Input]
        P1["PackedEmbedding\n(array 'f', 4-byte C array)"]
        P2["np.ndarray\n(float32 / float64)"]
        P3["list[float] / tuple\n(Raw API / JSON Import)"]
        P4["str\n(Literal passthrough)"]
    end

    subgraph FastPath [Zero-Copy Rust Ryu SIMD Engine]
        P1 -->|np.frombuffer\nzero-copy view| View["C-contiguous ndarray view"]
        P2 --> View
        View -->|orjson.OPT_SERIALIZE_NUMPY| Ryu["Rust Ryu float formatter\nDirect ASCII byte buffer"]
        P3 -->|orjson.dumps| Ryu
        P4 -->|Identity check| Out["Postgres literal '[0.1,0.2,...]'"]
        Ryu --> Out
    end

    subgraph Fallback [Safety Fallback]
        Ryu -->|Contains NaN / Inf / -Inf| Safe["_repr_literal fallback\n(IEEE 754 compliance)"]
        Safe --> Out
    end

Loading

Key Technical Details & Invariants:

1. Explicit Direct NumPy Dependency: Declares "numpy>=1.26.0" directly in hindsight-api-slim and hindsight-dev (eliminating brittle transitive dependency assumptions and ghost degradation paths).
2. Zero-Copy Memory Protocol: PackedEmbedding (stored as continuous 32-bit floats via array('f') per PR fix(retain): bound retain's memory by a budget instead of by the document (#3756) #3763 / issue Retain holds a whole document in memory: embeddings as list[float] are 74% of peak, and nothing bounds the pipeline by chunk count #3756) is converted to a NumPy buffer view via np.frombuffer(embedding, dtype=np.float32). Zero memory copies or allocations occur.
3. Shortest float32 Representation: The formatted string uses float32 shortest representation when serialized via the PackedEmbedding/NumPy fast-path (e.g. 0.1 vs legacy 0.10000000149011612). When parsed back by PostgreSQL as a 32-bit float vector, the stored bytes are bit-identical.
4. Isolated Safe Fallback (_repr_literal & _dumps_or_repr_fallback): Finite IEEE 754 float JSON formatting strictly consists of digits, decimal points, signs, exponents, commas, and brackets ([0-9.-e,[]]). The substring "null" can only appear if the vector contains NaN, Inf, or -Inf, which orjson encodes as null per JSON specification. When detected (or on type errors), it cleanly falls back to _repr_literal (matching legacy baseline byte-for-byte).
5. Streamlined Dispatcher: The entire embedding_to_pgvector function is a clean polymorphic dispatcher with no defensive dead code.

---

3. Comprehensive Benchmark & Throughput Results

Tested on Apple Silicon using the standardized microbenchmark suite (hindsight-dev/benchmarks/micro/vector_serialization.py, best of 5 repeats):

Workload	Batch & Dim	Floats	Baseline (ms)	Optimized (ms)	Speedup	Baseline Throughput	Optimized Throughput	Peak Heap Memory
single_bge_384	1 item, 384d	384	0.136 ms	0.040 ms	3.4x	2.91 Mfloat/s	8.72 Mfloat/s (+200%)	35.8 KiB $\rightarrow$ 12.6 KiB (-65%)
single_openai_1536	1 item, 1536d	1,536	0.489 ms	0.085 ms	5.8x	3.23 Mfloat/s	18.49 Mfloat/s (+472%)	143.5 KiB $\rightarrow$ 49.8 KiB (-65%)
batch_20_gemini_768	20 items, 768d	15,360	4.580 ms	0.514 ms	8.9x	3.46 Mfloat/s	30.30 Mfloat/s (+776%)	356.2 KiB $\rightarrow$ 184.8 KiB (-48%)
batch_200_openai_1536	200 items, 1536d	307,200	92.64 ms	9.45 ms	9.8x	3.32 Mfloat/s	32.50 Mfloat/s (+879%)	6.1 MiB $\rightarrow$ 3.4 MiB (-44%)
batch_500_large_doc	500 items, 1536d	768,000	221.78 ms	23.84 ms	9.3x	3.42 Mfloat/s	32.21 Mfloat/s (+842%)	15.0 MiB $\rightarrow$ 8.4 MiB (-44%)
batch_200_raw_list	200 items, 1536d (list)	307,200	87.70 ms	11.48 ms	7.6x	3.22 Mfloat/s	25.41 Mfloat/s (+689%)	6.1 MiB $\rightarrow$ 6.0 MiB (-2%)

Throughput Scaling Analysis:

• Single-item calls: Throughput expands from 2.9 ~ 3.2 Mfloat/s up to 8.7 ~ 18.5 Mfloat/s (+200% ~ +472%);
• Multi-item Retain batches (20~500 facts): Throughput expands from 3.2 ~ 3.5 Mfloat/s up to 30.3 ~ 32.5 Mfloat/s (+776% ~ +879%), reaching steady-state peak throughput of 32.50 Mfloat/s on batch_200_openai_1536.

---

4. Conformance & Test Verification

All edge cases and mathematical constraints are validated:

================================================================================
 CONFORMANCE SUITE VERIFICATION
================================================================================
 [PASS] float32_roundtrip_fidelity      Bit-identical roundtrip verified (array('f', parsed) == _VECTOR)
 [PASS] empty_vector_list               [] -> "[]"
 [PASS] empty_vector_packed             array('f', []) -> "[]"
 [PASS] raw_float_list                  list[float] parses to identical float32 vector
 [PASS] tuple_floats                    tuple[float] parses to identical float32 vector
 [PASS] numpy_float32                   np.ndarray(float32) parses to identical float32 vector
 [PASS] numpy_float64                   np.ndarray(float64) parses to identical float32 vector
 [PASS] non_finite_ieee754              [nan, inf, -inf, 1.0] correctly falls back to "[nan,inf,-inf,1.0]"
 [PASS] subnormal_small_floats          1e-35 preserved without precision loss
 [PASS] string_literal_passthrough      "[1.0, 2.0]" returned idempotently
 [PASS] custom_scalar_fallback          Custom object with __float__() falls back cleanly
 [PASS] non_f_array_fallthrough         array('d', ...) falls through to _repr_literal
================================================================================

Unit tests in hindsight-api-slim/tests/test_packed_embeddings.py (7/7 passed).

---

5. Changes Summary

• hindsight-api-slim/pyproject.toml:
  • Added "numpy>=1.26.0" to core dependencies.
• hindsight-api-slim/hindsight_api/engine/retain/types.py:
  • Direct top-level numpy/orjson imports and streamlined embedding_to_pgvector dispatcher with _dumps_or_repr_fallback and isolated _repr_literal.
• hindsight-api-slim/hindsight_api/engine/retain/link_utils.py:
  • Moved import numpy as np to top-level module import.
• hindsight-api-slim/tests/test_packed_embeddings.py:
  • Added semantic float32 parsed array assertions across all types, custom scalar objects, non-f array fallthrough, subnormals, and non-finites.
• hindsight-dev/pyproject.toml:
  • Added "numpy>=1.26.0" dependency and registered vector-serialization-bench CLI command.
• hindsight-dev/benchmarks/micro/vector_serialization.py:
  • Added standardized microbenchmark suite with 6 production workloads, timing metrics, throughput tracking, and rich formatting (following PR perf(tokenizer): replace tiktoken with quicktok and default to o200k_base #3788 pattern).
• scripts/benchmarks/run-vector-serialization-bench.sh:
  • Added runner script executable from repository root.

[strix-security]

Strix Security Review

Warning

This pull request has 1 commit after the last Strix review (e632cf9). Strix has not reviewed these changes.
Automatic review on push is off for this repository. To review the latest changes, tag @strix-security in a comment, or turn on re-review on push.

No security issues found.

_{Updated for e632cf9.}

---

Reviewed by Strix
_{Re-run review · Configure security review settings}