parquet_derive/
parquet_field.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
18use syn::ext::IdentExt;
19
20/// Parquet physical types used while selecting the code emitted by the macro.
21/// Keeping this local avoids compiling `parquet` for the proc-macro host.
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23enum PhysicalType {
24    Boolean,
25    Int32,
26    Int64,
27    Float,
28    Double,
29    ByteArray,
30    FixedLenByteArray,
31}
32
33#[derive(Debug, PartialEq)]
34pub struct Field {
35    ident: syn::Ident,
36    ty: Type,
37    is_a_byte_buf: bool,
38    third_party_type: Option<ThirdPartyType>,
39}
40
41/// Use third party libraries, detected
42/// at compile time. These libraries will
43/// be written to parquet as their preferred
44/// physical type.
45///
46///   ChronoNaiveDateTime is written as i64
47///   ChronoNaiveDate is written as i32
48#[derive(Debug, PartialEq)]
49enum ThirdPartyType {
50    ChronoNaiveDateTime,
51    ChronoNaiveDate,
52    Uuid,
53}
54
55impl Field {
56    pub fn from(f: &syn::Field) -> Self {
57        let ty = Type::from(f);
58        let is_a_byte_buf = ty.physical_type() == PhysicalType::ByteArray;
59
60        let third_party_type = match &ty.last_part()[..] {
61            "NaiveDateTime" => Some(ThirdPartyType::ChronoNaiveDateTime),
62            "NaiveDate" => Some(ThirdPartyType::ChronoNaiveDate),
63            "Uuid" => Some(ThirdPartyType::Uuid),
64            _ => None,
65        };
66
67        Field {
68            ident: f
69                .ident
70                .clone()
71                .expect("Only structs with named fields are currently supported"),
72            ty,
73            is_a_byte_buf,
74            third_party_type,
75        }
76    }
77
78    /// Takes the parsed field of the struct and emits a valid
79    /// column writer snippet. Should match exactly what you
80    /// would write by hand.
81    ///
82    /// Can only generate writers for basic structs, for example:
83    ///
84    /// struct Record {
85    ///   a_bool: bool,
86    ///   maybe_a_bool: `Option<bool>`
87    /// }
88    ///
89    /// but not
90    ///
91    /// struct UnsupportedNestedRecord {
92    ///   a_property: bool,
93    ///   nested_record: Record
94    /// }
95    ///
96    /// because this parsing logic is not sophisticated enough for definition
97    /// levels beyond 2.
98    pub fn writer_snippet(&self) -> proc_macro2::TokenStream {
99        let ident = &self.ident;
100        let column_writer = self.ty.column_writer();
101
102        let vals_builder = match &self.ty {
103            Type::TypePath(_) => self.copied_direct_vals(),
104            Type::Option(first_type) => match **first_type {
105                Type::TypePath(_) => self.option_into_vals(),
106                Type::Reference(_, ref second_type) => match **second_type {
107                    Type::TypePath(_) => self.option_into_vals(),
108                    _ => unimplemented!("Unsupported type encountered"),
109                },
110                Type::Vec(ref first_type) => match **first_type {
111                    Type::TypePath(_) => self.option_into_vals(),
112                    _ => unimplemented!("Unsupported type encountered"),
113                },
114                ref f => unimplemented!("Unsupported: {:#?}", f),
115            },
116            Type::Reference(_, first_type) => match **first_type {
117                Type::TypePath(_) => self.copied_direct_vals(),
118                Type::Option(ref second_type) => match **second_type {
119                    Type::TypePath(_) => self.option_into_vals(),
120                    Type::Reference(_, ref second_type) => match **second_type {
121                        Type::TypePath(_) => self.option_into_vals(),
122                        Type::Slice(ref second_type) => match **second_type {
123                            Type::TypePath(_) => self.option_into_vals(),
124                            ref f => unimplemented!("Unsupported: {:#?}", f),
125                        },
126                        _ => unimplemented!("Unsupported type encountered"),
127                    },
128                    Type::Vec(ref first_type) => match **first_type {
129                        Type::TypePath(_) => self.option_into_vals(),
130                        _ => unimplemented!("Unsupported type encountered"),
131                    },
132                    ref f => unimplemented!("Unsupported: {:#?}", f),
133                },
134                Type::Slice(ref second_type) => match **second_type {
135                    Type::TypePath(_) => self.copied_direct_vals(),
136                    ref f => unimplemented!("Unsupported: {:#?}", f),
137                },
138                ref f => unimplemented!("Unsupported: {:#?}", f),
139            },
140            Type::Vec(first_type) => match **first_type {
141                Type::TypePath(_) => self.copied_direct_vals(),
142                ref f => unimplemented!("Unsupported: {:#?}", f),
143            },
144            f => unimplemented!("Unsupported: {:#?}", f),
145        };
146
147        let definition_levels = match &self.ty {
148            Type::TypePath(_) => None,
149            Type::Option(first_type) => match **first_type {
150                Type::TypePath(_) => Some(self.optional_definition_levels()),
151                Type::Option(_) => unimplemented!("Unsupported nesting encountered"),
152                Type::Reference(_, ref second_type)
153                | Type::Vec(ref second_type)
154                | Type::Array(ref second_type, _)
155                | Type::Slice(ref second_type) => match **second_type {
156                    Type::TypePath(_) => Some(self.optional_definition_levels()),
157                    _ => unimplemented!("Unsupported nesting encountered"),
158                },
159            },
160            Type::Reference(_, first_type)
161            | Type::Vec(first_type)
162            | Type::Array(first_type, _)
163            | Type::Slice(first_type) => match **first_type {
164                Type::TypePath(_) => None,
165                Type::Vec(ref second_type)
166                | Type::Array(ref second_type, _)
167                | Type::Slice(ref second_type) => match **second_type {
168                    Type::TypePath(_) => None,
169                    Type::Reference(_, ref third_type) => match **third_type {
170                        Type::TypePath(_) => None,
171                        _ => unimplemented!("Unsupported definition encountered"),
172                    },
173                    _ => unimplemented!("Unsupported definition encountered"),
174                },
175                Type::Reference(_, ref second_type) | Type::Option(ref second_type) => {
176                    match **second_type {
177                        Type::TypePath(_) => Some(self.optional_definition_levels()),
178                        Type::Vec(ref third_type)
179                        | Type::Array(ref third_type, _)
180                        | Type::Slice(ref third_type) => match **third_type {
181                            Type::TypePath(_) => Some(self.optional_definition_levels()),
182                            Type::Reference(_, ref fourth_type) => match **fourth_type {
183                                Type::TypePath(_) => Some(self.optional_definition_levels()),
184                                _ => unimplemented!("Unsupported definition encountered"),
185                            },
186                            _ => unimplemented!("Unsupported definition encountered"),
187                        },
188                        Type::Reference(_, ref third_type) => match **third_type {
189                            Type::TypePath(_) => Some(self.optional_definition_levels()),
190                            Type::Slice(ref fourth_type) => match **fourth_type {
191                                Type::TypePath(_) => Some(self.optional_definition_levels()),
192                                _ => unimplemented!("Unsupported definition encountered"),
193                            },
194                            _ => unimplemented!("Unsupported definition encountered"),
195                        },
196                        _ => unimplemented!("Unsupported definition encountered"),
197                    }
198                }
199            },
200        };
201
202        // "vals" is the run of primitive data being written for the column
203        // "definition_levels" is a vector of bools which controls whether a value is missing or present
204        // this TokenStream is only one part of the code for writing a column and
205        // it relies on values calculated in prior code snippets, namely "definition_levels" and "vals_builder".
206        // All the context is put together in this functions final quote and
207        // this expression just switches between non-nullable and nullable write statements
208        let write_batch_expr = if definition_levels.is_some() {
209            quote! {
210                if let #column_writer(typed) = column_writer.untyped() {
211                    typed.write_batch(&vals[..], Some(&definition_levels[..]), None)?;
212                } else {
213                    panic!("Schema and struct disagree on type for {}", stringify!{#ident})
214                }
215            }
216        } else {
217            quote! {
218                if let #column_writer(typed) = column_writer.untyped() {
219                    typed.write_batch(&vals[..], None, None)?;
220                } else {
221                    panic!("Schema and struct disagree on type for {}", stringify!{#ident})
222                }
223            }
224        };
225
226        quote! {
227            {
228                #definition_levels
229
230                #vals_builder
231
232                #write_batch_expr
233            }
234        }
235    }
236
237    /// Takes the parsed field of the struct and emits a valid
238    /// column reader snippet. Should match exactly what you
239    /// would write by hand.
240    ///
241    /// Can only generate writers for basic structs, for example:
242    ///
243    /// struct Record {
244    ///   a_bool: bool
245    /// }
246    ///
247    /// but not
248    ///
249    /// struct UnsupportedNestedRecord {
250    ///   a_property: bool,
251    ///   nested_record: Record
252    /// }
253    ///
254    /// because this parsing logic is not sophisticated enough for definition
255    /// levels beyond 2.
256    ///
257    /// `Option` types and references not supported, but the column itself can be nullable
258    /// (i.e., def_level==1), as long as the values are all valid.
259    pub fn reader_snippet(&self) -> proc_macro2::TokenStream {
260        let ident = &self.ident;
261        let column_reader = self.ty.column_reader();
262
263        // generate the code to read the column into a vector `vals`
264        let write_batch_expr = quote! {
265            let mut vals = Vec::new();
266            if let #column_reader(mut typed) = column_reader {
267                let mut definition_levels = Vec::new();
268                let (total_num, valid_num, decoded_num) = typed.read_records(
269                    num_records, Some(&mut definition_levels), None, &mut vals)?;
270                if valid_num != decoded_num {
271                    panic!("Support only valid records, found {} null records in column type {}",
272                        decoded_num - valid_num, stringify!{#ident});
273                }
274            } else {
275                panic!("Schema and struct disagree on type for {}", stringify!{#ident});
276            }
277        };
278
279        // generate the code to convert each element of `vals` to the correct type and then write
280        // it to its field in the corresponding struct
281        let vals_writer = match &self.ty {
282            Type::TypePath(_) => self.copied_direct_fields(),
283            Type::Reference(_, first_type) => match **first_type {
284                Type::TypePath(_) => self.copied_direct_fields(),
285                Type::Slice(ref second_type) => match **second_type {
286                    Type::TypePath(_) => self.copied_direct_fields(),
287                    ref f => unimplemented!("Unsupported: {:#?}", f),
288                },
289                ref f => unimplemented!("Unsupported: {:#?}", f),
290            },
291            Type::Vec(first_type) => match **first_type {
292                Type::TypePath(_) => self.copied_direct_fields(),
293                ref f => unimplemented!("Unsupported: {:#?}", f),
294            },
295            f => unimplemented!("Unsupported: {:#?}", f),
296        };
297
298        quote! {
299            {
300                #write_batch_expr
301
302                #vals_writer
303            }
304        }
305    }
306
307    pub fn parquet_type(&self) -> proc_macro2::TokenStream {
308        // TODO: Support group types
309        // TODO: Add length if dealing with fixedlenbinary
310
311        // unraw the identifier, so a raw identifier like `r#type`
312        // becomes a column named `type` in the parquet schema
313        let field_name = self.ident.unraw().to_string();
314        let physical_type = match self.ty.physical_type() {
315            PhysicalType::Boolean => quote! {
316                ::parquet::basic::Type::BOOLEAN
317            },
318            PhysicalType::Int32 => quote! {
319                ::parquet::basic::Type::INT32
320            },
321            PhysicalType::Int64 => quote! {
322                ::parquet::basic::Type::INT64
323            },
324            PhysicalType::Float => quote! {
325                ::parquet::basic::Type::FLOAT
326            },
327            PhysicalType::Double => quote! {
328                ::parquet::basic::Type::DOUBLE
329            },
330            PhysicalType::ByteArray => quote! {
331                ::parquet::basic::Type::BYTE_ARRAY
332            },
333            PhysicalType::FixedLenByteArray => quote! {
334                ::parquet::basic::Type::FIXED_LEN_BYTE_ARRAY
335            },
336        };
337        let logical_type = self.ty.logical_type();
338        let repetition = self.ty.repetition();
339        let converted_type = self.ty.converted_type();
340        let length = self.ty.length();
341
342        let mut builder = quote! {
343            ParquetType::primitive_type_builder(#field_name, #physical_type)
344                .with_logical_type(#logical_type)
345                .with_repetition(#repetition)
346        };
347
348        if let Some(converted_type) = converted_type {
349            builder = quote! { #builder.with_converted_type(#converted_type) };
350        }
351
352        if let Some(length) = length {
353            builder = quote! { #builder.with_length(#length) };
354        }
355
356        quote! {  fields.push(#builder.build().unwrap().into()) }
357    }
358
359    fn option_into_vals(&self) -> proc_macro2::TokenStream {
360        let field_name = &self.ident;
361        let is_a_byte_buf = self.is_a_byte_buf;
362        let is_a_timestamp = self.third_party_type == Some(ThirdPartyType::ChronoNaiveDateTime);
363        let is_a_date = self.third_party_type == Some(ThirdPartyType::ChronoNaiveDate);
364        let is_a_uuid = self.third_party_type == Some(ThirdPartyType::Uuid);
365        let copy_to_vec = !matches!(
366            self.ty.physical_type(),
367            PhysicalType::ByteArray | PhysicalType::FixedLenByteArray
368        );
369
370        let binding = if copy_to_vec {
371            quote! { let Some(inner) = rec.#field_name }
372        } else {
373            quote! { let Some(inner) = &rec.#field_name }
374        };
375
376        let some = if is_a_timestamp {
377            quote! { Some(inner.timestamp_millis()) }
378        } else if is_a_date {
379            quote! { Some(inner.signed_duration_since(::chrono::NaiveDate::from_ymd(1970, 1, 1)).num_days() as i32)  }
380        } else if is_a_uuid {
381            quote! { Some((&inner.to_string()[..]).into()) }
382        } else if is_a_byte_buf {
383            quote! { Some((&inner[..]).into())}
384        } else {
385            // Type might need converting to a physical type
386            match self.ty.physical_type() {
387                PhysicalType::Int32 => quote! { Some(inner as i32) },
388                PhysicalType::Int64 => quote! { Some(inner as i64) },
389                _ => quote! { Some(inner) },
390            }
391        };
392
393        quote! {
394            let vals: Vec<_> = records.iter().filter_map(|rec| {
395                if #binding {
396                    #some
397                } else {
398                    None
399                }
400            }).collect();
401        }
402    }
403
404    // generates code to read `field_name` from each record into a vector `vals`
405    fn copied_direct_vals(&self) -> proc_macro2::TokenStream {
406        let field_name = &self.ident;
407
408        let access = match self.third_party_type {
409            Some(ThirdPartyType::ChronoNaiveDateTime) => {
410                quote! { rec.#field_name.timestamp_millis() }
411            }
412            Some(ThirdPartyType::ChronoNaiveDate) => {
413                quote! { rec.#field_name.signed_duration_since(::chrono::NaiveDate::from_ymd(1970, 1, 1)).num_days() as i32 }
414            }
415            Some(ThirdPartyType::Uuid) => {
416                quote! { rec.#field_name.as_bytes().to_vec().into() }
417            }
418            _ => {
419                if self.is_a_byte_buf {
420                    quote! { (&rec.#field_name[..]).into() }
421                } else {
422                    // Type might need converting to a physical type
423                    match self.ty.physical_type() {
424                        PhysicalType::Int32 => quote! { rec.#field_name as i32 },
425                        PhysicalType::Int64 => quote! { rec.#field_name as i64 },
426                        _ => quote! { rec.#field_name },
427                    }
428                }
429            }
430        };
431
432        quote! {
433            let vals: Vec<_> = records.iter().map(|rec| #access).collect();
434        }
435    }
436
437    // generates code to read a vector `records` into `field_name` for each record
438    fn copied_direct_fields(&self) -> proc_macro2::TokenStream {
439        let field_name = &self.ident;
440
441        let value = match self.third_party_type {
442            Some(ThirdPartyType::ChronoNaiveDateTime) => {
443                quote! { ::chrono::naive::NaiveDateTime::from_timestamp_millis(vals[i]).unwrap() }
444            }
445            Some(ThirdPartyType::ChronoNaiveDate) => {
446                // NaiveDateTime::UNIX_EPOCH.num_days_from_ce() == 719163
447                quote! {
448                    ::chrono::naive::NaiveDate::from_num_days_from_ce_opt(vals[i].saturating_add(719163)).unwrap()
449                }
450            }
451            Some(ThirdPartyType::Uuid) => {
452                quote! { ::uuid::Uuid::from_bytes(vals[i].data().try_into().unwrap()) }
453            }
454            _ => match &self.ty {
455                Type::TypePath(_) => match self.ty.last_part().as_str() {
456                    "String" => quote! { String::from(std::str::from_utf8(vals[i].data())
457                    .expect("invalid UTF-8 sequence")) },
458                    t => {
459                        let s: proc_macro2::TokenStream = t.parse().unwrap();
460                        quote! { vals[i] as #s }
461                    }
462                },
463                Type::Vec(_) => quote! { vals[i].data().to_vec() },
464                f => unimplemented!("Unsupported: {:#?}", f),
465            },
466        };
467
468        quote! {
469            for (i, r) in &mut records[..num_records].iter_mut().enumerate() {
470                r.#field_name = #value;
471            }
472        }
473    }
474
475    fn optional_definition_levels(&self) -> proc_macro2::TokenStream {
476        let field_name = &self.ident;
477
478        quote! {
479            let definition_levels: Vec<i16> = self
480              .iter()
481              .map(|rec| if rec.#field_name.is_some() { 1 } else { 0 })
482              .collect();
483        }
484    }
485}
486
487#[allow(clippy::enum_variant_names)]
488#[allow(clippy::large_enum_variant)]
489#[derive(Debug, PartialEq)]
490enum Type {
491    Array(Box<Type>, syn::Expr),
492    Option(Box<Type>),
493    Slice(Box<Type>),
494    Vec(Box<Type>),
495    TypePath(syn::Type),
496    Reference(Option<syn::Lifetime>, Box<Type>),
497}
498
499impl Type {
500    /// Takes a rust type and returns the appropriate
501    /// parquet-rs column writer
502    fn column_writer(&self) -> syn::TypePath {
503        match self.physical_type() {
504            PhysicalType::Boolean => {
505                syn::parse_quote!(ColumnWriter::BoolColumnWriter)
506            }
507            PhysicalType::Int32 => syn::parse_quote!(ColumnWriter::Int32ColumnWriter),
508            PhysicalType::Int64 => syn::parse_quote!(ColumnWriter::Int64ColumnWriter),
509            PhysicalType::Float => syn::parse_quote!(ColumnWriter::FloatColumnWriter),
510            PhysicalType::Double => syn::parse_quote!(ColumnWriter::DoubleColumnWriter),
511            PhysicalType::ByteArray => {
512                syn::parse_quote!(ColumnWriter::ByteArrayColumnWriter)
513            }
514            PhysicalType::FixedLenByteArray => {
515                syn::parse_quote!(ColumnWriter::FixedLenByteArrayColumnWriter)
516            }
517        }
518    }
519
520    /// Takes a rust type and returns the appropriate
521    /// parquet-rs column reader
522    fn column_reader(&self) -> syn::TypePath {
523        match self.physical_type() {
524            PhysicalType::Boolean => {
525                syn::parse_quote!(ColumnReader::BoolColumnReader)
526            }
527            PhysicalType::Int32 => syn::parse_quote!(ColumnReader::Int32ColumnReader),
528            PhysicalType::Int64 => syn::parse_quote!(ColumnReader::Int64ColumnReader),
529            PhysicalType::Float => syn::parse_quote!(ColumnReader::FloatColumnReader),
530            PhysicalType::Double => syn::parse_quote!(ColumnReader::DoubleColumnReader),
531            PhysicalType::ByteArray => {
532                syn::parse_quote!(ColumnReader::ByteArrayColumnReader)
533            }
534            PhysicalType::FixedLenByteArray => {
535                syn::parse_quote!(ColumnReader::FixedLenByteArrayColumnReader)
536            }
537        }
538    }
539
540    /// Helper to simplify a nested field definition to its leaf type
541    ///
542    /// Ex:
543    ///   `Option<&String>` => Type::TypePath(String)
544    ///   `&Option<i32>` => Type::TypePath(i32)
545    ///   `Vec<Vec<u8>>` => Type::Vec(u8)
546    ///
547    /// Useful in determining the physical type of a field and the
548    /// definition levels.
549    fn leaf_type_recursive(&self) -> &Type {
550        Type::leaf_type_recursive_helper(self, None)
551    }
552
553    fn leaf_type_recursive_helper<'a>(ty: &'a Type, parent_ty: Option<&'a Type>) -> &'a Type {
554        match ty {
555            Type::TypePath(_) => parent_ty.unwrap_or(ty),
556            Type::Option(first_type)
557            | Type::Vec(first_type)
558            | Type::Array(first_type, _)
559            | Type::Slice(first_type)
560            | Type::Reference(_, first_type) => {
561                Type::leaf_type_recursive_helper(first_type, Some(ty))
562            }
563        }
564    }
565
566    /// Helper method to further unwrap leaf_type() to get inner-most
567    /// type information, useful for determining the physical type
568    /// and normalizing the type paths.
569    fn inner_type(&self) -> &syn::Type {
570        let leaf_type = self.leaf_type_recursive();
571
572        match leaf_type {
573            Type::TypePath(type_) => type_,
574            Type::Option(first_type)
575            | Type::Vec(first_type)
576            | Type::Array(first_type, _)
577            | Type::Slice(first_type)
578            | Type::Reference(_, first_type) => match **first_type {
579                Type::TypePath(ref type_) => type_,
580                _ => unimplemented!("leaf_type() should only return shallow types"),
581            },
582        }
583    }
584
585    /// Helper to normalize a type path by extracting the
586    /// most identifiable part
587    ///
588    /// Ex:
589    ///   std::string::String => String
590    ///   `Vec<u8>` => `Vec<u8>`
591    ///   chrono::NaiveDateTime => NaiveDateTime
592    ///
593    /// Does run the risk of mis-identifying a type if import
594    /// rename is in play. Please note procedural macros always
595    /// run before type resolution so this is a risk the user
596    /// takes on when renaming imports.
597    fn last_part(&self) -> String {
598        let inner_type = self.inner_type();
599        let inner_type_str = (quote! { #inner_type }).to_string();
600
601        inner_type_str
602            .split("::")
603            .last()
604            .unwrap()
605            .trim()
606            .to_string()
607    }
608
609    /// Converts rust types to parquet physical types.
610    ///
611    /// Ex:
612    ///   [u8; 10] => FIXED_LEN_BYTE_ARRAY
613    ///   `Vec<u8>`  => BYTE_ARRAY
614    ///   String => BYTE_ARRAY
615    ///   i32 => INT32
616    fn physical_type(&self) -> PhysicalType {
617        let last_part = self.last_part();
618        let leaf_type = self.leaf_type_recursive();
619
620        match leaf_type {
621            Type::Array(first_type, _length) => {
622                if let Type::TypePath(_) = **first_type {
623                    if last_part == "u8" {
624                        return PhysicalType::FixedLenByteArray;
625                    }
626                }
627            }
628            Type::Vec(first_type) | Type::Slice(first_type) => {
629                if let Type::TypePath(_) = **first_type {
630                    if last_part == "u8" {
631                        return PhysicalType::ByteArray;
632                    }
633                }
634            }
635            _ => (),
636        }
637
638        match last_part.trim() {
639            "bool" => PhysicalType::Boolean,
640            "u8" | "u16" | "u32" => PhysicalType::Int32,
641            "i8" | "i16" | "i32" | "NaiveDate" => PhysicalType::Int32,
642            "u64" | "i64" | "NaiveDateTime" => PhysicalType::Int64,
643            "usize" | "isize" => {
644                if usize::BITS == 64 {
645                    PhysicalType::Int64
646                } else {
647                    PhysicalType::Int32
648                }
649            }
650            "f32" => PhysicalType::Float,
651            "f64" => PhysicalType::Double,
652            "String" | "str" | "Arc < str >" => PhysicalType::ByteArray,
653            "Uuid" => PhysicalType::FixedLenByteArray,
654            f => unimplemented!("{} currently is not supported", f),
655        }
656    }
657
658    fn length(&self) -> Option<syn::Expr> {
659        let last_part = self.last_part();
660        let leaf_type = self.leaf_type_recursive();
661
662        // `[u8; N]` => Some(N)
663        if let Type::Array(first_type, length) = leaf_type {
664            if let Type::TypePath(_) = **first_type {
665                if last_part == "u8" {
666                    return Some(length.clone());
667                }
668            }
669        }
670
671        match last_part.trim() {
672            // Uuid => [u8; 16] => Some(16)
673            "Uuid" => Some(syn::parse_quote!(16)),
674            _ => None,
675        }
676    }
677
678    fn logical_type(&self) -> proc_macro2::TokenStream {
679        let last_part = self.last_part();
680        let leaf_type = self.leaf_type_recursive();
681
682        match leaf_type {
683            Type::Array(first_type, _length) => {
684                if let Type::TypePath(_) = **first_type {
685                    if last_part == "u8" {
686                        return quote! { None };
687                    }
688                }
689            }
690            Type::Vec(first_type) | Type::Slice(first_type) => {
691                if let Type::TypePath(_) = **first_type {
692                    if last_part == "u8" {
693                        return quote! { None };
694                    }
695                }
696            }
697            _ => (),
698        }
699
700        match last_part.trim() {
701            "bool" => quote! { None },
702            "u8" => quote! { Some(LogicalType::integer(8, false)) },
703            "u16" => quote! { Some(LogicalType::integer(16, false)) },
704            "u32" => quote! { Some(LogicalType::integer(32, false)) },
705            "u64" => quote! { Some(LogicalType::integer(64, false)) },
706            "i8" => quote! { Some(LogicalType::integer(8, true)) },
707            "i16" => quote! { Some(LogicalType::integer(16, true)) },
708            "i32" | "i64" => quote! { None },
709            "usize" => {
710                quote! { Some(LogicalType::integer(usize::BITS as i8, false)) }
711            }
712            "isize" => {
713                quote! { Some(LogicalType::integer(usize::BITS as i8, true)) }
714            }
715            "NaiveDate" => quote! { Some(LogicalType::Date) },
716            "NaiveDateTime" => quote! { None },
717            "f32" | "f64" => quote! { None },
718            "String" | "str" | "Arc < str >" => quote! { Some(LogicalType::String) },
719            "Uuid" => quote! { Some(LogicalType::Uuid) },
720            f => unimplemented!("{} currently is not supported", f),
721        }
722    }
723
724    fn converted_type(&self) -> Option<proc_macro2::TokenStream> {
725        let last_part = self.last_part();
726
727        match last_part.trim() {
728            "NaiveDateTime" => Some(quote! { ::parquet::basic::ConvertedType::TIMESTAMP_MILLIS }),
729            _ => None,
730        }
731    }
732
733    fn repetition(&self) -> proc_macro2::TokenStream {
734        match self {
735            Type::Option(_) => quote! { ::parquet::basic::Repetition::OPTIONAL },
736            Type::Reference(_, ty) => ty.repetition(),
737            _ => quote! { ::parquet::basic::Repetition::REQUIRED },
738        }
739    }
740
741    /// Convert a parsed rust field AST in to a more easy to manipulate
742    /// parquet_derive::Field
743    fn from(f: &syn::Field) -> Self {
744        Type::from_type(f, &f.ty)
745    }
746
747    fn from_type(f: &syn::Field, ty: &syn::Type) -> Self {
748        match ty {
749            syn::Type::Path(p) => Type::from_type_path(f, p),
750            syn::Type::Reference(tr) => Type::from_type_reference(f, tr),
751            syn::Type::Array(ta) => Type::from_type_array(f, ta),
752            syn::Type::Slice(ts) => Type::from_type_slice(f, ts),
753            other => unimplemented!(
754                "Unable to derive {:?} - it is currently an unsupported type\n{:#?}",
755                f.ident.as_ref().unwrap(),
756                other
757            ),
758        }
759    }
760
761    fn from_type_path(f: &syn::Field, p: &syn::TypePath) -> Self {
762        let last_segment = p.path.segments.last().unwrap();
763
764        let is_vec = last_segment.ident == syn::Ident::new("Vec", proc_macro2::Span::call_site());
765        let is_option =
766            last_segment.ident == syn::Ident::new("Option", proc_macro2::Span::call_site());
767
768        if is_vec || is_option {
769            let generic_type = match &last_segment.arguments {
770                syn::PathArguments::AngleBracketed(angle_args) => {
771                    assert_eq!(angle_args.args.len(), 1);
772                    let first_arg = &angle_args.args[0];
773
774                    match first_arg {
775                        syn::GenericArgument::Type(typath) => typath.clone(),
776                        other => unimplemented!("Unsupported: {:#?}", other),
777                    }
778                }
779                other => unimplemented!("Unsupported: {:#?}", other),
780            };
781
782            if is_vec {
783                Type::Vec(Box::new(Type::from_type(f, &generic_type)))
784            } else {
785                Type::Option(Box::new(Type::from_type(f, &generic_type)))
786            }
787        } else {
788            Type::TypePath(syn::Type::Path(p.clone()))
789        }
790    }
791
792    fn from_type_reference(f: &syn::Field, tr: &syn::TypeReference) -> Self {
793        let lifetime = tr.lifetime.clone();
794        let inner_type = Type::from_type(f, tr.elem.as_ref());
795        Type::Reference(lifetime, Box::new(inner_type))
796    }
797
798    fn from_type_array(f: &syn::Field, ta: &syn::TypeArray) -> Self {
799        let inner_type = Type::from_type(f, ta.elem.as_ref());
800        Type::Array(Box::new(inner_type), ta.len.clone())
801    }
802
803    fn from_type_slice(f: &syn::Field, ts: &syn::TypeSlice) -> Self {
804        let inner_type = Type::from_type(f, ts.elem.as_ref());
805        Type::Slice(Box::new(inner_type))
806    }
807}
808
809#[cfg(test)]
810mod test {
811    use super::*;
812    use syn::{Data, DataStruct, DeriveInput};
813
814    fn extract_fields(input: proc_macro2::TokenStream) -> Vec<syn::Field> {
815        let input: DeriveInput = syn::parse2(input).unwrap();
816
817        let fields = match input.data {
818            Data::Struct(DataStruct { fields, .. }) => fields,
819            _ => panic!("Input must be a struct"),
820        };
821
822        fields.iter().map(|field| field.to_owned()).collect()
823    }
824
825    #[test]
826    fn test_generating_a_simple_writer_snippet() {
827        let snippet: proc_macro2::TokenStream = quote! {
828          struct ABoringStruct {
829            counter: usize,
830          }
831        };
832
833        let fields = extract_fields(snippet);
834        let counter = Field::from(&fields[0]);
835
836        let snippet = counter.writer_snippet().to_string();
837        assert_eq!(snippet,
838                   (quote!{
839                        {
840                            let vals : Vec < _ > = records . iter ( ) . map ( | rec | rec . counter as i64 ) . collect ( );
841
842                            if let ColumnWriter::Int64ColumnWriter ( typed ) = column_writer.untyped() {
843                                typed . write_batch ( & vals [ .. ] , None , None ) ?;
844                            }  else {
845                                panic!("Schema and struct disagree on type for {}" , stringify!{ counter } )
846                            }
847                        }
848                   }).to_string()
849        )
850    }
851
852    #[test]
853    fn test_generating_a_simple_reader_snippet() {
854        let snippet: proc_macro2::TokenStream = quote! {
855          struct ABoringStruct {
856            counter: usize,
857          }
858        };
859
860        let fields = extract_fields(snippet);
861        let counter = Field::from(&fields[0]);
862
863        let snippet = counter.reader_snippet().to_string();
864        assert_eq!(
865            snippet,
866            (quote! {
867                 {
868                    let mut vals = Vec::new();
869                    if let ColumnReader::Int64ColumnReader(mut typed) = column_reader {
870                        let mut definition_levels = Vec::new();
871                        let (total_num, valid_num, decoded_num) = typed.read_records(
872                            num_records, Some(&mut definition_levels), None, &mut vals)?;
873                        if valid_num != decoded_num {
874                            panic!("Support only valid records, found {} null records in column type {}",
875                                decoded_num - valid_num, stringify!{counter});
876                        }
877                    } else {
878                        panic!("Schema and struct disagree on type for {}", stringify!{counter});
879                    }
880                    for (i, r) in &mut records[..num_records].iter_mut().enumerate() {
881                        r.counter = vals[i] as usize;
882                    }
883                 }
884            })
885            .to_string()
886        )
887    }
888
889    #[test]
890    fn test_parquet_type_with_raw_identifier() {
891        let snippet: proc_macro2::TokenStream = quote! {
892          struct ABoringStruct {
893            r#type: i32,
894          }
895        };
896
897        let fields = extract_fields(snippet);
898        let r#type = Field::from(&fields[0]);
899
900        // the raw identifier `r#type` is named `type` in the parquet schema
901        let snippet = r#type.parquet_type().to_string();
902        assert!(
903            snippet.contains("primitive_type_builder (\"type\""),
904            "{snippet}"
905        );
906    }
907
908    #[test]
909    fn test_optional_to_writer_snippet() {
910        let struct_def: proc_macro2::TokenStream = quote! {
911          struct StringBorrower<'a> {
912            optional_str: Option<&'a str>,
913            optional_string: Option<&String>,
914            optional_dumb_int: Option<&i32>,
915          }
916        };
917
918        let fields = extract_fields(struct_def);
919
920        let optional = Field::from(&fields[0]);
921        let snippet = optional.writer_snippet();
922        assert_eq!(snippet.to_string(),
923          (quote! {
924          {
925                let definition_levels : Vec < i16 > = self . iter ( ) . map ( | rec | if rec . optional_str . is_some ( ) { 1 } else { 0 } ) . collect ( ) ;
926
927                let vals: Vec <_> = records.iter().filter_map( |rec| {
928                    if let Some ( inner ) = &rec . optional_str {
929                        Some ( (&inner[..]).into() )
930                    } else {
931                        None
932                    }
933                }).collect();
934
935                if let ColumnWriter::ByteArrayColumnWriter ( typed ) = column_writer.untyped() {
936                    typed . write_batch ( & vals [ .. ] , Some(&definition_levels[..]) , None ) ? ;
937                } else {
938                    panic!("Schema and struct disagree on type for {}" , stringify ! { optional_str } )
939                }
940           }
941            }
942          ).to_string());
943
944        let optional = Field::from(&fields[1]);
945        let snippet = optional.writer_snippet();
946        assert_eq!(snippet.to_string(),
947                   (quote!{
948                   {
949                        let definition_levels : Vec < i16 > = self . iter ( ) . map ( | rec | if rec . optional_string . is_some ( ) { 1 } else { 0 } ) . collect ( ) ;
950
951                        let vals: Vec <_> = records.iter().filter_map( |rec| {
952                            if let Some ( inner ) = &rec . optional_string {
953                                Some ( (&inner[..]).into() )
954                            } else {
955                                None
956                            }
957                        }).collect();
958
959                        if let ColumnWriter::ByteArrayColumnWriter ( typed ) = column_writer.untyped() {
960                            typed . write_batch ( & vals [ .. ] , Some(&definition_levels[..]) , None ) ? ;
961                        } else {
962                            panic!("Schema and struct disagree on type for {}" , stringify ! { optional_string } )
963                        }
964                    }
965        }).to_string());
966
967        let optional = Field::from(&fields[2]);
968        let snippet = optional.writer_snippet();
969        assert_eq!(snippet.to_string(),
970                   (quote!{
971                    {
972                        let definition_levels : Vec < i16 > = self . iter ( ) . map ( | rec | if rec . optional_dumb_int . is_some ( ) { 1 } else { 0 } ) . collect ( ) ;
973
974                        let vals: Vec <_> = records.iter().filter_map( |rec| {
975                            if let Some ( inner ) = rec . optional_dumb_int {
976                                Some ( inner as i32 )
977                            } else {
978                                None
979                            }
980                        }).collect();
981
982                        if let ColumnWriter::Int32ColumnWriter ( typed ) = column_writer.untyped() {
983                            typed . write_batch ( & vals [ .. ] , Some(&definition_levels[..]) , None ) ? ;
984                        }  else {
985                            panic!("Schema and struct disagree on type for {}" , stringify ! { optional_dumb_int } )
986                        }
987                    }
988        }).to_string());
989    }
990
991    #[test]
992    fn test_converting_to_column_writer_type() {
993        let snippet: proc_macro2::TokenStream = quote! {
994          struct ABasicStruct {
995            yes_no: bool,
996            name: String,
997          }
998        };
999
1000        let fields = extract_fields(snippet);
1001        let processed: Vec<_> = fields.iter().map(Field::from).collect();
1002
1003        let column_writers: Vec<_> = processed
1004            .iter()
1005            .map(|field| field.ty.column_writer())
1006            .collect();
1007
1008        assert_eq!(
1009            column_writers,
1010            vec![
1011                syn::parse_quote!(ColumnWriter::BoolColumnWriter),
1012                syn::parse_quote!(ColumnWriter::ByteArrayColumnWriter)
1013            ]
1014        );
1015    }
1016
1017    #[test]
1018    fn test_converting_to_column_reader_type() {
1019        let snippet: proc_macro2::TokenStream = quote! {
1020          struct ABasicStruct {
1021            yes_no: bool,
1022            name: String,
1023          }
1024        };
1025
1026        let fields = extract_fields(snippet);
1027        let processed: Vec<_> = fields.iter().map(Field::from).collect();
1028
1029        let column_readers: Vec<_> = processed
1030            .iter()
1031            .map(|field| field.ty.column_reader())
1032            .collect();
1033
1034        assert_eq!(
1035            column_readers,
1036            vec![
1037                syn::parse_quote!(ColumnReader::BoolColumnReader),
1038                syn::parse_quote!(ColumnReader::ByteArrayColumnReader)
1039            ]
1040        );
1041    }
1042
1043    #[test]
1044    fn convert_basic_struct() {
1045        let snippet: proc_macro2::TokenStream = quote! {
1046          struct ABasicStruct {
1047            yes_no: bool,
1048            name: String,
1049            length: usize
1050          }
1051        };
1052
1053        let fields = extract_fields(snippet);
1054        let processed: Vec<_> = fields.iter().map(Field::from).collect();
1055        assert_eq!(processed.len(), 3);
1056
1057        assert_eq!(
1058            processed,
1059            vec![
1060                Field {
1061                    ident: syn::Ident::new("yes_no", proc_macro2::Span::call_site()),
1062                    ty: Type::TypePath(syn::parse_quote!(bool)),
1063                    is_a_byte_buf: false,
1064                    third_party_type: None,
1065                },
1066                Field {
1067                    ident: syn::Ident::new("name", proc_macro2::Span::call_site()),
1068                    ty: Type::TypePath(syn::parse_quote!(String)),
1069                    is_a_byte_buf: true,
1070                    third_party_type: None,
1071                },
1072                Field {
1073                    ident: syn::Ident::new("length", proc_macro2::Span::call_site()),
1074                    ty: Type::TypePath(syn::parse_quote!(usize)),
1075                    is_a_byte_buf: false,
1076                    third_party_type: None,
1077                }
1078            ]
1079        )
1080    }
1081
1082    #[test]
1083    fn test_get_inner_type() {
1084        let snippet: proc_macro2::TokenStream = quote! {
1085          struct LotsOfInnerTypes {
1086            a_vec: Vec<u8>,
1087            a_option: ::std::option::Option<bool>,
1088            a_silly_string: ::std::string::String,
1089            a_complicated_thing: ::std::option::Option<::std::result::Result<(),()>>,
1090          }
1091        };
1092
1093        let fields = extract_fields(snippet);
1094        let converted_fields: Vec<_> = fields.iter().map(Type::from).collect();
1095        let inner_types: Vec<_> = converted_fields
1096            .iter()
1097            .map(|field| field.inner_type())
1098            .collect();
1099        let inner_types_strs: Vec<_> = inner_types
1100            .iter()
1101            .map(|ty| (quote! { #ty }).to_string())
1102            .collect();
1103
1104        assert_eq!(
1105            inner_types_strs,
1106            vec![
1107                "u8",
1108                "bool",
1109                ":: std :: string :: String",
1110                ":: std :: result :: Result < () , () >"
1111            ]
1112        )
1113    }
1114
1115    #[test]
1116    fn test_physical_type() {
1117        let snippet: proc_macro2::TokenStream = quote! {
1118          struct LotsOfInnerTypes {
1119            a_buf: ::std::vec::Vec<u8>,
1120            a_number: i32,
1121            a_verbose_option: ::std::option::Option<bool>,
1122            a_silly_string: String,
1123            a_fix_byte_buf: [u8; 10],
1124            a_complex_option: ::std::option::Option<&Vec<u8>>,
1125            a_complex_vec: &::std::vec::Vec<&Option<u8>>,
1126            a_uuid: ::uuid::Uuid,
1127          }
1128        };
1129
1130        let fields = extract_fields(snippet);
1131        let converted_fields: Vec<_> = fields.iter().map(Type::from).collect();
1132        let physical_types: Vec<_> = converted_fields
1133            .iter()
1134            .map(|ty| ty.physical_type())
1135            .collect();
1136
1137        assert_eq!(
1138            physical_types,
1139            vec![
1140                PhysicalType::ByteArray,
1141                PhysicalType::Int32,
1142                PhysicalType::Boolean,
1143                PhysicalType::ByteArray,
1144                PhysicalType::FixedLenByteArray,
1145                PhysicalType::ByteArray,
1146                PhysicalType::Int32,
1147                PhysicalType::FixedLenByteArray,
1148            ]
1149        )
1150    }
1151
1152    #[test]
1153    fn test_type_length() {
1154        let snippet: proc_macro2::TokenStream = quote! {
1155          struct LotsOfInnerTypes {
1156            a_buf: ::std::vec::Vec<u8>,
1157            a_number: i32,
1158            a_verbose_option: ::std::option::Option<bool>,
1159            a_silly_string: String,
1160            a_fix_byte_buf: [u8; 10],
1161            a_complex_option: ::std::option::Option<&Vec<u8>>,
1162            a_complex_vec: &::std::vec::Vec<&Option<u8>>,
1163            a_uuid: ::uuid::Uuid,
1164          }
1165        };
1166
1167        let fields = extract_fields(snippet);
1168        let converted_fields: Vec<_> = fields.iter().map(Type::from).collect();
1169        let lengths: Vec<_> = converted_fields.iter().map(|ty| ty.length()).collect();
1170
1171        assert_eq!(
1172            lengths,
1173            vec![
1174                None,
1175                None,
1176                None,
1177                None,
1178                Some(syn::parse_quote!(10)),
1179                None,
1180                None,
1181                Some(syn::parse_quote!(16)),
1182            ]
1183        )
1184    }
1185
1186    #[test]
1187    fn test_convert_comprehensive_owned_struct() {
1188        let snippet: proc_macro2::TokenStream = quote! {
1189          struct VecHolder {
1190            a_vec: ::std::vec::Vec<u8>,
1191            a_option: ::std::option::Option<bool>,
1192            a_silly_string: ::std::string::String,
1193            a_complicated_thing: ::std::option::Option<::std::result::Result<(),()>>,
1194          }
1195        };
1196
1197        let fields = extract_fields(snippet);
1198        let converted_fields: Vec<_> = fields.iter().map(Type::from).collect();
1199
1200        assert_eq!(
1201            converted_fields,
1202            vec![
1203                Type::Vec(Box::new(Type::TypePath(syn::parse_quote!(u8)))),
1204                Type::Option(Box::new(Type::TypePath(syn::parse_quote!(bool)))),
1205                Type::TypePath(syn::parse_quote!(::std::string::String)),
1206                Type::Option(Box::new(Type::TypePath(
1207                    syn::parse_quote!(::std::result::Result<(),()>)
1208                ))),
1209            ]
1210        );
1211    }
1212
1213    #[test]
1214    fn test_convert_borrowed_struct() {
1215        let snippet: proc_macro2::TokenStream = quote! {
1216          struct Borrower<'a> {
1217            a_str: &'a str,
1218            a_borrowed_option: &'a Option<bool>,
1219            so_many_borrows: &'a Option<&'a str>,
1220          }
1221        };
1222
1223        let fields = extract_fields(snippet);
1224        let types: Vec<_> = fields.iter().map(Type::from).collect();
1225
1226        assert_eq!(
1227            types,
1228            vec![
1229                Type::Reference(
1230                    Some(syn::Lifetime::new("'a", proc_macro2::Span::call_site())),
1231                    Box::new(Type::TypePath(syn::parse_quote!(str)))
1232                ),
1233                Type::Reference(
1234                    Some(syn::Lifetime::new("'a", proc_macro2::Span::call_site())),
1235                    Box::new(Type::Option(Box::new(Type::TypePath(syn::parse_quote!(
1236                        bool
1237                    )))))
1238                ),
1239                Type::Reference(
1240                    Some(syn::Lifetime::new("'a", proc_macro2::Span::call_site())),
1241                    Box::new(Type::Option(Box::new(Type::Reference(
1242                        Some(syn::Lifetime::new("'a", proc_macro2::Span::call_site())),
1243                        Box::new(Type::TypePath(syn::parse_quote!(str)))
1244                    ))))
1245                ),
1246            ]
1247        );
1248    }
1249
1250    #[test]
1251    fn test_chrono_timestamp_millis_write() {
1252        let snippet: proc_macro2::TokenStream = quote! {
1253          struct ATimestampStruct {
1254            henceforth: chrono::NaiveDateTime,
1255            maybe_happened: Option<&chrono::NaiveDateTime>,
1256          }
1257        };
1258
1259        let fields = extract_fields(snippet);
1260        let when = Field::from(&fields[0]);
1261        assert_eq!(when.writer_snippet().to_string(),(quote!{
1262            {
1263                let vals : Vec<_> = records.iter().map(|rec| rec.henceforth.timestamp_millis() ).collect();
1264                if let ColumnWriter::Int64ColumnWriter(typed) = column_writer.untyped() {
1265                    typed.write_batch(&vals[..], None, None) ?;
1266                } else {
1267                    panic!("Schema and struct disagree on type for {}" , stringify!{ henceforth })
1268                }
1269            }
1270        }).to_string());
1271
1272        let maybe_happened = Field::from(&fields[1]);
1273        assert_eq!(maybe_happened.writer_snippet().to_string(),(quote!{
1274            {
1275                let definition_levels : Vec<i16> = self.iter().map(|rec| if rec.maybe_happened.is_some() { 1 } else { 0 }).collect();
1276                let vals : Vec<_> = records.iter().filter_map(|rec| {
1277                    if let Some(inner) = rec.maybe_happened {
1278                        Some(inner.timestamp_millis())
1279                    } else {
1280                        None
1281                    }
1282                }).collect();
1283
1284                if let ColumnWriter::Int64ColumnWriter(typed) = column_writer.untyped() {
1285                    typed.write_batch(&vals[..], Some(&definition_levels[..]), None) ?;
1286                } else {
1287                    panic!("Schema and struct disagree on type for {}" , stringify!{ maybe_happened })
1288                }
1289            }
1290        }).to_string());
1291    }
1292
1293    #[test]
1294    fn test_chrono_timestamp_millis_read() {
1295        let snippet: proc_macro2::TokenStream = quote! {
1296          struct ATimestampStruct {
1297            henceforth: chrono::NaiveDateTime,
1298          }
1299        };
1300
1301        let fields = extract_fields(snippet);
1302        let when = Field::from(&fields[0]);
1303        assert_eq!(when.reader_snippet().to_string(),(quote!{
1304            {
1305                let mut vals = Vec::new();
1306                if let ColumnReader::Int64ColumnReader(mut typed) = column_reader {
1307                    let mut definition_levels = Vec::new();
1308                    let (total_num, valid_num, decoded_num) = typed.read_records(
1309                        num_records, Some(&mut definition_levels), None, &mut vals)?;
1310                    if valid_num != decoded_num {
1311                        panic!("Support only valid records, found {} null records in column type {}",
1312                            decoded_num - valid_num, stringify!{henceforth});
1313                    }
1314                } else {
1315                    panic!("Schema and struct disagree on type for {}", stringify!{ henceforth });
1316                }
1317                for (i, r) in &mut records[..num_records].iter_mut().enumerate() {
1318                    r.henceforth = ::chrono::naive::NaiveDateTime::from_timestamp_millis(vals[i]).unwrap();
1319                }
1320            }
1321        }).to_string());
1322    }
1323
1324    #[test]
1325    fn test_chrono_date_write() {
1326        let snippet: proc_macro2::TokenStream = quote! {
1327          struct ATimestampStruct {
1328            henceforth: chrono::NaiveDate,
1329            maybe_happened: Option<&chrono::NaiveDate>,
1330          }
1331        };
1332
1333        let fields = extract_fields(snippet);
1334        let when = Field::from(&fields[0]);
1335        assert_eq!(when.writer_snippet().to_string(),(quote!{
1336            {
1337                let vals : Vec<_> = records.iter().map(|rec| rec.henceforth.signed_duration_since(::chrono::NaiveDate::from_ymd(1970, 1, 1)).num_days() as i32).collect();
1338                if let ColumnWriter::Int32ColumnWriter(typed) = column_writer.untyped() {
1339                    typed.write_batch(&vals[..], None, None) ?;
1340                } else {
1341                    panic!("Schema and struct disagree on type for {}" , stringify!{ henceforth })
1342                }
1343            }
1344        }).to_string());
1345
1346        let maybe_happened = Field::from(&fields[1]);
1347        assert_eq!(maybe_happened.writer_snippet().to_string(),(quote!{
1348            {
1349                let definition_levels : Vec<i16> = self.iter().map(|rec| if rec.maybe_happened.is_some() { 1 } else { 0 }).collect();
1350                let vals : Vec<_> = records.iter().filter_map(|rec| {
1351                    if let Some(inner) = rec.maybe_happened {
1352                        Some(inner.signed_duration_since(::chrono::NaiveDate::from_ymd(1970, 1, 1)).num_days() as i32)
1353                    } else {
1354                        None
1355                    }
1356                }).collect();
1357
1358                if let ColumnWriter::Int32ColumnWriter(typed) = column_writer.untyped() {
1359                    typed.write_batch(&vals[..], Some(&definition_levels[..]), None) ?;
1360                } else {
1361                    panic!("Schema and struct disagree on type for {}" , stringify!{ maybe_happened })
1362                }
1363            }
1364        }).to_string());
1365    }
1366
1367    #[test]
1368    fn test_chrono_date_read() {
1369        let snippet: proc_macro2::TokenStream = quote! {
1370          struct ATimestampStruct {
1371            henceforth: chrono::NaiveDate,
1372          }
1373        };
1374
1375        let fields = extract_fields(snippet);
1376        let when = Field::from(&fields[0]);
1377        assert_eq!(when.reader_snippet().to_string(),(quote!{
1378            {
1379                let mut vals = Vec::new();
1380                if let ColumnReader::Int32ColumnReader(mut typed) = column_reader {
1381                    let mut definition_levels = Vec::new();
1382                    let (total_num, valid_num, decoded_num) = typed.read_records(
1383                        num_records, Some(&mut definition_levels), None, &mut vals)?;
1384                    if valid_num != decoded_num {
1385                        panic!("Support only valid records, found {} null records in column type {}",
1386                            decoded_num - valid_num, stringify!{henceforth});
1387                    }
1388                } else {
1389                    panic!("Schema and struct disagree on type for {}", stringify!{ henceforth });
1390                }
1391                for (i, r) in &mut records[..num_records].iter_mut().enumerate() {
1392                    r.henceforth = ::chrono::naive::NaiveDate::from_num_days_from_ce_opt(vals[i].saturating_add(719163)).unwrap();
1393                }
1394            }
1395        }).to_string());
1396    }
1397
1398    #[test]
1399    fn test_uuid_write() {
1400        let snippet: proc_macro2::TokenStream = quote! {
1401          struct AUuidStruct {
1402            unique_id: uuid::Uuid,
1403            maybe_unique_id: Option<&uuid::Uuid>,
1404          }
1405        };
1406
1407        let fields = extract_fields(snippet);
1408        let when = Field::from(&fields[0]);
1409        assert_eq!(when.writer_snippet().to_string(),(quote!{
1410            {
1411                let vals : Vec<_> = records.iter().map(|rec| rec.unique_id.as_bytes().to_vec().into() ).collect();
1412                if let ColumnWriter::FixedLenByteArrayColumnWriter(typed) = column_writer.untyped() {
1413                    typed.write_batch(&vals[..], None, None) ?;
1414                } else {
1415                    panic!("Schema and struct disagree on type for {}" , stringify!{ unique_id })
1416                }
1417            }
1418        }).to_string());
1419
1420        let maybe_happened = Field::from(&fields[1]);
1421        assert_eq!(maybe_happened.writer_snippet().to_string(),(quote!{
1422            {
1423                let definition_levels : Vec<i16> = self.iter().map(|rec| if rec.maybe_unique_id.is_some() { 1 } else { 0 }).collect();
1424                let vals : Vec<_> = records.iter().filter_map(|rec| {
1425                    if let Some(inner) = &rec.maybe_unique_id {
1426                        Some((&inner.to_string()[..]).into())
1427                    } else {
1428                        None
1429                    }
1430                }).collect();
1431
1432                if let ColumnWriter::FixedLenByteArrayColumnWriter(typed) = column_writer.untyped() {
1433                    typed.write_batch(&vals[..], Some(&definition_levels[..]), None) ?;
1434                } else {
1435                    panic!("Schema and struct disagree on type for {}" , stringify!{ maybe_unique_id })
1436                }
1437            }
1438        }).to_string());
1439    }
1440
1441    #[test]
1442    fn test_uuid_read() {
1443        let snippet: proc_macro2::TokenStream = quote! {
1444          struct AUuidStruct {
1445            unique_id: uuid::Uuid,
1446          }
1447        };
1448
1449        let fields = extract_fields(snippet);
1450        let when = Field::from(&fields[0]);
1451        assert_eq!(when.reader_snippet().to_string(),(quote!{
1452            {
1453                let mut vals = Vec::new();
1454                if let ColumnReader::FixedLenByteArrayColumnReader(mut typed) = column_reader {
1455                    let mut definition_levels = Vec::new();
1456                    let (total_num, valid_num, decoded_num) = typed.read_records(
1457                        num_records, Some(&mut definition_levels), None, &mut vals)?;
1458                    if valid_num != decoded_num {
1459                        panic!("Support only valid records, found {} null records in column type {}",
1460                            decoded_num - valid_num, stringify!{unique_id});
1461                    }
1462                } else {
1463                    panic!("Schema and struct disagree on type for {}", stringify!{ unique_id });
1464                }
1465                for (i, r) in &mut records[..num_records].iter_mut().enumerate() {
1466                    r.unique_id = ::uuid::Uuid::from_bytes(vals[i].data().try_into().unwrap());
1467                }
1468            }
1469        }).to_string());
1470    }
1471
1472    #[test]
1473    fn test_converted_type() {
1474        let snippet: proc_macro2::TokenStream = quote! {
1475          struct ATimeStruct {
1476            time: chrono::NaiveDateTime,
1477          }
1478        };
1479
1480        let fields = extract_fields(snippet);
1481
1482        let time = Field::from(&fields[0]);
1483
1484        let converted_type = time.ty.converted_type();
1485        assert_eq!(
1486            converted_type.unwrap().to_string(),
1487            quote! { ::parquet::basic::ConvertedType::TIMESTAMP_MILLIS }.to_string()
1488        );
1489    }
1490}
parquet_derive/parquet_field.rs

parquet_derive/
parquet_field.rs
