Skip to main content

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