read csv to parquet, no nuls

file1.csv

main.rs

extern crate bson;
extern crate parquet;
extern crate csv;
extern crate regex;
#[macro_use] extern crate lazy_static;

use std::any::{Any, TypeId};
use std::error::Error;
use std::fs::File;
use std::io;
use std::path::Path;
use std::process;
use std::rc::Rc;

use parquet::column::reader::ColumnReader;
use parquet::column::writer::ColumnWriter;
use parquet::file::properties::WriterProperties;
use parquet::file::reader::{FileReader, SerializedFileReader};
use parquet::file::writer::{FileWriter, SerializedFileWriter};
use parquet::schema::parser::parse_message_type;
use regex::Regex;

fn build_parquet_schema(messages: Vec<(String, String)>) -> String {
    let mut builder = String::from("message schema {");
    messages.iter().for_each(|(a, b)| {
        builder.push_str(&format!("REQUIRED {} {};", a, b));
    });
    builder.push_str("}");

    builder
}

// regex patterns
lazy_static! {
    static ref DECIMAL_RE: Regex = Regex::new(r"^(\d*.).\d+$").unwrap();
    static ref INTEGER_RE: Regex = Regex::new(r"^(\d*.)$").unwrap();
    static ref BOOLEAN_RE: Regex = Regex::new(r"^(true)$|^(false)$").unwrap();
    static ref UTC_TIMESTAMP_RE: Regex = Regex::new(r"^(-?(?:[1-9][0-9]*)?[0-9]{4})-(1[0-2]|0[1-9])-(3[01]|0[1-9]|[12][0-9])T(2[0-3]|[01][0-9]):([0-5][0-9]):([0-5][0-9])(\\.[0-9]+)?(Z)?$").unwrap();
}

fn write_parquet(schema: String, columns: Vec<Vec<Option<String>>>) -> Result<(), Box<Error>> {

    let schema = Rc::new(parse_message_type(&schema)?);
    let props = Rc::new(WriterProperties::builder().build());
    let file = File::create("./data/file1.parquet")?;
    let mut writer = SerializedFileWriter::new(file, schema, props)?;
    let mut row_group_writer = writer.next_row_group()?;
    let mut col_number = 0;
    while let Some(mut col_writer) = row_group_writer.next_column()? {
        let column = columns.get(col_number).unwrap();
        col_number = col_number + 1;
        match col_writer {
            ColumnWriter::ByteArrayColumnWriter(ref mut typed_writer) => {
                println!("writing a byte array");
                // I can remove this if-else when I start taking fn parameters of my schema and columns
                let mut vals = vec![];
                column.iter().for_each(|v| vals.push(parquet::data_type::ByteArray::from(v.clone().unwrap().as_str())));
                typed_writer.write_batch(
                    &vals, None, None
                )?;
            },
            ColumnWriter::Int32ColumnWriter(ref mut typed_writer) => {
                println!("writing an integer");
                let mut vals: Vec<i32> = vec![];
                column.iter().for_each(|mut v| {
                    let string = v.clone().unwrap_or(String::from("0"));
                    vals.push(string.parse::<i32>().unwrap())
                });
                typed_writer.write_batch(
                    &vals, None, None
                )?;
            },
            _ => {}
        }
        row_group_writer.close_column(col_writer)?;
    }
    writer.close_row_group(row_group_writer)?;
    writer.close()?;
    Ok(())
}

fn read_csv() -> Result<(), Box<Error>> {
    let file = File::open("./data/file1.csv")?;
    let mut rdr = csv::Reader::from_reader(file);
    let headers = rdr.headers()?;
    // println!("headers {:?}", headers);

    let headers_clone = headers.clone();
    let my_headers: Vec<&str> = headers_clone.iter().collect();
    // println!("header length {:?}", my_headers);
    let header_length = my_headers.len();
    let mut my_vals: Vec<Vec<Option<String>>> = vec![vec![]; header_length];

    // println!("header length {:?}", header_length);
    // println!("values {:?}", my_vals);

    // create a vector that will store the outputs of the CSV
    let mut current_record = 0;
    for result in rdr.records() {
        let record = result?;
        // println!("{:?}", record);
        for i in 0..header_length {
            let borrowed = record.clone();
            my_vals[i].push(borrowed.get(i).map(|s| s.to_string()));
        }
        current_record = current_record + 1;
    }

    // println!("values after creating {:?}", my_vals);

    // create schema
    let cloned_vals = my_vals.clone();
    let my_fields = cloned_vals.iter().map(|v| {
        // println!("pre-input schema {:?}", v);
        infer_field_schema(v)
    });

    let mut schema = String::from("message schema {");
    my_fields.zip(&my_headers).for_each(|((t, is_required), header_name)| {
        // println!("record {} {} {}", t, is_required, header_name);
        match is_required {
            true => {
                match t.as_str() {
                    "STRING" => {
                        schema.push_str("REQUIRED BYTE_ARRAY ");
                        schema.push_str(header_name);
                        schema.push_str(" (UTF8); ");
                    },
                    "BOOLEAN" => {
                        schema.push_str("REQUIRED BOOLEAN ");
                        schema.push_str(header_name);
                        schema.push_str("; ");
                    },
                    "TIMESTAMP" => {
                        schema.push_str("REQUIRED TIMESTAMP ");
                        schema.push_str(header_name);
                        schema.push_str("; ");
                    },
                    "INT32" => {
                        schema.push_str("REQUIRED INT32 ");
                        schema.push_str(header_name);
                        schema.push_str("; ");
                    },
                    "DOUBLE" => {
                        schema.push_str("REQUIRED DOUBLE ");
                        schema.push_str(header_name);
                        schema.push_str("; ");
                    },
                    _ => {
                        schema.push_str("REQUIRED BYTE_ARRAY ");
                        schema.push_str(header_name);
                        schema.push_str(" (UTF8); ");
                    }
                }
            },
            false => {
                match t.as_str() {
                    "STRING" => {
                        schema.push_str("OPTIONAL BYTE_ARRAY ");
                        schema.push_str(header_name);
                        schema.push_str(" (UTF8); ");
                    },
                    "BOOLEAN" => {
                        schema.push_str("OPTIONAL BOOLEAN ");
                        schema.push_str(header_name);
                        schema.push_str("; ");
                    },
                    "TIMESTAMP" => {
                        schema.push_str("OPTIONAL TIMESTAMP ");
                        schema.push_str(header_name);
                        schema.push_str("; ");
                    },
                    "INT32" => {
                        schema.push_str("OPTIONAL INT32 ");
                        schema.push_str(header_name);
                        schema.push_str("; ");
                    },
                    "DOUBLE" => {
                        schema.push_str("OPTIONAL DOUBLE ");
                        schema.push_str(header_name);
                        schema.push_str("; ");
                    },
                    _ => {
                        schema.push_str("OPTIONAL BYTE_ARRAY ");
                        schema.push_str(header_name);
                        schema.push_str(" (UTF8); ");
                    }
                }
            }
        };
        ()
    });
    schema.push_str("}");
    println!("schema {:?}", schema);

    // try create a parquet file from csv
    println!("{:?}", write_parquet(schema, my_vals));

    Ok(())
}

/// Try infer schema based on the number of fields that match a certain format.
/// 
/// Fall back to a string if schema can't be deferred. 
/// 
/// Mappings:
///     0: empty
///     1: BOOL
///     2: TIMESTAMP
///     3: DOUBLE
///     4: INT
///     5: STRING
fn infer_field_schema(vector: &Vec<Option<String>>) -> (String, bool) {
    // take the first n records, and infer schema from them
    let mut possibilities: Vec<i8> = vec![];
    for x in vector.iter().take(100) {
        match x {
            Some(v) => {
                // match regex in a particular order
                if BOOLEAN_RE.is_match(v) {
                    possibilities.push(1);
                } else if UTC_TIMESTAMP_RE.is_match(v) {
                    possibilities.push(2);
                } else if DECIMAL_RE.is_match(v) {
                    possibilities.push(3);
                } else if INTEGER_RE.is_match(v) {
                    possibilities.push(4);
                } else {
                    // we have a string, so the schema becomes a string
                    // possibilities.push(5);
                    return (String::from("STRING"), true)
                }
            },
            None => {
                // do nothing, we'll return an optional String
                possibilities.push(0);
            }
        }
    }
    possibilities.sort();
    possibilities.dedup();
    if possibilities.len() == 1 {
        match possibilities[0] {
            0 => { return (String::from("STRING"), true) },
            1 => { return (String::from("BOOLEAN"), true) },
            2 => { return (String::from("TIMESTAMP"), true) },
            3 => { return (String::from("DOUBLE"), true) },
            4 => { return (String::from("INT32"), true) },
            _ => { return (String::from("STRING"), true) },
        }
    } else if possibilities.len() > 2 {
        return (String::from("STRING"), false)
    } else {
        if possibilities.contains(&1) || possibilities.contains(&2) {
            // boolean or timestamp is incompatible with other types, return string
            return (String::from("STRING"), false)
        } else if possibilities.contains(&3) && possibilities.contains(&4) {
            // we have an integer and double, fall down to double
            return (String::from("DOUBLE"), false)
        } else {
            return (String::from("STRING"), false)
        }
    }
}

fn main() {
    let _ = read_csv();
}