Repository: jorgecarleitao/arrow2
Pull request #989 (2022-05-12 01:32:13)

Support for deserializing JSON from Utf8Array
This PR adds support for deserializing JSON for each row in a `Utf8Array`. It mimics the interface of `arrow2::io::json::read` and `arrow2::io::ndjson::read`, allowing you to apply the transformation on the full array. Currently `polars` supports JsonPath traversal of an underlying `Utf8Array` in [`DataFrame.str.json_path_match`](https://pola-rs.github.io/polars/py-polars/html/reference/api/polars.internals.expr.ExprStringNameSpace.json_path_match.html?highlight=json_path_match#polars.internals.expr.ExprStringNameSpace.json_path_match), but is not able to return a properly typed `Series` -- it only can handle the first element as a string type. This change opens support for getting the correct type. It also raises the question as to whether the reverse operations would be valuable. For now the scope of this PR is to support `Utf8Array => ArrayRef` of the appropriate underlying type based on the inferred or supplied schema.

In 7f00de16101845c7e1d897a248351bd6f2505384, I also added validity support for `StructArray` in the JSON deserializer. This is important to keep consistency on the null count.

I'm still getting familiar with Rust, so happy to hear any suggestions for improvements.

cjermain: I've fixed the previous tests by adding the correct validity. I also moved the tests to be more consistent with the rest of the package.
cjermain: @jorgecarleitao, @ritchie46, what do you think?
ritchie46: > @jorgecarleitao, @ritchie46, what do you think?

I am off for the weekend but cn review after. I like the feature so thanks for the PR @cjermain.
codecov[bot]: # [Codecov](https://codecov.io/gh/jorgecarleitao/arrow2/pull/989?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) Report
> Merging [#989](https://codecov.io/gh/jorgecarleitao/arrow2/pull/989?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) (6a8e2f6) into [main](https://codecov.io/gh/jorgecarleitao/arrow2/commit/b9aa8e8da7648559efdc78535085a509a281c55d?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) (b9aa8e8) will **increase** coverage by `0.02%`.
> The diff coverage is `92.00%`.

```diff
@@            Coverage Diff             @@
##             main     #989      +/-   ##
==========================================
+ Coverage   71.37%   71.39%   +0.02%     
==========================================
  Files         357      358       +1     
  Lines       19781    19805      +24     
==========================================
+ Hits        14118    14140      +22     
- Misses       5663     5665       +2     
```


| [Impacted Files](https://codecov.io/gh/jorgecarleitao/arrow2/pull/989?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) | Coverage Δ | |
|---|---|---|
| [src/array/utf8/mod.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/989/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2FycmF5L3V0ZjgvbW9kLnJz) | `70.65% <ø> (ø)` | |
| [src/array/utf8/json.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/989/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2FycmF5L3V0ZjgvanNvbi5ycw==) | `90.47% <90.47%> (ø)` | |
| [src/io/json/read/deserialize.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/989/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2lvL2pzb24vcmVhZC9kZXNlcmlhbGl6ZS5ycw==) | `68.18% <100.00%> (+0.63%)` | :arrow_up: |
| [src/compute/arithmetics/time.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/989/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2NvbXB1dGUvYXJpdGhtZXRpY3MvdGltZS5ycw==) | `25.68% <0.00%> (-0.92%)` | :arrow_down: |
| [src/bitmap/utils/slice\_iterator.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/989/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2JpdG1hcC91dGlscy9zbGljZV9pdGVyYXRvci5ycw==) | `87.93% <0.00%> (+1.72%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/jorgecarleitao/arrow2/pull/989?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/jorgecarleitao/arrow2/pull/989?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao). Last update [b9aa8e8...6a8e2f6](https://codecov.io/gh/jorgecarleitao/arrow2/pull/989?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao).

cjermain: Thanks @jorgecarleitao! Hope you are feeling better. I agree that moving it into `io::json::read` makes sense. There is a lot of overlap with `io::ndjson::read` (much of the code is based on that). I looked at adapting `io::ndjson::read::infer` to support iterators broadly, but was not sure how to have the `FallibleStreamingIterator` coexist with `Iterator` -- I'm still getting familiar with Rust. In C++ I would implement this with overloaded functions -- I'm not sure of the best solution in Rust. Do you have any suggestions for function signatures?
jorgecarleitao: So, this is what I was able to come up on the deserialization path (we need something similar for the inference path)

```rust
use std::sync::Arc;

use serde_json::Value;

use crate::array::Array;
use crate::datatypes::DataType;
use crate::error::ArrowError;

use super::super::super::json::read::_deserialize;

/// Deserializes rows into an [`Array`] of [`DataType`].
/// # Implementation
/// This function is CPU-bounded.
/// This function is guaranteed to return an array of length equal to `rows.len()`.
/// # Errors
/// This function errors iff any of the rows is not a valid JSON (i.e. the format is not valid NDJSON).
pub fn deserialize(rows: &[String], data_type: DataType) -> Result<Arc<dyn Array>, ArrowError> {
    deserialize_iter(rows.iter(), data_type)
}

/// Deserializes rows into an [`Array`] of [`DataType`].
/// # Implementation
/// This function is CPU-bounded.
/// This function is guaranteed to return an array of length equal to the leng
/// # Errors
/// This function errors iff any of the rows is not a valid JSON (i.e. the format is not valid NDJSON).
pub fn deserialize_iter<A: AsRef<str>>(
    rows: impl Iterator<Item = A>,
    data_type: DataType,
) -> Result<Arc<dyn Array>, ArrowError> {
    // deserialize strings to `Value`s
    let rows = rows
        .map(|row| serde_json::from_str(row.as_ref()).map_err(ArrowError::from))
        .collect::<Result<Vec<Value>, ArrowError>>()?;

    // deserialize &[Value] to Array
    Ok(_deserialize(&rows, data_type))
}


// ----------------------------
// test that it compiles for concrete Utf8Array ^_^; no need to expose it in the public API imo
// ----------------------------
fn des_array(
    array: &crate::array::Utf8Array<i32>,
    data_type: DataType,
) -> Result<Arc<dyn Array>, ArrowError> {
    deserialize_iter(array.iter().map(|x| x.unwrap_or("null")), data_type)
}

```

in other words, instead of requiring `&[String]`, which assumes a very concrete memory layout (a pointer of pointers to strings in C++), expose a more generic version based on a generic (template in C++) over an iterator, so that users can "templatise" it according to their needs.

----------------

On a separate note: a `StreamingIterator` is an iterator-like idiom where the item is only exposed to the consumers by reference. A common use-case we use it is in decompressing data. There we usually have `Fn(compressed: &[u8]) -> Vec<u8>`. When we need to decompress multiple items, this is an issue because we need to allocate one vec per item.

We can avoid multiple allocations of `Vec` by writing "something" where an item `compressed: &[u8]` is decompressed to an internal buffer `Vec<u8>` and we only expose a reference to the internal buffer. I.e. we need an iterator of `Option<&'a[u8]>`, where `'a` is "self". The signature of the `Iterator` trait does not support this atm in Rust (something solved by generic associated types / GAT). Thus, for now Rust de-facto uses a different trait for this (the `StreamingIterator`). `FallibleStreamingIterator` is the fallible version of `StreamingIterator`, which is a different trait because ... lack of GAT also does not allow to write a trait of `StreamingIterator` that supports `Result<'a ...>` xD

There are some aux methods to convert `Iterator -> StreamingIterator`; but I am not following - where do we use them here?
cjermain: Thanks, this gives me a clear idea of what to change. The `FallibleStreamingIterator` is used in the `FileReader` in `io::ndjson::read::infer`, which is used in a similar way to an iterator of `Option<&str>`.
https://github.com/jorgecarleitao/arrow2/blob/b9aa8e8da7648559efdc78535085a509a281c55d/src/io/ndjson/read/file.rs#L45-L58
https://github.com/jorgecarleitao/arrow2/blob/b9aa8e8da7648559efdc78535085a509a281c55d/src/io/ndjson/read/file.rs#L114-L118
cjermain: I've made the updates.

Patch:
diff --git a/src/array/primitive/mutable.rs b/src/array/primitive/mutable.rs
index d1fc8ae6ab..43f734f74b 100644
--- a/src/array/primitive/mutable.rs
+++ b/src/array/primitive/mutable.rs
@@ -234,12 +234,6 @@ impl<T: NativeType> MutablePrimitiveArray<T> {
         let len = self.len();
         if let Some(validity) = self.validity.as_mut() {
             validity.extend_constant(len - validity.len(), true);
-        } else {
-            let mut validity = MutableBitmap::new();
-            validity.extend_constant(len, true);
-            if validity.null_count() > 0 {
-                self.validity = Some(validity);
-            }
         }
     }
 
diff --git a/src/io/json/read/deserialize.rs b/src/io/json/read/deserialize.rs
index ee8068d081..8d6ddb160e 100644
--- a/src/io/json/read/deserialize.rs
+++ b/src/io/json/read/deserialize.rs
@@ -146,17 +146,21 @@ fn deserialize_struct<A: Borrow<Value>>(rows: &[A], data_type: DataType) -> Stru
         .map(|f| (&f.name, (f.data_type(), vec![])))
         .collect::<HashMap<_, _>>();
 
+    let mut validity = MutableBitmap::with_capacity(rows.len());
+
     rows.iter().for_each(|row| {
         match row.borrow() {
             Value::Object(value) => {
                 values
                     .iter_mut()
                     .for_each(|(s, (_, inner))| inner.push(value.get(*s).unwrap_or(&Value::Null)));
+                validity.push(true);
             }
             _ => {
                 values
                     .iter_mut()
                     .for_each(|(_, (_, inner))| inner.push(&Value::Null));
+                validity.push(false);
             }
         };
     });
@@ -166,7 +170,7 @@ fn deserialize_struct<A: Borrow<Value>>(rows: &[A], data_type: DataType) -> Stru
         .map(|(_, (data_type, values))| _deserialize(&values, data_type.clone()))
         .collect::<Vec<_>>();
 
-    StructArray::new(data_type, values, None)
+    StructArray::new(data_type, values, validity.into())
 }
 
 fn deserialize_dictionary<K: DictionaryKey, A: Borrow<Value>>(
diff --git a/src/io/json/write/serialize.rs b/src/io/json/write/serialize.rs
index 3b1f9e087b..fb90c17ad0 100644
--- a/src/io/json/write/serialize.rs
+++ b/src/io/json/write/serialize.rs
@@ -43,6 +43,29 @@ fn primitive_serializer<'a, T: NativeType + ToLexical>(
     ))
 }
 
+fn float_serializer<'a, T>(
+    array: &'a PrimitiveArray<T>,
+) -> Box<dyn StreamingIterator<Item = [u8]> + 'a + Send + Sync>
+where
+    T: num_traits::Float + NativeType + ToLexical,
+{
+    Box::new(BufStreamingIterator::new(
+        array.iter(),
+        |x, buf| {
+            if let Some(x) = x {
+                if T::is_nan(*x) {
+                    buf.extend(b"null")
+                } else {
+                    lexical_to_bytes_mut(*x, buf)
+                }
+            } else {
+                buf.extend(b"null")
+            }
+        },
+        vec![],
+    ))
+}
+
 fn utf8_serializer<'a, O: Offset>(
     array: &'a Utf8Array<O>,
 ) -> Box<dyn StreamingIterator<Item = [u8]> + 'a + Send + Sync> {
@@ -196,8 +219,8 @@ pub(crate) fn new_serializer<'a>(
         DataType::UInt16 => primitive_serializer::<u16>(array.as_any().downcast_ref().unwrap()),
         DataType::UInt32 => primitive_serializer::<u32>(array.as_any().downcast_ref().unwrap()),
         DataType::UInt64 => primitive_serializer::<u64>(array.as_any().downcast_ref().unwrap()),
-        DataType::Float32 => primitive_serializer::<f32>(array.as_any().downcast_ref().unwrap()),
-        DataType::Float64 => primitive_serializer::<f64>(array.as_any().downcast_ref().unwrap()),
+        DataType::Float32 => float_serializer::<f32>(array.as_any().downcast_ref().unwrap()),
+        DataType::Float64 => float_serializer::<f64>(array.as_any().downcast_ref().unwrap()),
         DataType::Utf8 => utf8_serializer::<i32>(array.as_any().downcast_ref().unwrap()),
         DataType::LargeUtf8 => utf8_serializer::<i64>(array.as_any().downcast_ref().unwrap()),
         DataType::Struct(_) => struct_serializer(array.as_any().downcast_ref().unwrap()),
diff --git a/src/io/ndjson/read/deserialize.rs b/src/io/ndjson/read/deserialize.rs
index e02927d024..5960b16498 100644
--- a/src/io/ndjson/read/deserialize.rs
+++ b/src/io/ndjson/read/deserialize.rs
@@ -22,9 +22,21 @@ pub fn deserialize(rows: &[String], data_type: DataType) -> Result<Arc<dyn Array
         ));
     }
 
+    deserialize_iter(rows.iter(), data_type)
+}
+
+/// Deserializes an iterator of rows into an [`Array`] of [`DataType`].
+/// # Implementation
+/// This function is CPU-bounded.
+/// This function is guaranteed to return an array of length equal to the leng
+/// # Errors
+/// This function errors iff any of the rows is not a valid JSON (i.e. the format is not valid NDJSON).
+pub fn deserialize_iter<A: AsRef<str>>(
+    rows: impl Iterator<Item = A>,
+    data_type: DataType,
+) -> Result<Arc<dyn Array>, ArrowError> {
     // deserialize strings to `Value`s
     let rows = rows
-        .iter()
         .map(|row| serde_json::from_str(row.as_ref()).map_err(ArrowError::from))
         .collect::<Result<Vec<Value>, ArrowError>>()?;
 
diff --git a/src/io/ndjson/read/file.rs b/src/io/ndjson/read/file.rs
index c45dc139e1..9bf6129543 100644
--- a/src/io/ndjson/read/file.rs
+++ b/src/io/ndjson/read/file.rs
@@ -125,3 +125,22 @@ pub fn infer<R: std::io::BufRead>(
     let v: Vec<&DataType> = data_types.iter().collect();
     Ok(coerce_data_type(&v))
 }
+
+/// Infers the [`DataType`] from an iterator of JSON strings. A limited number of
+/// rows can be used by passing `rows.take(number_of_rows)` as an input.
+///
+/// # Implementation
+/// This implementation infers each row by going through the entire iterator.
+pub fn infer_iter<A: AsRef<str>>(rows: impl Iterator<Item = A>) -> Result<DataType> {
+    let mut data_types = HashSet::new();
+    for row in rows {
+        let v: Value = serde_json::from_str(row.as_ref())?;
+        let data_type = infer_json(&v)?;
+        if data_type != DataType::Null {
+            data_types.insert(data_type);
+        }
+    }
+
+    let v: Vec<&DataType> = data_types.iter().collect();
+    Ok(coerce_data_type(&v))
+}
diff --git a/src/io/ndjson/read/mod.rs b/src/io/ndjson/read/mod.rs
index 5c52bd183f..6e7da6131b 100644
--- a/src/io/ndjson/read/mod.rs
+++ b/src/io/ndjson/read/mod.rs
@@ -4,5 +4,5 @@ pub use fallible_streaming_iterator::FallibleStreamingIterator;
 
 mod deserialize;
 mod file;
-pub use deserialize::deserialize;
-pub use file::{infer, FileReader};
+pub use deserialize::{deserialize, deserialize_iter};
+pub use file::{infer, infer_iter, FileReader};


Test patch:
diff --git a/tests/it/io/json/write.rs b/tests/it/io/json/write.rs
index 69e4bfa9d9..d79a5310f0 100644
--- a/tests/it/io/json/write.rs
+++ b/tests/it/io/json/write.rs
@@ -28,6 +28,24 @@ fn int32() -> Result<()> {
     test!(array, expected)
 }
 
+#[test]
+fn f32() -> Result<()> {
+    let array = Float32Array::from([Some(1.5), Some(2.5), Some(f32::NAN), None, Some(5.5)]);
+
+    let expected = r#"[1.5,2.5,null,null,5.5]"#;
+
+    test!(array, expected)
+}
+
+#[test]
+fn f64() -> Result<()> {
+    let array = Float64Array::from([Some(1.5), Some(2.5), Some(f64::NAN), None, Some(5.5)]);
+
+    let expected = r#"[1.5,2.5,null,null,5.5]"#;
+
+    test!(array, expected)
+}
+
 #[test]
 fn utf8() -> Result<()> {
     let array = Utf8Array::<i32>::from(&vec![Some("a"), Some("b"), Some("c"), Some("d"), None]);
diff --git a/tests/it/io/ndjson/mod.rs b/tests/it/io/ndjson/mod.rs
index 9f9cd0ab6b..1e25de7387 100644
--- a/tests/it/io/ndjson/mod.rs
+++ b/tests/it/io/ndjson/mod.rs
@@ -220,11 +220,19 @@ fn case_struct() -> (String, Arc<dyn Array>) {
 
     // build expected output
     let d = Utf8Array::<i32>::from(&vec![Some("text"), None, Some("text"), None]);
-    let c = StructArray::from_data(DataType::Struct(vec![d_field]), vec![Arc::new(d)], None);
+    let c = StructArray::from_data(
+        DataType::Struct(vec![d_field]),
+        vec![Arc::new(d)],
+        Some(Bitmap::from_u8_slice([0b11111101], 4)),
+    );
 
     let b = BooleanArray::from(vec![Some(true), Some(false), Some(true), None]);
     let inner = DataType::Struct(vec![Field::new("b", DataType::Boolean, true), c_field]);
-    let expected = StructArray::from_data(inner, vec![Arc::new(b), Arc::new(c)], None);
+    let expected = StructArray::from_data(
+        inner,
+        vec![Arc::new(b), Arc::new(c)],
+        Some(Bitmap::from_u8_slice([0b11110111], 4)),
+    );
 
     let data_type = DataType::Struct(fields);
 
@@ -268,7 +276,11 @@ fn case_nested_list() -> (String, Arc<dyn Array>) {
         None,
     ]);
 
-    let c = StructArray::from_data(DataType::Struct(vec![d_field]), vec![Arc::new(d)], None);
+    let c = StructArray::from_data(
+        DataType::Struct(vec![d_field]),
+        vec![Arc::new(d)],
+        Some(Bitmap::from_u8_slice([0b11111011], 6)),
+    );
 
     let b = BooleanArray::from(vec![
         Some(true),
diff --git a/tests/it/io/ndjson/read.rs b/tests/it/io/ndjson/read.rs
index 4b74248c56..5e19c1f351 100644
--- a/tests/it/io/ndjson/read.rs
+++ b/tests/it/io/ndjson/read.rs
@@ -273,3 +273,28 @@ fn skip_empty_lines() -> Result<()> {
     assert_eq!(3, arrays[0].len());
     Ok(())
 }
+
+#[test]
+fn utf8_array() -> Result<()> {
+    let array = Utf8Array::<i64>::from([
+        Some(r#"{"a": 1, "b": [{"c": 0}, {"c": 1}]}"#),
+        None,
+        Some(r#"{"a": 2, "b": [{"c": 2}, {"c": 5}]}"#),
+        None,
+    ]);
+    let data_type = ndjson_read::infer_iter(array.iter().map(|x| x.unwrap_or("null"))).unwrap();
+    let new_array =
+        ndjson_read::deserialize_iter(array.iter().map(|x| x.unwrap_or("null")), data_type)
+            .unwrap();
+
+    // Explicitly cast as StructArray
+    let new_array = new_array.as_any().downcast_ref::<StructArray>().unwrap();
+
+    assert_eq!(array.len(), new_array.len());
+    assert_eq!(array.null_count(), new_array.null_count());
+    assert_eq!(array.validity().unwrap(), new_array.validity().unwrap());
+
+    let field_names: Vec<String> = new_array.fields().iter().map(|f| f.name.clone()).collect();
+    assert_eq!(field_names, vec!["a".to_string(), "b".to_string()]);
+    Ok(())
+}