-
Notifications
You must be signed in to change notification settings - Fork 70
/
join.rs
460 lines (421 loc) · 17.1 KB
/
join.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
static USAGE: &str = r#"
Joins two sets of CSV data on the specified columns.
The default join operation is an 'inner' join. This corresponds to the
intersection of rows on the keys specified.
Joins are always done by ignoring leading and trailing whitespace. By default,
joins are done case sensitively, but this can be disabled with the --ignore-case
flag.
For examples, see https://github.com/jqnatividad/qsv/blob/master/tests/test_join.rs.
Usage:
qsv join [options] <columns1> <input1> <columns2> <input2>
qsv join --help
input arguments:
<input1> is the first CSV data set to join.
<input2> is the second CSV data set to join.
<columns1> & <columns2> are the columns to join on for each input.
The columns arguments specify the columns to join for each input. Columns can
be referenced by name or index, starting at 1. Specify multiple columns by
separating them with a comma. Specify a range of columns with `-`. Both
columns1 and columns2 must specify exactly the same number of columns.
(See 'qsv select --help' for the full syntax.)
For <input1> and <input2>, specifying `-` indicates reading from stdin.
e.g. 'qsv frequency -s Agency nyc311.csv | qsv join value - id nycagencyinfo.csv'
join options:
-i, --ignore-case When set, joins are done case insensitively.
--left Do a 'left outer' join. This returns all rows in
first CSV data set, including rows with no
corresponding row in the second data set. When no
corresponding row exists, it is padded out with
empty fields.
--left-anti Do a 'left anti' join. This returns all rows in
first CSV data set that has no match with the
second data set.
--left-semi Do a 'left semi' join. This returns all rows in
first CSV data set that has a match with the
second data set.
--right Do a 'right outer' join. This returns all rows in
second CSV data set, including rows with no
corresponding row in the first data set. When no
corresponding row exists, it is padded out with
empty fields. (This is the reverse of 'outer left'.)
--full Do a 'full outer' join. This returns all rows in
both data sets with matching records joined. If
there is no match, the missing side will be padded
out with empty fields. (This is the combination of
'outer left' and 'outer right'.)
--cross USE WITH CAUTION.
This returns the cartesian product of the CSV
data sets given. The number of rows return is
equal to N * M, where N and M correspond to the
number of rows in the given data sets, respectively.
--nulls When set, joins will work on empty fields.
Otherwise, empty fields are completely ignored.
(In fact, any row that has an empty field in the
key specified is ignored.)
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::{collections::hash_map::Entry, fmt, io, iter::repeat, str};
use ahash::AHashMap;
use byteorder::{BigEndian, WriteBytesExt};
use serde::Deserialize;
use crate::{
config::{Config, Delimiter, SeekRead},
index::Indexed,
select::{SelectColumns, Selection},
util,
util::ByteString,
CliResult,
};
#[derive(Deserialize)]
struct Args {
arg_columns1: SelectColumns,
arg_input1: String,
arg_columns2: SelectColumns,
arg_input2: String,
flag_left: bool,
flag_left_anti: bool,
flag_left_semi: bool,
flag_right: bool,
flag_full: bool,
flag_cross: bool,
flag_output: Option<String>,
flag_no_headers: bool,
flag_ignore_case: bool,
flag_nulls: bool,
flag_delimiter: Option<Delimiter>,
}
pub fn run(argv: &[&str]) -> CliResult<()> {
let args: Args = util::get_args(USAGE, argv)?;
let mut state = args.new_io_state()?;
match (
args.flag_left,
args.flag_left_anti,
args.flag_left_semi,
args.flag_right,
args.flag_full,
args.flag_cross,
) {
(true, false, false, false, false, false) => {
state.write_headers()?;
state.outer_join(false)
},
(false, true, false, false, false, false) => {
state.write_headers1()?;
state.left_join(true)
},
(false, false, true, false, false, false) => {
state.write_headers1()?;
state.left_join(false)
},
(false, false, false, true, false, false) => {
state.write_headers()?;
state.outer_join(true)
},
(false, false, false, false, true, false) => {
state.write_headers()?;
state.full_outer_join()
},
(false, false, false, false, false, true) => {
state.write_headers()?;
state.cross_join()
},
(false, false, false, false, false, false) => {
state.write_headers()?;
state.inner_join()
},
_ => fail_incorrectusage_clierror!("Please pick exactly one join operation."),
}
}
struct IoState<R, W: io::Write> {
wtr: csv::Writer<W>,
rdr1: csv::Reader<R>,
sel1: Selection,
rdr2: csv::Reader<R>,
sel2: Selection,
no_headers: bool,
casei: bool,
nulls: bool,
}
impl<R: io::Read + io::Seek, W: io::Write> IoState<R, W> {
fn write_headers(&mut self) -> CliResult<()> {
if !self.no_headers {
let mut headers = self.rdr1.byte_headers()?.clone();
headers.extend(self.rdr2.byte_headers()?.iter());
self.wtr.write_record(&headers)?;
}
Ok(())
}
fn write_headers1(&mut self) -> CliResult<()> {
if !self.no_headers {
let headers = self.rdr1.byte_headers()?;
self.wtr.write_record(headers)?;
}
Ok(())
}
fn inner_join(mut self) -> CliResult<()> {
let mut scratch = csv::ByteRecord::new();
let mut validx = ValueIndex::new(self.rdr2, &self.sel2, self.casei, self.nulls)?;
let mut row = csv::ByteRecord::new();
let mut key;
while self.rdr1.read_byte_record(&mut row)? {
key = get_row_key(&self.sel1, &row, self.casei);
if let Some(rows) = validx.values.get(&key) {
for &rowi in rows {
validx.idx.seek(rowi as u64)?;
validx.idx.read_byte_record(&mut scratch)?;
let combined = row.iter().chain(scratch.iter());
self.wtr.write_record(combined)?;
}
}
}
Ok(self.wtr.flush()?)
}
fn outer_join(mut self, right: bool) -> CliResult<()> {
if right {
::std::mem::swap(&mut self.rdr1, &mut self.rdr2);
::std::mem::swap(&mut self.sel1, &mut self.sel2);
}
let mut scratch = csv::ByteRecord::new();
let (_, pad2) = self.get_padding()?;
let mut validx = ValueIndex::new(self.rdr2, &self.sel2, self.casei, self.nulls)?;
let mut row = csv::ByteRecord::new();
let mut key;
while self.rdr1.read_byte_record(&mut row)? {
key = get_row_key(&self.sel1, &row, self.casei);
if let Some(rows) = validx.values.get(&key) {
for &rowi in rows {
validx.idx.seek(rowi as u64)?;
let row1 = row.iter();
validx.idx.read_byte_record(&mut scratch)?;
if right {
self.wtr.write_record(scratch.iter().chain(row1))?;
} else {
self.wtr.write_record(row1.chain(&scratch))?;
}
}
} else if right {
self.wtr.write_record(pad2.iter().chain(&row))?;
} else {
self.wtr.write_record(row.iter().chain(&pad2))?;
}
}
Ok(self.wtr.flush()?)
}
fn left_join(mut self, anti: bool) -> CliResult<()> {
let validx = ValueIndex::new(self.rdr2, &self.sel2, self.casei, self.nulls)?;
let mut row = csv::ByteRecord::new();
let mut key;
while self.rdr1.read_byte_record(&mut row)? {
key = get_row_key(&self.sel1, &row, self.casei);
if validx.values.get(&key).is_none() {
if anti {
self.wtr.write_record(&row)?;
}
} else if !anti {
self.wtr.write_record(&row)?;
}
}
Ok(self.wtr.flush()?)
}
fn full_outer_join(mut self) -> CliResult<()> {
let mut scratch = csv::ByteRecord::new();
let (pad1, pad2) = self.get_padding()?;
let mut validx = ValueIndex::new(self.rdr2, &self.sel2, self.casei, self.nulls)?;
// Keep track of which rows we've written from rdr2.
let mut rdr2_written: Vec<_> = repeat(false).take(validx.num_rows).collect();
let mut row1 = csv::ByteRecord::new();
let mut key;
while self.rdr1.read_byte_record(&mut row1)? {
key = get_row_key(&self.sel1, &row1, self.casei);
if let Some(rows) = validx.values.get(&key) {
for &rowi in rows {
rdr2_written[rowi] = true;
validx.idx.seek(rowi as u64)?;
validx.idx.read_byte_record(&mut scratch)?;
self.wtr.write_record(row1.iter().chain(&scratch))?;
}
} else {
self.wtr.write_record(row1.iter().chain(&pad2))?;
}
}
// OK, now write any row from rdr2 that didn't get joined with a row
// from rdr1.
for (i, &written) in rdr2_written.iter().enumerate() {
if !written {
validx.idx.seek(i as u64)?;
validx.idx.read_byte_record(&mut scratch)?;
self.wtr.write_record(pad1.iter().chain(&scratch))?;
}
}
Ok(self.wtr.flush()?)
}
fn cross_join(mut self) -> CliResult<()> {
let mut pos = csv::Position::new();
pos.set_byte(0);
let mut row2 = csv::ByteRecord::new();
let mut row1 = csv::ByteRecord::new();
let rdr2_has_headers = self.rdr2.has_headers();
while self.rdr1.read_byte_record(&mut row1)? {
self.rdr2.seek(pos.clone())?;
if rdr2_has_headers {
// Read and skip the header row, since CSV readers disable
// the header skipping logic after being seeked.
self.rdr2.read_byte_record(&mut row2)?;
}
while self.rdr2.read_byte_record(&mut row2)? {
self.wtr.write_record(row1.iter().chain(&row2))?;
}
}
Ok(self.wtr.flush()?)
}
fn get_padding(&mut self) -> CliResult<(csv::ByteRecord, csv::ByteRecord)> {
let len1 = self.rdr1.byte_headers()?.len();
let len2 = self.rdr2.byte_headers()?.len();
Ok((
repeat(b"").take(len1).collect(),
repeat(b"").take(len2).collect(),
))
}
}
impl Args {
fn new_io_state(
&self,
) -> CliResult<IoState<Box<dyn SeekRead + 'static>, Box<dyn io::Write + 'static>>> {
let rconf1 = Config::new(&Some(self.arg_input1.clone()))
.delimiter(self.flag_delimiter)
.no_headers(self.flag_no_headers)
.select(self.arg_columns1.clone());
let rconf2 = Config::new(&Some(self.arg_input2.clone()))
.delimiter(self.flag_delimiter)
.no_headers(self.flag_no_headers)
.select(self.arg_columns2.clone());
let mut rdr1 = rconf1.reader_file_stdin()?;
let mut rdr2 = rconf2.reader_file_stdin()?;
let (sel1, sel2) = self.get_selections(&rconf1, &mut rdr1, &rconf2, &mut rdr2)?;
Ok(IoState {
wtr: Config::new(&self.flag_output).writer()?,
rdr1,
sel1,
rdr2,
sel2,
no_headers: rconf1.no_headers,
casei: self.flag_ignore_case,
nulls: self.flag_nulls,
})
}
#[allow(clippy::unused_self)]
fn get_selections<R: io::Read>(
&self,
rconf1: &Config,
rdr1: &mut csv::Reader<R>,
rconf2: &Config,
rdr2: &mut csv::Reader<R>,
) -> CliResult<(Selection, Selection)> {
let headers1 = rdr1.byte_headers()?;
let headers2 = rdr2.byte_headers()?;
let select1 = rconf1.selection(headers1)?;
let select2 = rconf2.selection(headers2)?;
if select1.len() != select2.len() {
return fail_incorrectusage_clierror!(
"Column selections must have the same number of columns, but found column \
selections with {} and {} columns.",
select1.len(),
select2.len()
);
}
Ok((select1, select2))
}
}
struct ValueIndex<R> {
// This maps tuples of values to corresponding rows.
values: AHashMap<Vec<ByteString>, Vec<usize>>,
idx: Indexed<R, io::Cursor<Vec<u8>>>,
num_rows: usize,
}
impl<R: io::Read + io::Seek> ValueIndex<R> {
fn new(
mut rdr: csv::Reader<R>,
sel: &Selection,
casei: bool,
nulls: bool,
) -> CliResult<ValueIndex<R>> {
let mut val_idx = AHashMap::with_capacity(10000);
let mut row_idx = io::Cursor::new(Vec::with_capacity(8 * 10000));
let (mut rowi, mut count) = (0_usize, 0_usize);
// This logic is kind of tricky. Basically, we want to include
// the header row in the line index (because that's what csv::index
// does), but we don't want to include header values in the ValueIndex.
if rdr.has_headers() {
// ... so if there are headers, we make sure that we've parsed
// them, and write the offset of the header row to the index.
rdr.byte_headers()?;
row_idx.write_u64::<BigEndian>(0)?;
count += 1;
} else {
// ... and if there are no headers, we seek to the beginning and
// index everything.
let mut pos = csv::Position::new();
pos.set_byte(0);
rdr.seek(pos)?;
}
let mut row = csv::ByteRecord::new();
while rdr.read_byte_record(&mut row)? {
// This is a bit hokey. We're doing this manually instead of using
// the `csv-index` crate directly so that we can create both
// indexes in one pass.
row_idx.write_u64::<BigEndian>(row.position().unwrap().byte())?;
let fields: Vec<_> = sel
.select(&row)
.map(|v| util::transform(v, casei))
.collect();
if nulls || !fields.iter().any(std::vec::Vec::is_empty) {
match val_idx.entry(fields) {
Entry::Vacant(v) => {
let mut rows = Vec::with_capacity(4);
rows.push(rowi);
v.insert(rows);
},
Entry::Occupied(mut v) => {
v.get_mut().push(rowi);
},
}
}
rowi += 1;
count += 1;
}
row_idx.write_u64::<BigEndian>(count as u64)?;
let idx = Indexed::open(rdr, io::Cursor::new(row_idx.into_inner()))?;
Ok(ValueIndex {
values: val_idx,
idx,
num_rows: rowi,
})
}
}
impl<R> fmt::Debug for ValueIndex<R> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
// Sort the values by order of first appearance.
let mut kvs = self.values.iter().collect::<Vec<_>>();
kvs.sort_by(|&(_, v1), &(_, v2)| v1[0].cmp(&v2[0]));
for (keys, rows) in kvs {
// This is just for debugging, so assume Unicode for now.
let keys = keys
.iter()
.map(|k| String::from_utf8(k.clone()).unwrap())
.collect::<Vec<_>>();
writeln!(f, "({}) => {rows:?}", keys.join(", "))?;
}
Ok(())
}
}
#[inline]
fn get_row_key(sel: &Selection, row: &csv::ByteRecord, casei: bool) -> Vec<ByteString> {
sel.select(row).map(|v| util::transform(v, casei)).collect()
}