1mod vector_api;
4mod vector_codec;
5
6use crate::error::{Error, Result};
7use crate::types::DataType;
8use im::OrdMap;
9use serde::{Deserialize, Serialize};
10use serde_json::{Number, Value};
11use std::collections::BTreeMap;
12use std::sync::Arc;
13
14use vector_codec::validate_vector;
15
16pub(crate) use vector_codec::{f32_to_fp16, fp16_to_f32};
17
18#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
20#[serde(tag = "type", content = "value")]
21pub enum FieldValue {
22 Null,
23 Binary(Vec<u8>),
24 String(String),
25 Bool(bool),
26 Int32(i32),
27 Int64(i64),
28 Uint32(u32),
29 Uint64(u64),
30 Float(f32),
31 Double(f64),
32 ArrayBinary(Vec<Vec<u8>>),
33 ArrayString(Vec<String>),
34 ArrayBool(Vec<bool>),
35 ArrayInt32(Vec<i32>),
36 ArrayInt64(Vec<i64>),
37 ArrayUint32(Vec<u32>),
38 ArrayUint64(Vec<u64>),
39 ArrayFloat(Vec<f32>),
40 ArrayDouble(Vec<f64>),
41 Json(Value),
45}
46
47impl FieldValue {
48 pub fn data_type(&self) -> DataType {
49 match self {
50 Self::Null | Self::Json(_) => DataType::Undefined,
51 Self::Binary(_) => DataType::Binary,
52 Self::String(_) => DataType::String,
53 Self::Bool(_) => DataType::Bool,
54 Self::Int32(_) => DataType::Int32,
55 Self::Int64(_) => DataType::Int64,
56 Self::Uint32(_) => DataType::Uint32,
57 Self::Uint64(_) => DataType::Uint64,
58 Self::Float(_) => DataType::Float,
59 Self::Double(_) => DataType::Double,
60 Self::ArrayBinary(_) => DataType::ArrayBinary,
61 Self::ArrayString(_) => DataType::ArrayString,
62 Self::ArrayBool(_) => DataType::ArrayBool,
63 Self::ArrayInt32(_) => DataType::ArrayInt32,
64 Self::ArrayInt64(_) => DataType::ArrayInt64,
65 Self::ArrayUint32(_) => DataType::ArrayUint32,
66 Self::ArrayUint64(_) => DataType::ArrayUint64,
67 Self::ArrayFloat(_) => DataType::ArrayFloat,
68 Self::ArrayDouble(_) => DataType::ArrayDouble,
69 }
70 }
71
72 pub(crate) fn to_json(&self) -> Value {
73 match self {
74 Self::Null => Value::Null,
75 Self::Binary(bytes) => Value::String(base64_encode(bytes)),
76 Self::String(value) => Value::String(value.clone()),
77 Self::Bool(value) => Value::Bool(*value),
78 Self::Int32(value) => Value::Number((*value).into()),
79 Self::Int64(value) => Value::Number((*value).into()),
80 Self::Uint32(value) => Value::Number((*value).into()),
81 Self::Uint64(value) => Value::Number((*value).into()),
82 Self::Float(value) => number_from_f64(f64::from(*value)),
83 Self::Double(value) => number_from_f64(*value),
84 Self::ArrayBinary(values) => Value::Array(
85 values
86 .iter()
87 .map(|v| Value::String(base64_encode(v)))
88 .collect(),
89 ),
90 Self::ArrayString(values) => {
91 Value::Array(values.iter().cloned().map(Value::String).collect())
92 }
93 Self::ArrayBool(values) => {
94 Value::Array(values.iter().copied().map(Value::Bool).collect())
95 }
96 Self::ArrayInt32(values) => {
97 Value::Array(values.iter().map(|v| Value::Number((*v).into())).collect())
98 }
99 Self::ArrayInt64(values) => {
100 Value::Array(values.iter().map(|v| Value::Number((*v).into())).collect())
101 }
102 Self::ArrayUint32(values) => {
103 Value::Array(values.iter().map(|v| Value::Number((*v).into())).collect())
104 }
105 Self::ArrayUint64(values) => {
106 Value::Array(values.iter().map(|v| Value::Number((*v).into())).collect())
107 }
108 Self::ArrayFloat(values) => Value::Array(
109 values
110 .iter()
111 .map(|v| number_from_f64(f64::from(*v)))
112 .collect(),
113 ),
114 Self::ArrayDouble(values) => {
115 Value::Array(values.iter().map(|v| number_from_f64(*v)).collect())
116 }
117 Self::Json(value) => value.clone(),
118 }
119 }
120}
121
122#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
128#[serde(tag = "type", content = "value")]
129pub enum VectorValue {
130 Binary32(Vec<u8>),
131 Binary64(Vec<u8>),
132 Fp16(Vec<u16>),
133 Fp32(Vec<f32>),
134 Fp64(Vec<f64>),
135 Int4(Vec<i8>),
136 Int8(Vec<i8>),
137 Int16(Vec<i16>),
138 SparseFp16 { indices: Vec<u32>, values: Vec<u16> },
139 SparseFp32 { indices: Vec<u32>, values: Vec<f32> },
140}
141
142#[derive(Debug, PartialEq, Serialize, Deserialize)]
144pub struct Doc {
145 pk: Option<String>,
146 #[serde(default)]
147 score: f32,
148 #[serde(default)]
149 #[serde(rename = "doc_id")]
150 internal_id: Option<u64>,
151 #[serde(default)]
152 fields: BTreeMap<String, FieldValue>,
153 #[serde(default)]
154 vectors: BTreeMap<String, VectorValue>,
155}
156
157pub(crate) type DocumentMap = OrdMap<String, Arc<Doc>>;
161
162#[cfg(test)]
163thread_local! {
164 static DOC_BODY_CLONES: std::cell::Cell<u64> = const { std::cell::Cell::new(0) };
165}
166
167#[cfg(test)]
168pub(crate) fn reset_doc_body_clones() {
169 DOC_BODY_CLONES.with(|count| count.set(0));
170}
171
172#[cfg(test)]
173pub(crate) fn doc_body_clones() -> u64 {
174 DOC_BODY_CLONES.with(std::cell::Cell::get)
175}
176
177impl Clone for Doc {
178 fn clone(&self) -> Self {
179 #[cfg(test)]
180 DOC_BODY_CLONES.with(|count| count.set(count.get().saturating_add(1)));
181 Self {
182 pk: self.pk.clone(),
183 score: self.score,
184 internal_id: self.internal_id,
185 fields: self.fields.clone(),
186 vectors: self.vectors.clone(),
187 }
188 }
189}
190
191impl Default for Doc {
192 fn default() -> Self {
193 Self::new().unwrap_or_else(|_| Self {
194 pk: None,
195 score: 0.0,
196 internal_id: None,
197 fields: BTreeMap::new(),
198 vectors: BTreeMap::new(),
199 })
200 }
201}
202
203impl Doc {
204 pub(crate) fn from_persisted_parts(
205 pk: Option<String>,
206 score: f32,
207 internal_id: Option<u64>,
208 fields: BTreeMap<String, FieldValue>,
209 vectors: BTreeMap<String, VectorValue>,
210 ) -> Self {
211 Self {
212 pk,
213 score,
214 internal_id,
215 fields,
216 vectors,
217 }
218 }
219
220 pub fn new() -> Result<Self> {
221 Ok(Self {
222 pk: None,
223 score: 0.0,
224 internal_id: None,
225 fields: BTreeMap::new(),
226 vectors: BTreeMap::new(),
227 })
228 }
229
230 pub fn with_pk(pk: impl Into<String>) -> Result<Self> {
231 let mut doc = Self::new()?;
232 doc.set_pk(&pk.into());
233 Ok(doc)
234 }
235
236 pub fn set_pk(&mut self, pk: &str) {
239 self.pk = Some(pk.to_string());
240 }
241
242 pub fn get_pk(&self) -> Option<&str> {
243 self.pk.as_deref()
244 }
245
246 pub fn get_score(&self) -> f32 {
247 self.score
248 }
249
250 pub fn score(&self) -> f32 {
251 self.score
252 }
253
254 pub fn set_score(&mut self, score: f32) -> Result<()> {
255 if !score.is_finite() {
256 return Err(Error::invalid_argument("score must be finite"));
257 }
258 self.score = score;
259 Ok(())
260 }
261
262 pub fn doc_id(&self) -> Option<u64> {
266 self.internal_id
267 }
268
269 pub(crate) fn set_internal_id(&mut self, doc_id: Option<u64>) {
270 self.internal_id = doc_id;
271 }
272
273 pub fn field_count(&self) -> usize {
274 self.fields.len() + self.vectors.len()
275 }
276
277 pub fn is_empty(&self) -> bool {
278 self.fields.is_empty() && self.vectors.is_empty() && self.pk.is_none()
279 }
280
281 pub fn has_field(&self, name: &str) -> bool {
282 self.fields.contains_key(name) || self.vectors.contains_key(name)
283 }
284
285 pub fn is_field_null(&self, name: &str) -> bool {
286 matches!(self.fields.get(name), Some(FieldValue::Null))
287 }
288
289 pub fn field(&self, name: &str) -> Option<&FieldValue> {
290 self.fields.get(name)
291 }
292
293 pub fn vector(&self, name: &str) -> Option<&VectorValue> {
294 self.vectors.get(name)
295 }
296
297 pub fn fields(&self) -> &BTreeMap<String, FieldValue> {
298 &self.fields
299 }
300
301 pub fn vectors(&self) -> &BTreeMap<String, VectorValue> {
302 &self.vectors
303 }
304
305 pub fn set_field_value(&mut self, name: &str, value: FieldValue) -> Result<()> {
306 validate_name(name)?;
307 validate_field_finite(&value)?;
308 self.fields.insert(name.to_string(), value);
309 Ok(())
310 }
311
312 pub fn set_vector_value(&mut self, name: &str, value: VectorValue) -> Result<()> {
313 validate_name(name)?;
314 validate_vector(&value)?;
315 self.vectors.insert(name.to_string(), value);
316 Ok(())
317 }
318
319 pub fn add_string(&mut self, name: &str, value: &str) -> Result<()> {
320 self.set_field_value(name, FieldValue::String(value.to_string()))
321 }
322
323 pub fn add_bool(&mut self, name: &str, value: bool) -> Result<()> {
324 self.set_field_value(name, FieldValue::Bool(value))
325 }
326
327 pub fn add_i32(&mut self, name: &str, value: i32) -> Result<()> {
328 self.set_field_value(name, FieldValue::Int32(value))
329 }
330
331 pub fn add_i64(&mut self, name: &str, value: i64) -> Result<()> {
332 self.set_field_value(name, FieldValue::Int64(value))
333 }
334
335 pub fn add_u32(&mut self, name: &str, value: u32) -> Result<()> {
336 self.set_field_value(name, FieldValue::Uint32(value))
337 }
338
339 pub fn add_u64(&mut self, name: &str, value: u64) -> Result<()> {
340 self.set_field_value(name, FieldValue::Uint64(value))
341 }
342
343 pub fn add_f32(&mut self, name: &str, value: f32) -> Result<()> {
344 self.set_field_value(name, FieldValue::Float(value))
345 }
346
347 pub fn add_f64(&mut self, name: &str, value: f64) -> Result<()> {
348 self.set_field_value(name, FieldValue::Double(value))
349 }
350
351 pub fn add_binary(&mut self, name: &str, value: &[u8]) -> Result<()> {
352 self.set_field_value(name, FieldValue::Binary(value.to_vec()))
353 }
354
355 pub fn add_array_binary(&mut self, name: &str, values: &[Vec<u8>]) -> Result<()> {
356 self.set_field_value(name, FieldValue::ArrayBinary(values.to_vec()))
357 }
358
359 pub fn add_array_string(&mut self, name: &str, values: &[&str]) -> Result<()> {
360 self.set_field_value(
361 name,
362 FieldValue::ArrayString(values.iter().map(|v| (*v).to_string()).collect()),
363 )
364 }
365
366 pub fn add_array_i32(&mut self, name: &str, values: &[i32]) -> Result<()> {
367 self.set_field_value(name, FieldValue::ArrayInt32(values.to_vec()))
368 }
369
370 pub fn add_array_i64(&mut self, name: &str, values: &[i64]) -> Result<()> {
371 self.set_field_value(name, FieldValue::ArrayInt64(values.to_vec()))
372 }
373
374 pub fn add_array_u32(&mut self, name: &str, values: &[u32]) -> Result<()> {
375 self.set_field_value(name, FieldValue::ArrayUint32(values.to_vec()))
376 }
377
378 pub fn add_array_u64(&mut self, name: &str, values: &[u64]) -> Result<()> {
379 self.set_field_value(name, FieldValue::ArrayUint64(values.to_vec()))
380 }
381
382 pub fn add_array_f32(&mut self, name: &str, values: &[f32]) -> Result<()> {
383 self.set_field_value(name, FieldValue::ArrayFloat(values.to_vec()))
384 }
385
386 pub fn add_array_f64(&mut self, name: &str, values: &[f64]) -> Result<()> {
387 self.set_field_value(name, FieldValue::ArrayDouble(values.to_vec()))
388 }
389
390 pub fn add_array_bool(&mut self, name: &str, values: &[bool]) -> Result<()> {
391 self.set_field_value(name, FieldValue::ArrayBool(values.to_vec()))
392 }
393
394 pub fn set_field_null(&mut self, name: &str) -> Result<()> {
395 self.set_field_value(name, FieldValue::Null)
396 }
397
398 pub fn remove_field(&mut self, name: &str) -> Result<()> {
399 validate_name(name)?;
400 self.fields.remove(name);
401 self.vectors.remove(name);
402 Ok(())
403 }
404
405 pub fn clear(&mut self) {
406 self.fields.clear();
407 self.vectors.clear();
408 self.score = 0.0;
409 }
410
411 pub fn get_string(&self, name: &str) -> Result<Option<String>> {
412 Ok(match self.fields.get(name) {
413 Some(FieldValue::String(v)) => Some(v.clone()),
414 Some(FieldValue::Null) | None => None,
415 Some(_) => return Err(type_error(name, DataType::String)),
416 })
417 }
418
419 pub fn get_bool(&self, name: &str) -> Result<Option<bool>> {
420 self.get_scalar(
421 name,
422 |v| match v {
423 FieldValue::Bool(x) => Some(*x),
424 _ => None,
425 },
426 DataType::Bool,
427 )
428 }
429
430 pub fn get_i32(&self, name: &str) -> Result<Option<i32>> {
431 self.get_scalar(
432 name,
433 |v| match v {
434 FieldValue::Int32(x) => Some(*x),
435 _ => None,
436 },
437 DataType::Int32,
438 )
439 }
440
441 pub fn get_i64(&self, name: &str) -> Result<Option<i64>> {
442 self.get_scalar(
443 name,
444 |v| match v {
445 FieldValue::Int64(x) => Some(*x),
446 _ => None,
447 },
448 DataType::Int64,
449 )
450 }
451
452 pub fn get_u32(&self, name: &str) -> Result<Option<u32>> {
453 self.get_scalar(
454 name,
455 |v| match v {
456 FieldValue::Uint32(x) => Some(*x),
457 _ => None,
458 },
459 DataType::Uint32,
460 )
461 }
462
463 pub fn get_u64(&self, name: &str) -> Result<Option<u64>> {
464 self.get_scalar(
465 name,
466 |v| match v {
467 FieldValue::Uint64(x) => Some(*x),
468 _ => None,
469 },
470 DataType::Uint64,
471 )
472 }
473
474 pub fn get_f32(&self, name: &str) -> Result<Option<f32>> {
475 self.get_scalar(
476 name,
477 |v| match v {
478 FieldValue::Float(x) => Some(*x),
479 _ => None,
480 },
481 DataType::Float,
482 )
483 }
484
485 pub fn get_f64(&self, name: &str) -> Result<Option<f64>> {
486 self.get_scalar(
487 name,
488 |v| match v {
489 FieldValue::Double(x) => Some(*x),
490 _ => None,
491 },
492 DataType::Double,
493 )
494 }
495
496 pub fn get_binary(&self, name: &str) -> Result<Option<Vec<u8>>> {
497 self.get_scalar(
498 name,
499 |v| match v {
500 FieldValue::Binary(x) => Some(x.clone()),
501 _ => None,
502 },
503 DataType::Binary,
504 )
505 }
506
507 pub fn get_array_i32(&self, name: &str) -> Result<Option<Vec<i32>>> {
508 self.get_array(
509 name,
510 |v| match v {
511 FieldValue::ArrayInt32(x) => Some(x.clone()),
512 _ => None,
513 },
514 DataType::ArrayInt32,
515 )
516 }
517 pub fn get_array_i64(&self, name: &str) -> Result<Option<Vec<i64>>> {
518 self.get_array(
519 name,
520 |v| match v {
521 FieldValue::ArrayInt64(x) => Some(x.clone()),
522 _ => None,
523 },
524 DataType::ArrayInt64,
525 )
526 }
527 pub fn get_array_u32(&self, name: &str) -> Result<Option<Vec<u32>>> {
528 self.get_array(
529 name,
530 |v| match v {
531 FieldValue::ArrayUint32(x) => Some(x.clone()),
532 _ => None,
533 },
534 DataType::ArrayUint32,
535 )
536 }
537 pub fn get_array_u64(&self, name: &str) -> Result<Option<Vec<u64>>> {
538 self.get_array(
539 name,
540 |v| match v {
541 FieldValue::ArrayUint64(x) => Some(x.clone()),
542 _ => None,
543 },
544 DataType::ArrayUint64,
545 )
546 }
547 pub fn get_array_f32(&self, name: &str) -> Result<Option<Vec<f32>>> {
548 self.get_array(
549 name,
550 |v| match v {
551 FieldValue::ArrayFloat(x) => Some(x.clone()),
552 _ => None,
553 },
554 DataType::ArrayFloat,
555 )
556 }
557 pub fn get_array_f64(&self, name: &str) -> Result<Option<Vec<f64>>> {
558 self.get_array(
559 name,
560 |v| match v {
561 FieldValue::ArrayDouble(x) => Some(x.clone()),
562 _ => None,
563 },
564 DataType::ArrayDouble,
565 )
566 }
567 pub fn get_array_bool(&self, name: &str) -> Result<Option<Vec<bool>>> {
568 self.get_array(
569 name,
570 |v| match v {
571 FieldValue::ArrayBool(x) => Some(x.clone()),
572 _ => None,
573 },
574 DataType::ArrayBool,
575 )
576 }
577
578 pub fn project(&self, output_fields: Option<&[String]>, include_vector: bool) -> Self {
580 let mut out = self.clone();
581 if let Some(fields) = output_fields {
582 let wanted: std::collections::BTreeSet<&str> =
583 fields.iter().map(String::as_str).collect();
584 out.fields.retain(|k, _| wanted.contains(k.as_str()));
585 if include_vector {
586 out.vectors.retain(|k, _| wanted.contains(k.as_str()));
587 } else {
588 out.vectors.clear();
589 }
590 } else if !include_vector {
591 out.vectors.clear();
592 }
593 out
594 }
595
596 pub(crate) fn scalar_json(&self, name: &str) -> Option<Value> {
597 self.fields.get(name).map(FieldValue::to_json)
598 }
599
600 fn get_scalar<T, F>(&self, name: &str, f: F, expected: DataType) -> Result<Option<T>>
601 where
602 F: FnOnce(&FieldValue) -> Option<T>,
603 {
604 match self.fields.get(name) {
605 None | Some(FieldValue::Null) => Ok(None),
606 Some(value) => f(value).map(Some).ok_or_else(|| type_error(name, expected)),
607 }
608 }
609
610 fn get_array<T, F>(&self, name: &str, f: F, expected: DataType) -> Result<Option<T>>
611 where
612 F: FnOnce(&FieldValue) -> Option<T>,
613 {
614 self.get_scalar(name, f, expected)
615 }
616}
617
618fn validate_name(name: &str) -> Result<()> {
619 if name.is_empty() || name.contains('\0') {
620 return Err(Error::invalid_argument(
621 "field name must be non-empty and contain no NUL byte",
622 ));
623 }
624 Ok(())
625}
626
627fn validate_field_finite(value: &FieldValue) -> Result<()> {
628 let finite = match value {
629 FieldValue::Float(v) => v.is_finite(),
630 FieldValue::Double(v) => v.is_finite(),
631 FieldValue::ArrayFloat(v) => v.iter().all(|x| x.is_finite()),
632 FieldValue::ArrayDouble(v) => v.iter().all(|x| x.is_finite()),
633 _ => true,
634 };
635 finite
636 .then_some(())
637 .ok_or_else(|| Error::invalid_argument("floating-point field values must be finite"))
638}
639
640fn type_error(name: &str, expected: DataType) -> Error {
641 Error::invalid_argument(format!("field '{name}' is not of type {expected}"))
642}
643
644fn number_from_f64(value: f64) -> Value {
645 Number::from_f64(value).map_or(Value::Null, Value::Number)
646}
647
648fn base64_encode(bytes: &[u8]) -> String {
651 const TABLE: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
652 let mut out = String::with_capacity(bytes.len().div_ceil(3) * 4);
653 for chunk in bytes.chunks(3) {
654 let a = u32::from(chunk[0]);
655 let b = u32::from(chunk.get(1).copied().unwrap_or(0));
656 let c = u32::from(chunk.get(2).copied().unwrap_or(0));
657 out.push(TABLE[((a >> 2) & 63) as usize] as char);
658 out.push(TABLE[(((a << 4) | (b >> 4)) & 63) as usize] as char);
659 if chunk.len() > 1 {
660 out.push(TABLE[(((b << 2) | (c >> 6)) & 63) as usize] as char);
661 } else {
662 out.push('=');
663 }
664 if chunk.len() > 2 {
665 out.push(TABLE[(c & 63) as usize] as char);
666 } else {
667 out.push('=');
668 }
669 }
670 out
671}
672
673#[cfg(test)]
674mod tests {
675 use super::{Doc, FieldValue, VectorValue};
676 use serde_json::json;
677
678 #[test]
679 fn field_value_to_json_covers_every_variant() {
680 assert_eq!(FieldValue::Null.to_json(), json!(null));
681 assert_eq!(FieldValue::String("x".into()).to_json(), json!("x"));
682 assert_eq!(FieldValue::Bool(true).to_json(), json!(true));
683 assert_eq!(FieldValue::Int32(-1).to_json(), json!(-1));
684 assert_eq!(FieldValue::Int64(-2).to_json(), json!(-2));
685 assert_eq!(FieldValue::Uint32(3).to_json(), json!(3));
686 assert_eq!(FieldValue::Uint64(4).to_json(), json!(4));
687 assert_eq!(FieldValue::Float(1.5).to_json(), json!(1.5));
688 assert_eq!(FieldValue::Double(2.5).to_json(), json!(2.5));
689 assert_eq!(
690 FieldValue::Binary(vec![0, 255]).to_json(),
691 json!(super::base64_encode(&[0, 255]))
692 );
693 assert_eq!(
694 FieldValue::ArrayBinary(vec![vec![1], vec![2]]).to_json(),
695 json!([super::base64_encode(&[1]), super::base64_encode(&[2])])
696 );
697 assert_eq!(
698 FieldValue::ArrayString(vec!["a".into()]).to_json(),
699 json!(["a"])
700 );
701 assert_eq!(
702 FieldValue::ArrayBool(vec![true, false]).to_json(),
703 json!([true, false])
704 );
705 assert_eq!(
706 FieldValue::ArrayInt32(vec![1, -2]).to_json(),
707 json!([1, -2])
708 );
709 assert_eq!(
710 FieldValue::ArrayInt64(vec![3, -4]).to_json(),
711 json!([3, -4])
712 );
713 assert_eq!(FieldValue::ArrayUint32(vec![5]).to_json(), json!([5]));
714 assert_eq!(FieldValue::ArrayUint64(vec![6]).to_json(), json!([6]));
715 assert_eq!(FieldValue::ArrayFloat(vec![1.25]).to_json(), json!([1.25]));
716 assert_eq!(FieldValue::ArrayDouble(vec![2.5]).to_json(), json!([2.5]));
717 assert_eq!(FieldValue::Json(json!({"k": 1})).to_json(), json!({"k": 1}));
718 }
719
720 #[test]
721 fn scalar_json_and_projection_keep_requested_fields() {
722 let mut doc = Doc::with_pk("pk").expect("pk");
723 doc.set_score(1.25).expect("score");
724 doc.add_string("title", "hello").expect("string");
725 doc.add_i32("n", 7).expect("i32");
726 doc.add_vector_f32("embedding", &[1.0, 0.0])
727 .expect("vector");
728 assert_eq!(doc.scalar_json("title"), Some(json!("hello")));
729 assert_eq!(doc.scalar_json("missing"), None);
730 let projected = doc.project(Some(&["title".into()]), true);
731 assert!(projected.has_field("title"));
732 assert!(!projected.has_field("n"));
733 let _ = VectorValue::Fp32(vec![1.0]);
734 }
735
736 #[test]
737 fn typed_vector_getters_reject_mismatched_storage_variants() {
738 let mut doc = Doc::with_pk("pk").expect("pk");
739 doc.add_vector_f32("embedding", &[1.0, 0.0]).expect("f32");
740 assert!(doc.get_vector_f64("embedding").is_err());
741 assert!(doc.get_vector_fp16("embedding").is_err());
742 assert!(doc.get_vector_i4("embedding").is_err());
743 assert!(doc.get_vector_i8("embedding").is_err());
744 assert!(doc.get_vector_i16("embedding").is_err());
745 assert!(doc.get_vector_binary32("embedding").is_err());
746 assert!(doc.get_vector_binary64("embedding").is_err());
747 assert!(doc.get_sparse_vector_f32("embedding").is_err());
748 assert!(doc.get_sparse_vector_fp16("embedding").is_err());
749 assert_eq!(doc.get_vector_f32("missing").expect("ok"), None);
750 assert_eq!(doc.get_vector_f64("missing").expect("ok"), None);
751 assert_eq!(doc.get_vector_fp16("missing").expect("ok"), None);
752 assert_eq!(doc.get_vector_i4("missing").expect("ok"), None);
753 assert_eq!(doc.get_vector_i8("missing").expect("ok"), None);
754 assert_eq!(doc.get_vector_i16("missing").expect("ok"), None);
755 assert_eq!(doc.get_vector_binary32("missing").expect("ok"), None);
756 assert_eq!(doc.get_vector_binary64("missing").expect("ok"), None);
757 assert_eq!(doc.get_sparse_vector_f32("missing").expect("ok"), None);
758 assert_eq!(doc.get_sparse_vector_fp16("missing").expect("ok"), None);
759
760 doc.set_vector_value("bits", VectorValue::Binary32(vec![0xff; 4]))
761 .expect("bits");
762 assert_eq!(
763 doc.get_vector_binary32("bits").expect("ok"),
764 Some(vec![0xff; 4])
765 );
766 assert!(doc.get_vector_f32("bits").is_err());
767 doc.set_vector_value("bits64", VectorValue::Binary64(vec![1; 8]))
768 .expect("bits64");
769 assert_eq!(
770 doc.get_vector_binary64("bits64").expect("ok"),
771 Some(vec![1; 8])
772 );
773 doc.set_vector_value(
774 "sparse",
775 VectorValue::SparseFp32 {
776 indices: vec![0],
777 values: vec![1.0],
778 },
779 )
780 .expect("sparse");
781 assert_eq!(
782 doc.get_sparse_vector_f32("sparse").expect("ok"),
783 Some((vec![0], vec![1.0]))
784 );
785 doc.set_vector_value(
786 "sparse16",
787 VectorValue::SparseFp16 {
788 indices: vec![1],
789 values: vec![0x3c00],
790 },
791 )
792 .expect("sparse16");
793 assert_eq!(
794 doc.get_sparse_vector_fp16("sparse16").expect("ok"),
795 Some((vec![1], vec![0x3c00]))
796 );
797
798 let bad = VectorValue::SparseFp32 {
800 indices: vec![0, 1],
801 values: vec![1.0],
802 };
803 assert!(bad.to_sparse_f64().is_none());
804 let bad16 = VectorValue::SparseFp16 {
805 indices: vec![0],
806 values: vec![0x3c00, 0x4000],
807 };
808 assert!(bad16.to_sparse_f64().is_none());
809 }
810
811 #[test]
812 fn typed_scalar_getters_reject_mismatched_field_types() {
813 let mut doc = Doc::with_pk("pk").expect("pk");
814 doc.add_string("title", "hello").expect("string");
815 doc.add_bool("flag", true).expect("bool");
816 doc.add_i32("i32", 1).expect("i32");
817 doc.add_i64("i64", 2).expect("i64");
818 doc.add_u32("u32", 3).expect("u32");
819 doc.add_u64("u64", 4).expect("u64");
820 doc.add_f32("f32", 1.5).expect("f32");
821 doc.add_f64("f64", 2.5).expect("f64");
822 doc.add_binary("bin", &[1, 2]).expect("bin");
823 doc.add_array_i32("ai32", &[1]).expect("ai32");
824 doc.add_array_i64("ai64", &[2]).expect("ai64");
825 doc.add_array_u32("au32", &[3]).expect("au32");
826 doc.add_array_u64("au64", &[4]).expect("au64");
827 doc.add_array_f32("af32", &[1.0]).expect("af32");
828 doc.add_array_f64("af64", &[2.0]).expect("af64");
829 doc.add_array_bool("abool", &[true]).expect("abool");
830
831 assert!(doc.get_string("flag").is_err());
832 assert!(doc.get_bool("title").is_err());
833 assert!(doc.get_i32("title").is_err());
834 assert!(doc.get_i64("title").is_err());
835 assert!(doc.get_u32("title").is_err());
836 assert!(doc.get_u64("title").is_err());
837 assert!(doc.get_f32("title").is_err());
838 assert!(doc.get_f64("title").is_err());
839 assert!(doc.get_binary("title").is_err());
840 assert!(doc.get_array_i32("title").is_err());
841 assert!(doc.get_array_i64("title").is_err());
842 assert!(doc.get_array_u32("title").is_err());
843 assert!(doc.get_array_u64("title").is_err());
844 assert!(doc.get_array_f32("title").is_err());
845 assert!(doc.get_array_f64("title").is_err());
846 assert!(doc.get_array_bool("title").is_err());
847
848 assert_eq!(doc.get_string("title").expect("ok"), Some("hello".into()));
849 assert_eq!(doc.get_bool("flag").expect("ok"), Some(true));
850 assert_eq!(doc.get_i32("i32").expect("ok"), Some(1));
851 assert_eq!(doc.get_i64("i64").expect("ok"), Some(2));
852 assert_eq!(doc.get_u32("u32").expect("ok"), Some(3));
853 assert_eq!(doc.get_u64("u64").expect("ok"), Some(4));
854 assert_eq!(doc.get_f32("f32").expect("ok"), Some(1.5));
855 assert_eq!(doc.get_f64("f64").expect("ok"), Some(2.5));
856 assert_eq!(doc.get_binary("bin").expect("ok"), Some(vec![1, 2]));
857 assert_eq!(doc.get_array_i32("ai32").expect("ok"), Some(vec![1]));
858 assert_eq!(doc.get_array_i64("ai64").expect("ok"), Some(vec![2]));
859 assert_eq!(doc.get_array_u32("au32").expect("ok"), Some(vec![3]));
860 assert_eq!(doc.get_array_u64("au64").expect("ok"), Some(vec![4]));
861 assert_eq!(doc.get_array_f32("af32").expect("ok"), Some(vec![1.0]));
862 assert_eq!(doc.get_array_f64("af64").expect("ok"), Some(vec![2.0]));
863 assert_eq!(doc.get_array_bool("abool").expect("ok"), Some(vec![true]));
864 assert_eq!(doc.get_string("missing").expect("ok"), None);
865 let _ = Doc::default();
866 }
867}