Reimplement duplicate column handling while awake

I was tired and in a hurry when I implemented `--replace`, and I didn't
like the code. Here's a new version with a more flexible
`--duplicate-columns` flag and much cleaner code.

This also adds an integration test suite which actually runs the tool
from the command line and uses the SmartyStreets API.

Author: emk

src/addresses.rs

@@ -183,15 +183,15 @@ impl<Key: Default + Eq> AddressColumnSpec<Key> {
         self.address_columns_by_prefix.get(prefix)
     }
 
-    /// What column indices should we remove from the input records in order
+    /// What column should we remove from the input records in order
     /// to prevent duplicate columns?
     ///
-    /// These indices will always be returned in order.
-    pub fn column_indices_to_remove(
+    /// Returns the name and index of each column to remove, in order.
+    pub fn duplicate_columns<'header>(
         &self,
         structure: &Structure,
-        header: &StringRecord,
-    ) -> Result<Vec<usize>> {
+        header: &'header StringRecord,
+    ) -> Result<Vec<(&'header str, usize)>> {
         // Get all our column names for all prefixes, and insert them into a
         // hash table.
         let mut output_column_names = HashSet::new();
@@ -204,13 +204,13 @@ impl<Key: Default + Eq> AddressColumnSpec<Key> {
         }
 
         // Decide which columns of `header` need to be removed.
-        let mut indices_to_remove = vec![];
+        let mut duplicate_columns = vec![];
         for (i, col) in header.iter().enumerate() {
             if output_column_names.contains(col) {
-                indices_to_remove.push(i);
+                duplicate_columns.push((col, i));
             }
         }
-        Ok(indices_to_remove)
+        Ok(duplicate_columns)
     }
 }
 
@@ -235,8 +235,8 @@ fn find_columns_to_remove() {
     let structure = Structure::complete().unwrap();
     let header =
         StringRecord::from_iter(&["existing", "home_addressee", "work_addressee"]);
-    let indices = spec.column_indices_to_remove(&structure, &header).unwrap();
-    assert_eq!(indices, vec![1, 2]);
+    let indices = spec.duplicate_columns(&structure, &header).unwrap();
+    assert_eq!(indices, vec![("home_addressee", 1), ("work_addressee", 2)]);
 }
 
 impl AddressColumnSpec<String> {

src/geocoder.rs

@@ -10,6 +10,7 @@ use log::{debug, error, trace, warn};
 use std::{
     cmp::max, io, iter::FromIterator, sync::Arc, thread::sleep, time::Duration,
 };
+use strum_macros::EnumString;
 use tokio::{
     prelude::*,
     sync::mpsc::{self, Receiver, Sender},
@@ -32,6 +33,19 @@ const CONCURRENCY: usize = 48;
 /// The number of addresses to pass to SmartyStreets at one time.
 const GEOCODE_SIZE: usize = 72;
 
+/// What should we do if a geocoding output column has the same as a column in
+/// the input?
+#[derive(Debug, Clone, Copy, EnumString, Eq, PartialEq)]
+#[strum(serialize_all = "snake_case")]
+pub enum OnDuplicateColumns {
+    /// Fail with an error.
+    Error,
+    /// Replace existing columns with the same name.
+    Replace,
+    /// Leave the old columns in place and append the new ones.
+    Append,
+}
+
 /// Data about the CSV file that we include with every chunk to be geocoded.
 struct Shared {
     /// Which columns contain addresses that we need to geocode?
@@ -65,7 +79,7 @@ enum Message {
 pub async fn geocode_stdio(
     spec: AddressColumnSpec<String>,
     match_strategy: MatchStrategy,
-    replace_existing_columns: bool,
+    on_duplicate_columns: OnDuplicateColumns,
     structure: Structure,
 ) -> Result<()> {
     // Set up bounded channels for communication between the sync and async
@@ -76,7 +90,7 @@ pub async fn geocode_stdio(
     // Hook up our inputs and outputs, which are synchronous functions running
     // in their own threads.
     let read_fut = run_sync_fn_in_background("read CSV".to_owned(), move || {
-        read_csv_from_stdin(spec, structure, replace_existing_columns, in_tx)
+        read_csv_from_stdin(spec, structure, on_duplicate_columns, in_tx)
     });
     let write_fut = run_sync_fn_in_background("write CSV".to_owned(), move || {
         write_csv_to_stdout(out_rx)
@@ -152,7 +166,7 @@ pub async fn geocode_stdio(
 fn read_csv_from_stdin(
     spec: AddressColumnSpec<String>,
     structure: Structure,
-    replace_existing_columns: bool,
+    on_duplicate_columns: OnDuplicateColumns,
     mut tx: Sender<Message>,
 ) -> Result<()> {
     // Open up our CSV file and get the headers.
@@ -161,47 +175,56 @@ fn read_csv_from_stdin(
     let mut in_headers = rdr.headers()?.to_owned();
     debug!("input headers: {:?}", in_headers);
 
-    // Look for duplicate input columns, and decide what to do.
-    let column_indices_to_remove =
-        spec.column_indices_to_remove(&structure, &in_headers)?;
-    let remove_column_flags = if column_indices_to_remove.is_empty() {
-        // No columns to remove!
-        None
-    } else {
-        // We have to remove columns, so make a human-readable list...
-        let mut removed_names = vec![];
-        for i in &column_indices_to_remove {
-            removed_names.push(&in_headers[*i]);
-        }
-        let removed_names = removed_names.join(" ");
-
-        // And decide whether we're actually allowed to remove them or not.
-        if replace_existing_columns {
-            warn!("removing input columns: {}", removed_names);
+    // Figure out if we have any duplicate columns.
+    let (duplicate_column_indices, duplicate_column_names) = {
+        let duplicate_columns = spec.duplicate_columns(&structure, &in_headers)?;
+        let indices = duplicate_columns
+            .iter()
+            .map(|name_idx| name_idx.1)
+            .collect::<Vec<_>>();
+        let names = duplicate_columns
+            .iter()
+            .map(|name_idx| name_idx.0)
+            .collect::<Vec<_>>()
+            .join(", ");
+        (indices, names)
+    };
 
-            // Build the vector of bools specifying whether columns should
-            // stay or go.
-            let mut flags = vec![false; in_headers.len()];
-            for i in column_indices_to_remove {
-                flags[i] = true;
+    // If we do have duplicate columns, figure out what to do about it.
+    let mut should_remove_columns = false;
+    let mut remove_column_flags = vec![false; in_headers.len()];
+    if !duplicate_column_indices.is_empty() {
+        match on_duplicate_columns {
+            OnDuplicateColumns::Error => {
+                return Err(format_err!(
+                    "input columns would conflict with geocoding columns: {}",
+                    duplicate_column_names,
+                ));
+            }
+            OnDuplicateColumns::Replace => {
+                warn!("replacing input columns: {}", duplicate_column_names);
+                should_remove_columns = true;
+                for i in duplicate_column_indices.iter().cloned() {
+                    remove_column_flags[i] = true;
+                }
+            }
+            OnDuplicateColumns::Append => {
+                warn!(
+                    "output contains duplicate columns: {}",
+                    duplicate_column_names,
+                );
             }
-            Some(flags)
-        } else {
-            return Err(format_err!(
-                "input columns would conflict with geocoding columns: {}",
-                removed_names,
-            ));
         }
-    };
+    }
 
     // Remove any duplicate columns from our input headers.
-    if let Some(remove_column_flags) = &remove_column_flags {
+    if should_remove_columns {
         in_headers = remove_columns(&in_headers, &remove_column_flags);
     }
 
     // Convert our column spec from using header names to header indices.
     //
-    // This needs to use "post-removal" indices!
+    // This needs to happen _after_ `remove_columns` on our headers!
     let spec = spec.convert_to_indices_using_headers(&in_headers)?;
 
     // Decide how big to make our chunks. We want to geocode no more
@@ -228,9 +251,9 @@ fn read_csv_from_stdin(
     let mut rows = Vec::with_capacity(chunk_size);
     for row in rdr.records() {
         let mut row = row?;
-        if let Some(remove_column_flags) = &remove_column_flags {
+        if should_remove_columns {
             // Strip out any duplicate columns.
-            row = remove_columns(&row, remove_column_flags);
+            row = remove_columns(&row, &remove_column_flags);
         }
         rows.push(row);
         if rows.len() >= chunk_size {
@@ -270,6 +293,20 @@ fn read_csv_from_stdin(
     Ok(())
 }
 
+/// Remove columns from `row` if they're set to true in `remove_column_flags`.
+fn remove_columns(row: &StringRecord, remove_column_flags: &[bool]) -> StringRecord {
+    debug_assert_eq!(row.len(), remove_column_flags.len());
+    StringRecord::from_iter(row.iter().zip(remove_column_flags).filter_map(
+        |(value, &remove)| {
+            if remove {
+                None
+            } else {
+                Some(value.to_owned())
+            }
+        },
+    ))
+}
+
 /// Receive chunks of a CSV file from `rx` and write them to standard output.
 fn write_csv_to_stdout(mut rx: Receiver<Message>) -> Result<()> {
     let stdout = io::stdout();
@@ -410,15 +447,3 @@ async fn geocode_chunk(
     }
     Ok(chunk)
 }
-
-fn remove_columns(row: &StringRecord, remove_column_flags: &[bool]) -> StringRecord {
-    StringRecord::from_iter(row.iter().zip(remove_column_flags).filter_map(
-        |(value, &remove)| {
-            if remove {
-                None
-            } else {
-                Some(value.to_owned())
-            }
-        },
-    ))
-}

src/main.rs

@@ -18,7 +18,7 @@ mod structure;
 mod unpack_vec;
 
 use addresses::AddressColumnSpec;
-use geocoder::geocode_stdio;
+use geocoder::{geocode_stdio, OnDuplicateColumns};
 use smartystreets::MatchStrategy;
 use structure::Structure;
 
@@ -34,10 +34,10 @@ struct Opt {
     #[structopt(long = "match", default_value = "strict")]
     match_strategy: MatchStrategy,
 
-    /// Replace any existing columns with the same name as our geocoding
-    /// columns.
-    #[structopt(long = "replace")]
-    replace: bool,
+    /// What should we if geocoding output columns have the same names as input
+    /// columns? [error, replace, append]
+    #[structopt(long = "duplicate-columns", default_value = "error")]
+    on_duplicate_columns: OnDuplicateColumns,
 
     /// A JSON file describing what columns to geocode.
     #[structopt(long = "spec")]
@@ -58,7 +58,12 @@ fn run() -> Result<()> {
     let structure = Structure::complete()?;
 
     // Call our geocoder asynchronously.
-    let geocode_fut = geocode_stdio(spec, opt.match_strategy, opt.replace, structure);
+    let geocode_fut = geocode_stdio(
+        spec,
+        opt.match_strategy,
+        opt.on_duplicate_columns,
+        structure,
+    );
 
     // Pass our future to our async runtime.
     let mut runtime =

tests/about.rs

@@ -0,0 +1,19 @@
+//! Information about this CLI tool.
+
+use cli_test_dir::*;
+
+#[test]
+fn help_flag() {
+    let testdir = TestDir::new("geocode-csv", "help_flag");
+    let output = testdir.cmd().arg("--help").output().expect_success();
+    assert!(output.stdout_str().contains("geocode-csv"));
+    assert!(output.stdout_str().contains("--help"));
+}
+
+#[test]
+fn version_flag() {
+    let testdir = TestDir::new("geocode-csv", "version_flag");
+    let output = testdir.cmd().arg("--version").output().expect_success();
+    assert!(output.stdout_str().contains("geocode-csv"));
+    assert!(output.stdout_str().contains(env!("CARGO_PKG_VERSION")));
+}

tests/duplicate_columns.rs

@@ -0,0 +1,68 @@
+//! Specifying how to handle duplicate columns.
+
+use cli_test_dir::*;
+
+/// A CSV file to use for our tests.
+const SIMPLE_CSV: &str = r#"address,gc_addressee,zip
+20 W 34th St,,10118
+"#;
+
+/// A spec file to use for our tests.
+const SIMPLE_SPEC: &str = r#"{
+    "gc": {
+        "house_number_and_street": "address",
+        "postcode": "zip"
+    }
+}"#;
+
+#[test]
+#[ignore]
+fn duplicate_columns_error() {
+    let testdir = TestDir::new("geocode-csv", "duplicate_columns_error");
+
+    testdir.create_file("spec.json", SIMPLE_SPEC);
+    let output = testdir
+        .cmd()
+        .arg("--spec=spec.json")
+        .arg("--duplicate-columns=error")
+        .output_with_stdin(SIMPLE_CSV)
+        .expect("could not run geocode-csv");
+
+    assert!(!output.status.success());
+}
+
+#[test]
+#[ignore]
+fn duplicate_columns_replace() {
+    let testdir = TestDir::new("geocode-csv", "duplicate_columns_replace");
+
+    testdir.create_file("spec.json", SIMPLE_SPEC);
+    let output = testdir
+        .cmd()
+        .arg("--spec=spec.json")
+        .arg("--duplicate-columns=replace")
+        .output_with_stdin(SIMPLE_CSV)
+        .expect_success();
+
+    assert!(output.stdout_str().contains("address,zip,gc_addressee,"));
+    assert!(output.stdout_str().contains("Commercial"));
+}
+
+#[test]
+#[ignore]
+fn duplicate_columns_append() {
+    let testdir = TestDir::new("geocode-csv", "duplicate_columns_append");
+
+    testdir.create_file("spec.json", SIMPLE_SPEC);
+    let output = testdir
+        .cmd()
+        .arg("--spec=spec.json")
+        .arg("--duplicate-columns=append")
+        .output_with_stdin(SIMPLE_CSV)
+        .expect_success();
+
+    assert!(output
+        .stdout_str()
+        .contains("address,gc_addressee,zip,gc_addressee,"));
+    assert!(output.stdout_str().contains("Commercial"));
+}