Skip to main content

spg_sqlx/types/
vector.rs

1//! v7.17.0 Phase 3.P0-68 — sqlx bridge for pgvector `VECTOR(N)`
2//! and PG `TSVECTOR`.
3//!
4//! `Vec<f32>` bridges the dense pgvector surface (any storage
5//! encoding: default f32 / `USING SQ8` / `USING HALF` — quantised
6//! variants dequantise to f32 at the adapter boundary).
7//!
8//! `String` Decode on TsVector / Vector cells falls through to
9//! the canonical PG / pgvector external form, mirroring what
10//! pgwire ships on the wire. The text path in `types/text.rs`
11//! calls into the `try_*_as_string` helpers below.
12//!
13//! TsVector Encode is intentionally not implemented: clients
14//! build a tsvector via `to_tsvector(...)` in SQL, not by
15//! binding a raw lexeme list.
16
17use sqlx_core::decode::Decode;
18use sqlx_core::encode::{Encode, IsNull};
19use sqlx_core::error::BoxDynError;
20use sqlx_core::types::Type;
21
22use spg_embedded::Value as EngineValue;
23
24use crate::arguments::SpgArgumentValue;
25use crate::database::Spg;
26use crate::type_info::{Kind, SpgTypeInfo};
27use crate::value::SpgValueRef;
28
29// ---- text fallthrough helpers --------------------------------
30
31/// Try to render a Vector cell into pgvector's canonical
32/// external form (`[1, 2.5, -3]`). Returns `None` for non-
33/// vector cells so the caller can keep trying.
34pub(crate) fn try_vector_as_string(value: &EngineValue) -> Option<String> {
35    let v = match value {
36        EngineValue::Vector(v) => v.clone(),
37        EngineValue::Sq8Vector(q) => spg_storage::quantize::dequantize(q),
38        EngineValue::HalfVector(h) => h.to_f32_vec(),
39        _ => return None,
40    };
41    Some(format_vector(&v))
42}
43
44/// Try to render a TsVector cell into PG's canonical external
45/// form. Returns `None` for non-tsvector cells.
46pub(crate) fn try_tsvector_as_string(value: &EngineValue) -> Option<String> {
47    match value {
48        EngineValue::TsVector(lex) => Some(format_tsvector(lex)),
49        _ => None,
50    }
51}
52
53/// pgvector canonical external form. Matches what the engine's
54/// pgwire path sends so the adapter and wire agree.
55fn format_vector(v: &[f32]) -> String {
56    let mut out = String::with_capacity(v.len() * 6);
57    out.push('[');
58    for (i, x) in v.iter().enumerate() {
59        if i > 0 {
60            out.push_str(", ");
61        }
62        format_f32_into(&mut out, *x);
63    }
64    out.push(']');
65    out
66}
67
68fn format_f32_into(out: &mut String, x: f32) {
69    use core::fmt::Write;
70    // pgvector renders compact decimals: `1`, `2.5`, `-3` —
71    // integer values drop the trailing `.0`, fractions print
72    // their shortest round-trip form. f32 Display already does
73    // that on stable Rust.
74    let _ = write!(out, "{x}");
75}
76
77/// PG `tsvector` external form: `'word1':1 'word2':2,3A`.
78/// Mirrors `spg_engine::eval::format_tsvector` exactly so the
79/// adapter and pgwire agree byte-for-byte.
80fn format_tsvector(lexs: &[spg_storage::TsLexeme]) -> String {
81    let mut out = String::with_capacity(lexs.len() * 12);
82    for (i, l) in lexs.iter().enumerate() {
83        if i > 0 {
84            out.push(' ');
85        }
86        out.push('\'');
87        for c in l.word.chars() {
88            if c == '\'' {
89                out.push('\'');
90            }
91            out.push(c);
92        }
93        out.push('\'');
94        if !l.positions.is_empty() {
95            for (pi, p) in l.positions.iter().enumerate() {
96                out.push(if pi == 0 { ':' } else { ',' });
97                let _ = core::fmt::Write::write_fmt(&mut out, format_args!("{p}"));
98            }
99            match l.weight {
100                3 => out.push('A'),
101                2 => out.push('B'),
102                1 => out.push('C'),
103                _ => {}
104            }
105        }
106    }
107    out
108}
109
110// ---- Vec<f32> bridge -----------------------------------------
111
112impl Type<Spg> for Vec<f32> {
113    fn type_info() -> SpgTypeInfo {
114        SpgTypeInfo::of(Kind::Vector)
115    }
116
117    fn compatible(ty: &SpgTypeInfo) -> bool {
118        matches!(ty.kind(), Kind::Vector)
119    }
120}
121
122impl<'q> Encode<'q, Spg> for Vec<f32> {
123    fn encode_by_ref(&self, buf: &mut Vec<SpgArgumentValue<'q>>) -> Result<IsNull, BoxDynError> {
124        buf.push(SpgArgumentValue {
125            value: EngineValue::Vector(self.clone()),
126            type_info: Some(<Vec<f32> as Type<Spg>>::type_info()),
127            _phantom: core::marker::PhantomData,
128        });
129        Ok(IsNull::No)
130    }
131}
132
133impl<'r> Decode<'r, Spg> for Vec<f32> {
134    fn decode(value: SpgValueRef<'r>) -> Result<Self, BoxDynError> {
135        match value.engine() {
136            EngineValue::Vector(v) => Ok(v.clone()),
137            EngineValue::Sq8Vector(q) => Ok(spg_storage::quantize::dequantize(q)),
138            EngineValue::HalfVector(h) => Ok(h.to_f32_vec()),
139            other => Err(format!("cannot decode {other:?} as Vec<f32> / VECTOR").into()),
140        }
141    }
142}