Skip to main content

qdrant_edge/edge/
mod.rs

1pub mod bm25_embed;
2mod builders;
3pub mod config;
4mod count;
5mod facet;
6mod info;
7mod optimize;
8mod query;
9mod reexports;
10mod retrieve;
11mod scroll;
12mod search;
13mod snapshots;
14mod types;
15pub use types::*;
16mod update;
17
18use std::path::{Path, PathBuf};
19use std::sync::Arc;
20use std::sync::atomic::AtomicBool;
21use std::time::Duration;
22
23pub use builders::{EdgeConfigBuilder, EdgeSparseVectorParamsBuilder, EdgeVectorParamsBuilder};
24use crate::common::save_on_disk::SaveOnDisk;
25pub use config::optimizers::EdgeOptimizersConfig;
26pub use config::shard::EdgeConfig;
27pub use config::vectors::{EdgeSparseVectorParams, EdgeVectorParams};
28use fs_err as fs;
29pub use info::ShardInfo;
30use parking_lot::Mutex;
31pub use reexports::*;
32use crate::segment::entry::ReadSegmentEntry as _;
33use crate::segment::segment_constructor::{load_segment, normalize_segment_dir};
34use crate::shard::files::{PAYLOAD_INDEX_CONFIG_FILE, SEGMENTS_PATH};
35use crate::shard::operations::CollectionUpdateOperations;
36use crate::shard::segment_holder::SegmentHolder;
37use crate::shard::segment_holder::locked::LockedSegmentHolder;
38use crate::shard::wal::SerdeWal;
39
40use crate::edge::config::shard::EDGE_CONFIG_FILE;
41
42#[derive(Debug)]
43pub struct EdgeShard {
44    path: PathBuf,
45    config: SaveOnDisk<EdgeConfig>,
46    wal: Mutex<SerdeWal<CollectionUpdateOperations>>,
47    segments: LockedSegmentHolder,
48}
49
50const WAL_PATH: &str = "wal";
51impl EdgeShard {
52    /// Create a new edge shard at `path` with the given configuration.
53    ///
54    /// Fails if the shard already exists (i.e. the segments directory contains any segment).
55    /// Configuration is required and is persisted to `edge_config.json`. WAL
56    /// behavior follows `config.wal_options` (defaults to 32 MiB segments
57    /// when unset).
58    pub fn new(path: &Path, config: EdgeConfig) -> OperationResult<Self> {
59        if has_existing_segments(path) {
60            return Err(OperationError::service_error(
61                "cannot create edge shard: path already contains segment data",
62            ));
63        }
64
65        let wal_options = config.wal_options.clone().unwrap_or_default();
66        let (wal, segments_path) = ensure_dirs_and_open_wal(path, wal_options)?;
67        config.save(path)?;
68
69        let mut segments = SegmentHolder::default();
70        ensure_appendable_segment(&mut segments, path, &segments_path, &config)?;
71
72        let config_path = path.join(EDGE_CONFIG_FILE);
73        let config = SaveOnDisk::new(&config_path, config)
74            .map_err(|e| OperationError::service_error(e.to_string()))?;
75
76        Ok(Self {
77            path: path.into(),
78            config,
79            wal: parking_lot::Mutex::new(wal),
80            segments: LockedSegmentHolder::new(segments),
81        })
82    }
83
84    /// Load an edge shard from existing files at `path`.
85    ///
86    /// * If `config` is `Some`: check compatibility with loaded segments, then overwrite
87    ///   `edge_config.json` with it.
88    /// * If `config` is `None`: load config from `edge_config.json`, or infer from segments;
89    ///   check compatibility, then persist so future loads have it.
90    ///
91    /// Fails if no segments exist and no config can be loaded or inferred.
92    ///
93    /// To override WAL options (e.g. for embedded/mobile deployments where
94    /// the default 32 MiB segment capacity is too large), set
95    /// [`EdgeConfig::wal_options`] on the supplied config.
96    pub fn load(path: &Path, config: Option<EdgeConfig>) -> OperationResult<Self> {
97        let mut config = resolve_initial_config(path, config)?;
98
99        let wal_options = config
100            .as_ref()
101            .and_then(|c| c.wal_options.clone())
102            .unwrap_or_default();
103        let (wal, segments_path) = ensure_dirs_and_open_wal(path, wal_options)?;
104
105        let mut segments = load_segments(path, &segments_path, &mut config)?;
106
107        ensure_appendable_segment(
108            &mut segments,
109            path,
110            &segments_path,
111            config.as_ref().ok_or_else(|| {
112                OperationError::service_error(
113                    "edge config is not provided and no segments were loaded",
114                )
115            })?,
116        )?;
117
118        let config = config.ok_or_else(|| {
119            OperationError::service_error("edge config is not provided and no segments were loaded")
120        })?;
121
122        let config_path = path.join(EDGE_CONFIG_FILE);
123        let config = SaveOnDisk::new(&config_path, config)
124            .map_err(|e| OperationError::service_error(e.to_string()))?;
125
126        Ok(Self {
127            path: path.into(),
128            config,
129            wal: parking_lot::Mutex::new(wal),
130            segments: LockedSegmentHolder::new(segments),
131        })
132    }
133
134    pub fn config(&self) -> parking_lot::RwLockReadGuard<'_, EdgeConfig> {
135        self.config.read()
136    }
137
138    pub fn path(&self) -> &Path {
139        &self.path
140    }
141
142    /// Update global HNSW config and persist. Does not change per-vector HNSW.
143    pub fn set_hnsw_config(&self, hnsw_config: crate::segment::types::HnswConfig) -> OperationResult<()> {
144        self.config
145            .write(|cfg| cfg.set_hnsw_config(hnsw_config))
146            .map_err(|e| OperationError::service_error(e.to_string()))
147    }
148
149    /// Update HNSW config for a named vector and persist.
150    /// Fails if the vector does not exist. Immutable fields (e.g. size, distance) cannot be changed.
151    pub fn set_vector_hnsw_config(
152        &self,
153        vector_name: &str,
154        hnsw_config: crate::segment::types::HnswConfig,
155    ) -> OperationResult<()> {
156        let mut cfg = self.config.read().clone();
157        cfg.set_vector_hnsw_config(vector_name, hnsw_config)?;
158        self.config
159            .write(|c| *c = cfg)
160            .map_err(|e| OperationError::service_error(e.to_string()))
161    }
162
163    /// Update optimizer config and persist.
164    pub fn set_optimizers_config(&self, optimizers: EdgeOptimizersConfig) -> OperationResult<()> {
165        self.config
166            .write(|cfg| cfg.set_optimizers_config(optimizers))
167            .map_err(|e| OperationError::service_error(e.to_string()))
168    }
169
170    pub fn flush(&self) {
171        self.wal
172            .try_lock()
173            .expect("WAL lock acquired")
174            .flush()
175            .expect("WAL flushed");
176
177        self.segments
178            .try_read()
179            .expect("segment holder lock acquired")
180            .flush_all(true, true)
181            .expect("segments flushed");
182    }
183}
184
185impl Drop for EdgeShard {
186    fn drop(&mut self) {
187        self.flush();
188    }
189}
190
191fn has_existing_segments(path: &Path) -> bool {
192    let segments_path = path.join(SEGMENTS_PATH);
193    let Ok(entries) = fs::read_dir(&segments_path) else {
194        return false;
195    };
196    for entry in entries.flatten() {
197        let p = entry.path();
198        if !p.is_dir() {
199            continue;
200        }
201        if p.file_name()
202            .and_then(|n| n.to_str())
203            .is_some_and(|n| n.starts_with('.'))
204        {
205            continue;
206        }
207        if normalize_segment_dir(&p).ok().flatten().is_some() {
208            return true;
209        }
210    }
211    false
212}
213
214fn ensure_dirs_and_open_wal(
215    path: &Path,
216    wal_options: WalOptions,
217) -> OperationResult<(SerdeWal<CollectionUpdateOperations>, PathBuf)> {
218    let wal_path = path.join(WAL_PATH);
219    if !wal_path.exists() {
220        fs::create_dir(&wal_path).map_err(|err| {
221            OperationError::service_error(format!("failed to create WAL directory: {err}"))
222        })?;
223    }
224
225    let wal = SerdeWal::new(&wal_path, wal_options).map_err(|err| {
226        OperationError::service_error(format!("failed to open WAL {}: {err}", wal_path.display(),))
227    })?;
228
229    let segments_path = path.join(SEGMENTS_PATH);
230    if !segments_path.exists() {
231        fs::create_dir(&segments_path).map_err(|err| {
232            OperationError::service_error(format!("failed to create segments directory: {err}"))
233        })?;
234    }
235
236    Ok((wal, segments_path))
237}
238
239fn resolve_initial_config(
240    path: &Path,
241    config: Option<EdgeConfig>,
242) -> OperationResult<Option<EdgeConfig>> {
243    Ok(match config {
244        Some(c) => Some(c),
245        None => match EdgeConfig::load(path) {
246            Some(Ok(c)) => Some(c),
247            Some(Err(e)) => return Err(e),
248            None => None,
249        },
250    })
251}
252
253fn load_segments(
254    _path: &Path,
255    segments_path: &Path,
256    config: &mut Option<EdgeConfig>,
257) -> OperationResult<SegmentHolder> {
258    let segments_dir = fs::read_dir(segments_path).map_err(|err| {
259        OperationError::service_error(format!("failed to read segments directory: {err}"))
260    })?;
261
262    let mut segments = SegmentHolder::default();
263
264    for entry in segments_dir {
265        let entry = entry.map_err(|err| {
266            OperationError::service_error(format!(
267                "failed to read entry in segments directory: {err}",
268            ))
269        })?;
270
271        let segment_path = entry.path();
272
273        if !segment_path.is_dir() {
274            log::warn!(
275                "Skipping non-directory segment entry {}",
276                segment_path.display(),
277            );
278            continue;
279        }
280
281        if segment_path
282            .file_name()
283            .and_then(|n| n.to_str())
284            .is_some_and(|n| n.starts_with('.'))
285        {
286            log::warn!(
287                "Skipping hidden segment directory {}",
288                segment_path.display(),
289            );
290            continue;
291        }
292
293        let Some((segment_path, segment_uuid)) = normalize_segment_dir(&segment_path)? else {
294            continue;
295        };
296
297        let mut segment = load_segment(&segment_path, segment_uuid, None, &AtomicBool::new(false))
298            .map_err(|err| {
299                OperationError::service_error(format!(
300                    "failed to load segment {}: {err}",
301                    segment_path.display(),
302                ))
303            })?;
304
305        let segment_cfg = segment.config();
306        if let Some(cfg) = config.as_ref() {
307            cfg.check_compatible_with_segment_config(segment_cfg).map_err(
308                |err| OperationError::service_error(format!(
309                    "segment {} is incompatible with provided config or previously loaded segments: {err}",
310                    segment_path.display(),
311                ))
312            )?;
313        } else {
314            *config = Some(EdgeConfig::from_segment_config(segment_cfg));
315        }
316
317        segment.check_consistency_and_repair().map_err(|err| {
318            OperationError::service_error(format!(
319                "failed to repair segment {}: {err}",
320                segment_path.display(),
321            ))
322        })?;
323
324        segments.add_new(segment);
325    }
326
327    Ok(segments)
328}
329
330fn ensure_appendable_segment(
331    segments: &mut SegmentHolder,
332    path: &Path,
333    segments_path: &Path,
334    config: &EdgeConfig,
335) -> OperationResult<()> {
336    if segments.has_appendable_segment() {
337        return Ok(());
338    }
339
340    let payload_index_schema_path = path.join(PAYLOAD_INDEX_CONFIG_FILE);
341    let payload_index_schema = SaveOnDisk::load_or_init_default(&payload_index_schema_path)
342        .map_err(|err| {
343            OperationError::service_error(format!(
344                "failed to initialize payload index schema file {}: {err}",
345                payload_index_schema_path.display(),
346            ))
347        })?;
348
349    segments.create_appendable_segment(
350        segments_path,
351        config.plain_segment_config(),
352        Arc::new(payload_index_schema),
353        None,
354    )?;
355
356    debug_assert!(segments.has_appendable_segment());
357    Ok(())
358}
359
360// Default timeout of 1h used as a placeholder in Edge
361pub(crate) const DEFAULT_EDGE_TIMEOUT: Duration = Duration::from_secs(3600);