Skip to content

Commit

Permalink
Fix arrow-json writer empty
Browse files Browse the repository at this point in the history
The Writer output with the ArrayJson formatter was empty
when no record were in the output (no write or empty batch write).
It should output an empty JSON array `[]` instead.
  • Loading branch information
gwik committed Oct 22, 2024
1 parent 7e51d40 commit 081e350
Showing 1 changed file with 40 additions and 2 deletions.
42 changes: 40 additions & 2 deletions arrow-json/src/writer/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -367,10 +367,15 @@ where
/// all record batches have been produced. (e.g. producing the final `']'` if writing
/// arrays.
pub fn finish(&mut self) -> Result<(), ArrowError> {
if self.started && !self.finished {
if !self.started {
self.format.start_stream(&mut self.writer)?;
self.started = true;
}
if !self.finished {
self.format.end_stream(&mut self.writer)?;
self.finished = true;
}

Ok(())
}

Expand Down Expand Up @@ -1147,12 +1152,45 @@ mod tests {
}

#[test]
fn json_writer_empty() {
fn json_line_writer_empty() {
let mut writer = LineDelimitedWriter::new(vec![] as Vec<u8>);
writer.finish().unwrap();
assert_eq!(str::from_utf8(&writer.into_inner()).unwrap(), "");
}

#[test]
fn json_array_writer_empty() {
let mut writer = ArrayWriter::new(vec![] as Vec<u8>);
writer.finish().unwrap();
assert_eq!(str::from_utf8(&writer.into_inner()).unwrap(), "[]");
}

#[test]
fn json_line_writer_empty_batch() {
let mut writer = LineDelimitedWriter::new(vec![] as Vec<u8>);

let array = Int32Array::from(Vec::<i32>::new());
let schema = Schema::new(vec![Field::new("c", DataType::Int32, true)]);
let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(array)]).unwrap();

writer.write(&batch).unwrap();
writer.finish().unwrap();
assert_eq!(str::from_utf8(&writer.into_inner()).unwrap(), "");
}

#[test]
fn json_array_writer_empty_batch() {
let mut writer = ArrayWriter::new(vec![] as Vec<u8>);

let array = Int32Array::from(Vec::<i32>::new());
let schema = Schema::new(vec![Field::new("c", DataType::Int32, true)]);
let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(array)]).unwrap();

writer.write(&batch).unwrap();
writer.finish().unwrap();
assert_eq!(str::from_utf8(&writer.into_inner()).unwrap(), "[]");
}

#[test]
fn json_struct_array_nulls() {
let inner = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
Expand Down

0 comments on commit 081e350

Please sign in to comment.