laminar_sql/datafusion/
json_udaf.rs

1//! PostgreSQL-compatible JSON aggregate UDFs (F-SCHEMA-011).
2//!
3//! - [`JsonAgg`] — `json_agg(expr) -> jsonb` — collects values into a JSON array
4//! - [`JsonObjectAgg`] — `json_object_agg(key, value) -> jsonb` — collects
5//!   key-value pairs into a JSON object
6
7use std::hash::{Hash, Hasher};
8use std::sync::Arc;
9
10use arrow::datatypes::DataType;
11use arrow_array::{Array, ArrayRef, LargeBinaryArray, StringArray};
12use arrow_schema::Field;
13use datafusion_common::{Result, ScalarValue};
14use datafusion_expr::function::AccumulatorArgs;
15use datafusion_expr::{Accumulator, AggregateUDFImpl, Signature, TypeSignature, Volatility};
16
17use super::json_types;
18
19// ══════════════════════════════════════════════════════════════════
20// json_agg(expression) -> jsonb
21// ══════════════════════════════════════════════════════════════════
22
23/// `json_agg(expression) -> jsonb`
24///
25/// Collects all values of an expression into a JSON array.
26/// Executes in Ring 1 via the DataFusion aggregate bridge.
27#[derive(Debug)]
28pub struct JsonAgg {
29    signature: Signature,
30}
31
32impl JsonAgg {
33    /// Creates a new `json_agg` UDAF.
34    #[must_use]
35    pub fn new() -> Self {
36        Self {
37            signature: Signature::new(TypeSignature::Any(1), Volatility::Immutable),
38        }
39    }
40}
41
42impl Default for JsonAgg {
43    fn default() -> Self {
44        Self::new()
45    }
46}
47
48impl PartialEq for JsonAgg {
49    fn eq(&self, _other: &Self) -> bool {
50        true
51    }
52}
53
54impl Eq for JsonAgg {}
55
56impl Hash for JsonAgg {
57    fn hash<H: Hasher>(&self, state: &mut H) {
58        "json_agg".hash(state);
59    }
60}
61
62impl AggregateUDFImpl for JsonAgg {
63    fn as_any(&self) -> &dyn std::any::Any {
64        self
65    }
66
67    fn name(&self) -> &'static str {
68        "json_agg"
69    }
70
71    fn signature(&self) -> &Signature {
72        &self.signature
73    }
74
75    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
76        Ok(DataType::LargeBinary)
77    }
78
79    fn state_fields(
80        &self,
81        _args: datafusion_expr::function::StateFieldsArgs,
82    ) -> Result<Vec<Arc<Field>>> {
83        // State is a single LargeBinary holding concatenated JSONB values
84        Ok(vec![Arc::new(Field::new(
85            "json_agg_state",
86            DataType::LargeBinary,
87            true,
88        ))])
89    }
90
91    fn accumulator(&self, _args: AccumulatorArgs<'_>) -> Result<Box<dyn Accumulator>> {
92        Ok(Box::new(JsonAggAccumulator::new()))
93    }
94}
95
96/// Accumulator for `json_agg`.
97///
98/// Maintains a `Vec` of JSON values. On `evaluate`, serializes to JSONB array.
99#[derive(Debug)]
100struct JsonAggAccumulator {
101    values: Vec<serde_json::Value>,
102}
103
104impl JsonAggAccumulator {
105    fn new() -> Self {
106        Self { values: Vec::new() }
107    }
108}
109
110impl Accumulator for JsonAggAccumulator {
111    fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
112        let arr = &values[0];
113        for i in 0..arr.len() {
114            if arr.is_null(i) {
115                self.values.push(serde_json::Value::Null);
116            } else {
117                self.values.push(array_value_to_json(arr, i));
118            }
119        }
120        Ok(())
121    }
122
123    fn evaluate(&mut self) -> Result<ScalarValue> {
124        let json_arr = serde_json::Value::Array(self.values.clone());
125        let bytes = json_types::encode_jsonb(&json_arr);
126        Ok(ScalarValue::LargeBinary(Some(bytes)))
127    }
128
129    fn size(&self) -> usize {
130        std::mem::size_of::<Self>()
131            + self.values.capacity() * std::mem::size_of::<serde_json::Value>()
132    }
133
134    fn state(&mut self) -> Result<Vec<ScalarValue>> {
135        // Serialize current state as a JSONB array
136        let json_arr = serde_json::Value::Array(self.values.clone());
137        let bytes = json_types::encode_jsonb(&json_arr);
138        Ok(vec![ScalarValue::LargeBinary(Some(bytes))])
139    }
140
141    fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
142        let arr = states[0]
143            .as_any()
144            .downcast_ref::<LargeBinaryArray>()
145            .ok_or_else(|| {
146                datafusion_common::DataFusionError::Internal(
147                    "json_agg: merge state must be LargeBinary".into(),
148                )
149            })?;
150        for i in 0..arr.len() {
151            if !arr.is_null(i) {
152                let bytes = arr.value(i);
153                // Decode JSONB array and merge elements
154                if let Some(json_str) = json_types::jsonb_to_text(bytes) {
155                    if let Ok(serde_json::Value::Array(elems)) =
156                        serde_json::from_str::<serde_json::Value>(&json_str)
157                    {
158                        self.values.extend(elems);
159                    }
160                }
161            }
162        }
163        Ok(())
164    }
165}
166
167// ══════════════════════════════════════════════════════════════════
168// json_object_agg(key, value) -> jsonb
169// ══════════════════════════════════════════════════════════════════
170
171/// `json_object_agg(key, value) -> jsonb`
172///
173/// Collects key-value pairs into a JSON object. Duplicate keys use
174/// last-value-wins semantics (consistent with PostgreSQL).
175#[derive(Debug)]
176pub struct JsonObjectAgg {
177    signature: Signature,
178}
179
180impl JsonObjectAgg {
181    /// Creates a new `json_object_agg` UDAF.
182    #[must_use]
183    pub fn new() -> Self {
184        Self {
185            signature: Signature::new(TypeSignature::Any(2), Volatility::Immutable),
186        }
187    }
188}
189
190impl Default for JsonObjectAgg {
191    fn default() -> Self {
192        Self::new()
193    }
194}
195
196impl PartialEq for JsonObjectAgg {
197    fn eq(&self, _other: &Self) -> bool {
198        true
199    }
200}
201
202impl Eq for JsonObjectAgg {}
203
204impl Hash for JsonObjectAgg {
205    fn hash<H: Hasher>(&self, state: &mut H) {
206        "json_object_agg".hash(state);
207    }
208}
209
210impl AggregateUDFImpl for JsonObjectAgg {
211    fn as_any(&self) -> &dyn std::any::Any {
212        self
213    }
214
215    fn name(&self) -> &'static str {
216        "json_object_agg"
217    }
218
219    fn signature(&self) -> &Signature {
220        &self.signature
221    }
222
223    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
224        Ok(DataType::LargeBinary)
225    }
226
227    fn state_fields(
228        &self,
229        _args: datafusion_expr::function::StateFieldsArgs,
230    ) -> Result<Vec<Arc<Field>>> {
231        Ok(vec![Arc::new(Field::new(
232            "json_object_agg_state",
233            DataType::LargeBinary,
234            true,
235        ))])
236    }
237
238    fn accumulator(&self, _args: AccumulatorArgs<'_>) -> Result<Box<dyn Accumulator>> {
239        Ok(Box::new(JsonObjectAggAccumulator::new()))
240    }
241}
242
243/// Accumulator for `json_object_agg`.
244///
245/// Maintains an ordered map of key-value pairs.
246#[derive(Debug)]
247struct JsonObjectAggAccumulator {
248    entries: serde_json::Map<String, serde_json::Value>,
249}
250
251impl JsonObjectAggAccumulator {
252    fn new() -> Self {
253        Self {
254            entries: serde_json::Map::new(),
255        }
256    }
257}
258
259impl Accumulator for JsonObjectAggAccumulator {
260    fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
261        let key_arr = &values[0];
262        let val_arr = &values[1];
263
264        for i in 0..key_arr.len() {
265            if key_arr.is_null(i) {
266                continue; // Skip null keys (PostgreSQL behavior)
267            }
268            let key = array_value_to_string(key_arr, i)?;
269            let val = if val_arr.is_null(i) {
270                serde_json::Value::Null
271            } else {
272                array_value_to_json(val_arr, i)
273            };
274            self.entries.insert(key, val); // last-value-wins
275        }
276        Ok(())
277    }
278
279    fn evaluate(&mut self) -> Result<ScalarValue> {
280        let obj = serde_json::Value::Object(self.entries.clone());
281        let bytes = json_types::encode_jsonb(&obj);
282        Ok(ScalarValue::LargeBinary(Some(bytes)))
283    }
284
285    fn size(&self) -> usize {
286        std::mem::size_of::<Self>() + self.entries.len() * 64 // rough estimate
287    }
288
289    fn state(&mut self) -> Result<Vec<ScalarValue>> {
290        let obj = serde_json::Value::Object(self.entries.clone());
291        let bytes = json_types::encode_jsonb(&obj);
292        Ok(vec![ScalarValue::LargeBinary(Some(bytes))])
293    }
294
295    fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
296        let arr = states[0]
297            .as_any()
298            .downcast_ref::<LargeBinaryArray>()
299            .ok_or_else(|| {
300                datafusion_common::DataFusionError::Internal(
301                    "json_object_agg: merge state must be LargeBinary".into(),
302                )
303            })?;
304        for i in 0..arr.len() {
305            if !arr.is_null(i) {
306                let bytes = arr.value(i);
307                if let Some(json_str) = json_types::jsonb_to_text(bytes) {
308                    if let Ok(serde_json::Value::Object(map)) =
309                        serde_json::from_str::<serde_json::Value>(&json_str)
310                    {
311                        for (k, v) in map {
312                            self.entries.insert(k, v);
313                        }
314                    }
315                }
316            }
317        }
318        Ok(())
319    }
320}
321
322// ── Helpers ──────────────────────────────────────────────────────
323
324/// Convert an Arrow array element to a `serde_json::Value`.
325fn array_value_to_json(arr: &ArrayRef, row: usize) -> serde_json::Value {
326    if arr.is_null(row) {
327        return serde_json::Value::Null;
328    }
329    if let Some(a) = arr.as_any().downcast_ref::<StringArray>() {
330        return serde_json::Value::String(a.value(row).to_owned());
331    }
332    if let Some(a) = arr.as_any().downcast_ref::<arrow_array::Int64Array>() {
333        return serde_json::Value::Number(a.value(row).into());
334    }
335    if let Some(a) = arr.as_any().downcast_ref::<arrow_array::Int32Array>() {
336        return serde_json::Value::Number(i64::from(a.value(row)).into());
337    }
338    if let Some(a) = arr.as_any().downcast_ref::<arrow_array::Float64Array>() {
339        if let Some(n) = serde_json::Number::from_f64(a.value(row)) {
340            return serde_json::Value::Number(n);
341        }
342        return serde_json::Value::Null;
343    }
344    if let Some(a) = arr.as_any().downcast_ref::<arrow_array::BooleanArray>() {
345        return serde_json::Value::Bool(a.value(row));
346    }
347    // Fallback
348    let scalar = ScalarValue::try_from_array(arr, row).ok();
349    match scalar {
350        Some(s) => serde_json::Value::String(s.to_string()),
351        None => serde_json::Value::Null,
352    }
353}
354
355/// Extract a string key from an Arrow array.
356fn array_value_to_string(arr: &ArrayRef, row: usize) -> Result<String> {
357    if let Some(a) = arr.as_any().downcast_ref::<StringArray>() {
358        return Ok(a.value(row).to_owned());
359    }
360    // Fallback via ScalarValue display
361    let sv = ScalarValue::try_from_array(arr, row)?;
362    Ok(sv.to_string())
363}
364
365#[cfg(test)]
366mod tests {
367    use super::*;
368    use arrow_array::Int64Array;
369
370    fn make_string_array(vals: &[&str]) -> StringArray {
371        StringArray::from(vals.to_vec())
372    }
373
374    #[test]
375    fn test_json_agg_basic() {
376        let mut acc = JsonAggAccumulator::new();
377        let vals = Arc::new(Int64Array::from(vec![1, 2, 3])) as ArrayRef;
378        acc.update_batch(&[vals]).unwrap();
379        let result = acc.evaluate().unwrap();
380        match result {
381            ScalarValue::LargeBinary(Some(bytes)) => {
382                assert_eq!(json_types::jsonb_type_name(&bytes), Some("array"));
383                let e0 = json_types::jsonb_array_get(&bytes, 0).unwrap();
384                assert_eq!(json_types::jsonb_to_text(e0), Some("1".to_owned()));
385                let e2 = json_types::jsonb_array_get(&bytes, 2).unwrap();
386                assert_eq!(json_types::jsonb_to_text(e2), Some("3".to_owned()));
387            }
388            other => panic!("Expected LargeBinary, got {other:?}"),
389        }
390    }
391
392    #[test]
393    fn test_json_agg_strings() {
394        let mut acc = JsonAggAccumulator::new();
395        let vals = Arc::new(make_string_array(&["a", "b", "c"])) as ArrayRef;
396        acc.update_batch(&[vals]).unwrap();
397        let result = acc.evaluate().unwrap();
398        match result {
399            ScalarValue::LargeBinary(Some(bytes)) => {
400                let e0 = json_types::jsonb_array_get(&bytes, 0).unwrap();
401                assert_eq!(json_types::jsonb_to_text(e0), Some("a".to_owned()));
402            }
403            other => panic!("Expected LargeBinary, got {other:?}"),
404        }
405    }
406
407    #[test]
408    fn test_json_agg_multiple_batches() {
409        let mut acc = JsonAggAccumulator::new();
410        let v1 = Arc::new(Int64Array::from(vec![1, 2])) as ArrayRef;
411        let v2 = Arc::new(Int64Array::from(vec![3])) as ArrayRef;
412        acc.update_batch(&[v1]).unwrap();
413        acc.update_batch(&[v2]).unwrap();
414        let result = acc.evaluate().unwrap();
415        match result {
416            ScalarValue::LargeBinary(Some(bytes)) => {
417                // Should have 3 elements total
418                let text = json_types::jsonb_to_text(&bytes).unwrap();
419                assert_eq!(text, "[1,2,3]");
420            }
421            other => panic!("Expected LargeBinary, got {other:?}"),
422        }
423    }
424
425    #[test]
426    fn test_json_object_agg_basic() {
427        let mut acc = JsonObjectAggAccumulator::new();
428        let keys = Arc::new(make_string_array(&["a", "b", "c"])) as ArrayRef;
429        let vals = Arc::new(Int64Array::from(vec![1, 2, 3])) as ArrayRef;
430        acc.update_batch(&[keys, vals]).unwrap();
431        let result = acc.evaluate().unwrap();
432        match result {
433            ScalarValue::LargeBinary(Some(bytes)) => {
434                assert_eq!(json_types::jsonb_type_name(&bytes), Some("object"));
435                let a = json_types::jsonb_get_field(&bytes, "a").unwrap();
436                assert_eq!(json_types::jsonb_to_text(a), Some("1".to_owned()));
437                let c = json_types::jsonb_get_field(&bytes, "c").unwrap();
438                assert_eq!(json_types::jsonb_to_text(c), Some("3".to_owned()));
439            }
440            other => panic!("Expected LargeBinary, got {other:?}"),
441        }
442    }
443
444    #[test]
445    fn test_json_object_agg_last_value_wins() {
446        let mut acc = JsonObjectAggAccumulator::new();
447        let keys = Arc::new(make_string_array(&["a", "a"])) as ArrayRef;
448        let vals = Arc::new(Int64Array::from(vec![1, 2])) as ArrayRef;
449        acc.update_batch(&[keys, vals]).unwrap();
450        let result = acc.evaluate().unwrap();
451        match result {
452            ScalarValue::LargeBinary(Some(bytes)) => {
453                let a = json_types::jsonb_get_field(&bytes, "a").unwrap();
454                assert_eq!(json_types::jsonb_to_text(a), Some("2".to_owned()));
455            }
456            other => panic!("Expected LargeBinary, got {other:?}"),
457        }
458    }
459
460    #[test]
461    fn test_json_agg_state_merge() {
462        let mut acc1 = JsonAggAccumulator::new();
463        let v1 = Arc::new(Int64Array::from(vec![1, 2])) as ArrayRef;
464        acc1.update_batch(&[v1]).unwrap();
465        let state = acc1.state().unwrap();
466
467        let mut acc2 = JsonAggAccumulator::new();
468        let v2 = Arc::new(Int64Array::from(vec![3])) as ArrayRef;
469        acc2.update_batch(&[v2]).unwrap();
470
471        // Merge state from acc1 into acc2
472        let state_arr: ArrayRef = match &state[0] {
473            ScalarValue::LargeBinary(Some(b)) => {
474                Arc::new(LargeBinaryArray::from_iter_values(vec![b.as_slice()]))
475            }
476            _ => panic!("expected LargeBinary state"),
477        };
478        acc2.merge_batch(&[state_arr]).unwrap();
479
480        let result = acc2.evaluate().unwrap();
481        match result {
482            ScalarValue::LargeBinary(Some(bytes)) => {
483                let text = json_types::jsonb_to_text(&bytes).unwrap();
484                assert_eq!(text, "[3,1,2]");
485            }
486            other => panic!("Expected LargeBinary, got {other:?}"),
487        }
488    }
489
490    #[test]
491    fn test_udaf_registration() {
492        let json_agg = datafusion_expr::AggregateUDF::new_from_impl(JsonAgg::new());
493        assert_eq!(json_agg.name(), "json_agg");
494
495        let json_obj_agg = datafusion_expr::AggregateUDF::new_from_impl(JsonObjectAgg::new());
496        assert_eq!(json_obj_agg.name(), "json_object_agg");
497    }
498}
laminar_sql/datafusion/json_udaf.rs

laminar_sql/datafusion/
json_udaf.rs
