Merge pull request #436 from diffix/piotr/negative-summands

pdobacz · web-flow · commit 25ef33689611 · 2022-08-17T14:15:56.000+02:00
Support negative summands
diff --git a/pg_diffix/aggregation/contribution_tracker.h b/pg_diffix/aggregation/contribution_tracker.h
@@ -25,12 +25,16 @@ typedef contribution_t (*ContributionCombineFunc)(contribution_t x, contribution
 /* Casts x to double. */
 typedef double (*ContributionToDoubleFunc)(contribution_t x);
 
+/* Computes absolute value. */
+typedef contribution_t (*ContributionAbsFunc)(contribution_t x);
+
 typedef struct ContributionDescriptor
 {
   ContributionGreaterFunc contribution_greater;
   ContributionEqualFunc contribution_equal;
   ContributionCombineFunc contribution_combine;
   ContributionToDoubleFunc contribution_to_double;
+  ContributionAbsFunc contribution_abs;
   contribution_t contribution_initial; /* Initial or "zero" value for a contribution */
 } ContributionDescriptor;
 
diff --git a/src/aggregation/sum.c b/src/aggregation/sum.c
@@ -20,12 +20,15 @@ static double finalize_sum_result(const SummableResultAccumulator *accumulator)
  *-------------------------------------------------------------------------
  */
 
+typedef ContributionTrackerState *SumLeg;
+
 typedef struct SumState
 {
   AnonAggState base;
   int trackers_count;
   Oid summand_type;
-  ContributionTrackerState *trackers[FLEXIBLE_ARRAY_MEMBER];
+  SumLeg *positive;
+  SumLeg *negative;
 } SumState;
 
 static void sum_final_type(const ArgsDescriptor *args_desc, Oid *type, int32 *typmod, Oid *collid)
@@ -63,9 +66,11 @@ static AnonAggState *sum_create_state(MemoryContext memory_context, ArgsDescript
   MemoryContext old_context = MemoryContextSwitchTo(memory_context);
 
   int trackers_count = args_desc->num_args - SUM_AIDS_OFFSET;
-  SumState *state = palloc0(sizeof(SumState) + trackers_count * sizeof(ContributionTrackerState *));
+  SumState *state = palloc0(sizeof(SumState));
   state->trackers_count = trackers_count;
   state->summand_type = args_desc->args[SUM_VALUE_INDEX].type_oid;
+  state->positive = palloc0(trackers_count * sizeof(SumLeg));
+  state->negative = palloc0(trackers_count * sizeof(SumLeg));
   ContributionDescriptor typed_sum_descriptor = {0};
   switch (state->summand_type)
   {
@@ -87,36 +92,54 @@ static AnonAggState *sum_create_state(MemoryContext memory_context, ArgsDescript
   for (int i = 0; i < trackers_count; i++)
   {
     Oid aid_type = args_desc->args[i + SUM_AIDS_OFFSET].type_oid;
-    state->trackers[i] = contribution_tracker_new(get_aid_mapper(aid_type), &typed_sum_descriptor);
+    state->positive[i] = contribution_tracker_new(get_aid_mapper(aid_type), &typed_sum_descriptor);
+    state->negative[i] = contribution_tracker_new(get_aid_mapper(aid_type), &typed_sum_descriptor);
   }
 
   MemoryContextSwitchTo(old_context);
   return &state->base;
 }
 
-static SummableResultAccumulator sum_calculate_final(AnonAggState *base_state, Bucket *bucket, BucketDescriptor *bucket_desc)
+typedef struct SumResult
+{
+  bool not_enough_aid_values;
+  SummableResultAccumulator positive;
+  SummableResultAccumulator negative;
+} SumResult;
+
+static SumResult sum_calculate_final(AnonAggState *base_state, Bucket *bucket, BucketDescriptor *bucket_desc)
 {
   SumState *state = (SumState *)base_state;
-  SummableResultAccumulator result_accumulator = {0};
+  SummableResultAccumulator positive_result_accumulator = {0};
+  SummableResultAccumulator negative_result_accumulator = {0};
   seed_t bucket_seed = compute_bucket_seed(bucket, bucket_desc);
 
   for (int i = 0; i < state->trackers_count; i++)
   {
-    SummableResult result = calculate_result(bucket_seed, state->trackers[i]);
+    SummableResult positive_result = calculate_result(bucket_seed, state->positive[i]);
+    SummableResult negative_result = calculate_result(bucket_seed, state->negative[i]);
 
-    accumulate_result(&result_accumulator, &result);
-    if (result_accumulator.not_enough_aid_values)
-      break;
+    if (positive_result.not_enough_aid_values && negative_result.not_enough_aid_values)
+    {
+      return (SumResult){.not_enough_aid_values = true};
+    }
+    else
+    {
+      /* Unless both legs had `not_enough_aid_values` for given AID instance, we proceed. */
+      accumulate_result(&positive_result_accumulator, &positive_result);
+      accumulate_result(&negative_result_accumulator, &negative_result);
+    }
   }
-  return result_accumulator;
+  return (SumResult){.positive = positive_result_accumulator, .negative = negative_result_accumulator};
 }
 
 static Datum sum_finalize(AnonAggState *base_state, Bucket *bucket, BucketDescriptor *bucket_desc, bool *is_null)
 {
   SumState *state = (SumState *)base_state;
-  SummableResultAccumulator result_accumulator = sum_calculate_final(base_state, bucket, bucket_desc);
+  SumResult result = sum_calculate_final(base_state, bucket, bucket_desc);
 
-  if (result_accumulator.not_enough_aid_values)
+  /* We deliberately ignore the `not_enough_aid_values` fields in the `result.positive` and `negative`. */
+  if (result.not_enough_aid_values)
   {
     *is_null = true;
     switch (state->summand_type)
@@ -138,21 +161,22 @@ static Datum sum_finalize(AnonAggState *base_state, Bucket *bucket, BucketDescri
   }
   else
   {
+    double combined_result = finalize_sum_result(&result.positive) - finalize_sum_result(&result.negative);
     switch (state->summand_type)
     {
     case INT2OID:
     case INT4OID:
-      return Int64GetDatum((int64)round(finalize_sum_result(&result_accumulator)));
+      return Int64GetDatum((int64)round(combined_result));
     case INT8OID:
     case NUMERICOID:
-      return DirectFunctionCall1(float8_numeric, Float8GetDatum(finalize_sum_result(&result_accumulator)));
+      return DirectFunctionCall1(float8_numeric, Float8GetDatum(combined_result));
     case FLOAT4OID:
-      return Float4GetDatum((float4)finalize_sum_result(&result_accumulator));
+      return Float4GetDatum((float4)combined_result);
     case FLOAT8OID:
-      return Float8GetDatum(finalize_sum_result(&result_accumulator));
+      return Float8GetDatum(combined_result);
     default:
       Assert(false);
-      return Float8GetDatum(finalize_sum_result(&result_accumulator));
+      return Float8GetDatum(combined_result);
     }
   }
 }
@@ -163,7 +187,8 @@ static void sum_merge(AnonAggState *dst_base_state, const AnonAggState *src_base
   const SumState *src_state = (const SumState *)src_base_state;
 
   Assert(dst_state->summand_type == src_state->summand_type);
-  merge_trackers(dst_state->trackers_count, src_state->trackers_count, dst_state->trackers, src_state->trackers);
+  merge_trackers(dst_state->trackers_count, src_state->trackers_count, dst_state->positive, src_state->positive);
+  merge_trackers(dst_state->trackers_count, src_state->trackers_count, dst_state->negative, src_state->negative);
 }
 
 static contribution_t summand_to_contribution(Datum arg, Oid summand_type)
@@ -202,16 +227,26 @@ static void sum_transition(AnonAggState *base_state, int num_args, NullableDatum
     for (int i = 0; i < state->trackers_count; i++)
     {
       int aid_index = i + SUM_AIDS_OFFSET;
+      ContributionDescriptor descriptor = state->positive[i]->contribution_descriptor;
+      contribution_t abs_contribution = descriptor.contribution_abs(value_contribution);
+      ContributionCombineFunc combine = descriptor.contribution_combine;
+      ContributionGreaterFunc gt = descriptor.contribution_greater;
+      ContributionEqualFunc eq = descriptor.contribution_equal;
 
       if (!args[aid_index].isnull)
       {
-        aid_t aid = state->trackers[i]->aid_mapper(args[aid_index].value);
-        contribution_tracker_update_contribution(state->trackers[i], aid, value_contribution);
+        aid_t aid = state->positive[i]->aid_mapper(args[aid_index].value);
+        if (gt(value_contribution, descriptor.contribution_initial) || eq(value_contribution, descriptor.contribution_initial))
+          contribution_tracker_update_contribution(state->positive[i], aid, abs_contribution);
+        if (gt(descriptor.contribution_initial, value_contribution) || eq(value_contribution, descriptor.contribution_initial))
+          contribution_tracker_update_contribution(state->negative[i], aid, abs_contribution);
       }
       else
       {
-        ContributionCombineFunc combine = state->trackers[i]->contribution_descriptor.contribution_combine;
-        state->trackers[i]->unaccounted_for = combine(state->trackers[i]->unaccounted_for, value_contribution);
+        if (gt(value_contribution, descriptor.contribution_initial))
+          state->positive[i]->unaccounted_for = combine(state->positive[i]->unaccounted_for, abs_contribution);
+        if (gt(descriptor.contribution_initial, value_contribution))
+          state->negative[i]->unaccounted_for = combine(state->negative[i]->unaccounted_for, abs_contribution);
       }
     }
   }
@@ -240,15 +275,17 @@ static void sum_noise_final_type(const ArgsDescriptor *args_desc, Oid *type, int
 
 static Datum sum_noise_finalize(AnonAggState *base_state, Bucket *bucket, BucketDescriptor *bucket_desc, bool *is_null)
 {
-  SummableResultAccumulator result_accumulator = sum_calculate_final(base_state, bucket, bucket_desc);
-  if (result_accumulator.not_enough_aid_values)
+  SumResult result = sum_calculate_final(base_state, bucket, bucket_desc);
+
+  /* We deliberately ignore the `not_enough_aid_values` fields in the `result.positive` and `negative`. */
+  if (result.not_enough_aid_values)
   {
     *is_null = true;
     return Float8GetDatum(0.0);
   }
   else
   {
-    return Float8GetDatum(finalize_noise_result(&result_accumulator));
+    return Float8GetDatum(sqrt(pow(finalize_noise_result(&result.positive), 2) + pow(finalize_noise_result(&result.negative), 2)));
   }
 }
 
diff --git a/src/aggregation/summable.c b/src/aggregation/summable.c
@@ -27,11 +27,17 @@ static double integer_contribution_to_double(contribution_t x)
   return (double)x.integer;
 }
 
+static contribution_t integer_contribution_abs(contribution_t x)
+{
+  return (contribution_t){.integer = labs(x.integer)};
+}
+
 const ContributionDescriptor integer_descriptor = {
     .contribution_greater = integer_contribution_greater,
     .contribution_equal = integer_contribution_equal,
     .contribution_combine = integer_contribution_combine,
     .contribution_to_double = integer_contribution_to_double,
+    .contribution_abs = integer_contribution_abs,
     .contribution_initial = {.integer = 0},
 };
 static bool real_contribution_greater(contribution_t x, contribution_t y)
@@ -54,11 +60,17 @@ static double real_contribution_to_double(contribution_t x)
   return (double)x.real;
 }
 
+static contribution_t real_contribution_abs(contribution_t x)
+{
+  return (contribution_t){.real = fabs(x.real)};
+}
+
 const ContributionDescriptor real_descriptor = {
     .contribution_greater = real_contribution_greater,
     .contribution_equal = real_contribution_equal,
     .contribution_combine = real_contribution_combine,
     .contribution_to_double = real_contribution_to_double,
+    .contribution_abs = real_contribution_abs,
     .contribution_initial = {.real = 0.0},
 };
 
diff --git a/test/expected/noiseless.out b/test/expected/noiseless.out
@@ -6,6 +6,11 @@ SET pg_diffix.outlier_count_min = 1;
 SET pg_diffix.outlier_count_max = 1;
 SET pg_diffix.top_count_min = 3;
 SET pg_diffix.top_count_max = 3;
+-- Additional tables for SUM testing
+CREATE TABLE test_customers_negative AS SELECT id, city, -discount as discount, planet FROM test_customers;
+CREATE TABLE test_customers_mixed AS SELECT id, city, discount - 1.0 as discount, planet FROM test_customers;
+CALL diffix.mark_personal('public.test_customers_negative', 'id');
+CALL diffix.mark_personal('public.test_customers_mixed', 'id');
 SET ROLE diffix_test;
 SET pg_diffix.session_access_level = 'anonymized_trusted';
 ----------------------------------------------------------------
@@ -89,6 +94,18 @@ SELECT city, SUM(discount), diffix.sum_noise(discount) FROM test_customers GROUP
  Berlin |         9 |         0
 (3 rows)
 
+SELECT SUM(discount), diffix.sum_noise(discount) FROM test_customers_negative;
+ sum | sum_noise 
+-----+-----------
+ -19 |         0
+(1 row)
+
+SELECT SUM(discount), diffix.sum_noise(discount) FROM test_customers_mixed;
+ sum | sum_noise 
+-----+-----------
+ 2.5 |         0
+(1 row)
+
 -- sum supports numeric type
 SELECT city, SUM(discount::numeric), pg_typeof(SUM(discount::numeric)), diffix.sum_noise(discount::numeric)
 FROM test_customers
diff --git a/test/expected/validation.out b/test/expected/validation.out
@@ -294,11 +294,11 @@ SELECT * FROM diffix.show_settings() LIMIT 2;
  pg_diffix.default_access_level | direct  | Access level for unlabeled users.
 (2 rows)
 
-SELECT * FROM diffix.show_labels() WHERE objname LIKE 'public.test_customers%';
- objtype |         objname          |  label   
----------+--------------------------+----------
- table   | public.test_customers    | personal
- column  | public.test_customers.id | aid
+SELECT * FROM diffix.show_labels() WHERE objname LIKE 'public.empty_test_customers%';
+ objtype |            objname             |  label   
+---------+--------------------------------+----------
+ table   | public.empty_test_customers    | personal
+ column  | public.empty_test_customers.id | aid
 (2 rows)
 
 -- Allow prepared statements
diff --git a/test/sql/noiseless.sql b/test/sql/noiseless.sql
@@ -8,6 +8,12 @@ SET pg_diffix.outlier_count_max = 1;
 SET pg_diffix.top_count_min = 3;
 SET pg_diffix.top_count_max = 3;
 
+-- Additional tables for SUM testing
+CREATE TABLE test_customers_negative AS SELECT id, city, -discount as discount, planet FROM test_customers;
+CREATE TABLE test_customers_mixed AS SELECT id, city, discount - 1.0 as discount, planet FROM test_customers;
+CALL diffix.mark_personal('public.test_customers_negative', 'id');
+CALL diffix.mark_personal('public.test_customers_mixed', 'id');
+
 SET ROLE diffix_test;
 SET pg_diffix.session_access_level = 'anonymized_trusted';
 
@@ -41,6 +47,9 @@ SELECT SUM(discount), diffix.sum_noise(discount) FROM test_customers;
 SELECT city, SUM(id), diffix.sum_noise(id) FROM test_customers GROUP BY 1;
 SELECT city, SUM(discount), diffix.sum_noise(discount) FROM test_customers GROUP BY 1;
 
+SELECT SUM(discount), diffix.sum_noise(discount) FROM test_customers_negative;
+SELECT SUM(discount), diffix.sum_noise(discount) FROM test_customers_mixed;
+
 -- sum supports numeric type
 SELECT city, SUM(discount::numeric), pg_typeof(SUM(discount::numeric)), diffix.sum_noise(discount::numeric)
 FROM test_customers
diff --git a/test/sql/validation.sql b/test/sql/validation.sql
@@ -150,7 +150,7 @@ SELECT EXISTS (SELECT FROM Information_Schema.tables WHERE table_schema='public'
 
 -- Settings and labels UDFs work
 SELECT * FROM diffix.show_settings() LIMIT 2;
-SELECT * FROM diffix.show_labels() WHERE objname LIKE 'public.test_customers%';
+SELECT * FROM diffix.show_labels() WHERE objname LIKE 'public.empty_test_customers%';
 
 -- Allow prepared statements
 PREPARE prepared(float) AS SELECT discount, count(*) FROM empty_test_customers WHERE discount = $1 GROUP BY 1;
