parquet/arrow/async_reader/
spawn.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 std::future::Future;
19use std::ops::Range;
20use std::sync::Arc;
21
22use bytes::Bytes;
23use futures::future::BoxFuture;
24use futures::{FutureExt, TryFutureExt};
25use tokio::runtime::Handle;
26
27use crate::arrow::arrow_reader::ArrowReaderOptions;
28use crate::arrow::async_reader::{AsyncFileReader, MetadataSuffixFetch};
29use crate::errors::{ParquetError, Result};
30use crate::file::metadata::ParquetMetaData;
31
32/// An [`AsyncFileReader`] that performs I/O on a separate tokio runtime.
33///
34/// Tokio is a cooperative scheduler, and relies on tasks yielding in a timely
35/// manner to service IO. Therefore, running IO and CPU-bound tasks, such as
36/// parquet decoding, on the same tokio runtime can lead to degraded
37/// throughput, dropped connections and other issues. For more information see
38/// [here].
39///
40/// This wrapper spawns each operation of the inner reader onto the provided
41/// runtime [`Handle`], so that the runtime driving the parquet decoding does
42/// not also drive the I/O.
43///
44/// Note that [`Self::get_metadata`] spawns the entire metadata load, so the
45/// footer is also decoded on the provided runtime.
46///
47/// The inner reader must be [`Clone`] (typically an `Arc`'d handle to some
48/// shared resource) as each spawned task requires a `'static` copy of it.
49///
50/// [here]: https://www.influxdata.com/blog/using-rustlangs-async-tokio-runtime-for-cpu-bound-tasks/
51#[derive(Clone, Debug)]
52pub struct SpawnedReader<R> {
53    inner: R,
54    handle: Handle,
55}
56
57impl<R> SpawnedReader<R> {
58    /// Creates a new [`SpawnedReader`] that performs the I/O of `inner` on `handle`
59    pub fn new(inner: R, handle: Handle) -> Self {
60        Self { inner, handle }
61    }
62
63    /// Returns the inner reader
64    pub fn into_inner(self) -> R {
65        self.inner
66    }
67}
68
69/// Spawns `fut` on `handle`, propagating panics and mapping task cancellation
70/// to [`ParquetError::External`]
71fn spawn<T>(
72    handle: &Handle,
73    fut: impl Future<Output = Result<T>> + Send + 'static,
74) -> BoxFuture<'static, Result<T>>
75where
76    T: Send + 'static,
77{
78    handle
79        .spawn(fut)
80        .map_ok_or_else(
81            |e| match e.try_into_panic() {
82                Err(e) => Err(ParquetError::External(Box::new(e))),
83                Ok(p) => std::panic::resume_unwind(p),
84            },
85            |res| res,
86        )
87        .boxed()
88}
89
90impl<R> AsyncFileReader for SpawnedReader<R>
91where
92    R: AsyncFileReader + Clone + Send + 'static,
93{
94    fn get_bytes(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes>> {
95        let mut inner = self.inner.clone();
96        spawn(&self.handle, async move { inner.get_bytes(range).await })
97    }
98
99    fn get_byte_ranges(&mut self, ranges: Vec<Range<u64>>) -> BoxFuture<'_, Result<Vec<Bytes>>> {
100        let mut inner = self.inner.clone();
101        spawn(
102            &self.handle,
103            async move { inner.get_byte_ranges(ranges).await },
104        )
105    }
106
107    fn get_metadata<'a>(
108        &'a mut self,
109        options: Option<&'a ArrowReaderOptions>,
110    ) -> BoxFuture<'a, Result<Arc<ParquetMetaData>>> {
111        let mut inner = self.inner.clone();
112        let options = options.cloned();
113        spawn(&self.handle, async move {
114            inner.get_metadata(options.as_ref()).await
115        })
116    }
117}
118
119impl<R> MetadataSuffixFetch for &mut SpawnedReader<R>
120where
121    R: AsyncFileReader + Clone + Send + 'static,
122    for<'a> &'a mut R: MetadataSuffixFetch,
123{
124    fn fetch_suffix(&mut self, suffix: usize) -> BoxFuture<'_, Result<Bytes>> {
125        let mut inner = self.inner.clone();
126        spawn(&self.handle, async move {
127            (&mut inner).fetch_suffix(suffix).await
128        })
129    }
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135    use crate::arrow::ParquetRecordBatchStreamBuilder;
136    use crate::file::metadata::ParquetMetaDataReader;
137    use futures::TryStreamExt;
138    use std::thread::ThreadId;
139
140    /// An in-memory [`AsyncFileReader`] that records the thread each request ran on
141    #[derive(Clone)]
142    struct InMemoryReader {
143        data: Bytes,
144        threads: Arc<std::sync::Mutex<Vec<ThreadId>>>,
145    }
146
147    impl InMemoryReader {
148        fn new(data: Bytes) -> Self {
149            Self {
150                data,
151                threads: Default::default(),
152            }
153        }
154    }
155
156    impl AsyncFileReader for InMemoryReader {
157        fn get_bytes(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes>> {
158            self.threads
159                .lock()
160                .unwrap()
161                .push(std::thread::current().id());
162            let data = self.data.slice(range.start as usize..range.end as usize);
163            futures::future::ready(Ok(data)).boxed()
164        }
165
166        fn get_metadata<'a>(
167            &'a mut self,
168            options: Option<&'a ArrowReaderOptions>,
169        ) -> BoxFuture<'a, Result<Arc<ParquetMetaData>>> {
170            self.threads
171                .lock()
172                .unwrap()
173                .push(std::thread::current().id());
174            let metadata = ParquetMetaDataReader::new()
175                .with_arrow_reader_options(options)
176                .parse_and_finish(&self.data);
177            futures::future::ready(metadata.map(Arc::new)).boxed()
178        }
179    }
180
181    #[tokio::test]
182    async fn test_spawned_reader() {
183        let testdata = arrow::util::test_util::parquet_test_data();
184        let path = format!("{testdata}/alltypes_plain.parquet");
185        let data = Bytes::from(std::fs::read(path).unwrap());
186
187        let rt = tokio::runtime::Builder::new_multi_thread()
188            .worker_threads(1)
189            .build()
190            .unwrap();
191
192        let inner = InMemoryReader::new(data);
193        let threads = inner.threads.clone();
194        let reader = SpawnedReader::new(inner, rt.handle().clone());
195
196        let builder = ParquetRecordBatchStreamBuilder::new(reader).await.unwrap();
197        let batches: Vec<_> = builder.build().unwrap().try_collect().await.unwrap();
198
199        assert_eq!(batches.len(), 1);
200        assert_eq!(batches[0].num_rows(), 8);
201
202        // All I/O must have run on the spawned runtime, not the current one
203        let current_id = std::thread::current().id();
204        let threads = threads.lock().unwrap();
205        assert!(!threads.is_empty());
206        assert!(threads.iter().all(|id| *id != current_id));
207
208        // Runtimes have to be dropped in blocking contexts
209        tokio::runtime::Handle::current().spawn_blocking(move || drop(rt));
210    }
211
212    #[tokio::test]
213    async fn test_spawned_reader_fails_on_shutdown_runtime() {
214        let rt = tokio::runtime::Builder::new_multi_thread()
215            .worker_threads(1)
216            .build()
217            .unwrap();
218
219        let inner = InMemoryReader::new(Bytes::from_static(b"PAR1"));
220        let mut reader = SpawnedReader::new(inner, rt.handle().clone());
221
222        rt.shutdown_background();
223
224        let err = reader.get_bytes(0..1).await.unwrap_err().to_string();
225        assert!(err.contains("was cancelled"), "{err}");
226    }
227}
parquet/arrow/async_reader/spawn.rs

parquet/arrow/async_reader/
spawn.rs
