Skip to main content

arrow_avro/reader/async_reader/
store.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 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/// An implementation of an AsyncFileReader using the [`ObjectStore`] API.
31#[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    /// Creates a new [`Self`] from a store implementation and file location.
45    pub fn new(store: Arc<dyn ObjectStore>, path: Path) -> Self {
46        Self {
47            store,
48            path,
49            runtime: None,
50        }
51    }
52
53    /// Perform IO on the provided tokio runtime
54    ///
55    /// Tokio is a cooperative scheduler, and relies on tasks yielding in a timely manner
56    /// to service IO. Therefore, running IO and CPU-bound tasks, such as avro decoding,
57    /// on the same tokio runtime can lead to degraded throughput, dropped connections and
58    /// other issues. For more information see [here].
59    ///
60    /// [here]: https://www.influxdata.com/blog/using-rustlangs-async-tokio-runtime-for-cpu-bound-tasks/
61    #[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}