Skip to main content

qql_embed/
resolve.rs

1use qql_core::ast::{
2    EmbedKind, EmbeddingSpec, PointEntry, PointVectors, Stmt, UpsertStmt, VectorValue,
3};
4use qql_core::error::QqlError;
5
6use crate::embedder::Embedder;
7
8pub(crate) use super::resolve_query::ensure_batch_len;
9use super::resolve_query::resolve_query_embeddings;
10pub use super::resolve_query::{DENSE_VECTOR_NAME, SPARSE_VECTOR_NAME};
11
12/// Resolve text → vectors on a statement before routing/execution.
13///
14/// Every modality batches alike: dense, query-side sparse, multi, and image
15/// jobs are each collected in walk order and sent through the matching
16/// `*_batch` entry point grouped by model (one RPC per model). Sparse stays
17/// role-split via the embedder: queries embed with unit weights, documents
18/// with BM25 tf saturation (both wire-compatible with Qdrant's
19/// `qdrant/bm25`).
20pub async fn resolve_embeddings(stmt: &mut Stmt, embedder: &dyn Embedder) -> Result<(), QqlError> {
21    match stmt {
22        Stmt::Query(query) => resolve_query_embeddings(query, embedder).await?,
23        Stmt::Upsert(upsert) => resolve_upsert_embeddings(upsert, embedder).await?,
24        _ => {}
25    }
26    Ok(())
27}
28
29async fn resolve_upsert_embeddings(
30    upsert: &mut UpsertStmt,
31    embedder: &dyn Embedder,
32) -> Result<(), QqlError> {
33    if upsert.embedding.is_none() && upsert.embed.is_empty() {
34        let mut targets = Vec::new();
35        for (idx, point) in upsert.points.iter().enumerate() {
36            let PointEntry::Inline(inline) = point else {
37                continue;
38            };
39            if inline.vectors.is_none()
40                && let Some((_, qql_core::ast::Value::Str(text))) =
41                    inline.payload.iter().find(|(k, _)| {
42                        eq_lowered(k, "text") || eq_lowered(k, "body") || eq_lowered(k, "content")
43                    })
44                && !text.is_empty()
45            {
46                targets.push((idx, text.clone()));
47            }
48        }
49        if !targets.is_empty() {
50            // Topology-unaware fallback: dense only. Hybrid/sparse targets must
51            // be set by the executor (configure_upsert_embeddings) or explicit
52            // USING / EMBED directives before calling resolve_embeddings — so
53            // dense-only collections never receive orphan sparse vectors.
54            let (indices, texts): (Vec<usize>, Vec<String>) = targets.into_iter().unzip();
55            let dense_vecs = embedder.embed_dense_batch(&texts, "default").await?;
56            ensure_batch_len(dense_vecs.len(), indices.len(), "default")?;
57            for (idx, d_vec) in indices.into_iter().zip(dense_vecs) {
58                let point = &mut upsert.points[idx];
59                add_point_vector(point, DENSE_VECTOR_NAME, VectorValue::Dense(d_vec))?;
60            }
61        }
62    }
63
64    if let Some(spec) = upsert.embedding.clone() {
65        // Duplicate-target tracking stays a small `Vec`: a handful of names
66        // at most, so linear scan beats a per-op `HashSet` (no hashing).
67        let mut seen_vectors = Vec::new();
68        resolve_single_embedding_spec(upsert, &spec, embedder, &mut seen_vectors).await?;
69    }
70
71    for directive in &upsert.embed {
72        let field_name = &directive.source_field;
73        let target_vec_name = &directive.target_vector;
74        let mut targets = Vec::new();
75        for (idx, point) in upsert.points.iter().enumerate() {
76            let PointEntry::Inline(inline) = point else {
77                continue;
78            };
79            if let Some((_, qql_core::ast::Value::Str(text))) = inline
80                .payload
81                .iter()
82                .find(|(k, _)| k.eq_ignore_ascii_case(field_name))
83                && !text.is_empty()
84            {
85                targets.push((idx, text.clone()));
86            }
87        }
88
89        if !targets.is_empty() {
90            match &directive.kind {
91                EmbedKind::Dense { model } => {
92                    let m_name = model.as_deref().unwrap_or("default");
93                    let (indices, texts): (Vec<usize>, Vec<String>) = targets.into_iter().unzip();
94                    let vecs = embedder.embed_dense_batch(&texts, m_name).await?;
95                    ensure_batch_len(vecs.len(), indices.len(), m_name)?;
96                    for (idx, vec) in indices.into_iter().zip(vecs) {
97                        let point = &mut upsert.points[idx];
98                        add_point_vector(point, target_vec_name, VectorValue::Dense(vec))?;
99                    }
100                }
101                EmbedKind::Sparse { model } => {
102                    let m = model.as_deref().unwrap_or("default");
103                    let (indices, texts): (Vec<usize>, Vec<String>) = targets.into_iter().unzip();
104                    let vecs = embedder.embed_sparse_document_batch(&texts, m).await?;
105                    ensure_batch_len(vecs.len(), indices.len(), m)?;
106                    for (idx, s_vec) in indices.into_iter().zip(vecs) {
107                        let point = &mut upsert.points[idx];
108                        add_point_vector(
109                            point,
110                            target_vec_name,
111                            VectorValue::Sparse {
112                                indices: s_vec.indices,
113                                values: s_vec.values,
114                            },
115                        )?;
116                    }
117                }
118                EmbedKind::Multi { model } => {
119                    let m_name = model.as_deref().unwrap_or("default");
120                    let (indices, texts): (Vec<usize>, Vec<String>) = targets.into_iter().unzip();
121                    let bags = embedder.embed_multi_batch(&texts, m_name).await?;
122                    if bags.len() != indices.len() {
123                        return Err(QqlError::execution(
124                            "QQL-EMBEDDING-MULTI",
125                            format!(
126                                "embed_multi_batch returned {} bags for {} texts (model={m_name})",
127                                bags.len(),
128                                indices.len()
129                            ),
130                            None,
131                        ));
132                    }
133                    for (idx, rows) in indices.into_iter().zip(bags) {
134                        if rows.is_empty() {
135                            return Err(QqlError::execution(
136                                "QQL-EMBEDDING-MULTI",
137                                "embed_multi returned an empty multivector",
138                                None,
139                            ));
140                        }
141                        let point = &mut upsert.points[idx];
142                        add_point_vector(point, target_vec_name, VectorValue::MultiDense(rows))?;
143                    }
144                }
145                EmbedKind::Image { model } => {
146                    let m_name = model.as_deref().unwrap_or("default");
147                    let (indices, sources): (Vec<usize>, Vec<String>) = targets.into_iter().unzip();
148                    let vecs = embedder.embed_image_batch(&sources, m_name).await?;
149                    if vecs.len() != indices.len() {
150                        return Err(QqlError::execution(
151                            "QQL-EMBEDDING-IMAGE",
152                            format!(
153                                "embed_image_batch returned {} vectors for {} sources (model={m_name})",
154                                vecs.len(),
155                                indices.len()
156                            ),
157                            None,
158                        ));
159                    }
160                    for (idx, vec) in indices.into_iter().zip(vecs) {
161                        let point = &mut upsert.points[idx];
162                        add_point_vector(point, target_vec_name, VectorValue::Dense(vec))?;
163                    }
164                }
165            }
166        }
167    }
168
169    Ok(())
170}
171
172async fn resolve_single_embedding_spec(
173    upsert: &mut UpsertStmt,
174    spec: &EmbeddingSpec,
175    embedder: &dyn Embedder,
176    seen_vectors: &mut Vec<String>,
177) -> Result<(), QqlError> {
178    match spec {
179        EmbeddingSpec::Multi(specs) => {
180            for sub_spec in specs {
181                Box::pin(resolve_single_embedding_spec(
182                    upsert,
183                    sub_spec,
184                    embedder,
185                    seen_vectors,
186                ))
187                .await?;
188            }
189        }
190        EmbeddingSpec::Dense {
191            model,
192            vector,
193            field,
194        } => {
195            let model_name = model.as_deref().unwrap_or("default");
196            let vector_name = vector.as_deref().unwrap_or(DENSE_VECTOR_NAME);
197            check_and_insert_vector_name(seen_vectors, vector_name)?;
198
199            let targets = collect_text_targets(&upsert.points, field.as_deref());
200            validate_non_empty_targets(upsert, &targets, "DENSE", field.as_deref())?;
201
202            let (indices, texts): (Vec<usize>, Vec<String>) = targets.into_iter().unzip();
203            let vecs = embedder.embed_dense_batch(&texts, model_name).await?;
204            ensure_batch_len(vecs.len(), indices.len(), model_name)?;
205            for (idx, vec) in indices.into_iter().zip(vecs) {
206                let point = &mut upsert.points[idx];
207                add_point_vector(point, vector_name, VectorValue::Dense(vec))?;
208            }
209        }
210        EmbeddingSpec::Sparse {
211            model,
212            vector,
213            field,
214        } => {
215            let model_name = model.as_deref().unwrap_or("default");
216            let vector_name = vector.as_deref().unwrap_or(SPARSE_VECTOR_NAME);
217            check_and_insert_vector_name(seen_vectors, vector_name)?;
218
219            let targets = collect_text_targets(&upsert.points, field.as_deref());
220            validate_non_empty_targets(upsert, &targets, "SPARSE", field.as_deref())?;
221
222            let (indices, texts): (Vec<usize>, Vec<String>) = targets.into_iter().unzip();
223            let vecs = embedder
224                .embed_sparse_document_batch(&texts, model_name)
225                .await?;
226            ensure_batch_len(vecs.len(), indices.len(), model_name)?;
227            for (idx, sparse_vec) in indices.into_iter().zip(vecs) {
228                add_point_vector(
229                    &mut upsert.points[idx],
230                    vector_name,
231                    VectorValue::Sparse {
232                        indices: sparse_vec.indices,
233                        values: sparse_vec.values,
234                    },
235                )?;
236            }
237        }
238        EmbeddingSpec::Hybrid {
239            dense_model,
240            dense_vector,
241            dense_field,
242            sparse_model,
243            sparse_vector,
244            sparse_field,
245        } => {
246            let d_model = dense_model.as_deref().unwrap_or("default");
247            let s_model = sparse_model.as_deref().unwrap_or("default");
248            let d_vec_name = dense_vector.as_deref().unwrap_or(DENSE_VECTOR_NAME);
249            let s_vec_name = sparse_vector.as_deref().unwrap_or(SPARSE_VECTOR_NAME);
250
251            check_and_insert_vector_name(seen_vectors, d_vec_name)?;
252            check_and_insert_vector_name(seen_vectors, s_vec_name)?;
253
254            let dense_targets = collect_text_targets(&upsert.points, dense_field.as_deref());
255            let sparse_targets = collect_text_targets(&upsert.points, sparse_field.as_deref());
256
257            validate_non_empty_targets(upsert, &dense_targets, "DENSE", dense_field.as_deref())?;
258            validate_non_empty_targets(upsert, &sparse_targets, "SPARSE", sparse_field.as_deref())?;
259
260            let (indices, texts): (Vec<usize>, Vec<String>) = dense_targets.into_iter().unzip();
261            let dense_vecs = embedder.embed_dense_batch(&texts, d_model).await?;
262            ensure_batch_len(dense_vecs.len(), indices.len(), d_model)?;
263            for (idx, d_vec) in indices.into_iter().zip(dense_vecs) {
264                let point = &mut upsert.points[idx];
265                add_point_vector(point, d_vec_name, VectorValue::Dense(d_vec))?;
266            }
267
268            let (sparse_indices, sparse_texts): (Vec<usize>, Vec<String>) =
269                sparse_targets.into_iter().unzip();
270            let sparse_vecs = embedder
271                .embed_sparse_document_batch(&sparse_texts, s_model)
272                .await?;
273            ensure_batch_len(sparse_vecs.len(), sparse_indices.len(), s_model)?;
274            for (idx, sparse_vec) in sparse_indices.into_iter().zip(sparse_vecs) {
275                let point = &mut upsert.points[idx];
276                add_point_vector(
277                    point,
278                    s_vec_name,
279                    VectorValue::Sparse {
280                        indices: sparse_vec.indices,
281                        values: sparse_vec.values,
282                    },
283                )?;
284            }
285        }
286        EmbeddingSpec::MultiVector {
287            model,
288            vector,
289            field,
290        } => {
291            let model_name = model.as_deref().unwrap_or("default");
292            let vector_name = vector.as_deref().unwrap_or("colbert");
293            check_and_insert_vector_name(seen_vectors, vector_name)?;
294
295            let targets = collect_text_targets(&upsert.points, field.as_deref());
296            validate_non_empty_targets(upsert, &targets, "MULTI", field.as_deref())?;
297
298            let (indices, texts): (Vec<usize>, Vec<String>) = targets.into_iter().unzip();
299            let bags = embedder.embed_multi_batch(&texts, model_name).await?;
300            if bags.len() != indices.len() {
301                return Err(QqlError::execution(
302                    "QQL-EMBEDDING-MULTI",
303                    format!(
304                        "embed_multi_batch returned {} bags for {} texts (model={model_name})",
305                        bags.len(),
306                        indices.len()
307                    ),
308                    None,
309                ));
310            }
311            for (idx, rows) in indices.into_iter().zip(bags) {
312                if rows.is_empty() {
313                    return Err(QqlError::execution(
314                        "QQL-EMBEDDING-MULTI",
315                        "embed_multi returned an empty multivector",
316                        None,
317                    ));
318                }
319                add_point_vector(
320                    &mut upsert.points[idx],
321                    vector_name,
322                    VectorValue::MultiDense(rows),
323                )?;
324            }
325        }
326        EmbeddingSpec::Image {
327            model,
328            vector,
329            field,
330        } => {
331            let model_name = model.as_deref().unwrap_or("default");
332            let vector_name = vector.as_deref().unwrap_or("image");
333            check_and_insert_vector_name(seen_vectors, vector_name)?;
334
335            let targets = collect_image_targets(&upsert.points, field.as_deref());
336            validate_non_empty_targets(upsert, &targets, "IMAGE", field.as_deref())?;
337
338            let (indices, sources): (Vec<usize>, Vec<String>) = targets.into_iter().unzip();
339            let vecs = embedder.embed_image_batch(&sources, model_name).await?;
340            if vecs.len() != indices.len() {
341                return Err(QqlError::execution(
342                    "QQL-EMBEDDING-IMAGE",
343                    format!(
344                        "embed_image_batch returned {} vectors for {} sources (model={model_name})",
345                        vecs.len(),
346                        indices.len()
347                    ),
348                    None,
349                ));
350            }
351            for (idx, vec) in indices.into_iter().zip(vecs) {
352                add_point_vector(
353                    &mut upsert.points[idx],
354                    vector_name,
355                    VectorValue::Dense(vec),
356                )?;
357            }
358        }
359    }
360    Ok(())
361}
362
363fn validate_non_empty_targets(
364    upsert: &UpsertStmt,
365    targets: &[(usize, String)],
366    kind: &str,
367    field: Option<&str>,
368) -> Result<(), QqlError> {
369    if targets.is_empty() {
370        let actual_fields = upsert
371            .points
372            .first()
373            .and_then(|p| match p {
374                PointEntry::Inline(inline) => Some(inline),
375                PointEntry::Param(..) | PointEntry::PositionalParam(..) => None,
376            })
377            .map(|p| {
378                p.payload
379                    .iter()
380                    .map(|(k, _)| k.as_str())
381                    .collect::<Vec<_>>()
382                    .join(", ")
383            })
384            .unwrap_or_default();
385
386        let err_msg = if let Some(f) = field {
387            format!(
388                "USING {kind} MODEL specified with ON FIELD '{f}' but no matching text payload field found. Found fields: {actual_fields}"
389            )
390        } else {
391            format!(
392                "USING {kind} MODEL specified but no text payload field found. Expected one of: {}. Found fields: {actual_fields}",
393                DEFAULT_TEXT_FIELDS_ORDERED.join(", ")
394            )
395        };
396
397        return Err(QqlError::execution("QQL-EMBEDDING", err_msg, None));
398    }
399    Ok(())
400}
401
402fn check_and_insert_vector_name(
403    seen_vectors: &mut Vec<String>,
404    vector_name: &str,
405) -> Result<(), QqlError> {
406    if seen_vectors.iter().any(|name| name == vector_name) {
407        return Err(QqlError::execution(
408            "QQL-EMBEDDING",
409            format!("duplicate target vector '{vector_name}' in multi-spec embedding clause"),
410            None,
411        ));
412    }
413    seen_vectors.push(vector_name.to_string());
414    Ok(())
415}
416
417const DEFAULT_TEXT_FIELDS_ORDERED: &[&str] = &[
418    "text",
419    "body",
420    "content",
421    "title",
422    "description",
423    "name",
424    "summary",
425    "document",
426];
427
428const DEFAULT_IMAGE_FIELDS_ORDERED: &[&str] = &[
429    "image",
430    "image_path",
431    "image_url",
432    "photo",
433    "picture",
434    "img",
435    "path",
436    "url",
437];
438
439/// ASCII case-insensitive equality against an already-lowercased needle.
440///
441/// Callers lowercase the needle once outside the per-point loop, so each
442/// comparison folds only the haystack side. The exact-match fast path covers
443/// the common already-lowercase payload keys (`text`, `title`, …) with a
444/// single `memcmp`.
445fn eq_lowered(haystack: &str, needle_lower: &str) -> bool {
446    if haystack == needle_lower {
447        return true;
448    }
449    // ASCII lowercasing never changes byte length, so the length check is
450    // exact and short-circuits mismatched keys before folding.
451    haystack.len() == needle_lower.len()
452        && haystack
453            .as_bytes()
454            .iter()
455            .zip(needle_lower.as_bytes())
456            .all(|(a, b)| a.to_ascii_lowercase() == *b)
457}
458
459/// Collect image path/URL payload fields for IMAGE embedding specs.
460fn collect_image_targets(
461    points: &[PointEntry],
462    field_override: Option<&str>,
463) -> Vec<(usize, String)> {
464    if let Some(target_field) = field_override {
465        let target_lower = target_field.to_ascii_lowercase();
466        points
467            .iter()
468            .enumerate()
469            .filter_map(|(idx, point)| {
470                let PointEntry::Inline(inline) = point else {
471                    return None;
472                };
473                inline.payload.iter().find_map(|(key, value)| {
474                    if eq_lowered(key, &target_lower)
475                        && let qql_core::ast::Value::Str(source) = value
476                        && !source.is_empty()
477                    {
478                        return Some((idx, source.clone()));
479                    }
480                    None
481                })
482            })
483            .collect()
484    } else {
485        points
486            .iter()
487            .enumerate()
488            .filter_map(|(idx, point)| {
489                let PointEntry::Inline(inline) = point else {
490                    return None;
491                };
492                for &candidate in DEFAULT_IMAGE_FIELDS_ORDERED {
493                    if let Some((_, qql_core::ast::Value::Str(source))) = inline
494                        .payload
495                        .iter()
496                        .find(|(key, _)| eq_lowered(key, candidate))
497                        && !source.is_empty()
498                    {
499                        return Some((idx, source.clone()));
500                    }
501                }
502                None
503            })
504            .collect()
505    }
506}
507
508fn collect_text_targets(
509    points: &[PointEntry],
510    field_override: Option<&str>,
511) -> Vec<(usize, String)> {
512    if let Some(target_field) = field_override {
513        let target_lower = target_field.to_ascii_lowercase();
514        points
515            .iter()
516            .enumerate()
517            .filter_map(|(idx, point)| {
518                let PointEntry::Inline(inline) = point else {
519                    return None;
520                };
521                inline.payload.iter().find_map(|(key, value)| {
522                    if eq_lowered(key, &target_lower)
523                        && let qql_core::ast::Value::Str(text) = value
524                        && !text.is_empty()
525                    {
526                        return Some((idx, text.clone()));
527                    }
528                    None
529                })
530            })
531            .collect()
532    } else {
533        collect_default_text_targets(points)
534    }
535}
536
537fn collect_default_text_targets(points: &[PointEntry]) -> Vec<(usize, String)> {
538    points
539        .iter()
540        .enumerate()
541        .filter_map(|(idx, point)| {
542            let PointEntry::Inline(inline) = point else {
543                return None;
544            };
545            for &candidate in DEFAULT_TEXT_FIELDS_ORDERED {
546                if let Some((_, qql_core::ast::Value::Str(text))) = inline
547                    .payload
548                    .iter()
549                    .find(|(key, _)| eq_lowered(key, candidate))
550                    && !text.is_empty()
551                {
552                    return Some((idx, text.clone()));
553                }
554            }
555            None
556        })
557        .collect()
558}
559
560fn add_point_vector(
561    point: &mut PointEntry,
562    name: &str,
563    vector: VectorValue,
564) -> Result<(), QqlError> {
565    let point = match point {
566        PointEntry::Inline(inline) => inline,
567        // Unreachable via collect_*_targets (they skip placeholders), but
568        // embedding into unknown payload must never silently drop vectors.
569        PointEntry::Param(name, span) => {
570            return Err(QqlError::execution(
571                "QQL-EMBEDDING",
572                format!("cannot embed into unbound point parameter ':{name}'"),
573                span.as_deref().copied(),
574            ));
575        }
576        PointEntry::PositionalParam(idx, span) => {
577            return Err(QqlError::execution(
578                "QQL-EMBEDDING",
579                format!("cannot embed into unbound point parameter '?{}'", *idx + 1),
580                span.as_deref().copied(),
581            ));
582        }
583    };
584    if name.is_empty() {
585        return match &mut point.vectors {
586            Some(PointVectors::Unnamed(existing)) => {
587                *existing = vector;
588                Ok(())
589            }
590            Some(PointVectors::Named(list)) => {
591                if let Some(existing) = list.iter_mut().find(|(key, _)| key.is_empty()) {
592                    existing.1 = vector;
593                } else {
594                    list.push((String::new(), vector));
595                }
596                Ok(())
597            }
598            Some(PointVectors::Param(..)) | Some(PointVectors::PositionalParam(..)) => {
599                point.vectors = Some(PointVectors::Unnamed(vector));
600                Ok(())
601            }
602            None => {
603                point.vectors = Some(PointVectors::Unnamed(vector));
604                Ok(())
605            }
606        };
607    }
608    match &mut point.vectors {
609        Some(PointVectors::Named(list)) => {
610            if let Some(existing) = list.iter_mut().find(|(k, _)| k == name) {
611                existing.1 = vector;
612            } else {
613                list.push((name.to_string(), vector));
614            }
615            Ok(())
616        }
617        Some(PointVectors::Unnamed(_))
618        | Some(PointVectors::Param(..))
619        | Some(PointVectors::PositionalParam(..)) => Err(QqlError::execution(
620            "QQL-EMBEDDING",
621            format!(
622                "cannot add named vector '{name}' to a point that already has an unnamed vector; \
623                 provide an explicit named-vector topology or omit EMBED for this point"
624            ),
625            None,
626        )),
627        None => {
628            point.vectors = Some(PointVectors::Named(vec![(name.to_string(), vector)]));
629            Ok(())
630        }
631    }
632}