polars_io/ipc/
ipc_reader_async.rs

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

polars_io/ipc/
ipc_reader_async.rs
