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#![doc(
19    html_logo_url = "https://raw.githubusercontent.com/apache/parquet-format/25f05e73d8cd7f5c83532ce51cb4f4de8ba5f2a2/logo/parquet-logos_1.svg",
20    html_favicon_url = "https://raw.githubusercontent.com/apache/parquet-format/25f05e73d8cd7f5c83532ce51cb4f4de8ba5f2a2/logo/parquet-logos_1.svg"
21)]
22#![cfg_attr(docsrs, feature(doc_cfg))]
23#![allow(clippy::approx_constant)]
24
25use parquet_derive::{ParquetRecordReader, ParquetRecordWriter};
26use std::sync::Arc;
27
28#[derive(ParquetRecordWriter)]
29struct ACompleteRecord<'a> {
30    pub a_bool: bool,
31    pub a_str: &'a str,
32    pub a_string: String,
33    pub a_borrowed_string: &'a String,
34    pub a_arc_str: Arc<str>,
35    pub maybe_a_str: Option<&'a str>,
36    pub maybe_a_string: Option<String>,
37    pub maybe_a_arc_str: Option<Arc<str>>,
38    pub i16: i16,
39    pub i32: i32,
40    pub u64: u64,
41    pub maybe_u8: Option<u8>,
42    pub maybe_i16: Option<i16>,
43    pub maybe_u32: Option<u32>,
44    pub maybe_usize: Option<usize>,
45    pub isize: isize,
46    pub float: f32,
47    pub double: f64,
48    pub maybe_float: Option<f32>,
49    pub maybe_double: Option<f64>,
50    pub borrowed_maybe_a_string: &'a Option<String>,
51    pub borrowed_maybe_a_str: &'a Option<&'a str>,
52    pub now: chrono::NaiveDateTime,
53    pub uuid: uuid::Uuid,
54    pub byte_vec: Vec<u8>,
55    pub maybe_byte_vec: Option<Vec<u8>>,
56    pub borrowed_byte_vec: &'a [u8],
57    pub borrowed_maybe_byte_vec: &'a Option<Vec<u8>>,
58    pub borrowed_maybe_borrowed_byte_vec: &'a Option<&'a [u8]>,
59}
60
61#[derive(PartialEq, ParquetRecordWriter, ParquetRecordReader, Debug)]
62struct APartiallyCompleteRecord {
63    pub bool: bool,
64    pub string: String,
65    pub i16: i16,
66    pub i32: i32,
67    pub u64: u64,
68    pub isize: isize,
69    pub float: f32,
70    pub double: f64,
71    pub now: chrono::NaiveDateTime,
72    pub date: chrono::NaiveDate,
73    pub uuid: uuid::Uuid,
74    pub byte_vec: Vec<u8>,
75}
76
77// This struct has OPTIONAL columns
78// If these fields are guaranteed to be valid
79// we can load this struct into APartiallyCompleteRecord
80#[derive(PartialEq, ParquetRecordWriter, Debug)]
81struct APartiallyOptionalRecord {
82    pub bool: bool,
83    pub string: String,
84    pub i16: Option<i16>,
85    pub i32: Option<i32>,
86    pub u64: Option<u64>,
87    pub isize: isize,
88    pub float: f32,
89    pub double: f64,
90    pub now: chrono::NaiveDateTime,
91    pub date: chrono::NaiveDate,
92    pub uuid: uuid::Uuid,
93    pub byte_vec: Vec<u8>,
94}
95
96// This struct removes several fields from the "APartiallyCompleteRecord",
97// and it shuffles the fields.
98// we should still be able to load it from APartiallyCompleteRecord parquet file
99#[derive(PartialEq, ParquetRecordReader, Debug)]
100struct APrunedRecord {
101    pub bool: bool,
102    pub string: String,
103    pub byte_vec: Vec<u8>,
104    pub float: f32,
105    pub double: f64,
106    pub i16: i16,
107    pub i32: i32,
108    pub u64: u64,
109    pub isize: isize,
110}
111
112// This struct has a field declared with a raw identifier,
113// which maps to a parquet column named without the `r#` prefix
114// (e.g. a column named `type`, as written by other tools)
115#[derive(PartialEq, ParquetRecordWriter, ParquetRecordReader, Debug)]
116struct ARecordWithRawIdentifiers {
117    pub r#type: i32,
118    pub count: i32,
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124
125    use chrono::SubsecRound;
126    use std::{env, fs, io::Write, sync::Arc};
127
128    use parquet::{
129        file::writer::SerializedFileWriter,
130        record::{RecordReader, RecordWriter},
131        schema::parser::parse_message_type,
132    };
133
134    #[test]
135    fn test_parquet_derive_hello() {
136        let file = get_temp_file("test_parquet_derive_hello", &[]);
137
138        // The schema is not required, but this tests that the generated
139        // schema agrees with what one would write by hand.
140        let schema_str = "message rust_schema {
141            REQUIRED boolean         a_bool;
142            REQUIRED BINARY          a_str (STRING);
143            REQUIRED BINARY          a_string (STRING);
144            REQUIRED BINARY          a_borrowed_string (STRING);
145            REQUIRED BINARY          a_arc_str (STRING);
146            OPTIONAL BINARY          maybe_a_str (STRING);
147            OPTIONAL BINARY          maybe_a_string (STRING);
148            OPTIONAL BINARY          maybe_a_arc_str (STRING);
149            REQUIRED INT32           i16 (INTEGER(16,true));
150            REQUIRED INT32           i32;
151            REQUIRED INT64           u64 (INTEGER(64,false));
152            OPTIONAL INT32           maybe_u8 (INTEGER(8,false));
153            OPTIONAL INT32           maybe_i16 (INTEGER(16,true));
154            OPTIONAL INT32           maybe_u32 (INTEGER(32,false));
155            OPTIONAL INT64           maybe_usize (INTEGER(64,false));
156            REQUIRED INT64           isize (INTEGER(64,true));
157            REQUIRED FLOAT           float;
158            REQUIRED DOUBLE          double;
159            OPTIONAL FLOAT           maybe_float;
160            OPTIONAL DOUBLE          maybe_double;
161            OPTIONAL BINARY          borrowed_maybe_a_string (STRING);
162            OPTIONAL BINARY          borrowed_maybe_a_str (STRING);
163            REQUIRED INT64           now (TIMESTAMP_MILLIS);
164            REQUIRED FIXED_LEN_BYTE_ARRAY (16) uuid (UUID);
165            REQUIRED BINARY          byte_vec;
166            OPTIONAL BINARY          maybe_byte_vec;
167            REQUIRED BINARY          borrowed_byte_vec;
168            OPTIONAL BINARY          borrowed_maybe_byte_vec;
169            OPTIONAL BINARY          borrowed_maybe_borrowed_byte_vec;
170        }";
171
172        let schema = Arc::new(parse_message_type(schema_str).unwrap());
173
174        let a_str = "hello mother".to_owned();
175        let a_borrowed_string = "cool news".to_owned();
176        let a_arc_str: Arc<str> = "hello arc".into();
177        let maybe_a_string = Some("it's true, I'm a string".to_owned());
178        let maybe_a_str = Some(&a_str[..]);
179        let maybe_a_arc_str = Some(a_arc_str.clone());
180        let borrowed_byte_vec = vec![0x68, 0x69, 0x70];
181        let borrowed_maybe_byte_vec = Some(vec![0x71, 0x72]);
182        let borrowed_maybe_borrowed_byte_vec = Some(&borrowed_byte_vec[..]);
183
184        let drs: Vec<ACompleteRecord> = vec![ACompleteRecord {
185            a_bool: true,
186            a_str: &a_str[..],
187            a_string: "hello father".into(),
188            a_borrowed_string: &a_borrowed_string,
189            a_arc_str,
190            maybe_a_str: Some(&a_str[..]),
191            maybe_a_string: Some(a_str.clone()),
192            maybe_a_arc_str,
193            i16: -45,
194            i32: 456,
195            u64: 4563424,
196            maybe_u8: None,
197            maybe_i16: Some(3),
198            maybe_u32: None,
199            maybe_usize: Some(4456),
200            isize: -365,
201            float: 3.5,
202            double: f64::NAN,
203            maybe_float: None,
204            maybe_double: Some(f64::MAX),
205            borrowed_maybe_a_string: &maybe_a_string,
206            borrowed_maybe_a_str: &maybe_a_str,
207            now: chrono::Utc::now().naive_local(),
208            uuid: uuid::Uuid::new_v4(),
209            byte_vec: vec![0x65, 0x66, 0x67],
210            maybe_byte_vec: Some(vec![0x88, 0x89, 0x90]),
211            borrowed_byte_vec: &borrowed_byte_vec,
212            borrowed_maybe_byte_vec: &borrowed_maybe_byte_vec,
213            borrowed_maybe_borrowed_byte_vec: &borrowed_maybe_borrowed_byte_vec,
214        }];
215
216        let generated_schema = drs.as_slice().schema().unwrap();
217
218        assert_eq!(&schema, &generated_schema);
219
220        let props = Default::default();
221        let mut writer = SerializedFileWriter::new(file, generated_schema, props).unwrap();
222
223        let mut row_group = writer.next_row_group().unwrap();
224        drs.as_slice().write_to_row_group(&mut row_group).unwrap();
225        row_group.close().unwrap();
226        writer.close().unwrap();
227    }
228
229    #[test]
230    fn test_parquet_derive_read_write_combined() {
231        let file = get_temp_file("test_parquet_derive_combined", &[]);
232
233        let mut drs: Vec<APartiallyCompleteRecord> = vec![APartiallyCompleteRecord {
234            bool: true,
235            string: "a string".into(),
236            i16: -45,
237            i32: 456,
238            u64: 4563424,
239            isize: -365,
240            float: 3.5,
241            double: f64::NAN,
242            now: chrono::Utc::now().naive_local(),
243            date: chrono::naive::NaiveDate::from_ymd_opt(2015, 3, 14).unwrap(),
244            uuid: uuid::Uuid::new_v4(),
245            byte_vec: vec![0x65, 0x66, 0x67],
246        }];
247
248        let mut out: Vec<APartiallyCompleteRecord> = Vec::new();
249
250        use parquet::file::{reader::FileReader, serialized_reader::SerializedFileReader};
251
252        let generated_schema = drs.as_slice().schema().unwrap();
253
254        let props = Default::default();
255        let mut writer =
256            SerializedFileWriter::new(file.try_clone().unwrap(), generated_schema, props).unwrap();
257
258        let mut row_group = writer.next_row_group().unwrap();
259        drs.as_slice().write_to_row_group(&mut row_group).unwrap();
260        row_group.close().unwrap();
261        writer.close().unwrap();
262
263        let reader = SerializedFileReader::new(file).unwrap();
264
265        let mut row_group = reader.get_row_group(0).unwrap();
266        out.read_from_row_group(&mut *row_group, 1).unwrap();
267
268        // correct for rounding error when writing milliseconds
269
270        drs[0].now = drs[0].now.trunc_subsecs(3);
271
272        assert!(out[0].double.is_nan()); // these three lines are necessary because NAN != NAN
273        out[0].double = 0.;
274        drs[0].double = 0.;
275
276        assert_eq!(drs[0], out[0]);
277    }
278
279    #[test]
280    fn test_parquet_derive_read_optional_but_valid_column() {
281        let file = get_temp_file("test_parquet_derive_read_optional", &[]);
282        let drs = vec![APartiallyOptionalRecord {
283            bool: true,
284            string: "a string".into(),
285            i16: Some(-45),
286            i32: Some(456),
287            u64: Some(4563424),
288            isize: -365,
289            float: 3.5,
290            double: f64::NAN,
291            now: chrono::Utc::now().naive_local(),
292            date: chrono::naive::NaiveDate::from_ymd_opt(2015, 3, 14).unwrap(),
293            uuid: uuid::Uuid::new_v4(),
294            byte_vec: vec![0x65, 0x66, 0x67],
295        }];
296
297        let generated_schema = drs.as_slice().schema().unwrap();
298
299        let props = Default::default();
300        let mut writer =
301            SerializedFileWriter::new(file.try_clone().unwrap(), generated_schema, props).unwrap();
302
303        let mut row_group = writer.next_row_group().unwrap();
304        drs.as_slice().write_to_row_group(&mut row_group).unwrap();
305        row_group.close().unwrap();
306        writer.close().unwrap();
307
308        use parquet::file::{reader::FileReader, serialized_reader::SerializedFileReader};
309        let reader = SerializedFileReader::new(file).unwrap();
310        let mut out: Vec<APartiallyCompleteRecord> = Vec::new();
311
312        let mut row_group = reader.get_row_group(0).unwrap();
313        out.read_from_row_group(&mut *row_group, 1).unwrap();
314
315        assert_eq!(drs[0].i16.unwrap(), out[0].i16);
316        assert_eq!(drs[0].i32.unwrap(), out[0].i32);
317        assert_eq!(drs[0].u64.unwrap(), out[0].u64);
318    }
319
320    #[test]
321    fn test_parquet_derive_read_pruned_and_shuffled_columns() {
322        let file = get_temp_file("test_parquet_derive_read_pruned", &[]);
323        let drs = vec![APartiallyCompleteRecord {
324            bool: true,
325            string: "a string".into(),
326            i16: -45,
327            i32: 456,
328            u64: 4563424,
329            isize: -365,
330            float: 3.5,
331            double: f64::NAN,
332            now: chrono::Utc::now().naive_local(),
333            date: chrono::naive::NaiveDate::from_ymd_opt(2015, 3, 14).unwrap(),
334            uuid: uuid::Uuid::new_v4(),
335            byte_vec: vec![0x65, 0x66, 0x67],
336        }];
337
338        let generated_schema = drs.as_slice().schema().unwrap();
339
340        let props = Default::default();
341        let mut writer =
342            SerializedFileWriter::new(file.try_clone().unwrap(), generated_schema, props).unwrap();
343
344        let mut row_group = writer.next_row_group().unwrap();
345        drs.as_slice().write_to_row_group(&mut row_group).unwrap();
346        row_group.close().unwrap();
347        writer.close().unwrap();
348
349        use parquet::file::{reader::FileReader, serialized_reader::SerializedFileReader};
350        let reader = SerializedFileReader::new(file).unwrap();
351        let mut out: Vec<APrunedRecord> = Vec::new();
352
353        let mut row_group = reader.get_row_group(0).unwrap();
354        out.read_from_row_group(&mut *row_group, 1).unwrap();
355
356        assert_eq!(drs[0].bool, out[0].bool);
357        assert_eq!(drs[0].string, out[0].string);
358        assert_eq!(drs[0].byte_vec, out[0].byte_vec);
359        assert_eq!(drs[0].float, out[0].float);
360        assert!(drs[0].double.is_nan());
361        assert!(out[0].double.is_nan());
362        assert_eq!(drs[0].i16, out[0].i16);
363        assert_eq!(drs[0].i32, out[0].i32);
364        assert_eq!(drs[0].u64, out[0].u64);
365        assert_eq!(drs[0].isize, out[0].isize);
366    }
367
368    #[test]
369    fn test_parquet_derive_raw_identifiers() {
370        let file = get_temp_file("test_parquet_derive_raw_identifiers", &[]);
371        let drs = vec![ARecordWithRawIdentifiers {
372            r#type: 456,
373            count: 123,
374        }];
375
376        let generated_schema = drs.as_slice().schema().unwrap();
377
378        // raw identifiers are written without the `r#` prefix,
379        // while normal identifiers are unchanged
380        assert_eq!(
381            vec!["type", "count"],
382            generated_schema
383                .get_fields()
384                .iter()
385                .map(|field| field.name())
386                .collect::<Vec<_>>()
387        );
388
389        let props = Default::default();
390        let mut writer =
391            SerializedFileWriter::new(file.try_clone().unwrap(), generated_schema, props).unwrap();
392
393        let mut row_group = writer.next_row_group().unwrap();
394        drs.as_slice().write_to_row_group(&mut row_group).unwrap();
395        row_group.close().unwrap();
396        writer.close().unwrap();
397
398        use parquet::file::{reader::FileReader, serialized_reader::SerializedFileReader};
399        let reader = SerializedFileReader::new(file).unwrap();
400        let mut out: Vec<ARecordWithRawIdentifiers> = Vec::new();
401
402        let mut row_group = reader.get_row_group(0).unwrap();
403        out.read_from_row_group(&mut *row_group, 1).unwrap();
404
405        assert_eq!(drs, out);
406    }
407
408    #[test]
409    fn test_aliased_result() {
410        // Issue 7547, Where aliasing the `Result` led to
411        // a collision with the macro internals of derive ParquetRecordReader
412
413        mod aliased_result {
414            use parquet_derive::{ParquetRecordReader, ParquetRecordWriter};
415
416            // This is the normal pattern that raised this issue
417            // pub type Result = std::result::Result<(), Box<dyn std::error::Error>>;
418
419            // not an actual result type
420            // Used here only to make the harder.
421            pub type Result = ();
422
423            #[derive(ParquetRecordReader, ParquetRecordWriter, Debug)]
424            pub struct ARecord {
425                pub bool: bool,
426                pub string: String,
427            }
428
429            impl ARecord {
430                pub fn do_nothing(&self) -> Result {}
431                pub fn validate(&self) -> std::result::Result<(), Box<dyn std::error::Error>> {
432                    Ok(())
433                }
434            }
435        }
436
437        use aliased_result::ARecord;
438        let foo = ARecord {
439            bool: true,
440            string: "test".to_string(),
441        };
442        foo.do_nothing();
443        assert!(foo.validate().is_ok());
444    }
445
446    /// Returns file handle for a temp file in 'target' directory with a provided content
447    pub fn get_temp_file(file_name: &str, content: &[u8]) -> fs::File {
448        // build tmp path to a file in "target/debug/testdata"
449        let mut path_buf = env::current_dir().unwrap();
450        path_buf.push("target");
451        path_buf.push("debug");
452        path_buf.push("testdata");
453        fs::create_dir_all(&path_buf).unwrap();
454        path_buf.push(file_name);
455
456        // write file content
457        let mut tmp_file = fs::File::create(path_buf.as_path()).unwrap();
458        tmp_file.write_all(content).unwrap();
459        tmp_file.sync_all().unwrap();
460
461        // return file handle for both read and write
462        let file = fs::OpenOptions::new()
463            .read(true)
464            .write(true)
465            .open(path_buf.as_path());
466        assert!(file.is_ok());
467        file.unwrap()
468    }
469}
parquet_derive_test/lib.rs

parquet_derive_test/
lib.rs
