in arrow-json/src/writer/mod.rs [1422:1648]
fn test_writer_explicit_nulls() -> Result<(), ArrowError> {
fn nested_list() -> (Arc<ListArray>, Arc<Field>) {
let array = Arc::new(ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
Some(vec![None, None, None]),
Some(vec![Some(1), Some(2), Some(3)]),
None,
Some(vec![None, None, None]),
]));
let field = Arc::new(Field::new("list", array.data_type().clone(), true));
// [{"list":[null,null,null]},{"list":[1,2,3]},{"list":null},{"list":[null,null,null]}]
(array, field)
}
fn nested_dict() -> (Arc<DictionaryArray<Int32Type>>, Arc<Field>) {
let array = Arc::new(DictionaryArray::from_iter(vec![
Some("cupcakes"),
None,
Some("bear"),
Some("kuma"),
]));
let field = Arc::new(Field::new("dict", array.data_type().clone(), true));
// [{"dict":"cupcakes"},{"dict":null},{"dict":"bear"},{"dict":"kuma"}]
(array, field)
}
fn nested_map() -> (Arc<MapArray>, Arc<Field>) {
let string_builder = StringBuilder::new();
let int_builder = Int64Builder::new();
let mut builder = MapBuilder::new(None, string_builder, int_builder);
// [{"foo": 10}, null, {}, {"bar": 20, "baz": 30, "qux": 40}]
builder.keys().append_value("foo");
builder.values().append_value(10);
builder.append(true).unwrap();
builder.append(false).unwrap();
builder.append(true).unwrap();
builder.keys().append_value("bar");
builder.values().append_value(20);
builder.keys().append_value("baz");
builder.values().append_value(30);
builder.keys().append_value("qux");
builder.values().append_value(40);
builder.append(true).unwrap();
let array = Arc::new(builder.finish());
let field = Arc::new(Field::new("map", array.data_type().clone(), true));
(array, field)
}
fn root_list() -> (Arc<ListArray>, Field) {
let struct_array = StructArray::from(vec![
(
Arc::new(Field::new("utf8", DataType::Utf8, true)),
Arc::new(StringArray::from(vec![Some("a"), Some("b"), None, None])) as ArrayRef,
),
(
Arc::new(Field::new("int32", DataType::Int32, true)),
Arc::new(Int32Array::from(vec![Some(1), None, Some(5), None])) as ArrayRef,
),
]);
let field = Field::new_list(
"list",
Field::new("struct", struct_array.data_type().clone(), true),
true,
);
// [{"list":[{"int32":1,"utf8":"a"},{"int32":null,"utf8":"b"}]},{"list":null},{"list":[{int32":5,"utf8":null}]},{"list":null}]
let entry_offsets = Buffer::from([0, 2, 2, 3, 3].to_byte_slice());
let data = ArrayData::builder(field.data_type().clone())
.len(4)
.add_buffer(entry_offsets)
.add_child_data(struct_array.into_data())
.null_bit_buffer(Some([0b00000101].into()))
.build()
.unwrap();
let array = Arc::new(ListArray::from(data));
(array, field)
}
let (nested_list_array, nested_list_field) = nested_list();
let (nested_dict_array, nested_dict_field) = nested_dict();
let (nested_map_array, nested_map_field) = nested_map();
let (root_list_array, root_list_field) = root_list();
let schema = Schema::new(vec![
Field::new("date", DataType::Date32, true),
Field::new("null", DataType::Null, true),
Field::new_struct(
"struct",
vec![
Arc::new(Field::new("utf8", DataType::Utf8, true)),
nested_list_field.clone(),
nested_dict_field.clone(),
nested_map_field.clone(),
],
true,
),
root_list_field,
]);
let arr_date32 = Date32Array::from(vec![Some(0), None, Some(1), None]);
let arr_null = NullArray::new(4);
let arr_struct = StructArray::from(vec![
// [{"utf8":"a"},{"utf8":null},{"utf8":null},{"utf8":"b"}]
(
Arc::new(Field::new("utf8", DataType::Utf8, true)),
Arc::new(StringArray::from(vec![Some("a"), None, None, Some("b")])) as ArrayRef,
),
// [{"list":[null,null,null]},{"list":[1,2,3]},{"list":null},{"list":[null,null,null]}]
(nested_list_field, nested_list_array as ArrayRef),
// [{"dict":"cupcakes"},{"dict":null},{"dict":"bear"},{"dict":"kuma"}]
(nested_dict_field, nested_dict_array as ArrayRef),
// [{"foo": 10}, null, {}, {"bar": 20, "baz": 30, "qux": 40}]
(nested_map_field, nested_map_array as ArrayRef),
]);
let batch = RecordBatch::try_new(
Arc::new(schema),
vec![
// [{"date":"1970-01-01"},{"date":null},{"date":"1970-01-02"},{"date":null}]
Arc::new(arr_date32),
// [{"null":null},{"null":null},{"null":null},{"null":null}]
Arc::new(arr_null),
Arc::new(arr_struct),
// [{"list":[{"int32":1,"utf8":"a"},{"int32":null,"utf8":"b"}]},{"list":null},{"list":[{int32":5,"utf8":null}]},{"list":null}]
root_list_array,
],
)?;
let mut buf = Vec::new();
{
let mut writer = WriterBuilder::new()
.with_explicit_nulls(true)
.build::<_, JsonArray>(&mut buf);
writer.write_batches(&[&batch])?;
writer.finish()?;
}
let actual = serde_json::from_slice::<Vec<Value>>(&buf).unwrap();
let expected = serde_json::from_value::<Vec<Value>>(json!([
{
"date": "1970-01-01",
"list": [
{
"int32": 1,
"utf8": "a"
},
{
"int32": null,
"utf8": "b"
}
],
"null": null,
"struct": {
"dict": "cupcakes",
"list": [
null,
null,
null
],
"map": {
"foo": 10
},
"utf8": "a"
}
},
{
"date": null,
"list": null,
"null": null,
"struct": {
"dict": null,
"list": [
1,
2,
3
],
"map": null,
"utf8": null
}
},
{
"date": "1970-01-02",
"list": [
{
"int32": 5,
"utf8": null
}
],
"null": null,
"struct": {
"dict": "bear",
"list": null,
"map": {},
"utf8": null
}
},
{
"date": null,
"list": null,
"null": null,
"struct": {
"dict": "kuma",
"list": [
null,
null,
null
],
"map": {
"bar": 20,
"baz": 30,
"qux": 40
},
"utf8": "b"
}
}
]))
.unwrap();
assert_eq!(actual, expected);
Ok(())
}