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

polars_arrow/io/ipc/write/
mod.rs
