polars_io/parquet/read/
async_impl.rs

1//! Read parquet files in parallel from the Object Store without a third party crate.
2
3use object_store::path::Path as ObjectPath;
4use polars_arrow::datatypes::ArrowSchemaRef;
5use polars_buffer::Buffer;
6use polars_core::prelude::*;
7use polars_parquet::parquet::encryption::decrypt::FileDecryptionProperties;
8use polars_parquet::parquet::error::ParquetError;
9use polars_parquet::parquet::read::{deserialize_metadata, deserialize_num_rows};
10use polars_parquet::parquet::{ENCRYPTED_PARQUET_MAGIC, FOOTER_SIZE, PARQUET_MAGIC};
11use polars_utils::pl_path::PlRefPath;
12
13use crate::cloud::concurrency_config::FetchConfig;
14use crate::cloud::{
15    CloudLocation, CloudOptions, PolarsObjectStore, build_object_store, object_path_from_str,
16};
17use crate::configs::cloud_footer_read_size;
18use crate::parquet::encryption::{ParquetDecryption, file_decryption_properties};
19use crate::parquet::metadata::FileMetadataRef;
20
21pub struct ParquetObjectStore {
22    store: PolarsObjectStore,
23    path: ObjectPath,
24    metadata: Option<FileMetadataRef>,
25    schema: Option<ArrowSchemaRef>,
26    decryption_properties: Option<Arc<FileDecryptionProperties>>,
27}
28
29impl ParquetObjectStore {
30    pub async fn from_uri(
31        uri: PlRefPath,
32        options: Option<&CloudOptions>,
33        metadata: Option<FileMetadataRef>,
34    ) -> PolarsResult<Self> {
35        let (CloudLocation { prefix, .. }, store) = build_object_store(uri, options, false).await?;
36        let path = object_path_from_str(&prefix)?;
37
38        Ok(ParquetObjectStore {
39            store,
40            path,
41            metadata,
42            schema: None,
43            decryption_properties: None,
44        })
45    }
46
47    /// Set how to decrypt an encrypted parquet file.
48    pub fn with_decryption(mut self, decryption: Option<&ParquetDecryption>) -> PolarsResult<Self> {
49        self.decryption_properties = file_decryption_properties(decryption)?;
50        Ok(self)
51    }
52
53    /// Number of rows in the parquet file.
54    pub async fn num_rows(&mut self) -> PolarsResult<usize> {
55        let metadata = self.get_metadata().await?;
56        Ok(metadata.num_rows)
57    }
58
59    /// Fetch and memoize the metadata of the parquet file.
60    pub async fn get_metadata(&mut self) -> PolarsResult<&FileMetadataRef> {
61        if self.metadata.is_none() {
62            let footer = fetch_footer_bytes(&self.store, &self.path).await?;
63            self.metadata = Some(Arc::new(deserialize_metadata(
64                footer,
65                self.decryption_properties.as_ref(),
66            )?));
67        }
68        Ok(self.metadata.as_ref().unwrap())
69    }
70
71    /// Decode only `FileMetaData.num_rows` from the remote footer.
72    /// Not memoized. Used by `RowCounts` resolve mode.
73    pub async fn num_rows_only(&mut self) -> PolarsResult<i64> {
74        let footer = fetch_footer_bytes(&self.store, &self.path).await?;
75        Ok(deserialize_num_rows(
76            footer,
77            self.decryption_properties.as_ref(),
78        )?)
79    }
80
81    pub async fn schema(&mut self) -> PolarsResult<ArrowSchemaRef> {
82        self.schema = Some(match self.schema.as_ref() {
83            Some(schema) => Arc::clone(schema),
84            None => {
85                let metadata = self.get_metadata().await?;
86                let arrow_schema = polars_parquet::arrow::read::infer_schema(metadata)?;
87                Arc::new(arrow_schema)
88            },
89        });
90
91        Ok(self.schema.clone().unwrap())
92    }
93}
94
95fn read_n<const N: usize>(reader: &mut &[u8]) -> Option<[u8; N]> {
96    if N <= reader.len() {
97        let (head, tail) = reader.split_at(N);
98        *reader = tail;
99        Some(head.try_into().unwrap())
100    } else {
101        None
102    }
103}
104
105fn read_i32le(reader: &mut &[u8]) -> Option<i32> {
106    read_n(reader).map(i32::from_le_bytes)
107}
108
109/// Speculatively read `cloud_footer_read_size()` bytes from the tail. If the
110/// footer fits in the prefetch (the common case), we're done in one range
111/// request; otherwise re-fetch the full footer. Mirrors the sync
112/// `fetch_footer_buf` strategy.
113async fn fetch_footer_bytes(
114    store: &PolarsObjectStore,
115    path: &ObjectPath,
116) -> PolarsResult<Buffer<u8>> {
117    let out_of_spec = |msg: &str| ParquetError::OutOfSpec(msg.to_string());
118
119    let (prefetched, file_byte_length) = store
120        .get_suffix(path, cloud_footer_read_size(), FetchConfig::random_access())
121        .await?;
122
123    if prefetched.len() < FOOTER_SIZE as usize {
124        return Err(out_of_spec("not enough bytes to contain parquet footer").into());
125    }
126
127    // Trailing 8 bytes: footer size (i32 LE) + magic.
128    let footer_byte_length: usize = {
129        let tail_start = prefetched.len() - FOOTER_SIZE as usize;
130        let reader = &mut &prefetched.as_ref()[tail_start..];
131        let footer_byte_size = read_i32le(reader).unwrap();
132        let magic = read_n(reader).unwrap();
133        debug_assert!(reader.is_empty());
134        if magic != PARQUET_MAGIC && magic != ENCRYPTED_PARQUET_MAGIC {
135            return Err(out_of_spec("incorrect magic in parquet footer").into());
136        }
137        footer_byte_size
138            .try_into()
139            .map_err(|_| out_of_spec("negative footer byte length"))?
140    };
141
142    let footer_len = FOOTER_SIZE as usize + footer_byte_length;
143    if footer_len <= prefetched.len() {
144        // Common case: footer already in the prefetch; zero extra round trips.
145        let footer = prefetched.clone().sliced((prefetched.len() - footer_len)..);
146
147        // The footer is held for the lifetime of the metadata, so copy it out rather than
148        // pin the whole prefetch for a fraction of its bytes.
149        return Ok(if prefetched.len() >= 2 * footer_len {
150            Buffer::from_vec(footer.to_vec())
151        } else {
152            footer
153        });
154    }
155
156    // Fallback: footer larger than the prefetch; re-fetch the full footer.
157    store
158        .get_range(
159            path,
160            file_byte_length
161                .checked_sub(footer_len)
162                .ok_or_else(|| out_of_spec("not enough bytes to contain parquet footer"))?
163                ..file_byte_length,
164            FetchConfig::random_access(),
165        )
166        .await
167}
polars_io/parquet/read/async_impl.rs

polars_io/parquet/read/
async_impl.rs
