builtin function: json[b]_object_agg

Add basic support of json_object_agg. Current implementation of the
aggregator only allows agg over a single column. We implement
a workaround where we inject a call to `json_build_object` over the
key and value columns, hard-code an alias for that call,
and give that alias to the aggregator.

Release-Note-Core: Added support for builtin functions
  `json[b]_object_agg` and `json_objectagg`.
Fixes: REA-4369
Change-Id: I5c5548e0a491b1ba1adf49c0699f2308c8ab0acf
Reviewed-on: https://gerrit.readyset.name/c/readyset/+/8848
Reviewed-by: Vassili Zarouba <vassili@readyset.io>
Tested-by: Buildkite CI

Author: MohamedAbdeen-rs

Cargo.lock

@@ -19,7 +19,7 @@ macro_rules! non_null {
 }
 
 pub(crate) mod builtins;
-mod json;
+pub mod json;
 
 fn eval_binary_op(op: BinaryOperator, left: &DfValue, right: &DfValue) -> ReadySetResult<DfValue> {
     use BinaryOperator::*;
@@ -718,7 +718,7 @@ mod tests {
         assert_eq!(test_eval(&json_arr1, &json_obj1), expected);
 
         let expected = json!([99, 100]);
-        assert_eq!(test_eval(&json!(99), &json!(100)), expected);
+        assert_eq!(test_eval(&json!([99]), &json!([100])), expected);
 
         // Test error conditions
 

dataflow-expression/src/eval.rs

@@ -2551,11 +2551,6 @@ mod tests {
 
             test_identity("{}");
             test_identity("[]");
-            test_identity("1");
-            test_identity("2.0");
-            test_identity("true");
-            test_identity("false");
-            test_identity("null");
             test_identity("[null]");
         }
 

dataflow-expression/src/eval/builtins.rs

@@ -292,7 +292,7 @@ pub(crate) fn json_object_from_pairs(
     Ok(result.to_json_string()?.into())
 }
 
-pub(crate) fn json_object_from_keys_and_values(
+pub fn json_object_from_keys_and_values(
     keys: &Array,
     values: &Array,
     allow_duplicate_keys: bool,

dataflow-expression/src/eval/json.rs

@@ -1,5 +1,5 @@
 mod binary_operator;
-mod eval;
+pub mod eval;
 pub mod like;
 mod lower;
 mod promotion;

dataflow-expression/src/lib.rs

@@ -5,7 +5,7 @@ use std::sync::Arc;
 
 use partial_map::InsertionOrder;
 use readyset_data::DfValue;
-use readyset_errors::{internal, ReadySetResult};
+use readyset_errors::{internal, unsupported, ReadySetResult};
 use readyset_sql::ast::OrderType;
 use serde::{Deserialize, Serialize};
 
@@ -26,6 +26,8 @@ pub enum PostLookupAggregateFunction {
     Max,
     /// Take the minimum input value
     Min,
+    /// Use specified Key-value pair to build a JSON Object
+    JsonObjectAgg { allow_duplicate_keys: bool },
 }
 
 impl PostLookupAggregateFunction {
@@ -45,6 +47,9 @@ impl PostLookupAggregateFunction {
             .into()),
             PostLookupAggregateFunction::Max => Ok(cmp::max(val1, val2).clone()),
             PostLookupAggregateFunction::Min => Ok(cmp::min(val1, val2).clone()),
+            PostLookupAggregateFunction::JsonObjectAgg { .. } => {
+                unsupported!("JsonObjectAgg is not supported as a post-lookup aggregate")
+            }
         }
     }
 }

dataflow-expression/src/reader_processing.rs

@@ -0,0 +1,48 @@
+statement ok
+CREATE TABLE test (id INT, name TEXT PRIMARY KEY);
+
+statement ok
+INSERT INTO test (id, name) VALUES (1, 'a'), (2, 'b'), (3, 'c'), (2, 'd'), (3, 'e');
+
+query T
+SELECT json_object_agg(id, name) FROM test;
+----
+{"1":"a","2":"b","3":"c","2":"d","3":"e"}
+
+query T
+SELECT jsonb_object_agg(id, name) FROM test;
+----
+{"1":"a","2":"d","3":"e"}
+
+statement ok
+DELETE FROM test WHERE name = 'a';
+
+query T
+SELECT json_object_agg(id, name) FROM test;
+----
+{"2":"b","3":"c","2":"d","3":"e"}
+
+statement ok
+DELETE FROM test WHERE name = 'd';
+
+query T
+SELECT json_object_agg(id, name) FROM test;
+----
+{"2":"b","3":"c","3":"e"}
+
+query T
+SELECT jsonb_object_agg(id, name) FROM test;
+----
+{"2":"b","3":"e"}
+
+statement ok
+CREATE TABLE test_commas (name TEXT PRIMARY KEY, id INT);
+
+statement ok
+INSERT INTO test_commas (name, id) VALUES ('Hello, World', 1), (',', 2);
+
+query T
+SELECT json_object_agg(name, id) FROM test_commas;
+----
+{",":2,"Hello, World":1}
+

logictests/json_object_agg.test

@@ -374,6 +374,33 @@ pub fn function_expr(
                     separator,
                 },
             ),
+            map(
+                tuple((
+                    alt((
+                        map(tag_no_case("json_object_agg"), |_| true),   //psql
+                        map(tag_no_case("jsonb_object_agg"), |_| false), //psql
+                        map(tag_no_case("json_objectagg"), |_| false),   //mysql, jsonb_objectagg
+                                                                         //not supported
+                    )),
+                    preceded(
+                        whitespace0,
+                        delimited(
+                            terminated(tag("("), whitespace0),
+                            tuple((
+                                expression(dialect),
+                                preceded(whitespace0, tag(",")),
+                                expression(dialect),
+                            )),
+                            preceded(whitespace0, tag(")")),
+                        ),
+                    ),
+                )),
+                |(is_json, (key, _, value))| FunctionExpr::JsonObjectAgg {
+                    key: Box::new(key),
+                    value: Box::new(value),
+                    allow_duplicate_keys: is_json,
+                },
+            ),
             extract(dialect),
             lower(dialect),
             upper(dialect),

nom-sql/src/common.rs

@@ -1749,7 +1749,22 @@ impl From<JsonValue> for DfValue {
 
 impl From<&JsonValue> for DfValue {
     fn from(value: &JsonValue) -> Self {
-        value.to_string().into()
+        match value {
+            JsonValue::String(s) => s.clone().into(),
+            JsonValue::Null => DfValue::None,
+            JsonValue::Bool(true) => DfValue::UnsignedInt(1),
+            JsonValue::Bool(false) => DfValue::UnsignedInt(0),
+            JsonValue::Number(number) if number.is_i64() => DfValue::Int(number.as_i64().unwrap()),
+            JsonValue::Number(number) if number.is_u64() => {
+                DfValue::UnsignedInt(number.as_u64().unwrap())
+            }
+            JsonValue::Number(number) if number.is_f64() => {
+                DfValue::Double(number.as_f64().unwrap())
+            }
+            JsonValue::Number(number) => DfValue::from(number.to_string()),
+            // TODO: add DfValue::JSON
+            JsonValue::Object(_) | JsonValue::Array(_) => DfValue::from(value.to_string()),
+        }
     }
 }
 

readyset-data/src/lib.rs

@@ -2,6 +2,7 @@ use std::collections::hash_map::DefaultHasher;
 use std::collections::HashMap;
 use std::hash::{Hash, Hasher};
 
+use dataflow_expression::eval::json;
 use readyset_data::dialect::SqlEngine;
 use readyset_data::{Collation, DfType, Dialect};
 use readyset_errors::{invariant, ReadySetResult};
@@ -22,7 +23,13 @@ pub enum Aggregation {
     /// Average the value of the `over` column. Maintains count and sum in HashMap
     Avg,
     /// Concatenates using the given separator between values.
-    GroupConcat { separator: String },
+    GroupConcat {
+        separator: String,
+    },
+
+    JsonObjectAgg {
+        allow_duplicate_keys: bool,
+    },
 }
 
 impl Aggregation {
@@ -71,6 +78,7 @@ impl Aggregation {
                 }
             }
             Aggregation::GroupConcat { .. } => DfType::Text(/* TODO */ Collation::default()),
+            Aggregation::JsonObjectAgg { .. } => DfType::Text(Collation::default()),
         };
 
         Ok(GroupedOperator::new(
@@ -150,6 +158,12 @@ impl AverageDataPair {
 /// Auxiliary State for an Aggregator node, which is owned by a Domain
 pub struct AggregatorState {
     count_sum_map: HashMap<GroupHash, AverageDataPair>,
+
+    // Store all `json_object_agg` keys and values in vecs and compute the json from them
+    // on-the-fly. This allows for easier handling of distinct (jsonb) behaviour,
+    // especially with deletions.
+    json_agg_keys: Vec<DfValue>,
+    json_agg_vals: Vec<DfValue>,
 }
 
 impl Aggregator {
@@ -222,10 +236,12 @@ impl GroupedOperation for Aggregator {
             }
         };
 
-        let count_sum_map = match auxiliary_node_state {
-            Some(AuxiliaryNodeState::Aggregation(ref mut aggregator_state)) => {
-                &mut aggregator_state.count_sum_map
-            }
+        let (count_sum_map, json_agg_keys, json_agg_vals) = match auxiliary_node_state {
+            Some(AuxiliaryNodeState::Aggregation(ref mut aggregator_state)) => (
+                &mut aggregator_state.count_sum_map,
+                &mut aggregator_state.json_agg_keys,
+                &mut aggregator_state.json_agg_vals,
+            ),
             Some(_) => internal!("Incorrect auxiliary state for Aggregation node"),
             None => internal!("Missing auxiliary state for Aggregation node"),
         };
@@ -240,6 +256,42 @@ impl GroupedOperation for Aggregator {
                 .apply_diff(diff)
         };
 
+        let mut apply_json_object_agg =
+            |_curr, diff: Self::Diff, allow_dups| -> ReadySetResult<DfValue> {
+                let (key, value) = diff
+                    .value
+                    .to_json()?
+                    .as_object()
+                    .ok_or_else(|| {
+                        internal_err!("json_object_agg: json_object value is not an object")
+                    })?
+                    .iter()
+                    .next()
+                    .ok_or_else(|| internal_err!("json_object_agg: json_object is empty"))
+                    .map(|(k, v)| (DfValue::from(k.as_str()), DfValue::from(v)))?;
+
+                if diff.positive {
+                    json_agg_keys.push(key);
+                    json_agg_vals.push(value.clone());
+                } else if let Some(pos) = json_agg_keys
+                    .iter()
+                    .zip(json_agg_vals.iter_mut())
+                    .position(|(k, v)| k == &key && v == &value)
+                {
+                    json_agg_keys.remove(pos);
+                    json_agg_vals.remove(pos);
+                } else {
+                    internal!("json_object_agg: diff removed a non-existant key-value pair")
+                }
+
+                // TODO: Indent the output
+                json::json_object_from_keys_and_values(
+                    &json_agg_keys.clone().into(),
+                    &json_agg_vals.clone().into(),
+                    allow_dups,
+                )
+            };
+
         let apply_diff =
             |curr: ReadySetResult<DfValue>, diff: Self::Diff| -> ReadySetResult<DfValue> {
                 if diff.value.is_none() {
@@ -253,6 +305,9 @@ impl GroupedOperation for Aggregator {
                     Aggregation::GroupConcat { separator: _ } => internal!(
                         "GroupConcats are separate from the other aggregations in the dataflow."
                     ),
+                    Aggregation::JsonObjectAgg {
+                        allow_duplicate_keys,
+                    } => apply_json_object_agg(curr?, diff, allow_duplicate_keys),
                 }
             };
 
@@ -270,6 +325,15 @@ impl GroupedOperation for Aggregator {
                 Aggregation::GroupConcat { separator: ref s } => {
                     format!("||({})", s)
                 }
+                Aggregation::JsonObjectAgg {
+                    allow_duplicate_keys,
+                } => {
+                    if allow_duplicate_keys {
+                        "JsonObjectAgg".to_owned()
+                    } else {
+                        "JsonbObjectAgg".to_owned()
+                    }
+                }
             };
         }
 
@@ -278,6 +342,15 @@ impl GroupedOperation for Aggregator {
             Aggregation::Sum => format!("𝛴({})", self.over),
             Aggregation::Avg => format!("Avg({})", self.over),
             Aggregation::GroupConcat { separator: ref s } => format!("||({}, {})", s, self.over),
+            Aggregation::JsonObjectAgg {
+                allow_duplicate_keys,
+            } => {
+                if allow_duplicate_keys {
+                    format!("JsonObjectAgg({})", self.over)
+                } else {
+                    format!("JsonbObjectAgg({})", self.over)
+                }
+            }
         };
         let group_cols = self
             .group