Skip to content

Commit f01baf5

Browse files
committed
refactor: Simplify payload conversion
Signed-off-by: Anush008 <[email protected]>
1 parent 678f9fd commit f01baf5

File tree

1 file changed

+36
-64
lines changed

1 file changed

+36
-64
lines changed

src/ops/storages/qdrant.rs

+36-64
Original file line numberDiff line numberDiff line change
@@ -8,12 +8,14 @@ use crate::setup;
88
use anyhow::{bail, Result};
99
use derivative::Derivative;
1010
use futures::FutureExt;
11-
use qdrant_client::qdrant::value::Kind;
1211
use qdrant_client::qdrant::vectors_output::VectorsOptions;
13-
use qdrant_client::qdrant::{NamedVectors, PointStruct, UpsertPointsBuilder, Value as QdrantValue};
12+
use qdrant_client::qdrant::{
13+
NamedVectors, PointStruct, UpsertPointsBuilder, Value as QdrantValue,
14+
};
1415
use qdrant_client::qdrant::{Query, QueryPointsBuilder, ScoredPoint};
1516
use qdrant_client::Qdrant;
1617
use serde::Serialize;
18+
use serde_json::json;
1719

1820
fn key_value_fields_iter<'a>(
1921
key_fields_schema: &[FieldSchema],
@@ -71,7 +73,7 @@ impl ExportTargetExecutor for Executor {
7173
let mut points: Vec<PointStruct> = Vec::with_capacity(mutation.upserts.len());
7274
for upsert in mutation.upserts.iter() {
7375
let key_fields = key_value_fields_iter(&self.key_fields_schema, &upsert.key)?;
74-
let key_fields = key_values_to_payload(&key_fields, &self.key_fields_schema)?;
76+
let key_fields = key_values_to_payload(key_fields, &self.key_fields_schema)?;
7577
let (mut payload, vectors) =
7678
values_to_payload(&upsert.value.fields, &self.value_fields_schema)?;
7779
payload.extend(key_fields);
@@ -93,23 +95,16 @@ fn key_values_to_payload(
9395
let mut payload = HashMap::with_capacity(key_fields.len());
9496

9597
for (key_value, field_schema) in key_fields.iter().zip(schema.iter()) {
96-
let value = match key_value {
97-
KeyValue::Bytes(v) => QdrantValue {
98-
kind: Some(Kind::StringValue(String::from_utf8_lossy(v).into_owned())),
99-
},
100-
KeyValue::Str(v) => QdrantValue {
101-
kind: Some(Kind::StringValue(v.clone().to_string())),
102-
},
103-
KeyValue::Bool(v) => QdrantValue {
104-
kind: Some(Kind::BoolValue(*v)),
105-
},
106-
KeyValue::Int64(v) => QdrantValue {
107-
kind: Some(Kind::IntegerValue(*v)),
108-
},
109-
e => anyhow::bail!("Unsupported key value type {}", e),
98+
let json_value = match key_value {
99+
KeyValue::Bytes(v) => String::from_utf8_lossy(v).into(),
100+
KeyValue::Str(v) => v.to_string().into(),
101+
KeyValue::Bool(v) => (*v).into(),
102+
KeyValue::Int64(v) => (*v).into(),
103+
KeyValue::Uuid(v) => v.to_string().into(),
104+
KeyValue::Range(v) => json!({ "start": v.start, "end": v.end }),
105+
_ => bail!("Unsupported key value type"),
110106
};
111-
112-
payload.insert(field_schema.name.clone(), value);
107+
payload.insert(field_schema.name.clone(), json_value.into());
113108
}
114109

115110
Ok(payload)
@@ -124,59 +119,36 @@ fn values_to_payload(
124119

125120
for (value, field_schema) in value_fields.iter().zip(schema.iter()) {
126121
let field_name = &field_schema.name;
122+
127123
match value {
128-
Value::Basic(basic_value) => match basic_value {
129-
BasicValue::Bytes(v) => insert_qdrant_value(
130-
&mut payload,
131-
field_name,
132-
Kind::StringValue(String::from_utf8_lossy(v).into_owned()),
133-
),
134-
BasicValue::Str(v) => insert_qdrant_value(
135-
&mut payload,
136-
field_name,
137-
Kind::StringValue(v.clone().to_string()),
138-
),
139-
BasicValue::Bool(v) => {
140-
insert_qdrant_value(&mut payload, field_name, Kind::BoolValue(*v))
141-
}
142-
BasicValue::Int64(v) => {
143-
insert_qdrant_value(&mut payload, field_name, Kind::IntegerValue(*v))
144-
}
145-
BasicValue::Float32(v) => {
146-
insert_qdrant_value(&mut payload, field_name, Kind::DoubleValue(*v as f64))
147-
}
148-
BasicValue::Float64(v) => {
149-
insert_qdrant_value(&mut payload, field_name, Kind::DoubleValue(*v))
150-
}
151-
BasicValue::Range(v) => insert_qdrant_value(
152-
&mut payload,
153-
field_name,
154-
Kind::StringValue(format!("[{}, {})", v.start, v.end)),
155-
),
156-
BasicValue::Vector(v) => {
157-
let vector = convert_to_vector(v.to_vec());
158-
vectors = vectors.add_vector(field_name, vector);
159-
}
160-
_ => {
161-
bail!("Unsupported BasicValue type in Value::Basic");
162-
}
163-
},
164-
Value::Null => {
165-
payload.insert(field_schema.name.clone(), QdrantValue { kind: None });
124+
Value::Basic(basic_value) => {
125+
let json_value = match basic_value {
126+
BasicValue::Bytes(v) => String::from_utf8_lossy(v).into(),
127+
BasicValue::Str(v) => v.clone().to_string().into(),
128+
BasicValue::Bool(v) => (*v).into(),
129+
BasicValue::Int64(v) => (*v).into(),
130+
BasicValue::Float32(v) => (*v as f64).into(),
131+
BasicValue::Float64(v) => (*v).into(),
132+
BasicValue::Range(v) => json!([v.start, v.end]),
133+
BasicValue::Vector(v) => {
134+
let vector = convert_to_vector(v.to_vec());
135+
vectors = vectors.add_vector(field_name, vector);
136+
continue;
137+
}
138+
_ => bail!("Unsupported BasicValue type in Value::Basic"),
139+
};
140+
payload.insert(field_name.clone(), json_value.into());
166141
}
167-
_ => {
168-
bail!("Unsupported Value variant: {:?}", value);
142+
Value::Null => {
143+
payload.insert(field_name.clone(), QdrantValue { kind: None });
169144
}
145+
_ => bail!("Unsupported Value variant: {:?}", value),
170146
}
171147
}
172148

173149
Ok((payload, vectors))
174150
}
175151

176-
fn insert_qdrant_value(payload: &mut HashMap<String, QdrantValue>, field_name: &str, kind: Kind) {
177-
payload.insert(field_name.to_string(), QdrantValue { kind: Some(kind) });
178-
}
179-
180152
fn convert_to_vector(v: Vec<BasicValue>) -> Vec<f32> {
181153
v.iter()
182154
.filter_map(|elem| match elem {
@@ -393,7 +365,7 @@ impl StorageFactoryBase for Arc<Factory> {
393365
let setup_state = SetupState {};
394366
let executors = async move {
395367
let executor = Arc::new(Executor::new(
396-
&url,
368+
url,
397369
&collection_name,
398370
key_fields_schema,
399371
value_fields_schema,

0 commit comments

Comments
 (0)