Skip to main content

polars_io/ipc/
pl_ipc_metadata.rs

1use polars_utils::IdxSize;
2
3pub static POLARS_IPC_METADATA_KEY: &str = "__POLARS_IPC_METADATA";
4
5#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)]
6#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
7pub struct PlIpcMetadata {
8    /// Cumulative length including the current record batch.
9    pub record_batch_cum_len: Vec<IdxSize>,
10}
11
12impl PlIpcMetadata {
13    /// Reads the Polars metadata out of an already parsed IPC footer.
14    ///
15    /// `None` for a file that was not written by Polars, or whose Polars metadata
16    /// cannot be parsed or is inconsistent with the file.
17    pub fn from_ipc_footer(metadata: &polars_arrow::io::ipc::read::FileMetadata) -> Option<Self> {
18        #[cfg(feature = "serde")]
19        {
20            let raw = metadata
21                .custom_metadata
22                .as_ref()?
23                .get(POLARS_IPC_METADATA_KEY)?;
24            serde_json::from_str::<Self>(raw)
25                .ok()
26                // Older writers also counted dictionary batches, ignore such metadata.
27                .filter(|md| md.record_batch_cum_len.len() == metadata.blocks.len())
28        }
29        #[cfg(not(feature = "serde"))]
30        {
31            let _ = metadata;
32            None
33        }
34    }
35
36    /// Total number of rows over all record batches.
37    pub fn num_rows(&self) -> Option<IdxSize> {
38        self.record_batch_cum_len.last().copied()
39    }
40}