idhash 0.3.0

Calculate a Row-Invariant ID for Tabular Data
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118

use std::fs::File;

use std::ops::Deref;
use std::sync::Arc;

use arrow::array::Array;
use arrow::chunk::Chunk;
use arrow::datatypes::Schema;
use arrow::io::csv::read::{
    deserialize_batch, deserialize_column, infer, infer_schema, ByteRecord,
};
use arrow::io::csv::read::{Reader, ReaderBuilder};

pub struct CSVReader {
    reader: Reader<File>,
    buffer: Vec<ByteRecord>,
    pub batch_size: usize,
    pub schema: Arc<Schema>,
    exhausted: bool,
    line_number: usize,
}

impl CSVReader {
    pub fn new(file_path: String, lines_for_type_inference: usize, batch_size: usize) -> Self {
        let mut reader = ReaderBuilder::new()
            .has_headers(true)
            .from_path(file_path)
            .unwrap();

        let schema = Arc::new(Schema::from(
            infer_schema(&mut reader, Some(lines_for_type_inference), true, &infer)
                .unwrap()
                .0,
        ));

        CSVReader {
            reader,
            buffer: vec![ByteRecord::default(); batch_size],
            batch_size,
            line_number: 0,
            schema,
            exhausted: false,
        }
    }
}

pub struct ThreadArrayChunk {
    array: Arc<Box<dyn Array>>,
}

impl Deref for ThreadArrayChunk {
    type Target = Arc<Box<dyn Array>>;

    fn deref(&self) -> &Self::Target {
        &self.array
    }
}

impl<'a> AsRef<dyn Array + 'a> for ThreadArrayChunk {
    fn as_ref(&self) -> &(dyn Array + 'a) {
        &**self.array
    }
}

impl From<Arc<Box<dyn Array>>> for ThreadArrayChunk {
    fn from(item: Arc<Box<dyn Array>>) -> Self {
        ThreadArrayChunk { array: item }
    }
}

impl From<Arc<Arc<dyn Array>>> for ThreadArrayChunk {
    fn from(item: Arc<Arc<dyn Array>>) -> Self {
        let boxed_item: Box<dyn Array> = item.to_boxed();
        ThreadArrayChunk {
            array: Arc::new(boxed_item),
        }
    }
}

impl Iterator for CSVReader {
    type Item = Chunk<ThreadArrayChunk>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.exhausted {
            return None;
        }
        let mut row_count = 0;
        for slot in self.buffer.iter_mut() {
            match self.reader.read_byte_record(slot) {
                Ok(has_more) => {
                    if !has_more {
                        self.exhausted = true;
                        break;
                    }
                    row_count += 1;
                }
                Err(_) => break,
            }
        }
        self.line_number += row_count;
        match deserialize_batch(
            &self.buffer[..row_count],
            &self.schema.fields,
            None,
            self.line_number - row_count,
            deserialize_column,
        ) {
            Ok(chunk) => Some(Chunk::new(
                chunk
                    .into_arrays()
                    .into_iter()
                    .map(|x| ThreadArrayChunk::from(Arc::new(x)))
                    .collect::<Vec<ThreadArrayChunk>>(),
            )),
            Err(_) => panic!("Failed to Deserialize Batch"),
        }
    }
}