From b5448e43e3f60485ecf9e7382ec6e6e2adfc693c Mon Sep 17 00:00:00 2001
From: Ahmet Gedemenli <afgedemenli@gmail.com>
Date: Wed, 23 Mar 2022 13:42:03 +0300
Subject: [PATCH] Fix aggregate signature bug (#5854)

---
 src/backend/distributed/commands/function.c     | 16 +++++++++++++++-
 src/test/regress/expected/aggregate_support.out | 12 ++++++++++++
 src/test/regress/sql/aggregate_support.sql      |  8 ++++++++
 3 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/src/backend/distributed/commands/function.c b/src/backend/distributed/commands/function.c
index 83524625f..757c24fed 100644
--- a/src/backend/distributed/commands/function.c
+++ b/src/backend/distributed/commands/function.c
@@ -1454,7 +1454,21 @@ DefineAggregateStmtObjectAddress(Node *node, bool missing_ok)
 	}
 	else
 	{
-		objectWithArgs->objargs = list_make1(makeTypeName("anyelement"));
+		DefElem *defItem = NULL;
+		foreach_ptr(defItem, stmt->definition)
+		{
+			/*
+			 * If no explicit args are given, pg includes basetype in the signature.
+			 * If the basetype given is a type, like int4, we should include it in the
+			 * signature. In that case, defItem->arg would be a TypeName.
+			 * If the basetype given is a string, like "ANY", we shouldn't include it.
+			 */
+			if (strcmp(defItem->defname, "basetype") == 0 && IsA(defItem->arg, TypeName))
+			{
+				objectWithArgs->objargs = lappend(objectWithArgs->objargs,
+												  defItem->arg);
+			}
+		}
 	}
 
 	return FunctionToObjectAddress(OBJECT_AGGREGATE, objectWithArgs, missing_ok);
diff --git a/src/test/regress/expected/aggregate_support.out b/src/test/regress/expected/aggregate_support.out
index 5be6d26f9..2bf44bd1b 100644
--- a/src/test/regress/expected/aggregate_support.out
+++ b/src/test/regress/expected/aggregate_support.out
@@ -1236,5 +1236,17 @@ SELECT run_command_on_workers($$select aggfnoid from pg_aggregate where aggfnoid
  (localhost,57638,t,aggregate_support.dependent_agg)
 (2 rows)
 
+CREATE AGGREGATE newavg (
+   sfunc = int4_avg_accum, basetype = int4, stype = _int8,
+   finalfunc = int8_avg,
+   initcond1 = '{0,0}'
+);
+SELECT run_command_on_workers($$select aggfnoid from pg_aggregate where aggfnoid::text like '%newavg%';$$);
+               run_command_on_workers
+---------------------------------------------------------------------
+ (localhost,57637,t,aggregate_support.newavg)
+ (localhost,57638,t,aggregate_support.newavg)
+(2 rows)
+
 set client_min_messages to error;
 drop schema aggregate_support cascade;
diff --git a/src/test/regress/sql/aggregate_support.sql b/src/test/regress/sql/aggregate_support.sql
index 3862b15d6..b2cd063c7 100644
--- a/src/test/regress/sql/aggregate_support.sql
+++ b/src/test/regress/sql/aggregate_support.sql
@@ -638,5 +638,13 @@ RESET citus.create_object_propagation;
 
 SELECT run_command_on_workers($$select aggfnoid from pg_aggregate where aggfnoid::text like '%dependent_agg%';$$);
 
+CREATE AGGREGATE newavg (
+   sfunc = int4_avg_accum, basetype = int4, stype = _int8,
+   finalfunc = int8_avg,
+   initcond1 = '{0,0}'
+);
+
+SELECT run_command_on_workers($$select aggfnoid from pg_aggregate where aggfnoid::text like '%newavg%';$$);
+
 set client_min_messages to error;
 drop schema aggregate_support cascade;