Skip to main content

parquet/arrow/async_reader/
store.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::{ops::Range, sync::Arc};
6
7use crate::arrow::arrow_reader::ArrowReaderOptions;
8use crate::arrow::async_reader::{AsyncFileReader, MetadataSuffixFetch};
9use crate::errors::{ParquetError, Result};
10use crate::file::metadata::{PageIndexPolicy, ParquetMetaData, ParquetMetaDataReader};
11use bytes::Bytes;
12use futures::{FutureExt, TryFutureExt, future::BoxFuture};
13use object_store::ObjectStoreExt;
14use object_store::{GetOptions, GetRange};
15use object_store::{ObjectStore, path::Path};
16use tokio::runtime::Handle;
17/// Reads Parquet files in object storage using [`ObjectStore`].
18///
19/// ```no_run
20/// # use std::io::stdout;
21/// # use std::sync::Arc;
22/// # use object_store::azure::MicrosoftAzureBuilder;
23/// # use object_store::{ObjectStore, ObjectStoreExt};
24/// # use object_store::path::Path;
25/// # use parquet::arrow::async_reader::ParquetObjectReader;
26/// # use parquet::arrow::ParquetRecordBatchStreamBuilder;
27/// # use parquet::schema::printer::print_parquet_metadata;
28/// # async fn run() {
29/// // Populate configuration from environment
30/// let storage_container = Arc::new(MicrosoftAzureBuilder::from_env().build().unwrap());
31/// let location = Path::from("path/to/blob.parquet");
32/// let meta = storage_container.head(&location).await.unwrap();
33/// println!("Found Blob with {}B at {}", meta.size, meta.location);
34///
35/// // Show Parquet metadata
36/// let reader = ParquetObjectReader::new(storage_container, meta.location).with_file_size(meta.size);
37/// let builder = ParquetRecordBatchStreamBuilder::new(reader).await.unwrap();
38/// print_parquet_metadata(&mut stdout(), builder.metadata());
39/// # }
40/// ```
41#[deprecated(
42    since = "59.2.0",
43    note = "Implement `AsyncFileReader` directly instead; see the example on the `AsyncFileReader` trait documentation and `parquet/examples/object_store.rs`. Use `SpawnedReader` to perform I/O on a dedicated runtime. See https://github.com/apache/arrow-rs/issues/10308"
44)]
45#[derive(Clone, Debug)]
46pub struct ParquetObjectReader {
47    store: Arc<dyn ObjectStore>,
48    path: Path,
49    file_size: Option<u64>,
50    metadata_size_hint: Option<usize>,
51    preload_column_index: bool,
52    preload_offset_index: bool,
53    runtime: Option<Handle>,
54}
55
56#[expect(deprecated)]
57impl ParquetObjectReader {
58    /// Creates a new [`ParquetObjectReader`] for the provided [`ObjectStore`] and [`Path`].
59    pub fn new(store: Arc<dyn ObjectStore>, path: Path) -> Self {
60        Self {
61            store,
62            path,
63            file_size: None,
64            metadata_size_hint: None,
65            preload_column_index: false,
66            preload_offset_index: false,
67            runtime: None,
68        }
69    }
70
71    /// Provide a hint as to the size of the parquet file's footer,
72    /// see [`ParquetMetaDataReader::with_prefetch_hint`]
73    pub fn with_footer_size_hint(self, hint: usize) -> Self {
74        Self {
75            metadata_size_hint: Some(hint),
76            ..self
77        }
78    }
79
80    /// Provide the byte size of this file.
81    ///
82    /// If provided, the file size will ensure that only bounded range requests are used. If file
83    /// size is not provided, the reader will use suffix range requests to fetch the metadata.
84    ///
85    /// Providing this size up front is an important optimization to avoid extra calls when the
86    /// underlying store does not support suffix range requests.
87    ///
88    /// The file size can be obtained using [`ObjectStore::list`] or [`ObjectStoreExt::head`].
89    pub fn with_file_size(self, file_size: u64) -> Self {
90        Self {
91            file_size: Some(file_size),
92            ..self
93        }
94    }
95
96    /// Whether to load the Column Index as part of [`Self::get_metadata`]
97    ///
98    /// Note: This setting may be overridden by [`ArrowReaderOptions`] `page_index_policy`.
99    /// If `page_index_policy` is `Optional` or `Required`, it will take precedence
100    /// over this preload flag. When it is `Skip` (default), this flag is used.
101    pub fn with_preload_column_index(self, preload_column_index: bool) -> Self {
102        Self {
103            preload_column_index,
104            ..self
105        }
106    }
107
108    /// Whether to load the Offset Index as part of [`Self::get_metadata`]
109    ///
110    /// Note: This setting may be overridden by [`ArrowReaderOptions`] `page_index_policy`.
111    /// If `page_index_policy` is `Optional` or `Required`, it will take precedence
112    /// over this preload flag. When it is `Skip` (default), this flag is used.
113    pub fn with_preload_offset_index(self, preload_offset_index: bool) -> Self {
114        Self {
115            preload_offset_index,
116            ..self
117        }
118    }
119
120    /// Perform IO on the provided tokio runtime
121    ///
122    /// Tokio is a cooperative scheduler, and relies on tasks yielding in a timely manner
123    /// to service IO. Therefore, running IO and CPU-bound tasks, such as parquet decoding,
124    /// on the same tokio runtime can lead to degraded throughput, dropped connections and
125    /// other issues. For more information see [here].
126    ///
127    /// [here]: https://www.influxdata.com/blog/using-rustlangs-async-tokio-runtime-for-cpu-bound-tasks/
128    #[deprecated(
129        since = "59.2.0",
130        note = "Wrap the reader in a `SpawnedReader` instead, e.g. `SpawnedReader::new(reader, handle)`. See https://github.com/apache/arrow-rs/issues/10308"
131    )]
132    pub fn with_runtime(self, handle: Handle) -> Self {
133        Self {
134            runtime: Some(handle),
135            ..self
136        }
137    }
138
139    fn spawn<F, O, E>(&self, f: F) -> BoxFuture<'_, Result<O>>
140    where
141        F: for<'a> FnOnce(&'a Arc<dyn ObjectStore>, &'a Path) -> BoxFuture<'a, Result<O, E>>
142            + Send
143            + 'static,
144        O: Send + 'static,
145        E: Into<ParquetError> + Send + 'static,
146    {
147        match &self.runtime {
148            Some(handle) => {
149                let path = self.path.clone();
150                let store = Arc::clone(&self.store);
151                handle
152                    .spawn(async move { f(&store, &path).await })
153                    .map_ok_or_else(
154                        |e| match e.try_into_panic() {
155                            Err(e) => Err(ParquetError::External(Box::new(e))),
156                            Ok(p) => std::panic::resume_unwind(p),
157                        },
158                        |res| res.map_err(|e| e.into()),
159                    )
160                    .boxed()
161            }
162            None => f(&self.store, &self.path).map_err(|e| e.into()).boxed(),
163        }
164    }
165}
166
167#[expect(deprecated)]
168impl MetadataSuffixFetch for &mut ParquetObjectReader {
169    fn fetch_suffix(&mut self, suffix: usize) -> BoxFuture<'_, Result<Bytes>> {
170        let options = GetOptions {
171            range: Some(GetRange::Suffix(suffix as u64)),
172            ..Default::default()
173        };
174        self.spawn(|store, path| {
175            async move {
176                let resp = store.get_opts(path, options).await?;
177                Ok::<_, ParquetError>(resp.bytes().await?)
178            }
179            .boxed()
180        })
181    }
182}
183
184#[expect(deprecated)]
185impl AsyncFileReader for ParquetObjectReader {
186    fn get_bytes(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes>> {
187        self.spawn(|store, path| store.get_range(path, range).boxed())
188    }
189
190    fn get_byte_ranges(&mut self, ranges: Vec<Range<u64>>) -> BoxFuture<'_, Result<Vec<Bytes>>>
191    where
192        Self: Send,
193    {
194        self.spawn(|store, path| async move { store.get_ranges(path, &ranges).await }.boxed())
195    }
196
197    // This method doesn't directly call `self.spawn` because all of the IO that is done down the
198    // line due to this method call is done through `self.get_bytes` and/or `self.get_byte_ranges`.
199    // When `self` is passed into `ParquetMetaDataReader::load_and_finish`, it treats it as
200    // an `impl MetadataFetch` and calls those methods to get data from it. Due to `Self`'s impl of
201    // `AsyncFileReader`, the calls to `MetadataFetch::fetch` are just delegated to
202    // `Self::get_bytes`.
203    fn get_metadata<'a>(
204        &'a mut self,
205        options: Option<&'a ArrowReaderOptions>,
206    ) -> BoxFuture<'a, Result<Arc<ParquetMetaData>>> {
207        Box::pin(async move {
208            let metadata_opts = options.map(|o| o.metadata_options().clone());
209            let mut metadata = ParquetMetaDataReader::new()
210                .with_metadata_options(metadata_opts)
211                .with_column_index_policy(PageIndexPolicy::from(self.preload_column_index))
212                .with_offset_index_policy(PageIndexPolicy::from(self.preload_offset_index))
213                .with_prefetch_hint(self.metadata_size_hint);
214
215            #[cfg(feature = "encryption")]
216            if let Some(options) = options {
217                metadata = metadata.with_decryption_properties(
218                    options.file_decryption_properties.as_ref().map(Arc::clone),
219                );
220            }
221
222            // Override page index policies from ArrowReaderOptions if specified and not Skip.
223            // When page_index_policy is Skip (default), use the reader's preload flags.
224            // When page_index_policy is Optional or Required, override the preload flags
225            // to ensure the specified policy takes precedence.
226            if let Some(options) = options
227                && (options.column_index_policy() != PageIndexPolicy::Skip
228                    || options.offset_index_policy() != PageIndexPolicy::Skip)
229            {
230                metadata = metadata
231                    .with_column_index_policy(options.column_index_policy())
232                    .with_offset_index_policy(options.offset_index_policy());
233            }
234            // TODO: https://github.com/apache/arrow-rs/issues/11395 may deprecate this API. Decide
235            // at that time if masks should be added to ParquetObjectReader.
236            if let Some(options) = options {
237                if !options.column_index_mask().is_all() {
238                    metadata = metadata.with_column_index_mask(options.column_index_mask().clone());
239                }
240                if !options.offset_index_mask().is_all() {
241                    metadata = metadata.with_offset_index_mask(options.offset_index_mask().clone());
242                }
243            }
244
245            let metadata = if let Some(file_size) = self.file_size {
246                metadata.load_and_finish(self, file_size).await?
247            } else {
248                metadata.load_via_suffix_and_finish(self).await?
249            };
250
251            Ok(Arc::new(metadata))
252        })
253    }
254}
255
256#[cfg(test)]
257#[expect(deprecated)]
258mod tests {
259    use crate::arrow::async_reader::ArrowReaderOptions;
260    use crate::file::metadata::PageIndexPolicy;
261    use std::sync::{
262        Arc,
263        atomic::{AtomicUsize, Ordering},
264    };
265
266    use futures::TryStreamExt;
267
268    use crate::arrow::ParquetRecordBatchStreamBuilder;
269    use crate::arrow::async_reader::{AsyncFileReader, ParquetObjectReader};
270    use crate::errors::ParquetError;
271    use arrow::util::test_util::parquet_test_data;
272    use futures::FutureExt;
273    use object_store::local::LocalFileSystem;
274    use object_store::path::Path;
275    use object_store::{ObjectMeta, ObjectStore, ObjectStoreExt};
276
277    async fn get_meta_store() -> (ObjectMeta, Arc<dyn ObjectStore>) {
278        let res = parquet_test_data();
279        let store = LocalFileSystem::new_with_prefix(res).unwrap();
280
281        let meta = store
282            .head(&Path::from("alltypes_plain.parquet"))
283            .await
284            .unwrap();
285
286        (meta, Arc::new(store) as Arc<dyn ObjectStore>)
287    }
288
289    async fn get_meta_store_with_page_index() -> (ObjectMeta, Arc<dyn ObjectStore>) {
290        let res = parquet_test_data();
291        let store = LocalFileSystem::new_with_prefix(res).unwrap();
292
293        let meta = store
294            .head(&Path::from("alltypes_tiny_pages_plain.parquet"))
295            .await
296            .unwrap();
297
298        (meta, Arc::new(store) as Arc<dyn ObjectStore>)
299    }
300
301    #[tokio::test]
302    async fn test_simple() {
303        let (meta, store) = get_meta_store().await;
304        let object_reader =
305            ParquetObjectReader::new(store, meta.location).with_file_size(meta.size);
306
307        let builder = ParquetRecordBatchStreamBuilder::new(object_reader)
308            .await
309            .unwrap();
310        let batches: Vec<_> = builder.build().unwrap().try_collect().await.unwrap();
311
312        assert_eq!(batches.len(), 1);
313        assert_eq!(batches[0].num_rows(), 8);
314    }
315
316    #[tokio::test]
317    async fn test_simple_without_file_length() {
318        let (meta, store) = get_meta_store().await;
319        let object_reader = ParquetObjectReader::new(store, meta.location);
320
321        let builder = ParquetRecordBatchStreamBuilder::new(object_reader)
322            .await
323            .unwrap();
324        let batches: Vec<_> = builder.build().unwrap().try_collect().await.unwrap();
325
326        assert_eq!(batches.len(), 1);
327        assert_eq!(batches[0].num_rows(), 8);
328    }
329
330    #[tokio::test]
331    async fn test_not_found() {
332        let (mut meta, store) = get_meta_store().await;
333        meta.location = Path::from("I don't exist.parquet");
334
335        let object_reader =
336            ParquetObjectReader::new(store, meta.location).with_file_size(meta.size);
337        // Cannot use unwrap_err as ParquetRecordBatchStreamBuilder: !Debug
338        match ParquetRecordBatchStreamBuilder::new(object_reader).await {
339            Ok(_) => panic!("expected failure"),
340            Err(e) => {
341                let err = e.to_string();
342                assert!(err.contains("I don't exist.parquet not found:"), "{err}");
343            }
344        }
345    }
346
347    #[tokio::test]
348    async fn test_runtime_is_used() {
349        let num_actions = Arc::new(AtomicUsize::new(0));
350
351        let (a1, a2) = (num_actions.clone(), num_actions.clone());
352        let rt = tokio::runtime::Builder::new_multi_thread()
353            .on_thread_park(move || {
354                a1.fetch_add(1, Ordering::Relaxed);
355            })
356            .on_thread_unpark(move || {
357                a2.fetch_add(1, Ordering::Relaxed);
358            })
359            .build()
360            .unwrap();
361
362        let (meta, store) = get_meta_store().await;
363
364        let initial_actions = num_actions.load(Ordering::Relaxed);
365
366        let reader = ParquetObjectReader::new(store, meta.location)
367            .with_file_size(meta.size)
368            .with_runtime(rt.handle().clone());
369
370        let builder = ParquetRecordBatchStreamBuilder::new(reader).await.unwrap();
371        let batches: Vec<_> = builder.build().unwrap().try_collect().await.unwrap();
372
373        // Just copied these assert_eqs from the `test_simple` above
374        assert_eq!(batches.len(), 1);
375        assert_eq!(batches[0].num_rows(), 8);
376
377        assert!(num_actions.load(Ordering::Relaxed) - initial_actions > 0);
378
379        // Runtimes have to be dropped in blocking contexts, so we need to move this one to a new
380        // blocking thread to drop it.
381        tokio::runtime::Handle::current().spawn_blocking(move || drop(rt));
382    }
383
384    /// Unit test that `ParquetObjectReader::spawn`spawns on the provided runtime
385    #[tokio::test]
386    async fn test_runtime_thread_id_different() {
387        let rt = tokio::runtime::Builder::new_multi_thread()
388            .worker_threads(1)
389            .build()
390            .unwrap();
391
392        let (meta, store) = get_meta_store().await;
393
394        let reader = ParquetObjectReader::new(store, meta.location)
395            .with_file_size(meta.size)
396            .with_runtime(rt.handle().clone());
397
398        let current_id = std::thread::current().id();
399
400        let other_id = reader
401            .spawn(|_, _| async move { Ok::<_, ParquetError>(std::thread::current().id()) }.boxed())
402            .await
403            .unwrap();
404
405        assert_ne!(current_id, other_id);
406
407        tokio::runtime::Handle::current().spawn_blocking(move || drop(rt));
408    }
409
410    #[tokio::test]
411    async fn io_fails_on_shutdown_runtime() {
412        let rt = tokio::runtime::Builder::new_multi_thread()
413            .worker_threads(1)
414            .build()
415            .unwrap();
416
417        let (meta, store) = get_meta_store().await;
418
419        let mut reader = ParquetObjectReader::new(store, meta.location)
420            .with_file_size(meta.size)
421            .with_runtime(rt.handle().clone());
422
423        rt.shutdown_background();
424
425        let err = reader.get_bytes(0..1).await.unwrap_err().to_string();
426
427        assert!(err.contains("was cancelled"));
428    }
429
430    #[tokio::test]
431    async fn test_page_index_policy_skip_uses_preload_true() {
432        let (meta, store) = get_meta_store_with_page_index().await;
433
434        // Create reader with preload flags set to true
435        let mut reader = ParquetObjectReader::new(store.clone(), meta.location.clone())
436            .with_file_size(meta.size)
437            .with_preload_column_index(true)
438            .with_preload_offset_index(true);
439
440        // Create options with page_index_policy set to Skip (default)
441        let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Skip);
442
443        // Get metadata - Skip means use reader's preload flags (true)
444        let metadata = reader.get_metadata(Some(&options)).await.unwrap();
445
446        // With preload=true, indexes should be loaded since the test file has them
447        assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
448    }
449
450    #[tokio::test]
451    async fn test_page_index_policy_optional_overrides_preload_false() {
452        let (meta, store) = get_meta_store_with_page_index().await;
453
454        // Create reader with preload flags set to false
455        let mut reader = ParquetObjectReader::new(store.clone(), meta.location.clone())
456            .with_file_size(meta.size)
457            .with_preload_column_index(false)
458            .with_preload_offset_index(false);
459
460        // Create options with page_index_policy set to Optional
461        let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Optional);
462
463        // Get metadata - Optional overrides preload flags and attempts to load indexes
464        let metadata = reader.get_metadata(Some(&options)).await.unwrap();
465
466        // With Optional policy, it will TRY to load indexes but won't fail if they don't exist
467        // The test file has page indexes, so they will be some
468        assert!(metadata.page_index().is_some());
469    }
470
471    #[tokio::test]
472    async fn test_page_index_policy_optional_vs_skip() {
473        let (meta, store) = get_meta_store_with_page_index().await;
474
475        // Test 1: preload=false + Skip policy -> uses preload flags (false)
476        let mut reader1 = ParquetObjectReader::new(store.clone(), meta.location.clone())
477            .with_file_size(meta.size)
478            .with_preload_column_index(false)
479            .with_preload_offset_index(false);
480
481        let options1 = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Skip);
482        let metadata1 = reader1.get_metadata(Some(&options1)).await.unwrap();
483
484        // Test 2: preload=false + Optional policy -> overrides to try loading
485        let mut reader2 = ParquetObjectReader::new(store.clone(), meta.location.clone())
486            .with_file_size(meta.size)
487            .with_preload_column_index(false)
488            .with_preload_offset_index(false);
489
490        let options2 = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Optional);
491        let metadata2 = reader2.get_metadata(Some(&options2)).await.unwrap();
492
493        // Both should succeed (no panic/error)
494        // metadata1 (Skip) uses preload=false -> Skip policy
495        // metadata2 (Optional) overrides preload=false -> Optional policy
496        assert!(metadata1.page_index().is_none());
497        assert!(metadata2.page_index().is_some());
498    }
499
500    #[tokio::test]
501    async fn test_page_index_policy_no_options_uses_preload() {
502        let (meta, store) = get_meta_store_with_page_index().await;
503
504        // Create reader with preload flags set to true
505        let mut reader = ParquetObjectReader::new(store, meta.location)
506            .with_file_size(meta.size)
507            .with_preload_column_index(true)
508            .with_preload_offset_index(true);
509
510        // Get metadata without options - should use reader's preload flags
511        let metadata = reader.get_metadata(None).await.unwrap();
512
513        // With no options provided, preload flags (true) should be respected
514        // and converted to Optional policy internally (preload=true -> Optional)
515        // The test file has page indexes, so they will be some
516        assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
517    }
518}