Fix potential incorrect stddev(float8/4) in columnar pipeline (#10209)

Postgres uses the same float8_accum transition function for both avg()
and stddev() for these types. TimescaleDB used to use a simplified
implementation for avg(), because it doesn't require tracking sum of
squares. However, Postgres uses a shared transition state for aggregates
sharing a transition functions, as given by AggRef.aggtransno. This
leads to potentially wrong result when both avg() and stddev() of same
variable are computed in one query.

Change the aggregation states in columnar aggregation to follow
aggtransno. Remove the "float8 accum with no squares" implementation
because it cannot be safely used in general case, and checking for
special cases is complicated.

(cherry picked from commit 58e9840)

Author: akuzm

.unreleased/wrong-stddev

@@ -0,0 +1 @@
+Fixes: #10209 Potentially wrong result of stddev(float4/8) when used together with avg() in columnar query execution pipeline.

tsl/src/nodes/vector_agg/exec.c

@@ -711,55 +711,70 @@ vector_agg_begin(CustomScanState *node, EState *estate, int eflags)
 	const int tlist_length = list_length(aggregated_tlist);
 
 	/*
-	 * First, count how many grouping columns and aggregate functions we have.
+	 * First, count how many grouping columns and aggregate transition states we have.
 	 */
-	int agg_functions_counter = 0;
+	int max_aggtransno = -1;
 	int grouping_column_counter = 0;
 	for (int i = 0; i < tlist_length; i++)
 	{
 		TargetEntry *tlentry = list_nth_node(TargetEntry, aggregated_tlist, i);
 		if (IsA(tlentry->expr, Aggref))
 		{
-			agg_functions_counter++;
+			max_aggtransno = Max(max_aggtransno, castNode(Aggref, tlentry->expr)->aggtransno);
 		}
 		else
 		{
 			/* This is a grouping column. */
 			grouping_column_counter++;
 		}
 	}
-	Assert(agg_functions_counter + grouping_column_counter == tlist_length);
 
 	/*
 	 * Allocate the storage for definitions of aggregate function and grouping
 	 * columns.
 	 */
-	vector_agg_state->num_agg_defs = agg_functions_counter;
+	vector_agg_state->num_agg_defs = max_aggtransno + 1;
 	vector_agg_state->agg_defs =
 		palloc0(sizeof(*vector_agg_state->agg_defs) * vector_agg_state->num_agg_defs);
 
 	vector_agg_state->num_grouping_columns = grouping_column_counter;
 	vector_agg_state->grouping_columns = palloc0(sizeof(*vector_agg_state->grouping_columns) *
 												 vector_agg_state->num_grouping_columns);
 
+	/*
+	 * Some aggregates share transition state, but we shouldn't be able to get
+	 * more entries than the input targetlist.
+	 */
+	Assert(vector_agg_state->num_agg_defs + vector_agg_state->num_grouping_columns <= tlist_length);
+
 	/*
 	 * Loop through the aggregated targetlist again and fill the definitions.
 	 */
-	agg_functions_counter = 0;
 	grouping_column_counter = 0;
 	for (int i = 0; i < tlist_length; i++)
 	{
 		TargetEntry *tlentry = list_nth_node(TargetEntry, aggregated_tlist, i);
 		if (IsA(tlentry->expr, Aggref))
 		{
 			/* This is an aggregate function. */
-			VectorAggDef *def = &vector_agg_state->agg_defs[agg_functions_counter++];
-			def->output_offset = i;
-
 			Aggref *aggref = castNode(Aggref, tlentry->expr);
 
+			VectorAggDef *def = &vector_agg_state->agg_defs[aggref->aggtransno];
+
 			VectorAggFunctions *func = get_vector_aggregate(aggref->aggfnoid, aggref->inputcollid);
 			Assert(func != NULL);
+
+			if (def->func.agg_init != NULL)
+			{
+				/*
+				 * Already initialized for another aggregate sharing the same
+				 * transition state.
+				 */
+				Assert(def->func.agg_init == func->agg_init);
+				Assert(def->func.agg_vector == func->agg_vector);
+				continue;
+			}
+
 			def->func = *func;
 
 			if (list_length(aggref->args) > 0)
@@ -1108,9 +1123,11 @@ vector_agg_exec(CustomScanState *node)
 	/*
 	 * If we have more partial aggregation results, continue returning them.
 	 */
+	List *aggregated_tlist =
+		castNode(CustomScan, vector_agg_state->custom.ss.ps.plan)->custom_scan_tlist;
 	GroupingPolicy *grouping = vector_agg_state->grouping;
 	MemoryContext old_context = MemoryContextSwitchTo(econtext->ecxt_per_tuple_memory);
-	bool have_partial = grouping->gp_do_emit(grouping, aggregated_slot);
+	bool have_partial = grouping->gp_do_emit(grouping, aggregated_tlist, aggregated_slot);
 	MemoryContextSwitchTo(old_context);
 	if (have_partial)
 	{
@@ -1202,7 +1219,7 @@ vector_agg_exec(CustomScanState *node)
 	 * If we have partial aggregation results, start returning them.
 	 */
 	old_context = MemoryContextSwitchTo(econtext->ecxt_per_tuple_memory);
-	have_partial = grouping->gp_do_emit(grouping, aggregated_slot);
+	have_partial = grouping->gp_do_emit(grouping, aggregated_tlist, aggregated_slot);
 	MemoryContextSwitchTo(old_context);
 	if (have_partial)
 	{

tsl/src/nodes/vector_agg/exec.h

@@ -18,7 +18,6 @@ typedef struct VectorAggDef
 {
 	VectorAggFunctions func;
 	Expr *argument;
-	int output_offset;
 	List *filter_clauses;
 
 	/*
@@ -42,6 +41,10 @@ typedef struct VectorAggState
 {
 	CustomScanState custom;
 
+	/*
+	 * Postgres makes some Aggrefs share the transition state. This array tracks
+	 * the unique transition states indexed by Aggref.aggtransno.
+	 */
 	int num_agg_defs;
 	VectorAggDef *agg_defs;
 

tsl/src/nodes/vector_agg/function/float48_accum_single.c

@@ -6,8 +6,7 @@
 
 /*
  * Vectorized implementation of a Postgres float{4,8}_accum() transition
- * function for a single type. They use the same Youngs-Cramer state, but for
- * AVG we can skip calculating the Sxx variable.
+ * function for a single type.
  */
 
 #ifdef GENERATE_DISPATCH_TABLE
@@ -17,20 +16,16 @@
 extern VectorAggFunctions FUNCTION_NAME(argdef);
 
 /*
- * Helper macros to generate the cases for the given argument type. We support
- * different aggregate functions based on whether we calculate the Sxx variable.
+ * Helper macros to generate the cases for the given argument type.
  */
-#ifdef NEED_SXX
 #define ACCUM_CASE_HELPER(PG_TYPE)                                                                 \
 	case F_STDDEV_##PG_TYPE:                                                                       \
 	case F_STDDEV_SAMP_##PG_TYPE:                                                                  \
 	case F_STDDEV_POP_##PG_TYPE:                                                                   \
 	case F_VARIANCE_##PG_TYPE:                                                                     \
 	case F_VAR_SAMP_##PG_TYPE:                                                                     \
-	case F_VAR_POP_##PG_TYPE:
-#else
-#define ACCUM_CASE_HELPER(PG_TYPE) case F_AVG_##PG_TYPE:
-#endif
+	case F_VAR_POP_##PG_TYPE:                                                                      \
+	case F_AVG_##PG_TYPE:
 
 #define ACCUM_CASE(PG_TYPE) ACCUM_CASE_HELPER(PG_TYPE)
 
@@ -49,9 +44,7 @@ typedef struct
 {
 	double N;
 	double Sx;
-#ifdef NEED_SXX
 	double Sxx;
-#endif
 } FUNCTION_NAME(state);
 
 static void
@@ -84,17 +77,7 @@ FUNCTION_NAME(emit)(void *agg_state, Datum *out_result, bool *out_isnull)
 	 */
 	((float8 *) ARR_DATA_PTR(result))[0] = state->N;
 	((float8 *) ARR_DATA_PTR(result))[1] = state->Sx;
-	((float8 *) ARR_DATA_PTR(result))[2] =
-		/*
-		 * Sxx should be NaN if any of the inputs are infinite or NaN. This is
-		 * checked by float8_combine even if it's not used for the actual
-		 * calculations.
-		 */
-		0. * state->Sx
-#ifdef NEED_SXX
-		+ state->Sxx
-#endif
-		;
+	((float8 *) ARR_DATA_PTR(result))[2] = state->Sxx;
 
 	*out_result = PointerGetDatum(result);
 	*out_isnull = false;
@@ -104,12 +87,8 @@ FUNCTION_NAME(emit)(void *agg_state, Datum *out_result, bool *out_isnull)
  * Youngs-Cramer update for rows after the first.
  */
 static pg_attribute_always_inline void
-FUNCTION_NAME(update)(const uint64 *filter, const CTYPE *values, int row, double *N, double *Sx
-#ifdef NEED_SXX
-					  ,
-					  double *Sxx
-#endif
-)
+FUNCTION_NAME(update)(const uint64 *filter, const CTYPE *values, int row, double *N, double *Sx,
+					  double *Sxx)
 {
 	const CTYPE newval = values[row];
 	if (!arrow_row_is_valid(filter, row))
@@ -123,11 +102,9 @@ FUNCTION_NAME(update)(const uint64 *filter, const CTYPE *values, int row, double
 	 */
 	const double newN = *N + 1.0;
 	const double newSx = *Sx + newval;
-#ifdef NEED_SXX
 	Assert(*N > 0.0);
 	const double tmp = newval * newN - newSx;
 	*Sxx += tmp * tmp / (*N * newN);
-#endif
 
 	*N = newN;
 	*Sx = newSx;
@@ -137,67 +114,43 @@ FUNCTION_NAME(update)(const uint64 *filter, const CTYPE *values, int row, double
  * Combine two Youngs-Cramer states following the float8_combine() function.
  */
 static pg_attribute_always_inline void
-FUNCTION_NAME(combine)(double *inout_N, double *inout_Sx,
-#ifdef NEED_SXX
-					   double *inout_Sxx,
-#endif
-					   double N2, double Sx2
-#ifdef NEED_SXX
-					   ,
-					   double Sxx2
-#endif
-)
+FUNCTION_NAME(combine)(double *inout_N, double *inout_Sx, double *inout_Sxx, double N2, double Sx2,
+					   double Sxx2)
 {
 	const double N1 = *inout_N;
 	const double Sx1 = *inout_Sx;
-#ifdef NEED_SXX
 	const double Sxx1 = *inout_Sxx;
-#endif
 
 	if (unlikely(N1 == 0))
 	{
 		*inout_N = N2;
 		*inout_Sx = Sx2;
-#ifdef NEED_SXX
 		*inout_Sxx = Sxx2;
-#endif
 		return;
 	}
 
 	if (unlikely(N2 == 0))
 	{
 		*inout_N = N1;
 		*inout_Sx = Sx1;
-#ifdef NEED_SXX
 		*inout_Sxx = Sxx1;
-#endif
 		return;
 	}
 
 	const double combinedN = N1 + N2;
 	const double combinedSx = Sx1 + Sx2;
-#ifdef NEED_SXX
 	const double tmp = Sx1 / N1 - Sx2 / N2;
 	const double combinedSxx = Sxx1 + Sxx2 + N1 * N2 * tmp * tmp / combinedN;
-#endif
 
 	*inout_N = combinedN;
 	*inout_Sx = combinedSx;
-#ifdef NEED_SXX
 	*inout_Sxx = combinedSxx;
-#endif
 }
 
-#ifdef NEED_SXX
 #define UPDATE(filter, values, row, N, Sx, Sxx)                                                    \
 	FUNCTION_NAME(update)(filter, values, row, N, Sx, Sxx)
 #define COMBINE(inout_N, inout_Sx, inout_Sxx, N2, Sx2, Sxx2)                                       \
 	FUNCTION_NAME(combine)(inout_N, inout_Sx, inout_Sxx, N2, Sx2, Sxx2)
-#else
-#define UPDATE(filter, values, row, N, Sx, Sxx) FUNCTION_NAME(update)(filter, values, row, N, Sx)
-#define COMBINE(inout_N, inout_Sx, inout_Sxx, N2, Sx2, Sxx2)                                       \
-	FUNCTION_NAME(combine)(inout_N, inout_Sx, N2, Sx2)
-#endif
 
 static pg_attribute_always_inline void
 FUNCTION_NAME(vector_impl)(void *agg_state, size_t n, const CTYPE *values, const uint64 *filter,
@@ -214,13 +167,10 @@ FUNCTION_NAME(vector_impl)(void *agg_state, size_t n, const CTYPE *values, const
 	 */
 	double Narray[UNROLL_SIZE] = { 0 };
 	double Sxarray[UNROLL_SIZE] = { 0 };
-#ifdef NEED_SXX
 	double Sxxarray[UNROLL_SIZE] = { 0 };
-#endif
 
 	size_t row = 0;
 
-#ifdef NEED_SXX
 	/*
 	 * Initialize each state with the first matching row. We do this separately
 	 * to make the actual update function branchless, namely the computation of
@@ -251,7 +201,6 @@ FUNCTION_NAME(vector_impl)(void *agg_state, size_t n, const CTYPE *values, const
 	{
 		UPDATE(filter, values, row, &Narray[inner], &Sxarray[inner], &Sxxarray[inner]);
 	}
-#endif
 
 	/*
 	 * Unrolled loop.
@@ -300,7 +249,6 @@ FUNCTION_NAME(one)(void *restrict agg_state, const CTYPE value)
 	 */
 	const double newN = state->N + 1.0;
 	const double newSx = state->Sx + value;
-#ifdef NEED_SXX
 	if (state->N > 0.0)
 	{
 		const double tmp = value * newN - newSx;
@@ -310,7 +258,6 @@ FUNCTION_NAME(one)(void *restrict agg_state, const CTYPE value)
 	{
 		state->Sxx = 0 * value;
 	}
-#endif
 
 	state->N = newN;
 	state->Sx = newSx;

tsl/src/nodes/vector_agg/function/float48_accum_templates.c

@@ -20,19 +20,5 @@
 #include "template_helper.h"
 #include <compression/arrow_c_data_interface.h>
 
-#ifndef GENERATE_DISPATCH_TABLE
-
-#endif
-
-/*
- * Templated parts for vectorized avg(float).
- */
-#define AGG_NAME accum_no_squares
-#include "float48_accum_types.c"
-
-/*
- * Templated parts for vectorized functions that use the Sxx state (stddev etc).
- */
 #define AGG_NAME accum_with_squares
-#define NEED_SXX
 #include "float48_accum_types.c"

tsl/src/nodes/vector_agg/function/float48_accum_types.c

@@ -6,8 +6,7 @@
 
 /*
  * Functions handled by *accum() aggregate functions states, implementation for
- * all types. They use the same Youngs-Cramer state, but for AVG we can skip
- * calculating the Sxx variable.
+ * all types.
  */
 
 #define PG_TYPE FLOAT4
@@ -23,4 +22,3 @@
 #include "float48_accum_single.c"
 
 #undef AGG_NAME
-#undef NEED_SXX

tsl/src/nodes/vector_agg/grouping_policy.h

@@ -45,7 +45,7 @@ typedef struct GroupingPolicy
 	/*
 	 * Emit a partial aggregation result into the result slot.
 	 */
-	bool (*gp_do_emit)(GroupingPolicy *gp, TupleTableSlot *aggregated_slot);
+	bool (*gp_do_emit)(GroupingPolicy *gp, List *aggregated_tlist, TupleTableSlot *aggregated_slot);
 
 	/*
 	 * Destroy the grouping policy.