From 1f6959bd942547afc3845fd37ca4a3b97dc7002a Mon Sep 17 00:00:00 2001
From: Flavien David <flavien.david74@gmail.com>
Date: Thu, 23 May 2024 17:38:54 +0200
Subject: [PATCH] Create Qdrant indexes when creating collection (#5249)

---
 core/bin/qdrant/create_collection.rs | 69 +++++++++++++++++++++++++++-
 1 file changed, 68 insertions(+), 1 deletion(-)

diff --git a/core/bin/qdrant/create_collection.rs b/core/bin/qdrant/create_collection.rs
index 06a5206921b7..e72159acb18b 100644
--- a/core/bin/qdrant/create_collection.rs
+++ b/core/bin/qdrant/create_collection.rs
@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use anyhow::{anyhow, Result};
 use clap::Parser;
 use dust::{
@@ -8,7 +10,7 @@ use dust::{
     },
     utils,
 };
-use qdrant_client::qdrant;
+use qdrant_client::{client::QdrantClient, qdrant};
 use tokio;
 
 #[derive(Parser, Debug)]
@@ -27,6 +29,69 @@ struct Args {
     cluster: QdrantCluster,
 }
 
+async fn create_indexes_for_collection(
+    raw_client: &Arc<QdrantClient>,
+    cluster: &QdrantCluster,
+    collection_name: &String,
+) -> Result<()> {
+    let _ = raw_client
+        .create_field_index(
+            collection_name,
+            "document_id_hash",
+            qdrant::FieldType::Keyword,
+            None,
+            None,
+        )
+        .await?;
+
+    let _ = raw_client
+        .create_field_index(
+            collection_name,
+            "data_source_internal_id",
+            qdrant::FieldType::Keyword,
+            None,
+            None,
+        )
+        .await?;
+
+    let _ = raw_client
+        .create_field_index(
+            collection_name,
+            "tags",
+            qdrant::FieldType::Keyword,
+            None,
+            None,
+        )
+        .await?;
+
+    let _ = raw_client
+        .create_field_index(
+            collection_name,
+            "parents",
+            qdrant::FieldType::Keyword,
+            None,
+            None,
+        )
+        .await?;
+
+    let _ = raw_client
+        .create_field_index(
+            collection_name,
+            "timestamp",
+            qdrant::FieldType::Integer,
+            None,
+            None,
+        )
+        .await?;
+
+    println!(
+        "Done creating indexes for collection {} on cluster {}",
+        collection_name, cluster
+    );
+
+    Ok(())
+}
+
 async fn create_qdrant_collection(
     cluster: QdrantCluster,
     provider_id: ProviderID,
@@ -141,6 +206,8 @@ async fn create_qdrant_collection(
         }?;
     }
 
+    create_indexes_for_collection(&raw_client, &cluster, &collection_name).await?;
+
     Ok(())
 }