Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
pola-rs
GitHub Repository: pola-rs/polars
Path: blob/main/crates/polars-arrow/src/io/ipc/read/array/struct_.rs
8420 views
1
use std::collections::VecDeque;
2
use std::io::{Read, Seek};
3
4
use polars_error::{PolarsResult, polars_err};
5
use polars_utils::bool::UnsafeBool;
6
7
use super::super::super::IpcField;
8
use super::super::deserialize::{read, skip};
9
use super::super::read_basic::*;
10
use super::super::{Compression, Dictionaries, IpcBuffer, Node, Version};
11
use super::try_get_array_length;
12
use crate::array::StructArray;
13
use crate::datatypes::ArrowDataType;
14
use crate::io::ipc::read::array::try_get_field_node;
15
16
#[allow(clippy::too_many_arguments)]
17
pub fn read_struct<R: Read + Seek>(
18
field_nodes: &mut VecDeque<Node>,
19
variadic_buffer_counts: &mut VecDeque<usize>,
20
dtype: ArrowDataType,
21
ipc_field: &IpcField,
22
buffers: &mut VecDeque<IpcBuffer>,
23
reader: &mut R,
24
dictionaries: &Dictionaries,
25
block_offset: u64,
26
is_little_endian: bool,
27
compression: Option<Compression>,
28
limit: Option<usize>,
29
version: Version,
30
scratch: &mut Vec<u8>,
31
checked: UnsafeBool,
32
) -> PolarsResult<StructArray> {
33
let field_node = try_get_field_node(field_nodes, &dtype)?;
34
let length = try_get_array_length(field_node, limit)?;
35
36
let validity = read_validity(
37
buffers,
38
field_node,
39
reader,
40
block_offset,
41
is_little_endian,
42
compression,
43
limit,
44
scratch,
45
)?;
46
47
let fields = StructArray::get_fields(&dtype);
48
49
let values = fields
50
.iter()
51
.zip(ipc_field.fields.iter())
52
.map(|(field, ipc_field)| {
53
read(
54
field_nodes,
55
variadic_buffer_counts,
56
field,
57
ipc_field,
58
buffers,
59
reader,
60
dictionaries,
61
block_offset,
62
is_little_endian,
63
compression,
64
limit,
65
version,
66
scratch,
67
checked,
68
)
69
})
70
.collect::<PolarsResult<Vec<_>>>()?;
71
72
StructArray::try_new(dtype, length, values, validity)
73
}
74
75
pub fn skip_struct(
76
field_nodes: &mut VecDeque<Node>,
77
dtype: &ArrowDataType,
78
buffers: &mut VecDeque<IpcBuffer>,
79
variadic_buffer_counts: &mut VecDeque<usize>,
80
) -> PolarsResult<()> {
81
let _ = field_nodes.pop_front().ok_or_else(|| {
82
polars_err!(
83
oos = "IPC: unable to fetch the field for struct. The file or stream is corrupted."
84
)
85
})?;
86
87
let _ = buffers
88
.pop_front()
89
.ok_or_else(|| polars_err!(oos = "IPC: missing validity buffer."))?;
90
91
let fields = StructArray::get_fields(dtype);
92
93
fields
94
.iter()
95
.try_for_each(|field| skip(field_nodes, field.dtype(), buffers, variadic_buffer_counts))
96
}
97
98