-
Notifications
You must be signed in to change notification settings - Fork 70
/
python.rs
372 lines (303 loc) · 13.5 KB
/
python.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
static USAGE: &str = r#"
Create a new computed column or filter rows by evaluating a python expression on
every row of a CSV file.
The executed Python has 4 ways to reference cell values (as strings):
1. Directly by using column name (e.g. amount) as a local variable. If a column
name has spaces and other special characters, they are replaced with underscores
(e.g. "unit cost" -> unit_cost, "test-units/sec" -> test_units_sec)
2. Indexing cell value by column name as an attribute: row.amount
3. Indexing cell value by column name as a key: row["amount"]
4. Indexing cell value by column position: row[0]
Of course, if your input has no headers, then 4. will be the only available
option.
Some usage examples:
Sum numeric columns 'a' and 'b' and call new column 'c'
$ qsv py map c "int(a) + int(b)"
$ qsv py map c "int(col.a) + int(col['b'])"
$ qsv py map c "int(col[0]) + int(col[1])"
Use Python f-strings to calculate using multiple columns (qty, fruit & "unit cost")
and format into a new column 'formatted'
$ qsv py map formatted 'f"{qty} {fruit} cost ${(float(unit_cost) * float(qty)):.2f}"'
You can even have conditionals in your f-string:
$ qsv py map formatted 'f"""{qty} {fruit} cost ${(float(unit_cost) * float(qty)):.2f}. Its quite {"cheap" if ((float(unit_cost) * float(qty)) < 20.0) else "expensive"}!"""'
Note how we needed to use triple double quotes for the f-string, so we can use the literals
"cheap" and "expensive" in the f-string expression.
Strip and prefix cell values
$ qsv py map prefixed "'clean_' + a.strip()"
Filter some lines based on numerical filtering
$ qsv py filter "int(a) > 45"
Load helper file with function to compute Fibonacci sequence of the column "num_col"
$ qsv py map --helper fibonacci.py fib qsv_uh.fibonacci(num_col) data.csv
Below is a detailed example of the --helper option:
Use case:
Need to calculate checksum/md5sum of some columns. First column (c1) is "id", and do md5sum of
the rest of the columns (c2, c3 and c4).
Given test.csv:
c1,c2,c3,c4
1,a2,a3,a4
2,b2,b3,b4
3,c2,c3,c4
and hashhelper.py:
import hashlib
def md5hash (*args):
s = ",".join(args)
return(hashlib.md5(s.encode('utf-8')).hexdigest())
with the following command:
$ qsv py map --helper hashhelper.py hashcol 'qsv_uh.md5hash(c2,c3,c4)' test.csv
we get:
c1,c2,c3,c4,hashcol
1,a2,a3,a4,cb675342ed940908eef0844d17c35fab
2,b2,b3,b4,7d594b33f82bdcbc1cfa6f924a84c4cd
3,c2,c3,c4,6eabbfdbfd9ab6ae7737fb2b82f6a1af
NOTE: The prebuilt qsv binaries are linked against Python 3.10 and will require access
to the Python 3.10 shared libraries (libpython* on Linux/macOS, python*.dll on Windows)
during runtime for the py command to run.
Note that qsv with the `python` feature enabled will panic on startup even if you're not
using the `py` command if Python's shared libraries are not found.
If you wish qsv to use another Python version, you'll need to install/compile qsv from source.
Also, the following Python modules are automatically loaded and available to the user -
builtsin, math and random. The user can import additional modules with the --helper option.
With "py map", if a python expression is invalid, "<ERROR>" is returned.
With "py filter", if a python expression is invalid, no filtering is done.
Usage:
qsv py map [options] -n <script> [<input>]
qsv py map [options] <new-column> <script> [<input>]
qsv py map --helper <file> [options] <new-column> <script> [<input>]
qsv py filter [options] <script> [<input>]
qsv py map --help
qsv py filter --help
qsv py --help
py options:
-f, --helper <file> File containing Python code that's loaded into the
qsv_uh Python module. Functions with a return statement
in the file can be called with the prefix "qsv_uh".
The returned value is used in the map or filter operation.
-b, --batch <size> The number of rows per batch to process before
releasing memory and acquiring a new GILpool.
See https://pyo3.rs/v0.17.1/memory.html#gil-bound-memory
for more info. [default: 30000]
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. Namely, it will be sorted with the rest
of the rows. Otherwise, the first row will always
appear as the header row in the output.
-d, --delimiter <arg> The field delimiter for reading CSV data.
Must be a single character. (default: ,)
-p, --progressbar Show progress bars. Not valid for stdin.
"#;
use std::fs;
use indicatif::{ProgressBar, ProgressDrawTarget};
use log::{error, log_enabled, Level::Debug};
use pyo3::{intern, prelude::*, types::PyDict};
use serde::Deserialize;
use crate::{
config::{Config, Delimiter},
util, CliError, CliResult,
};
const HELPERS: &str = r#"
def cast_as_string(value):
if isinstance(value, str):
return value
return str(value)
def cast_as_bool(value):
return bool(value)
class QSVRow(object):
def __init__(self, headers):
self.__data = None
self.__headers = headers
self.__mapping = {h: i for i, h in enumerate(headers)}
def _update_underlying_data(self, row_data):
self.__data = row_data
def __getitem__(self, key):
if isinstance(key, int):
return self.__data[key]
return self.__data[self.__mapping[key]]
def __getattr__(self, key):
return self.__data[self.__mapping[key]]
"#;
#[derive(Deserialize)]
struct Args {
cmd_map: bool,
cmd_filter: bool,
arg_new_column: Option<String>,
arg_script: String,
flag_batch: u32,
flag_helper: Option<String>,
arg_input: Option<String>,
flag_output: Option<String>,
flag_no_headers: bool,
flag_delimiter: Option<Delimiter>,
flag_progressbar: bool,
}
impl From<PyErr> for CliError {
fn from(err: PyErr) -> CliError {
CliError::Other(err.to_string())
}
}
pub fn run(argv: &[&str]) -> CliResult<()> {
let args: Args = util::get_args(USAGE, argv)?;
let rconfig = Config::new(&args.arg_input)
.delimiter(args.flag_delimiter)
.no_headers(args.flag_no_headers);
let mut rdr = rconfig.reader()?;
let mut wtr = Config::new(&args.flag_output).writer()?;
if log_enabled!(Debug) {
Python::with_gil(|py| {
let msg = format!("Detected python={}", py.version());
winfo!("{msg}");
});
}
let mut helper_text = String::new();
if let Some(helper_file) = args.flag_helper {
helper_text = match fs::read_to_string(helper_file) {
Ok(helper_file) => helper_file,
Err(e) => return fail_clierror!("Cannot load python file: {e}"),
}
}
let mut headers = rdr.headers()?.clone();
let headers_len = headers.len();
if rconfig.no_headers {
headers = csv::StringRecord::new();
for i in 0..headers_len {
headers.push_field(&i.to_string());
}
} else {
if !args.cmd_filter {
let new_column = args
.arg_new_column
.as_ref()
.ok_or("Specify new column name")?;
headers.push_field(new_column);
}
wtr.write_record(&headers)?;
}
// prep progress bar
let show_progress =
(args.flag_progressbar || std::env::var("QSV_PROGRESSBAR").is_ok()) && !rconfig.is_stdin();
let progress = ProgressBar::with_draw_target(None, ProgressDrawTarget::stderr_with_hz(5));
if show_progress {
util::prep_progress(&progress, util::count_rows(&rconfig)?);
} else {
progress.set_draw_target(ProgressDrawTarget::hidden());
}
// ensure col/header names are valid and safe python variables
let header_vec = util::safe_header_names(&headers, true);
// amortize memory allocation by reusing record
#[allow(unused_assignments)]
let mut batch_record = csv::StringRecord::new();
// reuse batch buffers
let mut batch = Vec::with_capacity(args.flag_batch as usize);
// main loop to read CSV and construct batches.
// we batch python operations so that the GILPool does not get very large
// as we release the pool after each batch
// loop exits when batch is empty.
// see https://pyo3.rs/v0.17.1/memory.html#gil-bound-memory for more info.
loop {
for _ in 0..args.flag_batch {
match rdr.read_record(&mut batch_record) {
Ok(has_data) => {
if has_data {
batch.push(batch_record.clone());
} else {
// nothing else to add to batch
break;
}
}
Err(e) => {
return fail_clierror!("Error reading file: {e}");
}
}
}
if batch.is_empty() {
// break out of infinite loop when at EOF
break;
}
_ = Python::with_gil(|py| -> PyResult<()> {
{
let curr_batch = batch.clone();
let helpers = PyModule::from_code(py, HELPERS, "qsv_helpers.py", "qsv_helpers")?;
let batch_globals = PyDict::new(py);
let batch_locals = PyDict::new(py);
let user_helpers =
PyModule::from_code(py, &helper_text, "qsv_user_helpers.py", "qsv_uh")?;
batch_globals.set_item(intern!(py, "qsv_uh"), user_helpers)?;
// Global imports
let builtins = PyModule::import(py, "builtins")?;
let math_module = PyModule::import(py, "math")?;
let random_module = PyModule::import(py, "random")?;
batch_globals.set_item("__builtins__", builtins)?;
batch_globals.set_item("math", math_module)?;
batch_globals.set_item("random", random_module)?;
let py_row = helpers
.getattr("QSVRow")?
.call1((headers.iter().collect::<Vec<&str>>(),))?;
batch_locals.set_item("row", py_row)?;
let error_result = intern!(py, "<ERROR>");
for mut record in curr_batch {
// Initializing locals
let mut row_data: Vec<&str> = Vec::with_capacity(headers_len);
header_vec
.iter()
.enumerate()
.take(headers_len)
.for_each(|(i, key)| {
let cell_value = record.get(i).unwrap_or_default();
batch_locals
.set_item(key, cell_value)
.expect("cannot set_item");
row_data.push(cell_value);
});
py_row.call_method1(intern!(py, "_update_underlying_data"), (row_data,))?;
let result = py
.eval(&args.arg_script, Some(batch_globals), Some(batch_locals))
.map_err(|e| {
e.print_and_set_sys_last_vars(py);
if log_enabled!(Debug) {
error!("{e:?}");
}
"Evaluation of given expression failed with the above error!"
})
.unwrap_or(error_result);
if args.cmd_map {
let result = helpers
.getattr(intern!(py, "cast_as_string"))?
.call1((result,))?;
let value: String = result.extract()?;
record.push_field(&value);
if let Err(e) = wtr.write_record(&record) {
// we do this since we cannot use the ? operator here
// since this closure returns a PyResult
// this is converted to a CliError::Other anyway
return Err(pyo3::PyErr::new::<pyo3::exceptions::PyIOError, _>(
format!("cannot write record ({e})"),
));
}
} else if args.cmd_filter {
let result = helpers
.getattr(intern!(py, "cast_as_bool"))?
.call1((result,))?;
let value: bool = result.extract().unwrap_or(false);
if value {
if let Err(e) = wtr.write_record(&record) {
return Err(pyo3::PyErr::new::<pyo3::exceptions::PyIOError, _>(
format!("cannot write record ({e})"),
));
}
}
}
}
Ok(())
}
});
if show_progress {
progress.inc(batch.len() as u64);
}
batch.clear();
} // end loop
if show_progress {
util::finish_progress(&progress);
}
Ok(wtr.flush()?)
}