jsonl.rs

static USAGE: &str = r#"
Convert newline-delimited JSON (JSONL/NDJSON) to CSV.

The command tries to do its best but since it is not possible to
straightforwardly convert JSON lines to CSV, the process might lose some complex
fields from the input.

Also, it will fail if the JSON documents are not consistent with one another,
as the first JSON line will be used to infer the headers of the CSV output.

For examples, see https://github.com/dathere/qsv/blob/master/tests/test_jsonl.rs.
See also https://github.com/dathere/qsv/wiki/Conversion-and-IO#jsonl

Usage:
    qsv jsonl [options] [<input>]
    qsv jsonl --help

jsonl options:
    --ignore-errors        Skip malformed input lines.
    -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.
    -b, --batch <size>     The number of rows per batch to load into memory,
                           before running in parallel. Set to 0 to load all
                           rows in one batch. [default: 50000]

Common options:
    -h, --help             Display this message
    -o, --output <file>    Write output to <file> instead of stdout.
    -d, --delimiter <arg>  The delimiter to use when writing CSV data.
                           Must be a single character. [default: ,]
"#;

use std::{
    fs,
    io::{self, BufRead, BufReader},
};

use rayon::{
    iter::{IndexedParallelIterator, ParallelIterator},
    prelude::IntoParallelRefIterator,
};
use serde::Deserialize;
use serde_json::Value;

use crate::{
    CliResult,
    config::{Config, DEFAULT_RDR_BUFFER_CAPACITY, Delimiter},
    util,
};

#[derive(Deserialize)]
struct Args {
    arg_input:          Option<String>,
    flag_output:        Option<String>,
    flag_delimiter:     Option<Delimiter>,
    flag_ignore_errors: bool,
    flag_jobs:          Option<usize>,
    flag_batch:         usize,
}

fn recurse_to_infer_headers(value: &Value, headers: &mut Vec<Vec<String>>, path: &[String]) {
    match value {
        Value::Object(map) => {
            for (key, value) in map {
                match value {
                    Value::Null
                    | Value::Bool(_)
                    | Value::Number(_)
                    | Value::String(_)
                    | Value::Array(_) => {
                        let mut full_path = path.to_owned();
                        full_path.push(key.to_string());

                        headers.push(full_path);
                    },
                    Value::Object(_) => {
                        let mut new_path = path.to_owned();
                        new_path.push(key.to_string());

                        recurse_to_infer_headers(value, headers, &new_path);
                    },
                    #[allow(unreachable_patterns)]
                    _ => {},
                }
            }
        },
        _ => {
            // top-level non-Object root: emit a single column with an empty
            // path so `get_value_at_path` returns the root value as-is.
            // The empty path is rendered as "value" in the CSV header row.
            headers.push(Vec::new());
        },
    }
}

fn infer_headers(value: &Value) -> Vec<Vec<String>> {
    let mut headers: Vec<Vec<String>> = Vec::new();

    recurse_to_infer_headers(value, &mut headers, &Vec::new());

    headers
}

fn get_value_at_path(value: &Value, path: &[String]) -> Option<Value> {
    let mut current = value;

    for key in path {
        match current.get(key) {
            Some(new_value) => {
                current = new_value;
            },
            None => {
                return None;
            },
        }
    }

    Some(current.clone())
}

#[inline]
fn json_line_to_csv_record(value: &Value, headers: &[Vec<String>]) -> csv::StringRecord {
    let mut record = csv::StringRecord::new();

    for path in headers {
        let value = get_value_at_path(value, path);

        match value {
            Some(value) => {
                record.push_field(&match value {
                    Value::Bool(v) => {
                        if v {
                            String::from("true")
                        } else {
                            String::from("false")
                        }
                    },
                    Value::Number(v) => v.to_string(),
                    Value::String(v) => v,
                    Value::Array(v) => v
                        .iter()
                        .map(std::string::ToString::to_string)
                        .collect::<Vec<_>>()
                        .join(","),
                    _ => String::new(),
                });
            },
            _ => {
                record.push_field("");
            },
        }
    }

    record
}

pub fn run(argv: &[&str]) -> CliResult<()> {
    let args: Args = util::get_args(USAGE, argv)?;
    let mut wtr = Config::new(args.flag_output.as_ref())
        .delimiter(args.flag_delimiter)
        .writer()?;

    let mut is_stdin = false;
    let mut rdr: Box<dyn BufRead> = match args.arg_input.clone() {
        None => {
            is_stdin = true;
            Box::new(BufReader::new(io::stdin()))
        },
        Some(p) => Box::new(BufReader::with_capacity(
            DEFAULT_RDR_BUFFER_CAPACITY,
            fs::File::open(p)?,
        )),
    };

    let mut headers: Vec<Vec<String>> = Vec::new();
    let mut headers_emitted: bool = false;

    // amortize memory allocation by reusing record
    let mut batch_line = String::new();

    // reuse batch buffers
    let batchsize: usize = if args.flag_batch == 0 {
        if is_stdin {
            // if stdin, we don't know how many lines there are
            // so just make a reasonably big batch size
            1_000_000
        } else {
            // safety: `arg_input` is guaranteed `Some` here: the stdin branch
            // above is the only path where it stays `None`.
            util::count_lines_in_file(&args.arg_input.unwrap())? as usize
        }
    } else {
        args.flag_batch
    };
    // each batch entry carries its 1-based input-line number, so error messages
    // and header inference can reference original line numbers even when
    // --ignore-errors silently skips read errors.
    let mut batch: Vec<(u64, String)> = Vec::with_capacity(batchsize);
    let mut batch_results: Vec<(u64, Option<csv::StringRecord>)> = Vec::with_capacity(batchsize);

    util::njobs(args.flag_jobs);

    let mut input_line_idx: u64 = 0;

    'batch_loop: loop {
        for _ in 0..batchsize {
            batch_line.clear();
            match rdr.read_line(&mut batch_line) {
                Ok(0) => {
                    // EOF
                    break;
                },
                Ok(_) => {
                    input_line_idx += 1;
                    batch.push((input_line_idx, batch_line.clone()));
                },
                Err(e) => {
                    if args.flag_ignore_errors {
                        // best-effort: assume one read_line error == one line
                        // skipped. BufRead::read_line consumes bytes up to and
                        // including the next newline, so this holds for the
                        // common cases (e.g. invalid-UTF-8 errors are returned
                        // after the bytes have already been consumed). For
                        // pathological I/O errors mid-stream this may drift by
                        // one or more lines.
                        input_line_idx += 1;
                        continue;
                    }
                    return fail_clierror!(
                        r#"Could not read input line!: {e}
Use `--ignore-errors` option to skip malformed input lines.
Use `tojsonl` command to convert _to_ jsonl instead of _from_ jsonl."#,
                    );
                },
            }
        }

        if batch.is_empty() {
            break 'batch_loop; // EOF
        }

        if !headers_emitted {
            // Under --ignore-errors, scan forward for the first parseable line
            // in this batch instead of failing on a malformed first line.
            let header_value: Option<Value> = if args.flag_ignore_errors {
                batch
                    .iter()
                    .find_map(|(_, line)| serde_json::from_str::<Value>(line).ok())
            } else {
                match serde_json::from_str::<Value>(&batch[0].1) {
                    Ok(v) => Some(v),
                    Err(e) => {
                        return fail_clierror!(
                            r#"Could not parse first input line as JSON to infer headers: {e}
Use `--ignore-errors` option to skip malformed input lines.
Use `tojsonl` command to convert _to_ jsonl instead of _from_ jsonl."#,
                        );
                    },
                }
            };

            if let Some(value) = header_value {
                headers = infer_headers(&value);

                let headers_formatted = headers
                    .iter()
                    .map(|v| {
                        if v.is_empty() {
                            // top-level non-Object root: synthesize a "value" column name
                            String::from("value")
                        } else {
                            v.join(".")
                        }
                    })
                    .collect::<Vec<String>>();
                let headers_record = csv::StringRecord::from(headers_formatted);
                wtr.write_record(&headers_record)?;

                headers_emitted = true;
            } else {
                // --ignore-errors set and no parseable line in this batch;
                // discard it and try the next batch.
                batch.clear();
                continue 'batch_loop;
            }
        }

        // do actual work via rayon
        batch
            .par_iter()
            .map(
                |(line_no, json_line)| match serde_json::from_str(json_line) {
                    Ok(v) => (*line_no, Some(json_line_to_csv_record(&v, &headers))),
                    Err(e) => {
                        if !args.flag_ignore_errors {
                            log::error!("serde_json::from_str error: {e:#?}");
                        }
                        (*line_no, None)
                    },
                },
            )
            .collect_into_vec(&mut batch_results);

        // rayon collect() guarantees original order, so we can just append results of each batch
        for (line_no, result_record) in &batch_results {
            if let Some(record) = result_record {
                wtr.write_record(record)?;
            } else if !args.flag_ignore_errors {
                // there was an error parsing a json line
                return fail_clierror!(
                    r#"Could not parse input line {line_no} as JSON
Use `--ignore-errors` option to skip malformed input lines.
Use `tojsonl` command to convert _to_ jsonl instead of _from_ jsonl."#,
                );
            }
        }

        batch.clear();
    } // end batch loop

    // if --ignore-errors swallowed every line, the output is empty (no header,
    // no rows). Surface a warning so this isn't silently indistinguishable
    // from "input was empty".
    if !headers_emitted && input_line_idx > 0 {
        log::warn!(
            "All {input_line_idx} input line(s) were unparsable; --ignore-errors produced empty \
             output."
        );
    }

    Ok(wtr.flush()?)
}