Enable PPL JSON scalar functions on the analytics-engine route

sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ScalarFunction.java

@@ -168,7 +168,16 @@ public enum ScalarFunction {
     CURRENT_TIME(Category.SCALAR, SqlKind.OTHER_FUNCTION),
     CURTIME(Category.SCALAR, SqlKind.OTHER_FUNCTION),
     CONVERT_TZ(Category.SCALAR, SqlKind.OTHER_FUNCTION),
-    UNIX_TIMESTAMP(Category.SCALAR, SqlKind.OTHER_FUNCTION);
+    UNIX_TIMESTAMP(Category.SCALAR, SqlKind.OTHER_FUNCTION),
+
+    // ── JSON ────────────────────────────────────────────────────────
+    JSON_APPEND(Category.SCALAR, SqlKind.OTHER_FUNCTION),
+    JSON_ARRAY_LENGTH(Category.SCALAR, SqlKind.OTHER_FUNCTION),
+    JSON_DELETE(Category.SCALAR, SqlKind.OTHER_FUNCTION),
+    JSON_EXTEND(Category.SCALAR, SqlKind.OTHER_FUNCTION),
+    JSON_EXTRACT(Category.SCALAR, SqlKind.OTHER_FUNCTION),
+    JSON_KEYS(Category.SCALAR, SqlKind.OTHER_FUNCTION),
+    JSON_SET(Category.SCALAR, SqlKind.OTHER_FUNCTION);
 
     /**
      * Category of scalar function.

sandbox/plugins/analytics-backend-datafusion/rust/Cargo.toml

@@ -49,6 +49,26 @@ chrono-tz = "0.10"
 
 tokio-metrics = { workspace = true }
 
+# serde_json `preserve_order` — backs `Map<String,Value>` with `IndexMap`
+# instead of `BTreeMap` so json_keys / mutation UDFs see object keys in
+# insertion order (parity with legacy SQL-plugin's LinkedHashMap; required by
+# `testJsonKeysParityWithLegacy` + byte-for-byte json_extract fixtures).
+# Cargo's feature unification propagates this to every workspace member that
+# pulls in serde_json. Audit (2026-05-07): the five other consumers
+# (parquet-data-format, native-repository-{s3,gcs,azure,fs}) only call
+# `serde_json::from_str` into typed config structs, whose field layout is
+# fixed at the type level — `preserve_order` is inert for them, so the
+# feature is additive with no observable blast radius outside this crate.
+serde_json = { workspace = true, features = ["preserve_order"] }
+# jsonpath-rust 0.7 — JSONPath evaluator for json_extract. Published at
+# https://github.com/besok/jsonpath-rust (crates.io). We pin `0.7` (latest
+# `0.7.5`) rather than tracking the newer `1.0` release line because 0.7's
+# `JsonPathValue` enum exposes the Found/NoValue distinction json_extract
+# relies on to render missing-path matches as literal `null` elements in the
+# multi-path JSON-array output. Moving to 1.x is a follow-up once we can
+# reproduce that distinction against the new API surface.
+jsonpath-rust = "0.7"
+
 [dev-dependencies]
 criterion = { workspace = true }
 tempfile = { workspace = true }

sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_append.rs

@@ -0,0 +1,335 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+//! `json_append(value, path1, val1, [path2, val2, ...])` — push `valN` onto
+//! each path-matched array (parity with legacy `JsonAppendFunctionImpl`, which
+//! delegates to `JsonFunctions.jsonInsert` + `.meaningless_key` trick so Jayway
+//! routes to `Collection.add`). Non-array / missing targets are silent no-ops;
+//! any-NULL-arg / odd trailing arg / malformed-doc / malformed-path → NULL.
+//!
+//! Values always push as `Value::String` — every UDF arg is coerced to Utf8
+//! upstream, so nested `json_object` / `json_array` results arrive already
+//! stringified and append as strings, matching legacy.
+
+use std::any::Any;
+use std::sync::Arc;
+
+use datafusion::arrow::array::{Array, ArrayRef, StringBuilder};
+use datafusion::arrow::datatypes::DataType;
+use datafusion::common::ScalarValue;
+use datafusion::error::Result;
+use datafusion::execution::context::SessionContext;
+use datafusion::logical_expr::{
+    ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility,
+};
+use serde_json::Value;
+
+use super::json_common::{as_utf8_array, parse, parse_ppl_segments, walk_mut, Segment};
+use super::{coerce_slot, CoerceMode};
+
+const NAME: &str = "json_append";
+
+pub fn register_all(ctx: &SessionContext) {
+    ctx.register_udf(ScalarUDF::from(JsonAppendUdf::new()));
+}
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct JsonAppendUdf {
+    signature: Signature,
+}
+
+impl JsonAppendUdf {
+    pub fn new() -> Self {
+        Self {
+            signature: Signature::user_defined(Volatility::Immutable),
+        }
+    }
+}
+
+impl Default for JsonAppendUdf {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl ScalarUDFImpl for JsonAppendUdf {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+    fn name(&self) -> &str {
+        NAME
+    }
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+    fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+        Ok(DataType::Utf8)
+    }
+    fn coerce_types(&self, args: &[DataType]) -> Result<Vec<DataType>> {
+        args.iter()
+            .enumerate()
+            .map(|(i, ty)| coerce_slot(NAME, i, ty, CoerceMode::Utf8))
+            .collect()
+    }
+
+    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+        // Need doc + at least one (path, value) pair. Odd trailing arg mirrors
+        // the legacy `RuntimeException("needs corresponding path and values")`
+        // thrown by `JsonAppendFunctionImpl.eval`; we surface it as NULL to
+        // keep parity with the "malformed input → NULL" convention.
+        if args.args.len() < 3 || args.args.len().is_multiple_of(2) {
+            return Ok(ColumnarValue::Scalar(ScalarValue::Utf8(None)));
+        }
+        let n = args.number_rows;
+
+        if args
+            .args
+            .iter()
+            .all(|v| matches!(v, ColumnarValue::Scalar(_)))
+        {
+            let doc = scalar_utf8(&args.args[0]);
+            let rest: Vec<Option<&str>> = args.args[1..].iter().map(scalar_utf8).collect();
+            let out = append(doc, &rest);
+            return Ok(ColumnarValue::Scalar(ScalarValue::Utf8(out)));
+        }
+
+        let arrays: Vec<ArrayRef> = args
+            .args
+            .iter()
+            .map(|v| v.clone().into_array(n))
+            .collect::<Result<_>>()?;
+        let columns: Vec<&datafusion::arrow::array::StringArray> =
+            arrays.iter().map(as_utf8_array).collect::<Result<_>>()?;
+
+        let mut b = StringBuilder::with_capacity(n, n * 16);
+        let mut rest: Vec<Option<&str>> = Vec::with_capacity(columns.len() - 1);
+        for i in 0..n {
+            let doc = cell(columns[0], i);
+            rest.clear();
+            for col in &columns[1..] {
+                rest.push(cell(col, i));
+            }
+            match append(doc, &rest) {
+                Some(s) => b.append_value(&s),
+                None => b.append_null(),
+            }
+        }
+        Ok(ColumnarValue::Array(Arc::new(b.finish()) as ArrayRef))
+    }
+}
+
+fn scalar_utf8(v: &ColumnarValue) -> Option<&str> {
+    match v {
+        ColumnarValue::Scalar(
+            ScalarValue::Utf8(s) | ScalarValue::LargeUtf8(s) | ScalarValue::Utf8View(s),
+        ) => s.as_deref(),
+        _ => None,
+    }
+}
+
+fn cell(arr: &datafusion::arrow::array::StringArray, i: usize) -> Option<&str> {
+    if arr.is_null(i) {
+        None
+    } else {
+        Some(arr.value(i))
+    }
+}
+
+/// Apply each (path, value) pair to a fresh parse of `doc`. Push-only:
+/// non-array targets (scalar, object) are silent no-ops, matching legacy
+/// `jsonInsert`'s Collection-parent branch skip.
+fn append(doc: Option<&str>, rest: &[Option<&str>]) -> Option<String> {
+    let doc_str = doc?;
+    if rest.iter().any(|p| p.is_none()) {
+        return None;
+    }
+    let mut value = parse(doc_str)?;
+    for chunk in rest.chunks(2) {
+        let path = chunk[0].unwrap();
+        let new_val = chunk[1].unwrap();
+        let segments = parse_ppl_segments(path).ok()?;
+        if segments.is_empty() {
+            // Root-path is a no-op (legacy `ctx.set("$", v)` is silently
+            // discarded by Jayway for the same reason).
+            continue;
+        }
+        append_one(&mut value, &segments, new_val);
+    }
+    serde_json::to_string(&value).ok()
+}
+
+fn append_one(root: &mut Value, segments: &[Segment<'_>], new_val: &str) {
+    let item = Value::String(new_val.to_string());
+    walk_mut(root, segments, |parent, final_seg| {
+        match (parent, final_seg) {
+            // Push onto the matched array when the final segment names an
+            // existing array-valued field. Non-array / missing → no-op.
+            (Value::Object(map), Segment::Field(name)) => {
+                if let Some(Value::Array(arr)) = map.get_mut(*name) {
+                    arr.push(item.clone());
+                }
+            }
+            // Direct array-index / wildcard targets: push onto the *addressed*
+            // array element when that element is itself an array.
+            (Value::Array(arr), Segment::Index(i)) if *i < arr.len() => {
+                if let Value::Array(inner) = &mut arr[*i] {
+                    inner.push(item.clone());
+                }
+            }
+            (Value::Array(arr), Segment::Wildcard) => {
+                for slot in arr.iter_mut() {
+                    if let Value::Array(inner) = slot {
+                        inner.push(item.clone());
+                    }
+                }
+            }
+            _ => {}
+        }
+    });
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn single_value_appended_to_named_array() {
+        // testJsonAppend case b, single pair.
+        assert_eq!(
+            append(
+                Some(r#"{"teacher":["Alice"]}"#),
+                &[Some("teacher"), Some("Tom")],
+            )
+            .as_deref(),
+            Some(r#"{"teacher":["Alice","Tom"]}"#)
+        );
+    }
+
+    #[test]
+    fn multiple_pairs_append_sequentially() {
+        // testJsonAppend case b (multi-pair).
+        assert_eq!(
+            append(
+                Some(r#"{"teacher":["Alice"]}"#),
+                &[Some("teacher"), Some("Tom"), Some("teacher"), Some("Walt")],
+            )
+            .as_deref(),
+            Some(r#"{"teacher":["Alice","Tom","Walt"]}"#)
+        );
+    }
+
+    #[test]
+    fn nested_path_appends_to_inner_array() {
+        // testJsonAppend case c — a pre-stringified JSON array is appended as
+        // a single string element (legacy calls gson/jackson on the outer doc
+        // but NOT on the value; our Utf8-coerced arg arrives already
+        // stringified and is pushed as-is).
+        assert_eq!(
+            append(
+                Some(r#"{"school":{"teacher":["Alice"]}}"#),
+                &[Some("school.teacher"), Some(r#"["Tom","Walt"]"#)],
+            )
+            .as_deref(),
+            Some(r#"{"school":{"teacher":["Alice","[\"Tom\",\"Walt\"]"]}}"#)
+        );
+    }
+
+    #[test]
+    fn stringified_json_object_value_is_appended_as_single_string() {
+        // testJsonAppend case a — `json_object(...)` lowers to a string, so
+        // the element lands as a stringified object (legacy and Rust agree).
+        assert_eq!(
+            append(
+                Some(r#"{"student":[{"name":"Bob","rank":1}]}"#),
+                &[Some("student"), Some(r#"{"name":"Tomy","rank":5}"#)],
+            )
+            .as_deref(),
+            Some(r#"{"student":[{"name":"Bob","rank":1},"{\"name\":\"Tomy\",\"rank\":5}"]}"#)
+        );
+    }
+
+    #[test]
+    fn non_array_target_is_silent_noop() {
+        // teacher is a scalar here, not an array — legacy `Collection.add`
+        // branch skips; no-op is the observable parity.
+        assert_eq!(
+            append(
+                Some(r#"{"teacher":"Alice"}"#),
+                &[Some("teacher"), Some("Tom")],
+            )
+            .as_deref(),
+            Some(r#"{"teacher":"Alice"}"#)
+        );
+    }
+
+    #[test]
+    fn missing_path_is_silent_noop() {
+        assert_eq!(
+            append(
+                Some(r#"{"teacher":["Alice"]}"#),
+                &[Some("students"), Some("Tom")],
+            )
+            .as_deref(),
+            Some(r#"{"teacher":["Alice"]}"#)
+        );
+    }
+
+    #[test]
+    fn wildcard_path_appends_to_every_array_child() {
+        // Nested wildcard: every element of groups is an array; each receives
+        // the same appended scalar.
+        assert_eq!(
+            append(
+                Some(r#"{"groups":[["a"],["b","c"]]}"#),
+                &[Some("groups{}"), Some("x")],
+            )
+            .as_deref(),
+            Some(r#"{"groups":[["a","x"],["b","c","x"]]}"#)
+        );
+    }
+
+    #[test]
+    fn any_null_arg_returns_none() {
+        assert!(append(None, &[Some("a"), Some("v")]).is_none());
+        assert!(append(Some(r#"{"a":[1]}"#), &[None, Some("v")]).is_none());
+        assert!(append(Some(r#"{"a":[1]}"#), &[Some("a"), None]).is_none());
+    }
+
+    #[test]
+    fn malformed_doc_returns_none() {
+        assert!(append(Some("not-json"), &[Some("a"), Some("v")]).is_none());
+    }
+
+    #[test]
+    fn malformed_path_returns_none() {
+        assert!(append(Some(r#"{"a":[1]}"#), &[Some("a{"), Some("v")]).is_none());
+    }
+
+    #[test]
+    fn coerce_types_enforces_string_on_every_slot() {
+        let udf = JsonAppendUdf::new();
+        assert_eq!(
+            udf.coerce_types(&[DataType::Utf8, DataType::LargeUtf8, DataType::Utf8View])
+                .unwrap(),
+            vec![DataType::Utf8, DataType::Utf8, DataType::Utf8]
+        );
+        let err = udf
+            .coerce_types(&[DataType::Utf8, DataType::Int32, DataType::Utf8])
+            .unwrap_err()
+            .to_string();
+        assert!(err.contains("expected string"));
+    }
+
+    #[test]
+    fn return_type_is_utf8() {
+        assert_eq!(
+            JsonAppendUdf::new().return_type(&[DataType::Utf8]).unwrap(),
+            DataType::Utf8
+        );
+    }
+}