Skip to main content

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