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
use polars_utils::bool::UnsafeBool;
7

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

16
#[allow(clippy::too_many_arguments)]
17
pub fn read<R: Read + Seek>(
18
    field_nodes: &mut VecDeque<Node>,
19
    variadic_buffer_counts: &mut VecDeque<usize>,
20
    field: &Field,
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<BodyCompressionRef>,
28
    limit: Option<usize>,
29
    version: MetadataVersion,
30
    scratch: &mut Vec<u8>,
31
    checked: UnsafeBool,
32
) -> PolarsResult<Box<dyn Array>> {
33
    use PhysicalType::*;
34
    let dtype = field.dtype.clone();
35

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

289
pub fn skip(
290
    field_nodes: &mut VecDeque<Node>,
291
    dtype: &ArrowDataType,
292
    buffers: &mut VecDeque<IpcBuffer>,
293
    variadic_buffer_counts: &mut VecDeque<usize>,
294
) -> PolarsResult<()> {
295
    use PhysicalType::*;
296
    match dtype.to_physical_type() {
297
        Null => skip_null(field_nodes),
298
        Boolean => skip_boolean(field_nodes, buffers),
299
        Primitive(_) => skip_primitive(field_nodes, buffers),
300
        LargeBinary | Binary => skip_binary(field_nodes, buffers),
301
        LargeUtf8 | Utf8 => skip_utf8(field_nodes, buffers),
302
        FixedSizeBinary => skip_fixed_size_binary(field_nodes, buffers),
303
        List => skip_list::<i32>(field_nodes, dtype, buffers, variadic_buffer_counts),
304
        LargeList => skip_list::<i64>(field_nodes, dtype, buffers, variadic_buffer_counts),
305
        FixedSizeList => skip_fixed_size_list(field_nodes, dtype, buffers, variadic_buffer_counts),
306
        Struct => skip_struct(field_nodes, dtype, buffers, variadic_buffer_counts),
307
        Dictionary(_) => skip_dictionary(field_nodes, buffers),
308
        Union => skip_union(field_nodes, dtype, buffers, variadic_buffer_counts),
309
        Map => skip_map(field_nodes, dtype, buffers, variadic_buffer_counts),
310
        BinaryView | Utf8View => skip_binview(field_nodes, buffers, variadic_buffer_counts),
311
    }
312
}
313

314
Product

Resources

Company
