1use 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
17pub(crate) const EDGE_CONFIG_FILE: &str = "edge_config.json";
19
20#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
22#[serde(rename_all = "snake_case")]
23pub struct EdgeConfig {
24 #[serde(default = "default_on_disk_payload")]
26 pub on_disk_payload: bool,
27 #[serde(default)]
29 pub vectors: HashMap<VectorNameBuf, EdgeVectorParams>,
30 #[serde(default)]
32 pub sparse_vectors: HashMap<VectorNameBuf, EdgeSparseVectorParams>,
33 #[serde(default)]
35 pub hnsw_config: HnswConfig,
36 #[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 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 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 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 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 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 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}