parquet 58.1.0

Apache Parquet implementation in Rust
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

use crate::arrow::array_reader::{ArrayReader, ListArrayReader, StructArrayReader};
use crate::errors::Result;
use arrow_array::{Array, ArrayRef, MapArray};
use arrow_schema::DataType as ArrowType;
use std::any::Any;
use std::sync::Arc;

/// Implementation of a map array reader.
pub struct MapArrayReader {
    data_type: ArrowType,
    reader: ListArrayReader<i32>,
}

impl MapArrayReader {
    #[allow(rustdoc::private_intra_doc_links)]
    /// Creates a new [`MapArrayReader`] with a `def_level`, `rep_level` and `nullable`
    /// as defined on [`ParquetField`][crate::arrow::schema::ParquetField]
    pub fn new(
        key_reader: Box<dyn ArrayReader>,
        value_reader: Box<dyn ArrayReader>,
        data_type: ArrowType,
        def_level: i16,
        rep_level: i16,
        nullable: bool,
    ) -> Self {
        let struct_def_level = match nullable {
            true => def_level + 2,
            false => def_level + 1,
        };
        let struct_rep_level = rep_level + 1;

        let element = match &data_type {
            ArrowType::Map(element, _) => match element.data_type() {
                ArrowType::Struct(fields) if fields.len() == 2 => {
                    // Parquet cannot represent nullability at this level (#1697)
                    // and so encountering nullability here indicates some manner
                    // of schema inconsistency / inference bug
                    assert!(!element.is_nullable(), "map struct cannot be nullable");
                    element
                }
                _ => unreachable!("expected struct with two fields"),
            },
            _ => unreachable!("expected map type"),
        };

        let struct_reader = StructArrayReader::new(
            element.data_type().clone(),
            vec![key_reader, value_reader],
            struct_def_level,
            struct_rep_level,
            false,
        );

        let reader = ListArrayReader::new(
            Box::new(struct_reader),
            ArrowType::List(element.clone()),
            def_level,
            rep_level,
            nullable,
        );

        Self { data_type, reader }
    }
}

impl ArrayReader for MapArrayReader {
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn get_data_type(&self) -> &ArrowType {
        &self.data_type
    }

    fn read_records(&mut self, batch_size: usize) -> Result<usize> {
        self.reader.read_records(batch_size)
    }

    fn consume_batch(&mut self) -> Result<ArrayRef> {
        // A MapArray is just a ListArray with a StructArray child
        // we can therefore just alter the ArrayData
        let array = self.reader.consume_batch().unwrap();
        let data = array.to_data();
        let builder = data.into_builder().data_type(self.data_type.clone());

        // SAFETY - we can assume that ListArrayReader produces valid ListArray
        // of the expected type, and as such its output can be reinterpreted as
        // a MapArray without validation
        Ok(Arc::new(MapArray::from(unsafe {
            builder.build_unchecked()
        })))
    }

    fn skip_records(&mut self, num_records: usize) -> Result<usize> {
        self.reader.skip_records(num_records)
    }

    fn get_def_levels(&self) -> Option<&[i16]> {
        self.reader.get_def_levels()
    }

    fn get_rep_levels(&self) -> Option<&[i16]> {
        self.reader.get_rep_levels()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::arrow::ArrowWriter;
    use crate::arrow::arrow_reader::ParquetRecordBatchReader;
    use arrow::datatypes::{Field, Int32Type, Schema};
    use arrow_array::RecordBatch;
    use arrow_array::builder::{MapBuilder, PrimitiveBuilder, StringBuilder};
    use arrow_array::cast::*;
    use arrow_schema::Fields;
    use bytes::Bytes;

    #[test]
    // This test writes a parquet file with the following data:
    // +--------------------------------------------------------+
    // |map                                                     |
    // +--------------------------------------------------------+
    // |null                                                    |
    // |null                                                    |
    // |{three -> 3, four -> 4, five -> 5, six -> 6, seven -> 7}|
    // +--------------------------------------------------------+
    //
    // It then attempts to read the data back and checks that the third record
    // contains the expected values.
    fn read_map_array_column() {
        // Schema for single map of string to int32
        let schema = Schema::new(vec![Field::new(
            "map",
            ArrowType::Map(
                Arc::new(Field::new(
                    "entries",
                    ArrowType::Struct(Fields::from(vec![
                        Field::new("keys", ArrowType::Utf8, false),
                        Field::new("values", ArrowType::Int32, true),
                    ])),
                    false,
                )),
                false, // Map field not sorted
            ),
            true,
        )]);

        // Create builders for map
        let string_builder = StringBuilder::new();
        let ints_builder: PrimitiveBuilder<Int32Type> = PrimitiveBuilder::new();
        let mut map_builder = MapBuilder::new(None, string_builder, ints_builder);

        // Add two null records and one record with five entries
        map_builder.append(false).expect("adding null map entry");
        map_builder.append(false).expect("adding null map entry");
        map_builder.keys().append_value("three");
        map_builder.keys().append_value("four");
        map_builder.keys().append_value("five");
        map_builder.keys().append_value("six");
        map_builder.keys().append_value("seven");

        map_builder.values().append_value(3);
        map_builder.values().append_value(4);
        map_builder.values().append_value(5);
        map_builder.values().append_value(6);
        map_builder.values().append_value(7);
        map_builder.append(true).expect("adding map entry");

        // Create record batch
        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(map_builder.finish())])
            .expect("create record batch");

        // Write record batch to file
        let mut buffer = Vec::with_capacity(1024);
        let mut writer =
            ArrowWriter::try_new(&mut buffer, batch.schema(), None).expect("creat file writer");
        writer.write(&batch).expect("writing file");
        writer.close().expect("close writer");

        // Read file
        let reader = Bytes::from(buffer);
        let record_batch_reader = ParquetRecordBatchReader::try_new(reader, 1024).unwrap();
        for maybe_record_batch in record_batch_reader {
            let record_batch = maybe_record_batch.expect("Getting current batch");
            let col = record_batch.column(0);
            assert!(col.is_null(0));
            assert!(col.is_null(1));
            let map_entry = as_map_array(col).value(2);
            let struct_col = as_struct_array(&map_entry);
            let key_col = as_string_array(struct_col.column(0)); // Key column
            assert_eq!(key_col.value(0), "three");
            assert_eq!(key_col.value(1), "four");
            assert_eq!(key_col.value(2), "five");
            assert_eq!(key_col.value(3), "six");
            assert_eq!(key_col.value(4), "seven");
        }
    }
}