kosiew commented on code in PR #26187: URL: https://github.com/apache/datafusion/pull/26187#discussion_r4399586377
########## datafusion/spark/src/function/json/to_json.rs: ########## @@ -0,0 +1,617 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef, StringBuilder}; +use arrow::datatypes::{DataType, FieldRef}; +use arrow::json::writer::{EncoderOptions, make_encoder}; +use datafusion_common::utils::take_function_args; +use datafusion_common::{Column, Result, ScalarValue, exec_datafusion_err, plan_err}; +use datafusion_expr::{ + ColumnarValue, Expr, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, lit, +}; +use datafusion_functions::core::expr_fn::named_struct; + +/// Spark-compatible to_json expression +/// +/// <https://spark.apache.org/docs/latest/api/sql/index.html#to_json> +/// +/// Serializes a struct value to a JSON string, one JSON object per row. +/// +/// to_json(struct<...>) -> Utf8 +/// +/// Semantics (matching Spark's defaults): +/// - A NULL struct row produces NULL +/// - NULL fields are omitted from the object (`ignoreNullFields = true`) +/// - Nested structs, lists and maps are serialized recursively +/// +/// The serialization is done batch by batch with the arrow-json row encoder: +/// an encoder is built once per input array and each row is written straight +/// into a StringBuilder, so no intermediate RecordBatch or serde_json::Value +/// is created and nothing is collected. +/// +/// Not yet supported (Spark accepts these): map/array top-level inputs and the +/// options map argument (`ignoreNullFields`, timestampFormat, ...). + +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct ToJson { + signature: Signature, +} + +impl Default for ToJson { + fn default() -> Self { + Self::new() + } +} + +impl ToJson { + pub fn new() -> Self { + Self { + signature: Signature::any(1, Volatility::Immutable), + } + } +} + +impl ScalarUDFImpl for ToJson { + fn name(&self) -> &str { + "to_json" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> { + match arg_types { + [DataType::Struct(_)] => Ok(DataType::Utf8), + [other] => plan_err!("to_json expects a struct argument, got {other}"), + _ => plan_err!( + "to_json expects exactly one argument, got {}", + arg_types.len() + ), + } + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> { + let ScalarFunctionArgs { + args, arg_fields, .. + } = args; + let [input] = take_function_args(self.name(), args)?; + let [field] = take_function_args(self.name(), arg_fields)?; + + match input { + ColumnarValue::Array(array) => { + Ok(ColumnarValue::Array(struct_to_json(&field, &array)?)) + } + ColumnarValue::Scalar(scalar) => { + let array = scalar.to_array()?; + let result = struct_to_json(&field, &array)?; + Ok(ColumnarValue::Scalar(ScalarValue::try_from_array( + &result, 0, + )?)) + } + } + } +} + +/// Encodes every row of a struct array as a JSON object string. +/// +/// field describes array (its name and the nested field names/types are +/// what the encoder uses for the JSON keys). +fn struct_to_json(field: &FieldRef, array: &ArrayRef) -> Result<ArrayRef> { + // Spark default: ignoreNullFields = true -> omit null fields entirely + let options = EncoderOptions::default().with_explicit_nulls(false); + let mut encoder = make_encoder(field, array.as_ref(), &options)?; Review Comment: **[P2] Clarify supported nested map-key types.** `to_json(named_struct('m', map(make_array(1), make_array('v'))))` fails because Arrow's JSON writer requires UTF-8 map keys, while Spark returns `{"m":{"1":"v"}}`. Please either support integer keys recursively with a public-path regression or document the string-key restriction and test the unsupported-key error. Full map-key coverage and planning-time rejection are not required. ########## datafusion/spark/src/function/json/to_json.rs: ########## @@ -0,0 +1,617 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef, StringBuilder}; +use arrow::datatypes::{DataType, FieldRef}; +use arrow::json::writer::{EncoderOptions, make_encoder}; +use datafusion_common::utils::take_function_args; +use datafusion_common::{Column, Result, ScalarValue, exec_datafusion_err, plan_err}; +use datafusion_expr::{ + ColumnarValue, Expr, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, lit, +}; +use datafusion_functions::core::expr_fn::named_struct; + +/// Spark-compatible to_json expression +/// +/// <https://spark.apache.org/docs/latest/api/sql/index.html#to_json> +/// +/// Serializes a struct value to a JSON string, one JSON object per row. +/// +/// to_json(struct<...>) -> Utf8 +/// +/// Semantics (matching Spark's defaults): +/// - A NULL struct row produces NULL +/// - NULL fields are omitted from the object (`ignoreNullFields = true`) +/// - Nested structs, lists and maps are serialized recursively +/// +/// The serialization is done batch by batch with the arrow-json row encoder: +/// an encoder is built once per input array and each row is written straight +/// into a StringBuilder, so no intermediate RecordBatch or serde_json::Value +/// is created and nothing is collected. +/// +/// Not yet supported (Spark accepts these): map/array top-level inputs and the +/// options map argument (`ignoreNullFields`, timestampFormat, ...). + +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct ToJson { + signature: Signature, +} + +impl Default for ToJson { + fn default() -> Self { + Self::new() + } +} + +impl ToJson { + pub fn new() -> Self { + Self { + signature: Signature::any(1, Volatility::Immutable), + } + } +} + +impl ScalarUDFImpl for ToJson { + fn name(&self) -> &str { + "to_json" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> { + match arg_types { + [DataType::Struct(_)] => Ok(DataType::Utf8), + [other] => plan_err!("to_json expects a struct argument, got {other}"), + _ => plan_err!( + "to_json expects exactly one argument, got {}", + arg_types.len() + ), + } + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> { + let ScalarFunctionArgs { + args, arg_fields, .. + } = args; + let [input] = take_function_args(self.name(), args)?; + let [field] = take_function_args(self.name(), arg_fields)?; + + match input { + ColumnarValue::Array(array) => { + Ok(ColumnarValue::Array(struct_to_json(&field, &array)?)) + } + ColumnarValue::Scalar(scalar) => { + let array = scalar.to_array()?; + let result = struct_to_json(&field, &array)?; + Ok(ColumnarValue::Scalar(ScalarValue::try_from_array( + &result, 0, + )?)) + } + } + } +} + +/// Encodes every row of a struct array as a JSON object string. +/// +/// field describes array (its name and the nested field names/types are +/// what the encoder uses for the JSON keys). +fn struct_to_json(field: &FieldRef, array: &ArrayRef) -> Result<ArrayRef> { + // Spark default: ignoreNullFields = true -> omit null fields entirely + let options = EncoderOptions::default().with_explicit_nulls(false); Review Comment: **[P2] Preserve null-valued map entries.** `with_explicit_nulls(false)` also drops null map values: `to_json(named_struct('m', map(make_array('k'), make_array(CAST(NULL AS INT)))))` produces `{"m":{}}` instead of Spark's `{"m":{"k":null}}`. Please omit null struct fields while retaining null map entries, including nested maps, and add a regression covering both behaviors. ########## datafusion/spark/src/function/json/to_json.rs: ########## @@ -0,0 +1,617 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef, StringBuilder}; +use arrow::datatypes::{DataType, FieldRef}; +use arrow::json::writer::{EncoderOptions, make_encoder}; +use datafusion_common::utils::take_function_args; +use datafusion_common::{Column, Result, ScalarValue, exec_datafusion_err, plan_err}; +use datafusion_expr::{ + ColumnarValue, Expr, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, lit, +}; +use datafusion_functions::core::expr_fn::named_struct; + +/// Spark-compatible to_json expression +/// +/// <https://spark.apache.org/docs/latest/api/sql/index.html#to_json> +/// +/// Serializes a struct value to a JSON string, one JSON object per row. +/// +/// to_json(struct<...>) -> Utf8 +/// +/// Semantics (matching Spark's defaults): +/// - A NULL struct row produces NULL +/// - NULL fields are omitted from the object (`ignoreNullFields = true`) +/// - Nested structs, lists and maps are serialized recursively +/// +/// The serialization is done batch by batch with the arrow-json row encoder: +/// an encoder is built once per input array and each row is written straight +/// into a StringBuilder, so no intermediate RecordBatch or serde_json::Value +/// is created and nothing is collected. +/// +/// Not yet supported (Spark accepts these): map/array top-level inputs and the +/// options map argument (`ignoreNullFields`, timestampFormat, ...). + +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct ToJson { + signature: Signature, +} + +impl Default for ToJson { + fn default() -> Self { + Self::new() + } +} + +impl ToJson { + pub fn new() -> Self { + Self { + signature: Signature::any(1, Volatility::Immutable), + } + } +} + +impl ScalarUDFImpl for ToJson { + fn name(&self) -> &str { + "to_json" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> { + match arg_types { + [DataType::Struct(_)] => Ok(DataType::Utf8), + [other] => plan_err!("to_json expects a struct argument, got {other}"), + _ => plan_err!( + "to_json expects exactly one argument, got {}", + arg_types.len() + ), + } + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> { + let ScalarFunctionArgs { + args, arg_fields, .. + } = args; + let [input] = take_function_args(self.name(), args)?; + let [field] = take_function_args(self.name(), arg_fields)?; + + match input { + ColumnarValue::Array(array) => { + Ok(ColumnarValue::Array(struct_to_json(&field, &array)?)) + } + ColumnarValue::Scalar(scalar) => { + let array = scalar.to_array()?; + let result = struct_to_json(&field, &array)?; + Ok(ColumnarValue::Scalar(ScalarValue::try_from_array( + &result, 0, + )?)) + } + } + } +} + +/// Encodes every row of a struct array as a JSON object string. +/// +/// field describes array (its name and the nested field names/types are +/// what the encoder uses for the JSON keys). +fn struct_to_json(field: &FieldRef, array: &ArrayRef) -> Result<ArrayRef> { + // Spark default: ignoreNullFields = true -> omit null fields entirely + let options = EncoderOptions::default().with_explicit_nulls(false); + let mut encoder = make_encoder(field, array.as_ref(), &options)?; Review Comment: **[P2] Encode binary values as Base64 instead of hexadecimal.** `to_json(named_struct('b', arrow_cast('abc', 'Binary')))` currently produces `{"b":"616263"}`, whereas Spark produces `{"b":"YWJj"}`. Please override binary encoding recursively and add regressions for binary fields and binary values nested inside lists or structs. ########## datafusion/spark/src/function/json/to_json.rs: ########## @@ -0,0 +1,617 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef, StringBuilder}; +use arrow::datatypes::{DataType, FieldRef}; +use arrow::json::writer::{EncoderOptions, make_encoder}; +use datafusion_common::utils::take_function_args; +use datafusion_common::{Column, Result, ScalarValue, exec_datafusion_err, plan_err}; +use datafusion_expr::{ + ColumnarValue, Expr, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, lit, +}; +use datafusion_functions::core::expr_fn::named_struct; + +/// Spark-compatible to_json expression +/// +/// <https://spark.apache.org/docs/latest/api/sql/index.html#to_json> +/// +/// Serializes a struct value to a JSON string, one JSON object per row. +/// +/// to_json(struct<...>) -> Utf8 +/// +/// Semantics (matching Spark's defaults): +/// - A NULL struct row produces NULL +/// - NULL fields are omitted from the object (`ignoreNullFields = true`) +/// - Nested structs, lists and maps are serialized recursively +/// +/// The serialization is done batch by batch with the arrow-json row encoder: +/// an encoder is built once per input array and each row is written straight +/// into a StringBuilder, so no intermediate RecordBatch or serde_json::Value +/// is created and nothing is collected. +/// +/// Not yet supported (Spark accepts these): map/array top-level inputs and the +/// options map argument (`ignoreNullFields`, timestampFormat, ...). + +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct ToJson { + signature: Signature, +} + +impl Default for ToJson { + fn default() -> Self { + Self::new() + } +} + +impl ToJson { + pub fn new() -> Self { + Self { + signature: Signature::any(1, Volatility::Immutable), + } + } +} + +impl ScalarUDFImpl for ToJson { + fn name(&self) -> &str { + "to_json" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> { + match arg_types { + [DataType::Struct(_)] => Ok(DataType::Utf8), + [other] => plan_err!("to_json expects a struct argument, got {other}"), + _ => plan_err!( + "to_json expects exactly one argument, got {}", + arg_types.len() + ), + } + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> { + let ScalarFunctionArgs { + args, arg_fields, .. + } = args; + let [input] = take_function_args(self.name(), args)?; + let [field] = take_function_args(self.name(), arg_fields)?; + + match input { + ColumnarValue::Array(array) => { + Ok(ColumnarValue::Array(struct_to_json(&field, &array)?)) + } + ColumnarValue::Scalar(scalar) => { + let array = scalar.to_array()?; + let result = struct_to_json(&field, &array)?; + Ok(ColumnarValue::Scalar(ScalarValue::try_from_array( + &result, 0, + )?)) + } + } + } +} + +/// Encodes every row of a struct array as a JSON object string. +/// +/// field describes array (its name and the nested field names/types are +/// what the encoder uses for the JSON keys). +fn struct_to_json(field: &FieldRef, array: &ArrayRef) -> Result<ArrayRef> { + // Spark default: ignoreNullFields = true -> omit null fields entirely + let options = EncoderOptions::default().with_explicit_nulls(false); + let mut encoder = make_encoder(field, array.as_ref(), &options)?; + let mut builder = StringBuilder::with_capacity(array.len(), array.len() * 32); + let mut buf: Vec<u8> = Vec::with_capacity(128); + for row in 0..array.len() { + if encoder.is_null(row) { + builder.append_null(); + continue; + } + buf.clear(); + encoder.encode(row, &mut buf); + let json = std::str::from_utf8(&buf) + .map_err(|e| exec_datafusion_err!("to_json produced invalid UTF-8: {e}"))?; + builder.append_value(json); + } + + Ok(Arc::new(builder.finish())) +} + +/// Input accepted by the to_json expression builder. +/// +/// Implemented for an expression that evaluates to a struct, such as a struct column or a `named_struct` call, +/// which is used as-is; and for a slice or array of column names, which are packed into a struct first. +/// In the second form the JSON keys are the plain column names without any table qualifier. +pub trait IntoJsonStruct { + /// Converts self into the struct-valued [`Expr`] to serialize. + fn into_struct_expr(self) -> Expr; +} + +impl IntoJsonStruct for Expr { + fn into_struct_expr(self) -> Expr { + self + } +} + +impl<S: AsRef<str>> IntoJsonStruct for &[S] { + fn into_struct_expr(self) -> Expr { + let args = self + .iter() + .flat_map(|name| { + let column = Column::from_qualified_name(name.as_ref()); + [lit(column.name().to_owned()), Expr::Column(column)] + }) + .collect(); + named_struct(args) + } +} + +impl<S: AsRef<str>, const N: usize> IntoJsonStruct for &[S; N] { + fn into_struct_expr(self) -> Expr { + self.as_slice().into_struct_expr() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use arrow::array::{ + BooleanArray, Int32Array, Int64Array, ListBuilder, StringArray, StructArray, + }; + use arrow::buffer::NullBuffer; + use arrow::datatypes::{Field, Fields}; + use datafusion_common::cast::as_string_array; + use datafusion_common::config::ConfigOptions; + + /// Runs to_json over input (whose type is described by `field`) and + /// returns the resulting string array. + fn run(input: ColumnarValue, field: FieldRef) -> Result<ArrayRef> { + let number_rows = match &input { + ColumnarValue::Array(a) => a.len(), + ColumnarValue::Scalar(_) => 1, + }; + let args = ScalarFunctionArgs { + args: vec![input], + arg_fields: vec![field], + number_rows, + return_field: Field::new("to_json", DataType::Utf8, true).into(), + config_options: Arc::new(ConfigOptions::default()), + }; + ToJson::new().invoke_with_args(args)?.to_array(number_rows) + } + + fn sample_struct() -> StructArray { + let a = Arc::new(Int32Array::from(vec![Some(1), Some(2), Some(3)])) as ArrayRef; + let b = Arc::new(StringArray::from(vec![Some("x"), None, Some("z")])) as ArrayRef; + let c = + Arc::new(BooleanArray::from(vec![Some(true), Some(false), None])) as ArrayRef; + let fields = Fields::from(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Utf8, true), + Field::new("c", DataType::Boolean, true), + ]); + let nulls = NullBuffer::from(vec![true, true, false]); + StructArray::new(fields, vec![a, b, c], Some(nulls)) + } + + fn field_of(array: &StructArray) -> FieldRef { + Field::new("s", array.data_type().clone(), true).into() + } + + #[test] + fn test_return_type() { + let func = ToJson::new(); + let fields = Fields::from(vec![Field::new("a", DataType::Int32, true)]); + assert_eq!( + func.return_type(&[DataType::Struct(fields)]).unwrap(), + DataType::Utf8 + ); + } + + #[test] + fn test_return_type_rejects_non_struct() { + let err = ToJson::new().return_type(&[DataType::Utf8]).unwrap_err(); + assert!(err.to_string().contains("expects a struct argument")); + + let err = ToJson::new().return_type(&[]).unwrap_err(); + assert!(err.to_string().contains("exactly one argument")); + } + + #[test] + fn test_array_input() { + let input = sample_struct(); + let field = field_of(&input); + let out = run(ColumnarValue::Array(Arc::new(input)), field).unwrap(); + let out = as_string_array(&out).unwrap(); + + assert_eq!(out.len(), 3); + assert_eq!(out.value(0), r#"{"a":1,"b":"x","c":true}"#); + // NULL field is omitted (Spark ignoreNullFields = true) + assert_eq!(out.value(1), r#"{"a":2,"c":false}"#); + // NULL struct row -> NULL + assert!(out.is_null(2)); + } + + #[test] + fn test_scalar_input() { + let input = sample_struct(); + let field = field_of(&input); + let scalar = ScalarValue::Struct(Arc::new(input.slice(0, 1))); + let out = run(ColumnarValue::Scalar(scalar), field).unwrap(); + let out = as_string_array(&out).unwrap(); + + assert_eq!(out.len(), 1); + assert_eq!(out.value(0), r#"{"a":1,"b":"x","c":true}"#); + } + + #[test] + fn test_scalar_null_struct() { + let input = sample_struct(); + let field = field_of(&input); + let scalar = ScalarValue::Struct(Arc::new(input.slice(2, 1))); + let out = run(ColumnarValue::Scalar(scalar), field).unwrap(); + + assert_eq!(out.len(), 1); + assert!(out.is_null(0)); + } + + #[test] + fn test_nested_struct_and_list() { Review Comment: **Optional:** Consider extending the nested-struct/list test with a nonzero-offset slice of a multirow struct array containing nontrivial list offsets and a null parent row. Asserting both JSON values and output nullness would protect the row-index/child-offset boundary against future encoder or Arrow changes; no current slicing defect is established. ########## datafusion/spark/src/function/json/mod.rs: ########## @@ -16,23 +16,38 @@ // under the License. pub mod json_tuple; +pub mod to_json; use datafusion_expr::ScalarUDF; use datafusion_functions::make_udf_function; use std::sync::Arc; make_udf_function!(json_tuple::JsonTuple, json_tuple); +make_udf_function!(to_json::ToJson, to_json); pub mod expr_fn { + use datafusion_expr::Expr; use datafusion_functions::export_functions; + use crate::function::json::to_json::IntoJsonStruct; + export_functions!(( json_tuple, "Extracts top-level fields from a JSON string and returns them as a struct.", args, )); + + /// Serializes a struct value to a JSON string (Spark-compatible `to_json`). Review Comment: **Optional:** Could we add a runnable Rustdoc example showing both `to_json(&["foo", "bar"])` and `to_json(col("event"))`, including the import and expected JSON? This would make the two supported public helper forms easier to discover and use. ########## datafusion/spark/src/function/json/to_json.rs: ########## @@ -0,0 +1,617 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef, StringBuilder}; +use arrow::datatypes::{DataType, FieldRef}; +use arrow::json::writer::{EncoderOptions, make_encoder}; +use datafusion_common::utils::take_function_args; +use datafusion_common::{Column, Result, ScalarValue, exec_datafusion_err, plan_err}; +use datafusion_expr::{ + ColumnarValue, Expr, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, lit, +}; +use datafusion_functions::core::expr_fn::named_struct; + +/// Spark-compatible to_json expression +/// +/// <https://spark.apache.org/docs/latest/api/sql/index.html#to_json> +/// +/// Serializes a struct value to a JSON string, one JSON object per row. +/// +/// to_json(struct<...>) -> Utf8 +/// +/// Semantics (matching Spark's defaults): +/// - A NULL struct row produces NULL +/// - NULL fields are omitted from the object (`ignoreNullFields = true`) +/// - Nested structs, lists and maps are serialized recursively +/// +/// The serialization is done batch by batch with the arrow-json row encoder: +/// an encoder is built once per input array and each row is written straight +/// into a StringBuilder, so no intermediate RecordBatch or serde_json::Value +/// is created and nothing is collected. +/// +/// Not yet supported (Spark accepts these): map/array top-level inputs and the +/// options map argument (`ignoreNullFields`, timestampFormat, ...). + +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct ToJson { + signature: Signature, +} + +impl Default for ToJson { + fn default() -> Self { + Self::new() + } +} + +impl ToJson { + pub fn new() -> Self { + Self { + signature: Signature::any(1, Volatility::Immutable), + } + } +} + +impl ScalarUDFImpl for ToJson { + fn name(&self) -> &str { + "to_json" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> { + match arg_types { + [DataType::Struct(_)] => Ok(DataType::Utf8), + [other] => plan_err!("to_json expects a struct argument, got {other}"), + _ => plan_err!( + "to_json expects exactly one argument, got {}", + arg_types.len() + ), + } + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> { + let ScalarFunctionArgs { + args, arg_fields, .. + } = args; + let [input] = take_function_args(self.name(), args)?; + let [field] = take_function_args(self.name(), arg_fields)?; + + match input { + ColumnarValue::Array(array) => { + Ok(ColumnarValue::Array(struct_to_json(&field, &array)?)) + } + ColumnarValue::Scalar(scalar) => { + let array = scalar.to_array()?; + let result = struct_to_json(&field, &array)?; + Ok(ColumnarValue::Scalar(ScalarValue::try_from_array( + &result, 0, + )?)) + } + } + } +} + +/// Encodes every row of a struct array as a JSON object string. +/// +/// field describes array (its name and the nested field names/types are +/// what the encoder uses for the JSON keys). +fn struct_to_json(field: &FieldRef, array: &ArrayRef) -> Result<ArrayRef> { + // Spark default: ignoreNullFields = true -> omit null fields entirely + let options = EncoderOptions::default().with_explicit_nulls(false); + let mut encoder = make_encoder(field, array.as_ref(), &options)?; Review Comment: **[P2] Preserve non-finite floating-point values.** Arrow encodes NaN and infinities as JSON null, so `to_json(named_struct('x', CAST('NaN' AS DOUBLE)))` produces `{"x":null}` instead of Spark's `{"x":"NaN"}`. Please implement Spark-compatible Float32/Float64 encoding and test NaN, positive and negative infinity, including nested values. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
