arrow_avro/reader/async_reader/
store.rs1use crate::errors::AvroError;
19use crate::reader::async_reader::AsyncFileReader;
20use bytes::Bytes;
21use futures::future::BoxFuture;
22use futures::{FutureExt, TryFutureExt};
23use object_store::ObjectStore;
24use object_store::ObjectStoreExt;
25use object_store::path::Path;
26use std::ops::Range;
27use std::sync::Arc;
28use tokio::runtime::Handle;
29
30#[deprecated(
32 since = "59.2.0",
33 note = "Implement `AsyncFileReader` directly instead; see the example on the `AsyncFileReader` trait documentation and `arrow-avro/examples/object_store.rs`. Use `SpawnedReader` to perform I/O on a dedicated runtime. See also https://github.com/apache/arrow-rs/issues/10308"
34)]
35#[derive(Clone, Debug)]
36pub struct AvroObjectReader {
37 store: Arc<dyn ObjectStore>,
38 path: Path,
39 runtime: Option<Handle>,
40}
41
42#[expect(deprecated)]
43impl AvroObjectReader {
44 pub fn new(store: Arc<dyn ObjectStore>, path: Path) -> Self {
46 Self {
47 store,
48 path,
49 runtime: None,
50 }
51 }
52
53 #[deprecated(
62 since = "59.2.0",
63 note = "Wrap the reader in a `SpawnedReader` instead, e.g. `SpawnedReader::new(reader, handle)`. See also https://github.com/apache/arrow-rs/issues/10308"
64 )]
65 pub fn with_runtime(self, handle: Handle) -> Self {
66 Self {
67 runtime: Some(handle),
68 ..self
69 }
70 }
71
72 fn spawn<F, O, E>(&self, f: F) -> BoxFuture<'_, Result<O, AvroError>>
73 where
74 F: for<'a> FnOnce(&'a Arc<dyn ObjectStore>, &'a Path) -> BoxFuture<'a, Result<O, E>>
75 + Send
76 + 'static,
77 O: Send + 'static,
78 E: Into<AvroError> + Send + 'static,
79 {
80 match &self.runtime {
81 Some(handle) => {
82 let path = self.path.clone();
83 let store = Arc::clone(&self.store);
84 handle
85 .spawn(async move { f(&store, &path).await })
86 .map_ok_or_else(
87 |e| match e.try_into_panic() {
88 Err(e) => Err(AvroError::External(Box::new(e))),
89 Ok(p) => std::panic::resume_unwind(p),
90 },
91 |res| res.map_err(Into::into),
92 )
93 .boxed()
94 }
95 None => f(&self.store, &self.path).map_err(Into::into).boxed(),
96 }
97 }
98}
99
100#[expect(deprecated)]
101impl AsyncFileReader for AvroObjectReader {
102 fn get_bytes(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes, AvroError>> {
103 self.spawn(|store, path| async move { store.get_range(path, range).await }.boxed())
104 }
105
106 fn get_byte_ranges(
107 &mut self,
108 ranges: Vec<Range<u64>>,
109 ) -> BoxFuture<'_, Result<Vec<Bytes>, AvroError>>
110 where
111 Self: Send,
112 {
113 self.spawn(|store, path| async move { store.get_ranges(path, &ranges).await }.boxed())
114 }
115}
116
117#[cfg(test)]
118mod tests {
119 use super::*;
120 use object_store::memory::InMemory;
121
122 fn find_object_store_error(err: &AvroError) -> Option<&object_store::Error> {
123 let mut source: Option<&(dyn std::error::Error + 'static)> = Some(err);
124 while let Some(e) = source {
125 if let Some(os) = e.downcast_ref::<object_store::Error>() {
126 return Some(os);
127 }
128 source = e.source();
129 }
130 None
131 }
132
133 #[tokio::test]
134 async fn test_get_bytes_preserves_object_store_error_source() {
135 let store = Arc::new(InMemory::new());
136 #[expect(deprecated)]
137 let mut reader = AvroObjectReader::new(store, Path::from("missing.avro"));
138 let err = reader.get_bytes(0..10).await.unwrap_err();
139 assert!(
140 matches!(
141 find_object_store_error(&err),
142 Some(object_store::Error::NotFound { .. })
143 ),
144 "expected NotFound in source chain, got: {err}"
145 );
146 }
147
148 #[tokio::test]
149 async fn test_get_bytes_on_runtime_preserves_object_store_error_source() {
150 let store = Arc::new(InMemory::new());
151 #[expect(deprecated)]
152 let mut reader = AvroObjectReader::new(store, Path::from("missing.avro"))
153 .with_runtime(Handle::current());
154 let err = reader.get_bytes(0..10).await.unwrap_err();
155 assert!(
156 matches!(
157 find_object_store_error(&err),
158 Some(object_store::Error::NotFound { .. })
159 ),
160 "expected NotFound in source chain, got: {err}"
161 );
162 }
163}