From 269838385e183a1f221c6ec872230529ae957520 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 and avro example

---
 datafusion-cli/src/print_format.rs       | 25 +++++++++++-------------
 datafusion-examples/examples/avro_sql.rs |  4 ++--
 2 files changed, 13 insertions(+), 16 deletions(-)

diff --git a/datafusion-cli/src/print_format.rs b/datafusion-cli/src/print_format.rs
index b7de60ac858f..9ea811c3a92b 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)
diff --git a/datafusion-examples/examples/avro_sql.rs b/datafusion-examples/examples/avro_sql.rs
index 2489f3f42f81..b819f2b591bc 100644
--- a/datafusion-examples/examples/avro_sql.rs
+++ b/datafusion-examples/examples/avro_sql.rs
@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use datafusion::arrow::io::print;
+use datafusion::arrow_print;
 
 use datafusion::error::Result;
 use datafusion::prelude::*;
@@ -45,7 +45,7 @@ async fn main() -> Result<()> {
     let results = df.collect().await?;
 
     // print the results
-    print::print(&results);
+    println!("{}", arrow_print::write(&results));
 
     Ok(())
 }