lance_io/
object_reader.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use std::ops::Range;
5use std::sync::Arc;
6
7use async_trait::async_trait;
8use bytes::Bytes;
9use deepsize::DeepSizeOf;
10use futures::{
11    future::{BoxFuture, Shared},
12    FutureExt,
13};
14use lance_core::{error::CloneableError, Error, Result};
15use object_store::{path::Path, GetOptions, GetResult, ObjectStore, Result as OSResult};
16use tokio::sync::OnceCell;
17use tracing::instrument;
18
19use crate::{object_store::DEFAULT_CLOUD_IO_PARALLELISM, traits::Reader};
20
21/// Object Reader
22///
23/// Object Store + Base Path
24#[derive(Debug)]
25pub struct CloudObjectReader {
26    // Object Store.
27    pub object_store: Arc<dyn ObjectStore>,
28    // File path
29    pub path: Path,
30    // File size, if known.
31    size: OnceCell<usize>,
32
33    block_size: usize,
34    download_retry_count: usize,
35}
36
37impl DeepSizeOf for CloudObjectReader {
38    fn deep_size_of_children(&self, context: &mut deepsize::Context) -> usize {
39        // Skipping object_store because there is no easy way to do that and it shouldn't be too big
40        self.path.as_ref().deep_size_of_children(context)
41    }
42}
43
44impl CloudObjectReader {
45    /// Create an ObjectReader from URI
46    pub fn new(
47        object_store: Arc<dyn ObjectStore>,
48        path: Path,
49        block_size: usize,
50        known_size: Option<usize>,
51        download_retry_count: usize,
52    ) -> Result<Self> {
53        Ok(Self {
54            object_store,
55            path,
56            size: OnceCell::new_with(known_size),
57            block_size,
58            download_retry_count,
59        })
60    }
61
62    // Retries for the initial request are handled by object store, but
63    // there are no retries for failures that occur during the streaming
64    // of the response body. Thus we add an outer retry loop here.
65    async fn do_with_retry<'a, O>(
66        &self,
67        f: impl Fn() -> BoxFuture<'a, OSResult<O>>,
68    ) -> OSResult<O> {
69        let mut retries = 3;
70        loop {
71            match f().await {
72                Ok(val) => return Ok(val),
73                Err(err) => {
74                    if retries == 0 {
75                        return Err(err);
76                    }
77                    retries -= 1;
78                }
79            }
80        }
81    }
82
83    // We have a separate retry loop here.  This is because object_store does not
84    // attempt retries on downloads that fail during streaming of the response body.
85    //
86    // However, this failure is pretty common (e.g. timeout) and we want to retry in these
87    // situations.  In addition, we provide additional logging information in these
88    // failures cases.
89    async fn do_get_with_outer_retry<'a>(
90        &self,
91        f: impl Fn() -> BoxFuture<'a, OSResult<GetResult>> + Copy,
92        desc: impl Fn() -> String,
93    ) -> OSResult<Bytes> {
94        let mut retries = self.download_retry_count;
95        loop {
96            let get_result = self.do_with_retry(f).await?;
97            match get_result.bytes().await {
98                Ok(bytes) => return Ok(bytes),
99                Err(err) => {
100                    if retries == 0 {
101                        log::warn!("Failed to download {} from {} after {} attempts.  This may indicate that cloud storage is overloaded or your timeout settings are too restrictive.  Error details: {:?}", desc(), self.path, self.download_retry_count, err);
102                        return Err(err);
103                    }
104                    log::debug!(
105                        "Retrying {} from {} (remaining retries: {}).  Error details: {:?}",
106                        desc(),
107                        self.path,
108                        retries,
109                        err
110                    );
111                    retries -= 1;
112                }
113            }
114        }
115    }
116}
117
118#[async_trait]
119impl Reader for CloudObjectReader {
120    fn path(&self) -> &Path {
121        &self.path
122    }
123
124    fn block_size(&self) -> usize {
125        self.block_size
126    }
127
128    fn io_parallelism(&self) -> usize {
129        DEFAULT_CLOUD_IO_PARALLELISM
130    }
131
132    /// Object/File Size.
133    async fn size(&self) -> object_store::Result<usize> {
134        self.size
135            .get_or_try_init(|| async move {
136                let meta = self
137                    .do_with_retry(|| self.object_store.head(&self.path))
138                    .await?;
139                Ok(meta.size)
140            })
141            .await
142            .cloned()
143    }
144
145    #[instrument(level = "debug", skip(self))]
146    async fn get_range(&self, range: Range<usize>) -> OSResult<Bytes> {
147        self.do_get_with_outer_retry(
148            || {
149                let options = GetOptions {
150                    range: Some(range.clone().into()),
151                    ..Default::default()
152                };
153                self.object_store.get_opts(&self.path, options)
154            },
155            || format!("range {:?}", range),
156        )
157        .await
158    }
159
160    #[instrument(level = "debug", skip_all)]
161    async fn get_all(&self) -> OSResult<Bytes> {
162        self.do_get_with_outer_retry(
163            || {
164                self.object_store
165                    .get_opts(&self.path, GetOptions::default())
166            },
167            || "read_all".to_string(),
168        )
169        .await
170    }
171}
172
173/// A reader for a file so small, we just eagerly read it all into memory.
174///
175/// When created, it represents a future that will read the whole file into memory.
176///
177/// On the first read call, it will start the read. Multiple threads can call read at the same time.
178///
179/// Once the read is complete, any thread can call read again to get the result.
180#[derive(Debug)]
181pub struct SmallReader {
182    path: Path,
183    size: usize,
184    state: Arc<std::sync::Mutex<SmallReaderState>>,
185}
186
187enum SmallReaderState {
188    Loading(Shared<BoxFuture<'static, std::result::Result<Bytes, CloneableError>>>),
189    Finished(std::result::Result<Bytes, CloneableError>),
190}
191
192impl std::fmt::Debug for SmallReaderState {
193    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
194        match self {
195            Self::Loading(_) => write!(f, "Loading"),
196            Self::Finished(Ok(data)) => {
197                write!(f, "Finished({} bytes)", data.len())
198            }
199            Self::Finished(Err(err)) => {
200                write!(f, "Finished({})", err.0)
201            }
202        }
203    }
204}
205
206impl SmallReader {
207    pub fn new(
208        store: Arc<dyn ObjectStore>,
209        path: Path,
210        download_retry_count: usize,
211        size: usize,
212    ) -> Self {
213        let path_ref = path.clone();
214        let state = SmallReaderState::Loading(
215            Box::pin(async move {
216                let object_reader =
217                    CloudObjectReader::new(store, path_ref, 0, None, download_retry_count)
218                        .map_err(CloneableError)?;
219                object_reader
220                    .get_all()
221                    .await
222                    .map_err(|err| CloneableError(Error::from(err)))
223            })
224            .boxed()
225            .shared(),
226        );
227        Self {
228            path,
229            size,
230            state: Arc::new(std::sync::Mutex::new(state)),
231        }
232    }
233
234    async fn wait(&self) -> OSResult<Bytes> {
235        let future = {
236            let state = self.state.lock().unwrap();
237            match &*state {
238                SmallReaderState::Loading(future) => future.clone(),
239                SmallReaderState::Finished(result) => {
240                    return result.clone().map_err(|err| err.0.into());
241                }
242            }
243        };
244
245        let result = future.await;
246        let result_to_return = result.clone().map_err(|err| err.0.into());
247        let mut state = self.state.lock().unwrap();
248        if matches!(*state, SmallReaderState::Loading(_)) {
249            *state = SmallReaderState::Finished(result);
250        }
251        result_to_return
252    }
253}
254
255#[async_trait]
256impl Reader for SmallReader {
257    fn path(&self) -> &Path {
258        &self.path
259    }
260
261    fn block_size(&self) -> usize {
262        64 * 1024
263    }
264
265    fn io_parallelism(&self) -> usize {
266        1024
267    }
268
269    /// Object/File Size.
270    async fn size(&self) -> OSResult<usize> {
271        Ok(self.size)
272    }
273
274    async fn get_range(&self, range: Range<usize>) -> OSResult<Bytes> {
275        self.wait().await.and_then(|bytes| {
276            let start = range.start;
277            let end = range.end;
278            if start >= bytes.len() || end > bytes.len() {
279                return Err(object_store::Error::Generic {
280                    store: "memory",
281                    source: format!(
282                        "Invalid range {}..{} for object of size {} bytes",
283                        start,
284                        end,
285                        bytes.len()
286                    )
287                    .into(),
288                });
289            }
290            Ok(bytes.slice(range))
291        })
292    }
293
294    async fn get_all(&self) -> OSResult<Bytes> {
295        self.wait().await
296    }
297}
298
299impl DeepSizeOf for SmallReader {
300    fn deep_size_of_children(&self, context: &mut deepsize::Context) -> usize {
301        let mut size = self.path.as_ref().deep_size_of_children(context);
302
303        if let Ok(guard) = self.state.try_lock() {
304            if let SmallReaderState::Finished(Ok(data)) = &*guard {
305                size += data.len();
306            }
307        }
308
309        size
310    }
311}
lance_io/object_reader.rs

lance_io/
object_reader.rs
