Add emulation function for Snowflake's OBJECT_AGG() (GoogleCloudPlatform#539)

Author: Hackerpilot

udfs/migration/snowflake/README.md

@@ -56,3 +56,32 @@ Row|j
 ---|-
 1|{"frame":2,"xray":1}	
 
+### [object_agg(key STRING, value JSON)](object_agg.sqlx)
+Emulates the `OBJECT_AGG` function in Snowflake. [Snowflake docs](https://docs.snowflake.com/en/sql-reference/functions/object_agg)
+```sql
+SELECT object_agg(k, v)
+FROM
+  (
+    SELECT 'a' AS k, json '1' AS v
+    UNION ALL
+    SELECT 'b' as k, json '2' as v);
+```
+
+Row|f0_
+---|---
+1|{"b":2,"a":1}
+
+```sql
+SELECT object_agg(k, v)
+FROM
+  (
+    SELECT 'a' AS k, json '1' AS v
+    UNION ALL
+    SELECT 'a' AS k, json '3' AS b
+    UNION ALL
+    SELECT 'b' as k, json '2' as v);
+```
+
+Error: Duplicate field key 'a' at object_agg(STRING, JSON) line 7, columns 6-7
+
+

udfs/migration/snowflake/object_agg.sqlx

@@ -0,0 +1,61 @@
+config { hasOutput: true }
+/*
+ * Copyright 2026 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+CREATE OR REPLACE AGGREGATE FUNCTION ${self()}(
+  key STRING,
+  value JSON,
+  allow_duplicate_pairs BOOL NOT AGGREGATE)
+RETURNS JSON
+LANGUAGE js
+OPTIONS (
+  description = "Emulation of Snowflake's OBJECT_AGG function."
+) AS r"""
+  export function initialState() {
+    return {obj: {}, allow_duplicate_pairs: false};
+  }
+
+  export function aggregate(state, key, value, allow_duplicate_pairs) {
+    state.allow_duplicate_pairs = allow_duplicate_pairs;
+    if (key === null || value === null) {
+      return;
+    }
+    if (Object.hasOwn(state.obj, key) && (!allow_duplicate_pairs || state.obj[key] !== value)) {
+      throw new Error("Duplicate field key '" + key + "'");
+    }
+    state.obj[key] = value;
+  }
+
+  export function merge(state, partialState) {
+    const stateKeys = Object.keys(state.obj);
+    const partialStateKeys = Object.keys(partialState.obj);
+    const [smallerKeys, largerObj, smallerObj] = stateKeys.length < partialStateKeys.length
+      ? [stateKeys, partialState.obj, state.obj]
+      : [partialStateKeys, state.obj, partialState.obj];
+    for (let i = 0; i < smallerKeys.length; i++) {
+      const key = smallerKeys[i];
+      if (Object.hasOwn(largerObj, key) &&
+          (!state.allow_duplicate_pairs || largerObj[key] === smallerObj[key])) {
+        throw new Error("Duplicate field key '" + key + "'");
+      }
+    }
+    Object.assign(state.obj, partialState.obj);
+  }
+
+  export function finalize(state) {
+    return state.obj;
+  }
+""";

udfs/migration/snowflake/test_cases.js

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-const {generate_udf_test} = unit_test_utils;
+const {generate_udf_test, generate_udaf_test} = unit_test_utils;
 
 generate_udf_test("factorial", [
     {
@@ -108,3 +108,43 @@ generate_udf_test("json_ilike", [
     }
 ]);
 
+generate_udaf_test("object_agg",
+    {
+        input_columns: [`key`, `value`, `false NOT AGGREGATE`],
+        input_rows: `
+            select 'a' as key, JSON '1' as value
+            union all
+            select 'b' as key, JSON '2' as value
+        `,
+        expected_output: `JSON '{"a": 1, "b": 2}'`
+    }
+);
+
+generate_udaf_test("object_agg",
+    {
+        input_columns: [`key`, `value`, `true NOT AGGREGATE`],
+        input_rows: `
+            select 'a' as key, JSON '1' as value
+            union all
+            select null as key, JSON '3' as value
+            union all
+            select 'b' as key, JSON '2' as value
+            union all
+            select 'b' as key, null as value
+        `,
+        expected_output: `JSON '{"a": 1, "b": 2}'`
+    }
+);
+
+generate_udaf_test("object_agg",
+    {
+        input_columns: [`key`, `value`, `true NOT AGGREGATE`],
+        input_rows: `
+            select 'a' as key, JSON '1' as value
+            union all
+            select 'a' as key, JSON '1' as value
+        `,
+        expected_output: `JSON '{"a": 1}'`
+    }
+);
+