fix issue#2058 file_format/json.rs attempt to subtract with overflow (a…

…pache#2066)

* fix  issue#2058 file_format/json.rs attempt to subtract with overflow

* issue#2058 add infer_schema_with_limit test

Co-authored-by: p00512853 <[email protected]>

datafusion/src/datasource/file_format/json.rs

@@ -68,13 +68,15 @@ impl FileFormat for JsonFormat {
             let iter = ValueIter::new(&mut reader, None);
             let schema = infer_json_schema_from_iterator(iter.take_while(|_| {
                 let should_take = records_to_read > 0;
-                records_to_read -= 1;
+                if should_take {
+                    records_to_read -= 1;
+                }
                 should_take
             }))?;
+            schemas.push(schema);
             if records_to_read == 0 {
                 break;
             }
-            schemas.push(schema);
         }
 
         let schema = Schema::try_merge(schemas)?;
@@ -228,4 +230,20 @@ mod tests {
             .await?;
         Ok(exec)
     }
+
+    #[tokio::test]
+    async fn infer_schema_with_limit() {
+        let filename = "tests/jsons/schema_infer_limit.json";
+        let format = JsonFormat::default().with_schema_infer_max_rec(Some(3));
+        let file_schema = format
+            .infer_schema(local_object_reader_stream(vec![filename.to_owned()]))
+            .await
+            .expect("Schema inference");
+        let fields = file_schema
+            .fields()
+            .iter()
+            .map(|f| format!("{}: {:?}", f.name(), f.data_type()))
+            .collect::<Vec<_>>();
+        assert_eq!(vec!["a: Int64", "b: Float64", "c: Boolean"], fields);
+    }
 }

datafusion/tests/jsons/schema_infer_limit.json

@@ -0,0 +1,4 @@
+{"a":1}
+{"a":-10, "b":-3.5}
+{"a":2, "b":0.6, "c":false}
+{"a":1, "b":2.0, "c":false, "d":"4"}