提高PostgreSQL自定义聚合函数的性能

Question

我有一个自定义的聚合求和函数，可以接受布尔类型的数据：

create or replace function badd (bigint, boolean)
  returns bigint as
$body$
select $1 + case when $2 then 1 else 0 end;
$body$ language sql;

create aggregate sum(boolean) (
  sfunc=badd,
  stype=int8,
  initcond='0'
);

这个聚合应该计算有多少行是TRUE。例如，以下内容应该返回2（确实如此）：

with t (x) as 
   (values 
      (true::boolean), 
      (false::boolean), 
      (true::boolean),
      (null::boolean)      
      )
select sum(x) from t;

然而，它的性能相当糟糕，比使用整数转换慢了5.5倍：

with t as (select (gs > 0.5) as test_vector from generate_series(1,1000000,1) gs)
select sum(test_vector) from t; -- 52012ms

with t as (select (gs > 0.5) as test_vector from generate_series(1,1000000,1) gs)
select sum(test_vector::int) from t; -- 9484ms

想要改进这个聚合函数的唯一方法是编写一些新的 C 函数，例如在 src/backend/utils/adt/numeric.c 中的 int2_sum 函数的替代品吗？

- Tomas Greif

---

那不是一个int2-sum。它是一个bool-sum，更像是bit-1。而且它是错误的。应该是：$1 + $2 :: int，或者$1 + case $2 when true then 1 else 0 end，或类似的东西。 - Denis de Bernardy

我仍然不明白为什么需要自定义聚合函数，select sum(x::int) from t; 和 select sum(1) from t where x; 都是正确的且更快。 - foibs

@Denis 我更新了我的问题，我的聚合功能工作得很好，也达到了我想要的结果，但速度很慢。 - Tomas Greif

Answer 1

你的测试用例有误导性，你只计算了TRUE。你应该同时计算TRUE 和 FALSE - 或者甚至是NULL（如果适用）。 如@foibs已经解释的那样，一个人不会为此使用自定义聚合函数。内置的C函数更快且能够胜任工作。使用以下代码（还演示了更简单、更明智的测试）：

SELECT count(NULLIF(g%2 = 1, FALSE)) AS ct
FROM  generate_series(1,100000,1) g;

我会尽力为您翻译中文。以下是需要翻译的内容：

这是如何工作的？
在同一SELECT SQL查询中从SUM()计算百分比

在dba.SE上的相关答案下有几种快速简单的方法（以及基准测试）：
对于绝对性能，SUM还是COUNT更快？

或者更快的方式是，在WHERE子句中测试TRUE，如果可能的话：

SELECT count(*) AS ct
FROM   generate_series(1,100000,1) g;
WHERE  g%2 = 1             -- excludes FALSE and NULL !

---

如果你不得不为某些原因编写自定义聚合函数，那么这种形式会更加优越：

CREATE OR REPLACE FUNCTION test_sum_int8 (int8, boolean)
  RETURNS bigint as
'SELECT CASE WHEN $2 THEN $1 + 1 ELSE $1 END' LANGUAGE sql;

只有在必要的情况下才执行加法。您原来的代码会在`FALSE`情况下添加`0`。

更好的方法是使用`plpgsql`函数。它可以节省每次调用的一点开销，因为它工作方式类似于准备好的语句（查询不会重新计划）。对于被调用许多次的小型聚合函数，这会产生影响。

CREATE OR REPLACE FUNCTION test_sum_plpgsql (int8, boolean)
  RETURNS bigint AS
$func$
BEGIN
RETURN CASE WHEN $2 THEN $1 + 1 ELSE $1 END;
END
$func$ LANGUAGE plpgsql;

CREATE AGGREGATE test_sum_plpgsql(boolean) (
  sfunc = test_sum_plpgsql
 ,stype = int8
 ,initcond = '0'
);

比你之前使用的方法快，但比使用标准的 count() 函数要慢得多。而且比其他任何 C 函数都要慢。

->SQLfiddle

- Erwin Brandstetter

---

哦，你说得对，我的测试用例应该是：with t as (select (random() > 0.5) as test_vector from generate_series(1,1000000,1) gs) select sum(test_vector) from t; 我会尝试编写自己的 C 函数。 - Tomas Greif

Answer 2

我创建了一个用于布尔值的自定义 C 函数和聚合函数：

C 函数：

#include "postgres.h"
#include <fmgr.h>

#ifdef PG_MODULE_MAGIC
PG_MODULE_MAGIC;
#endif

int
bool_sum(int arg, bool tmp)
{
  if (tmp)
    {
    arg++;
    } 
    return arg;
}

转换和聚合函数：

-- transition function
create or replace function bool_sum(bigint, boolean)
  returns bigint
   AS '/usr/lib/postgresql/9.1/lib/bool_agg', 'bool_sum'
  language C strict
  cost 1;
alter function bool_sum(bigint, boolean) owner to postgres;

-- aggregate
create aggregate sum(boolean) (
  sfunc=bool_sum,
  stype=int8,
  initcond='0'
);
alter aggregate sum(boolean) owner to postgres;

性能测试：

-- Performance test - 10m rows

create table tmp_test as (select (case when random() <.3 then null when random() < .6 then true else false end) as test_vector from generate_series(1,10000000,1) gs);

-- Casting to integer
select sum(test_vector::int) from tmp_test;

-- Boolean sum
select sum(test_vector) from tmp_test;

现在，sum(boolean) 的速度与 sum(boolean::int) 相同。

更新：

事实证明，我可以直接调用现有的 C 转换函数，即使使用布尔数据类型也可以。它会在某种程度上神奇地转换为 0/1。因此，我当前的布尔求和和平均值解决方案如下：

create or replace function bool_sum(bigint, boolean)
  returns bigint as
'int2_sum'
  language internal immutable
  cost 1;


create aggregate sum(boolean) (
  sfunc=bool_sum,
  stype=int8
);

-- Average for boolean values (percentage of rows with TRUE)
create or replace function bool_avg_accum(bigint[], boolean)
  returns bigint[] as
'int2_avg_accum'
  language internal immutable strict
  cost 1;

create aggregate avg(boolean) (
  sfunc=bool_avg_accum,
  stype=int8[],
  finalfunc=int8_avg,
  initcond='{0,0}'
);

- Tomas Greif

---

+1 感谢您提供有趣的解决方案。即使对于所呈现的情况并非必需，它也可能对相关问题有用。 - Erwin Brandstetter

@ErwinBrandstetter 我已经更新了我的答案 - 不需要自定义 C 函数，仍然可以像内置聚合函数一样快速。 - Tomas Greif

Answer 3

我不认为真正的问题在这里。首先，将 sum 用作自定义聚合名称是错误的。当您使用将 test_vector 强制转换为 int 的方式调用 sum 时，嵌入式 postgres sum 被使用，而不是您的自定义函数，这就是为什么它更快的原因。C 函数总是会更快，但我不确定在这种情况下是否需要。

您可以轻松地删除 badd 函数和自定义的 sum，使用带有 where 子句的内置 sum。

with t as (select 1 as test_vector from generate_series(1,1000000,1) gs where gs > 0.5)
select sum(test_vector) from t;

编辑：

总的来说，优化自定义聚合的最佳方式是在不需要时将其删除。次佳的方式是编写一个 C 函数。

- foibs