snowflake-cloud-data-platform - 在雪花中创建 UDAF(不是 UDTF)

Question

Java UDFs return a scalar result. Java UDTFs are not currently supported. reference
也就是说，我创建了一个 Java UDF，如下所示

CREATE OR replace function MAP_COUNT(colValue String)
returns OBJECT 
language java
handler='Frequency.calculate'
target_path='@~/Frequency.jar'
as
$$
    import java.util.HashMap;
    import java.util.Map;
    import java.util.Optional;
    class Frequency {
        Map<String, Integer> frequencies = new HashMap<>();

        public Map<String, Integer> calculate(String colValue) {
            frequencies.putIfAbsent(colValue, 0);
            frequencies.computeIfPresent(colValue, (key, value) -> value + 1);
            return frequencies;
        }
    }
$$;

使用 MAP_COUNT查询中的 UDF 如下

with temp_1 as
(
    SELECT 'John' AS my_col, 27 as age
    UNION ALL
    SELECT 'John' AS my_col, 28 as age
    UNION ALL
    SELECT 'doe' AS my_col, 27 as age
    UNION ALL
    SELECT 'doe' AS my_col, 28 as age
)
select  MAP_COUNT(a.my_col) from temp_1 a;

我得到如下结果

|MAP_COUNT(A.MY_COL)            |
|-------------------------------|
|{  "John": "1" }               |
|{  "John": "2" }               |
|{ "John": "2",  "doe": "1" }   |
|{  "John": "2",  "doe": "2"}   |

我期望从我的 UDF 得到的结果如下

|MAP_COUNT(A.MY_COL)            |
|-------------------------------|
|{  "John": "2",  "doe": "2"}   |

在雪花中可能吗？
如果我有如下查询怎么办？

with temp_1 as
(
    SELECT 'John' AS my_col, 27 as age
    UNION ALL
    SELECT 'John' AS my_col, 28 as age
    UNION ALL
    SELECT 'doe' AS my_col, 27 as age
    UNION ALL
    SELECT 'doe' AS my_col, 28 as age
)
select  MAP_COUNT(a.my_col) as names, MAP_COUNT(a.age) as ages  from temp_1 a;

我期望从我的 UDF 得到的结果如下

|names                          ||AGES                           |
|-------------------------------||-------------------------------|
|{  "John": "2",  "doe": "2"}   ||{  "27": "2",  "28": "2"}      |

有一些方法可以通过简单地重组查询来实现这一点，但我想知道是否可以使用 MAP_COUNT 来做到这一点。功能类似于 OBJECT_AGG select 子句中的函数。

Answer 1

当您运行使用 UDF 的查询时，并非所有行都一定会转到 UDF 的同一实例。例如，假设您正在从表中进行选择，并且您执行以下操作:

SELECT MyUdf(x) FROM T

这里 T可能有多个 micro-partitions ，它的执行方式实际上类似于:

SELECT MyUdf(x) FROM T_part1 UNION ALL
SELECT MyUdf(x) FROM T_part2 UNION ALL
SELECT MyUdf(x) FROM T_part3 UNION ALL
SELECT MyUdf(x) FROM T_part4

这里有四个单独的 MyUdf 实例。 ，并且每个人只看到来自 T 的行的一个子集作为一个整体。
回到您的示例，您正在尝试模拟用户定义的聚合函数，其中 UDF 的特定实例查看每一行。保证这一点的方法是提前聚合，例如:

CREATE OR replace function MAP_COUNT(colValue array)
returns OBJECT 
language java
handler='Frequency.calculate'
target_path='@~/Frequency.jar'
as
$$
    import java.util.HashMap;
    import java.util.Map;
    import java.util.Optional;
    class Frequency {
        public Map<String, Integer> calculate(String[] colValues) {
            Map<String, Integer> frequencies = new HashMap<>();
            for (String colValue : colValues) {
                frequencies.putIfAbsent(colValue, 0);
                frequencies.computeIfPresent(colValue, (key, value) -> value + 1);
            }
            return frequencies;
        }
    }
$$;

(请注意，我将 UDF 和方法签名更改为分别使用 array 和 String[] 。)现在在查询中使用它:

with temp_1 as
(
    SELECT 'John' AS my_col, 27 as age
    UNION ALL
    SELECT 'John' AS my_col, 28 as age
    UNION ALL
    SELECT 'doe' AS my_col, 27 as age
    UNION ALL
    SELECT 'doe' AS my_col, 28 as age
)
select
  MAP_COUNT(ARRAY_AGG(a.my_col)) as names,
  MAP_COUNT(ARRAY_AGG(a.age)) as ages
from temp_1 a;

这给了我:

names                        ages
{ "John": "2", "doe": "2" }  { "27": "2", "28": "2" }

这里仍然存在两个问题，特别是:

这不能很好地扩展。如果任一数组的大小超过 16MB(最大值大小)，则查询将失败。

语法很笨拙。理想情况下，您只需像使用任何其他聚合函数一样使用 UDF，而不必将输入包装在 ARRAY_AGG 中。 .

好消息是，一旦 Java UDAF 在 future 某个时候可用，这两个问题都将得到解决。