Repository: qdrant/qdrant
File: lib/collection/src/shards/shard_holder/mod.rs (Rust)

Code before:
                        debug_assert!(false, "Shard id {shard_id} not found")
                    }
                }
            }
            ShardSelectorInternal::ShardId(shard_id) => {
                if let Some(replica_set) = self.shards.get(shard_id) {
                    res.push((replica_set, self.shard_id_to_key_mapping.get(shard_id)));
                } else {
                    return Err(shard_not_found_error(*shard_id));
                }
            }
        }
        Ok(res)
    }

    /// Common routing logic for reads and writes when using ShardKeyWithFallback
    ///
    /// If at least one of target shards is Active, use target shard. If not, redirect to fallback shard
    pub fn route_with_fallback<'a>(
        &self,
        key: &'a ShardKeyWithFallback,
    ) -> CollectionResult<(HashSet<ShardId>, &'a ShardKey)> {
        let shard_key_to_ids_mapping = self.get_shard_key_to_ids_mapping();

        let target_shard_ids = shard_key_to_ids_mapping.get(&key.target);
        let fallback_shard_ids = shard_key_to_ids_mapping.get(&key.fallback);

        if let Some(target_shard_ids) = target_shard_ids {
            let replicas = target_shard_ids
                .iter()
                .filter_map(|shard_id| self.shards.get(shard_id))
                .collect::<Vec<_>>();

            let target_shards_active = replicas
                .iter()
                .all(|replica_set| !replica_set.active_shards(false).is_empty());

            if !replicas.is_empty() && target_shards_active {
                // 1st condition is required to handle empty shard keys (2nd one returns true)
                Ok((target_shard_ids.clone(), &key.target))
            } else if let Some(fallback_shard_ids) = fallback_shard_ids {
                Ok((fallback_shard_ids.clone(), &key.fallback))
            } else {
                Err(CollectionError::shard_unavailable(format!(
                    "Neither target shard key {} nor fallback shard key {} have active replicas",
                    key.target, key.fallback
                )))
            }
        } else if let Some(fallback_shard_ids) = fallback_shard_ids {
            Ok((fallback_shard_ids.clone(), &key.fallback))
        } else {

Review comment (suggestion):
Maybe shortcut here for now? If either has `.len() != 1`, return early?

I think it's fine to leave the rest of the logic compatible with multiple shard IDs. We can use that once we reenable multiple shards.

Code after:
                    } else {
                        debug_assert!(false, "Shard id {shard_id} not found")
                    }
                }
            }
            ShardSelectorInternal::ShardId(shard_id) => {
                if let Some(replica_set) = self.shards.get(shard_id) {
                    res.push((replica_set, self.shard_id_to_key_mapping.get(shard_id)));
                } else {
                    return Err(shard_not_found_error(*shard_id));
                }
            }
        }
        Ok(res)
    }

    /// Common routing logic for reads when using ShardKeyWithFallback
    ///
    /// Example routing:
    ///
    /// request: {"target": "key1", "fallback": "default"}
    ///
    /// Situation 1:
    /// /// - key1 -> shard_ids {1, 2} (both active)
    /// Request is routed to shard_ids {1, 2} of target key1
    ///
    /// Situation 2:
    /// /// - key1 -> no shards found
    /// Request is routed to shard_ids of fallback key "default"
    ///
    /// Situation 3:
    /// /// - key1 -> shard_ids {1} and it is in Partial state (no active replicas)
    /// Request is routed to shard_ids of fallback key "default"
    ///
    /// Situation 4:
    /// /// - key1 -> shard_ids {1, 2} (shard 1 active, shard 2 partial)
    /// Request is routed to shard_ids of fallback key "default"
    ///
    ///
    /// If at least one of target shards is Active, use target shard. If not, redirect to fallback shard
    pub fn route_with_fallback_for_read<'a>(
        &self,
        key: &'a ShardKeyWithFallback,
    ) -> CollectionResult<(HashSet<ShardId>, &'a ShardKey)> {
        let mut shard_key_to_ids_mapping = self.get_shard_key_to_ids_mapping();

        let target_shard_ids = shard_key_to_ids_mapping.remove(&key.target);
        let fallback_shard_ids = shard_key_to_ids_mapping.remove(&key.fallback);

        if let Some(target_shard_ids) = target_shard_ids {
            let replicas = target_shard_ids