diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ScalarFunction.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ScalarFunction.java index ee32b35be4b9e..af6f15e84935b 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ScalarFunction.java +++ b/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. diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/Cargo.toml b/sandbox/plugins/analytics-backend-datafusion/rust/Cargo.toml index 63afe25b394f1..ba0f7a0c0226a 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/Cargo.toml +++ b/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` 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 } diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_append.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_append.rs new file mode 100644 index 0000000000000..8b3608c27fcb1 --- /dev/null +++ b/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 { + Ok(DataType::Utf8) + } + fn coerce_types(&self, args: &[DataType]) -> Result> { + args.iter() + .enumerate() + .map(|(i, ty)| coerce_slot(NAME, i, ty, CoerceMode::Utf8)) + .collect() + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + // 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> = args.args[1..].iter().map(scalar_utf8).collect(); + let out = append(doc, &rest); + return Ok(ColumnarValue::Scalar(ScalarValue::Utf8(out))); + } + + let arrays: Vec = args + .args + .iter() + .map(|v| v.clone().into_array(n)) + .collect::>()?; + let columns: Vec<&datafusion::arrow::array::StringArray> = + arrays.iter().map(as_utf8_array).collect::>()?; + + let mut b = StringBuilder::with_capacity(n, n * 16); + let mut rest: Vec> = 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 { + 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 + ); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_array_length.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_array_length.rs new file mode 100644 index 0000000000000..c8f60647d0ea0 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_array_length.rs @@ -0,0 +1,260 @@ +/* + * 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_array_length(value)` — length of a JSON array (parity with legacy +//! `JsonArrayLengthFunctionImpl`; verified by `CalcitePPLJsonBuiltinFunctionIT.testJsonArrayLength`). +//! NULL / non-array / malformed → NULL. Only plan-time arity / type failures +//! surface as `plan_err!`; runtime input of any content never errors. + +use std::any::Any; +use std::sync::Arc; + +use datafusion::arrow::array::{Array, ArrayRef, Int32Builder, StringArray}; +use datafusion::arrow::datatypes::DataType; +use datafusion::common::{plan_err, ScalarValue}; +use datafusion::error::{DataFusionError, Result}; +use datafusion::execution::context::SessionContext; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, +}; +use serde_json::Value; + +use super::{coerce_args, CoerceMode}; + +pub fn register_all(ctx: &SessionContext) { + ctx.register_udf(ScalarUDF::from(JsonArrayLengthUdf::new())); +} + +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct JsonArrayLengthUdf { + signature: Signature, +} + +impl JsonArrayLengthUdf { + pub fn new() -> Self { + // user_defined + coerce_types lets DF cast LargeUtf8 / Utf8View → Utf8. + Self { + signature: Signature::user_defined(Volatility::Immutable), + } + } +} + +impl Default for JsonArrayLengthUdf { + fn default() -> Self { + Self::new() + } +} + +impl ScalarUDFImpl for JsonArrayLengthUdf { + fn as_any(&self) -> &dyn Any { + self + } + fn name(&self) -> &str { + "json_array_length" + } + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result { + if arg_types.len() != 1 { + return plan_err!( + "json_array_length expects 1 argument, got {}", + arg_types.len() + ); + } + // Int32 to match PPL's INTEGER_FORCE_NULLABLE declaration. Returning + // Int64 here works for literal args (Calcite const-folds and inserts a + // narrowing CAST on the project) but leaks Int64 through the column + // path — caller sees Integer for literals, Long for column refs. + Ok(DataType::Int32) + } + + fn coerce_types(&self, arg_types: &[DataType]) -> Result> { + coerce_args("json_array_length", arg_types, &[CoerceMode::Utf8]) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + if args.args.len() != 1 { + return plan_err!( + "json_array_length expects 1 argument, got {}", + args.args.len() + ); + } + let n = args.number_rows; + + // Scalar fast-path: parse once, broadcast as scalar output. + if let ColumnarValue::Scalar(sv) = &args.args[0] { + let len = match sv { + ScalarValue::Utf8(opt) | ScalarValue::LargeUtf8(opt) | ScalarValue::Utf8View(opt) => { + opt.as_deref().and_then(json_array_len) + } + _ => None, + }; + return Ok(ColumnarValue::Scalar(ScalarValue::Int32(len))); + } + + let arr = args.args[0].clone().into_array(n)?; + let strings = arr.as_any().downcast_ref::().ok_or_else(|| { + DataFusionError::Internal(format!( + "json_array_length: expected Utf8, got {:?}", + arr.data_type() + )) + })?; + + let mut builder = Int32Builder::with_capacity(n); + for i in 0..n { + if strings.is_null(i) { + builder.append_null(); + continue; + } + match json_array_len(strings.value(i)) { + Some(len) => builder.append_value(len), + None => builder.append_null(), + } + } + Ok(ColumnarValue::Array(Arc::new(builder.finish()) as ArrayRef)) + } +} + +/// Returns the array length as i32, or None for malformed / non-array input. +/// i32 matches PPL's declared INTEGER return type; arrays exceeding i32::MAX +/// elements (>2B) saturate to NULL rather than silently truncating. +fn json_array_len(s: &str) -> Option { + serde_json::from_str::(s) + .ok() + .and_then(|v| v.as_array().map(|a| a.len())) + .and_then(|len| i32::try_from(len).ok()) +} + +// ─── tests ────────────────────────────────────────────────────────────────── + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::array::Int32Array; + use datafusion::arrow::datatypes::Field; + + #[test] + fn parses_array_returns_length() { + assert_eq!(json_array_len("[1,2,3]"), Some(3)); + assert_eq!(json_array_len("[]"), Some(0)); + assert_eq!(json_array_len("[\"a\",\"b\"]"), Some(2)); + // Heterogeneous array — parity with legacy Gson List parse. + assert_eq!(json_array_len("[1,\"x\",{\"k\":1}]"), Some(3)); + } + + #[test] + fn non_array_json_returns_none() { + assert_eq!(json_array_len("{\"k\":1}"), None); + assert_eq!(json_array_len("\"scalar\""), None); + assert_eq!(json_array_len("42"), None); + assert_eq!(json_array_len("null"), None); + } + + #[test] + fn malformed_json_returns_none() { + assert_eq!(json_array_len("not-json"), None); + assert_eq!(json_array_len("[1,2"), None); + assert_eq!(json_array_len(""), None); + } + + #[test] + fn coerce_types_accepts_string_variants() { + let udf = JsonArrayLengthUdf::new(); + for t in [DataType::Utf8, DataType::LargeUtf8, DataType::Utf8View] { + let out = udf.coerce_types(std::slice::from_ref(&t)).unwrap(); + assert_eq!(out, vec![DataType::Utf8], "input {t:?} should coerce to Utf8"); + } + } + + #[test] + fn coerce_types_rejects_non_string() { + let udf = JsonArrayLengthUdf::new(); + let err = udf.coerce_types(&[DataType::Int64]).unwrap_err(); + assert!(err.to_string().contains("expected string")); + } + + #[test] + fn coerce_types_rejects_wrong_arity() { + let udf = JsonArrayLengthUdf::new(); + assert!(udf.coerce_types(&[]).is_err()); + assert!(udf.coerce_types(&[DataType::Utf8, DataType::Utf8]).is_err()); + } + + #[test] + fn return_type_is_int32() { + let udf = JsonArrayLengthUdf::new(); + let out = udf.return_type(&[DataType::Utf8]).unwrap(); + assert_eq!(out, DataType::Int32); + } + + #[test] + fn invoke_handles_nulls_malformed_and_non_array() { + let udf = JsonArrayLengthUdf::new(); + let input = StringArray::from(vec![ + Some("[1,2,3]"), + None, + Some("{\"k\":1}"), + Some("not-json"), + Some("[]"), + ]); + let args = ScalarFunctionArgs { + args: vec![ColumnarValue::Array(Arc::new(input))], + number_rows: 5, + arg_fields: vec![], + return_field: Arc::new(Field::new("out", DataType::Int32, true)), + config_options: Arc::new(datafusion::config::ConfigOptions::new()), + }; + let out = udf.invoke_with_args(args).unwrap(); + let arr = match out { + ColumnarValue::Array(a) => a, + _ => panic!("expected array"), + }; + let arr = arr.as_any().downcast_ref::().unwrap(); + assert_eq!(arr.value(0), 3); + assert!(arr.is_null(1)); + assert!(arr.is_null(2)); + assert!(arr.is_null(3)); + assert_eq!(arr.value(4), 0); + } + + #[test] + fn invoke_scalar_input_produces_scalar_output() { + let udf = JsonArrayLengthUdf::new(); + let args = ScalarFunctionArgs { + args: vec![ColumnarValue::Scalar(ScalarValue::Utf8(Some("[1,2,3,4]".into())))], + number_rows: 1, + arg_fields: vec![], + return_field: Arc::new(Field::new("out", DataType::Int32, true)), + config_options: Arc::new(datafusion::config::ConfigOptions::new()), + }; + let out = udf.invoke_with_args(args).unwrap(); + match out { + ColumnarValue::Scalar(ScalarValue::Int32(Some(4))) => {} + other => panic!("expected Int32(Some(4)), got {other:?}"), + } + } + + #[test] + fn invoke_scalar_null_input_yields_scalar_null() { + let udf = JsonArrayLengthUdf::new(); + let args = ScalarFunctionArgs { + args: vec![ColumnarValue::Scalar(ScalarValue::Utf8(None))], + number_rows: 1, + arg_fields: vec![], + return_field: Arc::new(Field::new("out", DataType::Int32, true)), + config_options: Arc::new(datafusion::config::ConfigOptions::new()), + }; + let out = udf.invoke_with_args(args).unwrap(); + match out { + ColumnarValue::Scalar(ScalarValue::Int32(None)) => {} + other => panic!("expected Int32(None), got {other:?}"), + } + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_common.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_common.rs new file mode 100644 index 0000000000000..80e77431d9c3e --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_common.rs @@ -0,0 +1,294 @@ +/* + * 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. + */ + +//! Shared helpers for the PPL `json_*` UDFs: PPL-path parsing (to both JSONPath +//! strings and typed segment vectors), a segment-based mutation walker used by +//! the write UDFs, and a malformed-to-`None` JSON parser. + +use datafusion::arrow::array::{ArrayRef, StringArray}; +use datafusion::error::{DataFusionError, Result}; +use serde_json::Value; + +/// Convert a PPL-style path (`a.b{0}.c{}`) to a JSONPath expression +/// (`$.a.b[0].c[*]`). Empty input → `"$"` (document root), matching the +/// legacy contract. Returns a planning error for an unmatched `{`. +pub(crate) fn convert_ppl_path(input: &str) -> Result { + if input.is_empty() { + return Ok("$".into()); + } + let mut out = String::with_capacity(input.len() + 2); + out.push_str("$."); + let mut rest = input; + while !rest.is_empty() { + match rest.as_bytes()[0] { + b'{' => { + let end = rest.find('}').ok_or_else(|| { + datafusion::error::DataFusionError::Plan(format!( + "Unmatched '{{' in JSON path: {input}" + )) + })?; + let idx = rest[1..end].trim(); + if idx.is_empty() { + out.push_str("[*]"); + } else { + out.push('['); + out.push_str(idx); + out.push(']'); + } + rest = &rest[end + 1..]; + } + b'.' => { + out.push('.'); + rest = &rest[1..]; + } + _ => { + let cut = rest.find(['.', '{']).unwrap_or(rest.len()); + out.push_str(&rest[..cut]); + rest = &rest[cut..]; + } + } + } + Ok(out) +} + +/// Parse a JSON string; returns `None` on malformed input. Matches the +/// "malformed → NULL" convention across all json_* UDFs (see +/// `json_udf_legacy_semantics.md`). +pub(crate) fn parse(s: &str) -> Option { + serde_json::from_str(s).ok() +} + +/// One tokenised step of a PPL path. Mirrors the three cases `convert_ppl_path` +/// handles: bare identifier (field), `{n}` (array index), `{}` (array wildcard). +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum Segment<'a> { + Field(&'a str), + Index(usize), + Wildcard, +} + +/// Tokenise a PPL path into `Segment`s without allocating for field names. +/// Returns a planning error for unmatched `{` or a non-numeric index — the +/// same inputs `convert_ppl_path` rejects. +pub(crate) fn parse_ppl_segments(input: &str) -> Result>> { + let mut out = Vec::new(); + let mut rest = input; + while !rest.is_empty() { + match rest.as_bytes()[0] { + b'{' => { + let end = rest.find('}').ok_or_else(|| { + DataFusionError::Plan(format!("Unmatched '{{' in JSON path: {input}")) + })?; + let idx = rest[1..end].trim(); + if idx.is_empty() { + out.push(Segment::Wildcard); + } else { + let parsed = idx.parse::().map_err(|_| { + DataFusionError::Plan(format!( + "Non-numeric array index '{idx}' in JSON path: {input}" + )) + })?; + out.push(Segment::Index(parsed)); + } + rest = &rest[end + 1..]; + } + b'.' => rest = &rest[1..], + _ => { + let cut = rest.find(['.', '{']).unwrap_or(rest.len()); + if cut > 0 { + out.push(Segment::Field(&rest[..cut])); + } + rest = &rest[cut..]; + } + } + } + Ok(out) +} + +/// Drive `apply` against every terminal `(parent, final_segment)` reached by +/// `segments` inside `root`. Missing intermediate keys / out-of-range indices +/// are silently skipped (matching Jayway's `SUPPRESS_EXCEPTIONS` behaviour +/// that legacy mutation UDFs rely on). Wildcard segments fan out across every +/// element of the current array; descending through a non-container +/// short-circuits that branch. +/// +/// Empty `segments` is a no-op: PPL mutation UDFs reject a root-only path at +/// the call site before reaching the walker. +pub(crate) fn walk_mut(root: &mut Value, segments: &[Segment<'_>], mut apply: F) +where + F: FnMut(&mut Value, &Segment<'_>), +{ + if segments.is_empty() { + return; + } + walk_mut_inner(root, segments, &mut apply); +} + +fn walk_mut_inner(node: &mut Value, segments: &[Segment<'_>], apply: &mut F) +where + F: FnMut(&mut Value, &Segment<'_>), +{ + let (head, tail) = segments.split_first().expect("non-empty checked by caller"); + if tail.is_empty() { + // Parent is `node`; the final segment names the slot to mutate. + apply(node, head); + return; + } + match head { + Segment::Field(name) => { + if let Value::Object(map) = node { + if let Some(child) = map.get_mut(*name) { + walk_mut_inner(child, tail, apply); + } + } + } + Segment::Index(i) => { + if let Value::Array(arr) = node { + if let Some(child) = arr.get_mut(*i) { + walk_mut_inner(child, tail, apply); + } + } + } + Segment::Wildcard => { + if let Value::Array(arr) = node { + for child in arr.iter_mut() { + walk_mut_inner(child, tail, apply); + } + } + } + } +} + +/// Standard arity guard. +pub(crate) fn check_arity(udf: &str, observed: usize, expected: usize) -> Result<()> { + (observed == expected) + .then_some(()) + .ok_or_else(|| plan_err_msg(format!("{udf} expects {expected} arguments, got {observed}"))) +} + +fn plan_err_msg(msg: String) -> DataFusionError { + DataFusionError::Plan(msg) +} + +/// Downcast an `ArrayRef` to `StringArray`. `coerce_types` with `CoerceMode::Utf8` +/// canonicalizes every string input to `Utf8` before this point, so a failure +/// indicates a planner bug rather than bad user input. +pub(crate) fn as_utf8_array(arr: &ArrayRef) -> Result<&StringArray> { + arr.as_any().downcast_ref::().ok_or_else(|| { + DataFusionError::Internal(format!("expected Utf8, got {:?}", arr.data_type())) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn ppl_path_mirrors_legacy_convert_to_jsonpath() { + for (input, want) in [ + ("", "$"), + ("a", "$.a"), + ("a.b", "$.a.b"), + ("a{0}", "$.a[0]"), + ("a{}", "$.a[*]"), + ("a{0}.b{}.c", "$.a[0].b[*].c"), + ("a{ 2 }", "$.a[2]"), + ] { + assert_eq!(convert_ppl_path(input).unwrap(), want, "input={input}"); + } + assert!(convert_ppl_path("a{0").unwrap_err().to_string().contains("Unmatched")); + } + + #[test] + fn parse_handles_malformed_and_valid() { + assert!(parse("{not json").is_none()); + assert!(parse("[1,2,3]").is_some()); + } + + #[test] + fn arity_guards() { + assert!(check_arity("f", 1, 1).is_ok()); + assert!(check_arity("f", 2, 1).is_err()); + } + + #[test] + fn parse_ppl_segments_tokenises_field_index_and_wildcard() { + assert_eq!(parse_ppl_segments("").unwrap(), Vec::::new()); + assert_eq!(parse_ppl_segments("a").unwrap(), vec![Segment::Field("a")]); + assert_eq!( + parse_ppl_segments("a.b{0}.c{}").unwrap(), + vec![ + Segment::Field("a"), + Segment::Field("b"), + Segment::Index(0), + Segment::Field("c"), + Segment::Wildcard, + ] + ); + assert!(parse_ppl_segments("a{0").is_err()); + assert!(parse_ppl_segments("a{x}").is_err()); + } + + fn v(s: &str) -> Value { + serde_json::from_str(s).unwrap() + } + + #[test] + fn walk_mut_deletes_flat_key() { + let mut doc = v(r#"{"a":1,"b":2,"c":3}"#); + let segs = parse_ppl_segments("b").unwrap(); + walk_mut(&mut doc, &segs, |parent, seg| { + if let (Value::Object(map), Segment::Field(name)) = (parent, seg) { + map.shift_remove(*name); + } + }); + assert_eq!(serde_json::to_string(&doc).unwrap(), r#"{"a":1,"c":3}"#); + } + + #[test] + fn walk_mut_handles_missing_path_as_noop() { + let mut doc = v(r#"{"f1":"abc","f2":{"f3":"a"}}"#); + let segs = parse_ppl_segments("f2.nope").unwrap(); + walk_mut(&mut doc, &segs, |parent, seg| { + if let (Value::Object(map), Segment::Field(name)) = (parent, seg) { + map.shift_remove(*name); + } + }); + assert_eq!( + serde_json::to_string(&doc).unwrap(), + r#"{"f1":"abc","f2":{"f3":"a"}}"# + ); + } + + #[test] + fn walk_mut_wildcard_fans_out_across_array() { + let mut doc = v(r#"{"xs":[{"k":1,"v":10},{"k":2,"v":20}]}"#); + let segs = parse_ppl_segments("xs{}.v").unwrap(); + walk_mut(&mut doc, &segs, |parent, seg| { + if let (Value::Object(map), Segment::Field(name)) = (parent, seg) { + map.shift_remove(*name); + } + }); + assert_eq!( + serde_json::to_string(&doc).unwrap(), + r#"{"xs":[{"k":1},{"k":2}]}"# + ); + } + + #[test] + fn walk_mut_index_out_of_range_is_noop() { + let mut doc = v(r#"{"xs":[{"k":1}]}"#); + let segs = parse_ppl_segments("xs{5}.k").unwrap(); + walk_mut(&mut doc, &segs, |parent, seg| { + if let (Value::Object(map), Segment::Field(name)) = (parent, seg) { + map.shift_remove(*name); + } + }); + assert_eq!(serde_json::to_string(&doc).unwrap(), r#"{"xs":[{"k":1}]}"#); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_delete.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_delete.rs new file mode 100644 index 0000000000000..ff189a7bcc4bf --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_delete.rs @@ -0,0 +1,286 @@ +/* + * 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_delete(value, path1, [path2, ...])` — remove path-matched entries +//! from a JSON document (parity with legacy `JsonDeleteFunctionImpl` → Jayway +//! `JsonPath.delete` under `SUPPRESS_EXCEPTIONS`, applied per pathspec). +//! Missing paths are no-ops; any-NULL-arg / malformed-doc / malformed-path → +//! NULL. Output key order is preserved via `serde_json`'s `preserve_order` +//! feature (see `rust/Cargo.toml`). + +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_delete"; + +pub fn register_all(ctx: &SessionContext) { + ctx.register_udf(ScalarUDF::from(JsonDeleteUdf::new())); +} + +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct JsonDeleteUdf { + signature: Signature, +} + +impl JsonDeleteUdf { + pub fn new() -> Self { + Self { + signature: Signature::user_defined(Volatility::Immutable), + } + } +} + +impl Default for JsonDeleteUdf { + fn default() -> Self { + Self::new() + } +} + +impl ScalarUDFImpl for JsonDeleteUdf { + 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 { + Ok(DataType::Utf8) + } + fn coerce_types(&self, args: &[DataType]) -> Result> { + args.iter() + .enumerate() + .map(|(i, ty)| coerce_slot(NAME, i, ty, CoerceMode::Utf8)) + .collect() + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + if args.args.len() < 2 { + // Legacy `jsonRemove(doc)` with no pathspecs would return `doc` + // unchanged. Matching that as SQL NULL (not an error) keeps us + // consistent with the other json_* UDFs' any-NULL-arg convention. + 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 paths: Vec> = args.args[1..].iter().map(scalar_utf8).collect(); + let out = delete(doc, &paths); + return Ok(ColumnarValue::Scalar(ScalarValue::Utf8(out))); + } + + let arrays: Vec = args + .args + .iter() + .map(|v| v.clone().into_array(n)) + .collect::>()?; + let columns: Vec<&datafusion::arrow::array::StringArray> = + arrays.iter().map(as_utf8_array).collect::>()?; + + let mut b = StringBuilder::with_capacity(n, n * 16); + let mut path_buf: Vec> = Vec::with_capacity(columns.len() - 1); + for i in 0..n { + let doc = cell(columns[0], i); + path_buf.clear(); + for col in &columns[1..] { + path_buf.push(cell(col, i)); + } + match delete(doc, &path_buf) { + 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 every path's delete to a fresh parse of `doc`. Returns `None` for +/// any-NULL arg / malformed doc / malformed path; otherwise the mutated +/// document serialized back to a string. +fn delete(doc: Option<&str>, paths: &[Option<&str>]) -> Option { + let doc_str = doc?; + if paths.iter().any(|p| p.is_none()) { + return None; + } + let mut value = parse(doc_str)?; + for path in paths.iter().map(|p| p.unwrap()) { + let segments = parse_ppl_segments(path).ok()?; + if segments.is_empty() { + // Legacy `jsonRemove` on an empty path attempts to `ctx.read("$")` + // which returns the root; `ctx.delete("$")` is a Jayway no-op + // (root is indelible). Mirror that by skipping. + continue; + } + delete_one(&mut value, &segments); + } + serde_json::to_string(&value).ok() +} + +fn delete_one(root: &mut Value, segments: &[Segment<'_>]) { + walk_mut(root, segments, |parent, final_seg| { + match (parent, final_seg) { + (Value::Object(map), Segment::Field(name)) => { + map.shift_remove(*name); + } + (Value::Array(arr), Segment::Index(i)) if *i < arr.len() => { + arr.remove(*i); + } + (Value::Array(arr), Segment::Wildcard) => { + arr.clear(); + } + // Type mismatch between the container and the terminal segment is a + // silent no-op, matching Jayway's SUPPRESS_EXCEPTIONS. + _ => {} + } + }); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn flat_key_delete_matches_legacy_fixture() { + // testJsonDelete + assert_eq!( + delete( + Some(r#"{"account_number":1,"balance":39225,"age":32,"gender":"M"}"#), + &[Some("age"), Some("gender")], + ) + .as_deref(), + Some(r#"{"account_number":1,"balance":39225}"#) + ); + } + + #[test] + fn nested_key_delete_preserves_siblings() { + // testJsonDeleteWithNested + assert_eq!( + delete( + Some(r#"{"f1":"abc","f2":{"f3":"a","f4":"b"}}"#), + &[Some("f2.f3")], + ) + .as_deref(), + Some(r#"{"f1":"abc","f2":{"f4":"b"}}"#) + ); + } + + #[test] + fn missing_path_returns_document_unchanged() { + // testJsonDeleteWithNestedNothing + assert_eq!( + delete( + Some(r#"{"f1":"abc","f2":{"f3":"a","f4":"b"}}"#), + &[Some("f2.f100")], + ) + .as_deref(), + Some(r#"{"f1":"abc","f2":{"f3":"a","f4":"b"}}"#) + ); + } + + #[test] + fn wildcard_array_delete_matches_legacy_fixture() { + // testJsonDeleteWithNestedAndArray + assert_eq!( + delete( + Some( + r#"{"teacher":"Alice","student":[{"name":"Bob","rank":1},{"name":"Charlie","rank":2}]}"# + ), + &[Some("teacher"), Some("student{}.rank")], + ) + .as_deref(), + Some(r#"{"student":[{"name":"Bob"},{"name":"Charlie"}]}"#) + ); + } + + #[test] + fn any_null_arg_returns_none() { + assert!(delete(None, &[Some("a")]).is_none()); + assert!(delete(Some(r#"{"a":1}"#), &[None]).is_none()); + } + + #[test] + fn malformed_doc_returns_none() { + assert!(delete(Some("not-json"), &[Some("a")]).is_none()); + } + + #[test] + fn malformed_path_returns_none() { + assert!(delete(Some(r#"{"a":1}"#), &[Some("a{")]).is_none()); + } + + #[test] + fn less_than_two_args_returns_none_via_fast_path() { + // Exercised through invoke_with_args — the top-level guard returns + // Utf8(None) for <2 args, so helper-level coverage is enough. + assert_eq!( + delete(Some(r#"{"a":1}"#), &[Some("a")]).as_deref(), + Some(r#"{}"#) + ); + } + + #[test] + fn coerce_types_enforces_string_on_every_slot() { + let udf = JsonDeleteUdf::new(); + assert_eq!( + udf.coerce_types(&[DataType::LargeUtf8, DataType::Utf8View]) + .unwrap(), + vec![DataType::Utf8, DataType::Utf8] + ); + let err = udf + .coerce_types(&[DataType::Utf8, DataType::Int32]) + .unwrap_err() + .to_string(); + assert!(err.contains("expected string")); + } + + #[test] + fn return_type_is_utf8() { + assert_eq!( + JsonDeleteUdf::new().return_type(&[DataType::Utf8]).unwrap(), + DataType::Utf8 + ); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_extend.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_extend.rs new file mode 100644 index 0000000000000..d1f61289a61ed --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_extend.rs @@ -0,0 +1,348 @@ +/* + * 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_extend(value, path1, val1, [path2, val2, ...])` — spread-or-append: +//! if `valN` parses as a JSON array its elements are pushed individually onto +//! the path-matched target array; otherwise the whole value is pushed as a +//! single string element (parity with legacy `gson.fromJson(..., List.class)` +//! try/fall-back in `JsonExtendFunctionImpl`). Non-array / missing targets are +//! no-ops; any-NULL-arg / malformed-doc / malformed-path → NULL. +//! +//! Intentional divergence from legacy: spread preserves source numeric type +//! (`[1,2,3]` → `1,2,3`). Gson widens every number to `Double` (`1.0, 2.0, +//! 3.0`); no legacy IT covers this edge case, so every existing fixture still +//! passes. Tracked for follow-up cross-engine alignment. + +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_extend"; + +pub fn register_all(ctx: &SessionContext) { + ctx.register_udf(ScalarUDF::from(JsonExtendUdf::new())); +} + +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct JsonExtendUdf { + signature: Signature, +} + +impl JsonExtendUdf { + pub fn new() -> Self { + Self { + signature: Signature::user_defined(Volatility::Immutable), + } + } +} + +impl Default for JsonExtendUdf { + fn default() -> Self { + Self::new() + } +} + +impl ScalarUDFImpl for JsonExtendUdf { + 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 { + Ok(DataType::Utf8) + } + fn coerce_types(&self, args: &[DataType]) -> Result> { + args.iter() + .enumerate() + .map(|(i, ty)| coerce_slot(NAME, i, ty, CoerceMode::Utf8)) + .collect() + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + 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> = args.args[1..].iter().map(scalar_utf8).collect(); + let out = extend(doc, &rest); + return Ok(ColumnarValue::Scalar(ScalarValue::Utf8(out))); + } + + let arrays: Vec = args + .args + .iter() + .map(|v| v.clone().into_array(n)) + .collect::>()?; + let columns: Vec<&datafusion::arrow::array::StringArray> = + arrays.iter().map(as_utf8_array).collect::>()?; + + let mut b = StringBuilder::with_capacity(n, n * 16); + let mut rest: Vec> = 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 extend(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)) + } +} + +/// Classify the raw value string into the push-list the terminal closure +/// should apply. A successful JSON-array parse expands to the array +/// elements; anything else (scalar, object, malformed JSON, plain string) +/// expands to `[Value::String(value)]` — the legacy `gson.fromJson` +/// try/fall-back pattern. +fn spread(raw: &str) -> Vec { + if let Ok(Value::Array(items)) = serde_json::from_str::(raw) { + return items; + } + vec![Value::String(raw.to_string())] +} + +fn extend(doc: Option<&str>, rest: &[Option<&str>]) -> Option { + 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() { + continue; + } + let items = spread(new_val); + extend_one(&mut value, &segments, &items); + } + serde_json::to_string(&value).ok() +} + +fn extend_one(root: &mut Value, segments: &[Segment<'_>], items: &[Value]) { + walk_mut(root, segments, |parent, final_seg| { + match (parent, final_seg) { + (Value::Object(map), Segment::Field(name)) => { + if let Some(Value::Array(arr)) = map.get_mut(*name) { + arr.extend(items.iter().cloned()); + } + } + (Value::Array(arr), Segment::Index(i)) if *i < arr.len() => { + if let Value::Array(inner) = &mut arr[*i] { + inner.extend(items.iter().cloned()); + } + } + (Value::Array(arr), Segment::Wildcard) => { + for slot in arr.iter_mut() { + if let Value::Array(inner) = slot { + inner.extend(items.iter().cloned()); + } + } + } + _ => {} + } + }); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn json_array_value_is_spread_into_target_array() { + // testJsonExtend case c — the stringified json_array(...) value is a + // JSON array, so its elements are spread (contrast json_append's case + // c, which pushes the whole string as one element). + assert_eq!( + extend( + Some(r#"{"school":{"teacher":["Alice"]}}"#), + &[Some("school.teacher"), Some(r#"["Tom","Walt"]"#)], + ) + .as_deref(), + Some(r#"{"school":{"teacher":["Alice","Tom","Walt"]}}"#) + ); + } + + #[test] + fn non_array_value_falls_back_to_single_push() { + // testJsonExtend case a — json_object(...) stringifies to a JSON + // object, not an array. Parse-as-List fails → single element pushed. + assert_eq!( + extend( + Some(r#"{"student":[{"name":"Bob","rank":1}]}"#), + &[Some("student"), Some(r#"{"name":"Tommy","rank":5}"#)], + ) + .as_deref(), + Some(r#"{"student":[{"name":"Bob","rank":1},"{\"name\":\"Tommy\",\"rank\":5}"]}"#) + ); + } + + #[test] + fn plain_string_value_falls_back_to_single_push() { + // testJsonExtend case b — plain "Tom" / "Walt" strings are not JSON + // arrays, so each is pushed as a single element. + assert_eq!( + extend( + Some(r#"{"teacher":["Alice"]}"#), + &[Some("teacher"), Some("Tom"), Some("teacher"), Some("Walt")], + ) + .as_deref(), + Some(r#"{"teacher":["Alice","Tom","Walt"]}"#) + ); + } + + #[test] + fn non_array_target_is_silent_noop() { + assert_eq!( + extend( + Some(r#"{"teacher":"Alice"}"#), + &[Some("teacher"), Some(r#"["Tom"]"#)], + ) + .as_deref(), + Some(r#"{"teacher":"Alice"}"#) + ); + } + + #[test] + fn missing_path_is_silent_noop() { + assert_eq!( + extend( + Some(r#"{"teacher":["Alice"]}"#), + &[Some("students"), Some("Tom")], + ) + .as_deref(), + Some(r#"{"teacher":["Alice"]}"#) + ); + } + + #[test] + fn empty_json_array_value_is_a_noop_on_target() { + // Parses as an array of zero items → nothing to push. + assert_eq!( + extend( + Some(r#"{"teacher":["Alice"]}"#), + &[Some("teacher"), Some("[]")], + ) + .as_deref(), + Some(r#"{"teacher":["Alice"]}"#) + ); + } + + #[test] + fn mixed_type_json_array_elements_preserve_their_types() { + // Integers/booleans come through as JSON numbers/booleans, not as + // stringified elements — diverges from legacy Gson (which widens to + // Double) but no legacy IT asserts Gson's widening. See module-level + // docs for the rationale + tracking-issue pointer. + assert_eq!( + extend( + Some(r#"{"xs":[0]}"#), + &[Some("xs"), Some("[1,2,true,\"s\"]")], + ) + .as_deref(), + Some(r#"{"xs":[0,1,2,true,"s"]}"#) + ); + } + + #[test] + fn wildcard_path_extends_every_array_child() { + assert_eq!( + extend( + Some(r#"{"groups":[["a"],["b","c"]]}"#), + &[Some("groups{}"), Some(r#"["x","y"]"#)], + ) + .as_deref(), + Some(r#"{"groups":[["a","x","y"],["b","c","x","y"]]}"#) + ); + } + + #[test] + fn any_null_arg_returns_none() { + assert!(extend(None, &[Some("a"), Some("v")]).is_none()); + assert!(extend(Some(r#"{"a":[1]}"#), &[None, Some("v")]).is_none()); + assert!(extend(Some(r#"{"a":[1]}"#), &[Some("a"), None]).is_none()); + } + + #[test] + fn malformed_doc_returns_none() { + assert!(extend(Some("not-json"), &[Some("a"), Some("v")]).is_none()); + } + + #[test] + fn malformed_path_returns_none() { + assert!(extend(Some(r#"{"a":[1]}"#), &[Some("a{"), Some("v")]).is_none()); + } + + #[test] + fn coerce_types_enforces_string_on_every_slot() { + let udf = JsonExtendUdf::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!( + JsonExtendUdf::new().return_type(&[DataType::Utf8]).unwrap(), + DataType::Utf8 + ); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_extract.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_extract.rs new file mode 100644 index 0000000000000..e046fc759fa98 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_extract.rs @@ -0,0 +1,308 @@ +/* + * 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_extract(value, path1, [path2, ...])` — extract JSON value(s) by PPL-path +//! (parity with legacy `JsonExtractFunctionImpl` → Calcite `jsonQuery` / `jsonValue`). +//! Single-path: scalar → `.to_string()`, string → unquoted, object/array → +//! JSON-serialized, wildcard multi-match → JSON-array, miss/explicit-null → NULL. +//! Multi-path: per-path results (NULL → `null` element) wrapped in a JSON array. +//! `< 2` args / any-NULL-arg / malformed doc / malformed path → NULL. + +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 jsonpath_rust::{JsonPath, JsonPathValue}; +use serde_json::Value; + +use super::json_common::{as_utf8_array, convert_ppl_path, parse}; +use super::{coerce_slot, CoerceMode}; + +const NAME: &str = "json_extract"; + +pub fn register_all(ctx: &SessionContext) { + ctx.register_udf(ScalarUDF::from(JsonExtractUdf::new())); +} + +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct JsonExtractUdf { + signature: Signature, +} + +impl JsonExtractUdf { + pub fn new() -> Self { + Self { + signature: Signature::user_defined(Volatility::Immutable), + } + } +} + +impl Default for JsonExtractUdf { + fn default() -> Self { + Self::new() + } +} + +impl ScalarUDFImpl for JsonExtractUdf { + 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 { + Ok(DataType::Utf8) + } + fn coerce_types(&self, args: &[DataType]) -> Result> { + // Homogeneous string variadic — every slot canonicalizes to Utf8. + args.iter() + .enumerate() + .map(|(i, ty)| coerce_slot(NAME, i, ty, CoerceMode::Utf8)) + .collect() + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + if args.args.len() < 2 { + // Legacy short-circuit: < 2 args → NULL. Avoid plan_err so adapter + // mismatches surface as data NULL rather than query failure + // (matches JsonExtractFunctionImpl.eval's `if (args.length < 2)`). + return Ok(ColumnarValue::Scalar(ScalarValue::Utf8(None))); + } + let n = args.number_rows; + + // Scalar fast-path: every operand is Scalar → evaluate once. + if args + .args + .iter() + .all(|v| matches!(v, ColumnarValue::Scalar(_))) + { + let doc = scalar_utf8(&args.args[0]); + let paths: Vec> = args.args[1..].iter().map(scalar_utf8).collect(); + let out = extract(doc, &paths); + return Ok(ColumnarValue::Scalar(ScalarValue::Utf8(out))); + } + + // Columnar path: materialize each operand to a StringArray and walk + // row-by-row. This is the branch production traffic takes. + let arrays: Vec = args + .args + .iter() + .map(|v| v.clone().into_array(n)) + .collect::>()?; + let columns: Vec<&datafusion::arrow::array::StringArray> = + arrays.iter().map(as_utf8_array).collect::>()?; + + let mut b = StringBuilder::with_capacity(n, n * 16); + let mut path_buf: Vec> = Vec::with_capacity(columns.len() - 1); + for i in 0..n { + let doc = cell(columns[0], i); + path_buf.clear(); + for col in &columns[1..] { + path_buf.push(cell(col, i)); + } + match extract(doc, &path_buf) { + 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)) + } +} + +/// Core extraction. Returns `None` for the legacy NULL-producing cases +/// (any-null arg, malformed doc, malformed path, no match, explicit-null match) +/// and a `Some(String)` for every matched case. +fn extract(doc: Option<&str>, paths: &[Option<&str>]) -> Option { + let doc_str = doc?; + if paths.iter().any(|p| p.is_none()) { + return None; + } + let parsed = parse(doc_str)?; + let per_path: Vec = paths + .iter() + .map(|p| extract_one(&parsed, p.unwrap()).unwrap_or(Value::Null)) + .collect(); + if paths.len() == 1 { + // Single path: NULL-element collapses to a SQL NULL (legacy's + // `queryResult != null ? queryResult : valueResult` returns null). + // Scalar matches unwrap to their string form via `jsonize_single`. + match per_path.into_iter().next()? { + Value::Null => None, + v => Some(jsonize_single(v)), + } + } else { + // Multi-path: wrap in JSON array. NULL misses land as literal `null` + // elements (testJsonExtractMultiPathWithMissingPath). + serde_json::to_string(&Value::Array(per_path)).ok() + } +} + +/// Evaluate a single PPL path against a parsed document. Returns `None` for +/// malformed-path, no-match, and explicit-null matches (the three cases the +/// legacy Calcite pair resolves to SQL NULL). Single-match returns the raw +/// `Value`; multi-match returns `Value::Array(...)`. +fn extract_one(doc: &Value, path: &str) -> Option { + let jsonpath = convert_ppl_path(path).ok()?; + let compiled = JsonPath::try_from(jsonpath.as_str()).ok()?; + let slice = compiled.find_slice(doc); + let matches: Vec = slice + .into_iter() + .filter_map(|v| match v { + JsonPathValue::Slice(r, _) => Some(r.clone()), + _ => None, + }) + .collect(); + match matches.len() { + 0 => None, + 1 => match matches.into_iter().next().unwrap() { + Value::Null => None, + v => Some(v), + }, + _ => Some(Value::Array(matches)), + } +} + +/// Legacy `doJsonize` single-path output: strings emerge unquoted; every other +/// JSON value (numbers, bools, arrays, objects) is serialized. Matches the +/// legacy `isScalarObject` branch (`.toString()` on Java scalars → same bytes +/// as `serde_json::to_string` for numbers and booleans). +fn jsonize_single(v: Value) -> String { + match v { + Value::String(s) => s, + other => other.to_string(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn parsed(s: &str) -> Value { + serde_json::from_str(s).unwrap() + } + + #[test] + fn single_path_scalar_match_returns_tostring_form() { + let doc = parsed(r#"{"a":801.0,"b":"hi","c":true,"d":42}"#); + assert_eq!(extract_one(&doc, "a").map(jsonize_single).unwrap(), "801.0"); + assert_eq!(extract_one(&doc, "b").map(jsonize_single).unwrap(), "hi"); + assert_eq!(extract_one(&doc, "c").map(jsonize_single).unwrap(), "true"); + assert_eq!(extract_one(&doc, "d").map(jsonize_single).unwrap(), "42"); + } + + #[test] + fn single_path_container_match_is_jsonized() { + let doc = parsed(r#"{"a":{"x":1,"y":2}}"#); + assert_eq!( + extract_one(&doc, "a").map(jsonize_single).unwrap(), + r#"{"x":1,"y":2}"# + ); + } + + #[test] + fn wildcard_multi_match_wraps_in_array() { + let doc = parsed(r#"{"a":[{"t":"A"},{"t":"B"},{"t":"C"}]}"#); + let v = extract_one(&doc, "a{}.t").unwrap(); + assert_eq!(serde_json::to_string(&v).unwrap(), r#"["A","B","C"]"#); + } + + #[test] + fn missing_and_explicit_null_both_yield_none() { + let doc = parsed(r#"{"a":null}"#); + assert!(extract_one(&doc, "a").is_none()); + assert!(extract_one(&doc, "missing").is_none()); + } + + #[test] + fn multi_path_wraps_with_null_slots_for_misses() { + let doc = parsed(r#"{"name":"John"}"#); + let out = extract(Some(r#"{"name":"John"}"#), &[Some("name"), Some("age")]); + assert_eq!(out.as_deref(), Some(r#"["John",null]"#)); + // No-op parse path so the call shape matches the legacy IT input. + assert!(doc.is_object()); + } + + #[test] + fn less_than_two_args_returns_none_via_fast_path() { + // Exercised via the UDF entry point in the integration tests; the + // core `extract` helper is only called with ≥1 path, so we just + // verify the single-path path works end-to-end here. + let out = extract(Some(r#"{"a":1}"#), &[Some("a")]); + assert_eq!(out.as_deref(), Some("1")); + } + + #[test] + fn any_null_arg_returns_none() { + assert!(extract(None, &[Some("a")]).is_none()); + assert!(extract(Some(r#"{"a":1}"#), &[None]).is_none()); + } + + #[test] + fn malformed_document_returns_none() { + assert!(extract(Some("not-json"), &[Some("a")]).is_none()); + } + + #[test] + fn malformed_path_returns_none() { + // Unmatched `{` bubbles out as None (legacy would also emit NULL via + // the PLAN-error swallow in the stateful function). + assert!(extract(Some(r#"{"a":1}"#), &[Some("a{0")]).is_none()); + } + + #[test] + fn coerce_types_enforces_string_on_every_slot() { + let udf = JsonExtractUdf::new(); + assert_eq!( + udf.coerce_types(&[DataType::LargeUtf8, DataType::Utf8View]) + .unwrap(), + vec![DataType::Utf8, DataType::Utf8] + ); + let err = udf + .coerce_types(&[DataType::Utf8, DataType::Int32]) + .unwrap_err() + .to_string(); + assert!(err.contains("expected string")); + } + + #[test] + fn return_type_is_utf8() { + assert_eq!( + JsonExtractUdf::new() + .return_type(&[DataType::Utf8]) + .unwrap(), + DataType::Utf8 + ); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_keys.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_keys.rs new file mode 100644 index 0000000000000..5971f45aff3c9 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_keys.rs @@ -0,0 +1,155 @@ +/* + * 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_keys(value)` — top-level keys of a JSON object, encoded as a JSON-array +//! string (parity with legacy `JsonKeysFunctionImpl` → Calcite +//! `JsonFunctions.jsonKeys`). Non-object / malformed / NULL input → NULL. + +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, check_arity, parse}; +use super::{coerce_args, CoerceMode}; + +const NAME: &str = "json_keys"; + +pub fn register_all(ctx: &SessionContext) { + ctx.register_udf(ScalarUDF::from(JsonKeysUdf::new())); +} + +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct JsonKeysUdf { + signature: Signature, +} + +impl JsonKeysUdf { + pub fn new() -> Self { + Self { signature: Signature::user_defined(Volatility::Immutable) } + } +} + +impl Default for JsonKeysUdf { + fn default() -> Self { + Self::new() + } +} + +impl ScalarUDFImpl for JsonKeysUdf { + 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 { + Ok(DataType::Utf8) + } + fn coerce_types(&self, args: &[DataType]) -> Result> { + coerce_args(NAME, args, &[CoerceMode::Utf8]) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + check_arity(NAME, args.args.len(), 1)?; + let n = args.number_rows; + + if let ColumnarValue::Scalar(sv) = &args.args[0] { + let keys = match sv { + ScalarValue::Utf8(Some(s)) | ScalarValue::LargeUtf8(Some(s)) | ScalarValue::Utf8View(Some(s)) => { + json_keys(s) + } + _ => None, + }; + return Ok(ColumnarValue::Scalar(ScalarValue::Utf8(keys))); + } + + let arr = args.args[0].clone().into_array(n)?; + let strings = as_utf8_array(&arr)?; + let mut b = StringBuilder::with_capacity(n, n * 16); + for i in 0..n { + if strings.is_null(i) { + b.append_null(); + continue; + } + match json_keys(strings.value(i)) { + Some(s) => b.append_value(&s), + None => b.append_null(), + } + } + Ok(ColumnarValue::Array(Arc::new(b.finish()) as ArrayRef)) + } +} + +/// Returns the JSON-array-encoded list of top-level keys for an object input, +/// or `None` for malformed / non-object / scalar / array inputs. Matches the +/// legacy contract; `serde_json::Map` is order-preserving by default (via the +/// `preserve_order` feature disabled — insertion order on BTreeMap is +/// alphabetical. Tests assert the observed ordering rather than insertion +/// order to avoid coupling to a crate feature flag. +fn json_keys(s: &str) -> Option { + match parse(s)? { + Value::Object(map) => { + let keys: Vec<&String> = map.keys().collect(); + serde_json::to_string(&keys).ok() + } + _ => None, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn object_input_returns_jsonized_keys() { + assert_eq!( + json_keys(r#"{"f1":"abc","f2":{"f3":"a"}}"#).as_deref(), + Some(r#"["f1","f2"]"#) + ); + assert_eq!(json_keys(r#"{}"#).as_deref(), Some(r#"[]"#)); + } + + #[test] + fn non_object_returns_none() { + assert_eq!(json_keys(r#"[1,2,3]"#), None); + assert_eq!(json_keys(r#"42"#), None); + assert_eq!(json_keys(r#""scalar""#), None); + assert_eq!(json_keys(r#"null"#), None); + } + + #[test] + fn malformed_returns_none() { + assert_eq!(json_keys(""), None); + assert_eq!(json_keys("{not-json"), None); + } + + #[test] + fn return_type_is_utf8() { + assert_eq!(JsonKeysUdf::new().return_type(&[DataType::Utf8]).unwrap(), DataType::Utf8); + } + + #[test] + fn coerce_types_enforces_string_arity() { + let udf = JsonKeysUdf::new(); + assert_eq!(udf.coerce_types(&[DataType::LargeUtf8]).unwrap(), vec![DataType::Utf8]); + assert!(udf.coerce_types(&[DataType::Int64]).unwrap_err().to_string().contains("expected string")); + assert!(udf.coerce_types(&[]).is_err()); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_set.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_set.rs new file mode 100644 index 0000000000000..daf33e8d2f6bb --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_set.rs @@ -0,0 +1,281 @@ +/* + * 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_set(value, path1, val1, [path2, val2, ...])` — replace the value at +//! each path match (parity with legacy `JsonSetFunctionImpl` → Jayway +//! `ctx.set` guarded by `ctx.read != null`: *replace-only*, never inserts). +//! Missing paths are no-ops; any-NULL-arg / odd trailing arg / malformed-doc / +//! malformed-path → NULL. +//! +//! Values always store as JSON strings because every UDF arg is coerced to +//! Utf8 upstream — matching the legacy fixture `"b":"3"` (not `"b":3`). + +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_set"; + +pub fn register_all(ctx: &SessionContext) { + ctx.register_udf(ScalarUDF::from(JsonSetUdf::new())); +} + +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct JsonSetUdf { + signature: Signature, +} + +impl JsonSetUdf { + pub fn new() -> Self { + Self { + signature: Signature::user_defined(Volatility::Immutable), + } + } +} + +impl Default for JsonSetUdf { + fn default() -> Self { + Self::new() + } +} + +impl ScalarUDFImpl for JsonSetUdf { + 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 { + Ok(DataType::Utf8) + } + fn coerce_types(&self, args: &[DataType]) -> Result> { + args.iter() + .enumerate() + .map(|(i, ty)| coerce_slot(NAME, i, ty, CoerceMode::Utf8)) + .collect() + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + // Need doc + at least one (path, value) pair. Odd trailing arg (unpaired + // path) short-circuits to NULL — matches the legacy `for (i=1; i> = args.args[1..].iter().map(scalar_utf8).collect(); + let out = set(doc, &rest); + return Ok(ColumnarValue::Scalar(ScalarValue::Utf8(out))); + } + + let arrays: Vec = args + .args + .iter() + .map(|v| v.clone().into_array(n)) + .collect::>()?; + let columns: Vec<&datafusion::arrow::array::StringArray> = + arrays.iter().map(as_utf8_array).collect::>()?; + + let mut b = StringBuilder::with_capacity(n, n * 16); + let mut rest: Vec> = 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 set(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`. Replace-only: +/// missing paths are no-ops, matching legacy `jsonSet`'s `ctx.read != null` +/// guard. +fn set(doc: Option<&str>, rest: &[Option<&str>]) -> Option { + 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() { + // Setting the root is a Jayway no-op (root is indelible and + // unreplaceable via `ctx.set("$", v)`). Mirror that. + continue; + } + set_one(&mut value, &segments, new_val); + } + serde_json::to_string(&value).ok() +} + +fn set_one(root: &mut Value, segments: &[Segment<'_>], new_val: &str) { + let replacement = Value::String(new_val.to_string()); + walk_mut(root, segments, |parent, final_seg| { + match (parent, final_seg) { + // Replace-only: only overwrite if the key already exists. + (Value::Object(map), Segment::Field(name)) if map.contains_key(*name) => { + map.insert((*name).to_string(), replacement.clone()); + } + (Value::Array(arr), Segment::Index(i)) if *i < arr.len() => { + arr[*i] = replacement.clone(); + } + (Value::Array(arr), Segment::Wildcard) => { + for slot in arr.iter_mut() { + *slot = replacement.clone(); + } + } + // Type mismatch is a silent no-op (legacy SUPPRESS_EXCEPTIONS). + _ => {} + } + }); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn wildcard_replace_matches_legacy_fixture() { + // testJsonSet + assert_eq!( + set( + Some(r#"{"a":[{"b":1},{"b":2}]}"#), + &[Some("a{}.b"), Some("3")], + ) + .as_deref(), + Some(r#"{"a":[{"b":"3"},{"b":"3"}]}"#) + ); + } + + #[test] + fn wrong_path_leaves_input_unchanged() { + // testJsonSetWithWrongPath — 'a{}.b.d' doesn't exist (b is a scalar). + assert_eq!( + set( + Some(r#"{"a":[{"b":1},{"b":2}]}"#), + &[Some("a{}.b.d"), Some("3")], + ) + .as_deref(), + Some(r#"{"a":[{"b":1},{"b":2}]}"#) + ); + } + + #[test] + fn partial_wildcard_match_only_sets_where_path_exists() { + // testJsonSetPartialSet + assert_eq!( + set( + Some(r#"{"a":[{"b":1},{"b":{"c":2}}]}"#), + &[Some("a{}.b.c"), Some("3")], + ) + .as_deref(), + Some(r#"{"a":[{"b":1},{"b":{"c":"3"}}]}"#) + ); + } + + #[test] + fn multiple_path_value_pairs_apply_sequentially() { + assert_eq!( + set( + Some(r#"{"a":1,"b":2}"#), + &[Some("a"), Some("10"), Some("b"), Some("20")], + ) + .as_deref(), + Some(r#"{"a":"10","b":"20"}"#) + ); + } + + #[test] + fn any_null_arg_returns_none() { + assert!(set(None, &[Some("a"), Some("v")]).is_none()); + assert!(set(Some(r#"{"a":1}"#), &[None, Some("v")]).is_none()); + assert!(set(Some(r#"{"a":1}"#), &[Some("a"), None]).is_none()); + } + + #[test] + fn malformed_doc_returns_none() { + assert!(set(Some("not-json"), &[Some("a"), Some("v")]).is_none()); + } + + #[test] + fn malformed_path_returns_none() { + assert!(set(Some(r#"{"a":1}"#), &[Some("a{"), Some("v")]).is_none()); + } + + #[test] + fn coerce_types_enforces_string_on_every_slot() { + let udf = JsonSetUdf::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!( + JsonSetUdf::new().return_type(&[DataType::Utf8]).unwrap(), + DataType::Utf8 + ); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/mod.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/mod.rs index 2053625b544c0..cdbcb9db9bf83 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/mod.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/mod.rs @@ -7,11 +7,12 @@ */ //! OpenSearch scalar UDFs that aren't in DataFusion's built-in registry. Each -//! must have a matching YAML entry in `extensions/opensearch_scalar.yaml` so +//! must have a matching YAML entry in `opensearch_scalar_functions.yaml` so //! the substrait converter on the Java side can route to it by name. //! //! Functions registered here: //! - `convert_tz(ts, from_tz, to_tz)` — DST-aware timezone shift (chrono-tz) +//! - `json_array_length(value)` — length of a JSON array, NULL on malformed/non-array use datafusion::arrow::datatypes::{DataType, TimeUnit}; use datafusion::common::plan_err; @@ -113,14 +114,38 @@ pub(crate) fn coerce_args( } pub mod convert_tz; +pub mod json_append; +pub mod json_array_length; +pub(crate) mod json_common; +pub mod json_delete; +pub mod json_extend; +pub mod json_extract; +pub mod json_keys; +pub mod json_set; pub mod tonumber; pub mod tostring; +// Dev note: if a freshly added UDF here fails at runtime with +// "Unsupported function name: " despite the Java side being wired, the +// native dylib is stale. `sandbox/libs/dataformat-native/build.gradle` tracks +// only common/ + lib/ as inputs, so plugin-side Rust edits leave gradle +// UP-TO-DATE. Workaround: run +// `./gradlew :sandbox:libs:dataformat-native:buildRustLibrary --rerun-tasks` +// and restart the OpenSearch JVM (the loaded dylib is JVM-cached). pub fn register_all(ctx: &SessionContext) { convert_tz::register_all(ctx); + json_append::register_all(ctx); + json_array_length::register_all(ctx); + json_delete::register_all(ctx); + json_extend::register_all(ctx); + json_extract::register_all(ctx); + json_keys::register_all(ctx); + json_set::register_all(ctx); tonumber::register_all(ctx); tostring::register_all(ctx); - log::info!("OpenSearch UDF register_all: convert_tz, tonumber, tostring registered"); + log::info!( + "OpenSearch UDF register_all: convert_tz, json_append, json_array_length, json_delete, json_extend, json_extract, json_keys, json_set, tonumber, tostring registered" + ); } #[cfg(test)] diff --git a/sandbox/plugins/analytics-backend-datafusion/src/internalClusterTest/java/org/opensearch/be/datafusion/BaseScalarFunctionIT.java b/sandbox/plugins/analytics-backend-datafusion/src/internalClusterTest/java/org/opensearch/be/datafusion/BaseScalarFunctionIT.java index dd3dbcb1bc0bb..0c348bccfd9e1 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/internalClusterTest/java/org/opensearch/be/datafusion/BaseScalarFunctionIT.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/internalClusterTest/java/org/opensearch/be/datafusion/BaseScalarFunctionIT.java @@ -121,6 +121,12 @@ private void createBankIndex() throws Exception { .startObject("created_at") .field("type", "date") .endObject() + // json_str holds serialized JSON arrays/objects/malformed strings so + // scalar-JSON UDFs can be exercised on real column values (columnar + // UDF path), not just string literals (scalar fast-path). + .startObject("json_str") + .field("type", "keyword") + .endObject() .endObject() .endObject(); @@ -143,13 +149,38 @@ private void createBankIndex() throws Exception { } private void indexBankDocs() { + // Row 1 carries a 3-element JSON array in json_str; row 6 carries a JSON object. + // This lets scalar-JSON UDF tests assert both the happy path (row 1 → length 3) + // and the non-array → NULL path (row 6) from real column values. client().prepareIndex(BANK_INDEX) .setId("1") - .setSource("account_number", 1, "firstname", "Amber", "balance", 39225L, "created_at", "2024-06-15T10:30:00Z") + .setSource( + "account_number", + 1, + "firstname", + "Amber", + "balance", + 39225L, + "created_at", + "2024-06-15T10:30:00Z", + "json_str", + "[1,2,3]" + ) .get(); client().prepareIndex(BANK_INDEX) .setId("6") - .setSource("account_number", 6, "firstname", "Hattie", "balance", 5686L, "created_at", "2024-01-20T14:45:30Z") + .setSource( + "account_number", + 6, + "firstname", + "Hattie", + "balance", + 5686L, + "created_at", + "2024-01-20T14:45:30Z", + "json_str", + "{\"k\":1}" + ) .get(); } @@ -182,6 +213,32 @@ protected void assertScalarLong(String expr, long expected) { assertEquals(expr, expected, ((Number) cell).longValue()); } + /** + * Strict variant that asserts the cell is a {@link Long} (not just a {@link Number}). + * Use for functions whose on-wire BIGINT return type must not silently regress. + */ + protected void assertScalarLongStrict(String expr, long expected) { + Object cell = evalScalar(expr); + assertNotNull(expr + " result must not be null", cell); + assertTrue(expr + " result must be Long, got " + cell.getClass(), cell instanceof Long); + assertEquals(expr, expected, ((Long) cell).longValue()); + } + + /** + * Strict variant that asserts the cell is an {@link Integer}. Use for functions + * whose on-wire INTEGER return type must be preserved through the pipeline — + * e.g. PPL scalar UDFs declared as {@code INTEGER_FORCE_NULLABLE} whose Rust + * implementations return {@code Int64} but get narrowed via an implicit CAST + * on the enclosing Project. The non-strict {@link #assertScalarLong} silently + * accepts either width and would miss this contract regression. + */ + protected void assertScalarIntStrict(String expr, int expected) { + Object cell = evalScalar(expr); + assertNotNull(expr + " result must not be null", cell); + assertTrue(expr + " result must be Integer, got " + cell.getClass(), cell instanceof Integer); + assertEquals(expr, expected, ((Integer) cell).intValue()); + } + protected void assertScalarDouble(String expr, double expected, double delta) { Object cell = evalScalar(expr); assertNotNull(expr + " result must not be null", cell); diff --git a/sandbox/plugins/analytics-backend-datafusion/src/internalClusterTest/java/org/opensearch/be/datafusion/ScalarJsonFunctionIT.java b/sandbox/plugins/analytics-backend-datafusion/src/internalClusterTest/java/org/opensearch/be/datafusion/ScalarJsonFunctionIT.java new file mode 100644 index 0000000000000..67214bb09a724 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/internalClusterTest/java/org/opensearch/be/datafusion/ScalarJsonFunctionIT.java @@ -0,0 +1,180 @@ +/* + * 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. + */ + +package org.opensearch.be.datafusion; + +/** + * End-to-end smoke tests for PPL {@code json_*} scalar functions routed through + * PPL → Calcite → Substrait → DataFusion. One method per function for happy-path + * + column-valued coverage; {@code *ParityWithLegacy} methods replay the legacy + * SQL plugin's {@code CalcitePPLJsonBuiltinFunctionIT} fixtures verbatim. + * Edge cases are covered in Rust unit tests. + */ +public class ScalarJsonFunctionIT extends BaseScalarFunctionIT { + + /** Happy path + NULL-on-non-array/malformed (scalar fast-path) + column-valued (Arrow columnar path). */ + public void testJsonArrayLength() { + assertScalarIntStrict("json_array_length('[1,2,3]')", 3); + assertScalarIntStrict("json_array_length('[]')", 0); + assertScalarNull("json_array_length('{\"k\":1}')"); + assertScalarNull("json_array_length('not-json')"); + // Columnar path: bank fixture's json_str row 1 is '[1,2,3]'. + assertScalarIntStrict("json_array_length(json_str)", 3); + } + + /** Parity replay of {@code CalcitePPLJsonBuiltinFunctionIT.testJsonArrayLength}. */ + public void testJsonArrayLengthParityWithLegacy() { + assertScalarIntStrict("json_array_length('[1,2,3,4]')", 4); + assertScalarIntStrict("json_array_length('[1,2,3,{\"f1\":1,\"f2\":[5,6]},4]')", 5); + assertScalarNull("json_array_length('{\"key\": 1}')"); + } + + /** Parity replay of {@code CalcitePPLJsonBuiltinFunctionIT.testJsonKeys} — insertion order preserved via {@code serde_json} {@code preserve_order}. */ + public void testJsonKeysParityWithLegacy() { + assertScalarString("json_keys('{\"f1\":\"abc\",\"f2\":{\"f3\":\"a\",\"f4\":\"b\"}}')", "[\"f1\",\"f2\"]"); + assertScalarNull("json_keys('[1,2,3,{\"f1\":1,\"f2\":[5,6]},4]')"); + assertScalarNull("json_keys('not-json')"); + assertScalarNull("json_keys('42')"); + } + + /** Parity replay of {@code CalcitePPLJsonBuiltinFunctionIT.testJsonExtract*} — byte-for-byte match via {@code serde_json} {@code preserve_order} + no integer↔double coercion. */ + public void testJsonExtractParityWithLegacy() { + String candidate = "[{\"name\":\"London\",\"Bridges\":[{\"name\":\"Tower Bridge\",\"length\":801.0}," + + "{\"name\":\"Millennium Bridge\",\"length\":1066.0}]}," + + "{\"name\":\"Venice\",\"Bridges\":[{\"name\":\"Rialto Bridge\",\"length\":157.0}," + + "{\"type\":\"Bridge of Sighs\",\"length\":36.0}," + + "{\"type\":\"Ponte della Paglia\"}]}," + + "{\"name\":\"San Francisco\",\"Bridges\":[{\"name\":\"Golden Gate Bridge\",\"length\":8981.0}," + + "{\"name\":\"Bay Bridge\",\"length\":23556.0}]}]"; + + // Single-path, wildcard-at-root over top-level array → 3 matches wrapped + // in a JSON array. Round-tripped bytes equal the input because + // preserve_order + no numeric coercion. + assertScalarString("json_extract('" + candidate + "', '{}')", candidate); + + // Single-path scalar match — legacy `.toString()` on Double(8981.0). + assertScalarString("json_extract('" + candidate + "', '{2}.Bridges{0}.length')", "8981.0"); + + // Wildcard-over-wildcard-missing-key: only Venice entries without a + // `name` field expose a `type`, so two matches wrap into a JSON array. + assertScalarString("json_extract('" + candidate + "', '{}.Bridges{}.type')", "[\"Bridge of Sighs\",\"Ponte della Paglia\"]"); + + // Single-path object match — jsonized with insertion order preserved. + assertScalarString("json_extract('" + candidate + "', '{2}.Bridges{0}')", "{\"name\":\"Golden Gate Bridge\",\"length\":8981.0}"); + + // Multi-path with wildcard-multi + scalar-match → outer array wraps + // the two per-path results (array + scalar) as-is. + assertScalarString( + "json_extract('" + candidate + "', '{}.Bridges{}.type', '{2}.Bridges{0}.length')", + "[[\"Bridge of Sighs\",\"Ponte della Paglia\"],8981.0]" + ); + + // Missing path (empty object) and explicit-null both resolve to SQL NULL. + assertScalarNull("json_extract('{}', 'name')"); + assertScalarNull("json_extract('{\"name\": null}', 'name')"); + + // Multi-path with missing path yields literal `null` element in the + // outer JSON array. + assertScalarString("json_extract('{\"name\": \"John\"}', 'name', 'age')", "[\"John\",null]"); + } + + /** Parity replay of {@code CalcitePPLJsonBuiltinFunctionIT.testJsonSet*} — values stored as JSON strings matches legacy {@code "b":"3"} outputs (Utf8 arg coercion). */ + public void testJsonSetParityWithLegacy() { + // testJsonSet: wildcard replace across every array element. + assertScalarString("json_set('{\"a\":[{\"b\":1},{\"b\":2}]}', 'a{}.b', '3')", "{\"a\":[{\"b\":\"3\"},{\"b\":\"3\"}]}"); + + // testJsonSetWithWrongPath: 'a{}.b.d' doesn't exist — input unchanged. + assertScalarString("json_set('{\"a\":[{\"b\":1},{\"b\":2}]}', 'a{}.b.d', '3')", "{\"a\":[{\"b\":1},{\"b\":2}]}"); + + // testJsonSetPartialSet: wildcard where only one branch has the full + // path; only the matching branch is rewritten. + assertScalarString( + "json_set('{\"a\":[{\"b\":1},{\"b\":{\"c\":2}}]}', 'a{}.b.c', '3')", + "{\"a\":[{\"b\":1},{\"b\":{\"c\":\"3\"}}]}" + ); + } + + /** Parity replay of {@code CalcitePPLJsonBuiltinFunctionIT.testJsonAppend} — nested {@code json_object}/{@code json_array} constructors replaced with their stringified equivalents (same observable contract). */ + public void testJsonAppendParityWithLegacy() { + // Case a: pre-stringified json_object(...) appended as a single array element. + assertScalarString( + "json_append('{\"teacher\":[\"Alice\"],\"student\":[{\"name\":\"Bob\",\"rank\":1},{\"name\":\"Charlie\",\"rank\":2}]}'," + + " 'student', '{\"name\":\"Tomy\",\"rank\":5}')", + "{\"teacher\":[\"Alice\"],\"student\":[{\"name\":\"Bob\",\"rank\":1},{\"name\":\"Charlie\",\"rank\":2}," + + "\"{\\\"name\\\":\\\"Tomy\\\",\\\"rank\\\":5}\"]}" + ); + + // Case b: multi-pair append on the same target. + assertScalarString( + "json_append('{\"teacher\":[\"Alice\"],\"student\":[{\"name\":\"Bob\",\"rank\":1},{\"name\":\"Charlie\",\"rank\":2}]}'," + + " 'teacher', 'Tom', 'teacher', 'Walt')", + "{\"teacher\":[\"Alice\",\"Tom\",\"Walt\"],\"student\":[{\"name\":\"Bob\",\"rank\":1},{\"name\":\"Charlie\",\"rank\":2}]}" + ); + + // Case c: nested-path + pre-stringified json_array(...) appended as a single string element. + assertScalarString( + "json_append('{\"school\":{\"teacher\":[\"Alice\"],\"student\":[{\"name\":\"Bob\",\"rank\":1},{\"name\":\"Charlie\",\"rank\":2}]}}'," + + " 'school.teacher', '[\"Tom\",\"Walt\"]')", + "{\"school\":{\"teacher\":[\"Alice\",\"[\\\"Tom\\\",\\\"Walt\\\"]\"]," + + "\"student\":[{\"name\":\"Bob\",\"rank\":1},{\"name\":\"Charlie\",\"rank\":2}]}}" + ); + } + + /** Parity replay of {@code CalcitePPLJsonBuiltinFunctionIT.testJsonExtend} — case c diverges from append: a JSON-array value is spread (not pushed as a single element). */ + public void testJsonExtendParityWithLegacy() { + // Case a: stringified json_object value — not a JSON array → single push. + assertScalarString( + "json_extend('{\"teacher\":[\"Alice\"],\"student\":[{\"name\":\"Bob\",\"rank\":1},{\"name\":\"Charlie\",\"rank\":2}]}'," + + " 'student', '{\"name\":\"Tommy\",\"rank\":5}')", + "{\"teacher\":[\"Alice\"],\"student\":[{\"name\":\"Bob\",\"rank\":1},{\"name\":\"Charlie\",\"rank\":2}," + + "\"{\\\"name\\\":\\\"Tommy\\\",\\\"rank\\\":5}\"]}" + ); + + // Case b: plain strings — each fails List-parse → each pushed individually. + assertScalarString( + "json_extend('{\"teacher\":[\"Alice\"],\"student\":[{\"name\":\"Bob\",\"rank\":1},{\"name\":\"Charlie\",\"rank\":2}]}'," + + " 'teacher', 'Tom', 'teacher', 'Walt')", + "{\"teacher\":[\"Alice\",\"Tom\",\"Walt\"],\"student\":[{\"name\":\"Bob\",\"rank\":1},{\"name\":\"Charlie\",\"rank\":2}]}" + ); + + // Case c: stringified json_array — parses as JSON array → elements spread. + assertScalarString( + "json_extend('{\"school\":{\"teacher\":[\"Alice\"],\"student\":[{\"name\":\"Bob\",\"rank\":1},{\"name\":\"Charlie\",\"rank\":2}]}}'," + + " 'school.teacher', '[\"Tom\",\"Walt\"]')", + "{\"school\":{\"teacher\":[\"Alice\",\"Tom\",\"Walt\"]," + + "\"student\":[{\"name\":\"Bob\",\"rank\":1},{\"name\":\"Charlie\",\"rank\":2}]}}" + ); + } + + /** Parity replay of {@code CalcitePPLJsonBuiltinFunctionIT.testJsonDelete*} — output order preserved via {@code serde_json} {@code preserve_order}. */ + public void testJsonDeleteParityWithLegacy() { + // testJsonDelete: flat-key delete of two fields. + assertScalarString( + "json_delete('{\"account_number\":1,\"balance\":39225,\"age\":32,\"gender\":\"M\"}', 'age', 'gender')", + "{\"account_number\":1,\"balance\":39225}" + ); + + // testJsonDeleteWithNested: delete a single nested key. + assertScalarString( + "json_delete('{\"f1\":\"abc\",\"f2\":{\"f3\":\"a\",\"f4\":\"b\"}}', 'f2.f3')", + "{\"f1\":\"abc\",\"f2\":{\"f4\":\"b\"}}" + ); + + // testJsonDeleteWithNestedNothing: missing nested key leaves input unchanged. + assertScalarString( + "json_delete('{\"f1\":\"abc\",\"f2\":{\"f3\":\"a\",\"f4\":\"b\"}}', 'f2.f100')", + "{\"f1\":\"abc\",\"f2\":{\"f3\":\"a\",\"f4\":\"b\"}}" + ); + + // testJsonDeleteWithNestedAndArray: wildcard path drops one key from every array element. + assertScalarString( + "json_delete('{\"teacher\":\"Alice\",\"student\":[{\"name\":\"Bob\",\"rank\":1},{\"name\":\"Charlie\",\"rank\":2}]}', 'teacher', 'student{}.rank')", + "{\"student\":[{\"name\":\"Bob\"},{\"name\":\"Charlie\"}]}" + ); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/ConvertTzAdapter.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/ConvertTzAdapter.java index cb460333dbdaa..d123bf15e78f5 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/ConvertTzAdapter.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/ConvertTzAdapter.java @@ -33,7 +33,7 @@ import java.util.regex.Pattern; /** - * Cat-3b adapter for PPL's {@code CONVERT_TZ(ts, from_tz, to_tz)}. Two jobs in + * Adapter for PPL's {@code CONVERT_TZ(ts, from_tz, to_tz)}. Two jobs in * priority order: * *
    diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java index a7414eac65639..4e5b5955cb92c 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java @@ -93,9 +93,11 @@ public class DataFusionAnalyticsBackendPlugin implements AnalyticsSearchBackendP // path). COALESCE is the lowering target of PPL `fillnull`. CAST is required because // ReduceExpressionsRule.ProjectReduceExpressionsRule (in PlannerImpl) constant-folds field // references through equality filters into typed literals — e.g. after `where str0 = 'FURNITURE'`, - // the projection `fields str0` is rewritten to `CAST('FURNITURE' AS VARCHAR)`. CONCAT is the - // lowering target of PPL `eval`'s `+` for strings (Calcite emits `||`, resolved to CONCAT in - // ScalarFunction); SAFE_CAST covers PPL `eval`'s explicit nullable `CAST(... AS ...)` + // the projection `fields str0` is rewritten to `CAST('FURNITURE' AS VARCHAR)`. CAST is also the + // implicit result-type narrowing PPL inserts after a UDF call whose declared return type differs + // from the eval column's inferred type (e.g. JSON_ARRAY_LENGTH returns INTEGER_FORCE_NULLABLE). + // CONCAT is the lowering target of PPL `eval`'s `+` for strings (Calcite emits `||`, resolved to + // CONCAT in ScalarFunction); SAFE_CAST covers PPL `eval`'s explicit nullable `CAST(... AS ...)` // expressions. The remaining comparison / arithmetic / logical operators are project-capable // for eval-style projections. private static final Set STANDARD_PROJECT_OPS = Set.of( @@ -219,7 +221,14 @@ public class DataFusionAnalyticsBackendPlugin implements AnalyticsSearchBackendP ScalarFunction.STRCMP, ScalarFunction.TOSTRING, ScalarFunction.NUMBER_TO_STRING, - ScalarFunction.TONUMBER + ScalarFunction.TONUMBER, + ScalarFunction.JSON_APPEND, + ScalarFunction.JSON_ARRAY_LENGTH, + ScalarFunction.JSON_DELETE, + ScalarFunction.JSON_EXTEND, + ScalarFunction.JSON_EXTRACT, + ScalarFunction.JSON_KEYS, + ScalarFunction.JSON_SET ); private static final Set AGG_FUNCTIONS = Set.of( @@ -327,6 +336,13 @@ public Map scalarFunctionAdapters() { Map.entry(ScalarFunction.EXPM1, new Expm1Adapter()), Map.entry(ScalarFunction.HOUR, hour), Map.entry(ScalarFunction.HOUR_OF_DAY, hour), + Map.entry(ScalarFunction.JSON_APPEND, new JsonFunctionAdapters.JsonAppendAdapter()), + Map.entry(ScalarFunction.JSON_ARRAY_LENGTH, new JsonFunctionAdapters.JsonArrayLengthAdapter()), + Map.entry(ScalarFunction.JSON_DELETE, new JsonFunctionAdapters.JsonDeleteAdapter()), + Map.entry(ScalarFunction.JSON_EXTEND, new JsonFunctionAdapters.JsonExtendAdapter()), + Map.entry(ScalarFunction.JSON_EXTRACT, new JsonFunctionAdapters.JsonExtractAdapter()), + Map.entry(ScalarFunction.JSON_KEYS, new JsonFunctionAdapters.JsonKeysAdapter()), + Map.entry(ScalarFunction.JSON_SET, new JsonFunctionAdapters.JsonSetAdapter()), Map.entry(ScalarFunction.LIKE, new LikeAdapter()), Map.entry(ScalarFunction.LOCATE, new PositionAdapter()), Map.entry(ScalarFunction.MICROSECOND, DatePartAdapters.microsecond()), diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionFragmentConvertor.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionFragmentConvertor.java index 23f7012bb8208..a06d2a4bb20d3 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionFragmentConvertor.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionFragmentConvertor.java @@ -81,11 +81,29 @@ public class DataFusionFragmentConvertor implements FragmentConvertor { /** * Maps backend-specific Calcite operators to their Substrait extension names so Isthmus - * serializes them through our {@code SimpleExtension} catalog: + * serializes them through our {@code SimpleExtension} catalog. One entry per line so + * parallel per-UDF PRs append without hotspot conflicts. *
      *
    • {@link DelegatedPredicateFunction} → {@code delegated_predicate} (delegation to a peer backend).
      • *
    • {@link SqlLibraryOperators#ILIKE} → {@code ilike} (case-insensitive LIKE; resolved by * DataFusion's substrait consumer to a case-insensitive {@code LikeExpr}).
      • + *
    • {@link SqlLibraryOperators#DATE_PART} → {@code date_part} (target of YearAdapter's rewrite).
      • + *
    • {@link ConvertTzAdapter#LOCAL_CONVERT_TZ_OP} → {@code convert_tz} (Rust UDF).
      • + *
    • {@link UnixTimestampAdapter#LOCAL_TO_UNIXTIME_OP} → {@code to_unixtime} (DF native).
      • + *
    • {@link JsonFunctionAdapters.JsonAppendAdapter#LOCAL_JSON_APPEND_OP} → + * {@code json_append} (Rust UDF, homogeneous-string variadic path/value pairs).
      • + *
    • {@link JsonFunctionAdapters.JsonArrayLengthAdapter#LOCAL_JSON_ARRAY_LENGTH_OP} → + * {@code json_array_length} (Rust UDF).
      • + *
    • {@link JsonFunctionAdapters.JsonDeleteAdapter#LOCAL_JSON_DELETE_OP} → + * {@code json_delete} (Rust UDF, homogeneous-string variadic).
      • + *
    • {@link JsonFunctionAdapters.JsonExtendAdapter#LOCAL_JSON_EXTEND_OP} → + * {@code json_extend} (Rust UDF, homogeneous-string variadic path/value pairs).
      • + *
    • {@link JsonFunctionAdapters.JsonExtractAdapter#LOCAL_JSON_EXTRACT_OP} → + * {@code json_extract} (Rust UDF, homogeneous-string variadic).
      • + *
    • {@link JsonFunctionAdapters.JsonKeysAdapter#LOCAL_JSON_KEYS_OP} → + * {@code json_keys} (Rust UDF).
      • + *
    • {@link JsonFunctionAdapters.JsonSetAdapter#LOCAL_JSON_SET_OP} → + * {@code json_set} (Rust UDF, homogeneous-string variadic path/value pairs).
      • *
    • {@link SqlLibraryOperators#REGEXP_CONTAINS} → {@code regex_match} (boolean regex match; * resolved by DataFusion's substrait consumer to {@code Operator::RegexMatch}, the same * binary operator that backs PostgreSQL's {@code ~} regex match). Lowering target for PPL @@ -125,7 +143,17 @@ public class DataFusionFragmentConvertor implements FragmentConvertor { FunctionMappings.s(SqlStdOperatorTable.PI, "pi"), FunctionMappings.s(SqlStdOperatorTable.RAND, "random"), FunctionMappings.s(SqlLibraryOperators.LOG, "logb"), - FunctionMappings.s(SignumFunction.FUNCTION, SignumFunction.NAME) + FunctionMappings.s(SignumFunction.FUNCTION, SignumFunction.NAME), + FunctionMappings.s(JsonFunctionAdapters.JsonAppendAdapter.LOCAL_JSON_APPEND_OP, "json_append"), + FunctionMappings.s(JsonFunctionAdapters.JsonArrayLengthAdapter.LOCAL_JSON_ARRAY_LENGTH_OP, "json_array_length"), + FunctionMappings.s(JsonFunctionAdapters.JsonDeleteAdapter.LOCAL_JSON_DELETE_OP, "json_delete"), + FunctionMappings.s(JsonFunctionAdapters.JsonExtendAdapter.LOCAL_JSON_EXTEND_OP, "json_extend"), + FunctionMappings.s(JsonFunctionAdapters.JsonExtractAdapter.LOCAL_JSON_EXTRACT_OP, "json_extract"), + FunctionMappings.s(JsonFunctionAdapters.JsonKeysAdapter.LOCAL_JSON_KEYS_OP, "json_keys"), + FunctionMappings.s(JsonFunctionAdapters.JsonSetAdapter.LOCAL_JSON_SET_OP, "json_set"), + FunctionMappings.s(SqlLibraryOperators.REGEXP_CONTAINS, "regex_match"), + FunctionMappings.s(SqlStdOperatorTable.REPLACE, "replace"), + FunctionMappings.s(SqlLibraryOperators.REGEXP_REPLACE_3, "regexp_replace") ); private final SimpleExtension.ExtensionCollection extensions; diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/JsonFunctionAdapters.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/JsonFunctionAdapters.java new file mode 100644 index 0000000000000..9a416de26ae8f --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/JsonFunctionAdapters.java @@ -0,0 +1,159 @@ +/* + * 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. + */ + +package org.opensearch.be.datafusion; + +import org.apache.calcite.sql.SqlFunction; +import org.apache.calcite.sql.SqlFunctionCategory; +import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.sql.SqlOperator; +import org.apache.calcite.sql.type.OperandTypes; +import org.apache.calcite.sql.type.ReturnTypes; +import org.opensearch.analytics.spi.AbstractNameMappingAdapter; + +import java.util.List; + +/** + * Container for PPL JSON-function scalar adapters. Each inner class is a plain + * name-mapping rewrite from a Calcite call to a locally-declared + * {@link SqlOperator} whose name matches the corresponding Rust UDF at + * {@code rust/src/udf/.rs}. All validation (malformed JSON, malformed + * path, arity / pairing, any-NULL propagation) lives in the Rust UDF; the + * adapter does not inspect arguments. Return type is preserved from the + * original PPL call by {@link AbstractNameMappingAdapter#adapt}, matching the + * {@code *_FORCE_NULLABLE} declaration on the legacy {@code Json*FunctionImpl}. + * + *

      Each {@code LOCAL_*_OP} must also be registered in + * {@link DataFusionFragmentConvertor#ADDITIONAL_SCALAR_SIGS} via a + * {@code FunctionMappings.s(...)} entry keyed by the UDF's name. + * + * @opensearch.internal + */ +final class JsonFunctionAdapters { + + private JsonFunctionAdapters() {} + + /** {@code JSON_ARRAY_LENGTH(value)} → length of a JSON array; NULL on non-array / malformed input. */ + static class JsonArrayLengthAdapter extends AbstractNameMappingAdapter { + + static final SqlOperator LOCAL_JSON_ARRAY_LENGTH_OP = new SqlFunction( + "json_array_length", + SqlKind.OTHER_FUNCTION, + ReturnTypes.INTEGER_NULLABLE, + null, + OperandTypes.STRING, + SqlFunctionCategory.STRING + ); + + JsonArrayLengthAdapter() { + super(LOCAL_JSON_ARRAY_LENGTH_OP, List.of(), List.of()); + } + } + + /** {@code JSON_KEYS(value)} → JSON-array-encoded top-level keys; NULL on non-object / malformed input. */ + static class JsonKeysAdapter extends AbstractNameMappingAdapter { + + static final SqlOperator LOCAL_JSON_KEYS_OP = new SqlFunction( + "json_keys", + SqlKind.OTHER_FUNCTION, + ReturnTypes.VARCHAR_NULLABLE, + null, + OperandTypes.STRING, + SqlFunctionCategory.STRING + ); + + JsonKeysAdapter() { + super(LOCAL_JSON_KEYS_OP, List.of(), List.of()); + } + } + + /** {@code JSON_EXTRACT(value, path1, [path2, ...])} — single path → stringified match; multi-path → JSON-array wrap with {@code null} slots for misses. */ + static class JsonExtractAdapter extends AbstractNameMappingAdapter { + + static final SqlOperator LOCAL_JSON_EXTRACT_OP = new SqlFunction( + "json_extract", + SqlKind.OTHER_FUNCTION, + ReturnTypes.VARCHAR_NULLABLE, + null, + OperandTypes.VARIADIC, + SqlFunctionCategory.STRING + ); + + JsonExtractAdapter() { + super(LOCAL_JSON_EXTRACT_OP, List.of(), List.of()); + } + } + + /** {@code JSON_DELETE(value, path1, [path2, ...])} — remove PPL-path matches; missing paths are no-ops. */ + static class JsonDeleteAdapter extends AbstractNameMappingAdapter { + + static final SqlOperator LOCAL_JSON_DELETE_OP = new SqlFunction( + "json_delete", + SqlKind.OTHER_FUNCTION, + ReturnTypes.VARCHAR_NULLABLE, + null, + OperandTypes.VARIADIC, + SqlFunctionCategory.STRING + ); + + JsonDeleteAdapter() { + super(LOCAL_JSON_DELETE_OP, List.of(), List.of()); + } + } + + /** {@code JSON_SET(value, path1, val1, [path2, val2, ...])} — replace-only; missing paths are no-ops (parity with legacy {@code ctx.read != null} guard). */ + static class JsonSetAdapter extends AbstractNameMappingAdapter { + + static final SqlOperator LOCAL_JSON_SET_OP = new SqlFunction( + "json_set", + SqlKind.OTHER_FUNCTION, + ReturnTypes.VARCHAR_NULLABLE, + null, + OperandTypes.VARIADIC, + SqlFunctionCategory.STRING + ); + + JsonSetAdapter() { + super(LOCAL_JSON_SET_OP, List.of(), List.of()); + } + } + + /** {@code JSON_APPEND(value, path1, val1, [path2, val2, ...])} — push-only onto array-valued targets; non-array / missing targets are no-ops. */ + static class JsonAppendAdapter extends AbstractNameMappingAdapter { + + static final SqlOperator LOCAL_JSON_APPEND_OP = new SqlFunction( + "json_append", + SqlKind.OTHER_FUNCTION, + ReturnTypes.VARCHAR_NULLABLE, + null, + OperandTypes.VARIADIC, + SqlFunctionCategory.STRING + ); + + JsonAppendAdapter() { + super(LOCAL_JSON_APPEND_OP, List.of(), List.of()); + } + } + + /** {@code JSON_EXTEND(value, path1, val1, [path2, val2, ...])} — spread-or-append: JSON-array values are spread element-wise; otherwise the whole value is pushed as one string element. */ + static class JsonExtendAdapter extends AbstractNameMappingAdapter { + + static final SqlOperator LOCAL_JSON_EXTEND_OP = new SqlFunction( + "json_extend", + SqlKind.OTHER_FUNCTION, + ReturnTypes.VARCHAR_NULLABLE, + null, + OperandTypes.VARIADIC, + SqlFunctionCategory.STRING + ); + + JsonExtendAdapter() { + super(LOCAL_JSON_EXTEND_OP, List.of(), List.of()); + } + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/UnixTimestampAdapter.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/UnixTimestampAdapter.java index 2f7056ac92c55..7acdda227eb7e 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/UnixTimestampAdapter.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/UnixTimestampAdapter.java @@ -19,7 +19,7 @@ import java.util.List; /** - * Cat-3a rename adapter for PPL's {@code UNIX_TIMESTAMP(ts)}. Rewrites to a + * Rename adapter for PPL's {@code UNIX_TIMESTAMP(ts)}. Rewrites to a * locally-declared {@link SqlFunction} named {@code to_unixtime} — the name * DataFusion's substrait consumer recognizes for its native * {@code ToUnixtimeFunc} (no UDF registration required on the Rust side). diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/resources/opensearch_scalar_functions.yaml b/sandbox/plugins/analytics-backend-datafusion/src/main/resources/opensearch_scalar_functions.yaml index 4caf7c185dc96..d5dd9f2f213a6 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/resources/opensearch_scalar_functions.yaml +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/resources/opensearch_scalar_functions.yaml @@ -317,3 +317,53 @@ scalar_functions: - { name: base, value: i32 } nullability: DECLARED_OUTPUT return: fp64 + + # PPL json_* UDFs — Rust implementations under rust/src/udf/.rs, + # surfaced to Calcite via JsonFunctionAdapters. All return NULL on malformed + # input; per-function semantics are documented in the Rust module headers. + - name: "json_array_length" + description: "Length of a JSON array; NULL on malformed or non-array input." + impls: + - args: [{ value: string, name: "value" }] + return: any1 + + - name: "json_keys" + description: "Top-level keys of a JSON object, encoded as a JSON array string; NULL on non-object input." + impls: + - args: [{ value: string, name: "value" }] + return: any1 + + - name: "json_extract" + description: "Extract JSON value(s) at PPL path(s); single → stringified match, multi → JSON-array string." + impls: + - args: [{ value: string, name: "value" }, { value: string, name: "path" }] + variadic: { min: 1 } + return: string + + - name: "json_delete" + description: "Remove PPL-path matches from a JSON document; missing paths are no-ops." + impls: + - args: [{ value: string, name: "value" }, { value: string, name: "path" }] + variadic: { min: 1 } + return: string + + - name: "json_set" + description: "Replace values at PPL-path matches (replace-only; missing paths are no-ops)." + impls: + - args: [{ value: string, name: "value" }, { value: string, name: "path" }] + variadic: { min: 1 } + return: string + + - name: "json_append" + description: "Push values onto PPL-path-matched arrays; non-array / missing targets are no-ops." + impls: + - args: [{ value: string, name: "value" }, { value: string, name: "path" }] + variadic: { min: 1 } + return: string + + - name: "json_extend" + description: "Spread JSON-array values onto PPL-path-matched arrays; scalar values fall back to append." + impls: + - args: [{ value: string, name: "value" }, { value: string, name: "path" }] + variadic: { min: 1 } + return: string diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/JsonFunctionAdaptersTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/JsonFunctionAdaptersTests.java new file mode 100644 index 0000000000000..14e12b1a4694d --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/JsonFunctionAdaptersTests.java @@ -0,0 +1,118 @@ +/* + * 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. + */ + +package org.opensearch.be.datafusion; + +import org.apache.calcite.jdbc.JavaTypeFactoryImpl; +import org.apache.calcite.plan.RelOptCluster; +import org.apache.calcite.plan.hep.HepPlanner; +import org.apache.calcite.plan.hep.HepProgramBuilder; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.rex.RexCall; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.SqlFunction; +import org.apache.calcite.sql.SqlFunctionCategory; +import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.sql.type.OperandTypes; +import org.apache.calcite.sql.type.ReturnTypes; +import org.apache.calcite.sql.type.SqlTypeName; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.List; + +/** + * Unit tests for the JSON-function adapter inner classes in + * {@link JsonFunctionAdapters}. Each inner adapter gets its own test method + * (shape + {@code testAdaptedCallPreservesOriginalReturnType} regression + * guard). See {@link YearAdapterTests} for the regression-guard rationale. + */ +public class JsonFunctionAdaptersTests extends OpenSearchTestCase { + + private RelDataTypeFactory typeFactory; + private RexBuilder rexBuilder; + private RelOptCluster cluster; + + @Override + public void setUp() throws Exception { + super.setUp(); + typeFactory = new JavaTypeFactoryImpl(); + rexBuilder = new RexBuilder(typeFactory); + HepPlanner planner = new HepPlanner(new HepProgramBuilder().build()); + cluster = RelOptCluster.create(planner, rexBuilder); + } + + // ── JsonArrayLengthAdapter ──────────────────────────────────────────── + + public void testJsonArrayLengthRewritesToLocalOp() { + // Synthesize JSON_ARRAY_LENGTH(value) with a single VARCHAR operand. + RelDataType varcharNullable = typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.VARCHAR), true); + RelDataType integerNullable = typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.INTEGER), true); + SqlFunction pplJsonArrayLengthOp = new SqlFunction( + "JSON_ARRAY_LENGTH", + SqlKind.OTHER_FUNCTION, + ReturnTypes.explicit(integerNullable), + null, + OperandTypes.STRING, + SqlFunctionCategory.STRING + ); + RexNode valueRef = rexBuilder.makeInputRef(varcharNullable, 0); + RexCall original = (RexCall) rexBuilder.makeCall(pplJsonArrayLengthOp, List.of(valueRef)); + + RexNode adapted = new JsonFunctionAdapters.JsonArrayLengthAdapter().adapt(original, List.of(), cluster); + + assertTrue("adapted node must be a RexCall, got " + adapted.getClass(), adapted instanceof RexCall); + RexCall call = (RexCall) adapted; + assertSame( + "adapted call must target LOCAL_JSON_ARRAY_LENGTH_OP", + JsonFunctionAdapters.JsonArrayLengthAdapter.LOCAL_JSON_ARRAY_LENGTH_OP, + call.getOperator() + ); + assertEquals("json_array_length is unary — no prepend / append", 1, call.getOperands().size()); + assertSame("arg 0 must be the original value operand", valueRef, call.getOperands().get(0)); + } + + /** + * The adapter MUST preserve the Calcite {@link RelDataType} of the original call. + * PPL declares {@code JSON_ARRAY_LENGTH} with INTEGER_FORCE_NULLABLE; the + * locally-declared {@code LOCAL_JSON_ARRAY_LENGTH_OP} uses + * {@code ReturnTypes.INTEGER_NULLABLE} which would infer a different + * typeFactory type instance and trip {@code Project.isValid}'s + * {@code compatibleTypes} check during fragment conversion. See + * {@link YearAdapterTests#testAdaptedCallPreservesOriginalReturnType()} for + * the original incident. + */ + public void testJsonArrayLengthPreservesOriginalReturnType() { + RelDataType varcharNullable = typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.VARCHAR), true); + // Pick a type that specifically differs from what LOCAL_JSON_ARRAY_LENGTH_OP's + // ReturnTypes.INTEGER_NULLABLE would compute — BIGINT here — so the + // regression assertion actually distinguishes "preserve" from "infer". + RelDataType bigintNullable = typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.BIGINT), true); + SqlFunction pplJsonArrayLengthOp = new SqlFunction( + "JSON_ARRAY_LENGTH", + SqlKind.OTHER_FUNCTION, + ReturnTypes.explicit(bigintNullable), + null, + OperandTypes.STRING, + SqlFunctionCategory.STRING + ); + RexNode valueRef = rexBuilder.makeInputRef(varcharNullable, 0); + RexCall original = (RexCall) rexBuilder.makeCall(pplJsonArrayLengthOp, List.of(valueRef)); + assertEquals(bigintNullable, original.getType()); + + RexNode adapted = new JsonFunctionAdapters.JsonArrayLengthAdapter().adapt(original, List.of(), cluster); + + assertEquals( + "adapted call's return type must equal the original call's return type, " + + "otherwise the enclosing Project.rowType assertion fails in fragment conversion", + original.getType(), + adapted.getType() + ); + } +}