Merge pull request #437 from diffix/piotr/fix-avg-divide-by-zero

Fix avg/avg_noise when anonymized count is 0

Author: pdobacz

pg_diffix/oid_cache.h

@@ -56,6 +56,8 @@ typedef struct Oids
 
   /* Internal functions */
   Oid internal_qual_wrapper; /* diffix.internal_qual_wrapper(boolean) */
+
+  Oid op_int8eq;
 } Oids;
 
 /*

src/oid_cache.c

@@ -3,6 +3,7 @@
 #include "catalog/pg_type.h"
 #include "lib/stringinfo.h"
 #include "parser/parse_func.h"
+#include "parser/parse_oper.h"
 #include "utils/lsyscache.h"
 
 #include "pg_diffix/oid_cache.h"
@@ -66,6 +67,10 @@ void oid_cache_init(void)
 
   g_oid_cache.internal_qual_wrapper = lookup_function("diffix", "internal_qual_wrapper", 1, (Oid[]){BOOLOID});
 
+  List *eq_op_name = list_make1(makeString("="));
+  g_oid_cache.op_int8eq = LookupOperName(NULL, eq_op_name, INT8OID, INT8OID, false, -1);
+  list_free_deep(eq_op_name);
+
   g_loaded = true;
 }
 

src/query/anonymization.c

@@ -188,15 +188,30 @@ static void rewrite_count_histogram(Aggref *aggref, List *aid_refs)
   append_aid_args(aggref, aid_refs);
 }
 
-/* Intended for the denominator in the `avg(col)` rewritten to `sum(col) / count(col)`. */
-static Aggref *make_anon_count_value_aggref(const Aggref *source_aggref)
+/*
+ * Intended for the denominator in the `avg(col)` rewritten to `sum(col) / nullif(count(col), 0)`.
+ * `nullif` is necessary to handle cases where large negative noise brings `count` down to 0
+ * during global aggregation.
+ */
+static Expr *make_safe_anon_count_value(const Aggref *source_aggref)
 {
   Aggref *count_aggref = copyObjectImpl(source_aggref);
   count_aggref->aggfnoid = g_oid_cache.anon_count_value;
   count_aggref->aggtype = INT8OID;
   count_aggref->aggstar = false;
   count_aggref->aggdistinct = false;
-  return count_aggref;
+
+  NullIfExpr *nullif = makeNode(NullIfExpr);
+  nullif->opno = g_oid_cache.op_int8eq;
+  nullif->opfuncid = F_INT8EQ;
+  nullif->opresulttype = count_aggref->aggtype;
+  nullif->opretset = false;
+  nullif->opcollid = 0;
+  nullif->inputcollid = 0;
+  nullif->args = list_make2(count_aggref, make_const_int64(0L));
+  nullif->location = count_aggref->location;
+
+  return (Expr *)nullif;
 }
 
 static FuncExpr *make_i4tod(List *args)
@@ -254,11 +269,9 @@ static Node *rewrite_to_avg_aggregator(Aggref *aggref, List *aid_refs)
   aggref->aggfnoid = g_oid_cache.anon_sum;
   aggref->aggstar = false;
   aggref->aggdistinct = false;
-
-  Aggref *count_aggref = make_anon_count_value_aggref(aggref);
-
   append_aid_args(aggref, aid_refs);
-  append_aid_args(count_aggref, aid_refs);
+
+  Expr *count_aggref = make_safe_anon_count_value(aggref);
 
   FuncExpr *cast_sum;
   FuncExpr *cast_count;
@@ -310,11 +323,9 @@ static Node *rewrite_to_avg_noise_aggregator(Aggref *aggref, List *aid_refs)
   aggref->aggtype = FLOAT8OID;
   aggref->aggstar = false;
   aggref->aggdistinct = false;
-
-  Aggref *count_aggref = make_anon_count_value_aggref(aggref);
-
   append_aid_args(aggref, aid_refs);
-  append_aid_args(count_aggref, aid_refs);
+
+  Expr *count_aggref = make_safe_anon_count_value(aggref);
 
   FuncExpr *cast_count = make_i4tod(list_make1(count_aggref));
   FuncExpr *division = make_float8div(list_make2(aggref, cast_count));