parquet/arrow/async_reader/
spawn.rs

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

parquet/arrow/async_reader/
spawn.rs
