parquet_derive_test/
lib.rs

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

parquet_derive_test/
lib.rs
