Refactor contribution sampling to work with multiple features.

PiperOrigin-RevId: 811408440

Author: alexojica

pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/api/AggregationSpec.kt

@@ -148,44 +148,6 @@ internal fun AggregationSpec.getFeatureId(): String {
   }
 }
 
-internal fun List<AggregationSpec>.metrics(): List<MetricDefinition> = buildList {
-  for (aggregation in this@metrics) {
-    when (aggregation) {
-      // Count and PrivacyIdCount do not aggregate any specific value, therefore they are handled
-      // differently.
-      is PrivacyIdCount ->
-        add(
-          MetricDefinition(
-            MetricType.PRIVACY_ID_COUNT,
-            aggregation.budget?.toInternalBudgetPerOpSpec(),
-          )
-        )
-      is Count ->
-        add(MetricDefinition(MetricType.COUNT, aggregation.budget?.toInternalBudgetPerOpSpec()))
-      is ValueAggregations<*> -> {
-        for (valueAggregationSpec in aggregation.valueAggregationSpecs) {
-          add(
-            MetricDefinition(
-              valueAggregationSpec.metricType,
-              valueAggregationSpec.budget?.toInternalBudgetPerOpSpec(),
-            )
-          )
-        }
-      }
-      is VectorAggregations<*> -> {
-        for (vectorAggregationSpec in aggregation.vectorAggregationSpecs) {
-          add(
-            MetricDefinition(
-              vectorAggregationSpec.metricType,
-              vectorAggregationSpec.budget?.toInternalBudgetPerOpSpec(),
-            )
-          )
-        }
-      }
-    }
-  }
-}
-
 internal fun List<AggregationSpec>.outputColumnNamesWithMetricTypes():
   List<Pair<String, MetricType>> = buildList {
   for (aggregation in this@outputColumnNamesWithMetricTypes) {
@@ -227,3 +189,41 @@ internal fun List<AggregationSpec>.outputColumnNameToFeatureIdMap(): Map<String,
 
 internal fun List<AggregationSpec>.outputColumnNames(): List<String> =
   outputColumnNamesWithMetricTypes().map { it.first }
+
+internal fun AggregationSpec.toMetricDefinition():
+  com.google.privacy.differentialprivacy.pipelinedp4j.core.MetricDefinition {
+  val metricType =
+    when (this) {
+      is Count -> MetricType.COUNT
+      is PrivacyIdCount -> MetricType.PRIVACY_ID_COUNT
+      else ->
+        throw IllegalArgumentException("Unsupported AggregationSpec type for top-level metrics")
+    }
+  val budget =
+    when (this) {
+      is Count -> this.budget
+      is PrivacyIdCount -> this.budget
+      else ->
+        throw IllegalArgumentException("Unsupported AggregationSpec type for top-level metrics")
+    }
+  return com.google.privacy.differentialprivacy.pipelinedp4j.core.MetricDefinition(
+    metricType,
+    budget?.toInternalBudgetPerOpSpec(),
+  )
+}
+
+internal fun ValueAggregationSpec.toMetricDefinition():
+  com.google.privacy.differentialprivacy.pipelinedp4j.core.MetricDefinition {
+  return com.google.privacy.differentialprivacy.pipelinedp4j.core.MetricDefinition(
+    this.metricType,
+    this.budget?.toInternalBudgetPerOpSpec(),
+  )
+}
+
+internal fun VectorAggregationSpec.toMetricDefinition():
+  com.google.privacy.differentialprivacy.pipelinedp4j.core.MetricDefinition {
+  return com.google.privacy.differentialprivacy.pipelinedp4j.core.MetricDefinition(
+    this.metricType,
+    this.budget?.toInternalBudgetPerOpSpec(),
+  )
+}

pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/api/Query.kt

@@ -23,15 +23,15 @@ import com.google.privacy.differentialprivacy.pipelinedp4j.core.DpEngine
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.DpEngineBudgetSpec
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.Encoder
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.EncoderFactory
+import com.google.privacy.differentialprivacy.pipelinedp4j.core.FeatureSpec
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.FeatureValuesExtractor
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.FrameworkCollection
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.FrameworkTable
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.MetricType
+import com.google.privacy.differentialprivacy.pipelinedp4j.core.ScalarFeatureSpec
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.SelectPartitionsParams
 import com.google.privacy.differentialprivacy.pipelinedp4j.proto.DpAggregates
-import com.google.privacy.differentialprivacy.pipelinedp4j.proto.PerFeature
 import com.google.privacy.differentialprivacy.pipelinedp4j.proto.copy
-import com.google.privacy.differentialprivacy.pipelinedp4j.proto.dpAggregates
 import com.google.privacy.differentialprivacy.pipelinedp4j.proto.perFeature
 
 sealed interface Query<ReturnT> {
@@ -142,21 +142,6 @@ protected constructor(
         valueAndVectorAggs.map { it.getFeatureId() }
       }
     return aggResults
-      .zip(featureIdPerRun)
-      .map { (table, featureId) ->
-        table.mapValues("TagWithFeatureId", encoderFactory.protos(DpAggregates::class)) { _, agg ->
-          if (featureId == null) {
-            agg
-          } else {
-            val perFeature = constructPerFeature(agg, featureId)
-            dpAggregates {
-              count = agg.count
-              privacyIdCount = agg.privacyIdCount
-              this.perFeature += perFeature
-            }
-          }
-        }
-      }
       .reduce {
         acc: FrameworkTable<GroupKeysT, DpAggregates>,
         table: FrameworkTable<GroupKeysT, DpAggregates> ->
@@ -496,41 +481,59 @@ protected constructor(
   ): AggregationParams {
     val valueContributionBounds = valueAggregations?.contributionBounds
     val vectorContributionBounds = vectorAggregations?.vectorContributionBounds
+
+    val topLevelMetrics =
+      aggregationSpecs
+        .filter { it is Count || it is PrivacyIdCount }
+        .map { it.toMetricDefinition() }
+    val features =
+      mutableListOf<com.google.privacy.differentialprivacy.pipelinedp4j.core.FeatureSpec>()
+    if (valueAggregations != null) {
+      features.add(
+        com.google.privacy.differentialprivacy.pipelinedp4j.core.ScalarFeatureSpec(
+          featureId = valueAggregations.getFeatureId(),
+          metrics =
+            ImmutableList.copyOf(
+              valueAggregations.valueAggregationSpecs.map { it.toMetricDefinition() }
+            ),
+          minValue = valueContributionBounds?.valueBounds?.minValue,
+          maxValue = valueContributionBounds?.valueBounds?.maxValue,
+          minTotalValue = valueContributionBounds?.totalValueBounds?.minValue,
+          maxTotalValue = valueContributionBounds?.totalValueBounds?.maxValue,
+        )
+      )
+    }
+    if (vectorAggregations != null) {
+      checkNotNull(vectorContributionBounds)
+      features.add(
+        com.google.privacy.differentialprivacy.pipelinedp4j.core.VectorFeatureSpec(
+          featureId = vectorAggregations.getFeatureId(),
+          metrics =
+            ImmutableList.copyOf(
+              vectorAggregations.vectorAggregationSpecs.map { it.toMetricDefinition() }
+            ),
+          vectorSize =
+            checkNotNull(vectorAggregations.vectorSize) {
+              "vectorSize must be set for vector aggregations"
+            },
+          normKind = vectorContributionBounds.maxVectorTotalNorm.normKind.toInternalNormKind(),
+          vectorMaxTotalNorm = vectorContributionBounds.maxVectorTotalNorm.value,
+        )
+      )
+    }
+
     return AggregationParams(
-      metrics = ImmutableList.copyOf(aggregationSpecs.metrics()),
+      metrics = ImmutableList.copyOf(topLevelMetrics),
+      features = ImmutableList.copyOf(features),
       noiseKind =
         checkNotNull(noiseKind) { "noiseKind cannot be null if there are aggregations." }
           .toInternalNoiseKind(),
       maxPartitionsContributed = contributionBoundingLevel.getMaxPartitionsContributed(),
       maxContributionsPerPartition = contributionBoundingLevel.getMaxContributionsPerPartition(),
-      minValue = valueContributionBounds?.valueBounds?.minValue,
-      maxValue = valueContributionBounds?.valueBounds?.maxValue,
-      minTotalValue = valueContributionBounds?.totalValueBounds?.minValue,
-      maxTotalValue = valueContributionBounds?.totalValueBounds?.maxValue,
-      vectorNormKind = vectorContributionBounds?.maxVectorTotalNorm?.normKind?.toInternalNormKind(),
-      vectorMaxTotalNorm = vectorContributionBounds?.maxVectorTotalNorm?.value,
-      vectorSize = vectorAggregations?.vectorSize,
       partitionSelectionBudget = groupsType.getBudget()?.toInternalBudgetPerOpSpec(),
       preThreshold = groupsType.getPreThreshold(),
       contributionBoundingLevel = contributionBoundingLevel.toInternalContributionBoundingLevel(),
       partitionsBalance = groupByAdditionalParameters.groupsBalance.toPartitionsBalance(),
     )
   }
-
-  companion object {
-    private fun constructPerFeature(dpAggregates: DpAggregates, featureId: String): PerFeature {
-      return perFeature {
-        this.featureId = featureId
-        sum = dpAggregates.sum
-        mean = dpAggregates.mean
-        variance = dpAggregates.variance
-        if (dpAggregates.quantilesList.isNotEmpty()) {
-          quantiles += dpAggregates.quantilesList
-        }
-        if (dpAggregates.vectorSumList.isNotEmpty()) {
-          vectorSum += dpAggregates.vectorSumList
-        }
-      }
-    }
-  }
 }

pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/api/QueryPerGroupResult.kt

@@ -115,40 +115,23 @@ internal constructor(
             MetricType.PRIVACY_ID_COUNT -> put(outputColumnName, dpAggregates.privacyIdCount)
             MetricType.COUNT -> put(outputColumnName, dpAggregates.count)
             MetricType.SUM -> {
-              if (dpAggregates.perFeatureList.isNotEmpty()) {
-                val featureId = columnNameToFeatureIdMap[outputColumnName]!!
-                put(outputColumnName, featuresMap[featureId]!!.sum)
-              } else {
-                put(outputColumnName, dpAggregates.sum)
-              }
+              val featureId = columnNameToFeatureIdMap[outputColumnName]!!
+              put(outputColumnName, featuresMap[featureId]!!.sum)
             }
             MetricType.VECTOR_SUM -> {} // not processed in this function.
             MetricType.MEAN -> {
-              if (dpAggregates.perFeatureList.isNotEmpty()) {
-                val featureId = columnNameToFeatureIdMap[outputColumnName]!!
-                put(outputColumnName, featuresMap[featureId]!!.mean)
-              } else {
-                put(outputColumnName, dpAggregates.mean)
-              }
+              val featureId = columnNameToFeatureIdMap[outputColumnName]!!
+              put(outputColumnName, featuresMap[featureId]!!.mean)
             }
             MetricType.VARIANCE -> {
-              if (dpAggregates.perFeatureList.isNotEmpty()) {
-                val featureId = columnNameToFeatureIdMap[outputColumnName]!!
-                put(outputColumnName, featuresMap[featureId]!!.variance)
-              } else {
-                put(outputColumnName, dpAggregates.variance)
-              }
+              val featureId = columnNameToFeatureIdMap[outputColumnName]!!
+              put(outputColumnName, featuresMap[featureId]!!.variance)
             }
             is MetricType.QUANTILES -> {
               // TODO: consider creating a data class or resuing copy of
               // DpAggregates proto and not allowing outputColumnName.
-              val quantilesList =
-                if (dpAggregates.perFeatureList.isNotEmpty()) {
-                  val featureId = columnNameToFeatureIdMap[outputColumnName]!!
-                  featuresMap[featureId]!!.quantilesList
-                } else {
-                  dpAggregates.quantilesList
-                }
+              val featureId = columnNameToFeatureIdMap[outputColumnName]!!
+              val quantilesList = featuresMap[featureId]!!.quantilesList
               for ((rank, value) in metricType.sortedRanks.zip(quantilesList)) {
                 put(outputColumnName.withRank(rank), value)
               }
@@ -171,12 +154,8 @@ internal constructor(
             MetricType.COUNT -> {} // not processed in this function.
             MetricType.SUM -> {} // not processed in this function.
             MetricType.VECTOR_SUM -> {
-              if (dpAggregates.perFeatureList.isNotEmpty()) {
-                val featureId = colNameToFeatureIdMap[outputColumnName]!!
-                put(outputColumnName, featuresMap[featureId]!!.vectorSumList)
-              } else {
-                put(outputColumnName, dpAggregates.vectorSumList)
-              }
+              val featureId = colNameToFeatureIdMap[outputColumnName]!!
+              put(outputColumnName, featuresMap[featureId]!!.vectorSumList)
             }
             MetricType.MEAN -> {} // not processed in this function.
             MetricType.VARIANCE -> {} // not processed in this function.