From 25a35f87761bd364860264afbba56d33b3ce74d1 Mon Sep 17 00:00:00 2001
From: Jorge Leitao <jorgecarleitao@gmail.com>
Date: Wed, 29 Dec 2021 00:18:59 +0000
Subject: [PATCH] Fixed error in writing json (#720)

---
 src/io/json/write/serialize.rs |  2 +-
 tests/it/io/json/mod.rs        |  6 ++----
 tests/it/io/json/write.rs      | 35 +++++++++++++++++++++++++++-------
 3 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/src/io/json/write/serialize.rs b/src/io/json/write/serialize.rs
index c195e2f69f3..769d1e28551 100644
--- a/src/io/json/write/serialize.rs
+++ b/src/io/json/write/serialize.rs
@@ -222,7 +222,7 @@ where
                 let item = iter.next().unwrap();
                 record.push((name.as_ref(), item));
             });
-        serialize_item(buffer, &record, format, false);
+        serialize_item(buffer, &record, format, is_first_row);
         is_first_row = false;
     })
 }
diff --git a/tests/it/io/json/mod.rs b/tests/it/io/json/mod.rs
index df54fa0c404..e50abab78c3 100644
--- a/tests/it/io/json/mod.rs
+++ b/tests/it/io/json/mod.rs
@@ -22,9 +22,7 @@ fn read_batch(data: String, fields: Vec<Field>) -> Result<RecordBatch> {
     json_read::deserialize(rows, fields)
 }
 
-fn write_batch(batch: RecordBatch) -> Result<Vec<u8>> {
-    let format = json_write::LineDelimited::default();
-
+fn write_batch<F: json_write::JsonFormat>(batch: RecordBatch, format: F) -> Result<Vec<u8>> {
     let batches = vec![Ok(batch)].into_iter();
 
     let blocks = json_write::Serializer::new(batches, vec![], format);
@@ -41,7 +39,7 @@ fn round_trip(data: String) -> Result<()> {
 
     let batch = read_batch(data.clone(), fields)?;
 
-    let buf = write_batch(batch)?;
+    let buf = write_batch(batch, json_write::LineDelimited::default())?;
 
     let result = String::from_utf8(buf).unwrap();
     println!("{}", result);
diff --git a/tests/it/io/json/write.rs b/tests/it/io/json/write.rs
index 9bbec6dd155..c7fa9cc3771 100644
--- a/tests/it/io/json/write.rs
+++ b/tests/it/io/json/write.rs
@@ -23,7 +23,7 @@ fn write_simple_rows() -> Result<()> {
 
     let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)]).unwrap();
 
-    let buf = write_batch(batch)?;
+    let buf = write_batch(batch, json_write::LineDelimited::default())?;
 
     assert_eq!(
         String::from_utf8(buf).unwrap(),
@@ -37,6 +37,27 @@ fn write_simple_rows() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn write_simple_rows_array() -> Result<()> {
+    let schema = Schema::new(vec![
+        Field::new("c1", DataType::Int32, false),
+        Field::new("c2", DataType::Utf8, false),
+    ]);
+
+    let a = Int32Array::from([Some(1), Some(2), Some(3), None, Some(5)]);
+    let b = Utf8Array::<i32>::from(&vec![Some("a"), Some("b"), Some("c"), Some("d"), None]);
+
+    let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)]).unwrap();
+
+    let buf = write_batch(batch, json_write::JsonArray::default())?;
+
+    assert_eq!(
+        String::from_utf8(buf).unwrap(),
+        r#"[{"c1":1,"c2":"a"},{"c1":2,"c2":"b"},{"c1":3,"c2":"c"},{"c1":null,"c2":"d"},{"c1":5,"c2":null}]"#
+    );
+    Ok(())
+}
+
 #[test]
 fn write_nested_struct_with_validity() -> Result<()> {
     let inner = vec![
@@ -71,7 +92,7 @@ fn write_nested_struct_with_validity() -> Result<()> {
 
     let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(c1), Arc::new(c2)]).unwrap();
 
-    let buf = write_batch(batch)?;
+    let buf = write_batch(batch, json_write::LineDelimited::default())?;
 
     assert_eq!(
         String::from_utf8(buf).unwrap(),
@@ -116,7 +137,7 @@ fn write_nested_structs() -> Result<()> {
 
     let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(c1), Arc::new(c2)]).unwrap();
 
-    let buf = write_batch(batch)?;
+    let buf = write_batch(batch, json_write::LineDelimited::default())?;
 
     assert_eq!(
         String::from_utf8(buf).unwrap(),
@@ -153,7 +174,7 @@ fn write_struct_with_list_field() -> Result<()> {
 
     let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)]).unwrap();
 
-    let buf = write_batch(batch)?;
+    let buf = write_batch(batch, json_write::LineDelimited::default())?;
 
     assert_eq!(
         String::from_utf8(buf).unwrap(),
@@ -199,7 +220,7 @@ fn write_nested_list() -> Result<()> {
 
     let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(c1), Arc::new(c2)]).unwrap();
 
-    let buf = write_batch(batch)?;
+    let buf = write_batch(batch, json_write::LineDelimited::default())?;
 
     assert_eq!(
         String::from_utf8(buf).unwrap(),
@@ -259,7 +280,7 @@ fn write_list_of_struct() -> Result<()> {
 
     let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(c1), Arc::new(c2)]).unwrap();
 
-    let buf = write_batch(batch)?;
+    let buf = write_batch(batch, json_write::LineDelimited::default())?;
 
     assert_eq!(
         String::from_utf8(buf).unwrap(),
@@ -278,7 +299,7 @@ fn write_escaped_utf8() -> Result<()> {
 
     let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
 
-    let buf = write_batch(batch)?;
+    let buf = write_batch(batch, json_write::LineDelimited::default())?;
 
     assert_eq!(
         String::from_utf8(buf).unwrap().as_bytes(),