Skip to main content

diskann_providers/storage/
storage_provider.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5use std::io::{Read, Result, Seek, Write};
6
7/// This module provides traits and implementations to access a storage system, it could be
8/// a file system or an exchange store.
9///
10/// The `StorageReadProvider` trait is an abstraction for accessing a storage system in a read only manner. It could be
11/// implemented with a file system or an exchange store. It provides methods to open a storage
12/// for reading, get the length of the storage and check if a storage with a given
13/// identifier exists.
14///
15/// The method is defined in a way that it can work with both file system and BigStorageShim.
16///
17/// Implementations of this trait will be used by DiskIndexReader to access the storage system.
18///
19/// The internal Reader associated type of the `StorageReadProvider` trait is defined by the `StorageReader` trait.
20pub trait StorageReadProvider: Sync {
21    type Reader: Read + Seek;
22
23    /// Open a storage with the given identifier for read.
24    fn open_reader(&self, item_identifier: &str) -> Result<Self::Reader>;
25
26    /// Get the length of the storage with the given identifier.
27    fn get_length(&self, item_identifier: &str) -> Result<u64>;
28
29    /// Check if the storage with the given identifier exists.
30    fn exists(&self, item_identifier: &str) -> bool;
31}
32
33/// `StorageWriteProvider` is a trait that abstracts over the ability to write to a storage. Since the ANN algorithm only writes into file system,
34/// currently we only have one implementation for this trait based on file system.
35pub trait StorageWriteProvider: Sync {
36    type Writer: WriteSeek;
37
38    /// Open a storage with the given identifier for write.
39    fn open_writer(&self, item_identifier: &str) -> Result<Self::Writer>;
40
41    /// Create a storage with the given identifier for write.
42    fn create_for_write(&self, item_identifier: &str) -> Result<Self::Writer>;
43
44    // Deletes a storage item with the given identifier.
45    fn delete(&self, item_identifier: &str) -> Result<()>;
46}
47
48/// Trait alias for types that implement both `Write` and `Seek`.
49///
50/// Use this when an API needs a writer that can also move the cursor.
51/// Implemented for any type that implements `Write` and `Seek`.
52pub trait WriteSeek: Write + Seek {}
53impl<T> WriteSeek for T where T: Write + Seek {}
54
55/// Object safe interface for opening and creating writers without exposing a concrete provider type.
56///
57/// This is useful when passing a writer provider through trait objects or other dynamic
58/// boundaries. Methods return boxed writers so callers can use a single uniform interface.
59pub trait DynWriteProvider: Sync {
60    /// Open an existing item for writing.
61    ///
62    /// Returns a boxed writer positioned by the provider. Fails if the item does not exist.
63    fn open_writer(&self, item_identifier: &str) -> std::io::Result<Box<dyn WriteSeek + '_>>;
64
65    /// Create a new item for writing.
66    ///
67    /// Returns a boxed writer for a new item. Behavior if the item already exists depends on the provider.
68    fn create_for_write(&self, item_identifier: &str) -> std::io::Result<Box<dyn WriteSeek + '_>>;
69
70    /// Delete an item identified by `item_identifier`.
71    fn delete(&self, item_identifier: &str) -> std::io::Result<()>;
72}
73
74impl<T> DynWriteProvider for T
75where
76    T: StorageWriteProvider,
77{
78    fn open_writer(&self, item_identifier: &str) -> std::io::Result<Box<dyn WriteSeek + '_>> {
79        self.open_writer(item_identifier)
80            .map(|w| Box::new(w) as Box<dyn WriteSeek>)
81    }
82
83    fn create_for_write(&self, item_identifier: &str) -> std::io::Result<Box<dyn WriteSeek + '_>> {
84        self.create_for_write(item_identifier)
85            .map(|w| Box::new(w) as Box<dyn WriteSeek>)
86    }
87
88    fn delete(&self, item_identifier: &str) -> std::io::Result<()> {
89        self.delete(item_identifier)
90    }
91}
92
93/// Adapter that exposes a `&dyn DynWriteProvider` as a `StorageWriteProvider`.
94///
95/// Useful when an API is generic over `StorageWriteProvider` but the caller only has a dynamic
96/// provider. The wrapper forwards all calls to the inner provider and returns boxed writers
97/// tied to the wrapper lifetime `'a`.
98pub struct WriteProviderWrapper<'a> {
99    inner: &'a dyn DynWriteProvider,
100}
101
102impl<'a> WriteProviderWrapper<'a> {
103    /// Construct a new wrapper around the given dynamic provider reference.
104    pub const fn new(inner: &'a dyn DynWriteProvider) -> Self {
105        Self { inner }
106    }
107}
108
109impl<'a> StorageWriteProvider for WriteProviderWrapper<'a> {
110    type Writer = Box<dyn WriteSeek + 'a>;
111
112    fn open_writer(&self, item_identifier: &str) -> std::io::Result<Self::Writer> {
113        self.inner.open_writer(item_identifier)
114    }
115
116    fn create_for_write(&self, item_identifier: &str) -> std::io::Result<Self::Writer> {
117        self.inner.create_for_write(item_identifier)
118    }
119
120    fn delete(&self, item_identifier: &str) -> std::io::Result<()> {
121        self.inner.delete(item_identifier)
122    }
123}