arrow_avro/reader/async_reader/
async_file_reader.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use crate::errors::AvroError;
19use bytes::Bytes;
20use futures::FutureExt;
21use futures::future::BoxFuture;
22use std::io::SeekFrom;
23use std::ops::Range;
24use tokio::io::{AsyncRead, AsyncReadExt, AsyncSeek, AsyncSeekExt};
25
26/// The asynchronous interface used by [`super::AsyncAvroFileReader`] to read avro files
27///
28/// Notes:
29///
30/// 1. There is a default implementation for types that implement [`AsyncRead`]
31///    and [`AsyncSeek`], for example [`tokio::fs::File`].
32///
33/// 2. Implementations for remote storage, such as the [`object_store`] crate,
34///    can implement this interface directly, typically by pairing a store
35///    handle with an object path and delegating [`Self::get_bytes`] and
36///    [`Self::get_byte_ranges`] to ranged reads. [`super::SpawnedReader`] can
37///    wrap such a reader to perform its I/O on a dedicated tokio runtime.
38///
39/// [`object_store`]: https://crates.io/crates/object_store
40///
41/// # Example: implementing `AsyncFileReader` for the `object_store` crate
42///
43/// ```no_run
44/// # use std::ops::Range;
45/// # use std::sync::Arc;
46/// use arrow_avro::errors::AvroError;
47/// use arrow_avro::reader::AsyncFileReader;
48/// use bytes::Bytes;
49/// use futures::FutureExt;
50/// use futures::future::BoxFuture;
51/// use object_store::path::Path;
52/// use object_store::{ObjectStore, ObjectStoreExt};
53///
54/// #[derive(Clone, Debug)]
55/// struct ObjectStoreReader {
56///     store: Arc<dyn ObjectStore>,
57///     path: Path,
58/// }
59///
60/// impl AsyncFileReader for ObjectStoreReader {
61///     fn get_bytes(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes, AvroError>> {
62///         async move {
63///             self.store
64///                 .get_range(&self.path, range)
65///                 .await
66///                 .map_err(|e| AvroError::General(e.to_string()))
67///         }
68///         .boxed()
69///     }
70///
71///     fn get_byte_ranges(
72///         &mut self,
73///         ranges: Vec<Range<u64>>,
74///     ) -> BoxFuture<'_, Result<Vec<Bytes>, AvroError>> {
75///         async move {
76///             self.store
77///                 .get_ranges(&self.path, &ranges)
78///                 .await
79///                 .map_err(|e| AvroError::General(e.to_string()))
80///         }
81///         .boxed()
82///     }
83/// }
84/// ```
85///
86/// [`tokio::fs::File`]: https://docs.rs/tokio/latest/tokio/fs/struct.File.html
87pub trait AsyncFileReader: Send {
88    /// Retrieve the bytes in `range`
89    fn get_bytes(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes, AvroError>>;
90
91    /// Retrieve multiple byte ranges. The default implementation will call `get_bytes` sequentially
92    fn get_byte_ranges(
93        &mut self,
94        ranges: Vec<Range<u64>>,
95    ) -> BoxFuture<'_, Result<Vec<Bytes>, AvroError>> {
96        async move {
97            let mut result = Vec::with_capacity(ranges.len());
98
99            for range in ranges.into_iter() {
100                let data = self.get_bytes(range).await?;
101                result.push(data);
102            }
103
104            Ok(result)
105        }
106        .boxed()
107    }
108}
109
110/// This allows Box<dyn AsyncFileReader + '_> to be used as an AsyncFileReader,
111impl AsyncFileReader for Box<dyn AsyncFileReader + '_> {
112    fn get_bytes(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes, AvroError>> {
113        self.as_mut().get_bytes(range)
114    }
115
116    fn get_byte_ranges(
117        &mut self,
118        ranges: Vec<Range<u64>>,
119    ) -> BoxFuture<'_, Result<Vec<Bytes>, AvroError>> {
120        self.as_mut().get_byte_ranges(ranges)
121    }
122}
123
124impl<T: AsyncRead + AsyncSeek + Unpin + Send> AsyncFileReader for T {
125    fn get_bytes(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes, AvroError>> {
126        async move {
127            self.seek(SeekFrom::Start(range.start)).await?;
128
129            let to_read = range.end - range.start;
130            let mut buffer = Vec::with_capacity(to_read as usize);
131            let read = self.take(to_read).read_to_end(&mut buffer).await?;
132            if read as u64 != to_read {
133                return Err(AvroError::EOF(format!(
134                    "expected to read {} bytes, got {}",
135                    to_read, read
136                )));
137            }
138
139            Ok(buffer.into())
140        }
141        .boxed()
142    }
143}
arrow_avro/reader/async_reader/async_file_reader.rs

arrow_avro/reader/async_reader/
async_file_reader.rs
