fill.rs

static USAGE: &str = r#"
Fill empty fields in selected columns of a CSV.

This command fills empty fields in the selected column
using the last seen non-empty field in the CSV. This is
useful to forward-fill values which may only be included
the first time they are encountered.

The option `--default <value>` fills all empty values
in the selected columns with the provided default value.

The option `--first` fills empty values using the first
seen non-empty value in that column, instead of the most
recent non-empty value in that column.

The option `--backfill` fills empty values at the start of
the CSV with the first valid value in that column. This
requires buffering rows with empty values in the target
column which appear before the first valid value.

The option `--groupby` groups the rows by the specified
columns before filling in the empty values. Using this
option, empty values are only filled with values which
belong to the same group of rows, as determined by the
columns selected in the `--groupby` option.

When both `--groupby` and `--backfill` are specified, and the
CSV is not sorted by the `--groupby` columns, rows may be
re-ordered during output due to the buffering of rows
collected before the first valid value.

For examples, see https://github.com/jqnatividad/qsv/blob/master/tests/test_fill.rs.

Usage:
    qsv fill [options] [--] <selection> [<input>]
    qsv fill --help

fill options:
    -g --groupby <keys>    Group by specified columns.
    -f --first             Fill using the first valid value of a column, instead of the latest.
    -b --backfill          Fill initial empty values with the first valid value.
    -v --default <value>   Fill using this default value.

Common options:
    -h, --help             Display this message
    -o, --output <file>    Write output to <file> instead of stdout.
    -n, --no-headers       When set, the first row will not be interpreted
                           as headers. (i.e., They are not searched, analyzed,
                           sliced, etc.)
    -d, --delimiter <arg>  The field delimiter for reading CSV data.
                           Must be a single character. (default: ,)
"#;

use std::{io, iter, ops};

use ahash::AHashMap;
use serde::Deserialize;

use crate::{
    config::{Config, Delimiter},
    select::{SelectColumns, Selection},
    util,
    util::ByteString,
    CliResult,
};

type BoxedWriter = csv::Writer<Box<dyn io::Write + 'static>>;
type BoxedReader = csv::Reader<Box<dyn io::Read + Send + 'static>>;

#[derive(Deserialize)]
struct Args {
    arg_input:       Option<String>,
    arg_selection:   SelectColumns,
    flag_output:     Option<String>,
    flag_no_headers: bool,
    flag_delimiter:  Option<Delimiter>,
    flag_groupby:    Option<SelectColumns>,
    flag_first:      bool,
    flag_backfill:   bool,
    flag_default:    Option<String>,
}

pub fn run(argv: &[&str]) -> CliResult<()> {
    let args: Args = util::get_args(USAGE, argv)?;

    let rconfig = Config::new(args.arg_input.as_ref())
        .delimiter(args.flag_delimiter)
        .no_headers(args.flag_no_headers)
        .select(args.arg_selection);

    let wconfig = Config::new(args.flag_output.as_ref());

    let mut rdr = rconfig.reader()?;
    let mut wtr = wconfig.writer()?;

    let headers = rdr.byte_headers()?.clone();
    let select = rconfig.selection(&headers)?;
    let groupby = match args.flag_groupby {
        Some(value) => Some(value.selection(&headers, !rconfig.no_headers)?),
        None => None,
    };

    if !rconfig.no_headers {
        rconfig.write_headers(&mut rdr, &mut wtr)?;
    }

    let filler = Filler::new(groupby, select)
        .use_first_value(args.flag_first)
        .backfill_empty_values(args.flag_backfill)
        .use_default_value(args.flag_default);
    filler.fill(&mut rdr, &mut wtr)
}

#[derive(PartialEq, Eq, Hash, Clone, Debug)]
struct ByteRecord(Vec<ByteString>);

impl<'a> From<&'a csv::ByteRecord> for ByteRecord {
    fn from(record: &'a csv::ByteRecord) -> Self {
        ByteRecord(record.iter().map(<[u8]>::to_vec).collect())
    }
}

impl iter::FromIterator<ByteString> for ByteRecord {
    fn from_iter<T: IntoIterator<Item = ByteString>>(iter: T) -> Self {
        ByteRecord(Vec::from_iter(iter))
    }
}

impl iter::IntoIterator for ByteRecord {
    type IntoIter = ::std::vec::IntoIter<ByteString>;
    type Item = ByteString;

    fn into_iter(self) -> Self::IntoIter {
        self.0.into_iter()
    }
}

impl ops::Deref for ByteRecord {
    type Target = [ByteString];

    fn deref(&self) -> &[ByteString] {
        &self.0
    }
}

type GroupKey = Option<ByteRecord>;
type GroupBuffer = AHashMap<GroupKey, Vec<ByteRecord>>;
type Grouper = AHashMap<GroupKey, GroupValues>;
type GroupKeySelection = Option<Selection>;

trait GroupKeyConstructor {
    fn key(&self, record: &csv::ByteRecord) -> Result<GroupKey, String>;
}

impl GroupKeyConstructor for GroupKeySelection {
    fn key(&self, record: &csv::ByteRecord) -> Result<GroupKey, String> {
        match *self {
            Some(ref value) => Ok(Some(value.iter().map(|&i| record[i].to_vec()).collect())),
            None => Ok(None),
        }
    }
}

#[derive(Debug)]
struct GroupValues {
    map:     AHashMap<usize, ByteString>,
    default: Option<ByteString>,
}

impl GroupValues {
    fn new(default: Option<ByteString>) -> Self {
        Self {
            map: AHashMap::new(),
            default,
        }
    }
}

trait GroupMemorizer {
    fn fill(&self, selection: &Selection, record: ByteRecord) -> ByteRecord;
    fn memorize(&mut self, selection: &Selection, record: &csv::ByteRecord);
    fn memorize_first(&mut self, selection: &Selection, record: &csv::ByteRecord);
}

impl GroupMemorizer for GroupValues {
    fn memorize(&mut self, selection: &Selection, record: &csv::ByteRecord) {
        for &col in selection.iter().filter(|&col| !record[*col].is_empty()) {
            self.map.insert(col, record[col].to_vec());
        }
    }

    fn memorize_first(&mut self, selection: &Selection, record: &csv::ByteRecord) {
        for &col in selection.iter().filter(|&col| !record[*col].is_empty()) {
            self.map.entry(col).or_insert_with(|| record[col].to_vec());
        }
    }

    fn fill(&self, selection: &Selection, record: ByteRecord) -> ByteRecord {
        record
            .into_iter()
            .enumerate()
            .map_selected(selection, |(col, field)| {
                (
                    col,
                    if field.is_empty() {
                        self.default
                            .clone()
                            .or_else(|| self.map.get(&col).cloned())
                            .unwrap_or_else(|| field.clone())
                    } else {
                        field
                    },
                )
            })
            .map(|(_, field)| field)
            .collect()
    }
}

struct Filler {
    grouper:       Grouper,
    groupby:       GroupKeySelection,
    select:        Selection,
    buffer:        GroupBuffer,
    first:         bool,
    backfill:      bool,
    default_value: Option<ByteString>,
}

impl Filler {
    fn new(groupby: GroupKeySelection, select: Selection) -> Self {
        Self {
            grouper: Grouper::new(),
            groupby,
            select,
            buffer: GroupBuffer::new(),
            first: false,
            backfill: false,
            default_value: None,
        }
    }

    const fn use_first_value(mut self, first: bool) -> Self {
        self.first = first;
        self
    }

    const fn backfill_empty_values(mut self, backfill: bool) -> Self {
        self.backfill = backfill;
        self
    }

    fn use_default_value(mut self, value: Option<String>) -> Self {
        self.default_value = value.map(|v| v.as_bytes().to_vec());
        self
    }

    fn fill(mut self, rdr: &mut BoxedReader, wtr: &mut BoxedWriter) -> CliResult<()> {
        let mut record = csv::ByteRecord::new();

        while rdr.read_byte_record(&mut record)? {
            // Precompute groupby key
            let key = self.groupby.key(&record)?;

            // Record valid fields, and fill empty fields
            let default_value = self.default_value.clone();
            let group = self
                .grouper
                .entry(key.clone())
                .or_insert_with(|| GroupValues::new(default_value));

            match (self.default_value.is_some(), self.first) {
                (true, _) => {},
                (false, true) => group.memorize_first(&self.select, &record),
                (false, false) => group.memorize(&self.select, &record),
            };

            let row = group.fill(&self.select, ByteRecord::from(&record));

            // Handle buffering rows which still have nulls.
            if self.backfill && (self.select.iter().any(|&i| row[i] == b"")) {
                self.buffer.entry(key.clone()).or_default().push(row);
            } else {
                if let Some(rows) = self.buffer.remove(&key) {
                    for buffered_row in rows {
                        wtr.write_record(group.fill(&self.select, buffered_row).iter())?;
                    }
                }
                wtr.write_record(row.iter())?;
            }
        }

        // Ensure any remaining buffers are dumped at the end.
        for (key, rows) in self.buffer {
            let group = self.grouper.get(&key).unwrap();
            for buffered_row in rows {
                wtr.write_record(group.fill(&self.select, buffered_row).iter())?;
            }
        }

        wtr.flush()?;
        Ok(())
    }
}

struct MapSelected<I, F> {
    selection:       Vec<usize>,
    selection_index: usize,
    index:           usize,
    iterator:        I,
    predicate:       F,
}

impl<I: iter::Iterator, F> iter::Iterator for MapSelected<I, F>
where
    F: FnMut(I::Item) -> I::Item,
{
    type Item = I::Item;

    fn next(&mut self) -> Option<Self::Item> {
        let item = self.iterator.next()?;
        let result = match self.selection_index {
            ref mut sidx if (self.selection.get(*sidx) == Some(&self.index)) => {
                *sidx += 1;
                Some((self.predicate)(item))
            },
            _ => Some(item),
        };
        self.index += 1;
        result
    }
}

trait Selectable<B>
where
    Self: iter::Iterator<Item = B> + Sized,
{
    fn map_selected<F>(self, selector: &Selection, predicate: F) -> MapSelected<Self, F>
    where
        F: FnMut(B) -> B;
}

impl<B, C> Selectable<B> for C
where
    C: iter::Iterator<Item = B> + Sized,
{
    fn map_selected<F>(self, selector: &Selection, predicate: F) -> MapSelected<Self, F>
    where
        F: FnMut(B) -> B,
    {
        MapSelected {
            selection: selector.iter().copied().collect(),
            selection_index: 0,
            index: 0,
            iterator: self,
            predicate,
        }
    }
}