cleanup(qdrant): remove query related logic for qdrant

Author: badmonster0

examples/text_embedding_qdrant/main.py

@@ -26,7 +26,7 @@ def text_to_embedding(
 @cocoindex.flow_def(name="TextEmbeddingWithQdrant")
 def text_embedding_flow(
     flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.DataScope
-):
+) -> None:
     """
     Define an example flow that embeds text into a vector database.
     """
@@ -65,7 +65,7 @@ def text_embedding_flow(
     )
 
 
-def _main():
+def _main() -> None:
     # Initialize Qdrant client
     client = QdrantClient(url=QDRANT_GRPC_URL, prefer_grpc=True)
 
@@ -87,6 +87,8 @@ def _main():
         for result in search_results:
             score = result.score
             payload = result.payload
+            if payload is None:
+                continue
             print(f"[{score:.3f}] {payload['filename']}")
             print(f"    {payload['text']}")
             print("---")

src/ops/registration.rs

@@ -17,7 +17,7 @@ fn register_executor_factories(registry: &mut ExecutorFactoryRegistry) -> Result
     functions::extract_by_llm::Factory.register(registry)?;
 
     storages::postgres::Factory::default().register(registry)?;
-    Arc::new(storages::qdrant::Factory::default()).register(registry)?;
+    storages::qdrant::register(registry)?;
     storages::kuzu::register(registry, reqwest_client)?;
 
     storages::neo4j::Factory::new().register(registry)?;

src/ops/storages/qdrant.rs

@@ -3,28 +3,26 @@ use std::convert::Infallible;
 use std::fmt::Display;
 use std::sync::Arc;
 
-use crate::base::duration::parse_duration;
+use crate::ops::registry::ExecutorFactoryRegistry;
 use crate::ops::sdk::*;
 use crate::setup;
 use anyhow::{Result, bail};
 use futures::FutureExt;
 use qdrant_client::Qdrant;
-use qdrant_client::qdrant::vectors_output::VectorsOptions;
 use qdrant_client::qdrant::{
     DeletePointsBuilder, NamedVectors, PointId, PointStruct, PointsIdsList, UpsertPointsBuilder,
     Value as QdrantValue,
 };
-use qdrant_client::qdrant::{Query, QueryPointsBuilder, ScoredPoint};
 use serde_json::json;
 
 #[derive(Debug, Deserialize, Clone)]
-pub struct Spec {
+struct Spec {
     collection_name: String,
     grpc_url: String,
     api_key: Option<String>,
 }
 
-pub struct ExportContext {
+struct ExportContext {
     client: Qdrant,
     collection_name: String,
     value_fields_schema: Vec<FieldSchema>,
@@ -129,8 +127,8 @@ fn values_to_payload(
                     BasicValue::Range(v) => json!({ "start": v.start, "end": v.end }),
                     BasicValue::Uuid(v) => v.to_string().into(),
                     BasicValue::Date(v) => v.to_string().into(),
-                    BasicValue::LocalDateTime(v) => v.to_string().into(),
                     BasicValue::Time(v) => v.to_string().into(),
+                    BasicValue::LocalDateTime(v) => v.to_string().into(),
                     BasicValue::OffsetDateTime(v) => v.to_string().into(),
                     BasicValue::TimeDelta(v) => v.to_string().into(),
                     BasicValue::Json(v) => (**v).clone(),
@@ -163,167 +161,11 @@ fn convert_to_vector(v: Vec<BasicValue>) -> Vec<f32> {
         .collect()
 }
 
-fn into_value(point: &ScoredPoint, schema: &FieldSchema) -> Result<Value> {
-    let field_name = &schema.name;
-    let typ = schema.value_type.typ.clone();
-    let value = match typ {
-        ValueType::Basic(basic_type) => {
-            let basic_value = match basic_type {
-                BasicValueType::Str => point.payload.get(field_name).and_then(|v| {
-                    v.as_str()
-                        .map(|s| BasicValue::Str(Arc::from(s.to_string())))
-                }),
-                BasicValueType::Bool => point
-                    .payload
-                    .get(field_name)
-                    .and_then(|v| v.as_bool().map(BasicValue::Bool)),
-
-                BasicValueType::Int64 => point
-                    .payload
-                    .get(field_name)
-                    .and_then(|v| v.as_integer().map(BasicValue::Int64)),
-
-                BasicValueType::Float32 => point
-                    .payload
-                    .get(field_name)
-                    .and_then(|v| v.as_double().map(|f| BasicValue::Float32(f as f32))),
-
-                BasicValueType::Float64 => point
-                    .payload
-                    .get(field_name)
-                    .and_then(|v| v.as_double().map(BasicValue::Float64)),
-
-                BasicValueType::TimeDelta => point.payload.get(field_name).and_then(|v| {
-                    v.as_str()
-                        .and_then(|s| parse_duration(s).ok().map(BasicValue::TimeDelta))
-                }),
-
-                BasicValueType::Json => point
-                    .payload
-                    .get(field_name)
-                    .map(|v| BasicValue::Json(Arc::from(v.clone().into_json()))),
-
-                BasicValueType::Vector(_) => point
-                    .vectors
-                    .as_ref()
-                    .and_then(|v| v.vectors_options.as_ref())
-                    .and_then(|vectors_options| match vectors_options {
-                        VectorsOptions::Vector(vector) => {
-                            let values = vector
-                                .data
-                                .iter()
-                                .map(|f| BasicValue::Float32(*f))
-                                .collect::<Vec<_>>();
-                            Some(BasicValue::Vector(Arc::from(values)))
-                        }
-                        VectorsOptions::Vectors(vectors) => {
-                            vectors.vectors.get(field_name).map(|vector| {
-                                let values = vector
-                                    .data
-                                    .iter()
-                                    .map(|f| BasicValue::Float32(*f))
-                                    .collect::<Vec<_>>();
-                                BasicValue::Vector(Arc::from(values))
-                            })
-                        }
-                    }),
-
-                BasicValueType::Uuid => point
-                    .payload
-                    .get(field_name)
-                    .and_then(|v| v.as_str()?.parse().ok().map(BasicValue::Uuid)),
-
-                BasicValueType::Date => point
-                    .payload
-                    .get(field_name)
-                    .and_then(|v| v.as_str()?.parse().ok().map(BasicValue::Date)),
-
-                BasicValueType::Time => point
-                    .payload
-                    .get(field_name)
-                    .and_then(|v| v.as_str()?.parse().ok().map(BasicValue::Time)),
-
-                BasicValueType::LocalDateTime => point
-                    .payload
-                    .get(field_name)
-                    .and_then(|v| v.as_str()?.parse().ok().map(BasicValue::LocalDateTime)),
-
-                BasicValueType::OffsetDateTime => point
-                    .payload
-                    .get(field_name)
-                    .and_then(|v| v.as_str()?.parse().ok().map(BasicValue::OffsetDateTime)),
-
-                BasicValueType::Range => point.payload.get(field_name).and_then(|v| {
-                    v.as_struct().and_then(|s| {
-                        let start = s.fields.get("start").and_then(|f| f.as_integer());
-                        let end = s.fields.get("end").and_then(|f| f.as_integer());
-
-                        match (start, end) {
-                            (Some(start), Some(end)) => Some(BasicValue::Range(RangeValue {
-                                start: start as usize,
-                                end: end as usize,
-                            })),
-                            _ => None,
-                        }
-                    })
-                }),
-                _ => {
-                    anyhow::bail!("Unsupported value type")
-                }
-            };
-            basic_value.map(Value::Basic)
-        }
-        _ => point
-            .payload
-            .get(field_name)
-            .map(|v| Value::from_json(v.clone().into_json(), &typ))
-            .transpose()?,
-    };
-
-    let final_value = if let Some(v) = value { v } else { Value::Null };
-    Ok(final_value)
-}
-
-#[async_trait]
-impl QueryTarget for ExportContext {
-    async fn search(&self, query: VectorMatchQuery) -> Result<QueryResults> {
-        let points = self
-            .client
-            .query(
-                QueryPointsBuilder::new(&self.collection_name)
-                    .query(Query::new_nearest(query.vector))
-                    .limit(query.limit as u64)
-                    .using(query.vector_field_name)
-                    .with_payload(true)
-                    .with_vectors(true),
-            )
-            .await?
-            .result;
-
-        let results = points
-            .iter()
-            .map(|point| {
-                let score = point.score as f64;
-                let data = self
-                    .all_fields
-                    .iter()
-                    .map(|schema| into_value(point, schema))
-                    .collect::<Result<Vec<_>>>()?;
-                Ok(QueryResult { data, score })
-            })
-            .collect::<Result<Vec<QueryResult>>>()?;
-        Ok(QueryResults {
-            fields: self.all_fields.clone(),
-            results,
-        })
-    }
-}
-
 #[derive(Default)]
-pub struct Factory {}
+struct Factory {}
 
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
-pub struct CollectionId {
+struct CollectionId {
     collection_name: String,
 }
 
@@ -335,7 +177,7 @@ impl Display for CollectionId {
 }
 
 #[async_trait]
-impl StorageFactoryBase for Arc<Factory> {
+impl StorageFactoryBase for Factory {
     type Spec = Spec;
     type DeclarationSpec = ();
     type SetupState = ();
@@ -375,11 +217,10 @@ impl StorageFactoryBase for Arc<Factory> {
                     d.key_fields_schema,
                     d.value_fields_schema,
                 )?);
-                let query_target = export_context.clone();
                 let executors = async move {
                     Ok(TypedExportTargetExecutors {
                         export_context,
-                        query_target: Some(query_target as Arc<dyn QueryTarget>),
+                        query_target: None,
                     })
                 };
                 Ok(TypedExportDataCollectionBuildOutput {
@@ -435,3 +276,7 @@ impl StorageFactoryBase for Arc<Factory> {
         Err(anyhow!("Qdrant does not support setup changes"))
     }
 }
+
+pub fn register(registry: &mut ExecutorFactoryRegistry) -> Result<()> {
+    Factory {}.register(registry)
+}