//! This module defines helper methods to interact with a DataFrame
//! A DataFrame is a columnar representation of a SOR file and is
//! represented as a `Vec<Column>`
use crate::parsers::parse_line_with_schema;
use crate::schema::{infer_schema, DataType};
use std::fmt;
use std::fs::File;
use std::io::{BufRead, BufReader, Seek, SeekFrom};
use std::thread;
/// Represents a column in the DataFrame
#[derive(PartialEq, Clone, Debug)]
pub enum Column {
/// A Column consisting of either Ints or missing
Int(Vec<Option<i64>>),
/// A Column consisting of either Bool or missing
Bool(Vec<Option<bool>>),
/// A Column consisting of either Float or missing
Float(Vec<Option<f64>>),
/// A Column consisting of either String or missing
String(Vec<Option<String>>),
}
/// An enumeration of the possible `SoR` data types, that also contains the
/// data itself.
#[derive(PartialEq, Debug, Clone)]
pub enum Data {
/// A String cell
String(String),
/// A Int cell
Int(i64),
/// A Float Cell
Float(f64),
/// A Boolean Cell
Bool(bool),
/// A Missing Value
Null,
}
/// Represents a DataFrame which contains
/// [columnar](::crate::dataframe::Column) data and a
/// [Schema](::crate::parsers::Schema).
pub struct DataFrame {
/// The [Schema](::crate::parsers::Schema) of this DataFrame
pub schema: Vec<DataType>,
/// The [columnar](::crate::dataframe::Column) data of this DataFrame.
pub data: Vec<Column>,
}
/// Print the Data of a Data cell.
/// The number for Ints and floats
/// 0 for false
/// 1 for trues
/// a quotes string for Strings
/// and "Missing Value" for missing data
impl fmt::Display for Data {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self {
Data::String(s) => write!(f, "\"{}\"", s),
Data::Int(n) => write!(f, "{}", n),
Data::Float(fl) => write!(f, "{}", fl),
Data::Bool(true) => write!(f, "1"),
Data::Bool(false) => write!(f, "0"),
Data::Null => write!(f, "Missing Value"),
}
}
}
/// Implementations for a [crate::DataFrame::DataFrame]
impl DataFrame {
/// Generate an empty DataFrame for the given schema
pub fn init(schema: &Vec<DataType>) -> Self {
let mut result = Vec::with_capacity(schema.len() + 1);
for t in schema {
match t {
DataType::Bool => result.push(Column::Bool(Vec::new())),
DataType::Int => result.push(Column::Int(Vec::new())),
DataType::Float => result.push(Column::Float(Vec::new())),
DataType::String => result.push(Column::String(Vec::new())),
}
}
DataFrame {
data: result,
schema: schema.clone(),
}
}
/// Reads a file (even one too large to fit into memory) according to the given
/// `schema` and `options` and turns it into a columnar dataframe.
///
/// This is the top level function for using `SoRer` and the one you should be
/// using unless you are trying to extend `SoRer`. There are many intricate
/// facets to using `SoRer` so you *must* RTFM [here](../index.html)
pub fn from_file(file_path: String, from: u64, len: u64) -> Self {
// infer the schema
let f: File = File::open(file_path.clone()).unwrap();
let reader = BufReader::new(f);
let schema = infer_schema(reader);
// number of threads to use
let num_threads = 8;
// the total number of bytes to read
let num_chars = if len == std::u64::MAX {
(std::fs::metadata(file_path.clone()).unwrap().len() - from) as f64
} else {
len as f64
};
// each thread will parse this many characters +- some number
let step = (num_chars / num_threads as f64).ceil() as u64;
// setup the work array with the from / len for each thread
// each element in the work array is a tuple of (starting index, number of byte for this thread)
let f: File = File::open(file_path.clone()).unwrap();
let mut reader = BufReader::new(f);
let mut work: Vec<(u64, u64)> = Vec::with_capacity(num_threads + 1);
// add the first one separately since we want to access the previous thread's
// work when in the loop. Since the work of the first thread will call
// `read_file(schema, 0, step)` it will not throw away the first line
// since from is 0 and will throw away the last line since step > 0
work.push((from, step));
let mut so_far = from;
let mut buffer = Vec::new();
// This loop finds the byte offset for the start of a line
// by adding the length of the last line that a previous thread would've
// thrown away. The work gets added to the following thread so that
// each thread starts at a full line and reads only until the end of a line
for i in 1..num_threads {
so_far += step;
// advance the reader to this threads starting index then
// find the next newline character
reader.seek(SeekFrom::Start(so_far)).unwrap();
reader.read_until(b'\n', &mut buffer).unwrap();
work.push((so_far, step));
// Since the previous thread throws away the last line, add the length
// of the last line of prev thread to the work of this thread so that
// we read all lines.
work.get_mut(i - 1).unwrap().1 += buffer.len() as u64 + 1;
buffer.clear();
}
// initialize the threads with their own BufReader
let mut threads = Vec::new();
for w in work {
let new_schema = schema.clone();
let f: File = File::open(file_path.clone()).unwrap();
let mut r = BufReader::new(f);
// spawn the thread and give it a closure which calls `from_file`
// to parse the data into columnar format.
threads.push(thread::spawn(move || {
read_chunk(new_schema, &mut r, w.0, w.1)
}));
}
// initialize the resulting columnar data frame
let mut parsed_data: Vec<Column> = DataFrame::init(&schema).data;
// let all the threads finish then combine the parsed data into the
// columnar data frame
for t in threads {
let mut x: Vec<Column> = t.join().unwrap();
let iter = parsed_data.iter_mut().zip(x.iter_mut());
for (complete, partial) in iter {
match (complete, partial) {
(Column::Bool(c1), Column::Bool(c2)) => c1.append(c2),
(Column::Int(c1), Column::Int(c2)) => c1.append(c2),
(Column::Float(c1), Column::Float(c2)) => c1.append(c2),
(Column::String(c1), Column::String(c2)) => c1.append(c2),
_ => panic!("Unexpected result from thread"),
}
}
}
DataFrame {
data: parsed_data,
schema,
}
}
/// Get the (i,j) element from the DataFrame
pub fn get(&self, i: u64, j: u64) -> Data {
match &self.data[i as usize] {
Column::Bool(b) => {
if let Some(val) = &b[j as usize] {
Data::Bool(*val)
} else {
Data::Null
}
}
Column::Int(b) => {
if let Some(val) = &b[j as usize] {
Data::Int(*val)
} else {
Data::Null
}
}
Column::Float(b) => {
if let Some(val) = &b[j as usize] {
Data::Float(*val)
} else {
Data::Null
}
}
Column::String(b) => {
if let Some(val) = &b[j as usize] {
Data::String(val.clone())
} else {
Data::Null
}
}
}
}
}
fn read_chunk<T>(
schema: Vec<DataType>,
reader: &mut T,
from: u64,
len: u64,
) -> Vec<Column>
where
T: BufRead + Seek,
{
reader.seek(SeekFrom::Start(from)).unwrap();
let mut buffer = Vec::new();
let mut so_far = if from != 0 {
// throw away the first line
let l1_len = reader.read_until(b'\n', &mut buffer).unwrap();
buffer.clear();
l1_len as u64
} else {
0
};
let mut parsed_data = DataFrame::init(&schema).data;
loop {
let line_len = reader.read_until(b'\n', &mut buffer).unwrap();
so_far += line_len as u64;
if line_len == 0 || so_far >= len {
break;
}
// parse line with schema and place into the columnar vec here
match parse_line_with_schema(&buffer[..], &schema) {
None => {
buffer.clear();
continue;
}
Some(data) => {
let iter = data.iter().zip(parsed_data.iter_mut());
for (d, col) in iter {
match (d, col) {
(Data::Bool(b), Column::Bool(c)) => c.push(Some(*b)),
(Data::Int(i), Column::Int(c)) => c.push(Some(*i)),
(Data::Float(f), Column::Float(c)) => c.push(Some(*f)),
(Data::String(s), Column::String(c)) => {
c.push(Some(s.clone()))
}
(Data::Null, Column::Bool(c)) => c.push(None),
(Data::Null, Column::Int(c)) => c.push(None),
(Data::Null, Column::Float(c)) => c.push(None),
(Data::Null, Column::String(c)) => c.push(None),
_ => panic!("Parser Failed"),
}
}
}
}
buffer.clear();
}
parsed_data
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Cursor;
#[test]
fn test_read_file() {
let schema = vec![DataType::String, DataType::Bool];
let expected_col1 = Column::String(vec![
Some("1".to_string()),
Some("a".to_string()),
Some("1.2".to_string()),
]);
let expected_col2 = Column::Bool(vec![Some(true), Some(false), None]);
let expected = vec![expected_col1, expected_col2];
// Simple case : first nd last line are not discarded
let mut input = Cursor::new(b"<1><1>\n<a><0>\n<1.2><>");
let parsed1: Vec<Column> =
read_chunk(schema.clone(), &mut input, 0, 26);
assert_eq!(parsed1, expected.clone());
// last line is discarded
let mut larger_input = Cursor::new(b"<1><1>\n<a><0>\n<1.2><>\n<no><1>");
let parsed2: Vec<Column> =
read_chunk(schema.clone(), &mut larger_input, 0, 27);
assert_eq!(parsed2, expected.clone());
// first line is discarded
let mut input_skipped_l1 =
Cursor::new(b"<b><1>\n<1><1>\n<a><0>\n<1.2><>");
let parsed3: Vec<Column> =
read_chunk(schema.clone(), &mut input_skipped_l1, 3, 26);
assert_eq!(parsed3, expected.clone());
// Invalid line is discarded
// Note since parsed lines with schema is correctly tested we do not
// need to test every possible way a line can be invalid here
let mut input_with_invalid =
Cursor::new(b"<1><1>\n<a><0>\n<c><1.2>\n<1.2><>");
let parsed4: Vec<Column> =
read_chunk(schema.clone(), &mut input_with_invalid, 0, 32);
assert_eq!(parsed4, expected.clone());
}
}