fix: fix json decode number

Signed-off-by: fan <[email protected]>
apache · Nov 20, 2023 · 98c6e43 · 98c6e43
1 parent 61da64a
commit 98c6e43
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 2 deletions.
diff --git a/arrow-json/src/reader/mod.rs b/arrow-json/src/reader/mod.rs
@@ -717,7 +717,9 @@ mod tests {
 
     use arrow_array::cast::AsArray;
     use arrow_array::types::Int32Type;
-    use arrow_array::{make_array, Array, BooleanArray, ListArray, StringArray, StructArray};
+    use arrow_array::{
+        make_array, Array, BooleanArray, Float64Array, ListArray, StringArray, StructArray,
+    };
     use arrow_buffer::{ArrowNativeType, Buffer};
     use arrow_cast::display::{ArrayFormatter, FormatOptions};
     use arrow_data::ArrayDataBuilder;
@@ -2259,4 +2261,38 @@ mod tests {
             .values();
         assert_eq!(values, &[1699148028689, 2, 3, 4]);
     }
+
+    #[test]
+    fn test_coercing_primitive_into_string_decoder() {
+        let buf = r#"[
+            {"a": 1, "b": "A", "c": "T"},
+            {"a": 2, "b": "BB", "c": "F"},
+            {"a": 3, "b": 123, "c": false}
+        ]"#;
+        let schema = Schema::new(vec![
+            Field::new("a", DataType::Float64, true),
+            Field::new("b", DataType::Utf8, true),
+            Field::new("c", DataType::Utf8, true),
+        ]);
+        let json_array: Vec<serde_json::Value> = serde_json::from_str(buf).unwrap();
+        let schema_ref = Arc::new(schema);
+
+        // read record batches
+        let reader = ReaderBuilder::new(schema_ref.clone()).with_coerce_primitive(true);
+        let mut decoder = reader.build_decoder().unwrap();
+        decoder.serialize(json_array.as_slice()).unwrap();
+        let batch = decoder.flush().unwrap().unwrap();
+        assert_eq!(
+            batch,
+            RecordBatch::try_new(
+                schema_ref,
+                vec![
+                    Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0])),
+                    Arc::new(StringArray::from(vec!["A", "BB", "0"])),
+                    Arc::new(StringArray::from(vec!["T", "F", "false"])),
+                ]
+            )
+            .unwrap()
+        );
+    }
 }
diff --git a/arrow-json/src/reader/string_array.rs b/arrow-json/src/reader/string_array.rs
@@ -61,7 +61,15 @@ impl<O: OffsetSizeTrait> ArrayDecoder for StringArrayDecoder<O> {
                 TapeElement::Number(idx) if coerce_primitive => {
                     data_capacity += tape.get_string(idx).len();
                 }
-                _ => return Err(tape.error(*p, "string")),
+                TapeElement::I64(n) | TapeElement::I32(n) if coerce_primitive => {
+                    data_capacity += n.to_string().len();
+                }
+                TapeElement::F32(n) | TapeElement::F64(n) if coerce_primitive => {
+                    data_capacity += n.to_string().len();
+                }
+                _ => {
+                    return Err(tape.error(*p, "string"));
+                }
             }
         }
 
@@ -89,6 +97,12 @@ impl<O: OffsetSizeTrait> ArrayDecoder for StringArrayDecoder<O> {
                 TapeElement::Number(idx) if coerce_primitive => {
                     builder.append_value(tape.get_string(idx));
                 }
+                TapeElement::I64(n) | TapeElement::I32(n) if coerce_primitive => {
+                    builder.append_value(n.to_string());
+                }
+                TapeElement::F32(n) | TapeElement::F64(n) if coerce_primitive => {
+                    builder.append_value(n.to_string());
+                }
                 _ => unreachable!(),
             }
         }