vortex_io/
read_at.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::sync::Arc;
5use std::time::Instant;
6
7use futures::FutureExt;
8use futures::StreamExt;
9use futures::future::BoxFuture;
10use futures::stream;
11use futures::stream::BoxStream;
12use vortex_array::buffer::BufferHandle;
13use vortex_buffer::Alignment;
14use vortex_buffer::ByteBuffer;
15use vortex_error::VortexExpect;
16use vortex_error::VortexResult;
17use vortex_error::vortex_bail;
18use vortex_metrics::Counter;
19use vortex_metrics::Histogram;
20use vortex_metrics::Label;
21use vortex_metrics::MetricBuilder;
22use vortex_metrics::MetricsRegistry;
23use vortex_metrics::Timer;
24
25/// Configuration for coalescing nearby I/O requests into single operations.
26#[derive(Clone, Copy, Debug)]
27pub struct CoalesceConfig {
28    /// The maximum "empty" distance between two requests to consider them for coalescing.
29    pub distance: u64,
30    /// The maximum total size spanned by a coalesced request.
31    pub max_size: u64,
32}
33
34/// A positional read request used by [`VortexReadAt::read_ranges`].
35#[derive(Clone, Copy, Debug, PartialEq, Eq)]
36pub struct ReadAtRequest {
37    /// The byte offset at which to start reading.
38    pub offset: u64,
39    /// The exact number of bytes to read.
40    pub length: usize,
41    /// The required alignment of the returned buffer.
42    pub alignment: Alignment,
43}
44
45impl ReadAtRequest {
46    /// Creates a positional read request.
47    pub const fn new(offset: u64, length: usize, alignment: Alignment) -> Self {
48        Self {
49            offset,
50            length,
51            alignment,
52        }
53    }
54}
55
56/// A stream of positional read results, yielded as each request completes.
57pub type ReadAtStream = BoxStream<'static, (ReadAtRequest, VortexResult<BufferHandle>)>;
58
59impl CoalesceConfig {
60    /// Creates a new coalesce configuration.
61    pub const fn new(distance: u64, max_size: u64) -> Self {
62        Self { distance, max_size }
63    }
64
65    /// Configuration appropriate for in-memory / low-latency sources.
66    pub const fn in_memory() -> Self {
67        Self::new(8 * 1024, 8 * 1024) // 8KB
68    }
69
70    /// Configuration appropriate for local filesystem access.
71    pub const fn file() -> Self {
72        Self::new(1 << 20, 4 << 20) // 1MB distance, 4MB max
73    }
74
75    /// Configuration appropriate for object storage (S3, GCS, etc.).
76    pub const fn object_storage() -> Self {
77        Self::new(1 << 20, 16 << 20) // 1MB distance, 16MB max
78    }
79}
80
81/// The unified read trait for Vortex I/O sources.
82///
83/// This trait provides async positional reads to underlying storage and is used by the vortex-file
84/// crate to read data from files or object stores.
85pub trait VortexReadAt: Send + Sync + 'static {
86    /// URI for debugging/logging. Returns `None` for anonymous sources.
87    fn uri(&self) -> Option<&Arc<str>> {
88        None
89    }
90
91    /// Configuration for merging nearby I/O requests into fewer, larger reads.
92    fn coalesce_config(&self) -> Option<CoalesceConfig> {
93        None
94    }
95
96    /// Maximum number of concurrent I/O requests for that should be pulled from this source.
97    ///
98    /// This value is used to control how many [`VortexReadAt::read_at`] calls can
99    /// be in-flight simultaneously. Higher values allow more parallelism but consume
100    /// more resources (memory, file descriptors, network connections).
101    ///
102    /// Implementations should choose a value appropriate for their underlying storage
103    /// characteristics. Low-latency sources benefit less from high concurrency, while
104    /// high-latency sources (like remote storage) benefit significantly from issuing
105    /// many requests in parallel.
106    fn concurrency(&self) -> usize;
107
108    /// Asynchronously get the number of bytes of the underlying source.
109    fn size(&self) -> BoxFuture<'static, VortexResult<u64>>;
110
111    /// Request an asynchronous positional read. Results will be returned as a [`BufferHandle`].
112    ///
113    /// If the reader does not have the requested number of bytes, the returned Future will complete
114    /// with an [`UnexpectedEof`][std::io::ErrorKind::UnexpectedEof] error.
115    fn read_at(
116        &self,
117        offset: u64,
118        length: usize,
119        alignment: Alignment,
120    ) -> BoxFuture<'static, VortexResult<BufferHandle>>;
121
122    /// Request multiple asynchronous positional reads.
123    ///
124    /// Each item includes its request and result, and is yielded as soon as that read completes.
125    /// A failed request does not prevent other results from being yielded. Callers should submit
126    /// batches no larger than [`VortexReadAt::concurrency`].
127    fn read_ranges(&self, requests: Arc<[ReadAtRequest]>) -> ReadAtStream {
128        let reads = requests
129            .iter()
130            .copied()
131            .map(|request| {
132                let read = self.read_at(request.offset, request.length, request.alignment);
133                async move { (request, read.await) }
134            })
135            .collect::<Vec<_>>();
136        stream::iter(reads)
137            .buffer_unordered(self.concurrency().max(1))
138            .boxed()
139    }
140}
141
142impl VortexReadAt for Arc<dyn VortexReadAt> {
143    fn uri(&self) -> Option<&Arc<str>> {
144        self.as_ref().uri()
145    }
146
147    fn coalesce_config(&self) -> Option<CoalesceConfig> {
148        self.as_ref().coalesce_config()
149    }
150
151    fn concurrency(&self) -> usize {
152        self.as_ref().concurrency()
153    }
154
155    fn size(&self) -> BoxFuture<'static, VortexResult<u64>> {
156        self.as_ref().size()
157    }
158
159    fn read_at(
160        &self,
161        offset: u64,
162        length: usize,
163        alignment: Alignment,
164    ) -> BoxFuture<'static, VortexResult<BufferHandle>> {
165        self.as_ref().read_at(offset, length, alignment)
166    }
167
168    fn read_ranges(&self, requests: Arc<[ReadAtRequest]>) -> ReadAtStream {
169        self.as_ref().read_ranges(requests)
170    }
171}
172
173impl<R: VortexReadAt> VortexReadAt for Arc<R> {
174    fn uri(&self) -> Option<&Arc<str>> {
175        self.as_ref().uri()
176    }
177
178    fn coalesce_config(&self) -> Option<CoalesceConfig> {
179        self.as_ref().coalesce_config()
180    }
181
182    fn concurrency(&self) -> usize {
183        self.as_ref().concurrency()
184    }
185
186    fn size(&self) -> BoxFuture<'static, VortexResult<u64>> {
187        self.as_ref().size()
188    }
189
190    fn read_at(
191        &self,
192        offset: u64,
193        length: usize,
194        alignment: Alignment,
195    ) -> BoxFuture<'static, VortexResult<BufferHandle>> {
196        self.as_ref().read_at(offset, length, alignment)
197    }
198
199    fn read_ranges(&self, requests: Arc<[ReadAtRequest]>) -> ReadAtStream {
200        self.as_ref().read_ranges(requests)
201    }
202}
203
204impl VortexReadAt for ByteBuffer {
205    fn size(&self) -> BoxFuture<'static, VortexResult<u64>> {
206        let length = self.len() as u64;
207        async move { Ok(length) }.boxed()
208    }
209
210    fn concurrency(&self) -> usize {
211        16
212    }
213
214    fn read_at(
215        &self,
216        offset: u64,
217        length: usize,
218        alignment: Alignment,
219    ) -> BoxFuture<'static, VortexResult<BufferHandle>> {
220        let buffer = self.clone();
221        async move {
222            let start = usize::try_from(offset).vortex_expect("start too big for usize");
223            let end =
224                usize::try_from(offset + length as u64).vortex_expect("end too big for usize");
225            if end > buffer.len() {
226                vortex_bail!(
227                    "Requested range {}..{} out of bounds for buffer of length {}",
228                    start,
229                    end,
230                    buffer.len()
231                );
232            }
233            Ok(BufferHandle::new_host(
234                buffer.slice_unaligned(start..end).aligned(alignment),
235            ))
236        }
237        .boxed()
238    }
239}
240
241/// A wrapper that instruments a [`VortexReadAt`] with metrics.
242#[derive(Clone)]
243pub struct InstrumentedReadAt<T: VortexReadAt + Clone> {
244    read: T,
245    // We use `Arc` to take care of all the complexity that's potentially associated with reference counting
246    // and dropping
247    metrics: Arc<InnerMetrics>,
248}
249
250struct InnerMetrics {
251    sizes: Histogram,
252    total_size: Counter,
253    durations: Timer,
254}
255
256impl<T: VortexReadAt + Clone> InstrumentedReadAt<T> {
257    pub fn new(read: T, metrics_registry: &dyn MetricsRegistry) -> Self {
258        Self::new_with_labels(read, metrics_registry, Vec::<Label>::default())
259    }
260
261    pub fn new_with_labels<I, L>(read: T, metrics_registry: &dyn MetricsRegistry, labels: I) -> Self
262    where
263        I: IntoIterator<Item = L>,
264        L: Into<Label>,
265    {
266        let labels = labels.into_iter().map(|l| l.into()).collect::<Vec<Label>>();
267        let sizes = MetricBuilder::new(metrics_registry)
268            .add_labels(labels.clone())
269            .histogram("vortex.io.read.size");
270        let total_size = MetricBuilder::new(metrics_registry)
271            .add_labels(labels.clone())
272            .counter("vortex.io.read.total_size");
273        let durations = MetricBuilder::new(metrics_registry)
274            .add_labels(labels)
275            .timer("vortex.io.read.duration");
276
277        Self {
278            read,
279            metrics: Arc::new(InnerMetrics {
280                sizes,
281                total_size,
282                durations,
283            }),
284        }
285    }
286}
287
288impl InnerMetrics {
289    fn log_sizes(&self) {
290        tracing::debug!("Reads: {}", self.sizes.count());
291        if !self.sizes.is_empty() {
292            tracing::debug!(
293                "Read size: p50={} p95={} p99={} p999={}",
294                self.sizes.quantile(0.5).vortex_expect("must not be empty"),
295                self.sizes.quantile(0.95).vortex_expect("must not be empty"),
296                self.sizes.quantile(0.99).vortex_expect("must not be empty"),
297                self.sizes
298                    .quantile(0.999)
299                    .vortex_expect("must not be empty"),
300            );
301        }
302        tracing::debug!("Total read size: {}", self.total_size.value());
303    }
304
305    fn log_durations(&self) {
306        if !self.durations.is_empty() {
307            tracing::debug!(
308                "Read duration: p50={}ms p95={}ms p99={}ms p999={}ms",
309                self.durations
310                    .quantile(0.5)
311                    .vortex_expect("must not be empty")
312                    .as_millis(),
313                self.durations
314                    .quantile(0.95)
315                    .vortex_expect("must not be empty")
316                    .as_millis(),
317                self.durations
318                    .quantile(0.99)
319                    .vortex_expect("must not be empty")
320                    .as_millis(),
321                self.durations
322                    .quantile(0.999)
323                    .vortex_expect("must not be empty")
324                    .as_millis(),
325            );
326        }
327    }
328}
329
330// We implement drop for `InnerMetrics` so this will be logged only when we eventually drop the final instance of `InstrumentedRead`
331impl Drop for InnerMetrics {
332    fn drop(&mut self) {
333        self.log_sizes();
334        self.log_durations();
335    }
336}
337
338impl<T: VortexReadAt + Clone> VortexReadAt for InstrumentedReadAt<T> {
339    fn uri(&self) -> Option<&Arc<str>> {
340        self.read.uri()
341    }
342
343    fn coalesce_config(&self) -> Option<CoalesceConfig> {
344        self.read.coalesce_config()
345    }
346
347    fn concurrency(&self) -> usize {
348        self.read.concurrency()
349    }
350
351    fn size(&self) -> BoxFuture<'static, VortexResult<u64>> {
352        self.read.size()
353    }
354
355    fn read_at(
356        &self,
357        offset: u64,
358        length: usize,
359        alignment: Alignment,
360    ) -> BoxFuture<'static, VortexResult<BufferHandle>> {
361        let durations = self.metrics.durations.clone();
362        let sizes = self.metrics.sizes.clone();
363        let total_size = self.metrics.total_size.clone();
364
365        let read_fut = self.read.read_at(offset, length, alignment);
366        async move {
367            let _timer = durations.time();
368            let buf = read_fut.await;
369            sizes.update(length as f64);
370            total_size.add(length as u64);
371            buf
372        }
373        .boxed()
374    }
375
376    fn read_ranges(&self, requests: Arc<[ReadAtRequest]>) -> ReadAtStream {
377        let durations = self.metrics.durations.clone();
378        let sizes = self.metrics.sizes.clone();
379        let total_size = self.metrics.total_size.clone();
380        let start = Instant::now();
381        self.read
382            .read_ranges(requests)
383            .map(move |(request, result)| {
384                durations.update(start.elapsed());
385                sizes.update(request.length as f64);
386                total_size.add(request.length as u64);
387                (request, result)
388            })
389            .boxed()
390    }
391}
392
393#[cfg(test)]
394mod tests {
395    use std::sync::Arc;
396    use std::time::Duration;
397
398    use vortex_buffer::Alignment;
399    use vortex_buffer::ByteBuffer;
400
401    use super::*;
402
403    struct DelayedReadAt;
404
405    impl VortexReadAt for DelayedReadAt {
406        fn concurrency(&self) -> usize {
407            2
408        }
409
410        fn size(&self) -> BoxFuture<'static, VortexResult<u64>> {
411            async { Ok(2) }.boxed()
412        }
413
414        fn read_at(
415            &self,
416            offset: u64,
417            _length: usize,
418            _alignment: Alignment,
419        ) -> BoxFuture<'static, VortexResult<BufferHandle>> {
420            async move {
421                if offset == 0 {
422                    tokio::time::sleep(Duration::from_millis(50)).await;
423                }
424                Ok(BufferHandle::new_host(ByteBuffer::from(vec![
425                    u8::try_from(offset).vortex_expect("test offset fits in u8"),
426                ])))
427            }
428            .boxed()
429        }
430    }
431
432    #[test]
433    fn test_coalesce_config_in_memory() {
434        let config = CoalesceConfig::in_memory();
435        assert_eq!(config.distance, 8 * 1024);
436        assert_eq!(config.max_size, 8 * 1024);
437    }
438
439    #[test]
440    fn test_coalesce_config_file() {
441        let config = CoalesceConfig::file();
442        assert_eq!(config.distance, 1 << 20); // 1MB
443        assert_eq!(config.max_size, 4 << 20); // 4MB
444    }
445
446    #[test]
447    fn test_coalesce_config_object_storage() {
448        let config = CoalesceConfig::object_storage();
449        assert_eq!(config.distance, 1 << 20); // 1MB
450        assert_eq!(config.max_size, 16 << 20); // 16MB
451    }
452
453    #[tokio::test]
454    async fn test_byte_buffer_read_at() {
455        let data = ByteBuffer::from(vec![1, 2, 3, 4, 5]);
456
457        let result = data.read_at(1, 3, Alignment::none()).await.unwrap();
458        assert_eq!(result.to_host().await.as_ref(), &[2, 3, 4]);
459    }
460
461    #[tokio::test]
462    async fn test_byte_buffer_read_ranges() -> VortexResult<()> {
463        let data = ByteBuffer::from(vec![1, 2, 3, 4, 5, 6]);
464        let requests = Arc::from([
465            ReadAtRequest::new(4, 2, Alignment::none()),
466            ReadAtRequest::new(0, 1, Alignment::none()),
467            ReadAtRequest::new(2, 3, Alignment::none()),
468        ]);
469
470        let results = data.read_ranges(requests).collect::<Vec<_>>().await;
471        let mut offsets = Vec::new();
472        for (request, result) in results {
473            offsets.push(request.offset);
474            let expected: &[u8] = match request.offset {
475                0 => &[1],
476                2 => &[3, 4, 5],
477                4 => &[5, 6],
478                offset => panic!("unexpected offset: {offset}"),
479            };
480            assert_eq!(result?.to_host().await.as_ref(), expected);
481        }
482        offsets.sort_unstable();
483        assert_eq!(offsets, [0, 2, 4]);
484        Ok(())
485    }
486
487    #[tokio::test]
488    async fn test_read_ranges_keeps_streaming_after_an_error() -> VortexResult<()> {
489        let data = ByteBuffer::from(vec![1, 2, 3]);
490        let requests = Arc::from([
491            ReadAtRequest::new(100, 1, Alignment::none()),
492            ReadAtRequest::new(1, 2, Alignment::none()),
493        ]);
494
495        let results = data.read_ranges(requests).collect::<Vec<_>>().await;
496
497        assert_eq!(results.len(), 2);
498        assert!(results.iter().any(|(_, result)| result.is_err()));
499        let (_, valid) = results
500            .into_iter()
501            .find(|(request, _)| request.offset == 1)
502            .vortex_expect("valid request result is present");
503        assert_eq!(valid?.to_host().await.as_ref(), &[2, 3]);
504        Ok(())
505    }
506
507    #[tokio::test]
508    async fn test_read_ranges_yields_in_completion_order() -> VortexResult<()> {
509        let requests = Arc::from([
510            ReadAtRequest::new(0, 1, Alignment::none()),
511            ReadAtRequest::new(1, 1, Alignment::none()),
512        ]);
513        let mut results = DelayedReadAt.read_ranges(requests);
514
515        let (first_request, first_result) = results
516            .next()
517            .await
518            .vortex_expect("first result is present");
519
520        assert_eq!(first_request.offset, 1);
521        assert_eq!(first_result?.to_host().await.as_ref(), &[1]);
522        assert_eq!(results.count().await, 1);
523        Ok(())
524    }
525
526    #[tokio::test]
527    async fn test_byte_buffer_read_out_of_bounds() {
528        let data = ByteBuffer::from(vec![1, 2, 3]);
529
530        let result = data.read_at(1, 9, Alignment::none()).await;
531        assert!(result.is_err());
532    }
533
534    #[tokio::test]
535    async fn test_arc_read_at() {
536        let data = Arc::new(ByteBuffer::from(vec![1, 2, 3, 4, 5]));
537
538        let result = data.read_at(2, 3, Alignment::none()).await.unwrap();
539        assert_eq!(result.to_host().await.as_ref(), &[3, 4, 5]);
540
541        let size = data.size().await.unwrap();
542        assert_eq!(size, 5);
543    }
544}
vortex_io/read_at.rs

vortex_io/
read_at.rs
