parquet_derive/
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//! This crate provides a procedural macro to derive
19//! implementations of a RecordWriter and RecordReader
20
21#![warn(missing_docs)]
22#![recursion_limit = "128"]
23
24extern crate proc_macro;
25extern crate proc_macro2;
26extern crate syn;
27#[macro_use]
28extern crate quote;
29
30extern crate parquet;
31
32use ::syn::{parse_macro_input, Data, DataStruct, DeriveInput};
33
34mod parquet_field;
35
36/// Derive flat, simple RecordWriter implementations.
37///
38/// Works by parsing a struct tagged with `#[derive(ParquetRecordWriter)]` and emitting
39/// the correct writing code for each field of the struct. Column writers
40/// are generated in the order they are defined.
41///
42/// It is up to the programmer to keep the order of the struct
43/// fields lined up with the schema.
44///
45/// Example:
46///
47/// ```no_run
48/// use parquet_derive::ParquetRecordWriter;
49/// use std::io::{self, Write};
50/// use parquet::file::properties::WriterProperties;
51/// use parquet::file::writer::SerializedFileWriter;
52/// use parquet::record::RecordWriter;
53/// use std::fs::File;
54///
55/// use std::sync::Arc;
56///
57/// #[derive(ParquetRecordWriter)]
58/// struct ACompleteRecord<'a> {
59///   pub a_bool: bool,
60///   pub a_str: &'a str,
61/// }
62///
63/// pub fn write_some_records() {
64///   let samples = vec![
65///     ACompleteRecord {
66///       a_bool: true,
67///       a_str: "I'm true"
68///     },
69///     ACompleteRecord {
70///       a_bool: false,
71///       a_str: "I'm false"
72///     }
73///   ];
74///  let file = File::open("some_file.parquet").unwrap();
75///
76///  let schema = samples.as_slice().schema().unwrap();
77///
78///  let mut writer = SerializedFileWriter::new(file, schema, Default::default()).unwrap();
79///
80///  let mut row_group = writer.next_row_group().unwrap();
81///  samples.as_slice().write_to_row_group(&mut row_group).unwrap();
82///  row_group.close().unwrap();
83///  writer.close().unwrap();
84/// }
85/// ```
86///
87#[proc_macro_derive(ParquetRecordWriter)]
88pub fn parquet_record_writer(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
89    let input: DeriveInput = parse_macro_input!(input as DeriveInput);
90    let fields = match input.data {
91        Data::Struct(DataStruct { fields, .. }) => fields,
92        Data::Enum(_) => unimplemented!("Enum currently is not supported"),
93        Data::Union(_) => unimplemented!("Union currently is not supported"),
94    };
95
96    let field_infos: Vec<_> = fields.iter().map(parquet_field::Field::from).collect();
97
98    let writer_snippets: Vec<proc_macro2::TokenStream> =
99        field_infos.iter().map(|x| x.writer_snippet()).collect();
100
101    let derived_for = input.ident;
102    let generics = input.generics;
103
104    let field_types: Vec<proc_macro2::TokenStream> =
105        field_infos.iter().map(|x| x.parquet_type()).collect();
106
107    (quote! {
108    impl #generics ::parquet::record::RecordWriter<#derived_for #generics> for &[#derived_for #generics] {
109      fn write_to_row_group<W: ::std::io::Write + Send>(
110        &self,
111        row_group_writer: &mut ::parquet::file::writer::SerializedRowGroupWriter<'_, W>
112      ) -> Result<(), ::parquet::errors::ParquetError> {
113        use ::parquet::column::writer::ColumnWriter;
114
115        let mut row_group_writer = row_group_writer;
116        let records = &self; // Used by all the writer snippets to be more clear
117
118        #(
119          {
120              let mut some_column_writer = row_group_writer.next_column().unwrap();
121              if let Some(mut column_writer) = some_column_writer {
122                  #writer_snippets
123                  column_writer.close()?;
124              } else {
125                  return Err(::parquet::errors::ParquetError::General("Failed to get next column".into()))
126              }
127          }
128        );*
129
130        Ok(())
131      }
132
133      fn schema(&self) -> Result<::parquet::schema::types::TypePtr, ::parquet::errors::ParquetError> {
134        use ::parquet::schema::types::Type as ParquetType;
135        use ::parquet::schema::types::TypePtr;
136        use ::parquet::basic::LogicalType;
137
138        let mut fields: ::std::vec::Vec<TypePtr> = ::std::vec::Vec::new();
139        #(
140          #field_types
141        );*;
142        let group = ParquetType::group_type_builder("rust_schema")
143          .with_fields(fields)
144          .build()?;
145        Ok(group.into())
146      }
147    }
148  }).into()
149}
150
151/// Derive flat, simple RecordReader implementations.
152///
153/// Works by parsing a struct tagged with `#[derive(ParquetRecordReader)]` and emitting
154/// the correct writing code for each field of the struct. Column readers
155/// are generated by matching names in the schema to the names in the struct.
156///
157/// It is up to the programmer to ensure the names in the struct
158/// fields line up with the schema.
159///
160/// Example:
161///
162/// ```no_run
163/// use parquet::record::RecordReader;
164/// use parquet::file::{serialized_reader::SerializedFileReader, reader::FileReader};
165/// use parquet_derive::{ParquetRecordReader};
166/// use std::fs::File;
167///
168/// #[derive(ParquetRecordReader)]
169/// struct ACompleteRecord {
170///     pub a_bool: bool,
171///     pub a_string: String,
172/// }
173///
174/// pub fn read_some_records() -> Vec<ACompleteRecord> {
175///   let mut samples: Vec<ACompleteRecord> = Vec::new();
176///   let file = File::open("some_file.parquet").unwrap();
177///
178///   let reader = SerializedFileReader::new(file).unwrap();
179///   let mut row_group = reader.get_row_group(0).unwrap();
180///   samples.read_from_row_group(&mut *row_group, 1).unwrap();
181///   samples
182/// }
183/// ```
184///
185#[proc_macro_derive(ParquetRecordReader)]
186pub fn parquet_record_reader(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
187    let input: DeriveInput = parse_macro_input!(input as DeriveInput);
188    let fields = match input.data {
189        Data::Struct(DataStruct { fields, .. }) => fields,
190        Data::Enum(_) => unimplemented!("Enum currently is not supported"),
191        Data::Union(_) => unimplemented!("Union currently is not supported"),
192    };
193
194    let field_infos: Vec<_> = fields.iter().map(parquet_field::Field::from).collect();
195    let field_names: Vec<_> = fields.iter().map(|f| f.ident.clone()).collect();
196    let reader_snippets: Vec<proc_macro2::TokenStream> =
197        field_infos.iter().map(|x| x.reader_snippet()).collect();
198
199    let derived_for = input.ident;
200    let generics = input.generics;
201
202    (quote! {
203
204    impl #generics ::parquet::record::RecordReader<#derived_for #generics> for Vec<#derived_for #generics> {
205      fn read_from_row_group(
206        &mut self,
207        row_group_reader: &mut dyn ::parquet::file::reader::RowGroupReader,
208        num_records: usize,
209      ) -> Result<(), ::parquet::errors::ParquetError> {
210        use ::parquet::column::reader::ColumnReader;
211
212        let mut row_group_reader = row_group_reader;
213
214        // key: parquet file column name, value: column index
215        let mut name_to_index = std::collections::HashMap::new();
216        for (idx, col) in row_group_reader.metadata().schema_descr().columns().iter().enumerate() {
217            name_to_index.insert(col.name().to_string(), idx);
218        }
219
220        for _ in 0..num_records {
221          self.push(#derived_for {
222            #(
223              #field_names: Default::default()
224            ),*
225          })
226        }
227
228        let records = self; // Used by all the reader snippets to be more clear
229
230        #(
231          {
232              let idx: usize = match name_to_index.get(stringify!(#field_names)) {
233                Some(&col_idx) => col_idx,
234                None => {
235                  let error_msg = format!("column name '{}' is not found in parquet file!", stringify!(#field_names));
236                  return Err(::parquet::errors::ParquetError::General(error_msg));
237                }
238              };
239              if let Ok(mut column_reader) = row_group_reader.get_column_reader(idx) {
240                  #reader_snippets
241              } else {
242                  return Err(::parquet::errors::ParquetError::General("Failed to get next column".into()))
243              }
244          }
245        );*
246
247        Ok(())
248      }
249    }
250  }).into()
251}
parquet_derive/lib.rs

parquet_derive/
lib.rs
