laminar_sql/datafusion/json_extensions/
shape_ops.rs

1//! JSON flattening, reconstruction, column extraction, and schema inference.
2
3use std::collections::HashSet;
4use std::hash::{Hash, Hasher};
5use std::sync::Arc;
6
7use arrow::datatypes::DataType;
8use arrow_array::{
9    builder::{LargeBinaryBuilder, StringBuilder},
10    Array, LargeBinaryArray, StringArray,
11};
12use datafusion_common::Result;
13use datafusion_expr::{
14    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility,
15};
16
17use super::super::json_types;
18use super::super::json_udf::expand_args;
19
20const MAX_DEPTH: usize = 64;
21
22// ══════════════════════════════════════════════════════════════════
23// jsonb_flatten(jsonb, text) -> jsonb
24// ══════════════════════════════════════════════════════════════════
25
26/// `jsonb_flatten(jsonb, separator) -> jsonb`
27///
28/// Flattens a nested JSONB object into a single-level object with
29/// dot-path keys. Arrays are indexed numerically (e.g. `tags.0`).
30#[derive(Debug)]
31pub struct JsonbFlatten {
32    signature: Signature,
33}
34
35impl JsonbFlatten {
36    /// Creates a new `jsonb_flatten` UDF.
37    #[must_use]
38    pub fn new() -> Self {
39        Self {
40            signature: Signature::new(
41                TypeSignature::Exact(vec![DataType::LargeBinary, DataType::Utf8]),
42                Volatility::Immutable,
43            ),
44        }
45    }
46}
47
48impl Default for JsonbFlatten {
49    fn default() -> Self {
50        Self::new()
51    }
52}
53
54impl PartialEq for JsonbFlatten {
55    fn eq(&self, _other: &Self) -> bool {
56        true
57    }
58}
59
60impl Eq for JsonbFlatten {}
61
62impl Hash for JsonbFlatten {
63    fn hash<H: Hasher>(&self, state: &mut H) {
64        "jsonb_flatten".hash(state);
65    }
66}
67
68fn flatten_value(
69    val: &serde_json::Value,
70    prefix: &str,
71    sep: &str,
72    out: &mut serde_json::Map<String, serde_json::Value>,
73    depth: usize,
74) -> std::result::Result<(), String> {
75    if depth > MAX_DEPTH {
76        return Err("jsonb_flatten: max depth exceeded".into());
77    }
78    match val {
79        serde_json::Value::Object(obj) => {
80            for (k, v) in obj {
81                let new_key = if prefix.is_empty() {
82                    k.clone()
83                } else {
84                    format!("{prefix}{sep}{k}")
85                };
86                flatten_value(v, &new_key, sep, out, depth + 1)?;
87            }
88        }
89        serde_json::Value::Array(arr) => {
90            for (idx, v) in arr.iter().enumerate() {
91                let new_key = if prefix.is_empty() {
92                    idx.to_string()
93                } else {
94                    format!("{prefix}{sep}{idx}")
95                };
96                flatten_value(v, &new_key, sep, out, depth + 1)?;
97            }
98        }
99        _ => {
100            out.insert(prefix.to_owned(), val.clone());
101        }
102    }
103    Ok(())
104}
105
106impl ScalarUDFImpl for JsonbFlatten {
107    fn as_any(&self) -> &dyn std::any::Any {
108        self
109    }
110
111    fn name(&self) -> &'static str {
112        "jsonb_flatten"
113    }
114
115    fn signature(&self) -> &Signature {
116        &self.signature
117    }
118
119    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
120        Ok(DataType::LargeBinary)
121    }
122
123    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
124        let expanded = expand_args(&args.args)?;
125        let jsonb_arr = expanded[0]
126            .as_any()
127            .downcast_ref::<LargeBinaryArray>()
128            .ok_or_else(|| {
129                datafusion_common::DataFusionError::Internal(
130                    "jsonb_flatten: first arg must be LargeBinary".into(),
131                )
132            })?;
133        let sep_arr = expanded[1]
134            .as_any()
135            .downcast_ref::<StringArray>()
136            .ok_or_else(|| {
137                datafusion_common::DataFusionError::Internal(
138                    "jsonb_flatten: second arg must be Utf8".into(),
139                )
140            })?;
141
142        let mut builder = LargeBinaryBuilder::with_capacity(jsonb_arr.len(), 256);
143        for i in 0..jsonb_arr.len() {
144            if jsonb_arr.is_null(i) || sep_arr.is_null(i) {
145                builder.append_null();
146            } else {
147                let sep = sep_arr.value(i);
148                match json_types::jsonb_to_value(jsonb_arr.value(i)) {
149                    Some(val) => {
150                        let mut flat = serde_json::Map::new();
151                        flatten_value(&val, "", sep, &mut flat, 0)
152                            .map_err(datafusion_common::DataFusionError::Execution)?;
153                        builder.append_value(json_types::encode_jsonb(&serde_json::Value::Object(
154                            flat,
155                        )));
156                    }
157                    None => builder.append_null(),
158                }
159            }
160        }
161        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
162    }
163}
164
165// ══════════════════════════════════════════════════════════════════
166// jsonb_unflatten(jsonb, text) -> jsonb
167// ══════════════════════════════════════════════════════════════════
168
169/// `jsonb_unflatten(jsonb, separator) -> jsonb`
170///
171/// Rebuilds a nested JSONB object from a flat key-value structure.
172/// Keys are split by separator and nested accordingly. Numeric keys
173/// stay as object keys (not converted to arrays).
174#[derive(Debug)]
175pub struct JsonbUnflatten {
176    signature: Signature,
177}
178
179impl JsonbUnflatten {
180    /// Creates a new `jsonb_unflatten` UDF.
181    #[must_use]
182    pub fn new() -> Self {
183        Self {
184            signature: Signature::new(
185                TypeSignature::Exact(vec![DataType::LargeBinary, DataType::Utf8]),
186                Volatility::Immutable,
187            ),
188        }
189    }
190}
191
192impl Default for JsonbUnflatten {
193    fn default() -> Self {
194        Self::new()
195    }
196}
197
198impl PartialEq for JsonbUnflatten {
199    fn eq(&self, _other: &Self) -> bool {
200        true
201    }
202}
203
204impl Eq for JsonbUnflatten {}
205
206impl Hash for JsonbUnflatten {
207    fn hash<H: Hasher>(&self, state: &mut H) {
208        "jsonb_unflatten".hash(state);
209    }
210}
211
212fn unflatten_insert(root: &mut serde_json::Value, parts: &[&str], value: serde_json::Value) {
213    if parts.is_empty() {
214        return;
215    }
216    if parts.len() == 1 {
217        if let serde_json::Value::Object(obj) = root {
218            obj.insert(parts[0].to_owned(), value);
219        }
220        return;
221    }
222    if let serde_json::Value::Object(obj) = root {
223        let child = obj
224            .entry(parts[0].to_owned())
225            .or_insert_with(|| serde_json::Value::Object(serde_json::Map::new()));
226        unflatten_insert(child, &parts[1..], value);
227    }
228}
229
230impl ScalarUDFImpl for JsonbUnflatten {
231    fn as_any(&self) -> &dyn std::any::Any {
232        self
233    }
234
235    fn name(&self) -> &'static str {
236        "jsonb_unflatten"
237    }
238
239    fn signature(&self) -> &Signature {
240        &self.signature
241    }
242
243    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
244        Ok(DataType::LargeBinary)
245    }
246
247    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
248        let expanded = expand_args(&args.args)?;
249        let jsonb_arr = expanded[0]
250            .as_any()
251            .downcast_ref::<LargeBinaryArray>()
252            .ok_or_else(|| {
253                datafusion_common::DataFusionError::Internal(
254                    "jsonb_unflatten: first arg must be LargeBinary".into(),
255                )
256            })?;
257        let sep_arr = expanded[1]
258            .as_any()
259            .downcast_ref::<StringArray>()
260            .ok_or_else(|| {
261                datafusion_common::DataFusionError::Internal(
262                    "jsonb_unflatten: second arg must be Utf8".into(),
263                )
264            })?;
265
266        let mut builder = LargeBinaryBuilder::with_capacity(jsonb_arr.len(), 256);
267        for i in 0..jsonb_arr.len() {
268            if jsonb_arr.is_null(i) || sep_arr.is_null(i) {
269                builder.append_null();
270            } else {
271                let sep = sep_arr.value(i);
272                match json_types::jsonb_to_value(jsonb_arr.value(i)) {
273                    Some(serde_json::Value::Object(flat)) => {
274                        let mut root = serde_json::Value::Object(serde_json::Map::new());
275                        for (key, val) in flat {
276                            let parts: Vec<&str> = key.split(sep).collect();
277                            unflatten_insert(&mut root, &parts, val);
278                        }
279                        builder.append_value(json_types::encode_jsonb(&root));
280                    }
281                    Some(other) => {
282                        builder.append_value(json_types::encode_jsonb(&other));
283                    }
284                    None => builder.append_null(),
285                }
286            }
287        }
288        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
289    }
290}
291
292// ══════════════════════════════════════════════════════════════════
293// json_to_columns(jsonb, text) -> jsonb  (runtime fallback)
294// ══════════════════════════════════════════════════════════════════
295
296/// `json_to_columns(jsonb, type_spec) -> jsonb`
297///
298/// Runtime fallback for structured extraction. Parses the type_spec
299/// to determine field names, extracts each from the JSONB object,
300/// and returns the result as a new JSONB object containing only
301/// those fields. Full plan-time struct rewriting is deferred.
302#[derive(Debug)]
303pub struct JsonToColumns {
304    signature: Signature,
305}
306
307impl JsonToColumns {
308    /// Creates a new `json_to_columns` UDF.
309    #[must_use]
310    pub fn new() -> Self {
311        Self {
312            signature: Signature::new(
313                TypeSignature::Exact(vec![DataType::LargeBinary, DataType::Utf8]),
314                Volatility::Immutable,
315            ),
316        }
317    }
318}
319
320impl Default for JsonToColumns {
321    fn default() -> Self {
322        Self::new()
323    }
324}
325
326impl PartialEq for JsonToColumns {
327    fn eq(&self, _other: &Self) -> bool {
328        true
329    }
330}
331
332impl Eq for JsonToColumns {}
333
334impl Hash for JsonToColumns {
335    fn hash<H: Hasher>(&self, state: &mut H) {
336        "json_to_columns".hash(state);
337    }
338}
339
340/// Parse a type_spec like `"name VARCHAR, age BIGINT, active BOOLEAN"`
341/// into a list of field names.
342fn parse_type_spec_fields(spec: &str) -> Vec<String> {
343    spec.split(',')
344        .filter_map(|part| {
345            let trimmed = part.trim();
346            trimmed.split_whitespace().next().map(ToOwned::to_owned)
347        })
348        .collect()
349}
350
351impl ScalarUDFImpl for JsonToColumns {
352    fn as_any(&self) -> &dyn std::any::Any {
353        self
354    }
355
356    fn name(&self) -> &'static str {
357        "json_to_columns"
358    }
359
360    fn signature(&self) -> &Signature {
361        &self.signature
362    }
363
364    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
365        Ok(DataType::LargeBinary)
366    }
367
368    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
369        let expanded = expand_args(&args.args)?;
370        let jsonb_arr = expanded[0]
371            .as_any()
372            .downcast_ref::<LargeBinaryArray>()
373            .ok_or_else(|| {
374                datafusion_common::DataFusionError::Internal(
375                    "json_to_columns: first arg must be LargeBinary".into(),
376                )
377            })?;
378        let spec_arr = expanded[1]
379            .as_any()
380            .downcast_ref::<StringArray>()
381            .ok_or_else(|| {
382                datafusion_common::DataFusionError::Internal(
383                    "json_to_columns: second arg must be Utf8".into(),
384                )
385            })?;
386
387        let mut builder = LargeBinaryBuilder::with_capacity(jsonb_arr.len(), 256);
388        for i in 0..jsonb_arr.len() {
389            if jsonb_arr.is_null(i) || spec_arr.is_null(i) {
390                builder.append_null();
391            } else {
392                let fields = parse_type_spec_fields(spec_arr.value(i));
393                let field_set: HashSet<&str> = fields.iter().map(String::as_str).collect();
394                match json_types::jsonb_to_value(jsonb_arr.value(i)) {
395                    Some(serde_json::Value::Object(obj)) => {
396                        let picked: serde_json::Map<String, serde_json::Value> = obj
397                            .into_iter()
398                            .filter(|(k, _)| field_set.contains(k.as_str()))
399                            .collect();
400                        builder.append_value(json_types::encode_jsonb(&serde_json::Value::Object(
401                            picked,
402                        )));
403                    }
404                    Some(other) => {
405                        builder.append_value(json_types::encode_jsonb(&other));
406                    }
407                    None => builder.append_null(),
408                }
409            }
410        }
411        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
412    }
413}
414
415// ══════════════════════════════════════════════════════════════════
416// json_infer_schema(jsonb) -> text
417// ══════════════════════════════════════════════════════════════════
418
419/// `json_infer_schema(jsonb) -> text`
420///
421/// Infers the SQL schema of a JSONB value, returning a JSON object
422/// mapping field names to SQL type names.
423#[derive(Debug)]
424pub struct JsonInferSchema {
425    signature: Signature,
426}
427
428impl JsonInferSchema {
429    /// Creates a new `json_infer_schema` UDF.
430    #[must_use]
431    pub fn new() -> Self {
432        Self {
433            signature: Signature::new(
434                TypeSignature::Exact(vec![DataType::LargeBinary]),
435                Volatility::Immutable,
436            ),
437        }
438    }
439}
440
441impl Default for JsonInferSchema {
442    fn default() -> Self {
443        Self::new()
444    }
445}
446
447impl PartialEq for JsonInferSchema {
448    fn eq(&self, _other: &Self) -> bool {
449        true
450    }
451}
452
453impl Eq for JsonInferSchema {}
454
455impl Hash for JsonInferSchema {
456    fn hash<H: Hasher>(&self, state: &mut H) {
457        "json_infer_schema".hash(state);
458    }
459}
460
461fn infer_type(val: &serde_json::Value) -> String {
462    match val {
463        serde_json::Value::Null => "NULL".to_owned(),
464        serde_json::Value::Bool(_) => "BOOLEAN".to_owned(),
465        serde_json::Value::Number(n) => {
466            if n.is_i64() || n.is_u64() {
467                "BIGINT".to_owned()
468            } else {
469                "DOUBLE".to_owned()
470            }
471        }
472        serde_json::Value::String(_) => "VARCHAR".to_owned(),
473        serde_json::Value::Array(arr) => {
474            let inner = arr.first().map_or("NULL".to_owned(), infer_type);
475            format!("ARRAY<{inner}>")
476        }
477        serde_json::Value::Object(obj) => {
478            let fields: Vec<String> = obj
479                .iter()
480                .map(|(k, v)| format!("{k} {}", infer_type(v)))
481                .collect();
482            format!("STRUCT({})", fields.join(", "))
483        }
484    }
485}
486
487impl ScalarUDFImpl for JsonInferSchema {
488    fn as_any(&self) -> &dyn std::any::Any {
489        self
490    }
491
492    fn name(&self) -> &'static str {
493        "json_infer_schema"
494    }
495
496    fn signature(&self) -> &Signature {
497        &self.signature
498    }
499
500    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
501        Ok(DataType::Utf8)
502    }
503
504    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
505        let expanded = expand_args(&args.args)?;
506        let jsonb_arr = expanded[0]
507            .as_any()
508            .downcast_ref::<LargeBinaryArray>()
509            .ok_or_else(|| {
510                datafusion_common::DataFusionError::Internal(
511                    "json_infer_schema: arg must be LargeBinary".into(),
512                )
513            })?;
514
515        let mut builder = StringBuilder::with_capacity(jsonb_arr.len(), 256);
516        for i in 0..jsonb_arr.len() {
517            if jsonb_arr.is_null(i) {
518                builder.append_null();
519            } else {
520                match json_types::jsonb_to_value(jsonb_arr.value(i)) {
521                    Some(serde_json::Value::Object(obj)) => {
522                        let schema: serde_json::Map<String, serde_json::Value> = obj
523                            .iter()
524                            .map(|(k, v)| (k.clone(), serde_json::Value::String(infer_type(v))))
525                            .collect();
526                        builder.append_value(
527                            serde_json::to_string(&serde_json::Value::Object(schema))
528                                .unwrap_or_default(),
529                        );
530                    }
531                    Some(val) => {
532                        builder.append_value(infer_type(&val));
533                    }
534                    None => builder.append_null(),
535                }
536            }
537        }
538        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
539    }
540}
laminar_sql/datafusion/json_extensions/shape_ops.rs

laminar_sql/datafusion/json_extensions/
shape_ops.rs
