Skip to main content

qdrant_edge/edge/config/
shard.rs

1//! Edge shard configuration: user-facing params and conversion to/from SegmentConfig.
2
3use std::collections::HashMap;
4use std::path::Path;
5
6use crate::common::fs::{atomic_save_json, read_json};
7use crate::segment::common::operation_error::{OperationError, OperationResult};
8use crate::segment::types::{
9    HnswConfig, PayloadStorageType, QuantizationConfig, SegmentConfig, VectorNameBuf,
10};
11use serde::{Deserialize, Serialize};
12use crate::shard::operations::optimization::OptimizerThresholds;
13
14use super::optimizers::EdgeOptimizersConfig;
15use super::vectors::{EdgeSparseVectorParams, EdgeVectorParams};
16
17/// File name for the persisted edge shard config.
18pub(crate) const EDGE_CONFIG_FILE: &str = "edge_config.json";
19
20/// Full configuration for an edge shard.
21#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
22#[serde(rename_all = "snake_case")]
23pub struct EdgeConfig {
24    /// If true, payload is stored on disk (mmap); otherwise in RAM. Same as `CollectionParams::on_disk_payload`.
25    #[serde(default = "default_on_disk_payload")]
26    pub on_disk_payload: bool,
27    /// Dense vector params per vector name.
28    #[serde(default)]
29    pub vectors: HashMap<VectorNameBuf, EdgeVectorParams>,
30    /// Sparse vector params per vector name.
31    #[serde(default)]
32    pub sparse_vectors: HashMap<VectorNameBuf, EdgeSparseVectorParams>,
33    /// Global HNSW config; per-vector override is in `vectors[].hnsw_config`
34    #[serde(default)]
35    pub hnsw_config: HnswConfig,
36    /// Global quantization config for all vectors
37    /// Per-vector override in in `vectors[].quantization_config`
38    #[serde(default, skip_serializing_if = "Option::is_none")]
39    pub quantization_config: Option<QuantizationConfig>,
40    #[serde(default)]
41    pub optimizers: EdgeOptimizersConfig,
42}
43
44fn default_on_disk_payload() -> bool {
45    true
46}
47
48impl Default for EdgeConfig {
49    fn default() -> Self {
50        Self {
51            on_disk_payload: default_on_disk_payload(),
52            vectors: HashMap::new(),
53            sparse_vectors: HashMap::new(),
54            hnsw_config: HnswConfig::default(),
55            quantization_config: None,
56            optimizers: EdgeOptimizersConfig::default(),
57        }
58    }
59}
60
61impl EdgeConfig {
62    /// Build from existing segment config. Fills all parameters that can be inferred.
63    pub fn from_segment_config(segment: &SegmentConfig) -> Self {
64        let SegmentConfig {
65            vector_data,
66            sparse_vector_data,
67            payload_storage_type,
68        } = segment;
69
70        let vectors = vector_data
71            .iter()
72            .map(|(name, v)| (name.clone(), EdgeVectorParams::from_vector_data_config(v)))
73            .collect();
74
75        let sparse_vectors = sparse_vector_data
76            .iter()
77            .map(|(name, s)| {
78                (
79                    name.clone(),
80                    EdgeSparseVectorParams::from_sparse_vector_data_config(s),
81                )
82            })
83            .collect();
84
85        let on_disk_payload = payload_storage_type.is_on_disk();
86
87        // Infer global hnsw_config from per-vector HNSW configs when all agree
88        let hnsw_configs: Vec<HnswConfig> = vector_data
89            .values()
90            .filter_map(|v| match &v.index {
91                crate::segment::types::Indexes::Plain {} => None,
92                crate::segment::types::Indexes::Hnsw(h) => Some(*h),
93            })
94            .collect();
95        let hnsw_config = hnsw_configs
96            .first()
97            .and_then(|first| {
98                if hnsw_configs.iter().all(|h| h == first) {
99                    Some(*first)
100                } else {
101                    None
102                }
103            })
104            .unwrap_or_default();
105
106        Self {
107            on_disk_payload,
108            vectors,
109            sparse_vectors,
110            hnsw_config,
111            quantization_config: None,
112            optimizers: EdgeOptimizersConfig::default(),
113        }
114    }
115
116    /// Check compatibility with a segment config (e.g. loaded segment).
117    pub fn check_compatible_with_segment_config(
118        &self,
119        other: &SegmentConfig,
120    ) -> Result<(), String> {
121        self.plain_segment_config().check_compatible(other)
122    }
123
124    /// Segment config for creating appendable segments only.
125    /// Does not contain any HNSW configuration (plain index only).
126    pub fn plain_segment_config(&self) -> SegmentConfig {
127        let payload_storage_type = PayloadStorageType::from_on_disk_payload(self.on_disk_payload);
128        let vector_data = self
129            .vectors
130            .iter()
131            .map(|(name, p)| {
132                (
133                    name.clone(),
134                    p.to_plain_vector_data_config(self.quantization_config.as_ref()),
135                )
136            })
137            .collect();
138
139        let sparse_vector_data = self
140            .sparse_vectors
141            .iter()
142            .map(|(name, p)| (name.clone(), p.to_plain_sparse_vector_data_config()))
143            .collect();
144
145        SegmentConfig {
146            vector_data,
147            sparse_vector_data,
148            payload_storage_type,
149        }
150    }
151
152    /// Build segment optimizer config from this config (for blocking optimizers).
153    /// Use this instead of converting to SegmentConfig first.
154    pub fn segment_optimizer_config(&self) -> crate::shard::optimizers::config::SegmentOptimizerConfig {
155        use crate::shard::optimizers::config::SegmentOptimizerConfig;
156
157        let SegmentConfig {
158            vector_data: plain_dense_vector_config,
159            sparse_vector_data: plain_sparse_vector_config,
160            payload_storage_type,
161        } = self.plain_segment_config();
162
163        let dense_vector = self
164            .vectors
165            .iter()
166            .map(|(name, p)| {
167                (
168                    name.clone(),
169                    p.to_dense_vector_optimizer_config(
170                        &self.hnsw_config,
171                        self.quantization_config.as_ref(),
172                    ),
173                )
174            })
175            .collect();
176
177        let sparse_vector = self
178            .sparse_vectors
179            .iter()
180            .map(|(name, p)| (name.clone(), p.to_sparse_vector_optimizer_config()))
181            .collect();
182
183        SegmentOptimizerConfig {
184            payload_storage_type,
185            plain_dense_vector_config,
186            plain_sparse_vector_config,
187            dense_vector,
188            sparse_vector,
189        }
190    }
191
192    /// Return vector data config for a named vector (for read-only use, e.g. query).
193    /// Uses plain index; for optimizer/segment creation use segment_optimizer_config or plain_segment_config.
194    pub fn vector_data_config(
195        &self,
196        name: &VectorNameBuf,
197    ) -> Option<crate::segment::types::VectorDataConfig> {
198        self.vectors
199            .get(name)
200            .map(|p| p.to_plain_vector_data_config(self.quantization_config.as_ref()))
201    }
202
203    pub fn optimizer_thresholds(&self, num_indexing_threads: usize) -> OptimizerThresholds {
204        let indexing_threshold_kb = self.optimizers.get_indexing_threshold_kb();
205        OptimizerThresholds {
206            memmap_threshold_kb: usize::MAX,
207            indexing_threshold_kb,
208            max_segment_size_kb: self
209                .optimizers
210                .get_max_segment_size_kb(num_indexing_threads),
211            deferred_internal_id: None,
212        }
213    }
214
215    pub fn save(&self, path: &Path) -> OperationResult<()> {
216        let config_path = path.join(EDGE_CONFIG_FILE);
217        atomic_save_json(&config_path, self).map_err(|e| {
218            OperationError::service_error(format!(
219                "failed to write {}: {}",
220                config_path.display(),
221                e
222            ))
223        })
224    }
225
226    pub fn load(path: &Path) -> Option<OperationResult<Self>> {
227        let config_path = path.join(EDGE_CONFIG_FILE);
228        match fs_err::exists(&config_path) {
229            Ok(false) => return None,
230            Err(e) => return Some(Err(OperationError::from(e))),
231            Ok(true) => {}
232        }
233        Some(read_json(&config_path).map_err(OperationError::from))
234    }
235
236    pub fn set_hnsw_config(&mut self, hnsw_config: HnswConfig) {
237        self.hnsw_config = hnsw_config;
238    }
239
240    pub fn set_vector_hnsw_config(
241        &mut self,
242        vector_name: &str,
243        hnsw_config: HnswConfig,
244    ) -> OperationResult<()> {
245        let name = VectorNameBuf::from(vector_name);
246        let params = self
247            .vectors
248            .get_mut(&name)
249            .ok_or_else(|| OperationError::vector_name_not_exists(vector_name))?;
250        params.hnsw_config = Some(hnsw_config);
251        Ok(())
252    }
253
254    pub fn set_optimizers_config(&mut self, optimizers: EdgeOptimizersConfig) {
255        self.optimizers = optimizers;
256    }
257
258    pub fn optimizers_mut(&mut self) -> &mut EdgeOptimizersConfig {
259        &mut self.optimizers
260    }
261}