nodedb_sql/planner/
declared_vector_coerce.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Coerce an array literal bound for a declared `VECTOR(dim)` column into an
4//! array of floats.
5//!
6//! A fractional literal resolves to [`SqlValue::Decimal`], and the Origin
7//! planner serializes a non-integer decimal as a msgpack string. Left as is,
8//! `ARRAY[0.1, 0.2]` reaches every engine as an array of strings. The strict
9//! and columnar encoders refuse each string element, and the schemaless
10//! engine stores the strings. Converting each element here gives every
11//! engine and every write path the same array of floats.
12//!
13//! An integer, float, or decimal element becomes an `f64`. Any other element
14//! is refused with [`SqlError::VectorElementNotNumeric`]. A string element is
15//! refused even when its text spells a number, because the client wrote text.
16//! A number outside the `f32` range a vector element holds is refused with
17//! [`SqlError::FloatOutOfRange`].
18//!
19//! A value that is not an array passes through. The engines read a vector
20//! text literal or raw vector bytes by their own rule. The element count is
21//! checked by the engines, which own that error.
22
23use nodedb_types::columnar::ColumnType;
24use rust_decimal::prelude::ToPrimitive;
25
26use crate::error::{Result, SqlError};
27use crate::types::{ColumnInfo, SqlDataType, SqlValue};
28
29/// The declared type name a vector range error reports.
30const VECTOR_TYPE_NAME: &str = "VECTOR";
31
32/// The dimension of `column` when it is declared `VECTOR(dim)`.
33///
34/// A strict or key-value column carries the dimension in its
35/// [`SqlDataType::Vector`]. A schemaless or columnar-family column
36/// advertises a vector as text, so its dimension comes from the declared
37/// type text in [`ColumnInfo::raw_type`].
38pub(crate) fn declared_vector_dim(column: &ColumnInfo) -> Option<usize> {
39    if let SqlDataType::Vector(dim) = column.data_type {
40        return Some(dim);
41    }
42    match column
43        .raw_type
44        .as_deref()
45        .and_then(ColumnType::from_declared_type)
46    {
47        Some(ColumnType::Vector(dim)) => Some(dim as usize),
48        _ => None,
49    }
50}
51
52/// Coerce `value`, bound for the `VECTOR(dim)` column `column`.
53///
54/// An array becomes an array of [`SqlValue::Float`]. Every other value is
55/// returned unchanged.
56pub(crate) fn coerce_to_vector(column: &str, value: SqlValue, dim: usize) -> Result<SqlValue> {
57    let SqlValue::Array(elements) = value else {
58        return Ok(value);
59    };
60    elements
61        .into_iter()
62        .map(|element| vector_element(column, element, dim).map(SqlValue::Float))
63        .collect::<Result<Vec<_>>>()
64        .map(SqlValue::Array)
65}
66
67/// One vector element as the `f64` the column stores.
68fn vector_element(column: &str, element: SqlValue, dim: usize) -> Result<f64> {
69    let number = match element {
70        SqlValue::Int(i) => i as f64,
71        SqlValue::Float(f) => f,
72        SqlValue::Decimal(d) => d.to_f64().ok_or_else(|| SqlError::TypeMismatch {
73            detail: format!("column '{column}': '{d}' is not representable as VECTOR({dim})"),
74        })?,
75        other => {
76            return Err(SqlError::VectorElementNotNumeric {
77                column: column.to_string(),
78                dim,
79                element: element_description(&other),
80            });
81        }
82    };
83    if !number.is_finite() || (number as f32).is_infinite() {
84        return Err(SqlError::FloatOutOfRange {
85            column: column.to_string(),
86            value: number,
87            declared_type: VECTOR_TYPE_NAME,
88        });
89    }
90    Ok(number)
91}
92
93/// How a refused element is named: its kind, and its text for a string.
94fn element_description(element: &SqlValue) -> String {
95    match element {
96        SqlValue::String(s) => format!("text '{s}'"),
97        SqlValue::Null => "null".to_string(),
98        SqlValue::Bool(_) => "a boolean".to_string(),
99        SqlValue::Bytes(_) => "bytes".to_string(),
100        SqlValue::Array(_) => "a nested array".to_string(),
101        SqlValue::Timestamp(_) => "a timestamp".to_string(),
102        SqlValue::Timestamptz(_) => "a timestamptz".to_string(),
103        SqlValue::Int(_) => "an integer".to_string(),
104        SqlValue::Float(_) => "a float".to_string(),
105        SqlValue::Decimal(_) => "a decimal".to_string(),
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112
113    fn decimal(text: &str) -> SqlValue {
114        SqlValue::Decimal(text.parse().expect("decimal literal"))
115    }
116
117    fn column(data_type: SqlDataType, raw_type: Option<&str>) -> ColumnInfo {
118        ColumnInfo {
119            name: "embedding".to_string(),
120            data_type,
121            nullable: true,
122            is_primary_key: false,
123            default: None,
124            raw_type: raw_type.map(str::to_string),
125            int_width: None,
126            float_width: None,
127        }
128    }
129
130    #[test]
131    fn fractional_elements_become_floats() {
132        let value = SqlValue::Array(vec![decimal("0.1"), decimal("0.2"), decimal("0.3")]);
133        assert_eq!(
134            coerce_to_vector("embedding", value, 3).expect("coerces"),
135            SqlValue::Array(vec![
136                SqlValue::Float(0.1),
137                SqlValue::Float(0.2),
138                SqlValue::Float(0.3),
139            ])
140        );
141    }
142
143    #[test]
144    fn mixed_integer_and_fractional_elements_become_floats() {
145        let value = SqlValue::Array(vec![SqlValue::Int(1), decimal("0.5"), SqlValue::Float(2.0)]);
146        assert_eq!(
147            coerce_to_vector("embedding", value, 3).expect("coerces"),
148            SqlValue::Array(vec![
149                SqlValue::Float(1.0),
150                SqlValue::Float(0.5),
151                SqlValue::Float(2.0),
152            ])
153        );
154    }
155
156    #[test]
157    fn a_string_element_is_refused_naming_the_column() {
158        for text in ["abc", "0.1"] {
159            let value = SqlValue::Array(vec![decimal("0.1"), SqlValue::String(text.into())]);
160            let error = coerce_to_vector("embedding", value, 2).expect_err("refused");
161            assert_eq!(
162                error,
163                SqlError::VectorElementNotNumeric {
164                    column: "embedding".into(),
165                    dim: 2,
166                    element: format!("text '{text}'"),
167                }
168            );
169        }
170    }
171
172    #[test]
173    fn null_bool_and_nested_elements_are_refused() {
174        for element in [
175            SqlValue::Null,
176            SqlValue::Bool(true),
177            SqlValue::Array(vec![SqlValue::Int(1)]),
178        ] {
179            let value = SqlValue::Array(vec![element]);
180            assert!(matches!(
181                coerce_to_vector("embedding", value, 1),
182                Err(SqlError::VectorElementNotNumeric { .. })
183            ));
184        }
185    }
186
187    #[test]
188    fn an_element_past_the_f32_range_is_refused() {
189        let value = SqlValue::Array(vec![SqlValue::Float(1e39)]);
190        assert!(matches!(
191            coerce_to_vector("embedding", value, 1),
192            Err(SqlError::FloatOutOfRange { .. })
193        ));
194    }
195
196    #[test]
197    fn a_non_array_value_passes_through() {
198        for value in [
199            SqlValue::Null,
200            SqlValue::String("[0.1,0.2]".into()),
201            SqlValue::Bytes(vec![0; 8]),
202        ] {
203            assert_eq!(
204                coerce_to_vector("embedding", value.clone(), 2).expect("passes"),
205                value
206            );
207        }
208    }
209
210    #[test]
211    fn the_dimension_comes_from_the_type_or_the_declared_text() {
212        assert_eq!(
213            declared_vector_dim(&column(SqlDataType::Vector(3), None)),
214            Some(3)
215        );
216        assert_eq!(
217            declared_vector_dim(&column(SqlDataType::String, Some("VECTOR(4) NOT NULL"))),
218            Some(4)
219        );
220        assert_eq!(
221            declared_vector_dim(&column(SqlDataType::String, Some("TEXT"))),
222            None
223        );
224        assert_eq!(
225            declared_vector_dim(&column(SqlDataType::String, None)),
226            None
227        );
228    }
229}
nodedb_sql/planner/declared_vector_coerce.rs

nodedb_sql/planner/
declared_vector_coerce.rs
