polars_arrow/io/ipc/write/
mod.rs

1//! APIs to write to Arrow's IPC format.
2pub(crate) mod common;
3mod schema;
4mod serialize;
5mod stream;
6pub(crate) mod writer;
7
8pub use common::{
9    dictionaries_to_encode, encode_dictionary, encode_new_dictionaries, encode_record_batch,
10    Compression, DictionaryTracker, EncodedData, Record, WriteOptions,
11};
12pub use schema::schema_to_bytes;
13pub use serialize::write;
14use serialize::write_dictionary;
15pub use stream::StreamWriter;
16pub use writer::FileWriter;
17
18pub(crate) mod common_sync;
19
20use super::IpcField;
21use crate::datatypes::{ArrowDataType, Field};
22
23fn default_ipc_field(dtype: &ArrowDataType, current_id: &mut i64) -> IpcField {
24    use crate::datatypes::ArrowDataType::*;
25    match dtype.to_logical_type() {
26        // single child => recurse
27        Map(inner, ..) | FixedSizeList(inner, _) | LargeList(inner) | List(inner) => IpcField {
28            fields: vec![default_ipc_field(inner.dtype(), current_id)],
29            dictionary_id: None,
30        },
31        // multiple children => recurse
32        Struct(fields) => IpcField {
33            fields: fields
34                .iter()
35                .map(|f| default_ipc_field(f.dtype(), current_id))
36                .collect(),
37            dictionary_id: None,
38        },
39        // multiple children => recurse
40        Union(u) => IpcField {
41            fields: u
42                .fields
43                .iter()
44                .map(|f| default_ipc_field(f.dtype(), current_id))
45                .collect(),
46            dictionary_id: None,
47        },
48        // dictionary => current_id
49        Dictionary(_, dtype, _) => {
50            let dictionary_id = Some(*current_id);
51            *current_id += 1;
52            IpcField {
53                fields: vec![default_ipc_field(dtype, current_id)],
54                dictionary_id,
55            }
56        },
57        // no children => do nothing
58        _ => IpcField {
59            fields: vec![],
60            dictionary_id: None,
61        },
62    }
63}
64
65/// Assigns every dictionary field a unique ID
66pub fn default_ipc_fields<'a>(fields: impl ExactSizeIterator<Item = &'a Field>) -> Vec<IpcField> {
67    let mut dictionary_id = 0i64;
68    fields
69        .map(|field| default_ipc_field(field.dtype().to_logical_type(), &mut dictionary_id))
70        .collect()
71}