lance_io/
object_reader.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use std::fs::File;
5use std::ops::Range;
6use std::sync::Arc;
7
8use crate::local::join_local_io;
9#[cfg(windows)]
10use crate::local::read_exact_at;
11#[cfg(unix)]
12use std::os::unix::fs::FileExt;
13
14use bytes::Bytes;
15use futures::{
16    FutureExt,
17    future::{BoxFuture, Shared},
18    stream::{self, StreamExt},
19};
20use lance_core::deepsize::DeepSizeOf;
21use lance_core::{Error, Result, error::CloneableError};
22use object_store::ObjectStoreExt;
23use object_store::{GetOptions, GetResult, ObjectStore, Result as OSResult, path::Path};
24use tokio::sync::OnceCell;
25use tracing::instrument;
26
27use crate::{
28    object_store::DEFAULT_CLOUD_IO_PARALLELISM,
29    traits::{ByteStream, Reader},
30};
31
32trait StaticGetRange {
33    fn path(&self) -> &Path;
34    fn get_range(&self) -> BoxFuture<'static, OSResult<GetResult>>;
35}
36
37/// A wrapper around an object store and a path that implements a static
38/// get_range method by assuming self is stored in an Arc.
39struct GetRequest {
40    object_store: Arc<dyn ObjectStore>,
41    path: Path,
42    options: GetOptions,
43}
44
45impl StaticGetRange for Arc<GetRequest> {
46    fn path(&self) -> &Path {
47        &self.path
48    }
49
50    fn get_range(&self) -> BoxFuture<'static, OSResult<GetResult>> {
51        let store_and_path = self.clone();
52        Box::pin(async move {
53            store_and_path
54                .object_store
55                .get_opts(&store_and_path.path, store_and_path.options.clone())
56                .await
57        })
58    }
59}
60
61/// Object Reader
62///
63/// Object Store + Base Path
64#[derive(Debug)]
65pub struct CloudObjectReader {
66    // Object Store.
67    pub object_store: Arc<dyn ObjectStore>,
68    // File path
69    pub path: Path,
70    // File size, if known.
71    size: OnceCell<usize>,
72
73    block_size: usize,
74    io_parallelism: usize,
75    download_retry_count: usize,
76}
77
78impl DeepSizeOf for CloudObjectReader {
79    fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize {
80        // Skipping object_store because there is no easy way to do that and it shouldn't be too big
81        self.path.as_ref().deep_size_of_children(context)
82    }
83}
84
85impl CloudObjectReader {
86    /// Create an ObjectReader from URI
87    pub fn new(
88        object_store: Arc<dyn ObjectStore>,
89        path: Path,
90        block_size: usize,
91        known_size: Option<usize>,
92        download_retry_count: usize,
93    ) -> Result<Self> {
94        Ok(Self {
95            object_store,
96            path,
97            size: OnceCell::new_with(known_size),
98            block_size,
99            io_parallelism: DEFAULT_CLOUD_IO_PARALLELISM,
100            download_retry_count,
101        })
102    }
103
104    /// Override the I/O parallelism this reader advertises.
105    ///
106    /// `ObjectStore::open` / `open_with_size` pass their normalized effective
107    /// parallelism (`LANCE_IO_THREADS` override applied, at least 1) so
108    /// consumers that size concurrency windows off the reader honor the
109    /// configured request limit instead of the hardcoded cloud default.
110    pub fn with_io_parallelism(mut self, io_parallelism: usize) -> Self {
111        self.io_parallelism = io_parallelism;
112        self
113    }
114}
115
116// Retries for the initial request are handled by object store, but
117// there are no retries for failures that occur during the streaming
118// of the response body. Thus we add an outer retry loop here.
119async fn do_with_retry<'a, O>(f: impl Fn() -> BoxFuture<'a, OSResult<O>> + Clone) -> OSResult<O> {
120    let mut retries = 3;
121    loop {
122        let f = f.clone();
123        match f().await {
124            Ok(val) => return Ok(val),
125            Err(err) => {
126                if retries == 0 {
127                    return Err(err);
128                }
129                retries -= 1;
130            }
131        }
132    }
133}
134
135// We have a separate retry loop here.  This is because object_store does not
136// attempt retries on downloads that fail during streaming of the response body.
137//
138// However, this failure is pretty common (e.g. timeout) and we want to retry in these
139// situations.  In addition, we provide additional logging information in these
140// failures cases.
141async fn do_get_with_outer_retry(
142    download_retry_count: usize,
143    get_request: Arc<GetRequest>,
144    desc: impl Fn() -> String,
145) -> OSResult<Bytes> {
146    let mut retries = download_retry_count;
147    loop {
148        let get_request_clone = get_request.clone();
149        let get_result = do_with_retry(move || get_request_clone.get_range()).await?;
150        match get_result.bytes().await {
151            Ok(bytes) => return Ok(bytes),
152            Err(err) => {
153                if retries == 0 {
154                    log::warn!(
155                        "Failed to download {} from {} after {} attempts.  This may indicate that cloud storage is overloaded or your timeout settings are too restrictive.  Error details: {:?}",
156                        desc(),
157                        get_request.path(),
158                        download_retry_count,
159                        err
160                    );
161                    return Err(err);
162                }
163                log::debug!(
164                    "Retrying {} from {} (remaining retries: {}).  Error details: {:?}",
165                    desc(),
166                    get_request.path(),
167                    retries,
168                    err
169                );
170                retries -= 1;
171            }
172        }
173    }
174}
175
176impl Reader for CloudObjectReader {
177    fn path(&self) -> &Path {
178        &self.path
179    }
180
181    fn block_size(&self) -> usize {
182        self.block_size
183    }
184
185    fn io_parallelism(&self) -> usize {
186        self.io_parallelism
187    }
188
189    /// Object/File Size.
190    fn size(&self) -> BoxFuture<'_, object_store::Result<usize>> {
191        Box::pin(async move {
192            self.size
193                .get_or_try_init(|| async move {
194                    let meta =
195                        do_with_retry(|| Box::pin(self.object_store.head(&self.path))).await?;
196                    Ok(meta.size as usize)
197                })
198                .await
199                .cloned()
200        })
201    }
202
203    #[instrument(level = "debug", skip(self))]
204    fn get_range(&self, range: Range<usize>) -> BoxFuture<'static, OSResult<Bytes>> {
205        let range = Range {
206            start: range.start as u64,
207            end: range.end as u64,
208        };
209        // Downloading a range is the same two layers as downloading the whole
210        // object, so it goes through the same retry loop: `do_with_retry`
211        // retries the request, and the outer loop retries a body that stops
212        // partway.
213        let get_request = Arc::new(GetRequest {
214            object_store: self.object_store.clone(),
215            path: self.path.clone(),
216            options: GetOptions {
217                range: Some(range.clone().into()),
218                ..Default::default()
219            },
220        });
221        let download_retry_count = self.download_retry_count;
222        Box::pin(async move {
223            do_get_with_outer_retry(download_retry_count, get_request, move || {
224                format!("range {}..{}", range.start, range.end)
225            })
226            .await
227        })
228    }
229
230    #[instrument(level = "debug", skip_all)]
231    fn get_all(&self) -> BoxFuture<'_, OSResult<Bytes>> {
232        let get_request = Arc::new(GetRequest {
233            object_store: self.object_store.clone(),
234            path: self.path.clone(),
235            options: GetOptions::default(),
236        });
237        Box::pin(async move {
238            do_get_with_outer_retry(self.download_retry_count, get_request, || {
239                "read_all".to_string()
240            })
241            .await
242        })
243    }
244
245    fn get_stream(&self) -> BoxFuture<'_, OSResult<ByteStream>> {
246        let get_request = Arc::new(GetRequest {
247            object_store: self.object_store.clone(),
248            path: self.path.clone(),
249            options: GetOptions::default(),
250        });
251        Box::pin(async move {
252            let get_request_clone = get_request.clone();
253            let get_result = do_with_retry(move || get_request_clone.get_range()).await?;
254            Ok(get_result.into_stream())
255        })
256    }
257
258    fn get_range_stream(&self, range: Range<usize>) -> BoxFuture<'_, OSResult<ByteStream>> {
259        let get_request = Arc::new(GetRequest {
260            object_store: self.object_store.clone(),
261            path: self.path.clone(),
262            options: GetOptions {
263                range: Some(
264                    Range {
265                        start: range.start as u64,
266                        end: range.end as u64,
267                    }
268                    .into(),
269                ),
270                ..Default::default()
271            },
272        });
273        Box::pin(async move {
274            let get_request_clone = get_request.clone();
275            let get_result = do_with_retry(move || get_request_clone.get_range()).await?;
276            Ok(get_result.into_stream())
277        })
278    }
279}
280
281#[derive(Debug)]
282pub struct SmallReaderInner {
283    path: Path,
284    size: usize,
285    state: std::sync::Mutex<SmallReaderState>,
286}
287
288/// A reader for a file so small, we just eagerly read it all into memory.
289///
290/// When created, it represents a future that will read the whole file into memory.
291///
292/// On the first read call, it will start the read. Multiple threads can call read at the same time.
293///
294/// Once the read is complete, any thread can call read again to get the result.
295#[derive(Clone, Debug)]
296pub struct SmallReader {
297    inner: Arc<SmallReaderInner>,
298}
299
300enum SmallReaderState {
301    Loading(Shared<BoxFuture<'static, std::result::Result<Bytes, CloneableError>>>),
302    Finished(std::result::Result<Bytes, CloneableError>),
303}
304
305impl std::fmt::Debug for SmallReaderState {
306    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
307        match self {
308            Self::Loading(_) => write!(f, "Loading"),
309            Self::Finished(Ok(data)) => {
310                write!(f, "Finished({} bytes)", data.len())
311            }
312            Self::Finished(Err(err)) => {
313                write!(f, "Finished({})", err.inner())
314            }
315        }
316    }
317}
318
319impl SmallReader {
320    pub fn new(
321        store: Arc<dyn ObjectStore>,
322        path: Path,
323        download_retry_count: usize,
324        size: usize,
325    ) -> Self {
326        let path_ref = path.clone();
327        let state = SmallReaderState::Loading(
328            Box::pin(async move {
329                let object_reader =
330                    CloudObjectReader::new(store, path_ref, 0, None, download_retry_count)
331                        .map_err(CloneableError::new)?;
332                object_reader
333                    .get_all()
334                    .await
335                    .map_err(|err| CloneableError::new(Error::from(err)))
336            })
337            .boxed()
338            .shared(),
339        );
340        Self {
341            inner: Arc::new(SmallReaderInner {
342                path,
343                size,
344                state: std::sync::Mutex::new(state),
345            }),
346        }
347    }
348}
349
350impl SmallReaderInner {
351    async fn wait(&self) -> OSResult<Bytes> {
352        let future = {
353            let state = self.state.lock().unwrap();
354            match &*state {
355                SmallReaderState::Loading(future) => future.clone(),
356                SmallReaderState::Finished(result) => {
357                    return result.clone().map_err(|err| err.into_inner().into());
358                }
359            }
360        };
361
362        let result = future.await;
363        let result_to_return = result.clone().map_err(|err| err.into_inner().into());
364        let mut state = self.state.lock().unwrap();
365        if matches!(*state, SmallReaderState::Loading(_)) {
366            *state = SmallReaderState::Finished(result);
367        }
368        result_to_return
369    }
370}
371
372impl Reader for SmallReader {
373    fn path(&self) -> &Path {
374        &self.inner.path
375    }
376
377    fn block_size(&self) -> usize {
378        64 * 1024
379    }
380
381    fn io_parallelism(&self) -> usize {
382        1024
383    }
384
385    /// Object/File Size.
386    fn size(&self) -> BoxFuture<'_, OSResult<usize>> {
387        let size = self.inner.size;
388        Box::pin(async move { Ok(size) })
389    }
390
391    fn get_range(&self, range: Range<usize>) -> BoxFuture<'static, OSResult<Bytes>> {
392        let inner = self.inner.clone();
393        Box::pin(async move {
394            let bytes = inner.wait().await?;
395            let start = range.start;
396            let end = range.end;
397            if start >= bytes.len() || end > bytes.len() {
398                return Err(object_store::Error::Generic {
399                    store: "memory",
400                    source: format!(
401                        "Invalid range {}..{} for object of size {} bytes",
402                        start,
403                        end,
404                        bytes.len()
405                    )
406                    .into(),
407                });
408            }
409            Ok(bytes.slice(range))
410        })
411    }
412
413    fn get_all(&self) -> BoxFuture<'_, OSResult<Bytes>> {
414        Box::pin(async move { self.inner.wait().await })
415    }
416}
417
418pub(crate) fn stream_local_range(
419    file: Arc<File>,
420    path: Path,
421    io_tracker: Arc<crate::utils::tracking_store::IOTracker>,
422    range: Range<usize>,
423    chunk_size: usize,
424) -> ByteStream {
425    stream::try_unfold(
426        (file, path, io_tracker, range.start, range.end),
427        move |state| async move {
428            let (file, path, io_tracker, start, end) = state;
429            if start >= end {
430                return Ok(None);
431            }
432
433            let next = (start + chunk_size).min(end);
434            let file_clone = file.clone();
435            let path_clone = path.clone();
436            let num_bytes = (next - start) as u64;
437            let metrics = io_tracker.begin_io("get");
438            let result = join_local_io(tokio::task::spawn_blocking(move || {
439                let mut buf = bytes::BytesMut::with_capacity(next - start);
440                // Safety: buffer capacity matches the exact number of bytes we read below.
441                unsafe { buf.set_len(next - start) };
442                #[cfg(unix)]
443                file_clone.read_exact_at(buf.as_mut(), start as u64)?;
444                #[cfg(windows)]
445                read_exact_at(file_clone, buf.as_mut(), start as u64)?;
446                Ok::<_, std::io::Error>(buf.freeze())
447            }))
448            .await;
449            metrics.record(&result, num_bytes);
450            let bytes = result?;
451
452            io_tracker.record_read(
453                "get_range_stream",
454                path_clone,
455                num_bytes,
456                Some(start as u64..next as u64),
457            );
458
459            Ok(Some((bytes, (file, path, io_tracker, next, end))))
460        },
461    )
462    .boxed()
463}
464
465impl DeepSizeOf for SmallReader {
466    fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize {
467        let mut size = self.inner.path.as_ref().deep_size_of_children(context);
468
469        if let Ok(guard) = self.inner.state.try_lock()
470            && let SmallReaderState::Finished(Ok(data)) = &*guard
471        {
472            size += data.len();
473        }
474
475        size
476    }
477}
478
479#[cfg(test)]
480mod tests {
481    use std::fmt::{Display, Formatter};
482    use std::sync::Mutex;
483    use std::sync::atomic::{AtomicUsize, Ordering};
484
485    use futures::stream::BoxStream;
486    use lance_core::utils::testing::{ProxyObjectStore, ProxyObjectStorePolicy};
487    use object_store::memory::InMemory;
488    use object_store::{
489        Attributes, CopyOptions, GetResult, GetResultPayload, ListResult, MultipartUpload,
490        ObjectMeta, ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult,
491        Result as OSResult, path::Path,
492    };
493    use rstest::rstest;
494
495    use super::*;
496    use crate::utils::tracking_store::{IoStats, IoTrackingStore};
497
498    /// Answers every read with a body that fails partway: the request
499    /// succeeded and the download did not. Reads are all these tests do.
500    #[derive(Debug)]
501    struct PartialBodyStore;
502
503    impl Display for PartialBodyStore {
504        fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
505            write!(f, "PartialBodyStore")
506        }
507    }
508
509    #[async_trait::async_trait]
510    impl ObjectStore for PartialBodyStore {
511        async fn put_opts(
512            &self,
513            _location: &Path,
514            _payload: PutPayload,
515            _opts: PutOptions,
516        ) -> OSResult<PutResult> {
517            unimplemented!("put_opts is not used by these tests")
518        }
519        async fn put_multipart_opts(
520            &self,
521            _location: &Path,
522            _opts: PutMultipartOptions,
523        ) -> OSResult<Box<dyn MultipartUpload>> {
524            unimplemented!("put_multipart_opts is not used by these tests")
525        }
526        async fn get_opts(&self, location: &Path, options: GetOptions) -> OSResult<GetResult> {
527            let range = options
528                .range
529                .expect("the reader asks for a range")
530                .as_range(1024)
531                .expect("the range fits the object");
532            Ok(GetResult {
533                payload: GetResultPayload::Stream(
534                    futures::stream::once(async {
535                        Err(object_store::Error::Generic {
536                            store: "PartialBodyStore",
537                            source: "connection reset".into(),
538                        })
539                    })
540                    .boxed(),
541                ),
542                meta: ObjectMeta {
543                    location: location.clone(),
544                    last_modified: Default::default(),
545                    size: 1024,
546                    e_tag: None,
547                    version: None,
548                },
549                range,
550                attributes: Attributes::default(),
551                extensions: Default::default(),
552            })
553        }
554        fn delete_stream(
555            &self,
556            _locations: BoxStream<'static, OSResult<Path>>,
557        ) -> BoxStream<'static, OSResult<Path>> {
558            unimplemented!("delete_stream is not used by these tests")
559        }
560        fn list(&self, _prefix: Option<&Path>) -> BoxStream<'static, OSResult<ObjectMeta>> {
561            unimplemented!("list is not used by these tests")
562        }
563        async fn list_with_delimiter(&self, _prefix: Option<&Path>) -> OSResult<ListResult> {
564            unimplemented!("list_with_delimiter is not used by these tests")
565        }
566        async fn copy_opts(&self, _from: &Path, _to: &Path, _opts: CopyOptions) -> OSResult<()> {
567            unimplemented!("copy_opts is not used by these tests")
568        }
569    }
570
571    fn reader(store: Arc<dyn ObjectStore>, download_retry_count: usize) -> CloudObjectReader {
572        CloudObjectReader::new(
573            store,
574            Path::from("test/data.lance"),
575            4096,
576            None,
577            download_retry_count,
578        )
579        .unwrap()
580    }
581
582    /// A body that stops partway is what `download_retry_count` documents, and
583    /// `get_range` has to honor it like `get_all` already did.
584    #[rstest]
585    #[case::no_retries(0, 1)]
586    #[case::one_retry(1, 2)]
587    #[case::default(3, 4)]
588    #[case::raised(10, 11)]
589    #[tokio::test]
590    async fn get_range_retries_a_partial_download_as_configured(
591        #[case] download_retry_count: usize,
592        #[case] expected_requests: u64,
593    ) {
594        let stats = Arc::new(Mutex::new(IoStats::default()));
595        let store = IoTrackingStore::new(Arc::new(PartialBodyStore), stats.clone());
596        let reader = reader(Arc::new(store), download_retry_count);
597
598        reader.get_range(0..128).await.unwrap_err();
599
600        assert_eq!(stats.lock().unwrap().read_iops, expected_requests);
601    }
602
603    /// `do_with_retry` gives the request its own fixed retries, so the
604    /// configured download count does not stretch or shrink them.
605    #[tokio::test]
606    async fn get_range_leaves_the_request_retries_alone() {
607        // Counted by the policy: `IoTrackingStore` records only requests that
608        // succeed, and none of these do.
609        let requests = Arc::new(AtomicUsize::new(0));
610        let counted = requests.clone();
611        let mut policy = ProxyObjectStorePolicy::new();
612        policy.set_before_policy(
613            "fail_every_request",
614            Arc::new(move |_, _| {
615                counted.fetch_add(1, Ordering::Relaxed);
616                Err(Error::io("connection reset"))
617            }),
618        );
619        let store = ProxyObjectStore::new(Arc::new(InMemory::new()), Arc::new(Mutex::new(policy)));
620        let reader = reader(Arc::new(store), 10);
621
622        reader.get_range(0..128).await.unwrap_err();
623
624        assert_eq!(requests.load(Ordering::Relaxed), 4);
625    }
626}
lance_io/object_reader.rs

lance_io/
object_reader.rs
