From 17f59547e5d4e9b46b1778d0d58f2ca0b9eb497e Mon Sep 17 00:00:00 2001
From: Chengpeng Yan <chengpeng_yan@outlook.com>
Date: Tue, 21 Apr 2026 11:47:09 +0800
Subject: [PATCH v1] Avoid overflow and underflow in regr_r2()

regr_r2() evaluated (Sxy * Sxy) / (Sxx * Syy) directly.  At
extreme scales, either product can round to zero or infinity even
though the ratio itself is finite, yielding NaN for cases such as
perfectly correlated inputs.

corr() already has a stabilized calculation for the same Sxx * Syy
denominator scale.  Move that calculation into a helper, and let
regr_r2() use it as a fallback when one of its direct products has
rounded to zero or infinity.  Otherwise keep the existing direct formula,
so the common path and its rounding behavior stay close to the old
implementation.

This preserves regr_r2()'s existing SQL-level special cases.  Add
regression tests covering the new fallback path and nearby NaN behavior.
---
 src/backend/utils/adt/float.c            | 93 ++++++++++++++++--------
 src/test/regress/expected/aggregates.out | 39 ++++++++++
 src/test/regress/sql/aggregates.sql      |  9 +++
 3 files changed, 112 insertions(+), 29 deletions(-)

diff --git a/src/backend/utils/adt/float.c b/src/backend/utils/adt/float.c
index 362c29ab803..6594b1b4325 100644
--- a/src/backend/utils/adt/float.c
+++ b/src/backend/utils/adt/float.c
@@ -3854,6 +3854,43 @@ float8_covar_samp(PG_FUNCTION_ARGS)
 	PG_RETURN_FLOAT8(Sxy / (N - 1.0));
 }
 
+/*
+ * Compute a numerically stable correlation coefficient from summary terms.
+ * Caller must already have handled SQL-level special cases such as Sxx == 0
+ * or Syy == 0.
+ */
+static float8
+float8_corrcoef_stable(float8 Sxx, float8 Syy, float8 Sxy)
+{
+	float8		product;
+	float8		sqrtproduct;
+	float8		result;
+
+	/*
+	 * The product Sxx * Syy might underflow or overflow.  If so, we can
+	 * recover by computing sqrt(Sxx) * sqrt(Syy) instead of sqrt(Sxx * Syy).
+	 * However, the double sqrt() calculation is a bit slower and less
+	 * accurate, so don't do it if we don't have to.
+	 */
+	product = Sxx * Syy;
+	if (product == 0 || isinf(product))
+		sqrtproduct = sqrt(Sxx) * sqrt(Syy);
+	else
+		sqrtproduct = sqrt(product);
+	result = Sxy / sqrtproduct;
+
+	/*
+	 * Despite all these precautions, this formula can yield results outside
+	 * [-1, 1] due to roundoff error.  Clamp it to the expected range.
+	 */
+	if (result < -1)
+		result = -1;
+	else if (result > 1)
+		result = 1;
+
+	return result;
+}
+
 Datum
 float8_corr(PG_FUNCTION_ARGS)
 {
@@ -3862,10 +3899,7 @@ float8_corr(PG_FUNCTION_ARGS)
 	float8		N,
 				Sxx,
 				Syy,
-				Sxy,
-				product,
-				sqrtproduct,
-				result;
+				Sxy;
 
 	transvalues = check_float8_array(transarray, "float8_corr", 8);
 	N = transvalues[0];
@@ -3883,29 +3917,7 @@ float8_corr(PG_FUNCTION_ARGS)
 	if (Sxx == 0 || Syy == 0)
 		PG_RETURN_NULL();
 
-	/*
-	 * The product Sxx * Syy might underflow or overflow.  If so, we can
-	 * recover by computing sqrt(Sxx) * sqrt(Syy) instead of sqrt(Sxx * Syy).
-	 * However, the double sqrt() calculation is a bit slower and less
-	 * accurate, so don't do it if we don't have to.
-	 */
-	product = Sxx * Syy;
-	if (product == 0 || isinf(product))
-		sqrtproduct = sqrt(Sxx) * sqrt(Syy);
-	else
-		sqrtproduct = sqrt(product);
-	result = Sxy / sqrtproduct;
-
-	/*
-	 * Despite all these precautions, this formula can yield results outside
-	 * [-1, 1] due to roundoff error.  Clamp it to the expected range.
-	 */
-	if (result < -1)
-		result = -1;
-	else if (result > 1)
-		result = 1;
-
-	PG_RETURN_FLOAT8(result);
+	PG_RETURN_FLOAT8(float8_corrcoef_stable(Sxx, Syy, Sxy));
 }
 
 Datum
@@ -3916,7 +3928,10 @@ float8_regr_r2(PG_FUNCTION_ARGS)
 	float8		N,
 				Sxx,
 				Syy,
-				Sxy;
+				Sxy,
+				denominator,
+				numerator,
+				result;
 
 	transvalues = check_float8_array(transarray, "float8_regr_r2", 8);
 	N = transvalues[0];
@@ -3938,7 +3953,27 @@ float8_regr_r2(PG_FUNCTION_ARGS)
 	if (Syy == 0)
 		PG_RETURN_FLOAT8(1.0);
 
-	PG_RETURN_FLOAT8((Sxy * Sxy) / (Sxx * Syy));
+	if (Sxy == 0 && !isnan(Sxx) && !isnan(Syy))
+		PG_RETURN_FLOAT8(0.0);
+
+	numerator = Sxy * Sxy;
+	denominator = Sxx * Syy;
+
+	if (numerator != 0 && denominator != 0 &&
+		!isinf(numerator) && !isinf(denominator))
+	{
+		result = numerator / denominator;
+
+		/* Roundoff error can leave the direct result just above 1. */
+		if (result > 1)
+			result = 1;
+
+		PG_RETURN_FLOAT8(result);
+	}
+
+	result = float8_corrcoef_stable(Sxx, Syy, Sxy);
+
+	PG_RETURN_FLOAT8(result * result);
 }
 
 Datum
diff --git a/src/test/regress/expected/aggregates.out b/src/test/regress/expected/aggregates.out
index ff80869fb33..4ffea8793d9 100644
--- a/src/test/regress/expected/aggregates.out
+++ b/src/test/regress/expected/aggregates.out
@@ -544,6 +544,13 @@ SELECT corr(1e-100 + g * 1e-105, 1e-100 + g * 1e-105)
     1
 (1 row)
 
+SELECT regr_r2(1e-100 + g * 1e-105, 1e-100 + g * 1e-105)
+  FROM generate_series(1, 3) g;
+ regr_r2 
+---------
+       1
+(1 row)
+
 SELECT corr(1e-100 + g * 1e-105, 1e-100 + g * 1e-105)
   FROM generate_series(1, 30) g;
  corr 
@@ -551,6 +558,20 @@ SELECT corr(1e-100 + g * 1e-105, 1e-100 + g * 1e-105)
     1
 (1 row)
 
+SELECT regr_r2(1e-100 + g * 1e-105, 1e-100 + g * 1e-105)
+  FROM generate_series(1, 30) g;
+ regr_r2 
+---------
+       1
+(1 row)
+
+SELECT regr_r2(1e154::float8 * g, 1e154::float8 * g)
+  FROM generate_series(1, 2) g;
+ regr_r2 
+---------
+       1
+(1 row)
+
 -- these examples pose definitional questions for NaN inputs,
 -- which we resolve by saying that an all-NaN input column is not all equal
 SELECT corr(g, 'NaN') FROM generate_series(1, 30) g;
@@ -559,18 +580,36 @@ SELECT corr(g, 'NaN') FROM generate_series(1, 30) g;
   NaN
 (1 row)
 
+SELECT regr_r2(g, 'NaN') FROM generate_series(1, 30) g;
+ regr_r2 
+---------
+     NaN
+(1 row)
+
 SELECT corr(0.1, 'NaN') FROM generate_series(1, 30) g;
  corr 
 ------
      
 (1 row)
 
+SELECT regr_r2(0.1, 'NaN') FROM generate_series(1, 30) g;
+ regr_r2 
+---------
+       1
+(1 row)
+
 SELECT corr('NaN', 'NaN') FROM generate_series(1, 30) g;
  corr 
 ------
   NaN
 (1 row)
 
+SELECT regr_r2('NaN', 'NaN') FROM generate_series(1, 30) g;
+ regr_r2 
+---------
+     NaN
+(1 row)
+
 -- test accum and combine functions directly
 CREATE TABLE regr_test (x float8, y float8);
 INSERT INTO regr_test VALUES (10,150),(20,250),(30,350),(80,540),(100,200);
diff --git a/src/test/regress/sql/aggregates.sql b/src/test/regress/sql/aggregates.sql
index 89bb83718e0..fe155c351c7 100644
--- a/src/test/regress/sql/aggregates.sql
+++ b/src/test/regress/sql/aggregates.sql
@@ -149,14 +149,23 @@ SELECT corr(1.3 + g * 1e-16, 1.3 + g * 1e-16)
   FROM generate_series(1, 3) g;
 SELECT corr(1e-100 + g * 1e-105, 1e-100 + g * 1e-105)
   FROM generate_series(1, 3) g;
+SELECT regr_r2(1e-100 + g * 1e-105, 1e-100 + g * 1e-105)
+  FROM generate_series(1, 3) g;
 SELECT corr(1e-100 + g * 1e-105, 1e-100 + g * 1e-105)
   FROM generate_series(1, 30) g;
+SELECT regr_r2(1e-100 + g * 1e-105, 1e-100 + g * 1e-105)
+  FROM generate_series(1, 30) g;
+SELECT regr_r2(1e154::float8 * g, 1e154::float8 * g)
+  FROM generate_series(1, 2) g;
 
 -- these examples pose definitional questions for NaN inputs,
 -- which we resolve by saying that an all-NaN input column is not all equal
 SELECT corr(g, 'NaN') FROM generate_series(1, 30) g;
+SELECT regr_r2(g, 'NaN') FROM generate_series(1, 30) g;
 SELECT corr(0.1, 'NaN') FROM generate_series(1, 30) g;
+SELECT regr_r2(0.1, 'NaN') FROM generate_series(1, 30) g;
 SELECT corr('NaN', 'NaN') FROM generate_series(1, 30) g;
+SELECT regr_r2('NaN', 'NaN') FROM generate_series(1, 30) g;
 
 -- test accum and combine functions directly
 CREATE TABLE regr_test (x float8, y float8);
-- 
2.50.1 (Apple Git-155)

