From aaaed9b47af3600b0244c952b8b444a8b1627459 Mon Sep 17 00:00:00 2001
From: Qingping Hou <dave2008713@gmail.com>
Date: Wed, 19 Jan 2022 22:08:15 -0800
Subject: [PATCH] fix cli json print

---
 datafusion-cli/src/print_format.rs | 25 +++++++++++--------------
 1 file changed, 11 insertions(+), 14 deletions(-)

diff --git a/datafusion-cli/src/print_format.rs b/datafusion-cli/src/print_format.rs
index b7de60ac858f8..9ea811c3a92be 100644
--- a/datafusion-cli/src/print_format.rs
+++ b/datafusion-cli/src/print_format.rs
@@ -71,24 +71,21 @@ impl fmt::Display for PrintFormat {
 }
 
 fn print_batches_to_json<J: JsonFormat>(batches: &[RecordBatch]) -> Result<String> {
+    use arrow::io::json::write as json_write;
+
     if batches.is_empty() {
         return Ok("{}".to_string());
     }
     let mut bytes = vec![];
-    let schema = batches[0].schema();
-    let names = schema
-        .fields
-        .iter()
-        .map(|f| f.name.clone())
-        .collect::<Vec<String>>();
-    for batch in batches {
-        arrow::io::json::write::serialize(
-            &names,
-            batch.columns(),
-            J::default(),
-            &mut bytes,
-        );
-    }
+
+    let format = J::default();
+    let blocks = json_write::Serializer::new(
+        batches.iter().map(|r| Ok(r.clone())),
+        vec![],
+        format,
+    );
+    json_write::write(&mut bytes, format, blocks)?;
+
     let formatted = String::from_utf8(bytes)
         .map_err(|e| DataFusionError::Execution(e.to_string()))?;
     Ok(formatted)