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