From 116c9dd5b3fbb6626d81588c124cf8fdcb4185ce Mon Sep 17 00:00:00 2001
From: "Paul A. Jungwirth" <pj@illuminatedcomputing.com>
Date: Fri, 10 Dec 2021 16:04:57 -0800
Subject: [PATCH v1] Add range_agg with multirange inputs

---
 doc/src/sgml/func.sgml                        |  30 ++++++
 src/backend/utils/adt/multirangetypes.c       |  60 +++++++++++
 src/include/catalog/pg_aggregate.dat          |   3 +
 src/include/catalog/pg_proc.dat               |  11 ++
 src/test/regress/expected/multirangetypes.out | 100 ++++++++++++++++++
 src/test/regress/expected/opr_sanity.out      |   1 +
 src/test/regress/sql/multirangetypes.sql      |  21 ++++
 src/test/regress/sql/opr_sanity.sql           |   1 +
 8 files changed, 227 insertions(+)

diff --git a/doc/src/sgml/func.sgml b/doc/src/sgml/func.sgml
index 0a725a6711..697bf1924e 100644
--- a/doc/src/sgml/func.sgml
+++ b/doc/src/sgml/func.sgml
@@ -19951,6 +19951,21 @@ SELECT NULLIF(value, '(none)') ...
        <entry>No</entry>
       </row>
 
+      <row>
+       <entry role="func_table_entry"><para role="func_signature">
+        <indexterm>
+         <primary>range_agg</primary>
+        </indexterm>
+        <function>range_agg</function> ( <parameter>value</parameter>
+         <type>anymultirange</type> )
+        <returnvalue>anymultirange</returnvalue>
+       </para>
+       <para>
+        Computes the union of the non-null input values.
+       </para></entry>
+       <entry>No</entry>
+      </row>
+
       <row>
        <entry role="func_table_entry"><para role="func_signature">
         <indexterm>
@@ -19966,6 +19981,21 @@ SELECT NULLIF(value, '(none)') ...
        <entry>No</entry>
       </row>
 
+      <row>
+       <entry role="func_table_entry"><para role="func_signature">
+        <indexterm>
+         <primary>range_intersect_agg</primary>
+        </indexterm>
+        <function>range_intersect_agg</function> ( <parameter>value</parameter>
+         <type>anymultirange</type> )
+        <returnvalue>anymultirange</returnvalue>
+       </para>
+       <para>
+        Computes the intersection of the non-null input values.
+       </para></entry>
+       <entry>No</entry>
+      </row>
+
       <row>
        <entry role="func_table_entry"><para role="func_signature">
         <indexterm>
diff --git a/src/backend/utils/adt/multirangetypes.c b/src/backend/utils/adt/multirangetypes.c
index 7773215564..beaa843eab 100644
--- a/src/backend/utils/adt/multirangetypes.c
+++ b/src/backend/utils/adt/multirangetypes.c
@@ -1394,6 +1394,66 @@ range_agg_finalfn(PG_FUNCTION_ARGS)
 	PG_RETURN_MULTIRANGE_P(make_multirange(mltrngtypoid, typcache->rngtype, range_count, ranges));
 }
 
+/*
+ * multirange_agg_transfn: combine adjacent/overlapping multiranges.
+ *
+ * All we do here is gather the input multiranges' ranges into an array
+ * so that the finalfn can sort and combine them.
+ */
+Datum
+multirange_agg_transfn(PG_FUNCTION_ARGS)
+{
+	MemoryContext aggContext;
+	Oid			mltrngtypoid;
+	TypeCacheEntry *typcache;
+	TypeCacheEntry *rngtypcache;
+	ArrayBuildState *state;
+	MultirangeType *current;
+	int32		range_count;
+	RangeType **ranges;
+	int32		i;
+
+	if (!AggCheckCallContext(fcinfo, &aggContext))
+		elog(ERROR, "multirange_agg_transfn called in non-aggregate context");
+
+	mltrngtypoid = get_fn_expr_argtype(fcinfo->flinfo, 1);
+	if (!type_is_multirange(mltrngtypoid))
+		ereport(ERROR,
+				(errcode(ERRCODE_DATATYPE_MISMATCH),
+				 errmsg("range_agg must be called with a multirange")));
+
+	typcache = multirange_get_typcache(fcinfo, mltrngtypoid);
+	rngtypcache = typcache->rngtype;
+
+	if (PG_ARGISNULL(0))
+		state = initArrayResult(rngtypcache->type_id, aggContext, false);
+	else
+		state = (ArrayBuildState *) PG_GETARG_POINTER(0);
+
+	/* skip NULLs */
+	if (!PG_ARGISNULL(1))
+	{
+		current = PG_GETARG_MULTIRANGE_P(1);
+		multirange_deserialize(rngtypcache, current, &range_count, &ranges);
+		if (range_count == 0)
+		{
+			/*
+			 * Add an empty range so we get an empty result (not a null result).
+			 */
+			accumArrayResult(state,
+							 RangeTypePGetDatum(make_empty_range(rngtypcache)),
+							 false, rngtypcache->type_id, aggContext);
+		}
+		else
+		{
+			for (i = 0; i < range_count; i++)
+				accumArrayResult(state, RangeTypePGetDatum(ranges[i]), false, rngtypcache->type_id, aggContext);
+		}
+	}
+
+	PG_RETURN_POINTER(state);
+}
+
 Datum
 multirange_intersect_agg_transfn(PG_FUNCTION_ARGS)
 {
diff --git a/src/include/catalog/pg_aggregate.dat b/src/include/catalog/pg_aggregate.dat
index fc6d3bfd94..5a06583798 100644
--- a/src/include/catalog/pg_aggregate.dat
+++ b/src/include/catalog/pg_aggregate.dat
@@ -557,6 +557,9 @@
 { aggfnoid => 'range_agg(anyrange)', aggtransfn => 'range_agg_transfn',
   aggfinalfn => 'range_agg_finalfn', aggfinalextra => 't',
   aggtranstype => 'internal' },
+{ aggfnoid => 'range_agg(anymultirange)', aggtransfn => 'multirange_agg_transfn',
+  aggfinalfn => 'multirange_agg_finalfn', aggfinalextra => 't',
+  aggtranstype => 'internal' },
 
 # json
 { aggfnoid => 'json_agg', aggtransfn => 'json_agg_transfn',
diff --git a/src/include/catalog/pg_proc.dat b/src/include/catalog/pg_proc.dat
index 79d787cd26..01c5c52729 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -10619,6 +10619,17 @@
   proname => 'range_agg', prokind => 'a', proisstrict => 'f',
   prorettype => 'anymultirange', proargtypes => 'anyrange',
   prosrc => 'aggregate_dummy' },
+{ oid => '8000', descr => 'aggregate transition function',
+  proname => 'multirange_agg_transfn', proisstrict => 'f', prorettype => 'internal',
+  proargtypes => 'internal anymultirange', prosrc => 'multirange_agg_transfn' },
+{ oid => '8001', descr => 'combine aggregate input into a multirange',
+  proname => 'range_agg', prokind => 'a', proisstrict => 'f',
+  prorettype => 'anymultirange', proargtypes => 'anymultirange',
+  prosrc => 'aggregate_dummy' },
+{ oid => '8002', descr => 'aggregate final function',
+  proname => 'multirange_agg_finalfn', proisstrict => 'f',
+  prorettype => 'anymultirange', proargtypes => 'internal anymultirange',
+  prosrc => 'range_agg_finalfn' },
 { oid => '4388', descr => 'range aggregate by intersecting',
   proname => 'multirange_intersect_agg_transfn', prorettype => 'anymultirange',
   proargtypes => 'anymultirange anymultirange',
diff --git a/src/test/regress/expected/multirangetypes.out b/src/test/regress/expected/multirangetypes.out
index f5f089741c..7c91e3c5f2 100644
--- a/src/test/regress/expected/multirangetypes.out
+++ b/src/test/regress/expected/multirangetypes.out
@@ -2771,6 +2771,70 @@ FROM    (VALUES
  {[a,f],[g,j)}
 (1 row)
 
+-- range_agg with multirange inputs
+select range_agg(nmr) from nummultirange_test;
+ range_agg 
+-----------
+ {(,)}
+(1 row)
+
+select range_agg(nmr) from nummultirange_test where false;
+ range_agg 
+-----------
+ 
+(1 row)
+
+select range_agg(null::nummultirange) from nummultirange_test;
+ range_agg 
+-----------
+ 
+(1 row)
+
+select range_agg(nmr) from (values ('{}'::nummultirange)) t(nmr);
+ range_agg 
+-----------
+ {}
+(1 row)
+
+select range_agg(nmr) from (values ('{}'::nummultirange), ('{}'::nummultirange)) t(nmr);
+ range_agg 
+-----------
+ {}
+(1 row)
+
+select range_agg(nmr) from (values ('{[1,2]}'::nummultirange)) t(nmr);
+ range_agg 
+-----------
+ {[1,2]}
+(1 row)
+
+select range_agg(nmr) from (values ('{[1,2], [5,6]}'::nummultirange)) t(nmr);
+   range_agg   
+---------------
+ {[1,2],[5,6]}
+(1 row)
+
+select range_agg(nmr) from (values ('{[1,2], [2,3]}'::nummultirange)) t(nmr);
+ range_agg 
+-----------
+ {[1,3]}
+(1 row)
+
+select range_agg(nmr) from (values ('{[1,2]}'::nummultirange), ('{[5,6]}'::nummultirange)) t(nmr);
+   range_agg   
+---------------
+ {[1,2],[5,6]}
+(1 row)
+
+select range_agg(nmr) from (values ('{[1,2]}'::nummultirange), ('{[2,3]}'::nummultirange)) t(nmr);
+ range_agg 
+-----------
+ {[1,3]}
+(1 row)
+
+--
+-- range_intersect_agg function
+--
 select range_intersect_agg(nmr) from nummultirange_test;
  range_intersect_agg 
 ---------------------
@@ -2783,13 +2847,49 @@ select range_intersect_agg(nmr) from nummultirange_test where false;
  
 (1 row)
 
+select range_intersect_agg(null::nummultirange) from nummultirange_test;
+ range_intersect_agg 
+---------------------
+ 
+(1 row)
+
+select range_intersect_agg(nmr) from (values ('{[1,3]}'::nummultirange), ('{[6,12]}'::nummultirange)) t(nmr);
+ range_intersect_agg 
+---------------------
+ {}
+(1 row)
+
+select range_intersect_agg(nmr) from (values ('{[1,6]}'::nummultirange), ('{[3,12]}'::nummultirange)) t(nmr);
+ range_intersect_agg 
+---------------------
+ {[3,6]}
+(1 row)
+
+select range_intersect_agg(nmr) from (values ('{[1,6], [10,12]}'::nummultirange), ('{[4,14]}'::nummultirange)) t(nmr);
+ range_intersect_agg 
+---------------------
+ {[4,6],[10,12]}
+(1 row)
+
 -- test with just one input:
+select range_intersect_agg(nmr) from (values ('{}'::nummultirange)) t(nmr);
+ range_intersect_agg 
+---------------------
+ {}
+(1 row)
+
 select range_intersect_agg(nmr) from (values ('{[1,2]}'::nummultirange)) t(nmr);
  range_intersect_agg 
 ---------------------
  {[1,2]}
 (1 row)
 
+select range_intersect_agg(nmr) from (values ('{[1,6], [10,12]}'::nummultirange)) t(nmr);
+ range_intersect_agg 
+---------------------
+ {[1,6],[10,12]}
+(1 row)
+
 select range_intersect_agg(nmr) from nummultirange_test where nmr @> 4.0;
  range_intersect_agg 
 ---------------------
diff --git a/src/test/regress/expected/opr_sanity.out b/src/test/regress/expected/opr_sanity.out
index 562b586d8e..0561d2f705 100644
--- a/src/test/regress/expected/opr_sanity.out
+++ b/src/test/regress/expected/opr_sanity.out
@@ -193,6 +193,7 @@ WHERE p1.oid != p2.oid AND
     p2.prosrc NOT LIKE E'range\\_constructor_' AND
     p1.prosrc NOT LIKE E'multirange\\_constructor_' AND
     p2.prosrc NOT LIKE E'multirange\\_constructor_' AND
+    p2.prosrc NOT LIKE E'range\\_agg\\_finalfn' AND
     (p1.proargtypes[1] < p2.proargtypes[1])
 ORDER BY 1, 2;
          proargtypes         |       proargtypes        
diff --git a/src/test/regress/sql/multirangetypes.sql b/src/test/regress/sql/multirangetypes.sql
index ad50afc0f5..807a3a9b98 100644
--- a/src/test/regress/sql/multirangetypes.sql
+++ b/src/test/regress/sql/multirangetypes.sql
@@ -570,10 +570,31 @@ FROM    (VALUES
           ('[h,j)'::textrange)
         ) t(r);
 
+-- range_agg with multirange inputs
+select range_agg(nmr) from nummultirange_test;
+select range_agg(nmr) from nummultirange_test where false;
+select range_agg(null::nummultirange) from nummultirange_test;
+select range_agg(nmr) from (values ('{}'::nummultirange)) t(nmr);
+select range_agg(nmr) from (values ('{}'::nummultirange), ('{}'::nummultirange)) t(nmr);
+select range_agg(nmr) from (values ('{[1,2]}'::nummultirange)) t(nmr);
+select range_agg(nmr) from (values ('{[1,2], [5,6]}'::nummultirange)) t(nmr);
+select range_agg(nmr) from (values ('{[1,2], [2,3]}'::nummultirange)) t(nmr);
+select range_agg(nmr) from (values ('{[1,2]}'::nummultirange), ('{[5,6]}'::nummultirange)) t(nmr);
+select range_agg(nmr) from (values ('{[1,2]}'::nummultirange), ('{[2,3]}'::nummultirange)) t(nmr);
+
+--
+-- range_intersect_agg function
+--
 select range_intersect_agg(nmr) from nummultirange_test;
 select range_intersect_agg(nmr) from nummultirange_test where false;
+select range_intersect_agg(null::nummultirange) from nummultirange_test;
+select range_intersect_agg(nmr) from (values ('{[1,3]}'::nummultirange), ('{[6,12]}'::nummultirange)) t(nmr);
+select range_intersect_agg(nmr) from (values ('{[1,6]}'::nummultirange), ('{[3,12]}'::nummultirange)) t(nmr);
+select range_intersect_agg(nmr) from (values ('{[1,6], [10,12]}'::nummultirange), ('{[4,14]}'::nummultirange)) t(nmr);
 -- test with just one input:
+select range_intersect_agg(nmr) from (values ('{}'::nummultirange)) t(nmr);
 select range_intersect_agg(nmr) from (values ('{[1,2]}'::nummultirange)) t(nmr);
+select range_intersect_agg(nmr) from (values ('{[1,6], [10,12]}'::nummultirange)) t(nmr);
 select range_intersect_agg(nmr) from nummultirange_test where nmr @> 4.0;
 
 create table nummultirange_test2(nmr nummultirange);
diff --git a/src/test/regress/sql/opr_sanity.sql b/src/test/regress/sql/opr_sanity.sql
index 5a9c479692..bb129c42ff 100644
--- a/src/test/regress/sql/opr_sanity.sql
+++ b/src/test/regress/sql/opr_sanity.sql
@@ -153,6 +153,7 @@ WHERE p1.oid != p2.oid AND
     p2.prosrc NOT LIKE E'range\\_constructor_' AND
     p1.prosrc NOT LIKE E'multirange\\_constructor_' AND
     p2.prosrc NOT LIKE E'multirange\\_constructor_' AND
+    p2.prosrc NOT LIKE E'range\\_agg\\_finalfn' AND
     (p1.proargtypes[1] < p2.proargtypes[1])
 ORDER BY 1, 2;
 
-- 
2.25.1

