datui_lib/
csv_dialect.rs

1//! The parts of a CSV's dialect that Polars' reader does not have, named as the
2//! Frictionless Table Dialect names them.
3//!
4//! | Frictionless | Here | Done by |
5//! |---|---|---|
6//! | `commentChar` | `--comment` | Polars' `comment_prefix`, before the header and in the data |
7//! | `headerRows`, `headerJoin` | `--header-rows`, `header_join` | [`header_names`] reads those lines; Polars reads the rest without a header |
8//! | `skipInitialSpace` | `--skip-initial-space` | [`skip_initial_space`], lazy expressions over the text columns |
9//!
10//! Header names are trimmed whatever the dialect: [`shown_names`].
11
12use std::io::BufRead;
13
14use polars::prelude::*;
15
16/// What `header_join` is when nothing sets it: Frictionless' `headerJoin` default.
17pub const DEFAULT_HEADER_JOIN: &str = " ";
18
19pub use datui_cli::check_comment_char;
20
21/// The longest header line [`header_names`] reads: far wider than any real header,
22/// and a bound on what a file with no line breaks can make it hold.
23const MAX_HEADER_LINE: u64 = 16 << 20;
24
25/// The names of the columns, from the lines `rows` names (1-based, counted from the top
26/// of the file before anything is skipped), each split on `separator` and trimmed.
27///
28/// A column's name is its pieces from those lines, in the order `rows` gives them,
29/// joined with `join`; a blank piece adds nothing. A line that starts with `comment`
30/// is a header line all the same, since the user named it, and loses the prefix. A
31/// file that ends before the last line named is an error: it has no header there.
32pub fn header_names(
33    source: impl BufRead,
34    rows: &[usize],
35    join: &str,
36    separator: u8,
37    comment: Option<&str>,
38) -> color_eyre::Result<Vec<String>> {
39    let lines = named_lines(source, rows)?;
40    Ok(names_of(&lines, rows, join, separator, comment))
41}
42
43/// [`header_names`] from `lines`, the lines `rows` names as [`named_lines`] read them.
44pub fn names_of(
45    lines: &[Vec<u8>],
46    rows: &[usize],
47    join: &str,
48    separator: u8,
49    comment: Option<&str>,
50) -> Vec<String> {
51    let mut columns: Vec<Vec<String>> = Vec::new();
52    for (&row, line) in rows.iter().zip(lines) {
53        for (i, field) in header_fields(line, row, separator, comment)
54            .into_iter()
55            .enumerate()
56        {
57            if columns.len() <= i {
58                columns.resize_with(i + 1, Vec::new);
59            }
60            if !field.is_empty() {
61                columns[i].push(field);
62            }
63        }
64    }
65    columns
66        .into_iter()
67        .map(|pieces| pieces.join(join))
68        .collect()
69}
70
71/// The lines `rows` names (1-based, from the top of the file), in the order `rows`
72/// gives them, each with its line break. Only those lines are held, each up to a
73/// bound; a file that ends before the last of them is an error.
74pub fn named_lines(mut source: impl BufRead, rows: &[usize]) -> color_eyre::Result<Vec<Vec<u8>>> {
75    use std::io::Read;
76    let last = rows.iter().copied().max().unwrap_or(0);
77    // Only the named lines are kept; the others are passed over without being held.
78    let mut lines: Vec<Vec<u8>> = vec![Vec::new(); last];
79    // Whether every byte so far is blank: a file of nothing has no header, which a
80    // read of several files passes over, where a short file of text is an error.
81    let mut blank = true;
82    for (i, line) in lines.iter_mut().enumerate() {
83        let n = i + 1;
84        let read = if rows.contains(&n) {
85            let read = (&mut source)
86                .take(MAX_HEADER_LINE + 1)
87                .read_until(b'\n', line)?;
88            blank &= line.iter().all(u8::is_ascii_whitespace);
89            read
90        } else {
91            skip_line(&mut source, &mut blank)?
92        };
93        if read == 0 {
94            return Err(NoHeader { line: last, blank }.into());
95        }
96        if line.len() as u64 > MAX_HEADER_LINE {
97            return Err(color_eyre::eyre::eyre!(
98                "header line {n} is longer than {} MiB",
99                MAX_HEADER_LINE >> 20
100            ));
101        }
102    }
103    Ok(rows
104        .iter()
105        .map(|&row| {
106            row.checked_sub(1)
107                .and_then(|i| lines.get(i))
108                .cloned()
109                .unwrap_or_default()
110        })
111        .collect())
112}
113
114/// A file that ends before the header line a read needs. `blank` when all it holds
115/// is white space, or nothing.
116#[derive(Debug)]
117pub struct NoHeader {
118    pub line: usize,
119    pub blank: bool,
120}
121
122impl std::fmt::Display for NoHeader {
123    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
124        write!(f, "header line {} is past the end of the file", self.line)
125    }
126}
127
128impl std::error::Error for NoHeader {}
129
130/// Whether `e` says the file holds nothing but white space where its header should be.
131pub fn is_blank_file(e: &color_eyre::Report) -> bool {
132    e.chain()
133        .any(|cause| cause.downcast_ref::<NoHeader>().is_some_and(|h| h.blank))
134}
135
136/// Pass over one line of `source` without holding it, noting whether it is blank.
137fn skip_line(source: &mut impl BufRead, blank: &mut bool) -> std::io::Result<usize> {
138    let mut read = 0;
139    loop {
140        let buf = source.fill_buf()?;
141        if buf.is_empty() {
142            return Ok(read);
143        }
144        let (used, done) = match memchr::memchr(b'\n', buf) {
145            Some(at) => (at + 1, true),
146            None => (buf.len(), false),
147        };
148        *blank &= buf[..used].iter().all(u8::is_ascii_whitespace);
149        source.consume(used);
150        read += used;
151        if done {
152            return Ok(read);
153        }
154    }
155}
156
157/// The most bytes [`window`] reads: a bound on what a file of very long lines costs.
158const MAX_WINDOW_BYTES: u64 = 1 << 20;
159
160/// The first `rows` data lines of `source`, read on from where its header lines ended,
161/// each split on `separator` and trimmed: the lines a scan infers its types from. A
162/// line that starts with `comment`, or is blank, is not one. Stops at
163/// [`MAX_WINDOW_BYTES`].
164pub fn window(
165    source: impl BufRead,
166    rows: usize,
167    separator: u8,
168    comment: Option<&str>,
169) -> std::io::Result<Vec<Vec<String>>> {
170    Ok(window_of(source, rows, separator, comment)?.0)
171}
172
173/// [`window`], and whether a line it read holds bytes that are not UTF-8.
174pub fn window_of(
175    source: impl BufRead,
176    rows: usize,
177    separator: u8,
178    comment: Option<&str>,
179) -> std::io::Result<(Vec<Vec<String>>, bool)> {
180    let mut lossy = false;
181    let mut source = source.take(MAX_WINDOW_BYTES);
182    let comment = comment.filter(|c| !c.is_empty()).map(str::as_bytes);
183    let mut out = Vec::new();
184    let mut line = Vec::new();
185    while out.len() < rows {
186        line.clear();
187        if source.read_until(b'\n', &mut line)? == 0 {
188            break;
189        }
190        // A last line cut by the bound is not read: its last field may be cut too.
191        if !line.ends_with(b"\n") && source.limit() == 0 {
192            break;
193        }
194        let text = line.strip_suffix(b"\n").unwrap_or(&line);
195        let text = text.strip_suffix(b"\r").unwrap_or(text);
196        if text.iter().all(u8::is_ascii_whitespace) || comment.is_some_and(|c| text.starts_with(c))
197        {
198            continue;
199        }
200        lossy |= std::str::from_utf8(text).is_err();
201        out.push(
202            split_fields(text, separator)
203                .into_iter()
204                .map(|f| f.trim().to_string())
205                .collect(),
206        );
207    }
208    Ok((out, lossy))
209}
210
211/// Pass over `n` lines of `source`.
212pub fn skip_lines(source: &mut impl BufRead, n: usize) -> std::io::Result<()> {
213    let mut blank = true;
214    for _ in 0..n {
215        if skip_line(source, &mut blank)? == 0 {
216            break;
217        }
218    }
219    Ok(())
220}
221
222/// Header line `row`'s fields, trimmed: without a byte-order mark on line 1, its line
223/// break, or `comment`'s prefix, split on `separator`.
224pub fn header_fields(line: &[u8], row: usize, separator: u8, comment: Option<&str>) -> Vec<String> {
225    let mut line = line;
226    if row == 1 {
227        line = line.strip_prefix(b"\xEF\xBB\xBF").unwrap_or(line);
228    }
229    line = line.strip_suffix(b"\n").unwrap_or(line);
230    line = line.strip_suffix(b"\r").unwrap_or(line);
231    if let Some(prefix) = comment.filter(|c| !c.is_empty()) {
232        line = line.strip_prefix(prefix.as_bytes()).unwrap_or(line);
233    }
234    split_fields(line, separator)
235        .into_iter()
236        .map(|f| f.trim().to_string())
237        .collect()
238}
239
240/// One line's fields, split on `separator` outside double quotes, with the quotes
241/// removed and a doubled quote read as one. Padding before an opening quote is
242/// dropped, so `  "Lcl Date"` is one quoted field.
243fn split_fields(line: &[u8], separator: u8) -> Vec<String> {
244    let mut fields = Vec::new();
245    let mut field: Vec<u8> = Vec::new();
246    let mut quoted = false;
247    let mut i = 0;
248    while i < line.len() {
249        let b = line[i];
250        if quoted {
251            if b == b'"' {
252                if line.get(i + 1) == Some(&b'"') {
253                    field.push(b'"');
254                    i += 1;
255                } else {
256                    quoted = false;
257                }
258            } else {
259                field.push(b);
260            }
261        } else if b == separator {
262            fields.push(String::from_utf8_lossy(&field).into_owned());
263            field.clear();
264        } else if b == b'"' && field.iter().all(u8::is_ascii_whitespace) {
265            field.clear();
266            quoted = true;
267        } else {
268            field.push(b);
269        }
270        i += 1;
271    }
272    fields.push(String::from_utf8_lossy(&field).into_owned());
273    fields
274}
275
276/// The names a read shows for the columns Polars named `raw`: the header lines'
277/// names by position when `--header-rows` gave them, else Polars' own, trimmed. A
278/// blank name is `column_N`, as Polars names a headerless file's; a name already
279/// taken gets `_duplicated_K`, as Polars marks a repeated header.
280pub fn shown_names(raw: &[PlSmallStr], header: Option<&[String]>) -> Vec<String> {
281    let names: Vec<String> = raw
282        .iter()
283        .enumerate()
284        .map(|(i, name)| {
285            let name = match header {
286                Some(header) => header.get(i).map_or("", String::as_str),
287                None => name.as_str(),
288            }
289            .trim();
290            if name.is_empty() {
291                format!("column_{}", i + 1)
292            } else {
293                name.to_string()
294            }
295        })
296        .collect();
297    let mut taken: PlHashSet<String> = PlHashSet::with_capacity(names.len());
298    let mut seen: PlHashMap<String, usize> = PlHashMap::with_capacity(names.len());
299    let mut out = Vec::with_capacity(names.len());
300    for name in names {
301        let count = seen.entry(name.clone()).or_insert(0);
302        let mut candidate = name.clone();
303        while !taken.insert(candidate.clone()) {
304            candidate = format!("{name}_duplicated_{count}");
305            *count += 1;
306        }
307        out.push(candidate);
308    }
309    out
310}
311
312/// `lf` with its columns named as [`shown_names`] says. Nothing is read: Polars has
313/// the schema from the scan's own inference.
314///
315/// With `header`, a file with nothing after its header lines (a log with no rows yet)
316/// is a table with those columns and no rows, as Polars reads a file that is only a
317/// header; Polars itself finds no data past the lines it skipped.
318pub fn name_columns(mut lf: LazyFrame, header: Option<&[String]>) -> PolarsResult<LazyFrame> {
319    let schema = match (lf.collect_schema(), header) {
320        (Err(PolarsError::NoData(_)), Some(header)) => return header_only(header),
321        (Ok(schema), Some(header)) if schema.is_empty() => return header_only(header),
322        (schema, _) => schema?,
323    };
324    let raw: Vec<PlSmallStr> = schema.iter_names().cloned().collect();
325    let shown = shown_names(&raw, header);
326    if raw.iter().zip(&shown).all(|(r, s)| r.as_str() == s) {
327        return Ok(lf);
328    }
329    Ok(lf.rename(raw.iter().map(|s| s.as_str()), shown.iter(), true))
330}
331
332/// No rows, a text column for each name `header` gives.
333fn header_only(header: &[String]) -> PolarsResult<LazyFrame> {
334    let raw: Vec<PlSmallStr> = (1..=header.len().max(1))
335        .map(|i| format!("column_{i}").into())
336        .collect();
337    let columns: Vec<Column> = shown_names(&raw, Some(header))
338        .into_iter()
339        .map(|name| Column::new_empty(name.into(), &DataType::String))
340        .collect();
341    Ok(DataFrame::new(0, columns)?.lazy())
342}
343
344/// An eager read of the lines after `header`'s, with nothing there read as no
345/// columns, which [`name_columns`] makes the header's columns with no rows.
346pub fn read_after_header(
347    read: PolarsResult<DataFrame>,
348    header: Option<&[String]>,
349) -> PolarsResult<DataFrame> {
350    match read {
351        Err(PolarsError::NoData(_)) if header.is_some() => Ok(DataFrame::empty()),
352        read => read,
353    }
354}
355
356/// `skipInitialSpace`: the spaces after a delimiter are not part of a text value, so
357/// `"   152.6"` is `"152.6"` and a cell of spaces is empty, which reads as null like an
358/// empty cell. `nulls` are the null values this column takes, matched after the
359/// padding is gone, as they would be against the unpadded file.
360pub fn skip_initial_space(
361    mut lf: LazyFrame,
362    nulls: impl Fn(&str) -> Vec<String>,
363) -> PolarsResult<LazyFrame> {
364    let schema = lf.collect_schema()?;
365    let exprs: Vec<Expr> = schema
366        .iter()
367        .filter(|(_, dtype)| **dtype == DataType::String)
368        .map(|(name, _)| {
369            let stripped = col(name.clone())
370                .str()
371                .strip_chars_start(lit(PlSmallStr::from_static(" ")));
372            let null = nulls(name.as_str())
373                .into_iter()
374                .fold(stripped.clone().eq(lit("")), |any, value| {
375                    any.or(stripped.clone().eq(lit(value)))
376                });
377            when(null)
378                .then(Null {}.lit().cast(DataType::String))
379                .otherwise(stripped)
380                .alias(name.clone())
381        })
382        .collect();
383    if exprs.is_empty() {
384        return Ok(lf);
385    }
386    Ok(lf.with_columns(exprs))
387}
388
389#[cfg(test)]
390mod tests {
391    use super::*;
392
393    fn names(text: &str, rows: &[usize], comment: Option<&str>) -> Vec<String> {
394        header_names(text.as_bytes(), rows, " ", b',', comment).unwrap()
395    }
396
397    #[test]
398    fn one_header_line_is_split_and_trimmed() {
399        let text = "#info\n  Lcl Date, Lcl Time,     Latitude\n1,2,3\n";
400        assert_eq!(
401            names(text, &[2], None),
402            ["Lcl Date", "Lcl Time", "Latitude"]
403        );
404    }
405
406    #[test]
407    fn several_lines_join_in_the_order_given_and_skip_blank_pieces() {
408        let text = "#yyyy-mm-dd, hh:mm:ss, degrees\n  Lcl Date, Lcl Time, Latitude\n";
409        assert_eq!(
410            names(text, &[2, 1], Some("#")),
411            [
412                "Lcl Date yyyy-mm-dd",
413                "Lcl Time hh:mm:ss",
414                "Latitude degrees"
415            ]
416        );
417        let text = "a,,c\nx,y\n";
418        assert_eq!(
419            header_names(text.as_bytes(), &[1, 2], "_", b',', None).unwrap(),
420            ["a_x", "y", "c"]
421        );
422    }
423
424    #[test]
425    fn quotes_bom_and_carriage_returns() {
426        let text = "\u{FEFF}id,  \"last, first\",\"say \"\"hi\"\"\"\r\n";
427        assert_eq!(names(text, &[1], None), ["id", "last, first", "say \"hi\""]);
428    }
429
430    #[test]
431    fn a_file_that_ends_before_the_header_is_an_error() {
432        let err = header_names("a,b\n".as_bytes(), &[1, 5], " ", b',', None).unwrap_err();
433        assert!(err.to_string().contains("past the end"), "{err}");
434        assert!(header_names("".as_bytes(), &[1], " ", b',', None).is_err());
435        let blank = |text: &str| is_blank_file(&named_lines(text.as_bytes(), &[3]).unwrap_err());
436        assert!(blank(""), "empty");
437        assert!(blank(" \n\t\r\n"), "white space");
438        assert!(!blank("#a\n"), "text, too short");
439        assert!(!blank("\nx\n"), "text on a line passed over");
440        // The last line needs no line break.
441        assert_eq!(names("#u\na,b", &[2], None), ["a", "b"]);
442    }
443
444    #[test]
445    fn the_window_is_the_data_lines_after_the_header() {
446        let text = "a,b\n  1,  x\n#note\n\n  , 2.5\n3,4\n";
447        let mut source = text.as_bytes();
448        skip_lines(&mut source, 1).unwrap();
449        let rows = window(source, 2, b',', Some("#")).unwrap();
450        assert_eq!(rows, [vec!["1", "x"], vec!["", "2.5"]]);
451    }
452
453    #[test]
454    fn a_header_line_is_read_up_to_a_bound() {
455        let wide = "x".repeat(MAX_HEADER_LINE as usize + 1);
456        let err = header_names(wide.as_bytes(), &[1], " ", b',', None).unwrap_err();
457        assert!(err.to_string().contains("header line 1"), "{err}");
458        // A line that is not named is passed over however long it is.
459        let text = format!("{wide}\na,b\n");
460        assert_eq!(names(&text, &[2], None), ["a", "b"]);
461    }
462
463    #[test]
464    fn shown_names_trim_fill_and_deduplicate() {
465        let raw: Vec<PlSmallStr> = ["  a", "a", " ", "b"].map(PlSmallStr::from).to_vec();
466        assert_eq!(
467            shown_names(&raw, None),
468            ["a", "a_duplicated_0", "column_3", "b"]
469        );
470        let raw: Vec<PlSmallStr> = (1..=4).map(|i| format!("column_{i}").into()).collect();
471        let header = ["x".to_string(), String::new(), "column_1".to_string()];
472        assert_eq!(
473            shown_names(&raw, Some(&header)),
474            ["x", "column_2", "column_1", "column_4"]
475        );
476    }
477
478    #[test]
479    fn padding_is_skipped_and_blank_or_null_values_are_null() {
480        let df = df!(
481            "a" => ["   1.5", "    ", "  NA", " x y "],
482            "n" => [1i64, 2, 3, 4],
483        )
484        .unwrap();
485        let out = skip_initial_space(df.lazy(), |_| vec!["NA".into()])
486            .unwrap()
487            .collect()
488            .unwrap();
489        let a: Vec<Option<&str>> = out.column("a").unwrap().str().unwrap().iter().collect();
490        assert_eq!(a, [Some("1.5"), None, None, Some("x y ")]);
491        assert_eq!(out.column("n").unwrap().dtype(), &DataType::Int64);
492    }
493}
datui_lib/csv_dialect.rs

datui_lib/
csv_dialect.rs
