laminar_sql/datafusion/
json_path.rs

1//! SQL/JSON path query functions.
2//!
3//! Implements a basic JSON path compiler and two scalar UDFs:
4//!
5//! - `jsonb_path_exists(jsonb, path_text) → boolean`
6//! - `jsonb_path_match(jsonb, path_text) → boolean`
7//!
8//! The path language supports a PostgreSQL SQL/JSON subset:
9//! - `$` — root element
10//! - `.key` — object member access
11//! - `[n]` — array element by index
12//! - `[*]` — wildcard array elements
13
14use std::hash::{Hash, Hasher};
15use std::sync::Arc;
16
17use arrow::datatypes::DataType;
18use arrow_array::{Array, BooleanArray, LargeBinaryArray};
19use datafusion_common::Result;
20use datafusion_expr::{
21    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility,
22};
23use parking_lot::Mutex;
24
25use super::json_types;
26
27/// One-slot cache for compiled JSON paths.
28///
29/// Since a given UDF instance almost always evaluates the same constant path
30/// string, a single-entry cache avoids recompilation on every batch.
31type PathCache = Mutex<Option<(String, CompiledJsonPath)>>;
32
33// ── JSON Path Compiler ──────────────────────────────────────────────
34
35/// A single step in a compiled JSON path.
36#[derive(Debug, Clone, PartialEq)]
37pub enum JsonPathStep {
38    /// Root element (`$`).
39    Root,
40    /// Object member access (`.key`).
41    Member(String),
42    /// Array element by index (`[n]`).
43    ArrayIndex(i64),
44    /// Wildcard array element (`[*]`).
45    ArrayWildcard,
46}
47
48/// A compiled SQL/JSON path expression.
49///
50/// Path expressions are compiled from string form (e.g., `$.items[*].price`)
51/// into a sequence of steps for evaluation against JSONB binary data.
52#[derive(Debug, Clone, PartialEq)]
53pub struct CompiledJsonPath {
54    /// Sequence of path steps.
55    pub steps: Vec<JsonPathStep>,
56}
57
58impl CompiledJsonPath {
59    /// Compiles a SQL/JSON path string.
60    ///
61    /// # Supported syntax
62    ///
63    /// - `$` — root element (must appear first)
64    /// - `.key` — object member access
65    /// - `["key"]` — quoted member access
66    /// - `[n]` — array index (non-negative)
67    /// - `[*]` — wildcard array element
68    ///
69    /// # Errors
70    ///
71    /// Returns an error if the path syntax is invalid.
72    pub fn compile(path_str: &str) -> std::result::Result<Self, String> {
73        let path_str = path_str.trim();
74        if path_str.is_empty() {
75            return Err("empty path expression".into());
76        }
77
78        let mut steps = Vec::new();
79        let chars: Vec<char> = path_str.chars().collect();
80        let mut pos = 0;
81
82        // Must start with '$'
83        if chars.first() != Some(&'$') {
84            return Err(format!(
85                "path must start with '$', got '{}'",
86                chars.first().unwrap_or(&' ')
87            ));
88        }
89        steps.push(JsonPathStep::Root);
90        pos += 1;
91
92        while pos < chars.len() {
93            match chars[pos] {
94                '.' => {
95                    pos += 1;
96                    // Read member name
97                    let start = pos;
98                    while pos < chars.len()
99                        && chars[pos] != '.'
100                        && chars[pos] != '['
101                        && !chars[pos].is_whitespace()
102                    {
103                        pos += 1;
104                    }
105                    if pos == start {
106                        return Err(format!("empty member name at position {start}"));
107                    }
108                    let name: String = chars[start..pos].iter().collect();
109                    steps.push(JsonPathStep::Member(name));
110                }
111                '[' => {
112                    pos += 1;
113                    // Skip whitespace
114                    while pos < chars.len() && chars[pos].is_whitespace() {
115                        pos += 1;
116                    }
117                    if pos >= chars.len() {
118                        return Err("unclosed bracket".into());
119                    }
120                    if chars[pos] == '*' {
121                        steps.push(JsonPathStep::ArrayWildcard);
122                        pos += 1;
123                    } else if chars[pos] == '"' || chars[pos] == '\'' {
124                        // Quoted member access: ["key"] or ['key']
125                        let quote = chars[pos];
126                        pos += 1;
127                        let start = pos;
128                        while pos < chars.len() && chars[pos] != quote {
129                            pos += 1;
130                        }
131                        if pos >= chars.len() {
132                            return Err("unclosed quoted member".into());
133                        }
134                        let name: String = chars[start..pos].iter().collect();
135                        steps.push(JsonPathStep::Member(name));
136                        pos += 1; // skip closing quote
137                    } else {
138                        // Array index
139                        let start = pos;
140                        let mut negative = false;
141                        if pos < chars.len() && chars[pos] == '-' {
142                            negative = true;
143                            pos += 1;
144                        }
145                        while pos < chars.len() && chars[pos].is_ascii_digit() {
146                            pos += 1;
147                        }
148                        if pos == start || (negative && pos == start + 1) {
149                            return Err(format!("expected array index or '*' at position {start}"));
150                        }
151                        let idx_str: String = chars[start..pos].iter().collect();
152                        let idx: i64 = idx_str
153                            .parse()
154                            .map_err(|_| format!("invalid array index: '{idx_str}'"))?;
155                        steps.push(JsonPathStep::ArrayIndex(idx));
156                    }
157                    // Skip whitespace and expect ']'
158                    while pos < chars.len() && chars[pos].is_whitespace() {
159                        pos += 1;
160                    }
161                    if pos >= chars.len() || chars[pos] != ']' {
162                        return Err(format!("expected ']' at position {pos}"));
163                    }
164                    pos += 1;
165                }
166                c if c.is_whitespace() => {
167                    pos += 1;
168                }
169                c => {
170                    return Err(format!("unexpected character '{c}' at position {pos}"));
171                }
172            }
173        }
174
175        Ok(Self { steps })
176    }
177
178    /// Evaluates the compiled path against JSONB binary data.
179    ///
180    /// Returns all matched sub-values as byte slices into the original data.
181    #[must_use]
182    pub fn evaluate<'a>(&self, jsonb: &'a [u8]) -> Vec<&'a [u8]> {
183        if jsonb.is_empty() {
184            return Vec::new();
185        }
186        let mut current = vec![jsonb];
187
188        for step in &self.steps {
189            match step {
190                JsonPathStep::Root => {
191                    // Already at root
192                }
193                JsonPathStep::Member(name) => {
194                    let mut next = Vec::new();
195                    for data in &current {
196                        if let Some(val) = json_types::jsonb_get_field(data, name) {
197                            next.push(val);
198                        }
199                    }
200                    current = next;
201                }
202                JsonPathStep::ArrayIndex(idx) => {
203                    let mut next = Vec::new();
204                    for data in &current {
205                        if *idx >= 0 {
206                            if let Ok(i) = usize::try_from(*idx) {
207                                if let Some(val) = json_types::jsonb_array_get(data, i) {
208                                    next.push(val);
209                                }
210                            }
211                        }
212                    }
213                    current = next;
214                }
215                JsonPathStep::ArrayWildcard => {
216                    let mut next = Vec::new();
217                    for data in &current {
218                        if !data.is_empty() && data[0] == 0x06 {
219                            // ARRAY tag
220                            if data.len() >= 5 {
221                                let count = u32::from_le_bytes([data[1], data[2], data[3], data[4]])
222                                    as usize;
223                                for i in 0..count {
224                                    if let Some(val) = json_types::jsonb_array_get(data, i) {
225                                        next.push(val);
226                                    }
227                                }
228                            }
229                        }
230                    }
231                    current = next;
232                }
233            }
234        }
235
236        current
237    }
238
239    /// Returns `true` if evaluating this path against the JSONB data
240    /// produces at least one match.
241    #[must_use]
242    pub fn exists(&self, jsonb: &[u8]) -> bool {
243        !self.evaluate(jsonb).is_empty()
244    }
245}
246
247// ── jsonb_path_exists UDF ───────────────────────────────────────────
248
249/// `jsonb_path_exists(jsonb, path_text) → boolean`
250///
251/// Looks up or compiles a JSON path, caching the result.
252fn compile_cached(
253    cache: &PathCache,
254    path_str: &str,
255) -> std::result::Result<CompiledJsonPath, String> {
256    let mut guard = cache.lock();
257    if let Some((cached_str, cached_path)) = guard.as_ref() {
258        if cached_str == path_str {
259            return Ok(cached_path.clone());
260        }
261    }
262    let compiled = CompiledJsonPath::compile(path_str)?;
263    *guard = Some((path_str.to_owned(), compiled.clone()));
264    Ok(compiled)
265}
266
267/// Returns `true` if the JSON path matches any element in the JSONB value.
268#[derive(Debug)]
269pub struct JsonbPathExistsUdf {
270    signature: Signature,
271    path_cache: PathCache,
272}
273
274impl Default for JsonbPathExistsUdf {
275    fn default() -> Self {
276        Self::new()
277    }
278}
279
280impl JsonbPathExistsUdf {
281    /// Creates a new `jsonb_path_exists` UDF.
282    #[must_use]
283    pub fn new() -> Self {
284        Self {
285            signature: Signature::new(
286                TypeSignature::Exact(vec![DataType::LargeBinary, DataType::Utf8]),
287                Volatility::Immutable,
288            ),
289            path_cache: Mutex::new(None),
290        }
291    }
292}
293
294impl PartialEq for JsonbPathExistsUdf {
295    fn eq(&self, _other: &Self) -> bool {
296        true
297    }
298}
299
300impl Eq for JsonbPathExistsUdf {}
301
302impl Hash for JsonbPathExistsUdf {
303    fn hash<H: Hasher>(&self, state: &mut H) {
304        "jsonb_path_exists".hash(state);
305    }
306}
307
308impl ScalarUDFImpl for JsonbPathExistsUdf {
309    fn as_any(&self) -> &dyn std::any::Any {
310        self
311    }
312
313    fn name(&self) -> &'static str {
314        "jsonb_path_exists"
315    }
316
317    fn signature(&self) -> &Signature {
318        &self.signature
319    }
320
321    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
322        Ok(DataType::Boolean)
323    }
324
325    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
326        let first = &args.args[0];
327        let second = &args.args[1];
328
329        match (first, second) {
330            (
331                ColumnarValue::Array(bin_arr),
332                ColumnarValue::Scalar(datafusion_common::ScalarValue::Utf8(Some(path_str))),
333            ) => {
334                let compiled = compile_cached(&self.path_cache, path_str)
335                    .map_err(datafusion_common::DataFusionError::Execution)?;
336                let binary = bin_arr
337                    .as_any()
338                    .downcast_ref::<LargeBinaryArray>()
339                    .ok_or_else(|| {
340                        datafusion_common::DataFusionError::Execution(
341                            "expected LargeBinary array".into(),
342                        )
343                    })?;
344                let results: BooleanArray = (0..binary.len())
345                    .map(|i| {
346                        if binary.is_null(i) {
347                            None
348                        } else {
349                            Some(compiled.exists(binary.value(i)))
350                        }
351                    })
352                    .collect();
353                Ok(ColumnarValue::Array(Arc::new(results)))
354            }
355            (
356                ColumnarValue::Scalar(datafusion_common::ScalarValue::LargeBinary(Some(bytes))),
357                ColumnarValue::Scalar(datafusion_common::ScalarValue::Utf8(Some(path_str))),
358            ) => {
359                let compiled = compile_cached(&self.path_cache, path_str)
360                    .map_err(datafusion_common::DataFusionError::Execution)?;
361                let result = compiled.exists(bytes);
362                Ok(ColumnarValue::Scalar(
363                    datafusion_common::ScalarValue::Boolean(Some(result)),
364                ))
365            }
366            _ => Ok(ColumnarValue::Scalar(
367                datafusion_common::ScalarValue::Boolean(None),
368            )),
369        }
370    }
371}
372
373// ── jsonb_path_match UDF ────────────────────────────────────────────
374
375/// `jsonb_path_match(jsonb, path_text) → boolean`
376///
377/// Returns the boolean result of evaluating a JSON path. The path must
378/// yield exactly one boolean value. Returns NULL if the path matches
379/// no values or the matched value is not boolean.
380#[derive(Debug)]
381pub struct JsonbPathMatchUdf {
382    signature: Signature,
383    path_cache: PathCache,
384}
385
386impl Default for JsonbPathMatchUdf {
387    fn default() -> Self {
388        Self::new()
389    }
390}
391
392impl JsonbPathMatchUdf {
393    /// Creates a new `jsonb_path_match` UDF.
394    #[must_use]
395    pub fn new() -> Self {
396        Self {
397            signature: Signature::new(
398                TypeSignature::Exact(vec![DataType::LargeBinary, DataType::Utf8]),
399                Volatility::Immutable,
400            ),
401            path_cache: Mutex::new(None),
402        }
403    }
404}
405
406impl PartialEq for JsonbPathMatchUdf {
407    fn eq(&self, _other: &Self) -> bool {
408        true
409    }
410}
411
412impl Eq for JsonbPathMatchUdf {}
413
414impl Hash for JsonbPathMatchUdf {
415    fn hash<H: Hasher>(&self, state: &mut H) {
416        "jsonb_path_match".hash(state);
417    }
418}
419
420/// Checks if a JSONB binary is a boolean true value.
421fn jsonb_is_true(data: &[u8]) -> Option<bool> {
422    if data.is_empty() {
423        return None;
424    }
425    match data[0] {
426        0x01 => Some(false), // BOOL_FALSE
427        0x02 => Some(true),  // BOOL_TRUE
428        _ => None,           // not a boolean
429    }
430}
431
432impl ScalarUDFImpl for JsonbPathMatchUdf {
433    fn as_any(&self) -> &dyn std::any::Any {
434        self
435    }
436
437    fn name(&self) -> &'static str {
438        "jsonb_path_match"
439    }
440
441    fn signature(&self) -> &Signature {
442        &self.signature
443    }
444
445    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
446        Ok(DataType::Boolean)
447    }
448
449    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
450        let first = &args.args[0];
451        let second = &args.args[1];
452
453        match (first, second) {
454            (
455                ColumnarValue::Array(bin_arr),
456                ColumnarValue::Scalar(datafusion_common::ScalarValue::Utf8(Some(path_str))),
457            ) => {
458                let compiled = compile_cached(&self.path_cache, path_str)
459                    .map_err(datafusion_common::DataFusionError::Execution)?;
460                let binary = bin_arr
461                    .as_any()
462                    .downcast_ref::<LargeBinaryArray>()
463                    .ok_or_else(|| {
464                        datafusion_common::DataFusionError::Execution(
465                            "expected LargeBinary array".into(),
466                        )
467                    })?;
468                let results: BooleanArray = (0..binary.len())
469                    .map(|i| {
470                        if binary.is_null(i) {
471                            None
472                        } else {
473                            let matched = compiled.evaluate(binary.value(i));
474                            if matched.len() == 1 {
475                                jsonb_is_true(matched[0])
476                            } else {
477                                None
478                            }
479                        }
480                    })
481                    .collect();
482                Ok(ColumnarValue::Array(Arc::new(results)))
483            }
484            (
485                ColumnarValue::Scalar(datafusion_common::ScalarValue::LargeBinary(Some(bytes))),
486                ColumnarValue::Scalar(datafusion_common::ScalarValue::Utf8(Some(path_str))),
487            ) => {
488                let compiled = compile_cached(&self.path_cache, path_str)
489                    .map_err(datafusion_common::DataFusionError::Execution)?;
490                let matched = compiled.evaluate(bytes);
491                let result = if matched.len() == 1 {
492                    jsonb_is_true(matched[0])
493                } else {
494                    None
495                };
496                Ok(ColumnarValue::Scalar(
497                    datafusion_common::ScalarValue::Boolean(result),
498                ))
499            }
500            _ => Ok(ColumnarValue::Scalar(
501                datafusion_common::ScalarValue::Boolean(None),
502            )),
503        }
504    }
505}
506
507#[cfg(test)]
508mod tests {
509    use super::*;
510
511    // ── Path compiler tests ──
512
513    #[test]
514    fn test_compile_root_only() {
515        let path = CompiledJsonPath::compile("$").unwrap();
516        assert_eq!(path.steps, vec![JsonPathStep::Root]);
517    }
518
519    #[test]
520    fn test_compile_member_access() {
521        let path = CompiledJsonPath::compile("$.name").unwrap();
522        assert_eq!(
523            path.steps,
524            vec![JsonPathStep::Root, JsonPathStep::Member("name".into()),]
525        );
526    }
527
528    #[test]
529    fn test_compile_nested_members() {
530        let path = CompiledJsonPath::compile("$.user.address.city").unwrap();
531        assert_eq!(
532            path.steps,
533            vec![
534                JsonPathStep::Root,
535                JsonPathStep::Member("user".into()),
536                JsonPathStep::Member("address".into()),
537                JsonPathStep::Member("city".into()),
538            ]
539        );
540    }
541
542    #[test]
543    fn test_compile_array_index() {
544        let path = CompiledJsonPath::compile("$.items[0]").unwrap();
545        assert_eq!(
546            path.steps,
547            vec![
548                JsonPathStep::Root,
549                JsonPathStep::Member("items".into()),
550                JsonPathStep::ArrayIndex(0),
551            ]
552        );
553    }
554
555    #[test]
556    fn test_compile_wildcard() {
557        let path = CompiledJsonPath::compile("$.items[*].price").unwrap();
558        assert_eq!(
559            path.steps,
560            vec![
561                JsonPathStep::Root,
562                JsonPathStep::Member("items".into()),
563                JsonPathStep::ArrayWildcard,
564                JsonPathStep::Member("price".into()),
565            ]
566        );
567    }
568
569    #[test]
570    fn test_compile_quoted_member() {
571        let path = CompiledJsonPath::compile("$[\"spaced key\"]").unwrap();
572        assert_eq!(
573            path.steps,
574            vec![
575                JsonPathStep::Root,
576                JsonPathStep::Member("spaced key".into()),
577            ]
578        );
579    }
580
581    #[test]
582    fn test_compile_empty_path_error() {
583        assert!(CompiledJsonPath::compile("").is_err());
584    }
585
586    #[test]
587    fn test_compile_no_root_error() {
588        assert!(CompiledJsonPath::compile("name").is_err());
589    }
590
591    // ── Path evaluation tests ──
592
593    #[test]
594    fn test_evaluate_root() {
595        let json: serde_json::Value = serde_json::from_str(r#"{"a": 1}"#).unwrap();
596        let jsonb = json_types::encode_jsonb(&json);
597        let path = CompiledJsonPath::compile("$").unwrap();
598        let results = path.evaluate(&jsonb);
599        assert_eq!(results.len(), 1);
600    }
601
602    #[test]
603    fn test_evaluate_member() {
604        let json: serde_json::Value = serde_json::from_str(r#"{"name": "Alice"}"#).unwrap();
605        let jsonb = json_types::encode_jsonb(&json);
606        let path = CompiledJsonPath::compile("$.name").unwrap();
607        let results = path.evaluate(&jsonb);
608        assert_eq!(results.len(), 1);
609        // Should be the string "Alice" in JSONB format
610        assert_eq!(
611            json_types::jsonb_to_text(results[0]),
612            Some("Alice".to_string())
613        );
614    }
615
616    #[test]
617    fn test_evaluate_missing_member() {
618        let json: serde_json::Value = serde_json::from_str(r#"{"name": "Alice"}"#).unwrap();
619        let jsonb = json_types::encode_jsonb(&json);
620        let path = CompiledJsonPath::compile("$.age").unwrap();
621        let results = path.evaluate(&jsonb);
622        assert!(results.is_empty());
623    }
624
625    #[test]
626    fn test_evaluate_array_index() {
627        let json: serde_json::Value = serde_json::from_str(r#"{"items": [10, 20, 30]}"#).unwrap();
628        let jsonb = json_types::encode_jsonb(&json);
629        let path = CompiledJsonPath::compile("$.items[1]").unwrap();
630        let results = path.evaluate(&jsonb);
631        assert_eq!(results.len(), 1);
632        assert_eq!(
633            json_types::jsonb_to_text(results[0]),
634            Some("20".to_string())
635        );
636    }
637
638    #[test]
639    fn test_evaluate_wildcard() {
640        let json: serde_json::Value =
641            serde_json::from_str(r#"{"items": [{"price": 10}, {"price": 20}]}"#).unwrap();
642        let jsonb = json_types::encode_jsonb(&json);
643        let path = CompiledJsonPath::compile("$.items[*].price").unwrap();
644        let results = path.evaluate(&jsonb);
645        assert_eq!(results.len(), 2);
646    }
647
648    #[test]
649    fn test_exists_true() {
650        let json: serde_json::Value = serde_json::from_str(r#"{"users": [{"age": 30}]}"#).unwrap();
651        let jsonb = json_types::encode_jsonb(&json);
652        let path = CompiledJsonPath::compile("$.users[0].age").unwrap();
653        assert!(path.exists(&jsonb));
654    }
655
656    #[test]
657    fn test_exists_false() {
658        let json: serde_json::Value = serde_json::from_str(r#"{"users": [{"age": 30}]}"#).unwrap();
659        let jsonb = json_types::encode_jsonb(&json);
660        let path = CompiledJsonPath::compile("$.users[0].email").unwrap();
661        assert!(!path.exists(&jsonb));
662    }
663
664    // ── UDF tests ──
665
666    #[test]
667    fn test_path_exists_udf_scalar() {
668        let udf = JsonbPathExistsUdf::new();
669        assert_eq!(udf.name(), "jsonb_path_exists");
670        assert_eq!(udf.return_type(&[]).unwrap(), DataType::Boolean);
671    }
672
673    #[test]
674    fn test_path_match_udf_scalar() {
675        let udf = JsonbPathMatchUdf::new();
676        assert_eq!(udf.name(), "jsonb_path_match");
677        assert_eq!(udf.return_type(&[]).unwrap(), DataType::Boolean);
678    }
679}
laminar_sql/datafusion/json_path.rs

laminar_sql/datafusion/
json_path.rs
