diff --git a/datafusion/datasource-json/src/file_format.rs b/datafusion/datasource-json/src/file_format.rs index aef89560a3b5f..377beecf347fa 100644 --- a/datafusion/datasource-json/src/file_format.rs +++ b/datafusion/datasource-json/src/file_format.rs @@ -26,7 +26,7 @@ use std::sync::Arc; use crate::source::JsonSource; use arrow::array::RecordBatch; -use arrow::datatypes::{Schema, SchemaRef}; +use arrow::datatypes::{DataType, Schema, SchemaRef}; use arrow::error::ArrowError; use arrow::json; use arrow::json::reader::{ValueIter, infer_json_schema_from_iterator}; @@ -386,6 +386,40 @@ impl Default for JsonSerializer { } } +/// Returns true if `data_type` is, or contains at any depth, a `Map`. +fn contains_map(data_type: &DataType) -> bool { + match data_type { + DataType::Map(_, _) => true, + DataType::List(f) + | DataType::LargeList(f) + | DataType::FixedSizeList(f, _) + | DataType::ListView(f) + | DataType::LargeListView(f) => contains_map(f.data_type()), + DataType::Struct(fields) => fields.iter().any(|f| contains_map(f.data_type())), + DataType::Dictionary(_, v) => contains_map(v), + DataType::RunEndEncoded(_, v) => contains_map(v.data_type()), + _ => false, + } +} + +/// Creates a newline-delimited JSON writer for `schema`. +/// +/// By default arrow-json omits keys whose value is NULL. For a `Map` that drops +/// the entry itself, losing the key and changing the map's cardinality (see +/// ). So when the schema +/// contains a `Map` at any depth, NULLs are written explicitly. This is a +/// writer-wide setting, so every other NULL in the same output is written as +/// `null` too. Schemas without a `Map` keep the default. +pub(crate) fn ndjson_writer( + schema: &Schema, + sink: W, +) -> json::Writer { + let explicit_nulls = schema.fields().iter().any(|f| contains_map(f.data_type())); + json::writer::WriterBuilder::new() + .with_explicit_nulls(explicit_nulls) + .build::<_, json::writer::LineDelimited>(sink) +} + /// Define a struct for serializing Json records to a stream pub struct JsonSerializer {} @@ -399,7 +433,7 @@ impl JsonSerializer { impl BatchSerializer for JsonSerializer { fn serialize(&self, batch: RecordBatch, _initial: bool) -> Result { let mut buffer = Vec::with_capacity(4096); - let mut writer = json::LineDelimitedWriter::new(&mut buffer); + let mut writer = ndjson_writer(batch.schema_ref(), &mut buffer); writer.write(&batch)?; Ok(Bytes::from(buffer)) } @@ -658,3 +692,76 @@ impl From<&JsonFormatFactory> for datafusion_proto_models::protobuf::JsonOptions } } } + +#[cfg(test)] +mod tests { + use super::*; + + use arrow::array::{Array, ArrayRef, Int32Array, MapBuilder, StringBuilder}; + use arrow::datatypes::{Field, Fields}; + + fn map_type() -> DataType { + DataType::Map( + Arc::new(Field::new( + "entries", + DataType::Struct(Fields::from(vec![ + Field::new("keys", DataType::Utf8, false), + Field::new("values", DataType::Utf8, true), + ])), + false, + )), + false, + ) + } + + #[test] + fn contains_map_looks_through_nested_types() { + let map = map_type(); + let list_of_struct_of_map = DataType::new_list( + DataType::Struct(Fields::from(vec![Field::new("m", map.clone(), true)])), + true, + ); + let dictionary = + DataType::Dictionary(Box::new(DataType::Int32), Box::new(map.clone())); + let run_end_encoded = DataType::RunEndEncoded( + Arc::new(Field::new("run_ends", DataType::Int32, false)), + Arc::new(Field::new("values", map.clone(), true)), + ); + + assert!(contains_map(&map)); + assert!(contains_map(&list_of_struct_of_map)); + assert!(contains_map(&dictionary)); + assert!(contains_map(&run_end_encoded)); + assert!(!contains_map(&DataType::new_list(DataType::Utf8, true))); + } + + #[test] + fn serialize_writes_explicit_nulls_only_when_schema_has_map() -> Result<()> { + let ids: ArrayRef = Arc::new(Int32Array::from(vec![None::])); + + // Without a Map, NULLs are omitted as before. + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, true)])); + let batch = RecordBatch::try_new(schema, vec![Arc::clone(&ids)])?; + let bytes = JsonSerializer::new().serialize(batch, true)?; + assert_eq!(std::str::from_utf8(&bytes).unwrap(), "{}\n"); + + // With a Map, NULL map values and every other NULL are written explicitly. + let mut builder = + MapBuilder::new(None, StringBuilder::new(), StringBuilder::new()); + builder.keys().append_value("k"); + builder.values().append_null(); + builder.append(true)?; + let map = Arc::new(builder.finish()); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, true), + Field::new("m", map.data_type().clone(), true), + ])); + let batch = RecordBatch::try_new(schema, vec![ids, map])?; + let bytes = JsonSerializer::new().serialize(batch, true)?; + assert_eq!( + std::str::from_utf8(&bytes).unwrap(), + "{\"id\":null,\"m\":{\"k\":null}}\n" + ); + Ok(()) + } +} diff --git a/datafusion/datasource-json/src/source.rs b/datafusion/datasource-json/src/source.rs index f4c7021c86bdd..15579efdbb9c4 100644 --- a/datafusion/datasource-json/src/source.rs +++ b/datafusion/datasource-json/src/source.rs @@ -39,8 +39,8 @@ use datafusion_physical_plan::projection::ProjectionExprs; use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; use arrow::array::RecordBatch; +use arrow::datatypes::SchemaRef; use arrow::json::ReaderBuilder; -use arrow::{datatypes::SchemaRef, json}; use datafusion_datasource::file::FileSource; use datafusion_datasource::file_scan_config::FileScanConfig; use datafusion_execution::TaskContext; @@ -584,7 +584,8 @@ pub async fn plan_to_json( let mut buffer = Vec::with_capacity(1024); while let Some(batch) = stream.next().await.transpose()? { - let mut writer = json::LineDelimitedWriter::new(buffer); + let mut writer = + crate::file_format::ndjson_writer(batch.schema_ref(), buffer); writer.write(&batch)?; buffer = writer.into_inner(); buf_writer.write_all(&buffer).await?; diff --git a/datafusion/sqllogictest/test_files/copy.slt b/datafusion/sqllogictest/test_files/copy.slt index 6f62827412738..e3114c3c35d0d 100644 --- a/datafusion/sqllogictest/test_files/copy.slt +++ b/datafusion/sqllogictest/test_files/copy.slt @@ -815,3 +815,50 @@ physical_plan statement count 0 drop table t; + +# JSON output must keep MAP entries whose value is NULL (the entry is the data) +statement ok +create table map_null_src as +select 1 as id, map(['color', 'size', 'note'], ['red', NULL, NULL]) as attrs +union all +select 2 as id, map(['flag'], [NULL]) as attrs; + +query I +copy map_null_src to 'test_files/scratch/copy/map_null_values.json' stored as json; +---- +2 + +statement ok +create external table map_null_back (id bigint, attrs map(varchar, varchar)) +stored as json location 'test_files/scratch/copy/map_null_values.json'; + +query I?I +select id, attrs, cardinality(attrs) from map_null_back order by id; +---- +1 {color: red, size: NULL, note: NULL} 3 +2 {flag: NULL} 1 + +# MAP nested in a struct / list (no top-level MAP) must also keep NULL values +query I +copy (select named_struct('m', map(['k'], [NULL])) as s, [map(['k'], [NULL])] as l) +to 'test_files/scratch/copy/map_null_nested.json' stored as json; +---- +1 + +statement ok +create external table map_null_nested (s struct, l map(varchar, varchar)[]) +stored as json location 'test_files/scratch/copy/map_null_nested.json'; + +query ?I?I +select s, cardinality(s['m']), l, cardinality(l[1]) from map_null_nested; +---- +{m: {k: NULL}} 1 [{k: NULL}] 1 + +statement ok +drop table map_null_nested; + +statement ok +drop table map_null_back; + +statement ok +drop table map_null_src; diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md index e0f1aff1a7ca1..0dd0bec9d1c4e 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -673,6 +673,22 @@ should treat an empty list as the absent-key case instead. See [issue #24981](https://github.com/apache/datafusion/issues/24981) and [issue #24983](https://github.com/apache/datafusion/issues/24983) for details. +### JSON output writes explicit nulls when the schema contains a `Map` + +Writing JSON (for example `COPY ... STORED AS JSON` or `INSERT INTO` a JSON +table) previously dropped every `MAP` entry whose value is `NULL`, so +`{color: red, size: NULL}` was written as `{"color":"red"}`. These entries are +now kept and written as `{"color":"red","size":null}`. + +arrow-json's `explicit_nulls` is a writer-wide setting, so when the schema +contains a `Map` at any depth, every other `NULL` in the same file is also +written as `"field":null` instead of being omitted. Reading the file back with +an explicit schema gives the same result. With schema inference, a column or +struct field that is `NULL` in every sampled row is now inferred as type `Null` +instead of being missing. Files whose schema has no `Map` are unchanged. + +See [issue #26493](https://github.com/apache/datafusion/issues/26493) for details. + ### `date_bin` rejects strides that are not positive `date_bin` now returns an error when its stride is not greater than zero,