cj-zhukov commented on code in PR #26187: URL: https://github.com/apache/datafusion/pull/26187#discussion_r4410979213
########## 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: Agreed -- 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]
