CoCalc -- deserialize.rs

GitHub Repository: pola-rs/polars
Path: blob/main/crates/polars-arrow/src/io/ipc/read/deserialize.rs
⁶⁹⁴⁰ views
1
use std::collections::VecDeque;
2
use std::io::{Read, Seek};
3

4
use arrow_format::ipc::{BodyCompressionRef, MetadataVersion};
5
use polars_error::PolarsResult;
6

7
use super::array::*;
8
use super::{Dictionaries, IpcBuffer, Node};
9
use crate::array::*;
10
use crate::datatypes::{ArrowDataType, Field, PhysicalType};
11
use crate::io::ipc::IpcField;
12
use crate::{match_integer_type, with_match_primitive_type_full};
13

14
#[allow(clippy::too_many_arguments)]
15
pub fn read<R: Read + Seek>(
16
    field_nodes: &mut VecDeque<Node>,
17
    variadic_buffer_counts: &mut VecDeque<usize>,
18
    field: &Field,
19
    ipc_field: &IpcField,
20
    buffers: &mut VecDeque<IpcBuffer>,
21
    reader: &mut R,
22
    dictionaries: &Dictionaries,
23
    block_offset: u64,
24
    is_little_endian: bool,
25
    compression: Option<BodyCompressionRef>,
26
    limit: Option<usize>,
27
    version: MetadataVersion,
28
    scratch: &mut Vec<u8>,
29
) -> PolarsResult<Box<dyn Array>> {
30
    use PhysicalType::*;
31
    let dtype = field.dtype.clone();
32

33
    match dtype.to_physical_type() {
34
        Null => read_null(field_nodes, dtype, limit).map(|x| x.boxed()),
35
        Boolean => read_boolean(
36
            field_nodes,
37
            dtype,
38
            buffers,
39
            reader,
40
            block_offset,
41
            is_little_endian,
42
            compression,
43
            limit,
44
            scratch,
45
        )
46
        .map(|x| x.boxed()),
47
        Primitive(primitive) => with_match_primitive_type_full!(primitive, |$T| {
48
            read_primitive::<$T, _>(
49
                field_nodes,
50
                dtype,
51
                buffers,
52
                reader,
53
                block_offset,
54
                is_little_endian,
55
                compression,
56
                limit,
57
                scratch,
58
            )
59
            .map(|x| x.boxed())
60
        }),
61
        Binary => read_binary::<i32, _>(
62
            field_nodes,
63
            dtype,
64
            buffers,
65
            reader,
66
            block_offset,
67
            is_little_endian,
68
            compression,
69
            limit,
70
            scratch,
71
        )
72
        .map(|x| x.boxed()),
73
        LargeBinary => read_binary::<i64, _>(
74
            field_nodes,
75
            dtype,
76
            buffers,
77
            reader,
78
            block_offset,
79
            is_little_endian,
80
            compression,
81
            limit,
82
            scratch,
83
        )
84
        .map(|x| x.boxed()),
85
        FixedSizeBinary => read_fixed_size_binary(
86
            field_nodes,
87
            dtype,
88
            buffers,
89
            reader,
90
            block_offset,
91
            is_little_endian,
92
            compression,
93
            limit,
94
            scratch,
95
        )
96
        .map(|x| x.boxed()),
97
        Utf8 => read_utf8::<i32, _>(
98
            field_nodes,
99
            dtype,
100
            buffers,
101
            reader,
102
            block_offset,
103
            is_little_endian,
104
            compression,
105
            limit,
106
            scratch,
107
        )
108
        .map(|x| x.boxed()),
109
        LargeUtf8 => read_utf8::<i64, _>(
110
            field_nodes,
111
            dtype,
112
            buffers,
113
            reader,
114
            block_offset,
115
            is_little_endian,
116
            compression,
117
            limit,
118
            scratch,
119
        )
120
        .map(|x| x.boxed()),
121
        List => read_list::<i32, _>(
122
            field_nodes,
123
            variadic_buffer_counts,
124
            dtype,
125
            ipc_field,
126
            buffers,
127
            reader,
128
            dictionaries,
129
            block_offset,
130
            is_little_endian,
131
            compression,
132
            limit,
133
            version,
134
            scratch,
135
        )
136
        .map(|x| x.boxed()),
137
        LargeList => read_list::<i64, _>(
138
            field_nodes,
139
            variadic_buffer_counts,
140
            dtype,
141
            ipc_field,
142
            buffers,
143
            reader,
144
            dictionaries,
145
            block_offset,
146
            is_little_endian,
147
            compression,
148
            limit,
149
            version,
150
            scratch,
151
        )
152
        .map(|x| x.boxed()),
153
        FixedSizeList => read_fixed_size_list(
154
            field_nodes,
155
            variadic_buffer_counts,
156
            dtype,
157
            ipc_field,
158
            buffers,
159
            reader,
160
            dictionaries,
161
            block_offset,
162
            is_little_endian,
163
            compression,
164
            limit,
165
            version,
166
            scratch,
167
        )
168
        .map(|x| x.boxed()),
169
        Struct => read_struct(
170
            field_nodes,
171
            variadic_buffer_counts,
172
            dtype,
173
            ipc_field,
174
            buffers,
175
            reader,
176
            dictionaries,
177
            block_offset,
178
            is_little_endian,
179
            compression,
180
            limit,
181
            version,
182
            scratch,
183
        )
184
        .map(|x| x.boxed()),
185
        Dictionary(key_type) => {
186
            match_integer_type!(key_type, |$T| {
187
                read_dictionary::<$T, _>(
188
                    field_nodes,
189
                    dtype,
190
                    ipc_field.dictionary_id,
191
                    buffers,
192
                    reader,
193
                    dictionaries,
194
                    block_offset,
195
                    compression,
196
                    limit,
197
                    is_little_endian,
198
                    scratch,
199
                )
200
                .map(|x| x.boxed())
201
            })
202
        },
203
        Union => read_union(
204
            field_nodes,
205
            variadic_buffer_counts,
206
            dtype,
207
            ipc_field,
208
            buffers,
209
            reader,
210
            dictionaries,
211
            block_offset,
212
            is_little_endian,
213
            compression,
214
            limit,
215
            version,
216
            scratch,
217
        )
218
        .map(|x| x.boxed()),
219
        Map => read_map(
220
            field_nodes,
221
            variadic_buffer_counts,
222
            dtype,
223
            ipc_field,
224
            buffers,
225
            reader,
226
            dictionaries,
227
            block_offset,
228
            is_little_endian,
229
            compression,
230
            limit,
231
            version,
232
            scratch,
233
        )
234
        .map(|x| x.boxed()),
235
        Utf8View => read_binview::<str, _>(
236
            field_nodes,
237
            variadic_buffer_counts,
238
            dtype,
239
            buffers,
240
            reader,
241
            block_offset,
242
            is_little_endian,
243
            compression,
244
            limit,
245
            scratch,
246
        ),
247
        BinaryView => read_binview::<[u8], _>(
248
            field_nodes,
249
            variadic_buffer_counts,
250
            dtype,
251
            buffers,
252
            reader,
253
            block_offset,
254
            is_little_endian,
255
            compression,
256
            limit,
257
            scratch,
258
        ),
259
    }
260
}
261

262
pub fn skip(
263
    field_nodes: &mut VecDeque<Node>,
264
    dtype: &ArrowDataType,
265
    buffers: &mut VecDeque<IpcBuffer>,
266
    variadic_buffer_counts: &mut VecDeque<usize>,
267
) -> PolarsResult<()> {
268
    use PhysicalType::*;
269
    match dtype.to_physical_type() {
270
        Null => skip_null(field_nodes),
271
        Boolean => skip_boolean(field_nodes, buffers),
272
        Primitive(_) => skip_primitive(field_nodes, buffers),
273
        LargeBinary | Binary => skip_binary(field_nodes, buffers),
274
        LargeUtf8 | Utf8 => skip_utf8(field_nodes, buffers),
275
        FixedSizeBinary => skip_fixed_size_binary(field_nodes, buffers),
276
        List => skip_list::<i32>(field_nodes, dtype, buffers, variadic_buffer_counts),
277
        LargeList => skip_list::<i64>(field_nodes, dtype, buffers, variadic_buffer_counts),
278
        FixedSizeList => skip_fixed_size_list(field_nodes, dtype, buffers, variadic_buffer_counts),
279
        Struct => skip_struct(field_nodes, dtype, buffers, variadic_buffer_counts),
280
        Dictionary(_) => skip_dictionary(field_nodes, buffers),
281
        Union => skip_union(field_nodes, dtype, buffers, variadic_buffer_counts),
282
        Map => skip_map(field_nodes, dtype, buffers, variadic_buffer_counts),
283
        BinaryView | Utf8View => skip_binview(field_nodes, buffers, variadic_buffer_counts),
284
    }
285
}
286

287
Product

Resources

Company
