parquet_derive_test/
lib.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18#![allow(clippy::approx_constant)]
19
20use parquet_derive::{ParquetRecordReader, ParquetRecordWriter};
21
22#[derive(ParquetRecordWriter)]
23struct ACompleteRecord<'a> {
24    pub a_bool: bool,
25    pub a_str: &'a str,
26    pub a_string: String,
27    pub a_borrowed_string: &'a String,
28    pub maybe_a_str: Option<&'a str>,
29    pub maybe_a_string: Option<String>,
30    pub i16: i16,
31    pub i32: i32,
32    pub u64: u64,
33    pub maybe_u8: Option<u8>,
34    pub maybe_i16: Option<i16>,
35    pub maybe_u32: Option<u32>,
36    pub maybe_usize: Option<usize>,
37    pub isize: isize,
38    pub float: f32,
39    pub double: f64,
40    pub maybe_float: Option<f32>,
41    pub maybe_double: Option<f64>,
42    pub borrowed_maybe_a_string: &'a Option<String>,
43    pub borrowed_maybe_a_str: &'a Option<&'a str>,
44    pub now: chrono::NaiveDateTime,
45    pub uuid: uuid::Uuid,
46    pub byte_vec: Vec<u8>,
47    pub maybe_byte_vec: Option<Vec<u8>>,
48    pub borrowed_byte_vec: &'a [u8],
49    pub borrowed_maybe_byte_vec: &'a Option<Vec<u8>>,
50    pub borrowed_maybe_borrowed_byte_vec: &'a Option<&'a [u8]>,
51}
52
53#[derive(PartialEq, ParquetRecordWriter, ParquetRecordReader, Debug)]
54struct APartiallyCompleteRecord {
55    pub bool: bool,
56    pub string: String,
57    pub i16: i16,
58    pub i32: i32,
59    pub u64: u64,
60    pub isize: isize,
61    pub float: f32,
62    pub double: f64,
63    pub now: chrono::NaiveDateTime,
64    pub date: chrono::NaiveDate,
65    pub uuid: uuid::Uuid,
66    pub byte_vec: Vec<u8>,
67}
68
69// This struct has OPTIONAL columns
70// If these fields are guaranteed to be valid
71// we can load this struct into APartiallyCompleteRecord
72#[derive(PartialEq, ParquetRecordWriter, Debug)]
73struct APartiallyOptionalRecord {
74    pub bool: bool,
75    pub string: String,
76    pub i16: Option<i16>,
77    pub i32: Option<i32>,
78    pub u64: Option<u64>,
79    pub isize: isize,
80    pub float: f32,
81    pub double: f64,
82    pub now: chrono::NaiveDateTime,
83    pub date: chrono::NaiveDate,
84    pub uuid: uuid::Uuid,
85    pub byte_vec: Vec<u8>,
86}
87
88// This struct removes several fields from the "APartiallyCompleteRecord",
89// and it shuffles the fields.
90// we should still be able to load it from APartiallyCompleteRecord parquet file
91#[derive(PartialEq, ParquetRecordReader, Debug)]
92struct APrunedRecord {
93    pub bool: bool,
94    pub string: String,
95    pub byte_vec: Vec<u8>,
96    pub float: f32,
97    pub double: f64,
98    pub i16: i16,
99    pub i32: i32,
100    pub u64: u64,
101    pub isize: isize,
102}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107
108    use chrono::SubsecRound;
109    use std::{env, fs, io::Write, sync::Arc};
110
111    use parquet::{
112        file::writer::SerializedFileWriter,
113        record::{RecordReader, RecordWriter},
114        schema::parser::parse_message_type,
115    };
116
117    #[test]
118    fn test_parquet_derive_hello() {
119        let file = get_temp_file("test_parquet_derive_hello", &[]);
120
121        // The schema is not required, but this tests that the generated
122        // schema agrees with what one would write by hand.
123        let schema_str = "message rust_schema {
124            REQUIRED boolean         a_bool;
125            REQUIRED BINARY          a_str (STRING);
126            REQUIRED BINARY          a_string (STRING);
127            REQUIRED BINARY          a_borrowed_string (STRING);
128            OPTIONAL BINARY          maybe_a_str (STRING);
129            OPTIONAL BINARY          maybe_a_string (STRING);
130            REQUIRED INT32           i16 (INTEGER(16,true));
131            REQUIRED INT32           i32;
132            REQUIRED INT64           u64 (INTEGER(64,false));
133            OPTIONAL INT32           maybe_u8 (INTEGER(8,false));
134            OPTIONAL INT32           maybe_i16 (INTEGER(16,true));
135            OPTIONAL INT32           maybe_u32 (INTEGER(32,false));
136            OPTIONAL INT64           maybe_usize (INTEGER(64,false));
137            REQUIRED INT64           isize (INTEGER(64,true));
138            REQUIRED FLOAT           float;
139            REQUIRED DOUBLE          double;
140            OPTIONAL FLOAT           maybe_float;
141            OPTIONAL DOUBLE          maybe_double;
142            OPTIONAL BINARY          borrowed_maybe_a_string (STRING);
143            OPTIONAL BINARY          borrowed_maybe_a_str (STRING);
144            REQUIRED INT64           now (TIMESTAMP_MILLIS);
145            REQUIRED FIXED_LEN_BYTE_ARRAY (16) uuid (UUID);
146            REQUIRED BINARY          byte_vec;
147            OPTIONAL BINARY          maybe_byte_vec;
148            REQUIRED BINARY          borrowed_byte_vec;
149            OPTIONAL BINARY          borrowed_maybe_byte_vec;
150            OPTIONAL BINARY          borrowed_maybe_borrowed_byte_vec;
151        }";
152
153        let schema = Arc::new(parse_message_type(schema_str).unwrap());
154
155        let a_str = "hello mother".to_owned();
156        let a_borrowed_string = "cool news".to_owned();
157        let maybe_a_string = Some("it's true, I'm a string".to_owned());
158        let maybe_a_str = Some(&a_str[..]);
159        let borrowed_byte_vec = vec![0x68, 0x69, 0x70];
160        let borrowed_maybe_byte_vec = Some(vec![0x71, 0x72]);
161        let borrowed_maybe_borrowed_byte_vec = Some(&borrowed_byte_vec[..]);
162
163        let drs: Vec<ACompleteRecord> = vec![ACompleteRecord {
164            a_bool: true,
165            a_str: &a_str[..],
166            a_string: "hello father".into(),
167            a_borrowed_string: &a_borrowed_string,
168            maybe_a_str: Some(&a_str[..]),
169            maybe_a_string: Some(a_str.clone()),
170            i16: -45,
171            i32: 456,
172            u64: 4563424,
173            maybe_u8: None,
174            maybe_i16: Some(3),
175            maybe_u32: None,
176            maybe_usize: Some(4456),
177            isize: -365,
178            float: 3.5,
179            double: f64::NAN,
180            maybe_float: None,
181            maybe_double: Some(f64::MAX),
182            borrowed_maybe_a_string: &maybe_a_string,
183            borrowed_maybe_a_str: &maybe_a_str,
184            now: chrono::Utc::now().naive_local(),
185            uuid: uuid::Uuid::new_v4(),
186            byte_vec: vec![0x65, 0x66, 0x67],
187            maybe_byte_vec: Some(vec![0x88, 0x89, 0x90]),
188            borrowed_byte_vec: &borrowed_byte_vec,
189            borrowed_maybe_byte_vec: &borrowed_maybe_byte_vec,
190            borrowed_maybe_borrowed_byte_vec: &borrowed_maybe_borrowed_byte_vec,
191        }];
192
193        let generated_schema = drs.as_slice().schema().unwrap();
194
195        assert_eq!(&schema, &generated_schema);
196
197        let props = Default::default();
198        let mut writer = SerializedFileWriter::new(file, generated_schema, props).unwrap();
199
200        let mut row_group = writer.next_row_group().unwrap();
201        drs.as_slice().write_to_row_group(&mut row_group).unwrap();
202        row_group.close().unwrap();
203        writer.close().unwrap();
204    }
205
206    #[test]
207    fn test_parquet_derive_read_write_combined() {
208        let file = get_temp_file("test_parquet_derive_combined", &[]);
209
210        let mut drs: Vec<APartiallyCompleteRecord> = vec![APartiallyCompleteRecord {
211            bool: true,
212            string: "a string".into(),
213            i16: -45,
214            i32: 456,
215            u64: 4563424,
216            isize: -365,
217            float: 3.5,
218            double: f64::NAN,
219            now: chrono::Utc::now().naive_local(),
220            date: chrono::naive::NaiveDate::from_ymd_opt(2015, 3, 14).unwrap(),
221            uuid: uuid::Uuid::new_v4(),
222            byte_vec: vec![0x65, 0x66, 0x67],
223        }];
224
225        let mut out: Vec<APartiallyCompleteRecord> = Vec::new();
226
227        use parquet::file::{reader::FileReader, serialized_reader::SerializedFileReader};
228
229        let generated_schema = drs.as_slice().schema().unwrap();
230
231        let props = Default::default();
232        let mut writer =
233            SerializedFileWriter::new(file.try_clone().unwrap(), generated_schema, props).unwrap();
234
235        let mut row_group = writer.next_row_group().unwrap();
236        drs.as_slice().write_to_row_group(&mut row_group).unwrap();
237        row_group.close().unwrap();
238        writer.close().unwrap();
239
240        let reader = SerializedFileReader::new(file).unwrap();
241
242        let mut row_group = reader.get_row_group(0).unwrap();
243        out.read_from_row_group(&mut *row_group, 1).unwrap();
244
245        // correct for rounding error when writing milliseconds
246
247        drs[0].now = drs[0].now.trunc_subsecs(3);
248
249        assert!(out[0].double.is_nan()); // these three lines are necessary because NAN != NAN
250        out[0].double = 0.;
251        drs[0].double = 0.;
252
253        assert_eq!(drs[0], out[0]);
254    }
255
256    #[test]
257    fn test_parquet_derive_read_optional_but_valid_column() {
258        let file = get_temp_file("test_parquet_derive_read_optional", &[]);
259        let drs = vec![APartiallyOptionalRecord {
260            bool: true,
261            string: "a string".into(),
262            i16: Some(-45),
263            i32: Some(456),
264            u64: Some(4563424),
265            isize: -365,
266            float: 3.5,
267            double: f64::NAN,
268            now: chrono::Utc::now().naive_local(),
269            date: chrono::naive::NaiveDate::from_ymd_opt(2015, 3, 14).unwrap(),
270            uuid: uuid::Uuid::new_v4(),
271            byte_vec: vec![0x65, 0x66, 0x67],
272        }];
273
274        let generated_schema = drs.as_slice().schema().unwrap();
275
276        let props = Default::default();
277        let mut writer =
278            SerializedFileWriter::new(file.try_clone().unwrap(), generated_schema, props).unwrap();
279
280        let mut row_group = writer.next_row_group().unwrap();
281        drs.as_slice().write_to_row_group(&mut row_group).unwrap();
282        row_group.close().unwrap();
283        writer.close().unwrap();
284
285        use parquet::file::{reader::FileReader, serialized_reader::SerializedFileReader};
286        let reader = SerializedFileReader::new(file).unwrap();
287        let mut out: Vec<APartiallyCompleteRecord> = Vec::new();
288
289        let mut row_group = reader.get_row_group(0).unwrap();
290        out.read_from_row_group(&mut *row_group, 1).unwrap();
291
292        assert_eq!(drs[0].i16.unwrap(), out[0].i16);
293        assert_eq!(drs[0].i32.unwrap(), out[0].i32);
294        assert_eq!(drs[0].u64.unwrap(), out[0].u64);
295    }
296
297    #[test]
298    fn test_parquet_derive_read_pruned_and_shuffled_columns() {
299        let file = get_temp_file("test_parquet_derive_read_pruned", &[]);
300        let drs = vec![APartiallyCompleteRecord {
301            bool: true,
302            string: "a string".into(),
303            i16: -45,
304            i32: 456,
305            u64: 4563424,
306            isize: -365,
307            float: 3.5,
308            double: f64::NAN,
309            now: chrono::Utc::now().naive_local(),
310            date: chrono::naive::NaiveDate::from_ymd_opt(2015, 3, 14).unwrap(),
311            uuid: uuid::Uuid::new_v4(),
312            byte_vec: vec![0x65, 0x66, 0x67],
313        }];
314
315        let generated_schema = drs.as_slice().schema().unwrap();
316
317        let props = Default::default();
318        let mut writer =
319            SerializedFileWriter::new(file.try_clone().unwrap(), generated_schema, props).unwrap();
320
321        let mut row_group = writer.next_row_group().unwrap();
322        drs.as_slice().write_to_row_group(&mut row_group).unwrap();
323        row_group.close().unwrap();
324        writer.close().unwrap();
325
326        use parquet::file::{reader::FileReader, serialized_reader::SerializedFileReader};
327        let reader = SerializedFileReader::new(file).unwrap();
328        let mut out: Vec<APrunedRecord> = Vec::new();
329
330        let mut row_group = reader.get_row_group(0).unwrap();
331        out.read_from_row_group(&mut *row_group, 1).unwrap();
332
333        assert_eq!(drs[0].bool, out[0].bool);
334        assert_eq!(drs[0].string, out[0].string);
335        assert_eq!(drs[0].byte_vec, out[0].byte_vec);
336        assert_eq!(drs[0].float, out[0].float);
337        assert!(drs[0].double.is_nan());
338        assert!(out[0].double.is_nan());
339        assert_eq!(drs[0].i16, out[0].i16);
340        assert_eq!(drs[0].i32, out[0].i32);
341        assert_eq!(drs[0].u64, out[0].u64);
342        assert_eq!(drs[0].isize, out[0].isize);
343    }
344
345    /// Returns file handle for a temp file in 'target' directory with a provided content
346    pub fn get_temp_file(file_name: &str, content: &[u8]) -> fs::File {
347        // build tmp path to a file in "target/debug/testdata"
348        let mut path_buf = env::current_dir().unwrap();
349        path_buf.push("target");
350        path_buf.push("debug");
351        path_buf.push("testdata");
352        fs::create_dir_all(&path_buf).unwrap();
353        path_buf.push(file_name);
354
355        // write file content
356        let mut tmp_file = fs::File::create(path_buf.as_path()).unwrap();
357        tmp_file.write_all(content).unwrap();
358        tmp_file.sync_all().unwrap();
359
360        // return file handle for both read and write
361        let file = fs::OpenOptions::new()
362            .read(true)
363            .write(true)
364            .open(path_buf.as_path());
365        assert!(file.is_ok());
366        file.unwrap()
367    }
368}
parquet_derive_test/lib.rs

parquet_derive_test/
lib.rs
