micromegas_datafusion_extensions/jsonb/
extract.rs

1//! Shared JSONB extraction helpers, used by both the `jsonb_each` / `jsonb_array_elements`
2//! UDTFs and the list-returning scalar UDFs (`jsonb_entries`, `jsonb_elements`,
3//! `jsonb_path_elements`).
4//!
5//! These functions are pure — no DataFusion array/UDF plumbing — and return `Ok(None)` for a
6//! "wrong shape" input (e.g. an object passed where an array is expected) rather than an error,
7//! so callers can decide whether that should surface as NULL (the scalar UDFs) or as an error
8//! (the UDTFs, which map `Ok(None)` back to their existing error text).
9
10use datafusion::common::Result;
11use datafusion::error::DataFusionError;
12use jsonb::RawJsonb;
13use jsonb::jsonpath::JsonPath;
14
15/// `(key, value)` pairs, one per JSONB object field or array element.
16pub type EntryList = Vec<(String, Vec<u8>)>;
17
18/// Object → `(field name, value)`; array → `(index as string, value)`. `Ok(None)` if the input
19/// is a JSON scalar (neither an object nor an array).
20pub fn object_or_array_entries(jsonb_bytes: &[u8]) -> Result<Option<EntryList>> {
21    let jsonb = RawJsonb::new(jsonb_bytes);
22    match jsonb.object_each() {
23        Ok(Some(entries)) => {
24            return Ok(Some(
25                entries
26                    .into_iter()
27                    .map(|(k, v)| (k, v.as_ref().to_vec()))
28                    .collect(),
29            ));
30        }
31        Ok(None) => {}
32        Err(e) => return Err(DataFusionError::External(e.into())),
33    }
34    match jsonb.array_values() {
35        Ok(Some(values)) => Ok(Some(
36            values
37                .into_iter()
38                .enumerate()
39                .map(|(i, v)| (i.to_string(), v.as_ref().to_vec()))
40                .collect(),
41        )),
42        Ok(None) => Ok(None),
43        Err(e) => Err(DataFusionError::External(e.into())),
44    }
45}
46
47/// Array → elements. `Ok(None)` if the input is not an array.
48pub fn array_elements(jsonb_bytes: &[u8]) -> Result<Option<Vec<Vec<u8>>>> {
49    let jsonb = RawJsonb::new(jsonb_bytes);
50    match jsonb.array_values() {
51        Ok(Some(values)) => Ok(Some(
52            values.into_iter().map(|v| v.as_ref().to_vec()).collect(),
53        )),
54        Ok(None) => Ok(None),
55        Err(e) => Err(DataFusionError::External(e.into())),
56    }
57}
58
59/// All matches of a parsed JSONPath, as separate values — not wrapped in one JSONB array (that's
60/// `select_array_by_path`, used by `jsonb_path_query`). An empty result means zero matches, not
61/// an error.
62pub fn path_select_all<'a>(jsonb_bytes: &[u8], path: &'a JsonPath<'a>) -> Result<Vec<Vec<u8>>> {
63    let jsonb = RawJsonb::new(jsonb_bytes);
64    match jsonb.select_by_path(path) {
65        Ok(values) => Ok(values.into_iter().map(|v| v.as_ref().to_vec()).collect()),
66        Err(e) => Err(DataFusionError::External(e.into())),
67    }
68}
micromegas_datafusion_extensions/jsonb/extract.rs

micromegas_datafusion_extensions/jsonb/
extract.rs
