Split operations into update_all/update_existing during resharding (#4928)

Co-authored-by: generall <andrey@vasnetsov.com>

lib/collection/src/collection/mod.rs

@@ -419,18 +419,17 @@ impl Collection {
 
         // Abort resharding, if resharding shard is marked as `Dead`.
         //
-        // This branch should only be triggered, if resharding is currently at
-        // `ReshardStage::MigratingPoints` stage, because resharding shard should be marked as
-        // `Active` when all resharding transfers are successfully completed, and so the check
-        // *right above* this one should be triggered.
+        // This branch should only be triggered, if resharding is currently at `MigratingPoints`
+        // stage, because target shard should be marked as `Active`, when all resharding transfers
+        // are successfully completed, and so the check *right above* this one would be triggered.
         //
-        // So, if resharding reached `ReshardingStage::ReadHashRingCommitted`, this branch *won't*
-        // be triggered, and in this case, resharding *won't* be cancelled. Though, the update
-        // request should *fail* with "failed to update all replicas of a shard" error.
+        // So, if resharding reached `ReadHashRingCommitted`, this branch *won't* be triggered,
+        // and resharding *won't* be cancelled. The update request should *fail* with "failed to
+        // update all replicas of a shard" error.
         //
-        // If resharding reached `ReshardingStage::WriteHashRingCommitted`, and this branch is
-        // triggered *somehow*, then `Collection::abort_resharding` call should return an error,
-        // so no special handling is needed for `ReshardingStage::WriteHashRingCommitted`.
+        // If resharding reached `ReadHashRingCommitted`, and this branch is triggered *somehow*,
+        // then `Collection::abort_resharding` call should return an error, so no special handling
+        // is needed.
         if current_state == Some(ReplicaState::Resharding) && state == ReplicaState::Dead {
             let shard_key = shard_holder
                 .get_shard_id_to_key_mapping()

lib/collection/src/collection/point_ops.rs

@@ -108,7 +108,7 @@ impl Collection {
                     }
 
                     shard
-                        .update_with_consistency(operation.operation, wait, ordering)
+                        .update_with_consistency(operation.operation, wait, ordering, false)
                         .await
                         .map(Some)
                 }
@@ -145,13 +145,42 @@ impl Collection {
         let mut results = tokio::task::spawn(async move {
             let _update_lock = update_lock;
 
-            let updates: FuturesUnordered<_> = shard_holder
-                .split_by_shard(operation, &shard_keys_selection)?
-                .into_iter()
-                .map(move |(shard, operation)| {
-                    shard.update_with_consistency(operation, wait, ordering)
-                })
-                .collect();
+            let updates = FuturesUnordered::new();
+            let operations = shard_holder.split_by_shard(operation, &shard_keys_selection)?;
+
+            for (shard, operation) in operations {
+                let operation = shard_holder.split_by_mode(shard.shard_id, operation);
+
+                updates.push(async move {
+                    let mut result = UpdateResult {
+                        operation_id: None,
+                        status: UpdateStatus::Acknowledged,
+                        clock_tag: None,
+                    };
+
+                    for operation in operation.update_all {
+                        result = shard
+                            .update_with_consistency(operation, wait, ordering, false)
+                            .await?;
+                    }
+
+                    for operation in operation.update_only_existing {
+                        let res = shard
+                            .update_with_consistency(operation, wait, ordering, true)
+                            .await;
+
+                        if let Err(err) = &res {
+                            if err.is_missing_point() {
+                                continue;
+                            }
+                        }
+
+                        result = res?;
+                    }
+
+                    CollectionResult::Ok(result)
+                });
+            }
 
             let results: Vec<_> = updates.collect().await;
 

lib/collection/src/operations/mod.rs

@@ -21,7 +21,7 @@ pub mod verification;
 use std::collections::HashMap;
 
 use segment::json_path::JsonPath;
-use segment::types::{ExtendedPointId, PayloadFieldSchema};
+use segment::types::{ExtendedPointId, PayloadFieldSchema, PointIdType};
 use serde::{Deserialize, Serialize};
 use strum::{EnumDiscriminants, EnumIter};
 use validator::Validate;
@@ -146,6 +146,43 @@ pub enum CollectionUpdateOperations {
     FieldIndexOperation(FieldIndexOperations),
 }
 
+impl CollectionUpdateOperations {
+    pub fn is_upsert_points(&self) -> bool {
+        matches!(
+            self,
+            Self::PointOperation(point_ops::PointOperations::UpsertPoints(_))
+        )
+    }
+
+    pub fn is_delete_points(&self) -> bool {
+        matches!(
+            self,
+            Self::PointOperation(point_ops::PointOperations::DeletePoints { .. })
+        )
+    }
+
+    pub fn point_ids(&self) -> Vec<PointIdType> {
+        match self {
+            Self::PointOperation(op) => op.point_ids(),
+            Self::VectorOperation(op) => op.point_ids(),
+            Self::PayloadOperation(op) => op.point_ids(),
+            Self::FieldIndexOperation(_) => Vec::new(),
+        }
+    }
+
+    pub fn retain_point_ids<F>(&mut self, filter: F)
+    where
+        F: Fn(&PointIdType) -> bool,
+    {
+        match self {
+            Self::PointOperation(op) => op.retain_point_ids(filter),
+            Self::VectorOperation(op) => op.retain_point_ids(filter),
+            Self::PayloadOperation(op) => op.retain_point_ids(filter),
+            Self::FieldIndexOperation(_) => (),
+        }
+    }
+}
+
 /// A mapping of operation to shard.
 /// Is a result of splitting one operation into several shards by corresponding PointIds
 pub enum OperationToShard<O> {

lib/collection/src/operations/payload_ops.rs

@@ -160,6 +160,38 @@ impl PayloadOps {
             PayloadOps::OverwritePayload(_) => true,
         }
     }
+
+    pub fn point_ids(&self) -> Vec<PointIdType> {
+        match self {
+            Self::SetPayload(op) => op.points.clone().unwrap_or(Vec::new()),
+            Self::DeletePayload(op) => op.points.clone().unwrap_or(Vec::new()),
+            Self::ClearPayload { points } => points.clone(),
+            Self::ClearPayloadByFilter(_) => Vec::new(),
+            Self::OverwritePayload(op) => op.points.clone().unwrap_or(Vec::new()),
+        }
+    }
+
+    pub fn retain_point_ids<F>(&mut self, filter: F)
+    where
+        F: Fn(&PointIdType) -> bool,
+    {
+        match self {
+            Self::SetPayload(op) => retain_opt(op.points.as_mut(), filter),
+            Self::DeletePayload(op) => retain_opt(op.points.as_mut(), filter),
+            Self::ClearPayload { points } => points.retain(filter),
+            Self::ClearPayloadByFilter(_) => (),
+            Self::OverwritePayload(op) => retain_opt(op.points.as_mut(), filter),
+        }
+    }
+}
+
+fn retain_opt<T, F>(vec: Option<&mut Vec<T>>, filter: F)
+where
+    F: Fn(&T) -> bool,
+{
+    if let Some(vec) = vec {
+        vec.retain(filter);
+    }
 }
 
 impl Validate for PayloadOps {

lib/collection/src/operations/point_ops.rs

@@ -1,5 +1,6 @@
 use std::borrow::Cow;
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
+use std::iter;
 
 use api::rest::{BatchVectorStruct, ShardKeySelector, VectorStruct};
 use itertools::izip;
@@ -12,9 +13,15 @@ use serde::{Deserialize, Serialize};
 use strum::{EnumDiscriminants, EnumIter};
 use validator::Validate;
 
-use super::{point_to_shards, split_iter_by_shard, OperationToShard, SplitByShard};
+use super::payload_ops::SetPayloadOp;
+use super::vector_ops::{PointVectors, UpdateVectorsOp};
+use super::{
+    point_to_shards, split_iter_by_shard, CollectionUpdateOperations, OperationToShard,
+    SplitByShard,
+};
 use crate::hash_ring::HashRingRouter;
 use crate::operations::types::Record;
+use crate::operations::{payload_ops, vector_ops};
 use crate::shards::shard::ShardId;
 
 /// Defines write ordering guarantees for collection operations
@@ -216,6 +223,200 @@ pub enum PointInsertOperationsInternal {
     PointsList(Vec<PointStruct>),
 }
 
+impl PointInsertOperationsInternal {
+    pub fn point_ids(&self) -> Vec<PointIdType> {
+        match self {
+            Self::PointsBatch(batch) => batch.ids.clone(),
+            Self::PointsList(points) => points.iter().map(|point| point.id).collect(),
+        }
+    }
+
+    pub fn retain_point_ids<F>(&mut self, filter: F)
+    where
+        F: Fn(&PointIdType) -> bool,
+    {
+        match self {
+            Self::PointsBatch(batch) => {
+                let mut retain_indices = HashSet::new();
+
+                retain_with_index(&mut batch.ids, |index, id| {
+                    if filter(id) {
+                        retain_indices.insert(index);
+                        true
+                    } else {
+                        false
+                    }
+                });
+
+                match &mut batch.vectors {
+                    BatchVectorStruct::Single(vectors) => {
+                        retain_with_index(vectors, |index, _| retain_indices.contains(&index));
+                    }
+
+                    BatchVectorStruct::MultiDense(vectors) => {
+                        retain_with_index(vectors, |index, _| retain_indices.contains(&index));
+                    }
+
+                    BatchVectorStruct::Named(vectors) => {
+                        for (_, vectors) in vectors.iter_mut() {
+                            retain_with_index(vectors, |index, _| retain_indices.contains(&index));
+                        }
+                    }
+
+                    BatchVectorStruct::Document(documents) => {
+                        retain_with_index(documents, |index, _| retain_indices.contains(&index));
+                    }
+                }
+
+                if let Some(payload) = &mut batch.payloads {
+                    retain_with_index(payload, |index, _| retain_indices.contains(&index));
+                }
+            }
+
+            Self::PointsList(points) => points.retain(|point| filter(&point.id)),
+        }
+    }
+
+    pub fn into_update_only(self) -> Vec<CollectionUpdateOperations> {
+        let mut operations = Vec::new();
+
+        match self {
+            Self::PointsBatch(batch) => {
+                let mut update_vectors = UpdateVectorsOp { points: Vec::new() };
+
+                match batch.vectors {
+                    BatchVectorStruct::Single(vectors) => {
+                        let ids = batch.ids.iter().copied();
+                        let vectors = vectors.into_iter().map(VectorStruct::Single);
+
+                        update_vectors.points = ids
+                            .zip(vectors)
+                            .map(|(id, vector)| PointVectors { id, vector })
+                            .collect();
+                    }
+
+                    BatchVectorStruct::MultiDense(vectors) => {
+                        let ids = batch.ids.iter().copied();
+                        let vectors = vectors.into_iter().map(VectorStruct::MultiDense);
+
+                        update_vectors.points = ids
+                            .zip(vectors)
+                            .map(|(id, vector)| PointVectors { id, vector })
+                            .collect();
+                    }
+
+                    BatchVectorStruct::Named(batch_vectors) => {
+                        let ids = batch.ids.iter().copied();
+
+                        let mut batch_vectors: HashMap<_, _> = batch_vectors
+                            .into_iter()
+                            .map(|(name, vectors)| (name, vectors.into_iter()))
+                            .collect();
+
+                        let vectors = iter::repeat(()).filter_map(move |_| {
+                            let mut point_vectors =
+                                HashMap::with_capacity(batch_vectors.capacity());
+
+                            for (vector_name, vectors) in batch_vectors.iter_mut() {
+                                point_vectors.insert(vector_name.clone(), vectors.next()?);
+                            }
+
+                            Some(VectorStruct::Named(point_vectors))
+                        });
+
+                        update_vectors.points = ids
+                            .zip(vectors)
+                            .map(|(id, vector)| PointVectors { id, vector })
+                            .collect();
+                    }
+
+                    BatchVectorStruct::Document(documents) => {
+                        let ids = batch.ids.iter().copied();
+                        let documents = documents.into_iter().map(VectorStruct::Document);
+
+                        update_vectors.points = ids
+                            .zip(documents)
+                            .map(|(id, vector)| PointVectors { id, vector })
+                            .collect();
+                    }
+                }
+
+                let update_vectors = vector_ops::VectorOperations::UpdateVectors(update_vectors);
+                let update_vectors = CollectionUpdateOperations::VectorOperation(update_vectors);
+
+                operations.push(update_vectors);
+
+                if let Some(payloads) = batch.payloads {
+                    let ids = batch.ids.iter().copied();
+
+                    for (id, payload) in ids.zip(payloads) {
+                        if let Some(payload) = payload {
+                            let set_payload = SetPayloadOp {
+                                points: Some(vec![id]),
+                                payload,
+                                filter: None,
+                                key: None,
+                            };
+
+                            let set_payload =
+                                payload_ops::PayloadOps::OverwritePayload(set_payload);
+                            let set_payload =
+                                CollectionUpdateOperations::PayloadOperation(set_payload);
+
+                            operations.push(set_payload);
+                        }
+                    }
+                }
+            }
+
+            Self::PointsList(points) => {
+                let mut update_vectors = UpdateVectorsOp { points: Vec::new() };
+
+                for point in points {
+                    update_vectors.points.push(PointVectors {
+                        id: point.id,
+                        vector: point.vector,
+                    });
+
+                    if let Some(payload) = point.payload {
+                        let set_payload = SetPayloadOp {
+                            points: Some(vec![point.id]),
+                            payload,
+                            filter: None,
+                            key: None,
+                        };
+
+                        let set_payload = payload_ops::PayloadOps::OverwritePayload(set_payload);
+                        let set_payload = CollectionUpdateOperations::PayloadOperation(set_payload);
+
+                        operations.push(set_payload);
+                    }
+                }
+
+                let update_vectors = vector_ops::VectorOperations::UpdateVectors(update_vectors);
+                let update_vectors = CollectionUpdateOperations::VectorOperation(update_vectors);
+
+                operations.insert(0, update_vectors);
+            }
+        }
+
+        operations
+    }
+}
+
+fn retain_with_index<T, F>(vec: &mut Vec<T>, mut filter: F)
+where
+    F: FnMut(usize, &T) -> bool,
+{
+    let mut index = 0;
+
+    vec.retain(|item| {
+        let retain = filter(index, item);
+        index += 1;
+        retain
+    });
+}
+
 impl Validate for PointInsertOperationsInternal {
     fn validate(&self) -> Result<(), validator::ValidationErrors> {
         match self {
@@ -355,6 +556,27 @@ impl PointOperations {
             PointOperations::SyncPoints(_) => true,
         }
     }
+
+    pub fn point_ids(&self) -> Vec<PointIdType> {
+        match self {
+            Self::UpsertPoints(op) => op.point_ids(),
+            Self::DeletePoints { ids } => ids.clone(),
+            Self::DeletePointsByFilter(_) => Vec::new(),
+            Self::SyncPoints(op) => op.points.iter().map(|point| point.id).collect(),
+        }
+    }
+
+    pub fn retain_point_ids<F>(&mut self, filter: F)
+    where
+        F: Fn(&PointIdType) -> bool,
+    {
+        match self {
+            Self::UpsertPoints(op) => op.retain_point_ids(filter),
+            Self::DeletePoints { ids } => ids.retain(filter),
+            Self::DeletePointsByFilter(_) => (),
+            Self::SyncPoints(op) => op.points.retain(|point| filter(&point.id)),
+        }
+    }
 }
 
 impl Validate for PointOperations {

lib/collection/src/operations/types.rs

@@ -1019,6 +1019,18 @@ impl CollectionError {
             Self::StrictMode { .. } => false,
         }
     }
+
+    pub fn is_pre_condition_failed(&self) -> bool {
+        matches!(self, Self::PreConditionFailed { .. })
+    }
+
+    pub fn is_missing_point(&self) -> bool {
+        match self {
+            CollectionError::NotFound { what } => what.contains("No point with id"),
+            CollectionError::PointNotFound { .. } => true,
+            _ => false,
+        }
+    }
 }
 
 impl From<SystemTimeError> for CollectionError {

lib/collection/src/operations/vector_ops.rs

@@ -88,6 +88,25 @@ impl VectorOperations {
             VectorOperations::DeleteVectorsByFilter(..) => false,
         }
     }
+
+    pub fn point_ids(&self) -> Vec<PointIdType> {
+        match self {
+            Self::UpdateVectors(op) => op.points.iter().map(|point| point.id).collect(),
+            Self::DeleteVectors(points, _) => points.points.clone(),
+            Self::DeleteVectorsByFilter(_, _) => Vec::new(),
+        }
+    }
+
+    pub fn retain_point_ids<F>(&mut self, filter: F)
+    where
+        F: Fn(&PointIdType) -> bool,
+    {
+        match self {
+            Self::UpdateVectors(op) => op.points.retain(|point| filter(&point.id)),
+            Self::DeleteVectors(points, _) => points.points.retain(filter),
+            Self::DeleteVectorsByFilter(_, _) => (),
+        }
+    }
 }
 
 impl Validate for VectorOperations {

lib/collection/src/shards/replica_set/update.rs

@@ -71,6 +71,7 @@ impl ShardReplicaSet {
         operation: CollectionUpdateOperations,
         wait: bool,
         ordering: WriteOrdering,
+        update_only_existing: bool,
     ) -> CollectionResult<UpdateResult> {
         // `ShardReplicaSet::update` is not cancel safe, so this method is not cancel safe.
 
@@ -91,7 +92,7 @@ impl ShardReplicaSet {
                 WriteOrdering::Weak => None,
             };
 
-            self.update(operation, wait).await
+            self.update(operation, wait, update_only_existing).await
         } else {
             // Forward the update to the designated leader
             self.forward_update(leader_peer, operation, wait, ordering)
@@ -128,7 +129,7 @@ impl ShardReplicaSet {
 
         peer_ids
             .into_iter()
-            .filter(|peer_id| self.peer_is_active(peer_id)) // re-acquire replica_state read lock
+            .filter(|peer_id| self.peer_is_active_or_resharding(peer_id)) // re-acquire replica_state read lock
             .max()
     }
 
@@ -143,6 +144,7 @@ impl ShardReplicaSet {
         &self,
         operation: CollectionUpdateOperations,
         wait: bool,
+        update_only_existing: bool,
     ) -> CollectionResult<UpdateResult> {
         // `ShardRepilcaSet::update_impl` is not cancel safe, so this method is not cancel safe.
 
@@ -158,7 +160,7 @@ impl ShardReplicaSet {
             let is_non_zero_tick = clock.current_tick().is_some();
 
             let res = self
-                .update_impl(operation.clone(), wait, &mut clock)
+                .update_impl(operation.clone(), wait, &mut clock, update_only_existing)
                 .await?;
 
             if let Some(res) = res {
@@ -190,6 +192,7 @@ impl ShardReplicaSet {
         operation: CollectionUpdateOperations,
         wait: bool,
         clock: &mut clock_set::ClockGuard,
+        update_only_existing: bool,
     ) -> CollectionResult<Option<UpdateResult>> {
         // `LocalShard::update` is not guaranteed to be cancel safe and it's impossible to cancel
         // multiple parallel updates in a way that is *guaranteed* not to introduce inconsistencies
@@ -330,8 +333,11 @@ impl ShardReplicaSet {
         };
 
         if successes.len() >= minimal_success_count {
-            let wait_for_deactivation =
-                self.handle_failed_replicas(&failures, &self.replica_state.read());
+            let wait_for_deactivation = self.handle_failed_replicas(
+                &failures,
+                &self.replica_state.read(),
+                update_only_existing,
+            );
 
             // report all failing peers to consensus
             if wait && wait_for_deactivation && !failures.is_empty() {
@@ -373,6 +379,7 @@ impl ShardReplicaSet {
                     .iter()
                     .filter(|(peer_id, _)| self.peer_is_resharding(peer_id)),
                 &self.replica_state.read(),
+                update_only_existing,
             );
 
             // completely failed - report error to user
@@ -431,6 +438,7 @@ impl ShardReplicaSet {
         &self,
         failures: impl IntoIterator<Item = &'a (PeerId, CollectionError)>,
         state: &ReplicaSetState,
+        update_only_existing: bool,
     ) -> bool {
         let mut wait_for_deactivation = false;
 
@@ -450,14 +458,18 @@ impl ShardReplicaSet {
                 _ => continue,
             }
 
-            if peer_state == ReplicaState::Partial
-                && matches!(err, CollectionError::PreConditionFailed { .. })
+            if matches!(peer_state, ReplicaState::Partial | ReplicaState::Resharding)
+                && err.is_pre_condition_failed()
             {
                 // Handles a special case where transfer receiver haven't created a shard yet.
                 // In this case update should be handled by source shard and forward proxy.
                 continue;
             }
 
+            if update_only_existing && err.is_missing_point() {
+                continue;
+            }
+
             if err.is_transient() || peer_state == ReplicaState::Initializing {
                 // If the error is transient, we should not deactivate the peer
                 // before allowing other operations to continue.

lib/collection/src/shards/resharding/stage_migrate_points.rs

@@ -353,7 +353,7 @@ async fn drive_down(
             // Wait on all updates here, not just the last batch
             // If we don't wait on all updates it somehow results in inconsistent results
             target_replica_set
-                .update_with_consistency(operation, true, WriteOrdering::Weak)
+                .update_with_consistency(operation, true, WriteOrdering::Weak, false)
                 .await?;
 
             if offset.is_none() {

lib/collection/src/shards/resharding/stage_propagate_deletes.rs

@@ -95,7 +95,7 @@ pub(super) async fn drive(
             // Wait on all updates here, not just the last batch
             // If we don't wait on all updates it somehow results in inconsistent deletes
             replica_set
-                .update_with_consistency(operation, true, WriteOrdering::Weak)
+                .update_with_consistency(operation, true, WriteOrdering::Weak, false)
                 .await?;
 
             if offset.is_none() {

lib/collection/src/shards/shard_holder/resharding.rs

@@ -1,18 +1,23 @@
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::fmt;
 use std::ops::Deref as _;
 
-use segment::types::ShardKey;
+use segment::types::{CustomIdCheckerCondition as _, ShardKey};
 
 use super::ShardHolder;
 use crate::hash_ring::{self, HashRingRouter};
 use crate::operations::cluster_ops::ReshardingDirection;
 use crate::operations::types::{CollectionError, CollectionResult};
+use crate::operations::{point_ops, CollectionUpdateOperations};
 use crate::shards::replica_set::{ReplicaState, ShardReplicaSet};
 use crate::shards::resharding::{ReshardKey, ReshardStage, ReshardState};
 use crate::shards::shard::ShardId;
 
 impl ShardHolder {
+    pub fn resharding_state(&self) -> Option<ReshardState> {
+        self.resharding_state.read().clone()
+    }
+
     pub fn check_start_resharding(&mut self, resharding_key: &ReshardKey) -> CollectionResult<()> {
         let ReshardKey {
             direction,
@@ -364,6 +369,117 @@ impl ShardHolder {
         Ok(())
     }
 
+    /// Split collection update operation by "update mode":
+    /// - update all:
+    ///   - "regular" operation
+    ///   - `upsert` inserts new points and updates existing ones
+    ///   - other update operations return error, if a point does not exist in collection
+    /// - update existing:
+    ///   - `upsert` does *not* insert new points, only updates existing ones
+    ///   - other update operations ignore points that do not exist in collection
+    ///
+    /// Depends on the current resharding state. If resharding is not active operations are not split.
+    pub fn split_by_mode(
+        &self,
+        shard_id: ShardId,
+        operation: CollectionUpdateOperations,
+    ) -> OperationsByMode {
+        let Some(state) = self.resharding_state() else {
+            return OperationsByMode::from(operation);
+        };
+
+        // Resharding *UP*
+        // ┌────────────┐   ┌──────────┐
+        // │            │   │          │
+        // │ Shard 1    │   │ Shard 2  │
+        // │ Non-Target ├──►│ Target   │
+        // │ Sender     │   │ Receiver │
+        // │            │   │          │
+        // └────────────┘   └──────────┘
+        //
+        // Resharding *DOWN*
+        // ┌────────────┐   ┌──────────┐
+        // │            │   │          │
+        // │ Shard 1    │   │ Shard 2  │
+        // │ Non-Target │◄──┤ Target   │
+        // │ Receiver   │   │ Sender   │
+        // │            │   │          │
+        // └────────────┘   └──────────┘
+
+        // Target shard of the resharding operation. This is the shard that:
+        //
+        // - *created* during resharding *up*
+        // - *deleted* during resharding *down*
+        let is_target_shard = shard_id == state.shard_id;
+
+        // Shard that will be *receiving* migrated points during resharding:
+        //
+        // - *target* shard during resharding *up*
+        // - *non* target shards during resharding *down*
+        let is_receiver_shard = match state.direction {
+            ReshardingDirection::Up => is_target_shard,
+            ReshardingDirection::Down => !is_target_shard,
+        };
+
+        // Shard that will be *sending* migrated points during resharding:
+        //
+        // - *non* target shards during resharding *up*
+        // - *target* shard during resharding *down*
+        let is_sender_shard = !is_receiver_shard;
+
+        // We split update operations:
+        //
+        // - on *receiver* shards during `MigratingPoints` stage (for all operations except `upsert`)
+        // - and on *sender* shards during `ReadHashRingCommitted` stage when resharding *up*
+
+        let should_split_receiver = is_receiver_shard
+            && state.stage == ReshardStage::MigratingPoints
+            && !operation.is_upsert_points();
+
+        let should_split_sender = is_sender_shard
+            && state.stage >= ReshardStage::ReadHashRingCommitted
+            && state.direction == ReshardingDirection::Up;
+
+        if !should_split_receiver && !should_split_sender {
+            return OperationsByMode::from(operation);
+        }
+
+        // There's no point splitting delete operations
+        if operation.is_delete_points() {
+            return OperationsByMode::from(operation);
+        }
+
+        let Some(filter) = self.resharding_filter() else {
+            return OperationsByMode::from(operation);
+        };
+
+        let point_ids = operation.point_ids();
+
+        if point_ids.is_empty() {
+            return OperationsByMode::from(operation);
+        }
+
+        let target_point_ids: HashSet<_> = point_ids
+            .iter()
+            .copied()
+            .filter(|&point_id| filter.check(point_id))
+            .collect();
+
+        if target_point_ids.is_empty() {
+            OperationsByMode::from(operation)
+        } else if target_point_ids.len() == point_ids.len() {
+            OperationsByMode::default().with_update_only_existing(operation)
+        } else {
+            let mut update_all = operation.clone();
+            update_all.retain_point_ids(|point_id| !target_point_ids.contains(point_id));
+
+            let mut update_only_existing = operation;
+            update_only_existing.retain_point_ids(|point_id| target_point_ids.contains(point_id));
+
+            OperationsByMode::from(update_all).with_update_only_existing(update_only_existing)
+        }
+    }
+
     pub fn resharding_filter(&self) -> Option<hash_ring::HashRingFilter> {
         let shard_id = self.resharding_state.read().as_ref()?.shard_id;
         self.hash_ring_filter(shard_id)
@@ -391,6 +507,39 @@ impl ShardHolder {
     }
 }
 
+#[derive(Clone, Debug, Default)]
+pub struct OperationsByMode {
+    pub update_all: Vec<CollectionUpdateOperations>,
+    pub update_only_existing: Vec<CollectionUpdateOperations>,
+}
+
+impl OperationsByMode {
+    pub fn with_update_only_existing(mut self, operation: CollectionUpdateOperations) -> Self {
+        match operation {
+            CollectionUpdateOperations::PointOperation(
+                point_ops::PointOperations::UpsertPoints(operation),
+            ) => {
+                self.update_only_existing = operation.into_update_only();
+            }
+
+            operation => {
+                self.update_only_existing = vec![operation];
+            }
+        }
+
+        self
+    }
+}
+
+impl From<CollectionUpdateOperations> for OperationsByMode {
+    fn from(operation: CollectionUpdateOperations) -> Self {
+        Self {
+            update_all: vec![operation],
+            update_only_existing: Vec::new(),
+        }
+    }
+}
+
 fn get_ring<'a>(
     rings: &'a mut HashMap<Option<ShardKey>, HashRingRouter>,
     shard_key: &'_ Option<ShardKey>,

src/main.rs

@@ -348,9 +348,11 @@ fn main() -> anyhow::Result<()> {
             }
         });
 
-        runtime_handle.block_on(async {
-            toc_arc.resume_resharding_tasks().await;
-        });
+        // TODO(resharding): Remove resharding driver?
+        //
+        // runtime_handle.block_on(async {
+        //     toc_arc.resume_resharding_tasks().await;
+        // });
 
         let collections_to_recover_in_consensus = if is_new_deployment {
             let existing_collections =