micromegas_datafusion_extensions/jsonb/
path_query.rs

1use crate::binary_column_accessor::create_binary_accessor;
2use datafusion::arrow::array::{Array, BinaryDictionaryBuilder, StringArray};
3use datafusion::arrow::datatypes::{DataType, Int32Type};
4use datafusion::common::{Result, internal_err};
5use datafusion::error::DataFusionError;
6use datafusion::logical_expr::{
7    ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility,
8};
9use jsonb::RawJsonb;
10use jsonb::jsonpath::parse_json_path;
11use std::collections::HashMap;
12use std::sync::Arc;
13
14#[derive(Clone, Copy)]
15enum PathQueryMode {
16    First,
17    All,
18}
19
20fn eval_jsonb_path_query(
21    func_name: &str,
22    args: ScalarFunctionArgs,
23    mode: PathQueryMode,
24) -> Result<ColumnarValue> {
25    let args = ColumnarValue::values_to_arrays(&args.args)?;
26    if args.len() != 2 {
27        return internal_err!("wrong number of arguments to {func_name}");
28    }
29
30    let accessor = create_binary_accessor(&args[0]).map_err(|e| {
31        DataFusionError::Execution(format!(
32            "Invalid input type for {func_name}: {e}. Expected Binary or Dictionary<Int32, Binary>"
33        ))
34    })?;
35
36    let paths = args[1]
37        .as_any()
38        .downcast_ref::<StringArray>()
39        .ok_or_else(|| {
40            DataFusionError::Execution(format!("second argument to {func_name} must be a string"))
41        })?;
42
43    let mut builder = BinaryDictionaryBuilder::<Int32Type>::new();
44    let mut path_cache: HashMap<&str, _> = HashMap::new();
45
46    for i in 0..accessor.len() {
47        if accessor.is_null(i) || paths.is_null(i) {
48            builder.append_null();
49        } else {
50            let path_str = paths.value(i);
51            if !path_cache.contains_key(path_str) {
52                let parsed = parse_json_path(path_str.as_bytes()).map_err(|e| {
53                    DataFusionError::Execution(format!(
54                        "{func_name}: invalid JSONPath '{path_str}': {e}"
55                    ))
56                })?;
57                path_cache.insert(path_str, parsed);
58            }
59            let json_path = path_cache.get(path_str).expect("just inserted");
60            let raw = RawJsonb::new(accessor.value(i));
61            match mode {
62                PathQueryMode::First => match raw.select_first_by_path(json_path) {
63                    Ok(Some(value)) => builder.append_value(value.as_ref()),
64                    Ok(None) => builder.append_null(),
65                    Err(e) => return Err(DataFusionError::External(e.into())),
66                },
67                PathQueryMode::All => match raw.select_array_by_path(json_path) {
68                    Ok(value) => builder.append_value(value.as_ref()),
69                    Err(e) => return Err(DataFusionError::External(e.into())),
70                },
71            }
72        }
73    }
74
75    Ok(ColumnarValue::Array(Arc::new(builder.finish())))
76}
77
78/// A scalar UDF that returns the first match of a JSONPath expression on a JSONB value.
79///
80/// Accepts both Binary and Dictionary<Int32, Binary> inputs for the JSONB argument.
81/// The path argument is Utf8. Returns Dictionary<Int32, Binary> for memory efficiency,
82/// or NULL if no match is found.
83#[derive(Debug, PartialEq, Eq, Hash)]
84pub struct JsonbPathQueryFirst {
85    signature: Signature,
86}
87
88impl JsonbPathQueryFirst {
89    pub fn new() -> Self {
90        Self {
91            signature: Signature::any(2, Volatility::Immutable),
92        }
93    }
94}
95
96impl Default for JsonbPathQueryFirst {
97    fn default() -> Self {
98        Self::new()
99    }
100}
101
102impl ScalarUDFImpl for JsonbPathQueryFirst {
103    fn name(&self) -> &str {
104        "jsonb_path_query_first"
105    }
106
107    fn signature(&self) -> &Signature {
108        &self.signature
109    }
110
111    fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
112        Ok(DataType::Dictionary(
113            Box::new(DataType::Int32),
114            Box::new(DataType::Binary),
115        ))
116    }
117
118    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
119        eval_jsonb_path_query("jsonb_path_query_first", args, PathQueryMode::First)
120    }
121}
122
123/// Creates a user-defined function to extract the first JSONPath match from a JSONB value.
124pub fn make_jsonb_path_query_first_udf() -> ScalarUDF {
125    ScalarUDF::new_from_impl(JsonbPathQueryFirst::new())
126}
127
128/// A scalar UDF that returns all matches of a JSONPath expression on a JSONB value as a JSONB array.
129///
130/// Accepts both Binary and Dictionary<Int32, Binary> inputs for the JSONB argument.
131/// The path argument is Utf8. Returns Dictionary<Int32, Binary> containing a JSONB array
132/// of all matched values.
133#[derive(Debug, PartialEq, Eq, Hash)]
134pub struct JsonbPathQuery {
135    signature: Signature,
136}
137
138impl JsonbPathQuery {
139    pub fn new() -> Self {
140        Self {
141            signature: Signature::any(2, Volatility::Immutable),
142        }
143    }
144}
145
146impl Default for JsonbPathQuery {
147    fn default() -> Self {
148        Self::new()
149    }
150}
151
152impl ScalarUDFImpl for JsonbPathQuery {
153    fn name(&self) -> &str {
154        "jsonb_path_query"
155    }
156
157    fn signature(&self) -> &Signature {
158        &self.signature
159    }
160
161    fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
162        Ok(DataType::Dictionary(
163            Box::new(DataType::Int32),
164            Box::new(DataType::Binary),
165        ))
166    }
167
168    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
169        eval_jsonb_path_query("jsonb_path_query", args, PathQueryMode::All)
170    }
171}
172
173/// Creates a user-defined function to extract all JSONPath matches from a JSONB value as a JSONB array.
174pub fn make_jsonb_path_query_udf() -> ScalarUDF {
175    ScalarUDF::new_from_impl(JsonbPathQuery::new())
176}
micromegas_datafusion_extensions/jsonb/path_query.rs

micromegas_datafusion_extensions/jsonb/
path_query.rs
