#![recursion_limit = "128"]
extern crate proc_macro;
extern crate proc_macro2;
extern crate syn;
#[macro_use]
extern crate quote;
extern crate parquet;
use ::syn::{parse_macro_input, Data, DataStruct, DeriveInput};
mod parquet_field;
#[proc_macro_derive(ParquetRecordWriter)]
pub fn parquet_record_writer(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
let input: DeriveInput = parse_macro_input!(input as DeriveInput);
let fields = match input.data {
Data::Struct(DataStruct { fields, .. }) => fields,
Data::Enum(_) => unimplemented!("Enum currently is not supported"),
Data::Union(_) => unimplemented!("Union currently is not supported"),
};
let field_infos: Vec<_> = fields.iter().map(parquet_field::Field::from).collect();
let writer_snippets: Vec<proc_macro2::TokenStream> =
field_infos.iter().map(|x| x.writer_snippet()).collect();
let derived_for = input.ident;
let generics = input.generics;
let field_types: Vec<proc_macro2::TokenStream> =
field_infos.iter().map(|x| x.parquet_type()).collect();
(quote! {
impl #generics ::parquet::record::RecordWriter<#derived_for #generics> for &[#derived_for #generics] {
fn write_to_row_group<W: ::std::io::Write>(
&self,
row_group_writer: &mut ::parquet::file::writer::SerializedRowGroupWriter<'_, W>
) -> Result<(), ::parquet::errors::ParquetError> {
use ::parquet::column::writer::ColumnWriter;
let mut row_group_writer = row_group_writer;
let records = &self;
#(
{
let mut some_column_writer = row_group_writer.next_column().unwrap();
if let Some(mut column_writer) = some_column_writer {
#writer_snippets
column_writer.close()?;
} else {
return Err(::parquet::errors::ParquetError::General("Failed to get next column".into()))
}
}
);*
Ok(())
}
fn schema(&self) -> Result<::parquet::schema::types::TypePtr, ::parquet::errors::ParquetError> {
use ::parquet::schema::types::Type as ParquetType;
use ::parquet::schema::types::TypePtr;
use ::parquet::basic::LogicalType;
let mut fields: ::std::vec::Vec<TypePtr> = ::std::vec::Vec::new();
#(
#field_types
);*;
let group = ParquetType::group_type_builder("rust_schema")
.with_fields(&mut fields)
.build()?;
Ok(group.into())
}
}
}).into()
}