Merge pull request #1875 from jqnatividad/1871-select-sort-random

`select`: `--sort` & `--random` options now work with the initial selection, not just the entire CSV

src/cmd/select.rs

@@ -27,13 +27,19 @@ selected using regular expressions.
  Reverse the order of columns:
  $ qsv select _-1
 
- Sort the columns lexicographically. Note that you must provide a dummy selector:
- $ qsv select 1 --sort
+ Sort the columns lexicographically (i.e. by their byte values)
+ $ qsv select 1- --sort
+
+ Select some columns and then sort them:
+ $ qsv select 1,4,5-7 --sort
 
  Randomly shuffle the columns:
- $ qsv select 1 --random
+ $ qsv select 1- --random
  # with a seed
- $ qsv select 1 --random --seed 42
+ $ qsv select 1- --random --seed 42
+
+ Select some columns and then shuffle them with a seed:
+ $ qsv select 1,4,5-7 --random --seed 42
 
  Select columns using a regex using '/<regex>/':
  $ qsv select /^a/
@@ -55,13 +61,12 @@ Usage:
 
 select options:
 These options only apply to the `select` command, not the `--select` flag in other commands.
-Be sure to provide a dummy selector (e.g. '1') to avoid command-line parsing errors.
 
- -R, --random Randomly reorder the columns. 
+ -R, --random Randomly shuffle the columns in the selection.
  --seed <number> Seed for the random number generator.
 
- -S, --sort Sort the columns lexicographically, i.e. by their
- byte values.
+ -S, --sort Sort the selected columns lexicographically,
+ i.e. by their byte values.
 
 Common options:
  -h, --help Display this message
@@ -95,78 +100,79 @@ struct Args {
 }
 
 pub fn run(argv: &[&str]) -> CliResult<()> {
- let mut args: Args = util::get_args(USAGE, argv)?;
+ let args: Args = util::get_args(USAGE, argv)?;
 
  if args.flag_sort && args.flag_random {
  return fail_clierror!("Cannot use both --random and --sort flags.");
  }
 
- if args.flag_random {
- // get the number of columns
- let num_cols = Config::new(&args.arg_input)
- .delimiter(args.flag_delimiter)
- .no_headers(true)
- .reader()?
- .byte_headers()?
- .len();
+ let rconfig = Config::new(&args.arg_input)
+ .delimiter(args.flag_delimiter)
+ .no_headers(args.flag_no_headers)
+ .select(args.arg_selection);
 
-  // make a vector of the column indices (1-indexed).
-  let mut original_selection: Vec<usize> = (1..=num_cols).collect();
+ let mut rdr = rconfig.reader()?;
+ let mut wtr = Config::new(&args.flag_output).writer()?;
 
- // Use seed if it is provided.
+ let headers = rdr.byte_headers()?.clone();
+ let sel = if args.flag_random {
+ // Use seed if it is provided when initializing the random number generator.
  let mut rng = if let Some(seed) = args.flag_seed {
  rand::rngs::StdRng::seed_from_u64(seed) // DevSkim: ignore DS148264
  } else {
  rand::rngs::StdRng::from_entropy()
  };
 
- // Shuffle the vector of column indices.
- original_selection.shuffle(&mut rng);
+ let initial_selection = rconfig.selection(&headers)?;
+
+ // make a vector of the column indices (1-indexed).
+ let mut shuffled_selection_vec: Vec<usize> =
+ initial_selection.iter().map(|&i| i + 1).collect();
+
+ shuffled_selection_vec.shuffle(&mut rng);
 
  // Convert the shuffled indices into a comma-separated string.
- let randomized_selection = original_selection
+ let shuffled_selection_string = shuffled_selection_vec
  .into_iter()
  .map(|i| i.to_string())
  .collect::<Vec<String>>()
  .join(",");
 
  // Parse the shuffled string into a SelectColumns object.
- args.arg_selection = SelectColumns::parse(&randomized_selection)?;
- }
-
- if args.flag_sort {
- // get the headers
- let headers = Config::new(&args.arg_input)
- .delimiter(args.flag_delimiter)
- .reader()?
- .byte_headers()?
- .clone();
+ let shuffled_selection = SelectColumns::parse(&shuffled_selection_string)?;
+
+ rconfig
+ .clone()
+ .select(shuffled_selection)
+ .selection(&headers)?
+ } else if args.flag_sort {
+ // get the headers from the initial selection
+ let initial_selection = rconfig.selection(&headers)?;
+ let mut initial_headers_vec = initial_selection
+ .iter()
+ .map(|&i| &headers[i])
+ .collect::<Vec<&[u8]>>();
 
  // sort the headers lexicographically
- let mut sorted_headers = Vec::with_capacity(headers.len());
- sorted_headers.extend(headers.iter().map(<[u8]>::to_vec));
- sorted_headers.sort_unstable();
+ initial_headers_vec.sort_unstable();
 
  // make a comma-separated string of the sorted, quoted headers
- let sorted_selection = sorted_headers
+ let sorted_selection_string = initial_headers_vec
  .iter()
  .map(|h| format!("\"{}\"", String::from_utf8_lossy(h)))
  .collect::<Vec<String>>()
  .join(",");
 
- args.arg_selection = SelectColumns::parse(&sorted_selection)?;
- }
+ // Parse the sorted selection string into a SelectColumns object.
+  let sorted_selection = SelectColumns::parse(&sorted_selection_string)?;
 
- let rconfig = Config::new(&args.arg_input)
- .delimiter(args.flag_delimiter)
- .no_headers(args.flag_no_headers)
- .select(args.arg_selection);
-
- let mut rdr = rconfig.reader()?;
- let mut wtr = Config::new(&args.flag_output).writer()?;
-
- let headers = rdr.byte_headers()?.clone();
- let sel = rconfig.selection(&headers)?;
+ rconfig
+ .clone()
+ .select(sorted_selection)
+ .selection(&headers)?
+ } else {
+ rconfig.selection(&headers)?
+ };
 
  if !rconfig.no_headers {
  wtr.write_record(sel.iter().map(|&i| &headers[i]))?;

tests/test_select.rs

@@ -231,7 +231,7 @@ fn test_select_sort() {
  let wrk = Workdir::new("test_select_sort");
  wrk.create("data.csv", unsorted_data(true));
  let mut cmd = wrk.command("select");
- cmd.arg("1").arg("--sort").arg("data.csv");
+ cmd.arg("1-").arg("--sort").arg("data.csv");
  let got: Vec<Vec<String>> = wrk.read_stdout(&mut cmd);
 
  let expected = vec![
@@ -252,12 +252,29 @@ fn test_select_sort() {
  assert_eq!(got, expected);
 }
 
+#[test]
+fn test_select_sort_subset() {
+ let wrk = Workdir::new("test_select_sort_subset");
+ wrk.create("data.csv", unsorted_data(true));
+ let mut cmd = wrk.command("select");
+ cmd.arg("4,7-9").arg("--sort").arg("data.csv");
+ let got: Vec<Vec<String>> = wrk.read_stdout(&mut cmd);
+
+ let expected = vec![
+ svec!["Bob", "Jürgen", "Đan", "İbrahim"],
+ svec!["value8", "value9", "value4", "value7"],
+ svec!["8", "9", "4", "7"],
+ svec!["value3", "value2", "value7", "value4"],
+ ];
+ assert_eq!(got, expected);
+}
+
 #[test]
 fn test_select_random_seeded() {
  let wrk = Workdir::new("test_select_random_seeded");
  wrk.create("data.csv", unsorted_data(true));
  let mut cmd = wrk.command("select");
- cmd.arg("1")
+ cmd.arg("1-")
  .arg("--random")
  .args(["--seed", "42"])
  .arg("data.csv");
@@ -280,3 +297,23 @@ fn test_select_random_seeded() {
  ];
  assert_eq!(got, expected);
 }
+
+#[test]
+fn test_select_random_seeded_subset() {
+ let wrk = Workdir::new("test_select_random_seeded_subset");
+ wrk.create("data.csv", unsorted_data(true));
+ let mut cmd = wrk.command("select");
+ cmd.arg("4,7-9")
+ .arg("--random")
+ .args(["--seed", "42"])
+ .arg("data.csv");
+ let got: Vec<Vec<String>> = wrk.read_stdout(&mut cmd);
+
+ let expected = vec![
+ svec!["Jürgen", "İbrahim", "Đan", "Bob"],
+ svec!["value9", "value7", "value4", "value8"],
+ svec!["9", "7", "4", "8"],
+ svec!["value2", "value4", "value7", "value3"],
+ ];
+ assert_eq!(got, expected);
+}