Merge pull request #417 from diffix/edon/shared-agg-state

Support state sharing across aggregates

Author: edongashi

pg_diffix/aggregation/bucket_scan.h

@@ -21,4 +21,9 @@ extern Plan *make_bucket_scan(Plan *left_tree, AnonymizationContext *anon_contex
  */
 extern bool is_bucket_scan(Plan *plan);
 
+/*
+ * Returns true if another aggregate in the bucket does identical transitions as the given Aggref.
+ */
+extern bool aggref_shares_state(Aggref *aggref);
+
 #endif /* PG_DIFFIX_BUCKET_SCAN_H */

pg_diffix/aggregation/common.h

@@ -28,6 +28,8 @@
  * The `finalize` function derives the final value (of type `final_type`) of the aggregator.
  * Temporary and return data should not be allocated in the state's memory context but in
  * the current memory context which is shorter lived. See below for information about memory.
+ * Because state might be borrowed from another aggregator, `finalize` must be idempotent,
+ * meaning multiple executions against the same state have to return the same result.
  *
  * The `explain` function returns a human-readable representation of the aggregator state.
  * As with `finalize`, the current memory context should be used for temporary and return values.
@@ -76,6 +78,8 @@
  *-------------------------------------------------------------------------
  */
 
+#define AGG_STATE_REDIRECTED NULL
+
 /* Describes a single function call argument. */
 typedef struct ArgDescriptor
 {
@@ -125,9 +129,10 @@ typedef struct BucketAttribute
   BucketAttributeTag tag; /* Label or aggregate? */
   struct
   {
-    Oid fn_oid;                /* Agg function OID */
+    Aggref *aggref;            /* Expr of aggregate */
     ArgsDescriptor *args_desc; /* Agg arguments descriptor */
     const AnonAggFuncs *funcs; /* Agg funcs if tag=BUCKET_ANON_AGG */
+    int redirect_to;           /* If shared, points to attribute that owns the state */
   } agg;                       /* Populated if tag!=BUCKET_LABEL */
   int typ_len;                 /* Data type length */
   bool typ_byval;              /* Data type is by value? */

src/aggregation/bucket_scan.c

@@ -36,9 +36,9 @@
  *   Aggregates can be found in tlist and qual. We need to export both in Agg's tlist because we
  *   move the actual projection and qual to BucketScan. TLEs n+1..n+m will be the aggregates.
  *   When rewriting expressions for proj/qual, we do a simple equality-based deduplication to
- *   minimize aggregates in tlist. It is not very important to be smart about optimizing at this
- *   stage because ExecInitAgg will take care of sharing aggregation state during execution.
- *   Arguments to aggregates are untouched because they do not leave the node.
+ *   minimize aggregates in tlist. During execution, anonymizing aggregators will reuse state
+ *   if conditions in `can_share_agg_state` are met. Arguments to aggregates are untouched
+ *   because they do not leave the node.
  *
  *   Projection/filtering:
  *
@@ -89,14 +89,60 @@ static inline bool has_star_bucket(BucketScanState *bucket_state)
   return linitial(bucket_state->buckets) != NULL;
 }
 
-/* Memory context of currently executing BucketScan node. */
-MemoryContext g_current_bucket_context = NULL;
+/* State of currently executing bucket scan. */
+static BucketScanState *g_current_bucket_scan = NULL;
+
+MemoryContext get_current_bucket_context(void);
+bool aggref_shares_state(Aggref *aggref);
+
+/* Used by common.c to locate the bucket memory context. */
+MemoryContext get_current_bucket_context(void)
+{
+  return g_current_bucket_scan != NULL
+             ? g_current_bucket_scan->bucket_context
+             : NULL;
+}
+
+/* Used by common.c to check if an agg has redirected state. */
+bool aggref_shares_state(Aggref *aggref)
+{
+  if (g_current_bucket_scan == NULL)
+    return false;
+
+  BucketDescriptor *bucket_desc = g_current_bucket_scan->bucket_desc;
+  int num_atts = bucket_num_atts(bucket_desc);
+  for (int i = bucket_desc->num_labels; i < num_atts; i++)
+  {
+    BucketAttribute *att = &bucket_desc->attrs[i];
+    /* We use reference comparison to pinpoint exact position of aggregate. */
+    if (att->agg.aggref == aggref)
+      return i != att->agg.redirect_to;
+  }
+
+  /* This should not happen, but we can't guarantee that the Aggref was not copied somewhere. */
+  return false;
+}
 
 /*-------------------------------------------------------------------------
  * CustomExecMethods
  *-------------------------------------------------------------------------
  */
 
+/*
+ * Returns true if aggregates can share the same agg state.
+ * This is possible when args, initial state, transition, and merge functions are identical.
+ */
+static bool can_share_agg_state(BucketAttribute *agg1, BucketAttribute *agg2)
+{
+  const AnonAggFuncs *funcs1 = agg1->agg.funcs;
+  const AnonAggFuncs *funcs2 = agg2->agg.funcs;
+
+  return funcs1->create_state == funcs2->create_state &&
+         funcs1->transition == funcs2->transition &&
+         funcs1->merge == funcs2->merge &&
+         equal(agg1->agg.aggref->args, agg2->agg.aggref->args);
+}
+
 /*
  * Populates `bucket_desc` field with type metadata.
  */
@@ -131,16 +177,32 @@ static void init_bucket_descriptor(BucketScanState *bucket_state)
     {
       Aggref *aggref = castNode(Aggref, tle->expr);
       agg_funcs = find_agg_funcs(aggref->aggfnoid);
-      att->agg.fn_oid = aggref->aggfnoid;
+      att->agg.aggref = aggref;
       att->agg.funcs = agg_funcs;
       att->agg.args_desc = build_args_desc(aggref);
+      att->agg.redirect_to = i; /* Pointing to itself means state is not shared. */
       att->tag = agg_funcs != NULL ? BUCKET_ANON_AGG : BUCKET_REGULAR_AGG;
     }
 
     if (agg_funcs != NULL)
     {
       /* For anonymizing aggregators we describe finalized type. */
       agg_funcs->final_type(att->agg.args_desc, &att->final_type, &att->final_typmod, &att->final_collid);
+
+      /* Look back to check if there is a compatible agg which we can share state with. */
+      for (int j = plan_data->num_labels; j < i; j++)
+      {
+        BucketAttribute *other_att = &bucket_desc->attrs[j];
+        if (other_att->agg.funcs == NULL)
+          continue;
+
+        if (can_share_agg_state(att, other_att))
+        {
+          Assert(i != plan_data->low_count_index); /* low_count is always unique. */
+          att->agg.redirect_to = j;
+          break;
+        }
+      }
     }
     else
     {
@@ -182,13 +244,13 @@ static void bucket_begin_scan(CustomScanState *css, EState *estate, int eflags)
 
 static void fill_bucket_list(BucketScanState *bucket_state)
 {
-  MemoryContext old_bucket_context = g_current_bucket_context;
-  MemoryContext bucket_context = bucket_state->bucket_context;
+  BucketScanState *old_bucket_scan = g_current_bucket_scan;
 
   ExprContext *econtext = bucket_state->css.ss.ps.ps_ExprContext;
   MemoryContext per_tuple_memory = econtext->ecxt_per_tuple_memory;
   PlanState *outer_plan_state = outerPlanState(bucket_state);
 
+  MemoryContext bucket_context = bucket_state->bucket_context;
   BucketDescriptor *bucket_desc = bucket_state->bucket_desc;
   int num_atts = bucket_num_atts(bucket_desc);
   int low_count_index = bucket_desc->low_count_index;
@@ -201,7 +263,7 @@ static void fill_bucket_list(BucketScanState *bucket_state)
   {
     CHECK_FOR_INTERRUPTS();
 
-    g_current_bucket_context = bucket_context;
+    g_current_bucket_scan = bucket_state;
     TupleTableSlot *outer_slot = ExecProcNode(outer_plan_state);
 
     if (TupIsNull(outer_slot))
@@ -247,8 +309,8 @@ static void fill_bucket_list(BucketScanState *bucket_state)
   bucket_state->buckets = buckets;
   bucket_state->input_done = true;
 
-  /* Restore previous bucket context. */
-  g_current_bucket_context = old_bucket_context;
+  /* Restore previous bucket scan context. */
+  g_current_bucket_scan = old_bucket_scan;
 }
 
 static void run_hooks(BucketScanState *bucket_state)
@@ -289,9 +351,9 @@ static void finalize_bucket(Bucket *bucket, BucketDescriptor *bucket_desc, ExprC
     BucketAttribute *att = &bucket_desc->attrs[i];
     if (att->tag == BUCKET_ANON_AGG)
     {
-      AnonAggState *agg_state = (AnonAggState *)DatumGetPointer(bucket->values[i]);
+      int state_source_index = att->agg.redirect_to; /* If shared, points to some other non-NULL state. */
+      AnonAggState *agg_state = (AnonAggState *)DatumGetPointer(bucket->values[state_source_index]);
       Assert(agg_state != NULL);
-      Assert(agg_state->agg_funcs == att->agg.funcs);
       is_null[i] = false;
       values[i] = att->agg.funcs->finalize(agg_state, bucket, bucket_desc, &is_null[i]);
     }

src/aggregation/common.c

@@ -15,8 +15,9 @@
 #define PG_GET_AGG_STATE(index) ((AnonAggState *)PG_GETARG_POINTER(index))
 #define PG_RETURN_AGG_STATE(state) PG_RETURN_POINTER(state)
 
-/* Memory context of currently executing BucketScan node (if any). */
-extern MemoryContext g_current_bucket_context;
+/* Functions declared in bucket_scan.c. Depend on global state and should not be public API. */
+extern MemoryContext get_current_bucket_context(void);
+extern bool aggref_shares_state(Aggref *aggref);
 
 PG_FUNCTION_INFO_V1(anon_agg_state_input);
 PG_FUNCTION_INFO_V1(anon_agg_state_output);
@@ -95,11 +96,16 @@ void merge_bucket(Bucket *destination, Bucket *source, BucketDescriptor *bucket_
   for (int i = bucket_desc->num_labels; i < num_atts; i++)
   {
     BucketAttribute *att = &bucket_desc->attrs[i];
-    if (att->tag == BUCKET_ANON_AGG)
+    if (att->tag == BUCKET_ANON_AGG &&
+        i == att->agg.redirect_to /* Shared states need to be merged only once. */)
     {
       Assert(!source->is_null[i]);
       Assert(!destination->is_null[i]);
-      att->agg.funcs->merge((AnonAggState *)destination->values[i], (AnonAggState *)source->values[i]);
+      AnonAggState *dst_state = (AnonAggState *)destination->values[i];
+      AnonAggState *src_state = (AnonAggState *)source->values[i];
+      Assert(dst_state != AGG_STATE_REDIRECTED);
+      Assert(src_state != AGG_STATE_REDIRECTED);
+      att->agg.funcs->merge(dst_state, src_state);
     }
   }
 }
@@ -113,10 +119,16 @@ static AnonAggState *get_agg_state(PG_FUNCTION_ARGS)
   if (AggCheckCallContext(fcinfo, &bucket_context) != AGG_CONTEXT_AGGREGATE)
     FAILWITH("Aggregate called in non-aggregate context");
 
-  if (g_current_bucket_context != NULL)
-    bucket_context = g_current_bucket_context;
-
   Aggref *aggref = AggGetAggref(fcinfo);
+
+  if (get_current_bucket_context() != NULL)
+  {
+    if (aggref_shares_state(aggref))
+      return AGG_STATE_REDIRECTED;
+
+    bucket_context = get_current_bucket_context();
+  }
+
   const AnonAggFuncs *agg_funcs = find_agg_funcs(aggref->aggfnoid);
 
   if (unlikely(agg_funcs == NULL))
@@ -134,20 +146,23 @@ Datum anon_agg_state_input(PG_FUNCTION_ARGS)
 Datum anon_agg_state_output(PG_FUNCTION_ARGS)
 {
   AnonAggState *state = PG_GET_AGG_STATE(0);
+  Assert(state != AGG_STATE_REDIRECTED); /* Won't happen outside of a BucketScan context. */
   const char *str = state->agg_funcs->explain(state);
   PG_RETURN_CSTRING(str);
 }
 
 Datum anon_agg_state_transfn(PG_FUNCTION_ARGS)
 {
   AnonAggState *state = get_agg_state(fcinfo);
-  state->agg_funcs->transition(state, PG_NARGS(), fcinfo->args);
+  /* AGG_STATE_REDIRECTED means the owning aggregator will handle transitions. */
+  if (state != AGG_STATE_REDIRECTED)
+    state->agg_funcs->transition(state, PG_NARGS(), fcinfo->args);
   PG_RETURN_AGG_STATE(state);
 }
 
 /*
  * This finalfunc is a dummy version which does nothing.
- * It only ensures that state is not null for empty buckets.
+ * It only ensures that state is initialized for empty buckets.
  */
 Datum anon_agg_state_finalfn(PG_FUNCTION_ARGS)
 {

src/aggregation/star_bucket.c

@@ -48,10 +48,12 @@ Bucket *star_bucket_hook(List *buckets, BucketDescriptor *bucket_desc)
     BucketAttribute *att = &bucket_desc->attrs[i];
     if (att->tag == BUCKET_ANON_AGG)
       /* Create an empty anon agg state and merge buckets into it. */
-      star_bucket->values[i] = PointerGetDatum(create_anon_agg_state(att->agg.funcs, bucket_context, att->agg.args_desc));
+      star_bucket->values[i] = PointerGetDatum(i != att->agg.redirect_to
+                                                   ? AGG_STATE_REDIRECTED
+                                                   : create_anon_agg_state(att->agg.funcs, bucket_context, att->agg.args_desc));
     else if (att->tag == BUCKET_LABEL)
       set_text_label(star_bucket, i, att->final_type, bucket_context);
-    else if (att->agg.fn_oid == g_oid_cache.is_suppress_bin)
+    else if (att->agg.aggref->aggfnoid == g_oid_cache.is_suppress_bin)
       star_bucket->values[i] = BoolGetDatum(true);
     else
       star_bucket->is_null[i] = true;