NdJsonReadOptions add array json support

Author: zhuqi-lucas

datafusion/core/src/datasource/file_format/json.rs

@@ -47,6 +47,7 @@ mod tests {
     use datafusion_common::stats::Precision;
 
     use datafusion_common::Result;
+    use datafusion_datasource::file_compression_type::FileCompressionType;
     use futures::StreamExt;
     use insta::assert_snapshot;
     use object_store::local::LocalFileSystem;
@@ -578,4 +579,89 @@ mod tests {
 
         Ok(())
     }
+
+    #[tokio::test]
+    async fn test_ndjson_read_options_format_array() -> Result<()> {
+        let ctx = SessionContext::new();
+
+        // Create a temporary file with JSON array format
+        let tmp_dir = tempfile::TempDir::new()?;
+        let path = format!("{}/array.json", tmp_dir.path().to_string_lossy());
+        std::fs::write(
+            &path,
+            r#"[
+            {"a": 1, "b": "hello"},
+            {"a": 2, "b": "world"},
+            {"a": 3, "b": "test"}
+        ]"#,
+        )?;
+
+        // Use NdJsonReadOptions with format_array = true
+        let options = NdJsonReadOptions::default().format_array(true);
+
+        ctx.register_json("json_array_table", &path, options)
+            .await?;
+
+        let result = ctx
+            .sql("SELECT a, b FROM json_array_table ORDER BY a")
+            .await?
+            .collect()
+            .await?;
+
+        assert_snapshot!(batches_to_string(&result), @r"
+    +---+-------+
+    | a | b     |
+    +---+-------+
+    | 1 | hello |
+    | 2 | world |
+    | 3 | test  |
+    +---+-------+
+    ");
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn test_ndjson_read_options_format_array_with_compression() -> Result<()> {
+        use flate2::Compression;
+        use flate2::write::GzEncoder;
+        use std::io::Write;
+
+        let ctx = SessionContext::new();
+
+        // Create a temporary gzip compressed JSON array file
+        let tmp_dir = tempfile::TempDir::new()?;
+        let path = format!("{}/array.json.gz", tmp_dir.path().to_string_lossy());
+
+        let json_content = r#"[{"a": 1, "b": "hello"}, {"a": 2, "b": "world"}]"#;
+        let file = std::fs::File::create(&path)?;
+        let mut encoder = GzEncoder::new(file, Compression::default());
+        encoder.write_all(json_content.as_bytes())?;
+        encoder.finish()?;
+
+        // Use NdJsonReadOptions with format_array and GZIP compression
+        let options = NdJsonReadOptions::default()
+            .format_array(true)
+            .file_compression_type(FileCompressionType::GZIP)
+            .file_extension(".json.gz");
+
+        ctx.register_json("json_array_gzip", &path, options).await?;
+
+        let result = ctx
+            .sql("SELECT a, b FROM json_array_gzip ORDER BY a")
+            .await?
+            .collect()
+            .await?;
+
+        assert_snapshot!(batches_to_string(&result), @r"
+    +---+-------+
+    | a | b     |
+    +---+-------+
+    | 1 | hello |
+    | 2 | world |
+    +---+-------+
+    ");
+
+        Ok(())
+    }
 }

datafusion/core/src/datasource/file_format/options.rs

@@ -465,6 +465,9 @@ pub struct NdJsonReadOptions<'a> {
     pub infinite: bool,
     /// Indicates how the file is sorted
     pub file_sort_order: Vec<Vec<SortExpr>>,
+    /// Whether the JSON file is in array format `[{...}, {...}]` instead of
+    /// line-delimited format. Defaults to `false`.
+    pub format_array: bool,
 }
 
 impl Default for NdJsonReadOptions<'_> {
@@ -477,6 +480,7 @@ impl Default for NdJsonReadOptions<'_> {
             file_compression_type: FileCompressionType::UNCOMPRESSED,
             infinite: false,
             file_sort_order: vec![],
+            format_array: false,
         }
     }
 }
@@ -529,6 +533,13 @@ impl<'a> NdJsonReadOptions<'a> {
         self.schema_infer_max_records = schema_infer_max_records;
         self
     }
+
+    /// Specify whether the JSON file is in array format `[{...}, {...}]`
+    /// instead of line-delimited format.
+    pub fn format_array(mut self, format_array: bool) -> Self {
+        self.format_array = format_array;
+        self
+    }
 }
 
 #[async_trait]
@@ -663,7 +674,8 @@ impl ReadOptions<'_> for NdJsonReadOptions<'_> {
         let file_format = JsonFormat::default()
             .with_options(table_options.json)
             .with_schema_infer_max_rec(self.schema_infer_max_records)
-            .with_file_compression_type(self.file_compression_type.to_owned());
+            .with_file_compression_type(self.file_compression_type.to_owned())
+            .with_format_array(self.format_array);
 
         ListingOptions::new(Arc::new(file_format))
             .with_file_extension(self.file_extension)