polars_io/ipc/
ipc_reader_async.rs

1use std::sync::Arc;
2
3use object_store::path::Path;
4use polars_arrow::io::ipc::read::{FileMetadata, OutOfSpecKind, get_row_count};
5use polars_buffer::Buffer;
6use polars_core::datatypes::IDX_DTYPE;
7use polars_core::frame::DataFrame;
8use polars_core::runtime::ASYNC;
9use polars_core::schema::{Schema, SchemaExt};
10use polars_error::{PolarsResult, polars_bail, polars_ensure, polars_err, to_compute_err};
11use polars_utils::mmap::MMapSemaphore;
12use polars_utils::pl_path::PlRefPath;
13use polars_utils::pl_str::PlSmallStr;
14
15use crate::RowIndex;
16use crate::cloud::concurrency_config::FetchConfig;
17use crate::cloud::{
18    CloudLocation, CloudOptions, PolarsObjectStore, build_object_store, object_path_from_str,
19};
20use crate::configs::cloud_footer_read_size;
21use crate::file_cache::{FileCacheEntry, init_entries_from_uri_list};
22use crate::predicates::PhysicalIoExpr;
23use crate::prelude::{IpcReader, materialize_projection};
24use crate::shared::SerReader;
25
26/// An Arrow IPC reader implemented on top of PolarsObjectStore.
27pub struct IpcReaderAsync {
28    store: PolarsObjectStore,
29    cache_entry: Arc<FileCacheEntry>,
30    path: Path,
31}
32
33#[derive(Default, Clone)]
34pub struct IpcReadOptions {
35    // Names of the columns to include in the output.
36    projection: Option<Arc<[PlSmallStr]>>,
37
38    // The maximum number of rows to include in the output.
39    row_limit: Option<usize>,
40
41    // Include a column with the row number under the provided name  starting at the provided index.
42    row_index: Option<RowIndex>,
43
44    // Only include rows that pass this predicate.
45    predicate: Option<Arc<dyn PhysicalIoExpr>>,
46}
47
48impl IpcReadOptions {
49    pub fn with_projection(mut self, projection: Option<Arc<[PlSmallStr]>>) -> Self {
50        self.projection = projection;
51        self
52    }
53
54    pub fn with_row_limit(mut self, row_limit: impl Into<Option<usize>>) -> Self {
55        self.row_limit = row_limit.into();
56        self
57    }
58
59    pub fn with_row_index(mut self, row_index: impl Into<Option<RowIndex>>) -> Self {
60        self.row_index = row_index.into();
61        self
62    }
63
64    pub fn with_predicate(mut self, predicate: impl Into<Option<Arc<dyn PhysicalIoExpr>>>) -> Self {
65        self.predicate = predicate.into();
66        self
67    }
68}
69
70impl IpcReaderAsync {
71    pub async fn from_uri(
72        uri: PlRefPath,
73        cloud_options: Option<&CloudOptions>,
74    ) -> PolarsResult<IpcReaderAsync> {
75        let cache_entry =
76            init_entries_from_uri_list([uri.clone()].into_iter(), cloud_options).await?[0].clone();
77        let (CloudLocation { prefix, .. }, store) =
78            build_object_store(uri, cloud_options, false).await?;
79
80        let path = object_path_from_str(&prefix)?;
81
82        Ok(Self {
83            store,
84            cache_entry,
85            path,
86        })
87    }
88
89    pub async fn metadata(&self) -> PolarsResult<FileMetadata> {
90        // A suffix request returns the file size along with the tail, saving a round trip to
91        // request the size separately. The tail usually contains the whole footer.
92        let (tail, file_size) = self
93            .store
94            .get_suffix(
95                &self.path,
96                cloud_footer_read_size(),
97                FetchConfig::random_access(),
98            )
99            .await?;
100
101        polars_ensure!(
102            tail.len() >= FOOTER_METADATA_SIZE,
103            ComputeError: "ipc file size is smaller than the minimum"
104        );
105
106        let footer_size = deserialize_footer_metadata(
107            tail[tail.len() - FOOTER_METADATA_SIZE..]
108                .try_into()
109                .map_err(to_compute_err)?,
110        )?;
111        let footer_len = FOOTER_METADATA_SIZE + footer_size;
112
113        polars_ensure!(
114            file_size >= footer_len,
115            ComputeError: "invalid ipc footer metadata: footer size too large"
116        );
117
118        let footer = if tail.len() >= footer_len {
119            let offset = tail.len() - footer_len;
120            tail.sliced(offset..)
121        } else {
122            // Only fetch the part of the footer that is not in the tail.
123            let missing = self
124                .store
125                .get_range(
126                    &self.path,
127                    file_size - footer_len..file_size - tail.len(),
128                    FetchConfig::random_access(),
129                )
130                .await?;
131            let mut footer = Vec::with_capacity(footer_len);
132            footer.extend_from_slice(&missing);
133            footer.extend_from_slice(&tail);
134            Buffer::from_vec(footer)
135        };
136
137        polars_arrow::io::ipc::read::deserialize_footer(
138            footer.as_ref(),
139            footer_size.try_into().map_err(to_compute_err)?,
140        )
141    }
142
143    pub async fn data(
144        &self,
145        metadata: Option<&FileMetadata>,
146        options: IpcReadOptions,
147        verbose: bool,
148    ) -> PolarsResult<DataFrame> {
149        // TODO: Only download what is needed rather than the entire file by
150        // making use of the projection, row limit, predicate and such.
151        let file = ASYNC.block_in_place(|| self.cache_entry.try_open_check_latest())?;
152        let bytes = MMapSemaphore::new_from_file(&file).unwrap();
153
154        let projection = match options.projection.as_deref() {
155            Some(projection) => {
156                fn prepare_schema(mut schema: Schema, row_index: Option<&RowIndex>) -> Schema {
157                    if let Some(rc) = row_index {
158                        let _ = schema.insert_at_index(0, rc.name.clone(), IDX_DTYPE);
159                    }
160                    schema
161                }
162
163                // Retrieve the metadata for the schema so we can map column names to indices.
164                let fetched_metadata;
165                let metadata = if let Some(metadata) = metadata {
166                    metadata
167                } else {
168                    // This branch is  happens when _metadata is None, which can happen if we Deserialize the execution plan.
169                    fetched_metadata = self.metadata().await?;
170                    &fetched_metadata
171                };
172
173                let schema = prepare_schema(
174                    Schema::from_arrow_schema(metadata.schema.as_ref()),
175                    options.row_index.as_ref(),
176                );
177
178                let hive_partitions = None;
179
180                materialize_projection(
181                    Some(projection),
182                    &schema,
183                    hive_partitions,
184                    options.row_index.is_some(),
185                )
186            },
187            None => None,
188        };
189
190        let reader = <IpcReader<_> as SerReader<_>>::new(std::io::Cursor::new(bytes.as_ref()))
191            .with_row_index(options.row_index)
192            .with_n_rows(options.row_limit)
193            .with_projection(projection);
194        reader.finish_with_scan_ops(options.predicate, verbose)
195    }
196
197    pub async fn count_rows(&self, _metadata: Option<&FileMetadata>) -> PolarsResult<i64> {
198        // TODO: Only download what is needed rather than the entire file by
199        // making use of the projection, row limit, predicate and such.
200        let file = ASYNC.block_in_place(|| self.cache_entry.try_open_check_latest())?;
201        let bytes = MMapSemaphore::new_from_file(&file).unwrap();
202        get_row_count(&mut std::io::Cursor::new(bytes.as_ref()))
203    }
204}
205
206const FOOTER_METADATA_SIZE: usize = 10;
207
208// TODO: Move to polars-arrow and deduplicate parsing of footer metadata in
209// sync and async readers.
210fn deserialize_footer_metadata(bytes: [u8; FOOTER_METADATA_SIZE]) -> PolarsResult<usize> {
211    let footer_size: usize =
212        i32::from_le_bytes(bytes[0..4].try_into().unwrap_or_else(|_| unreachable!()))
213            .try_into()
214            .map_err(|_| polars_err!(oos = OutOfSpecKind::NegativeFooterLength))?;
215
216    if &bytes[4..] != b"ARROW1" {
217        polars_bail!(oos = OutOfSpecKind::InvalidFooter);
218    }
219
220    Ok(footer_size)
221}
polars_io/ipc/ipc_reader_async.rs

polars_io/ipc/
ipc_reader_async.rs
