vortex_io/object_store/
read_at.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::io;
5use std::sync::Arc;
6
7use futures::FutureExt;
8use futures::StreamExt;
9use futures::future::BoxFuture;
10use object_store::GetOptions;
11use object_store::GetRange;
12use object_store::GetResultPayload;
13use object_store::ObjectStore;
14use object_store::ObjectStoreExt;
15use object_store::path::Path as ObjectPath;
16use vortex_array::buffer::BufferHandle;
17use vortex_buffer::Alignment;
18use vortex_buffer::ByteBufferMut;
19use vortex_error::VortexError;
20use vortex_error::VortexResult;
21use vortex_error::vortex_ensure;
22
23use crate::CoalesceConfig;
24use crate::VortexReadAt;
25use crate::runtime::Handle;
26#[cfg(not(target_arch = "wasm32"))]
27use crate::std_file::read_exact_at;
28
29/// Default number of concurrent requests to allow.
30pub const DEFAULT_CONCURRENCY: usize = 192;
31
32/// An object store backed I/O source.
33pub struct ObjectStoreReadAt {
34    store: Arc<dyn ObjectStore>,
35    path: ObjectPath,
36    uri: Arc<str>,
37    handle: Handle,
38    concurrency: usize,
39    coalesce_config: Option<CoalesceConfig>,
40}
41
42impl ObjectStoreReadAt {
43    /// Create a new object store source.
44    pub fn new(store: Arc<dyn ObjectStore>, path: ObjectPath, handle: Handle) -> Self {
45        let uri = Arc::from(path.to_string());
46        Self {
47            store,
48            path,
49            uri,
50            handle,
51            concurrency: DEFAULT_CONCURRENCY,
52            coalesce_config: Some(CoalesceConfig::object_storage()),
53        }
54    }
55
56    /// Set the concurrency for this source.
57    pub fn with_concurrency(mut self, concurrency: usize) -> Self {
58        self.concurrency = concurrency;
59        self
60    }
61
62    /// Set the coalesce config for this source.
63    pub fn with_coalesce_config(mut self, config: CoalesceConfig) -> Self {
64        self.coalesce_config = Some(config);
65        self
66    }
67}
68
69impl VortexReadAt for ObjectStoreReadAt {
70    fn uri(&self) -> Option<&Arc<str>> {
71        Some(&self.uri)
72    }
73
74    fn coalesce_config(&self) -> Option<CoalesceConfig> {
75        self.coalesce_config
76    }
77
78    fn concurrency(&self) -> usize {
79        self.concurrency
80    }
81
82    fn size(&self) -> BoxFuture<'static, VortexResult<u64>> {
83        let store = self.store.clone();
84        let path = self.path.clone();
85        async move {
86            store
87                .head(&path)
88                .await
89                .map(|h| h.size)
90                .map_err(VortexError::from)
91        }
92        .boxed()
93    }
94
95    fn read_at(
96        &self,
97        offset: u64,
98        length: usize,
99        alignment: Alignment,
100    ) -> BoxFuture<'static, VortexResult<BufferHandle>> {
101        let store = self.store.clone();
102        let path = self.path.clone();
103        let handle = self.handle.clone();
104        let range = offset..(offset + length as u64);
105
106        async move {
107            let mut buffer = ByteBufferMut::with_capacity_aligned(length, alignment);
108
109            let response = store
110                .get_opts(
111                    &path,
112                    GetOptions {
113                        range: Some(GetRange::Bounded(range.clone())),
114                        ..Default::default()
115                    },
116                )
117                .await?;
118
119            let buffer = match response.payload {
120                #[cfg(not(target_arch = "wasm32"))]
121                GetResultPayload::File(file, _) => {
122                    unsafe { buffer.set_len(length) };
123
124                    handle
125                        .spawn_blocking(move || {
126                            read_exact_at(&file, &mut buffer, range.start)?;
127                            Ok::<_, io::Error>(buffer)
128                        })
129                        .await
130                        .map_err(io::Error::other)?
131                }
132                #[cfg(target_arch = "wasm32")]
133                GetResultPayload::File(..) => {
134                    unreachable!("File payload not supported on wasm32")
135                }
136                GetResultPayload::Stream(mut byte_stream) => {
137                    while let Some(bytes) = byte_stream.next().await {
138                        buffer.extend_from_slice(&bytes?);
139                    }
140
141                    vortex_ensure!(
142                        buffer.len() == length,
143                        "Object store stream returned {} bytes but expected {} bytes (range: {:?})",
144                        buffer.len(),
145                        length,
146                        range
147                    );
148
149                    buffer
150                }
151            };
152
153            Ok(BufferHandle::new_host(buffer.freeze()))
154        }
155        .boxed()
156    }
157}
vortex_io/object_store/read_at.rs

vortex_io/object_store/
read_at.rs
