polars_io/parquet/read/
async_impl.rs

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

polars_io/parquet/read/
async_impl.rs
