blob: b61a3172fbd7e91b7c562a80a7b92a25b748a92a [file]
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use arrow::array::{Array, ArrayRef, StringArray, StringBuilder};
use datafusion::common::{
cast::as_generic_string_array, exec_err, Result as DataFusionResult, ScalarValue,
};
use datafusion::logical_expr::ColumnarValue;
use serde::de::{DeserializeSeed, IgnoredAny, MapAccess, SeqAccess, Visitor};
use serde::{Deserialize, Deserializer};
use serde_json::Value;
use std::fmt;
use std::sync::Arc;
/// Extracts a string from a ScalarValue, returning Ok(None) for null values.
fn scalar_to_str(scalar: &ScalarValue, arg_name: &str) -> DataFusionResult<Option<String>> {
match scalar {
ScalarValue::Utf8(s) | ScalarValue::LargeUtf8(s) => Ok(s.clone()),
_ => exec_err!("get_json_object {arg_name} must be a string"),
}
}
/// Spark-compatible `get_json_object` function.
///
/// Extracts a JSON value from a JSON string using a JSONPath expression.
/// Returns the result as a string, or null if the path doesn't match or input is invalid.
///
/// Supported JSONPath syntax:
/// - `$` — root element
/// - `.name` or `['name']` — named child
/// - `[n]` — array index (0-based)
/// - `[*]` — array wildcard (iterates over array elements)
pub fn spark_get_json_object(args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {
if args.len() != 2 {
return exec_err!(
"get_json_object expects 2 arguments (json, path), got {}",
args.len()
);
}
match (&args[0], &args[1]) {
// Column json, scalar path (most common case)
(ColumnarValue::Array(json_array), ColumnarValue::Scalar(path_scalar)) => {
let path_str = match scalar_to_str(path_scalar, "path")? {
Some(p) => p,
None => {
let null_array: ArrayRef = Arc::new(StringArray::new_null(json_array.len()));
return Ok(ColumnarValue::Array(null_array));
}
};
let parsed_path = match parse_json_path(&path_str) {
Some(p) => p,
None => {
let null_array: ArrayRef = Arc::new(StringArray::new_null(json_array.len()));
return Ok(ColumnarValue::Array(null_array));
}
};
let json_strings = as_generic_string_array::<i32>(json_array)?;
let mut builder = StringBuilder::with_capacity(json_strings.len(), 0);
for i in 0..json_strings.len() {
if json_strings.is_null(i) {
builder.append_null();
} else {
let json_str = json_strings.value(i);
match evaluate_path(json_str, &parsed_path) {
Some(result) => builder.append_value(&result),
None => builder.append_null(),
}
}
}
Ok(ColumnarValue::Array(Arc::new(builder.finish())))
}
// Scalar json, scalar path
(ColumnarValue::Scalar(json_scalar), ColumnarValue::Scalar(path_scalar)) => {
let json_str = match scalar_to_str(json_scalar, "json")? {
Some(s) => s,
None => return Ok(ColumnarValue::Scalar(ScalarValue::Utf8(None))),
};
let path_str = match scalar_to_str(path_scalar, "path")? {
Some(p) => p,
None => return Ok(ColumnarValue::Scalar(ScalarValue::Utf8(None))),
};
let parsed_path = match parse_json_path(&path_str) {
Some(p) => p,
None => return Ok(ColumnarValue::Scalar(ScalarValue::Utf8(None))),
};
let result = evaluate_path(&json_str, &parsed_path);
Ok(ColumnarValue::Scalar(ScalarValue::Utf8(result)))
}
// Column json, column path
(ColumnarValue::Array(json_array), ColumnarValue::Array(path_array)) => {
let json_strings = as_generic_string_array::<i32>(json_array)?;
let path_strings = as_generic_string_array::<i32>(path_array)?;
let mut builder = StringBuilder::with_capacity(json_strings.len(), 0);
for i in 0..json_strings.len() {
if json_strings.is_null(i) || path_strings.is_null(i) {
builder.append_null();
} else {
let json_str = json_strings.value(i);
let path_str = path_strings.value(i);
match parse_json_path(path_str) {
Some(parsed_path) => match evaluate_path(json_str, &parsed_path) {
Some(result) => builder.append_value(&result),
None => builder.append_null(),
},
None => builder.append_null(),
}
}
}
Ok(ColumnarValue::Array(Arc::new(builder.finish())))
}
_ => exec_err!("get_json_object: unsupported argument types"),
}
}
/// A parsed JSONPath segment.
#[derive(Debug, Clone)]
enum PathSegment {
/// Named field: `.name` or `['name']`
Field(String),
/// Array index: `[n]`
Index(usize),
/// Wildcard: `[*]` (iterates over array elements)
Wildcard,
}
/// A parsed JSONPath expression with precomputed metadata.
struct ParsedPath {
segments: Vec<PathSegment>,
has_wildcard: bool,
}
/// Parse a Spark-compatible JSONPath expression.
/// Returns None for invalid paths.
fn parse_json_path(path: &str) -> Option<ParsedPath> {
let mut chars = path.chars().peekable();
// Must start with '$'
if chars.next()? != '$' {
return None;
}
let mut segments = Vec::new();
let mut has_wildcard = false;
while chars.peek().is_some() {
match chars.peek()? {
'.' => {
chars.next();
if chars.peek() == Some(&'.') {
// Recursive descent not supported
return None;
}
if chars.peek() == Some(&'*') {
chars.next();
segments.push(PathSegment::Wildcard);
has_wildcard = true;
} else {
// Read field name
let mut name = String::new();
while let Some(&c) = chars.peek() {
if c == '.' || c == '[' {
break;
}
name.push(c);
chars.next();
}
if name.is_empty() {
return None;
}
segments.push(PathSegment::Field(name));
}
}
'[' => {
chars.next();
if chars.peek() == Some(&'\'') {
// Bracket notation with quotes: ['name'] or ['*']
chars.next();
let mut name = String::new();
loop {
match chars.next()? {
'\'' => break,
c => name.push(c),
}
}
if chars.next()? != ']' {
return None;
}
if name == "*" {
segments.push(PathSegment::Wildcard);
has_wildcard = true;
} else {
segments.push(PathSegment::Field(name));
}
} else if chars.peek() == Some(&'*') {
// [*]
chars.next();
if chars.next()? != ']' {
return None;
}
segments.push(PathSegment::Wildcard);
has_wildcard = true;
} else {
// [n] — numeric index
let mut num_str = String::new();
while let Some(&c) = chars.peek() {
if c == ']' {
break;
}
num_str.push(c);
chars.next();
}
if chars.next()? != ']' {
return None;
}
let idx: usize = num_str.parse().ok()?;
segments.push(PathSegment::Index(idx));
}
}
_ => {
// Unexpected character
return None;
}
}
}
Some(ParsedPath {
segments,
has_wildcard,
})
}
/// Evaluate a parsed JSONPath against a JSON string.
/// Returns the result as a string, or None if no match.
fn evaluate_path(json_str: &str, path: &ParsedPath) -> Option<String> {
if !path.has_wildcard {
return value_into_string(extract_no_wildcard(json_str, &path.segments)?);
}
let value: Value = serde_json::from_str(json_str).ok()?;
// Wildcard path: may return multiple results
let results = evaluate_with_wildcard(&value, &path.segments);
match results.len() {
0 => None,
1 => {
// Single wildcard match: Spark preserves JSON serialization format
// (strings keep their quotes, numbers don't)
if results[0].is_null() {
None
} else {
serde_json::to_string(results[0]).ok()
}
}
// Multiple results: wrap in JSON array. A slice of `&Value` serializes
// as a JSON array, so no clone into an owned `Value::Array` is needed.
_ => serde_json::to_string(&results).ok(),
}
}
/// Evaluation for paths without wildcards.
///
/// Descends into the document while it is being parsed, so only the matched
/// subtree is materialized as a `Value`; everything else is skipped by the
/// parser without allocating. The whole document is still consumed, so
/// malformed JSON anywhere in the input yields no match, as a full parse would.
fn extract_no_wildcard(json_str: &str, segments: &[PathSegment]) -> Option<Value> {
let mut de = serde_json::Deserializer::from_str(json_str);
let found = PathSeed { segments }.deserialize(&mut de).ok()?;
de.end().ok()?;
found
}
/// Deserializes the value at `segments`, discarding everything else.
struct PathSeed<'a> {
segments: &'a [PathSegment],
}
impl<'de> DeserializeSeed<'de> for PathSeed<'_> {
type Value = Option<Value>;
fn deserialize<D: Deserializer<'de>>(self, deserializer: D) -> Result<Self::Value, D::Error> {
if self.segments.is_empty() {
return Value::deserialize(deserializer).map(Some);
}
deserializer.deserialize_any(SegmentVisitor {
segments: self.segments,
})
}
}
/// Applies `segments[0]` to the value being visited. A value whose shape does
/// not match the segment (an index into an object, say) is skipped and reported
/// as no match rather than as an error, matching a lookup on a parsed document.
struct SegmentVisitor<'a> {
segments: &'a [PathSegment],
}
impl<'de> Visitor<'de> for SegmentVisitor<'_> {
type Value = Option<Value>;
fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("a JSON value")
}
fn visit_bool<E>(self, _: bool) -> Result<Self::Value, E> {
Ok(None)
}
fn visit_i64<E>(self, _: i64) -> Result<Self::Value, E> {
Ok(None)
}
fn visit_u64<E>(self, _: u64) -> Result<Self::Value, E> {
Ok(None)
}
fn visit_f64<E>(self, _: f64) -> Result<Self::Value, E> {
Ok(None)
}
fn visit_str<E>(self, _: &str) -> Result<Self::Value, E> {
Ok(None)
}
fn visit_unit<E>(self) -> Result<Self::Value, E> {
Ok(None)
}
fn visit_map<A: MapAccess<'de>>(self, mut map: A) -> Result<Self::Value, A::Error> {
let PathSegment::Field(name) = &self.segments[0] else {
IgnoredAny.visit_map(map)?;
return Ok(None);
};
let mut found = None;
// Every entry is visited so that a duplicated key resolves to its last
// occurrence, as it would in a parsed object.
while let Some(matched) = map.next_key_seed(KeySeed(name))? {
if matched {
found = map.next_value_seed(PathSeed {
segments: &self.segments[1..],
})?;
} else {
map.next_value::<IgnoredAny>()?;
}
}
Ok(found)
}
fn visit_seq<A: SeqAccess<'de>>(self, mut seq: A) -> Result<Self::Value, A::Error> {
let PathSegment::Index(idx) = &self.segments[0] else {
IgnoredAny.visit_seq(seq)?;
return Ok(None);
};
for _ in 0..*idx {
if seq.next_element::<IgnoredAny>()?.is_none() {
return Ok(None);
}
}
let found = seq
.next_element_seed(PathSeed {
segments: &self.segments[1..],
})?
.flatten();
// The remaining elements are still visited, so that a malformed element
// after the match yields no match, as a full parse would.
IgnoredAny.visit_seq(seq)?;
Ok(found)
}
}
/// Compares an object key against a field name without allocating it.
struct KeySeed<'a>(&'a str);
impl<'de> DeserializeSeed<'de> for KeySeed<'_> {
type Value = bool;
fn deserialize<D: Deserializer<'de>>(self, deserializer: D) -> Result<bool, D::Error> {
deserializer.deserialize_str(self)
}
}
impl<'de> Visitor<'de> for KeySeed<'_> {
type Value = bool;
fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("an object key")
}
fn visit_str<E>(self, key: &str) -> Result<bool, E> {
Ok(key == self.0)
}
}
/// Evaluation for paths containing wildcards.
/// Returns references to all matching values.
fn evaluate_with_wildcard<'a>(value: &'a Value, segments: &[PathSegment]) -> Vec<&'a Value> {
if segments.is_empty() {
return vec![value];
}
let rest = &segments[1..];
match &segments[0] {
PathSegment::Field(name) => match value {
Value::Object(map) => match map.get(name) {
Some(v) => evaluate_with_wildcard(v, rest),
None => vec![],
},
_ => vec![],
},
PathSegment::Index(idx) => match value {
Value::Array(arr) => match arr.get(*idx) {
Some(v) => evaluate_with_wildcard(v, rest),
None => vec![],
},
_ => vec![],
},
PathSegment::Wildcard => match value {
Value::Array(arr) => arr
.iter()
.flat_map(|v| evaluate_with_wildcard(v, rest))
.collect(),
_ => vec![],
},
}
}
/// Convert a JSON value to its string representation matching Spark behavior.
/// - Strings are returned without quotes
/// - null returns None
/// - Numbers, booleans, objects, arrays are serialized as JSON
fn value_into_string(value: Value) -> Option<String> {
match value {
Value::Null => None,
Value::String(s) => Some(s),
other => Some(other.to_string()),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_parse_json_path() {
// Root only
let path = parse_json_path("$").unwrap();
assert!(path.segments.is_empty());
assert!(!path.has_wildcard);
// Simple field
let path = parse_json_path("$.name").unwrap();
assert!(matches!(&path.segments[0], PathSegment::Field(n) if n == "name"));
assert!(!path.has_wildcard);
// Array index
let path = parse_json_path("$[0]").unwrap();
assert!(matches!(&path.segments[0], PathSegment::Index(0)));
// Bracket notation
let path = parse_json_path("$['key with spaces']").unwrap();
assert!(matches!(&path.segments[0], PathSegment::Field(n) if n == "key with spaces"));
// Wildcard
let path = parse_json_path("$[*]").unwrap();
assert!(matches!(&path.segments[0], PathSegment::Wildcard));
assert!(path.has_wildcard);
let path = parse_json_path("$.*").unwrap();
assert!(matches!(&path.segments[0], PathSegment::Wildcard));
assert!(path.has_wildcard);
// Recursive descent not supported
assert!(parse_json_path("$..name").is_none());
// Must start with $
assert!(parse_json_path("name").is_none());
assert!(parse_json_path("[0]").is_none());
}
#[test]
fn test_evaluate_simple_field() {
let path = parse_json_path("$.name").unwrap();
assert_eq!(
evaluate_path(r#"{"name":"John","age":30}"#, &path),
Some("John".to_string())
);
let path = parse_json_path("$.age").unwrap();
assert_eq!(
evaluate_path(r#"{"name":"John","age":30}"#, &path),
Some("30".to_string())
);
}
#[test]
fn test_evaluate_nested() {
let json = r#"{"user":{"profile":{"name":"Alice"}}}"#;
let path = parse_json_path("$.user.profile.name").unwrap();
assert_eq!(evaluate_path(json, &path), Some("Alice".to_string()));
}
#[test]
fn test_evaluate_array_index() {
let path = parse_json_path("$[0]").unwrap();
assert_eq!(evaluate_path(r#"[1,2,3]"#, &path), Some("1".to_string()));
let path = parse_json_path("$[3]").unwrap();
assert_eq!(evaluate_path(r#"[1,2,3]"#, &path), None);
}
#[test]
fn test_evaluate_root() {
let path = parse_json_path("$").unwrap();
assert_eq!(
evaluate_path(r#"{"a":"b"}"#, &path),
Some(r#"{"a":"b"}"#.to_string())
);
}
#[test]
fn test_evaluate_null_value() {
let path = parse_json_path("$.a").unwrap();
assert_eq!(evaluate_path(r#"{"a":null}"#, &path), None);
}
#[test]
fn test_null_midpath_is_null() {
// A null encountered before the path is exhausted has no child to
// recurse into, so the lookup yields no match (SegmentVisitor::visit_unit).
let path = parse_json_path("$.a.b").unwrap();
assert_eq!(evaluate_path(r#"{"a":null}"#, &path), None);
}
#[test]
fn test_match_then_trailing_garbage_is_null() {
// The match is found early, but trailing content makes the document
// malformed; the full document must still be validated (de.end()).
let path = parse_json_path("$.a").unwrap();
assert_eq!(evaluate_path(r#"{"a":1} garbage"#, &path), None);
}
#[test]
fn test_match_then_malformed_sibling_is_null() {
// "a" matches early, but the "b" sibling is malformed; visit_map must
// keep draining entries after a match so the parse still rejects this.
let path = parse_json_path("$.a").unwrap();
assert_eq!(evaluate_path(r#"{"a":1,"b":}"#, &path), None);
}
#[test]
fn test_array_match_then_malformed_element_is_null() {
// The element at index 0 matches, but a later element is malformed;
// visit_seq must keep draining elements after a match.
let path = parse_json_path("$[0]").unwrap();
assert_eq!(evaluate_path(r#"[1,,]"#, &path), None);
}
#[test]
fn test_duplicate_key_last_wins() {
// Every entry is visited, so a duplicated key resolves to its last
// occurrence, matching serde_json's preserve_order overwrite behavior.
let path = parse_json_path("$.a").unwrap();
assert_eq!(
evaluate_path(r#"{"a":1,"a":2}"#, &path),
Some("2".to_string())
);
}
#[test]
fn test_evaluate_missing_field() {
let path = parse_json_path("$.c").unwrap();
assert_eq!(evaluate_path(r#"{"a":"b"}"#, &path), None);
}
#[test]
fn test_evaluate_invalid_json() {
let path = parse_json_path("$.a").unwrap();
assert_eq!(evaluate_path("not json", &path), None);
}
#[test]
fn test_evaluate_wildcard() {
let json = r#"[{"a":"b"},{"a":"c"}]"#;
let path = parse_json_path("$[*].a").unwrap();
assert_eq!(evaluate_path(json, &path), Some(r#"["b","c"]"#.to_string()));
}
#[test]
fn test_evaluate_string_unquoted() {
// Strings should be returned without quotes for non-wildcard paths
let path = parse_json_path("$[1]").unwrap();
assert_eq!(evaluate_path(r#"["a","b"]"#, &path), Some("b".to_string()));
}
#[test]
fn test_evaluate_nested_array_field() {
let json = r#"{"items":["apple","banana","cherry"]}"#;
let path = parse_json_path("$.items[1]").unwrap();
assert_eq!(evaluate_path(json, &path), Some("banana".to_string()));
}
#[test]
fn test_evaluate_bracket_notation_with_spaces() {
let json = r#"{"key with spaces":"it works"}"#;
let path = parse_json_path("$['key with spaces']").unwrap();
assert_eq!(evaluate_path(json, &path), Some("it works".to_string()));
}
#[test]
fn test_evaluate_boolean_and_nested_object() {
let json = r#"{"a":true,"b":{"c":1}}"#;
let path_a = parse_json_path("$.a").unwrap();
assert_eq!(evaluate_path(json, &path_a), Some("true".to_string()));
let path_b = parse_json_path("$.b").unwrap();
assert_eq!(evaluate_path(json, &path_b), Some(r#"{"c":1}"#.to_string()));
}
#[test]
fn test_object_key_order_preserved() {
// Depends on serde_json "preserve_order" feature (see Cargo.toml)
let json = r#"{"z":1,"a":2}"#;
let path = parse_json_path("$").unwrap();
assert_eq!(
evaluate_path(json, &path),
Some(r#"{"z":1,"a":2}"#.to_string())
);
}
#[test]
fn test_wildcard_single_match() {
// Single wildcard match on string: Spark preserves JSON quotes
let json = r#"[{"a":"only"}]"#;
let path = parse_json_path("$[*].a").unwrap();
assert_eq!(evaluate_path(json, &path), Some(r#""only""#.to_string()));
// Single wildcard match on number: no quotes
let json = r#"[{"a":42}]"#;
assert_eq!(evaluate_path(json, &path), Some("42".to_string()));
}
#[test]
fn test_wildcard_missing_fields() {
// Wildcard should skip elements where the field is missing
let json = r#"[{"a":1},{"b":2},{"a":3}]"#;
let path = parse_json_path("$[*].a").unwrap();
assert_eq!(evaluate_path(json, &path), Some("[1,3]".to_string()));
}
#[test]
fn test_field_with_colon() {
let json = r#"{"fb:testid":"123"}"#;
let path = parse_json_path("$.fb:testid").unwrap();
assert_eq!(evaluate_path(json, &path), Some("123".to_string()));
}
#[test]
fn test_dot_bracket_invalid() {
// $.[0] is not valid path syntax in Spark
assert!(parse_json_path("$.[0]").is_none());
}
#[test]
fn test_object_wildcard() {
// Spark returns null for $.* on objects (wildcard only works in array contexts)
let json = r#"{"a":1,"b":2,"c":3}"#;
let path = parse_json_path("$.*").unwrap();
assert_eq!(evaluate_path(json, &path), None);
}
#[test]
fn test_unicode_field_names() {
let json = r#"{"名前":"太郎","年齢":25}"#;
let path = parse_json_path("$.名前").unwrap();
assert_eq!(evaluate_path(json, &path), Some("太郎".to_string()));
}
#[test]
fn test_unicode_values() {
let json = r#"{"greeting":"こんにちは世界"}"#;
let path = parse_json_path("$.greeting").unwrap();
assert_eq!(
evaluate_path(json, &path),
Some("こんにちは世界".to_string())
);
}
#[test]
fn test_unicode_emoji() {
let json = r#"{"emoji":"🎉🚀","nested":{"flag":"🇺🇸"}}"#;
let path = parse_json_path("$.emoji").unwrap();
assert_eq!(evaluate_path(json, &path), Some("🎉🚀".to_string()));
let path = parse_json_path("$.nested.flag").unwrap();
assert_eq!(evaluate_path(json, &path), Some("🇺🇸".to_string()));
}
#[test]
fn test_unicode_bracket_notation() {
let json = r#"{"":""}"#;
let path = parse_json_path("$['键']").unwrap();
assert_eq!(evaluate_path(json, &path), Some("值".to_string()));
}
#[test]
fn test_unicode_mixed_scripts() {
let json = r#"{"data":"café résumé naïve"}"#;
let path = parse_json_path("$.data").unwrap();
assert_eq!(
evaluate_path(json, &path),
Some("café résumé naïve".to_string())
);
}
#[test]
fn test_unicode_wildcard() {
let json = r#"[{"":"Alice"},{"":"太郎"}]"#;
let path = parse_json_path("$[*].名").unwrap();
assert_eq!(
evaluate_path(json, &path),
Some(r#"["Alice","太郎"]"#.to_string())
);
}
}