parquet/arrow/async_reader/
store.rs1use 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#[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 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 pub fn with_footer_size_hint(self, hint: usize) -> Self {
74 Self {
75 metadata_size_hint: Some(hint),
76 ..self
77 }
78 }
79
80 pub fn with_file_size(self, file_size: u64) -> Self {
90 Self {
91 file_size: Some(file_size),
92 ..self
93 }
94 }
95
96 pub fn with_preload_column_index(self, preload_column_index: bool) -> Self {
102 Self {
103 preload_column_index,
104 ..self
105 }
106 }
107
108 pub fn with_preload_offset_index(self, preload_offset_index: bool) -> Self {
114 Self {
115 preload_offset_index,
116 ..self
117 }
118 }
119
120 #[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 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 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 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 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 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 tokio::runtime::Handle::current().spawn_blocking(move || drop(rt));
382 }
383
384 #[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 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 let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Skip);
442
443 let metadata = reader.get_metadata(Some(&options)).await.unwrap();
445
446 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 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 let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Optional);
462
463 let metadata = reader.get_metadata(Some(&options)).await.unwrap();
465
466 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 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 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 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 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 let metadata = reader.get_metadata(None).await.unwrap();
512
513 assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
517 }
518}