Return allocated budget details from DPEngine.done()

pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/core/BUILD.bazel

@@ -115,6 +115,7 @@ kt_jvm_library(
         ":framework_collections",
         "//main/com/google/privacy/differentialprivacy/pipelinedp4j/core/budget:allocated_budget",
         "//main/com/google/privacy/differentialprivacy/pipelinedp4j/core/budget:budget_accountant",
+        "//main/com/google/privacy/differentialprivacy/pipelinedp4j/core/budget:budget_allocation_details",
         "//main/com/google/privacy/differentialprivacy/pipelinedp4j/core/budget:budget_spec",
         "//main/com/google/privacy/differentialprivacy/pipelinedp4j/dplibrary:noise_factories",
         "//main/com/google/privacy/differentialprivacy/pipelinedp4j/dplibrary:pre_aggregation_partition_selection_factory",

pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/core/DpEngine.kt

@@ -16,6 +16,7 @@
 
 package com.google.privacy.differentialprivacy.pipelinedp4j.core
 
+import com.google.errorprone.annotations.CanIgnoreReturnValue
 import com.google.privacy.differentialprivacy.Noise
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.MetricType.COUNT
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.MetricType.MEAN
@@ -35,6 +36,7 @@ import com.google.privacy.differentialprivacy.pipelinedp4j.core.budget.BudgetAcc
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.budget.BudgetAccountantFactory
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.budget.BudgetAccountingStrategy
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.budget.BudgetAccountingStrategy.NAIVE
+import com.google.privacy.differentialprivacy.pipelinedp4j.core.budget.BudgetAllocationDetails
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.budget.BudgetPerOpSpec
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.budget.BudgetRequest
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.budget.RelativeBudgetPerOpSpec
@@ -204,13 +206,26 @@ internal constructor(
   }
 
   /**
-   * Allocates privacy budgets to the metrics whose computation has been requested by calling
-   * [aggregate]. This method must be called once per [DpEngine] instance.
+   * Allocates privacy budgets to privacy-preserving operations in [aggregate] and
+   * [selectPartitions] calls.
+   *
+   * Privacy-preserving operations are various aggregation metrics, like COUNT or SUM, and partition
+   * selection. There might be multiple privacy-preserving operations in a single [DpEngine]
+   * instance.
+   *
+   * This method must be called once per [DpEngine] instance.
+   *
+   * @return a list of [BudgetAllocationDetails] for each privacy-preserving operation. This reports
+   *   the actual budgets used during computation, which may include budgets for operations that
+   *   were not directly requested (e.g., for a MEAN aggregation, budget details for both SUM and
+   *   COUNT will be returned).
+   * @throws IllegalStateException if [done] has already been called on this instance.
    */
-  fun done() {
+  @CanIgnoreReturnValue
+  fun done(): List<BudgetAllocationDetails> {
     throwIfDoneWasCalled()
     doneCalled = true
-    budgetAccountant.allocateBudgets()
+    return budgetAccountant.allocateBudgets()
   }
 
   private fun throwIfDoneWasCalled() {

pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/core/budget/BUILD.bazel

@@ -25,6 +25,11 @@ package(
     ],
 )
 
+kt_jvm_library(
+    name = "budget_allocation_details",
+    srcs = ["BudgetAllocationDetails.kt"],
+)
+
 kt_jvm_library(
     name = "budget_spec",
     srcs = ["BudgetSpec.kt"],
@@ -43,6 +48,8 @@ kt_jvm_library(
     srcs = ["BudgetAccountant.kt"],
     deps = [
         ":allocated_budget",
+        ":budget_allocation_details",
         ":budget_spec",
+        "@maven//:com_google_errorprone_error_prone_annotations",
     ],
 )

pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/core/budget/BudgetAccountant.kt

@@ -16,6 +16,7 @@
 
 package com.google.privacy.differentialprivacy.pipelinedp4j.core.budget
 
+import com.google.errorprone.annotations.CanIgnoreReturnValue
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.budget.BudgetAccountingStrategy.NAIVE
 import java.lang.IllegalArgumentException
 import java.lang.IllegalStateException
@@ -40,12 +41,15 @@ interface BudgetAccountant {
   fun requestBudget(budgetRequest: BudgetRequest): AllocatedBudget
 
   /**
-   * Allocates budgets to all previously recorded [BudgetRequest]s. This method should only be
-   * called once.
+   * Allocates budgets to all previously recorded [BudgetRequest]s.
    *
+   * This method should only be called once.
+   *
+   * @return the allocation details for each requested budget (e.g. very useful to understand how
+   *   much budget was allocated to relative budget requests).
    * @throws IllegalStateException if budgets have already been allocated.
    */
-  fun allocateBudgets()
+  @CanIgnoreReturnValue fun allocateBudgets(): List<BudgetAllocationDetails>
 }
 
 /**
@@ -100,7 +104,8 @@ class NaiveBudgetAccountant(private val totalBudget: TotalBudget) : BudgetAccoun
     return allocatedBudget
   }
 
-  override fun allocateBudgets() {
+  @CanIgnoreReturnValue
+  override fun allocateBudgets(): List<BudgetAllocationDetails> {
     if (budgetsAllocated) {
       throw IllegalStateException("Budgets have already been allocated.")
     }
@@ -119,8 +124,11 @@ class NaiveBudgetAccountant(private val totalBudget: TotalBudget) : BudgetAccoun
     checkEnoughAbsoluteBudget(totalRequestedEpsilon, totalRequestedDelta)
     checkEnoughRelativeBudget(remainingEpsilon, remainingDelta)
 
-    allocateAbsoluteBudgets()
-    allocateRelativeBudgets(remainingEpsilon, remainingDelta)
+    initializeAbsoluteBudgets()
+    initializeRelativeBudgets(remainingEpsilon, remainingDelta)
+
+    return absoluteBudgets.map { it.toBudgetAllocationDetails() } +
+      relativeBudgets.map { it.toBudgetAllocationDetails() }
   }
 
   private fun checkEnoughAbsoluteBudget(requestedEpsilon: Double, requestedDelta: Double) {
@@ -147,7 +155,7 @@ class NaiveBudgetAccountant(private val totalBudget: TotalBudget) : BudgetAccoun
     return Math.abs(diff) > remaining / FLOATING_POINT_ARITHMETICS_TOLERANCE
   }
 
-  private fun allocateAbsoluteBudgets() {
+  private fun initializeAbsoluteBudgets() {
     for (requestedAndAllocated in absoluteBudgets) {
       val budgetSpec = requestedAndAllocated.requested.budgetSpec as AbsoluteBudgetPerOpSpec
       requestedAndAllocated.allocated.initialize(
@@ -157,7 +165,7 @@ class NaiveBudgetAccountant(private val totalBudget: TotalBudget) : BudgetAccoun
     }
   }
 
-  private fun allocateRelativeBudgets(remainingEpsilon: Double, remainingDelta: Double) {
+  private fun initializeRelativeBudgets(remainingEpsilon: Double, remainingDelta: Double) {
     var totalEpsilonWeight = 0.0
     var totalDeltaWeight = 0.0
     for (requestedAndAllocated in relativeBudgets) {
@@ -200,11 +208,13 @@ class NaiveBudgetAccountant(private val totalBudget: TotalBudget) : BudgetAccoun
     }
   }
 
-  private fun relativeEpsilonRequested(): Boolean =
-    relativeBudgets.any { it.requested.mechanism.usesEpsilon }
+  private fun relativeEpsilonRequested(): Boolean = relativeBudgets.any {
+    it.requested.mechanism.usesEpsilon
+  }
 
-  private fun relativeDeltaRequested(): Boolean =
-    relativeBudgets.any { it.requested.mechanism.usesDelta }
+  private fun relativeDeltaRequested(): Boolean = relativeBudgets.any {
+    it.requested.mechanism.usesDelta
+  }
 }
 
 /**
@@ -243,7 +253,23 @@ enum class AccountedMechanism(val usesEpsilon: Boolean, val usesDelta: Boolean)
 internal data class RequestedAndAllocatedBudget(
   val requested: BudgetRequest,
   val allocated: AllocatedBudget,
-)
+) {
+  /** Converts a [RequestedAndAllocatedBudget] to a [BudgetAllocationDetails]. */
+  fun toBudgetAllocationDetails(): BudgetAllocationDetails {
+    val epsilon = allocated.epsilon()
+    val delta = allocated.delta()
+    return when (requested.mechanism) {
+      AccountedMechanism.GAUSSIAN_NOISE ->
+        BudgetAllocationDetails.GaussianAggregationAllocation(epsilon, delta)
+      AccountedMechanism.LAPLACE_NOISE ->
+        BudgetAllocationDetails.LaplaceAggregationAllocation(epsilon)
+      AccountedMechanism.PREAGGREGATED_PARTITION_SELECTION ->
+        BudgetAllocationDetails.PreaggregatedPartitionSelectionAllocation(epsilon, delta)
+      AccountedMechanism.POSTAGGREGATED_PARTITION_SELECTION ->
+        BudgetAllocationDetails.PostaggregatedPartitionSelectionAllocation(delta)
+    }
+  }
+}
 
 enum class BudgetAccountingStrategy {
   NAIVE

pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/core/budget/BudgetAllocationDetails.kt

@@ -0,0 +1,46 @@
+package com.google.privacy.differentialprivacy.pipelinedp4j.core.budget
+
+/**
+ * Sealed class representing the details of a budget allocation for different differential privacy
+ * mechanisms.
+ *
+ * Each subclass contains only the parameters relevant to that mechanism.
+ *
+ * This class is used to pass budget allocation details from [DpEngine] to API implementations
+ * (e.g., BeamApi), but it is not intended for direct use by end-users of the library.
+ *
+ * Extend this class if you need to propagate more details about the budget allocation from DPEngine
+ * to the backend-specific API implementations in the API package (e.g. BeamApi, etc.).
+ */
+sealed class BudgetAllocationDetails {
+  /**
+   * Budget allocation details for Gaussian mechanism used for aggregation.
+   *
+   * Uses both epsilon and delta.
+   */
+  data class GaussianAggregationAllocation(val epsilon: Double, val delta: Double) :
+    BudgetAllocationDetails()
+
+  /**
+   * Budget allocation details for Laplace mechanism used for aggregation.
+   *
+   * Uses only epsilon.
+   */
+  data class LaplaceAggregationAllocation(val epsilon: Double) : BudgetAllocationDetails()
+
+  /**
+   * Budget allocation details for pre-aggregated partition selection.
+   *
+   * Uses both epsilon and delta.
+   */
+  data class PreaggregatedPartitionSelectionAllocation(val epsilon: Double, val delta: Double) :
+    BudgetAllocationDetails()
+
+  /**
+   * Budget allocation details for post-aggregated partition selection.
+   *
+   * Only uses delta.
+   */
+  data class PostaggregatedPartitionSelectionAllocation(val delta: Double) :
+    BudgetAllocationDetails()
+}

pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/core/BUILD.bazel

@@ -58,6 +58,7 @@ kt_jvm_test(
         "//main/com/google/privacy/differentialprivacy/pipelinedp4j/core:framework_collections",
         "//main/com/google/privacy/differentialprivacy/pipelinedp4j/core/budget:allocated_budget",
         "//main/com/google/privacy/differentialprivacy/pipelinedp4j/core/budget:budget_accountant",
+        "//main/com/google/privacy/differentialprivacy/pipelinedp4j/core/budget:budget_allocation_details",
         "//main/com/google/privacy/differentialprivacy/pipelinedp4j/core/budget:budget_spec",
         "//main/com/google/privacy/differentialprivacy/pipelinedp4j/dplibrary:noise_factories",
         "//main/com/google/privacy/differentialprivacy/pipelinedp4j/dplibrary:pre_aggregation_partition_selection_factory",

pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/core/DpEngineTest.kt

@@ -34,6 +34,7 @@ import com.google.privacy.differentialprivacy.pipelinedp4j.core.NoiseKind.GAUSSI
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.NoiseKind.LAPLACE
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.budget.AbsoluteBudgetPerOpSpec
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.budget.BudgetAccountingStrategy.NAIVE
+import com.google.privacy.differentialprivacy.pipelinedp4j.core.budget.BudgetAllocationDetails
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.budget.BudgetPerOpSpec
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.budget.RelativeBudgetPerOpSpec
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.budget.TotalBudget
@@ -96,6 +97,36 @@ class DpEngineTest {
     assertThat(e).hasMessageThat().contains("done() has already been called")
   }
 
+  @Test
+  fun done_returnsBudgetAllocationDetails() {
+    val dpEngine = DpEngine.createForTesting(LOCAL_EF, LARGE_BUDGET_SPEC, ZeroNoiseFactory())
+    val params =
+      AggregationParams(
+        metrics = ImmutableList.of(MetricDefinition(COUNT, AbsoluteBudgetPerOpSpec(1.0, 1e-5))),
+        noiseKind = GAUSSIAN,
+        maxPartitionsContributed = 1,
+        maxContributionsPerPartition = 1,
+        partitionSelectionBudget = AbsoluteBudgetPerOpSpec(2.0, 2e-5),
+      )
+
+    val unused =
+      dpEngine.aggregate(
+        LocalCollection(sequenceOf(TestDataRow("Alice", "US", 1.0))),
+        params,
+        testDataExtractors,
+      )
+    val details = dpEngine.done()
+
+    assertThat(details)
+      .containsExactly(
+        BudgetAllocationDetails.GaussianAggregationAllocation(epsilon = 1.0, delta = 1e-5),
+        BudgetAllocationDetails.PreaggregatedPartitionSelectionAllocation(
+          epsilon = 2.0,
+          delta = 2e-5,
+        ),
+      )
+  }
+
   @Test
   fun aggregate_incorrectAggregateParams_throws() {
     val e =

pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/core/budget/BUILD.bazel

@@ -27,6 +27,7 @@ kt_jvm_test(
     test_class = "com.google.privacy.differentialprivacy.pipelinedp4j.core.budget.BudgetTests",
     deps = [
         "//main/com/google/privacy/differentialprivacy/pipelinedp4j/core/budget:budget_accountant",
+        "//main/com/google/privacy/differentialprivacy/pipelinedp4j/core/budget:budget_allocation_details",
         "//main/com/google/privacy/differentialprivacy/pipelinedp4j/core/budget:budget_spec",
         "@maven//:com_google_testparameterinjector_test_parameter_injector",
         "@maven//:com_google_truth_truth",