PostgreSQL SRF (set-returning functions) 函数where, index实现

背景

PostgreSQL 函数支持返回多行，定义returns setof type即可。PG内部也内置了蛮多这样的函数，例如

select proname from pg_proc where proretset;  
  
generate_series  
unnest  
 json_array_elements  
 json_array_elements_text  
 json_object_keys  
 json_each  
 json_each_text  
 json_to_recordset  
 unnest  
 ts_stat  
 ts_stat  
 ts_token_type  
 ts_token_type  
 ts_parse  
 ts_parse  
 jsonb_array_elements  
 jsonb_array_elements_text  
 jsonb_object_keys  
 jsonb_each  
 jsonb_each_text  
 jsonb_populate_recordset  
 jsonb_to_recordset  
............  

自定义SRF方法如下：

https://www.postgresql.org/docs/11/static/plpgsql-control-structures.html

对于返回多行的函数，如果在where条件中执行过滤，会返回如下错误。

ERROR:  0A000: set-returning functions are not allowed in WHERE  

如果要创建基于SRF函数的表达式索引，会报如下错误：

ERROR:  0A000: set-returning functions are not allowed in index expressions  

而实际上，可能在场景中有查询srf = ?这样的需求。

如何达到这样的需求呢？

1、 可以将SRF再转换为ARRAY，然后构建ARRAY的GIN索引。

例如这是个例子  
  
postgres=# create or replace function gen_rand(int,int) returns int[] as $$  
select array(select (random()*$1)::int from generate_series(1,$2));  
$$ language sql strict;  
CREATE FUNCTION  
  
postgres=# select gen_rand(100,10);  
           gen_rand              
-------------------------------  
 {17,5,33,70,54,10,89,96,5,94}  
(1 row)  

例子

1、建表

create table a (id int, js jsonb);  

2、写入测试数据

postgres=# insert into a select id, jsonb_build_array((random()*1000)::int, (random()*1000)::int,(random()*1000)::int,(random()*1000)::int,(random()*1000)::int,(random()*1000)::int,(random()*1000)::int,(random()*1000)::int,(random()*1000)::int,(random()*1000)::int) from generate_series(1,10000) t(id);  
INSERT 0 10000  

3、数据样本

postgres=# select * from  a limit 2;  
 id |                        js                           
----+---------------------------------------------------  
  1 | [774, 509, 740, 813, 925, 67, 539, 730, 598, 465]  
  2 | [232, 153, 119, 895, 962, 501, 998, 609, 10, 410]  
(2 rows)  

4、某个JSON类型的SRF函数调用如下

postgres=# select jsonb_array_elements(js) from  a where id=1;  
 jsonb_array_elements   
----------------------  
 774  
 509  
 740  
 813  
 925  
 67  
 539  
 730  
 598  
 465  
(10 rows)  

在WHERE中过滤，创建表达式索引，均报错

postgres=# \set VERBOSITY verbose  
postgres=# select * from a where jsonb_array_elements(js) =1;  
ERROR:  0A000: set-returning functions are not allowed in WHERE  
LINE 1: select * from a where jsonb_array_elements(js) =1;  
                              ^  
LOCATION:  check_srf_call_placement, parse_func.c:2258  
  
  
  
postgres=# create index idx_a_1 on a (jsonb_array_elements(js));  
ERROR:  0A000: set-returning functions are not allowed in index expressions  
LINE 1: create index idx_a_1 on a (jsonb_array_elements(js));  
                                   ^  
LOCATION:  check_srf_call_placement, parse_func.c:2251  

5、创建一个UDF，将SRF结果转换为ARRAY

postgres=# create or replace function srf_to_arr(jsonb) returns text[] as $$  
  select array(select jsonb_array_elements($1)::text);  
$$ language sql strict immutable;  
CREATE FUNCTION  

postgres=# select id,js,srf_to_arr(js) from a limit 1;  
 id |                        js                         |                srf_to_arr                  
----+---------------------------------------------------+------------------------------------------  
  1 | [774, 509, 740, 813, 925, 67, 539, 730, 598, 465] | {774,509,740,813,925,67,539,730,598,465}  
(1 row)  

6、WHERE srf = ? 条件变成如下

postgres=# select id,js,srf_to_arr(js) from a where srf_to_arr(js) @> array['1'];  
  id  |                        js                        |               srf_to_arr                  
------+--------------------------------------------------+-----------------------------------------  
   18 | [96, 4, 447, 177, 53, 550, 768, 27, 1, 280]      | {96,4,447,177,53,550,768,27,1,280}  
   67 | [402, 1, 841, 834, 462, 769, 247, 568, 114, 690] | {402,1,841,834,462,769,247,568,114,690}  
  102 | [555, 599, 389, 719, 1, 916, 910, 637, 566, 36]  | {555,599,389,719,1,916,910,637,566,36}  
  162 | [687, 1, 628, 851, 20, 522, 883, 814, 874, 938]  | {687,1,628,851,20,522,883,814,874,938}  
.....  

7、创建表达式索引

postgres=# create index idx_a_1 on a using gin(srf_to_arr(js));  
CREATE INDEX  

用上了这个GIN倒排索引

postgres=# explain select id,js,srf_to_arr(js) from a where srf_to_arr(js) @> array['1'];  
                              QUERY PLAN                                 
-----------------------------------------------------------------------  
 Bitmap Heap Scan on a  (cost=3.69..76.70 rows=50 width=160)  
   Recheck Cond: (srf_to_arr(js) @> '{1}'::text[])  
   ->  Bitmap Index Scan on idx_a_1  (cost=0.00..3.67 rows=50 width=0)  
         Index Cond: (srf_to_arr(js) @> '{1}'::text[])  
(4 rows)  

参考

https://www.postgresql.org/docs/11/static/plpgsql-control-structures.html

digoal’s 大量PostgreSQL文章入口