(arrow-rs) branch master updated: Fix arrow-json writer empty (#6614)

[tustvold]

This is an automated email from the ASF dual-hosted git repository.

tustvold pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git


The following commit(s) were added to refs/heads/master by this push:
     new f86668c5f9 Fix arrow-json writer empty (#6614)
f86668c5f9 is described below

commit f86668c5f972bd3c3d8f7740f21f3d1c3b783f8b
Author: Antonin Amand <antonin.am...@gmail.com>
AuthorDate: Wed Oct 23 18:39:21 2024 +0200

    Fix arrow-json writer empty (#6614)
    
    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.
---
 arrow-json/src/writer/mod.rs | 42 ++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 40 insertions(+), 2 deletions(-)

diff --git a/arrow-json/src/writer/mod.rs b/arrow-json/src/writer/mod.rs
index df18a9ad5f..77e17d89f6 100644
--- a/arrow-json/src/writer/mod.rs
+++ b/arrow-json/src/writer/mod.rs
@@ -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(())
     }
 
@@ -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![