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}