native/spark-expr/src/json_funcs/to_json.rs

// 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.

// TODO upstream this to DataFusion as long as we have a way to specify all
// of the Spark-specific compatibility features that we need (including
// being able to specify Spark-compatible cast from all types to string)

use crate::SparkCastOptions;
use crate::{spark_cast, EvalMode};
use arrow::array::builder::StringBuilder;
use arrow::array::{Array, ArrayRef, RecordBatch, StringArray, StructArray};
use arrow::datatypes::{DataType, Schema};
use datafusion::common::Result;
use datafusion::physical_expr::PhysicalExpr;
use datafusion::physical_plan::ColumnarValue;
use std::any::Any;
use std::borrow::Cow;
use std::fmt::{Debug, Display, Formatter};
use std::hash::Hash;
use std::sync::Arc;

/// to_json function
#[derive(Debug, Eq)]
pub struct ToJson {
    /// The input to convert to JSON
    expr: Arc<dyn PhysicalExpr>,
    /// Timezone to use when converting timestamps to JSON
    timezone: String,
    ignore_null_fields: bool,
}

impl Hash for ToJson {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.expr.hash(state);
        self.timezone.hash(state);
        self.ignore_null_fields.hash(state);
    }
}
impl PartialEq for ToJson {
    fn eq(&self, other: &Self) -> bool {
        self.expr.eq(&other.expr)
            && self.timezone.eq(&other.timezone)
            && self.ignore_null_fields.eq(&other.ignore_null_fields)
    }
}

impl ToJson {
    pub fn new(expr: Arc<dyn PhysicalExpr>, timezone: &str, ignore_null_fields: bool) -> Self {
        Self {
            expr,
            timezone: timezone.to_owned(),
            ignore_null_fields,
        }
    }
}

impl Display for ToJson {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "to_json({}, timezone={}, ignore_null_fields={})",
            self.expr, self.timezone, self.ignore_null_fields
        )
    }
}

impl PartialEq<dyn Any> for ToJson {
    fn eq(&self, other: &dyn Any) -> bool {
        if let Some(other) = other.downcast_ref::<ToJson>() {
            self == other
        } else {
            false
        }
    }
}

impl PhysicalExpr for ToJson {
    fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        Display::fmt(self, f)
    }

    fn data_type(&self, _: &Schema) -> Result<DataType> {
        Ok(DataType::Utf8)
    }

    fn nullable(&self, input_schema: &Schema) -> Result<bool> {
        self.expr.nullable(input_schema)
    }

    fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
        let input = self.expr.evaluate(batch)?.into_array(batch.num_rows())?;
        Ok(ColumnarValue::Array(array_to_json_string(
            &input,
            &self.timezone,
            self.ignore_null_fields,
        )?))
    }

    fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
        vec![&self.expr]
    }

    fn with_new_children(
        self: Arc<Self>,
        children: Vec<Arc<dyn PhysicalExpr>>,
    ) -> Result<Arc<dyn PhysicalExpr>> {
        assert!(children.len() == 1);
        Ok(Arc::new(Self::new(
            Arc::clone(&children[0]),
            &self.timezone,
            self.ignore_null_fields,
        )))
    }
}

/// Convert an array into a JSON value string representation
fn array_to_json_string(
    arr: &Arc<dyn Array>,
    timezone: &str,
    ignore_null_fields: bool,
) -> Result<ArrayRef> {
    if let Some(struct_array) = arr.as_any().downcast_ref::<StructArray>() {
        struct_to_json(struct_array, timezone, ignore_null_fields)
    } else {
        spark_cast(
            ColumnarValue::Array(Arc::clone(arr)),
            &DataType::Utf8,
            &SparkCastOptions::new(EvalMode::Legacy, timezone, false),
        )?
        .into_array(arr.len())
    }
}

/// Returns the JSON escape sequence for a byte that requires escaping, or `None`
/// otherwise. Every escaped character is ASCII, so scanning the input as bytes is
/// safe: a UTF-8 continuation or lead byte (>= 0x80) can never match here, so run
/// boundaries always fall on `char` boundaries.
#[inline]
fn escape_replacement(b: u8) -> Option<&'static str> {
    match b {
        b'"' => Some("\\\""),
        b'\\' => Some("\\\\"),
        b'\t' => Some("\\t"),
        b'\r' => Some("\\r"),
        b'\n' => Some("\\n"),
        0x0C => Some("\\f"),
        0x08 => Some("\\b"),
        _ => None,
    }
}

/// Escapes the characters that Spark's JSON writer escapes (`"`, `\`, and the
/// `\t`/`\r`/`\n`/`\f`/`\b` control characters). Returns the input unchanged
/// (without allocating) when nothing needs escaping, which is the common case,
/// and otherwise copies unescaped runs in bulk rather than character by character.
fn escape_string(input: &str) -> Cow<'_, str> {
    let bytes = input.as_bytes();
    let first = match bytes.iter().position(|&b| escape_replacement(b).is_some()) {
        None => return Cow::Borrowed(input),
        Some(pos) => pos,
    };

    let mut escaped = String::with_capacity(input.len() + 8);
    let mut run_start = 0;
    for i in first..bytes.len() {
        if let Some(replacement) = escape_replacement(bytes[i]) {
            escaped.push_str(&input[run_start..i]);
            escaped.push_str(replacement);
            run_start = i + 1;
        }
    }
    escaped.push_str(&input[run_start..]);
    Cow::Owned(escaped)
}

fn struct_to_json(
    array: &StructArray,
    timezone: &str,
    ignore_null_fields: bool,
) -> Result<ArrayRef> {
    // get field names and escape any quotes
    let field_names: Vec<String> = array
        .fields()
        .iter()
        .map(|f| escape_string(f.name().as_str()).into_owned())
        .collect();
    // determine which fields need to have their values quoted
    let is_string: Vec<bool> = array
        .fields()
        .iter()
        .map(|f| match f.data_type() {
            DataType::Utf8 | DataType::LargeUtf8 => true,
            DataType::Dictionary(_, dt) => {
                matches!(dt.as_ref(), DataType::Utf8 | DataType::LargeUtf8)
            }
            _ => false,
        })
        .collect();
    // create JSON string representation of each column
    let string_arrays: Vec<ArrayRef> = array
        .columns()
        .iter()
        .map(|arr| array_to_json_string(arr, timezone, ignore_null_fields))
        .collect::<Result<Vec<_>>>()?;
    let string_arrays: Vec<&StringArray> = string_arrays
        .iter()
        .map(|arr| {
            arr.as_any()
                .downcast_ref::<StringArray>()
                .expect("string array")
        })
        .collect();
    // build the JSON string containing entries in the format `"field_name":field_value`
    let mut builder = StringBuilder::with_capacity(array.len(), array.len() * 16);
    let mut json = String::with_capacity(array.len() * 16);
    for row_index in 0..array.len() {
        if array.is_null(row_index) {
            builder.append_null();
        } else {
            json.clear();
            let mut any_fields_written = false;
            json.push('{');
            for col_index in 0..string_arrays.len() {
                let is_null = string_arrays[col_index].is_null(row_index);
                if is_null && ignore_null_fields {
                    continue;
                }
                if any_fields_written {
                    json.push(',');
                }
                // quoted field name
                json.push('"');
                json.push_str(&field_names[col_index]);
                json.push_str("\":");
                if is_null {
                    json.push_str("null");
                } else {
                    // value
                    let string_value = string_arrays[col_index].value(row_index);
                    if is_string[col_index] || is_infinity(string_value) || is_nan(string_value) {
                        json.push('"');
                        json.push_str(escape_string(string_value).as_ref());
                        json.push('"');
                    } else {
                        json.push_str(string_value);
                    }
                }
                any_fields_written = true;
            }
            json.push('}');
            builder.append_value(&json);
        }
    }
    Ok(Arc::new(builder.finish()))
}

fn is_infinity(input: &str) -> bool {
    input == "Infinity" || input == "-Infinity"
}

fn is_nan(input: &str) -> bool {
    input == "NaN"
}

#[cfg(test)]
mod test {
    use crate::json_funcs::to_json::{struct_to_json, ToJson};
    use arrow::array::types::Int32Type;
    use arrow::array::{Array, PrimitiveArray, StringArray};
    use arrow::array::{ArrayRef, BooleanArray, Int32Array, StructArray};
    use arrow::datatypes::{DataType, Field};
    use datafusion::common::Result;
    use datafusion::physical_plan::expressions::Column;
    use std::any::Any;
    use std::sync::Arc;

    #[test]
    fn test_primitives() -> Result<()> {
        let bools: ArrayRef = create_bools();
        let ints: ArrayRef = create_ints();
        let strings: ArrayRef = create_strings();
        let struct_array = StructArray::from(vec![
            (Arc::new(Field::new("a", DataType::Boolean, true)), bools),
            (Arc::new(Field::new("b", DataType::Int32, true)), ints),
            (Arc::new(Field::new("c", DataType::Utf8, true)), strings),
        ]);
        let json = struct_to_json(&struct_array, "UTC", true)?;
        let json = json
            .as_any()
            .downcast_ref::<StringArray>()
            .expect("string array");
        assert_eq!(4, json.len());
        assert_eq!(r#"{"b":123}"#, json.value(0));
        assert_eq!(r#"{"a":true,"c":"foo"}"#, json.value(1));
        assert_eq!(r#"{"a":false,"b":2147483647,"c":"bar"}"#, json.value(2));
        assert_eq!(r#"{"a":false,"b":-2147483648,"c":""}"#, json.value(3));
        Ok(())
    }

    #[test]
    fn test_nested_struct() -> Result<()> {
        let bools: ArrayRef = create_bools();
        let ints: ArrayRef = create_ints();

        // create first struct array
        let struct_fields = vec![
            Arc::new(Field::new("a", DataType::Boolean, true)),
            Arc::new(Field::new("b", DataType::Int32, true)),
        ];
        let struct_values = vec![bools, ints];
        let struct_array = StructArray::from(
            struct_fields
                .clone()
                .into_iter()
                .zip(struct_values)
                .collect::<Vec<_>>(),
        );

        // create second struct array containing the first struct array
        let struct_fields2 = vec![Arc::new(Field::new(
            "a",
            DataType::Struct(struct_fields.into()),
            true,
        ))];
        let struct_values2: Vec<ArrayRef> = vec![Arc::new(struct_array.clone())];
        let struct_array2 = StructArray::from(
            struct_fields2
                .into_iter()
                .zip(struct_values2)
                .collect::<Vec<_>>(),
        );

        let json = struct_to_json(&struct_array2, "UTC", true)?;
        let json = json
            .as_any()
            .downcast_ref::<StringArray>()
            .expect("string array");
        assert_eq!(4, json.len());
        assert_eq!(r#"{"a":{"b":123}}"#, json.value(0));
        assert_eq!(r#"{"a":{"a":true}}"#, json.value(1));
        assert_eq!(r#"{"a":{"a":false,"b":2147483647}}"#, json.value(2));
        assert_eq!(r#"{"a":{"a":false,"b":-2147483648}}"#, json.value(3));
        Ok(())
    }

    fn create_ints() -> Arc<PrimitiveArray<Int32Type>> {
        Arc::new(Int32Array::from(vec![
            Some(123),
            None,
            Some(i32::MAX),
            Some(i32::MIN),
        ]))
    }

    fn create_bools() -> Arc<BooleanArray> {
        Arc::new(BooleanArray::from(vec![
            None,
            Some(true),
            Some(false),
            Some(false),
        ]))
    }

    fn create_strings() -> Arc<StringArray> {
        Arc::new(StringArray::from(vec![
            None,
            Some("foo"),
            Some("bar"),
            Some(""),
        ]))
    }

    fn make_to_json(timezone: &str, ignore_null_fields: bool) -> ToJson {
        ToJson::new(Arc::new(Column::new("x", 0)), timezone, ignore_null_fields)
    }

    #[test]
    fn test_partial_eq_same() {
        let a = make_to_json("UTC", true);
        let b = make_to_json("UTC", true);
        assert_eq!(a, b);
        assert!(<ToJson as PartialEq<dyn Any>>::eq(&a, &b as &dyn Any));
    }

    #[test]
    fn test_partial_eq_dyn_any_differs_on_timezone() {
        let a = make_to_json("UTC", true);
        let b = make_to_json("America/New_York", true);
        assert_ne!(a, b);
        assert!(!<ToJson as PartialEq<dyn Any>>::eq(&a, &b as &dyn Any));
    }

    #[test]
    fn test_partial_eq_dyn_any_differs_on_ignore_null_fields() {
        let a = make_to_json("UTC", true);
        let b = make_to_json("UTC", false);
        assert_ne!(a, b);
        assert!(!<ToJson as PartialEq<dyn Any>>::eq(&a, &b as &dyn Any));
    }
}