diff --git a/udfs/migration/snowflake/README.md b/udfs/migration/snowflake/README.md index b95025a78..71a9bf725 100644 --- a/udfs/migration/snowflake/README.md +++ b/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 + + diff --git a/udfs/migration/snowflake/object_agg.sqlx b/udfs/migration/snowflake/object_agg.sqlx new file mode 100644 index 000000000..1f673ca2f --- /dev/null +++ b/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; + } +"""; diff --git a/udfs/migration/snowflake/test_cases.js b/udfs/migration/snowflake/test_cases.js index 3500ce565..d23ad24a1 100644 --- a/udfs/migration/snowflake/test_cases.js +++ b/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}'` + } +); +