connectorx/sources/csv/
mod.rs

1//! Source implementation for CSV files.
2
3mod errors;
4mod typesystem;
5
6pub use self::errors::CSVSourceError;
7pub use self::typesystem::CSVTypeSystem;
8use super::{PartitionParser, Produce, Source, SourcePartition};
9use crate::{data_order::DataOrder, errors::ConnectorXError, sql::CXQuery};
10use anyhow::anyhow;
11use chrono::{DateTime, Utc};
12use fehler::{throw, throws};
13#[cfg(feature = "src_csv")]
14use regex::{Regex, RegexBuilder};
15use std::collections::HashSet;
16use std::fs::File;
17
18pub struct CSVSource {
19    schema: Vec<CSVTypeSystem>,
20    files: Vec<CXQuery<String>>,
21    names: Vec<String>,
22}
23
24impl CSVSource {
25    pub fn new(schema: &[CSVTypeSystem]) -> Self {
26        CSVSource {
27            schema: schema.to_vec(),
28            files: vec![],
29            names: vec![],
30        }
31    }
32
33    #[throws(CSVSourceError)]
34    pub fn infer_schema(&mut self) -> Vec<CSVTypeSystem> {
35        // regular expressions for infer CSVTypeSystem from string
36        let decimal_re: Regex = Regex::new(r"^-?(\d+\.\d+)$")?;
37        let integer_re: Regex = Regex::new(r"^-?(\d+)$")?;
38        let boolean_re: Regex = RegexBuilder::new(r"^(true)$|^(false)$")
39            .case_insensitive(true)
40            .build()?;
41        let datetime_re: Regex = Regex::new(r"^\d{4}-\d\d-\d\dT\d\d:\d\d:\d\d$")?;
42
43        // read max_records rows to infer possible CSVTypeSystems for each field
44        let mut reader = csv::ReaderBuilder::new()
45            .has_headers(true)
46            .from_reader(File::open(self.files[0].as_str())?);
47
48        let max_records_to_read = 50;
49        let num_cols = self.names.len();
50
51        let mut column_types: Vec<HashSet<CSVTypeSystem>> = vec![HashSet::new(); num_cols];
52        let mut nulls: Vec<bool> = vec![false; num_cols];
53
54        let mut record = csv::StringRecord::new();
55
56        for _record_counter in 0..max_records_to_read {
57            if !reader.read_record(&mut record)? {
58                break;
59            }
60            for field_counter in 0..num_cols {
61                if let Some(string) = record.get(field_counter) {
62                    if string.is_empty() {
63                        nulls[field_counter] = true;
64                    } else {
65                        let dt: CSVTypeSystem = if string.starts_with('"') {
66                            CSVTypeSystem::String(false)
67                        } else if boolean_re.is_match(string) {
68                            CSVTypeSystem::Bool(false)
69                        } else if decimal_re.is_match(string) {
70                            CSVTypeSystem::F64(false)
71                        } else if integer_re.is_match(string) {
72                            CSVTypeSystem::I64(false)
73                        } else if datetime_re.is_match(string) {
74                            CSVTypeSystem::DateTime(false)
75                        } else {
76                            CSVTypeSystem::String(false)
77                        };
78                        column_types[field_counter].insert(dt);
79                    }
80                }
81            }
82        }
83
84        // determine CSVTypeSystem based on possible candidates
85        let mut schema = vec![];
86
87        for field_counter in 0..num_cols {
88            let possibilities = &column_types[field_counter];
89            let has_nulls = nulls[field_counter];
90
91            match possibilities.len() {
92                1 => {
93                    for dt in possibilities.iter() {
94                        match *dt {
95                            CSVTypeSystem::I64(false) => {
96                                schema.push(CSVTypeSystem::I64(has_nulls));
97                            }
98                            CSVTypeSystem::F64(false) => {
99                                schema.push(CSVTypeSystem::F64(has_nulls));
100                            }
101                            CSVTypeSystem::Bool(false) => {
102                                schema.push(CSVTypeSystem::Bool(has_nulls));
103                            }
104                            CSVTypeSystem::String(false) => {
105                                schema.push(CSVTypeSystem::String(has_nulls));
106                            }
107                            CSVTypeSystem::DateTime(false) => {
108                                schema.push(CSVTypeSystem::DateTime(has_nulls));
109                            }
110                            _ => {}
111                        }
112                    }
113                }
114                2 if possibilities.contains(&CSVTypeSystem::I64(false))
115                    && possibilities.contains(&CSVTypeSystem::F64(false)) =>
116                {
117                    // Integer && Float -> Float
118                    schema.push(CSVTypeSystem::F64(has_nulls));
119                }
120                _ => {
121                    // Conflicting CSVTypeSystems -> String
122                    schema.push(CSVTypeSystem::String(has_nulls));
123                }
124            }
125        }
126        schema
127    }
128}
129
130impl Source for CSVSource {
131    const DATA_ORDERS: &'static [DataOrder] = &[DataOrder::RowMajor];
132    type Partition = CSVSourcePartition;
133    type TypeSystem = CSVTypeSystem;
134    type Error = CSVSourceError;
135
136    #[throws(CSVSourceError)]
137    fn set_data_order(&mut self, data_order: DataOrder) {
138        if !matches!(data_order, DataOrder::RowMajor) {
139            throw!(ConnectorXError::UnsupportedDataOrder(data_order))
140        }
141    }
142
143    fn set_queries<Q: ToString>(&mut self, queries: &[CXQuery<Q>]) {
144        self.files = queries.iter().map(|q| q.map(Q::to_string)).collect();
145    }
146
147    fn set_origin_query(&mut self, _query: Option<String>) {}
148
149    #[throws(CSVSourceError)]
150    fn fetch_metadata(&mut self) {
151        let mut reader = csv::ReaderBuilder::new()
152            .has_headers(true)
153            .from_reader(File::open(self.files[0].as_str())?);
154        let header = reader.headers()?;
155
156        self.names = header.iter().map(|s| s.to_string()).collect();
157
158        if self.schema.is_empty() {
159            self.schema = self.infer_schema()?;
160        }
161
162        assert_eq!(header.len(), self.schema.len());
163    }
164
165    #[throws(CSVSourceError)]
166    fn result_rows(&mut self) -> Option<usize> {
167        None
168    }
169
170    fn names(&self) -> Vec<String> {
171        self.names.clone()
172    }
173
174    fn schema(&self) -> Vec<Self::TypeSystem> {
175        self.schema.clone()
176    }
177
178    #[throws(CSVSourceError)]
179    fn partition(self) -> Vec<Self::Partition> {
180        let mut partitions = vec![];
181        for file in self.files {
182            partitions.push(CSVSourcePartition::new(file)?);
183        }
184        partitions
185    }
186}
187
188pub struct CSVSourcePartition {
189    records: Vec<csv::StringRecord>,
190    counter: usize,
191    nrows: usize,
192    ncols: usize,
193}
194
195impl CSVSourcePartition {
196    #[throws(CSVSourceError)]
197    pub fn new(fname: CXQuery<String>) -> Self {
198        let reader = csv::ReaderBuilder::new()
199            .has_headers(true)
200            .from_reader(File::open(fname.as_str())?);
201        let mut records = vec![];
202        reader
203            .into_records()
204            .try_for_each(|v| -> Result<(), CSVSourceError> {
205                records.push(v.map_err(|e| anyhow!(e))?);
206                Ok(())
207            })?;
208
209        let nrows = records.len();
210        let ncols = if nrows > 0 { records[0].len() } else { 0 };
211
212        Self {
213            records,
214            counter: 0,
215            nrows,
216            ncols,
217        }
218    }
219}
220
221impl SourcePartition for CSVSourcePartition {
222    type TypeSystem = CSVTypeSystem;
223    type Parser<'a> = CSVSourcePartitionParser<'a>;
224    type Error = CSVSourceError;
225
226    /// The parameter `query` is the path of the csv file
227    #[throws(CSVSourceError)]
228    fn result_rows(&mut self) {}
229
230    fn nrows(&self) -> usize {
231        self.nrows
232    }
233
234    fn ncols(&self) -> usize {
235        self.ncols
236    }
237
238    #[throws(CSVSourceError)]
239    fn parser(&mut self) -> Self::Parser<'_> {
240        CSVSourcePartitionParser {
241            records: &mut self.records,
242            counter: &mut self.counter,
243            ncols: self.ncols,
244        }
245    }
246}
247
248pub struct CSVSourcePartitionParser<'a> {
249    records: &'a mut [csv::StringRecord],
250    counter: &'a mut usize,
251    ncols: usize,
252}
253
254impl<'a> CSVSourcePartitionParser<'a> {
255    fn next_val(&mut self) -> &str {
256        let v: &str = self.records[*self.counter / self.ncols][*self.counter % self.ncols].as_ref();
257        *self.counter += 1;
258
259        v
260    }
261}
262
263impl<'a> PartitionParser<'a> for CSVSourcePartitionParser<'a> {
264    type TypeSystem = CSVTypeSystem;
265    type Error = CSVSourceError;
266
267    #[throws(CSVSourceError)]
268    fn fetch_next(&mut self) -> (usize, bool) {
269        (self.records.len(), true)
270    }
271}
272
273impl<'r, 'a> Produce<'r, i64> for CSVSourcePartitionParser<'a> {
274    type Error = CSVSourceError;
275
276    #[throws(CSVSourceError)]
277    fn produce(&mut self) -> i64 {
278        let v = self.next_val();
279        v.parse()
280            .map_err(|_| ConnectorXError::cannot_produce::<i64>(Some(v.into())))?
281    }
282}
283
284impl<'r, 'a> Produce<'r, Option<i64>> for CSVSourcePartitionParser<'a> {
285    type Error = CSVSourceError;
286
287    #[throws(CSVSourceError)]
288    fn produce(&mut self) -> Option<i64> {
289        let v = self.next_val();
290        if v.is_empty() {
291            return None;
292        }
293        let v = v
294            .parse()
295            .map_err(|_| ConnectorXError::cannot_produce::<Option<i64>>(Some(v.into())))?;
296
297        Some(v)
298    }
299}
300
301impl<'r, 'a> Produce<'r, f64> for CSVSourcePartitionParser<'a> {
302    type Error = CSVSourceError;
303
304    #[throws(CSVSourceError)]
305    fn produce(&mut self) -> f64 {
306        let v = self.next_val();
307        v.parse()
308            .map_err(|_| ConnectorXError::cannot_produce::<f64>(Some(v.into())))?
309    }
310}
311
312impl<'r, 'a> Produce<'r, Option<f64>> for CSVSourcePartitionParser<'a> {
313    type Error = CSVSourceError;
314
315    #[throws(CSVSourceError)]
316    fn produce(&mut self) -> Option<f64> {
317        let v = self.next_val();
318        if v.is_empty() {
319            return None;
320        }
321        let v = v
322            .parse()
323            .map_err(|_| ConnectorXError::cannot_produce::<Option<f64>>(Some(v.into())))?;
324
325        Some(v)
326    }
327}
328
329impl<'r, 'a> Produce<'r, bool> for CSVSourcePartitionParser<'a> {
330    type Error = CSVSourceError;
331
332    #[throws(CSVSourceError)]
333    fn produce(&mut self) -> bool {
334        let v = self.next_val();
335        v.parse()
336            .map_err(|_| ConnectorXError::cannot_produce::<bool>(Some(v.into())))?
337    }
338}
339
340impl<'r, 'a> Produce<'r, Option<bool>> for CSVSourcePartitionParser<'a> {
341    type Error = CSVSourceError;
342
343    #[throws(CSVSourceError)]
344    fn produce(&mut self) -> Option<bool> {
345        let v = self.next_val();
346        if v.is_empty() {
347            return None;
348        }
349        let v = v
350            .parse()
351            .map_err(|_| ConnectorXError::cannot_produce::<Option<bool>>(Some(v.into())))?;
352
353        Some(v)
354    }
355}
356
357impl<'r, 'a> Produce<'r, String> for CSVSourcePartitionParser<'a> {
358    type Error = CSVSourceError;
359
360    #[throws(CSVSourceError)]
361    fn produce(&mut self) -> String {
362        let v = self.next_val();
363        String::from(v)
364    }
365}
366
367impl<'a, 'r> Produce<'r, Option<String>> for CSVSourcePartitionParser<'a> {
368    type Error = CSVSourceError;
369
370    #[throws(CSVSourceError)]
371    fn produce(&'r mut self) -> Option<String> {
372        let v = self.next_val();
373
374        Some(String::from(v))
375    }
376}
377
378impl<'r, 'a> Produce<'r, DateTime<Utc>> for CSVSourcePartitionParser<'a> {
379    type Error = CSVSourceError;
380
381    #[throws(CSVSourceError)]
382    fn produce(&mut self) -> DateTime<Utc> {
383        let v = self.next_val();
384        v.parse()
385            .map_err(|_| ConnectorXError::cannot_produce::<DateTime<Utc>>(Some(v.into())))?
386    }
387}
388
389impl<'r, 'a> Produce<'r, Option<DateTime<Utc>>> for CSVSourcePartitionParser<'a> {
390    type Error = CSVSourceError;
391
392    #[throws(CSVSourceError)]
393    fn produce(&mut self) -> Option<DateTime<Utc>> {
394        let v = self.next_val();
395        if v.is_empty() {
396            return None;
397        }
398        let v = v
399            .parse()
400            .map_err(|_| ConnectorXError::cannot_produce::<DateTime<Utc>>(Some(v.into())))?;
401        Some(v)
402    }
403}
404
405#[cfg(test)]
406mod tests {
407    use super::{CSVSource, CSVTypeSystem};
408    use crate::{sources::Source, sql::CXQuery};
409    use std::{fs, path::PathBuf};
410
411    #[test]
412    fn infers_schema_from_numeric_boolean_datetime_and_string_values() {
413        let path: PathBuf =
414            std::env::temp_dir().join(format!("connectorx-csv-infer-{}.csv", std::process::id()));
415        fs::write(
416            &path,
417            "integer,float,mixed,boolean,datetime,string,empty\n\
418             1,1.5,1,true,2020-01-02T03:04:05,hello,\n\
419             ,2.5,2.5,,2020-01-02T03:04:05,world,\n",
420        )
421        .unwrap();
422
423        let mut source = CSVSource::new(&[]);
424        source.set_queries(&[CXQuery::naked(path.to_string_lossy().into_owned())]);
425        source.fetch_metadata().unwrap();
426
427        assert_eq!(
428            source.schema(),
429            vec![
430                CSVTypeSystem::I64(true),
431                CSVTypeSystem::F64(false),
432                CSVTypeSystem::F64(false),
433                CSVTypeSystem::Bool(true),
434                CSVTypeSystem::DateTime(false),
435                CSVTypeSystem::String(false),
436                CSVTypeSystem::String(true),
437            ]
438        );
439        assert_eq!(
440            source.names(),
441            vec![
442                "integer".to_owned(),
443                "float".to_owned(),
444                "mixed".to_owned(),
445                "boolean".to_owned(),
446                "datetime".to_owned(),
447                "string".to_owned(),
448                "empty".to_owned(),
449            ]
450        );
451
452        fs::remove_file(path).unwrap();
453    }
454}
connectorx/sources/csv/mod.rs

connectorx/sources/csv/
mod.rs
