xsv: add --drop flag to partition command

The --drop flag drops the partition column itself.

PR #135

src/cmd/partition.rs

@@ -31,6 +31,7 @@ partition options:
  -p, --prefix-length <n> Truncate the partition column after the
  specified number of bytes when creating the
  output file.
+ --drop Drop the partition column from results.
 
 Common options:
  -h, --help Display this message
@@ -48,6 +49,7 @@ struct Args {
  arg_outdir: String,
  flag_filename: FilenameTemplate,
  flag_prefix_length: Option<usize>,
+ flag_drop: bool,
  flag_no_headers: bool,
  flag_delimiter: Option<Delimiter>,
 }
@@ -112,12 +114,22 @@ impl Args {
  // We have a new key, so make a new writer.
  let mut wtr = gen.writer(&*self.arg_outdir, key)?;
  if !rconfig.no_headers {
- wtr.write_record(&headers)?;
+ if self.flag_drop {
+ wtr.write_record(headers.iter().enumerate()
+ .filter_map(|(i, e)| if i != key_col { Some(e) } else { None }))?;
+ } else {
+ wtr.write_record(&headers)?;
+ }
  }
  vacant.insert(wtr)
  }
  };
- wtr.write_byte_record(&row)?;
+ if self.flag_drop {
+ wtr.write_record(row.iter().enumerate()
+ .filter_map(|(i, e)| if i != key_col { Some(e) } else { None }))?;
+ } else {
+ wtr.write_byte_record(&row)?;
+ }
  }
  Ok(())
  }

tests/test_partition.rs

@@ -46,6 +46,31 @@ TX,Fort Worth
 ");
 }
 
+#[test]
+fn partition_drop() {
+ let wrk = Workdir::new("partition");
+ wrk.create("in.csv", data(true));
+
+ let mut cmd = wrk.command("partition");
+ cmd.arg("--drop").arg("state").arg(&wrk.path(".")).arg("in.csv");
+ wrk.run(&mut cmd);
+
+ part_eq!(wrk, "CA.csv", "\
+city
+San Francisco
+");
+ part_eq!(wrk, "NY.csv", "\
+city
+Manhatten
+Buffalo
+");
+ part_eq!(wrk, "TX.csv", "\
+city
+Dallas
+Fort Worth
+");
+}
+
 #[test]
 fn partition_without_headers() {
  let wrk = Workdir::new("partition_without_headers");
@@ -68,6 +93,28 @@ TX,Fort Worth
 ");
 }
 
+#[test]
+fn partition_drop_without_headers() {
+ let wrk = Workdir::new("partition_without_headers");
+ wrk.create("in.csv", data(false));
+
+ let mut cmd = wrk.command("partition");
+ cmd.arg("--drop").arg("--no-headers").arg("1").arg(&wrk.path(".")).arg("in.csv");
+ wrk.run(&mut cmd);
+
+ part_eq!(wrk, "CA.csv", "\
+San Francisco
+");
+ part_eq!(wrk, "NY.csv", "\
+Manhatten
+Buffalo
+");
+ part_eq!(wrk, "TX.csv", "\
+Dallas
+Fort Worth
+");
+}
+
 #[test]
 fn partition_into_new_directory() {
  let wrk = Workdir::new("partition_into_new_directory");