changed weighted_avg to ignore rows were value is NULL, added tests f…

…or this case
gropaul · Jan 8, 2025 · 26ffde8 · 26ffde8
1 parent af0f93b
commit 26ffde8
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 2 deletions.
diff --git a/src/catalog/default/default_functions.cpp b/src/catalog/default/default_functions.cpp
@@ -110,7 +110,7 @@ static const DefaultMacro internal_macros[] = {
 	{DEFAULT_SCHEMA, "geomean", {"x", nullptr}, {{nullptr, nullptr}}, "exp(avg(ln(x)))"},
 	{DEFAULT_SCHEMA, "geometric_mean", {"x", nullptr}, {{nullptr, nullptr}}, "geomean(x)"},
 
-	{DEFAULT_SCHEMA, "weighted_avg", {"value", "weight", nullptr}, {{nullptr, nullptr}}, "SUM(value * weight) / SUM(weight)"},
+	{DEFAULT_SCHEMA, "weighted_avg", {"value", "weight", nullptr}, {{nullptr, nullptr}}, "SUM(value * weight) / SUM(CASE WHEN value IS NOT NULL THEN weight ELSE 0 END)"},
 	{DEFAULT_SCHEMA, "wavg", {"value", "weight", nullptr}, {{nullptr, nullptr}}, "weighted_avg(value, weight)"},
 
     {DEFAULT_SCHEMA, "list_reverse", {"l", nullptr}, {{nullptr, nullptr}}, "l[:-:-1]"},

diff --git a/test/sql/aggregate/aggregates/test_weighted_avg.test b/test/sql/aggregate/aggregates/test_weighted_avg.test
@@ -49,4 +49,24 @@ query II
 SELECT name, weighted_avg(grade, etcs) FROM students GROUP BY name ORDER BY name
 ----
 Alice	7.5
-Bob		6.5
+Bob		6.5
+
+# weighted_avg skips rows were the weight is NULL, so adding a row with NULL weight should not change the result
+statement ok
+INSERT INTO students VALUES ('Alice', 42, NULL);
+
+query II
+SELECT name, weighted_avg(grade, etcs) FROM students GROUP BY name ORDER BY name
+----
+Alice	7.5
+Bob		6.5
+
+# weighted_avg skips rows were the value is NULL, so adding a row with NULL value should not change the result
+statement ok
+INSERT INTO students VALUES ('Alice', NULL, 42);
+
+query II
+SELECT name, weighted_avg(grade, etcs) FROM students GROUP BY name ORDER BY name
+----
+Alice	7.5
+Bob		6.5