WIP: collection init from other collection

qdrant · generall · Jan 22, 2023 · Jan 12, 2023 · Jan 17, 2023 · Jan 17, 2023
commit a619bff9f89b190599e0ae5eb1e65c9729ae5039
diff --git a/docs/grpc/docs.md b/docs/grpc/docs.md
@@ -333,6 +333,7 @@
 | vectors_config | [VectorsConfig](#qdrant-VectorsConfig) | optional | Configuration for vectors |
 | replication_factor | [uint32](#uint32) | optional | Number of replicas of each shard that network tries to maintain, default = 1 |
 | write_consistency_factor | [uint32](#uint32) | optional | How many replicas should apply the operation for us to consider it successful, default = 1 |
+| init_from_collection | [string](#string) | optional | Specify name of the other collection to copy data from |
 
 
 

diff --git a/docs/redoc/master/openapi.json b/docs/redoc/master/openapi.json
@@ -4261,6 +4261,18 @@
                 "nullable": true
               }
             ]
+          },
+          "init_from": {
+            "description": "Specify other collection to copy data from.",
+            "default": null,
+            "anyOf": [
+              {
+                "$ref": "#/components/schemas/InitFrom"
+              },
+              {
+                "nullable": true
+              }
+            ]
           }
         }
       },
@@ -4391,6 +4403,18 @@
           }
         }
       },
+      "InitFrom": {
+        "description": "Operation for creating new collection and (optionally) specify index params",
+        "type": "object",
+        "required": [
+          "collection"
+        ],
+        "properties": {
+          "collection": {
+            "type": "string"
+          }
+        }
+      },
       "UpdateCollection": {
         "description": "Operation for updating parameters of the existing collection",
         "type": "object",

diff --git a/lib/collection/src/collection.rs b/lib/collection/src/collection.rs
@@ -21,6 +21,7 @@ use tokio::runtime::Handle;
 use tokio::sync::{Mutex, RwLock};
 
 use crate::collection_state::{ShardInfo, State};
+use crate::common::is_ready::IsReady;
 use crate::config::CollectionConfig;
 use crate::hash_ring::HashRing;
 use crate::operations::config_diff::{CollectionParamsDiff, DiffConfig, OptimizersConfigDiff};
@@ -81,6 +82,9 @@ pub struct Collection {
     #[allow(dead_code)] //Might be useful in case of repartition implementation
     notify_peer_failure_cb: OnPeerFailure,
     init_time: Duration,
+    // One-way boolean flag that is set to true when the collection is fully initialized
+    // i.e. all shards are activated for the first time.
+    is_initialized: Arc<IsReady>,
 }
 
 impl Collection {
@@ -151,6 +155,7 @@ impl Collection {
             request_shard_transfer_cb: request_shard_transfer.clone(),
             notify_peer_failure_cb: on_replica_failure.clone(),
             init_time: start_time.elapsed(),
+            is_initialized: Arc::new(Default::default()),
         })
     }
 
@@ -250,6 +255,7 @@ impl Collection {
             request_shard_transfer_cb: request_shard_transfer.clone(),
             notify_peer_failure_cb: on_replica_failure,
             init_time: start_time.elapsed(),
+            is_initialized: Arc::new(Default::default()),
         }
     }
 
@@ -328,6 +334,25 @@ impl Collection {
             }
         }
 
+        if !self.is_initialized.check_ready() {
+            // If not initialized yet, we need to check if it was initialized by this call
+            let state = self.state().await;
+            let mut is_fully_active = true;
+            for (_shard_id, shard_info) in state.shards {
+                if shard_info
+                    .replicas
+                    .into_iter()
+                    .any(|(_peer_id, state)| state != ReplicaState::Active)
+                {
+                    is_fully_active = false;
+                    break;
+                }
+            }
+            if is_fully_active {
+                self.is_initialized.make_ready();
+            }
+        }
+
         // Try to request shard transfer if replicas on the current peer are dead
         if state == ReplicaState::Dead && self.this_peer_id == peer_id {
             let transfer_from = replica_set
@@ -1418,6 +1443,10 @@ impl Collection {
 
         Ok(())
     }
+
+    pub fn wait_collection_initiated(&self, timeout: Duration) -> bool {
+        self.is_initialized.await_ready_for_timeout(timeout)
+    }
 }
 
 impl Drop for Collection {

diff --git a/...src/content_manager/consensus/is_ready.rs → lib/collection/src/common/is_ready.rs b/...src/content_manager/consensus/is_ready.rs → lib/collection/src/common/is_ready.rs
diff --git a/lib/collection/src/common/mod.rs b/lib/collection/src/common/mod.rs
@@ -1,2 +1,3 @@
 pub mod stoppable_task;
 pub mod stoppable_task_async;
+pub mod is_ready;
diff --git a/lib/collection/src/lib.rs b/lib/collection/src/lib.rs
@@ -3,7 +3,7 @@
 pub mod collection;
 pub mod collection_manager;
 pub mod collection_state;
-mod common;
+pub mod common;
 pub mod config;
 pub mod hash_ring;
 pub mod operations;

diff --git a/lib/storage/src/content_manager/collections_ops.rs b/lib/storage/src/content_manager/collections_ops.rs
@@ -2,10 +2,11 @@ use std::collections::HashMap;
 
 use async_trait::async_trait;
 use collection::collection::Collection;
+use collection::shards::CollectionId;
 
 use crate::content_manager::errors::StorageError;
 
-pub type Collections = HashMap<String, Collection>;
+pub type Collections = HashMap<CollectionId, Collection>;
 
 #[async_trait]
 pub trait Checker {

diff --git a/lib/storage/src/content_manager/consensus/mod.rs b/lib/storage/src/content_manager/consensus/mod.rs
@@ -1,5 +1,4 @@
 pub mod consensus_wal;
 pub mod entry_queue;
-pub mod is_ready;
 pub mod operation_sender;
 pub mod persistent;
diff --git a/lib/storage/src/content_manager/consensus_manager.rs b/lib/storage/src/content_manager/consensus_manager.rs
@@ -27,7 +27,7 @@ use super::errors::StorageError;
 use super::CollectionContainer;
 use crate::content_manager::consensus::consensus_wal::ConsensusOpWal;
 use crate::content_manager::consensus::entry_queue::EntryId;
-use crate::content_manager::consensus::is_ready::IsReady;
+use collection::common::is_ready::IsReady;
 use crate::content_manager::consensus::operation_sender::OperationSender;
 use crate::content_manager::consensus::persistent::Persistent;
 use crate::types::{

diff --git a/lib/storage/src/content_manager/migration.rs b/lib/storage/src/content_manager/migration.rs
@@ -1,9 +1,20 @@
-use std::collections::HashMap;
+use std::sync::Arc;
+use std::time::Duration;
 
 use collection::collection::Collection;
-use collection::operations::types::{CollectionError, CollectionResult};
+use collection::operations::point_ops::{PointInsertOperations, PointOperations, PointStruct};
+use collection::operations::types::{CollectionError, CollectionResult, ScrollRequest};
+use collection::operations::CollectionUpdateOperations;
 use collection::shards::replica_set::ReplicaState;
 use collection::shards::shard::{PeerId, ShardId};
+use collection::shards::CollectionId;
+use segment::types::{WithPayloadInterface, WithVector};
+use tokio::sync::RwLock;
+
+use crate::content_manager::collections_ops::Collections;
+
+const MIGRATION_BATCH_SIZE: usize = 1000;
+const COLLECTION_INITIATION_TIMEOUT: Duration = Duration::from_secs(60);
 
 /// Handlers for migration data from one collection into another within single cluster
 
@@ -24,7 +35,7 @@ async fn get_local_source_shards(
         let responsible_shard_opt = shard_info
             .replicas
             .iter()
-            .filter(|(_, replica_state)| replica_state == ReplicaState::Active)
+            .filter(|(_, replica_state)| **replica_state == ReplicaState::Active)
             .max_by_key(|(peer_id, _)| *peer_id)
             .map(|(peer_id, _)| *peer_id);
 
@@ -46,14 +57,102 @@ async fn get_local_source_shards(
     Ok(local_responsible_shards)
 }
 
+fn handle_get_collection(collection: Option<&Collection>) -> CollectionResult<&Collection> {
+    match collection {
+        Some(collection) => Ok(collection),
+        None => Err(CollectionError::service_error(
+            "Collection is not found".to_string(),
+        )),
+    }
+}
+
+async fn migrate_shard(
+    collections: Arc<RwLock<Collections>>,
+    source_collection: &CollectionId,
+    target_collection: &CollectionId,
+    shard_id: ShardId,
+) -> CollectionResult<()> {
+    let mut offset = None;
+    let limit = MIGRATION_BATCH_SIZE;
+
+    loop {
+        let request = ScrollRequest {
+            offset,
+            limit: Some(limit),
+            filter: None,
+            with_payload: Some(WithPayloadInterface::Bool(true)),
+            with_vector: WithVector::Bool(true),
+        };
+
+        let scroll_result = {
+            let collections_read = collections.read().await;
+            let collection = handle_get_collection(collections_read.get(source_collection))?;
+            collection.scroll_by(request, Some(shard_id)).await?
+        };
+
+        offset = scroll_result.next_page_offset;
+
+        if offset.is_none() {
+            break;
+        }
+
+        let records = scroll_result
+            .points
+            .into_iter()
+            .map(|point| PointStruct {
+                id: point.id,
+                vector: point.vector.unwrap(),
+                payload: point.payload,
+            })
+            .collect();
+
+        let upsert_request = CollectionUpdateOperations::PointOperation(
+            PointOperations::UpsertPoints(PointInsertOperations::PointsList(records)),
+        );
+
+        let collections_read = collections.read().await;
+        let collection = handle_get_collection(collections_read.get(target_collection))?;
+
+        collection.update_from_client(upsert_request, false).await?;
+    }
+    Ok(())
+}
+
 /// Spawns a task which will retrieve data from appropriate local shards of the `source` collection
 /// into target collection.
 pub async fn migrate(
-    source: &Collection,
-    target: &Collection,
+    collections: Arc<RwLock<Collections>>,
+    source_collection: CollectionId,
+    target_collection: CollectionId,
     this_peer_id: PeerId,
 ) -> CollectionResult<()> {
-    let local_responsible_shards = get_local_source_shards(source, this_peer_id).await?;
+    let collections_read = collections.read().await;
+    let collection = handle_get_collection(collections_read.get(&source_collection))?;
+    let local_responsible_shards = get_local_source_shards(collection, this_peer_id).await?;
+
+    // Wait for all shards to be active
+    {
+        let collections_read = collections.read().await;
+        let collection = handle_get_collection(collections_read.get(&target_collection))?;
+        let is_initialized = collection.wait_collection_initiated(COLLECTION_INITIATION_TIMEOUT);
+        if !is_initialized {
+            return Err(CollectionError::service_error(format!(
+                "Collection {} was not initialized within {} sec timeout",
+                target_collection,
+                COLLECTION_INITIATION_TIMEOUT.as_secs()
+            )));
+        }
+    }
+
+    for shard_id in local_responsible_shards {
+        migrate_shard(
+            collections.clone(),
+            &source_collection,
+            &target_collection,
+            shard_id,
+        )
+        .await?;
+    }
 
     Ok(())
 }
diff --git a/lib/storage/src/content_manager/toc.rs b/lib/storage/src/content_manager/toc.rs
@@ -48,6 +48,7 @@ use crate::content_manager::collection_meta_ops::{
 use crate::content_manager::collections_ops::{Checker, Collections};
 use crate::content_manager::consensus::operation_sender::OperationSender;
 use crate::content_manager::errors::StorageError;
+use crate::content_manager::migration::migrate;
 use crate::content_manager::shard_distribution::ShardDistributionProposal;
 use crate::types::{PeerAddressById, StorageConfig};
 use crate::ConsensusOperations;
@@ -258,7 +259,7 @@ impl TableOfContent {
             optimizers_config: optimizers_config_diff,
             replication_factor,
             write_consistency_factor,
-            init_from: _init_from,
+            init_from,
         } = operation;
 
         self.collections
@@ -356,9 +357,26 @@ impl TableOfContent {
             self.on_peer_created(collection_name.to_string(), self.this_peer_id, shard_id)
                 .await?;
         }
+
+        if let Some(init_from) = init_from {
+            self.run_migration(init_from.collection, collection_name.to_string()).await;
+        }
+
         Ok(true)
     }
 
+    pub async fn run_migration(&self, from_collection: CollectionId, to_collection: CollectionId) {
+        let collections = self.collections.clone();
+        let this_peer_id = self.this_peer_id;
+        // There should be no search requests during initial update, so we can search runtime.
+        self.search_runtime.spawn(async move {
+            match migrate(collections, from_collection, to_collection, this_peer_id).await {
+                Ok(_) => log::info!("Migration completed"),
+                Err(err) => log::error!("Migration failed: {}", err),
+            }
+        });
+    }
+
     async fn on_peer_created(
         &self,
         collection_name: String,

diff --git a/lib/storage/tests/alias_tests.rs b/lib/storage/tests/alias_tests.rs
@@ -84,6 +84,7 @@ mod tests {
                             on_disk_payload: None,
                             replication_factor: None,
                             write_consistency_factor: None,
+                            init_from: None,
                         },
                     )),
                     None,

diff --git a/src/common/helpers.rs b/src/common/helpers.rs
@@ -37,7 +37,7 @@ mod tests {
     use std::thread::sleep;
     use std::time::Duration;
 
-    use storage::content_manager::consensus::is_ready::IsReady;
+    use collection::common::is_ready::IsReady;
 
     #[test]
     fn test_is_ready() {

diff --git a/src/consensus.rs b/src/consensus.rs
@@ -921,6 +921,7 @@ mod tests {
                             on_disk_payload: None,
                             replication_factor: None,
                             write_consistency_factor: None,
+                            init_from: None,
                         },
                     )),
                     None,