static USAGE: &str = r#"
Sort an arbitrarily large CSV/text file using a multithreaded external sort algorithm.
This command has TWO modes of operation.
* CSV MODE
when --select is set, it sorts based on the given column/s. Requires an index.
See `qsv select --help` for select syntax details.
* LINE MODE
when --select is NOT set, it sorts any input text file (not just CSVs) on a
line-by-line basis. If sorting a non-CSV file, be sure to set --no-headers,
otherwise, the first line will not be included in the external sort.
Usage:
qsv extsort [options] [<input>] [<output>]
qsv extsort --help
External sort option:
-s, --select <arg> Select a subset of columns to sort (CSV MODE).
Note that the outputs will remain at the full width of the CSV.
If --select is NOT set, extsort will work in LINE MODE, sorting
the input as a text file on a line-by-line basis.
-R, --reverse Reverse order
--memory-limit <arg> The maximum amount of memory to buffer the external merge sort.
If less than 50, this is a percentage of total memory.
If more than 50, this is the memory in MB to allocate, capped
at 90 percent of total memory.
[default: 20]
--tmp-dir <arg> The directory to use for externally sorting file segments.
[default: ./]
-j, --jobs <arg> The number of jobs to run in parallel.
When not set, the number of jobs is set to the
number of CPUs detected.
Common options:
CSV MODE ONLY:
-d, --delimiter <arg> The field delimiter for reading CSV data.
Must be a single character. (default: ,)
-h, --help Display this message
-n, --no-headers When set, the first row will not be interpreted
as headers and will be sorted with the rest
of the rows. Otherwise, the first row will always
appear as the header row in the output.
"#;
use std::{
fs,
io::{self, BufRead, Write, stdin, stdout},
path,
};
use ext_sort::{ExternalSorter, ExternalSorterBuilder, LimitedBufferBuilder};
use serde::Deserialize;
use crate::{
CliResult,
cmd::extdedup::calculate_memory_limit,
config,
config::{Config, Delimiter},
select::SelectColumns,
util,
};
#[derive(Deserialize)]
struct Args {
arg_input: Option<String>,
arg_output: Option<String>,
flag_select: Option<SelectColumns>,
flag_reverse: bool,
flag_delimiter: Option<Delimiter>,
flag_jobs: Option<usize>,
flag_memory_limit: Option<u64>,
flag_tmp_dir: Option<String>,
flag_no_headers: bool,
}
const RW_BUFFER_CAPACITY: usize = 1_000_000;
pub fn run(argv: &[&str]) -> CliResult<()> {
let args: Args = util::get_args(USAGE, argv)?;
let tmp_dir = match &args.flag_tmp_dir {
Some(tmp_dir) => {
if !path::Path::new(tmp_dir).exists() {
return fail_clierror!("tmp-dir '{tmp_dir}' does not exist");
}
tmp_dir.to_string()
},
None => "./".to_string(),
};
let mem_limited_buffer_bytes = calculate_memory_limit(args.flag_memory_limit);
log::info!("{mem_limited_buffer_bytes} bytes used for in memory mergesort buffer...");
let sorter: ExternalSorter<String, io::Error, LimitedBufferBuilder> =
match ExternalSorterBuilder::new()
.with_tmp_dir(path::Path::new(&tmp_dir))
.with_buffer(LimitedBufferBuilder::new(
mem_limited_buffer_bytes as usize,
true,
))
.with_rw_buf_size(RW_BUFFER_CAPACITY)
.with_threads_number(util::njobs(args.flag_jobs))
.build()
{
Ok(sorter) => sorter,
Err(e) => {
return fail_clierror!("cannot create external sorter: {e}");
},
};
if args.flag_select.is_some() {
sort_csv(&args, &tmp_dir, &sorter)
} else {
sort_lines(&args, &sorter)
}
}
fn sort_csv(
args: &Args,
tmp_dir: &str,
sorter: &ExternalSorter<String, io::Error, LimitedBufferBuilder>,
) -> Result<(), crate::clitypes::CliError> {
let rconfig = Config::new(args.arg_input.as_ref())
.delimiter(args.flag_delimiter)
.no_headers_flag(args.flag_no_headers)
.select(args.flag_select.clone().unwrap());
let mut idxfile = match rconfig.indexed() {
Ok(idx) => {
if idx.is_none() {
return fail_incorrectusage_clierror!("extsort CSV mode requires an index");
}
idx.unwrap()
},
_ => {
return fail_incorrectusage_clierror!("extsort CSV mode requires an index");
},
};
let mut input_rdr = rconfig.reader()?;
let linewtr_tfile = tempfile::NamedTempFile::new_in(tmp_dir)?;
let mut line_wtr = io::BufWriter::with_capacity(RW_BUFFER_CAPACITY, linewtr_tfile.as_file());
let headers = input_rdr.byte_headers()?.clone();
let sel = rconfig.selection(&headers)?;
let mut sort_key = String::with_capacity(20);
let mut utf8_string = String::with_capacity(20);
let mut curr_row = csv::ByteRecord::new();
let rowcount = idxfile.count();
let width = rowcount.to_string().len();
for row in input_rdr.byte_records() {
curr_row.clone_from(&row?);
sort_key.clear();
for field in sel.select(&curr_row) {
if let Ok(s_utf8) = simdutf8::basic::from_utf8(field) {
sort_key.push_str(s_utf8);
} else {
utf8_string.clear();
utf8_string.push_str(&String::from_utf8_lossy(field));
sort_key.push_str(&utf8_string);
}
}
let idx_position = curr_row.position().unwrap();
writeln!(line_wtr, "{sort_key}|{:01$}", idx_position.line(), width)?;
}
line_wtr.flush()?;
let line_rdr = io::BufReader::with_capacity(
RW_BUFFER_CAPACITY,
std::fs::File::open(linewtr_tfile.path())?,
);
let reverse_flag = args.flag_reverse;
let compare = |a: &String, b: &String| {
if reverse_flag {
a.cmp(b).reverse()
} else {
a.cmp(b)
}
};
let sorted = match sorter.sort_by(line_rdr.lines(), compare) {
Ok(sorted) => sorted,
Err(e) => {
return fail!(format!("cannot do external sort: {e:?}"));
},
};
let sorted_tfile = tempfile::NamedTempFile::new_in(tmp_dir)?;
let mut sorted_line_wtr =
io::BufWriter::with_capacity(RW_BUFFER_CAPACITY, sorted_tfile.as_file());
for item in sorted.map(Result::unwrap) {
sorted_line_wtr.write_all(format!("{item}\n").as_bytes())?;
}
sorted_line_wtr.flush()?;
drop(line_wtr);
linewtr_tfile.close()?;
let sorted_lines = std::fs::File::open(sorted_tfile.path())?;
let sorted_line_rdr = io::BufReader::with_capacity(RW_BUFFER_CAPACITY, sorted_lines);
let mut sorted_csv_wtr = Config::new(args.arg_output.as_ref()).writer()?;
let position_delta: u64 = if args.flag_no_headers {
1
} else {
sorted_csv_wtr.write_byte_record(&headers)?;
2
};
let mut record_wrk = csv::ByteRecord::new();
let mut line = String::new();
for l in sorted_line_rdr.lines() {
line.clone_from(&l?);
let Ok(position) = atoi_simd::parse::<u64>(&line.as_bytes()[line.len() - width..]) else {
return fail!("Failed to retrieve position: invalid integer");
};
idxfile.seek(position.saturating_sub(position_delta))?;
idxfile.read_byte_record(&mut record_wrk)?;
sorted_csv_wtr.write_byte_record(&record_wrk)?;
}
sorted_csv_wtr.flush()?;
drop(sorted_line_wtr);
sorted_tfile.close()?;
Ok(())
}
fn sort_lines(
args: &Args,
sorter: &ExternalSorter<String, io::Error, LimitedBufferBuilder>,
) -> Result<(), crate::clitypes::CliError> {
let mut input_rdr: Box<dyn BufRead> = match &args.arg_input {
Some(input_path) => {
if input_path.to_lowercase().ends_with(".sz") {
return fail_incorrectusage_clierror!(
"Input file cannot be a .sz file. Use 'qsv snappy decompress' first."
);
}
let file = fs::File::open(input_path)?;
Box::new(io::BufReader::with_capacity(
config::DEFAULT_RDR_BUFFER_CAPACITY,
file,
))
},
None => Box::new(io::BufReader::new(stdin().lock())),
};
let mut output_wtr: Box<dyn Write> = match &args.arg_output {
Some(output_path) => {
if output_path.to_lowercase().ends_with(".sz") {
return fail_clierror!(
"Output file cannot be a .sz file. Compress it after sorting with 'qsv snappy \
compress'."
);
}
Box::new(io::BufWriter::with_capacity(
RW_BUFFER_CAPACITY,
fs::File::create(output_path)?,
))
},
None => Box::new(io::BufWriter::with_capacity(
RW_BUFFER_CAPACITY,
stdout().lock(),
)),
};
let mut header = String::new();
if !args.flag_no_headers {
input_rdr.read_line(&mut header)?;
}
let reverse_flag = args.flag_reverse;
let compare = |a: &String, b: &String| {
if reverse_flag {
a.cmp(b).reverse()
} else {
a.cmp(b)
}
};
let sorted = match sorter.sort_by(input_rdr.lines(), compare) {
Ok(sorted) => sorted,
Err(e) => {
return fail!(format!("cannot do external sort: {e:?}"));
},
};
if !header.is_empty() {
output_wtr.write_all(format!("{}\n", header.trim_end()).as_bytes())?;
}
for item in sorted.map(Result::unwrap) {
output_wtr.write_all(format!("{item}\n").as_bytes())?;
}
output_wtr.flush()?;
Ok(())
}