laminar_sql/datafusion/json_udf/
construction.rs

1//! JSON type inspection, construction, and scalar conversion UDFs.
2
3use std::hash::{Hash, Hasher};
4use std::sync::Arc;
5
6use arrow::datatypes::DataType;
7use arrow_array::{
8    builder::{LargeBinaryBuilder, StringBuilder},
9    Array, ArrayRef, BooleanArray, LargeBinaryArray, StringArray,
10};
11use datafusion_common::Result;
12use datafusion_expr::{
13    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility,
14};
15
16use super::{expand_args, json_types};
17
18/// Returns the outermost JSON type name by reading its tag byte.
19#[derive(Debug)]
20pub struct JsonTypeof {
21    signature: Signature,
22}
23
24impl JsonTypeof {
25    /// Creates a new `json_typeof` UDF.
26    #[must_use]
27    pub fn new() -> Self {
28        Self {
29            signature: Signature::new(
30                TypeSignature::Exact(vec![DataType::LargeBinary]),
31                Volatility::Immutable,
32            ),
33        }
34    }
35}
36
37impl Default for JsonTypeof {
38    fn default() -> Self {
39        Self::new()
40    }
41}
42
43impl PartialEq for JsonTypeof {
44    fn eq(&self, _other: &Self) -> bool {
45        true
46    }
47}
48
49impl Eq for JsonTypeof {}
50
51impl Hash for JsonTypeof {
52    fn hash<H: Hasher>(&self, state: &mut H) {
53        "json_typeof".hash(state);
54    }
55}
56
57impl ScalarUDFImpl for JsonTypeof {
58    fn as_any(&self) -> &dyn std::any::Any {
59        self
60    }
61
62    fn name(&self) -> &'static str {
63        "json_typeof"
64    }
65
66    fn signature(&self) -> &Signature {
67        &self.signature
68    }
69
70    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
71        Ok(DataType::Utf8)
72    }
73
74    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
75        let expanded = expand_args(&args.args)?;
76        let jsonb_arr = expanded[0]
77            .as_any()
78            .downcast_ref::<LargeBinaryArray>()
79            .ok_or_else(|| {
80                datafusion_common::DataFusionError::Internal(
81                    "json_typeof: arg must be LargeBinary".into(),
82                )
83            })?;
84
85        let mut builder = StringBuilder::with_capacity(jsonb_arr.len(), jsonb_arr.len() * 8);
86        for i in 0..jsonb_arr.len() {
87            if jsonb_arr.is_null(i) {
88                builder.append_null();
89            } else {
90                match json_types::jsonb_type_name(jsonb_arr.value(i)) {
91                    Some(name) => builder.append_value(name),
92                    None => builder.append_null(),
93                }
94            }
95        }
96        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
97    }
98}
99
100// ══════════════════════════════════════════════════════════════════
101// json_build_object(k1, v1, k2, v2, ...) -> jsonb
102// ══════════════════════════════════════════════════════════════════
103
104/// `json_build_object(key1, value1, key2, value2, ...) -> jsonb`
105///
106/// Constructs a JSONB object from alternating key-value pairs.
107/// Executes in Ring 1 (allocates JSONB binary buffer).
108#[derive(Debug)]
109pub struct JsonBuildObject {
110    signature: Signature,
111}
112
113impl JsonBuildObject {
114    /// Creates a new `json_build_object` UDF.
115    #[must_use]
116    pub fn new() -> Self {
117        Self {
118            signature: Signature::new(TypeSignature::VariadicAny, Volatility::Immutable),
119        }
120    }
121}
122
123impl Default for JsonBuildObject {
124    fn default() -> Self {
125        Self::new()
126    }
127}
128
129impl PartialEq for JsonBuildObject {
130    fn eq(&self, _other: &Self) -> bool {
131        true
132    }
133}
134
135impl Eq for JsonBuildObject {}
136
137impl Hash for JsonBuildObject {
138    fn hash<H: Hasher>(&self, state: &mut H) {
139        "json_build_object".hash(state);
140    }
141}
142
143impl ScalarUDFImpl for JsonBuildObject {
144    fn as_any(&self) -> &dyn std::any::Any {
145        self
146    }
147
148    fn name(&self) -> &'static str {
149        "json_build_object"
150    }
151
152    fn signature(&self) -> &Signature {
153        &self.signature
154    }
155
156    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
157        Ok(DataType::LargeBinary)
158    }
159
160    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
161        if !args.args.len().is_multiple_of(2) {
162            return Err(datafusion_common::DataFusionError::Execution(
163                "json_build_object requires an even number of arguments".into(),
164            ));
165        }
166
167        let expanded = expand_args(&args.args)?;
168        let len = expanded.first().map_or(1, Array::len);
169
170        let mut builder = LargeBinaryBuilder::with_capacity(len, 256);
171        for row in 0..len {
172            let mut obj = serde_json::Map::new();
173            let mut is_null = false;
174            for pair in expanded.chunks(2) {
175                let key_arr = &pair[0];
176                let val_arr = &pair[1];
177                if key_arr.is_null(row) {
178                    is_null = true;
179                    break;
180                }
181                let key = scalar_to_json_key(key_arr, row)?;
182                let val = scalar_to_json_value(val_arr, row);
183                obj.insert(key, val);
184            }
185            if is_null {
186                builder.append_null();
187            } else {
188                let jsonb = json_types::encode_jsonb(&serde_json::Value::Object(obj));
189                builder.append_value(&jsonb);
190            }
191        }
192        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
193    }
194}
195
196// ══════════════════════════════════════════════════════════════════
197// json_build_array(v1, v2, ...) -> jsonb
198// ══════════════════════════════════════════════════════════════════
199
200/// `json_build_array(v1, v2, ...) -> jsonb`
201///
202/// Constructs a JSONB array from the given values.
203/// Executes in Ring 1 (allocates JSONB binary buffer).
204#[derive(Debug)]
205pub struct JsonBuildArray {
206    signature: Signature,
207}
208
209impl JsonBuildArray {
210    /// Creates a new `json_build_array` UDF.
211    #[must_use]
212    pub fn new() -> Self {
213        Self {
214            signature: Signature::new(TypeSignature::VariadicAny, Volatility::Immutable),
215        }
216    }
217}
218
219impl Default for JsonBuildArray {
220    fn default() -> Self {
221        Self::new()
222    }
223}
224
225impl PartialEq for JsonBuildArray {
226    fn eq(&self, _other: &Self) -> bool {
227        true
228    }
229}
230
231impl Eq for JsonBuildArray {}
232
233impl Hash for JsonBuildArray {
234    fn hash<H: Hasher>(&self, state: &mut H) {
235        "json_build_array".hash(state);
236    }
237}
238
239impl ScalarUDFImpl for JsonBuildArray {
240    fn as_any(&self) -> &dyn std::any::Any {
241        self
242    }
243
244    fn name(&self) -> &'static str {
245        "json_build_array"
246    }
247
248    fn signature(&self) -> &Signature {
249        &self.signature
250    }
251
252    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
253        Ok(DataType::LargeBinary)
254    }
255
256    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
257        let expanded = expand_args(&args.args)?;
258        let len = expanded.first().map_or(1, Array::len);
259
260        let mut builder = LargeBinaryBuilder::with_capacity(len, 256);
261        for row in 0..len {
262            let mut arr = Vec::with_capacity(expanded.len());
263            for col in &expanded {
264                arr.push(scalar_to_json_value(col, row));
265            }
266            let jsonb = json_types::encode_jsonb(&serde_json::Value::Array(arr));
267            builder.append_value(&jsonb);
268        }
269        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
270    }
271}
272
273// ══════════════════════════════════════════════════════════════════
274// to_jsonb(any) -> jsonb
275// ══════════════════════════════════════════════════════════════════
276
277/// `to_jsonb(value) -> jsonb`
278///
279/// Converts any SQL value to JSONB binary format.
280#[derive(Debug)]
281pub struct ToJsonb {
282    signature: Signature,
283}
284
285impl ToJsonb {
286    /// Creates a new `to_jsonb` UDF.
287    #[must_use]
288    pub fn new() -> Self {
289        Self {
290            signature: Signature::new(TypeSignature::Any(1), Volatility::Immutable),
291        }
292    }
293}
294
295impl Default for ToJsonb {
296    fn default() -> Self {
297        Self::new()
298    }
299}
300
301impl PartialEq for ToJsonb {
302    fn eq(&self, _other: &Self) -> bool {
303        true
304    }
305}
306
307impl Eq for ToJsonb {}
308
309impl Hash for ToJsonb {
310    fn hash<H: Hasher>(&self, state: &mut H) {
311        "to_jsonb".hash(state);
312    }
313}
314
315impl ScalarUDFImpl for ToJsonb {
316    fn as_any(&self) -> &dyn std::any::Any {
317        self
318    }
319
320    fn name(&self) -> &'static str {
321        "to_jsonb"
322    }
323
324    fn signature(&self) -> &Signature {
325        &self.signature
326    }
327
328    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
329        Ok(DataType::LargeBinary)
330    }
331
332    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
333        let expanded = expand_args(&args.args)?;
334        let arr = &expanded[0];
335        let len = arr.len();
336
337        let mut builder = LargeBinaryBuilder::with_capacity(len, 64);
338        for row in 0..len {
339            if arr.is_null(row) {
340                let jsonb = json_types::encode_jsonb(&serde_json::Value::Null);
341                builder.append_value(&jsonb);
342            } else {
343                let val = scalar_to_json_value(arr, row);
344                let jsonb = json_types::encode_jsonb(&val);
345                builder.append_value(&jsonb);
346            }
347        }
348        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
349    }
350}
351
352// ── Scalar conversion helpers ────────────────────────────────────
353
354/// Convert an Arrow array value at `row` to a JSON key string.
355fn scalar_to_json_key(arr: &ArrayRef, row: usize) -> Result<String> {
356    if let Some(s) = arr.as_any().downcast_ref::<StringArray>() {
357        return Ok(s.value(row).to_owned());
358    }
359    // Fallback: convert to string representation
360    Err(datafusion_common::DataFusionError::Execution(
361        "json_build_object keys must be text".into(),
362    ))
363}
364
365/// Convert an Arrow array value at `row` to a `serde_json::Value`.
366fn scalar_to_json_value(arr: &ArrayRef, row: usize) -> serde_json::Value {
367    if arr.is_null(row) {
368        return serde_json::Value::Null;
369    }
370
371    // Try common types
372    if let Some(a) = arr.as_any().downcast_ref::<StringArray>() {
373        return serde_json::Value::String(a.value(row).to_owned());
374    }
375    if let Some(a) = arr.as_any().downcast_ref::<arrow_array::Int64Array>() {
376        return serde_json::Value::Number(a.value(row).into());
377    }
378    if let Some(a) = arr.as_any().downcast_ref::<arrow_array::Int32Array>() {
379        return serde_json::Value::Number(i64::from(a.value(row)).into());
380    }
381    if let Some(a) = arr.as_any().downcast_ref::<arrow_array::Float64Array>() {
382        if let Some(n) = serde_json::Number::from_f64(a.value(row)) {
383            return serde_json::Value::Number(n);
384        }
385        return serde_json::Value::Null;
386    }
387    if let Some(a) = arr.as_any().downcast_ref::<BooleanArray>() {
388        return serde_json::Value::Bool(a.value(row));
389    }
390    if let Some(a) = arr.as_any().downcast_ref::<LargeBinaryArray>() {
391        // Already JSONB — convert to JSON value for re-encoding
392        let bytes = a.value(row);
393        if let Some(text) = json_types::jsonb_to_text(bytes) {
394            // Try to parse the text as JSON
395            if let Ok(val) = serde_json::from_str::<serde_json::Value>(&text) {
396                return val;
397            }
398            return serde_json::Value::String(text);
399        }
400        return serde_json::Value::Null;
401    }
402
403    // Fallback: use display
404    let scalar = datafusion_common::ScalarValue::try_from_array(arr, row).ok();
405    match scalar {
406        Some(s) => serde_json::Value::String(s.to_string()),
407        None => serde_json::Value::Null,
408    }
409}
410
411// ══════════════════════════════════════════════════════════════════
412// Tests
413// ══════════════════════════════════════════════════════════════════
laminar_sql/datafusion/json_udf/construction.rs

laminar_sql/datafusion/json_udf/
construction.rs
