1use std::io::Write;
2
3use arrow_array::RecordBatch;
4use arrow_schema::SchemaRef;
5use serde_json::{Map as JsonMap, Value};
6
7use crate::Result;
8use crate::cli::BinaryFormat;
9use crate::output::{RowWriter, value};
10
11pub struct JsonlRowWriter<W: Write> {
12 writer: W,
13 field_names: Vec<String>,
14 binary_format: BinaryFormat,
15}
16
17impl<W: Write> JsonlRowWriter<W> {
18 pub fn new(writer: W, binary_format: BinaryFormat) -> Self {
19 Self {
20 writer,
21 field_names: Vec::new(),
22 binary_format,
23 }
24 }
25}
26
27impl<W: Write> RowWriter for JsonlRowWriter<W> {
28 fn start(&mut self, schema: &SchemaRef) -> Result<()> {
29 self.field_names = schema.fields().iter().map(|f| f.name().clone()).collect();
30 Ok(())
31 }
32
33 fn write_batch(&mut self, batch: &RecordBatch) -> Result<()> {
34 debug_assert_eq!(
35 self.field_names.len(),
36 batch.num_columns(),
37 "start() must be called with the batch's schema"
38 );
39 let num_rows = batch.num_rows();
40 let num_cols = batch.num_columns();
41 for row in 0..num_rows {
42 let mut obj = JsonMap::with_capacity(num_cols);
43 for col in 0..num_cols {
44 let arr = batch.column(col);
45 let v = value::json_value(arr.as_ref(), row, self.binary_format)?;
46 obj.insert(self.field_names[col].clone(), v);
47 }
48 serde_json::to_writer(&mut self.writer, &Value::Object(obj))?;
49 self.writer.write_all(b"\n")?;
50 }
51 Ok(())
52 }
53
54 fn finish(&mut self) -> Result<()> {
55 self.writer.flush()?;
56 Ok(())
57 }
58}