Bugfixes & improvement for the new privacy budget API in Privacy on Beam

Privacy on Beam:
- Add a check for disallowing beam universal types in pbeam.ParDo
- Add check for partition selection budget being set when the pre-threshold is set
- Rename test mode constants in order to be more readable when called outside the package
- Refer to the new privacy budget API in documentation and codelab
- fix bugs with the new privacy budget API in DistinctPerKey

Python Accounting:
- Increment patch version of DP accounting library for PyPi release, reflecting changes in requirements

Change-Id: I4cc7eb22625cae4d71fd5dd1098197dfb7b32ea0
GitOrigin-RevId: 9ca4d7d9aa388880d7663979fd0c9acfaf8fb4a1

privacy-on-beam/codelab/BUILD.bazel

@@ -37,5 +37,6 @@ go_library(
         "@com_github_apache_beam_sdks_v2//go/pkg/beam:go_default_library",
         "@com_github_apache_beam_sdks_v2//go/pkg/beam/register:go_default_library",
         "@com_github_apache_beam_sdks_v2//go/pkg/beam/transforms/stats:go_default_library",
+        "@com_github_golang_glog//:go_default_library",
     ],
 )

privacy-on-beam/codelab/count.go

@@ -20,6 +20,7 @@ package codelab
 import (
 	"math"
 
+	log "github.com/golang/glog"
 	"github.com/google/differential-privacy/privacy-on-beam/v2/pbeam"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/register"
@@ -54,7 +55,14 @@ func extractVisitHourFn(v Visit) int {
 func PrivateCountVisitsPerHour(s beam.Scope, col beam.PCollection) beam.PCollection {
 	s = s.Scope("PrivateCountVisitsPerHour")
 	// Create a Privacy Spec and convert col into a PrivatePCollection
-	spec := pbeam.NewPrivacySpec(epsilon, delta)
+	spec, err := pbeam.NewPrivacySpecTemp(pbeam.PrivacySpecParams{
+		AggregationEpsilon:        epsilon / 2,
+		PartitionSelectionEpsilon: epsilon / 2,
+		PartitionSelectionDelta:   delta,
+	})
+	if err != nil {
+		log.Fatalf("Couldn't create a PrivacySpec: %v", err)
+	}
 	pCol := pbeam.MakePrivateFromStruct(s, col, spec, "VisitorID")
 
 	visitHours := pbeam.ParDo(s, extractVisitHourFn, pCol)

privacy-on-beam/codelab/mean.go

@@ -17,6 +17,7 @@
 package codelab
 
 import (
+	log "github.com/golang/glog"
 	"github.com/google/differential-privacy/privacy-on-beam/v2/pbeam"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/register"
@@ -47,7 +48,10 @@ func extractVisitHourAndTimeSpentFn(v Visit) (int, int) {
 func PrivateMeanTimeSpent(s beam.Scope, col beam.PCollection) beam.PCollection {
 	s = s.Scope("PrivateMeanTimeSpent")
 	// Create a Privacy Spec and convert col into a PrivatePCollection.
-	spec := pbeam.NewPrivacySpec(epsilon /* delta */, 0)
+	spec, err := pbeam.NewPrivacySpecTemp(pbeam.PrivacySpecParams{AggregationEpsilon: epsilon})
+	if err != nil {
+		log.Fatalf("Couldn't create a PrivacySpec: %v", err)
+	}
 	pCol := pbeam.MakePrivateFromStruct(s, col, spec, "VisitorID")
 
 	// Create a PCollection of output partitions, i.e. restaurant's work hours (from 9 am till 9pm (exclusive)).

privacy-on-beam/codelab/multiple.go

@@ -17,6 +17,7 @@
 package codelab
 
 import (
+	log "github.com/golang/glog"
 	"github.com/google/differential-privacy/privacy-on-beam/v2/pbeam"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam"
 )
@@ -28,24 +29,26 @@ import (
 func ComputeCountMeanSum(s beam.Scope, col beam.PCollection) (visitsPerHour, meanTimeSpent, revenues beam.PCollection) {
 	s = s.Scope("ComputeCountMeanSum")
 	// Create a Privacy Spec and convert col into a PrivatePCollection.
-	spec := pbeam.NewPrivacySpec(epsilon /* delta */, 0) // Shared by count, mean and sum.
+	spec, err := pbeam.NewPrivacySpecTemp(pbeam.PrivacySpecParams{AggregationEpsilon: epsilon}) // Shared by count, mean and sum.
+	if err != nil {
+		log.Fatalf("Couldn't create a PrivacySpec: %v", err)
+	}
 	pCol := pbeam.MakePrivateFromStruct(s, col, spec, "VisitorID")
 
 	// Create a PCollection of output partitions, i.e. restaurant's work hours (from 9 am till 9pm (exclusive)).
 	hours := beam.CreateList(s, [12]int{9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20})
 
 	visitHours := pbeam.ParDo(s, extractVisitHourFn, pCol)
 	visitsPerHour = pbeam.Count(s, visitHours, pbeam.CountParams{
-		Epsilon:                  epsilon / 3,
-		Delta:                    0,
+		AggregationEpsilon:       epsilon / 3,
 		MaxPartitionsContributed: 1,     // Visitors can visit the restaurant once (one hour) a day
 		MaxValue:                 1,     // Visitors can visit the restaurant once within an hour
 		PublicPartitions:         hours, // Visitors only visit during work hours
 	})
 
 	hourToTimeSpent := pbeam.ParDo(s, extractVisitHourAndTimeSpentFn, pCol)
 	meanTimeSpent = pbeam.MeanPerKey(s, hourToTimeSpent, pbeam.MeanParams{
-		Epsilon:                      epsilon / 3,
+		AggregationEpsilon:           epsilon / 3,
 		Delta:                        0,
 		MaxPartitionsContributed:     1,     // Visitors can visit the restaurant once (one hour) a day
 		MaxContributionsPerPartition: 1,     // Visitors can visit the restaurant once within an hour
@@ -56,8 +59,7 @@ func ComputeCountMeanSum(s beam.Scope, col beam.PCollection) (visitsPerHour, mea
 
 	hourToMoneySpent := pbeam.ParDo(s, extractVisitHourAndMoneySpentFn, pCol)
 	revenues = pbeam.SumPerKey(s, hourToMoneySpent, pbeam.SumParams{
-		Epsilon:                  epsilon / 3,
-		Delta:                    0,
+		AggregationEpsilon:       epsilon / 3,
 		MaxPartitionsContributed: 1,     // Visitors can visit the restaurant once (one hour) a day
 		MinValue:                 0,     // Minimum money spent per user (in euros)
 		MaxValue:                 40,    // Maximum money spent per user (in euros)

privacy-on-beam/codelab/public_partitions.go

@@ -17,6 +17,7 @@
 package codelab
 
 import (
+	log "github.com/golang/glog"
 	"github.com/google/differential-privacy/privacy-on-beam/v2/pbeam"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam"
 )
@@ -27,7 +28,10 @@ import (
 func PrivateCountVisitsPerHourWithPublicPartitions(s beam.Scope, col beam.PCollection) beam.PCollection {
 	s = s.Scope("PrivateCountVisitsPerHourWithPublicPartitions")
 	// Create a Privacy Spec and convert col into a PrivatePCollection.
-	spec := pbeam.NewPrivacySpec(epsilon /* delta */, 0)
+	spec, err := pbeam.NewPrivacySpecTemp(pbeam.PrivacySpecParams{AggregationEpsilon: epsilon})
+	if err != nil {
+		log.Fatalf("Couldn't create a PrivacySpec: %v", err)
+	}
 	pCol := pbeam.MakePrivateFromStruct(s, col, spec, "VisitorID")
 
 	// Create a PCollection of output partitions, i.e. restaurant's work hours (from 9 am till 9pm (exclusive)).

privacy-on-beam/codelab/sum.go

@@ -17,6 +17,7 @@
 package codelab
 
 import (
+	log "github.com/golang/glog"
 	"github.com/google/differential-privacy/privacy-on-beam/v2/pbeam"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/register"
@@ -46,7 +47,10 @@ func extractVisitHourAndMoneySpentFn(v Visit) (int, int) {
 func PrivateRevenuePerHour(s beam.Scope, col beam.PCollection) beam.PCollection {
 	s = s.Scope("PrivateRevenuePerHour")
 	// Create a Privacy Spec and convert col into a PrivatePCollection.
-	spec := pbeam.NewPrivacySpec(epsilon /* delta */, 0)
+	spec, err := pbeam.NewPrivacySpecTemp(pbeam.PrivacySpecParams{AggregationEpsilon: epsilon})
+	if err != nil {
+		log.Fatalf("Couldn't create a PrivacySpec: %v", err)
+	}
 	pCol := pbeam.MakePrivateFromStruct(s, col, spec, "VisitorID")
 
 	// Create a PCollection of output partitions, i.e. restaurant's work hours (from 9 am till 9pm (exclusive)).

privacy-on-beam/pbeam/aggregations.go

@@ -392,7 +392,7 @@ func (fn *boundedSumInt64Fn) Setup() {
 }
 
 func (fn *boundedSumInt64Fn) CreateAccumulator() (boundedSumAccumInt64, error) {
-	if fn.TestMode == NoNoiseWithoutContributionBounding {
+	if fn.TestMode == TestModeWithoutContributionBounding {
 		fn.Lower = math.MinInt64
 		fn.Upper = math.MaxInt64
 	}
@@ -557,7 +557,7 @@ func (fn *boundedSumFloat64Fn) Setup() {
 }
 
 func (fn *boundedSumFloat64Fn) CreateAccumulator() (boundedSumAccumFloat64, error) {
-	if fn.TestMode == NoNoiseWithoutContributionBounding {
+	if fn.TestMode == TestModeWithoutContributionBounding {
 		fn.Lower = math.Inf(-1)
 		fn.Upper = math.Inf(1)
 	}

privacy-on-beam/pbeam/aggregations_test.go

@@ -87,7 +87,7 @@ func TestNewBoundedSumFn(t *testing.T) {
 				PublicPartitions:          false,
 			}},
 	} {
-		got, err := newBoundedSumFn(1, 1e-5, 17, 0, 10, tc.noiseKind, tc.vKind, false, Disabled)
+		got, err := newBoundedSumFn(1, 1e-5, 17, 0, 10, tc.noiseKind, tc.vKind, false, TestModeDisabled)
 		if err != nil {
 			t.Fatalf("Couldn't get boundedSumFn: %v", err)
 		}
@@ -195,7 +195,7 @@ func TestNewBoundedSumFnTemp(t *testing.T) {
 		{"lower > upper", noise.GaussianNoise, reflect.Int64, 0.5, 1e-5, 0.5, 1e-5, 0, 10, 0, true, nil},
 		{"Float64 bounds that overflow when converted to int64", noise.GaussianNoise, reflect.Int64, 0.5, 1e-5, 0.5, 1e-5, 0, 0, 1e100, true, nil},
 	} {
-		got, err := newBoundedSumFnTemp(PrivacySpec{preThreshold: tc.preThreshold, testMode: Disabled},
+		got, err := newBoundedSumFnTemp(PrivacySpec{preThreshold: tc.preThreshold, testMode: TestModeDisabled},
 			SumParams{
 				AggregationEpsilon:       tc.aggregationEpsilon,
 				AggregationDelta:         tc.aggregationDelta,
@@ -221,7 +221,7 @@ func TestBoundedSumFloat64FnSetup(t *testing.T) {
 	}{
 		{"Laplace noise kind", noise.LaplaceNoise, noise.Laplace()},
 		{"Gaussian noise kind", noise.GaussianNoise, noise.Gaussian()}} {
-		got, err := newBoundedSumFloat64Fn(1, 1e-5, 17, 0, 10, tc.noiseKind, false, Disabled)
+		got, err := newBoundedSumFloat64Fn(1, 1e-5, 17, 0, 10, tc.noiseKind, false, TestModeDisabled)
 		if err != nil {
 			t.Fatalf("Couldn't get boundedSumFloat64Fn: %v", err)
 		}
@@ -240,7 +240,7 @@ func TestBoundedSumInt64FnSetup(t *testing.T) {
 	}{
 		{"Laplace noise kind", noise.LaplaceNoise, noise.Laplace()},
 		{"Gaussian noise kind", noise.GaussianNoise, noise.Gaussian()}} {
-		got, err := newBoundedSumInt64Fn(1, 1e-5, 17, 0, 10, tc.noiseKind, false, Disabled)
+		got, err := newBoundedSumInt64Fn(1, 1e-5, 17, 0, 10, tc.noiseKind, false, TestModeDisabled)
 		if err != nil {
 			t.Fatalf("Couldn't get boundedSumInf64Fn: %v", err)
 		}
@@ -255,7 +255,7 @@ func TestBoundedSumInt64FnAddInput(t *testing.T) {
 	// Since δ=0.5 and 2 entries are added, PreAggPartitionSelection always emits.
 	// Since ε=1e100, the noise is added with probability in the order of exp(-1e100),
 	// which means we don't have to worry about tolerance/flakiness calculations.
-	fn, err := newBoundedSumInt64Fn(1e100, 0.5, 1, 0, 2, noise.LaplaceNoise, false, Disabled)
+	fn, err := newBoundedSumInt64Fn(1e100, 0.5, 1, 0, 2, noise.LaplaceNoise, false, TestModeDisabled)
 	if err != nil {
 		t.Fatalf("Couldn't get boundedSumInt64Fn: %v", err)
 	}
@@ -285,7 +285,7 @@ func TestBoundedSumInt64FnMergeAccumulators(t *testing.T) {
 	//
 	// Since ε=1e100, the noise is added with probability in the order of exp(-1e100),
 	// which means we don't have to worry about tolerance/flakiness calculations.
-	fn, err := newBoundedSumInt64Fn(1e100, 0.5, 1, 0, 2, noise.LaplaceNoise, false, Disabled)
+	fn, err := newBoundedSumInt64Fn(1e100, 0.5, 1, 0, 2, noise.LaplaceNoise, false, TestModeDisabled)
 	if err != nil {
 		t.Fatalf("Couldn't get boundedSumInt64Fn: %v", err)
 	}
@@ -323,7 +323,7 @@ func TestBoundedSumInt64FnExtractOutputReturnsNilForSmallPartitions(t *testing.T
 		// The probability of keeping a partition with 1 privacy unit is equal to δ=1e-23 which results in a flakiness of 10⁻²³.
 		{"Input with 1 privacy unit", 1}} {
 
-		fn, err := newBoundedSumInt64Fn(1, 1e-23, 1, 0, 2, noise.LaplaceNoise, false, Disabled)
+		fn, err := newBoundedSumInt64Fn(1, 1e-23, 1, 0, 2, noise.LaplaceNoise, false, TestModeDisabled)
 		if err != nil {
 			t.Fatalf("Couldn't get boundedSumInt64Fn: %v", err)
 		}
@@ -358,7 +358,7 @@ func TestBoundedSumInt64FnExtractOutputWithPublicPartitionsDoesNotThreshold(t *t
 		{"Input with 10 users", 10},
 		{"Input with 100 users", 100}} {
 
-		fn, err := newBoundedSumInt64Fn(1, 0, 1, 0, 2, noise.LaplaceNoise, true, Disabled)
+		fn, err := newBoundedSumInt64Fn(1, 0, 1, 0, 2, noise.LaplaceNoise, true, TestModeDisabled)
 		if err != nil {
 			t.Fatalf("Couldn't get boundedSumInt64Fn: %v", err)
 		}
@@ -385,7 +385,7 @@ func TestBoundedSumInt64FnExtractOutputWithPublicPartitionsDoesNotThreshold(t *t
 func TestBoundedSumFloat64FnAddInput(t *testing.T) {
 	// Since δ=0.5 and 2 entries are added, PreAggPartitionSelection always emits.
 	// Since ε=1e100, added noise is negligible.
-	fn, err := newBoundedSumFloat64Fn(1e100, 0.5, 1, 0, 2, noise.LaplaceNoise, false, Disabled)
+	fn, err := newBoundedSumFloat64Fn(1e100, 0.5, 1, 0, 2, noise.LaplaceNoise, false, TestModeDisabled)
 	if err != nil {
 		t.Fatalf("Couldn't get boundedSumFloat64Fn: %v", err)
 	}
@@ -414,7 +414,7 @@ func TestBoundedSumFloat64FnMergeAccumulators(t *testing.T) {
 	// accumulators is also effecting our partition selection outcome.
 	//
 	// Since ε=1e100, added noise is negligible.
-	fn, err := newBoundedSumFloat64Fn(1e100, 0.5, 1, 0, 2, noise.LaplaceNoise, false, Disabled)
+	fn, err := newBoundedSumFloat64Fn(1e100, 0.5, 1, 0, 2, noise.LaplaceNoise, false, TestModeDisabled)
 	if err != nil {
 		t.Fatalf("Couldn't get boundedSumFloat64Fn: %v", err)
 	}
@@ -452,7 +452,7 @@ func TestBoundedSumFloat64FnExtractOutputReturnsNilForSmallPartitions(t *testing
 		// The probability of keeping a partition with 1 privacy unit is equal to δ=1e-23 which results in a flakiness of 10⁻²³.
 		{"Input with 1 privacy unit", 1}} {
 
-		fn, err := newBoundedSumFloat64Fn(1, 1e-23, 1, 0, 2, noise.LaplaceNoise, false, Disabled)
+		fn, err := newBoundedSumFloat64Fn(1, 1e-23, 1, 0, 2, noise.LaplaceNoise, false, TestModeDisabled)
 		if err != nil {
 			t.Fatalf("Couldn't get boundedSumFloat64Fn: %v", err)
 		}
@@ -487,7 +487,7 @@ func TestBoundedSumFloat64FnExtractOutputWithPublicPartitionsDoesNotThreshold(t
 		{"Input with 10 users", 10},
 		{"Input with 100 users", 100}} {
 		publicPartitions := true
-		fn, err := newBoundedSumFloat64Fn(1, 0, 1, 0, 2, noise.LaplaceNoise, publicPartitions, Disabled)
+		fn, err := newBoundedSumFloat64Fn(1, 0, 1, 0, 2, noise.LaplaceNoise, publicPartitions, TestModeDisabled)
 		if err != nil {
 			t.Fatalf("Couldn't get boundedSumFloat64Fn: %v", err)
 		}

privacy-on-beam/pbeam/count.go

@@ -171,7 +171,7 @@ func Count(s beam.Scope, pcol PrivatePCollection, params CountParams) beam.PColl
 	counts64 := beam.ParDo(s, convertIntToInt64, kvCounts)
 	rekeyed := beam.ParDo(s, rekeyInt64, counts64)
 	// Second, do cross-partition contribution bounding if not in test mode without contribution bounding.
-	if spec.testMode != NoNoiseWithoutContributionBounding {
+	if spec.testMode != TestModeWithoutContributionBounding {
 		rekeyed = boundContributions(s, rekeyed, params.MaxPartitionsContributed)
 	}
 	// Third, now that contribution bounding is done, remove the privacy keys,

privacy-on-beam/pbeam/distinct_id.go

@@ -174,7 +174,7 @@ func DistinctPrivacyID(s beam.Scope, pcol PrivatePCollection, params DistinctPri
 		beam.TypeDefinition{Var: beam.TType, T: idT.Type()},
 		beam.TypeDefinition{Var: beam.VType, T: partitionT.Type()})
 	// Second, do cross-partition contribution bounding if not in test mode without contribution bounding.
-	if spec.testMode != NoNoiseWithoutContributionBounding {
+	if spec.testMode != TestModeWithoutContributionBounding {
 		decoded = boundContributions(s, decoded, params.MaxPartitionsContributed)
 	}
 	// Third, now that KV pairs are deduplicated and contribution bounding is

privacy-on-beam/pbeam/distinct_id_test.go

@@ -610,7 +610,7 @@ func TestNewCountFn(t *testing.T) {
 				NoiseKind:                noise.GaussianNoise,
 			}},
 	} {
-		got, err := newCountFn(1, 1e-5, 17, tc.noiseKind, false, Disabled)
+		got, err := newCountFn(1, 1e-5, 17, tc.noiseKind, false, TestModeDisabled)
 		if err != nil {
 			t.Fatalf("Couldn't get countFn: %v", err)
 		}
@@ -658,7 +658,7 @@ func TestNewCountFnTemp(t *testing.T) {
 				NoiseKind:                noise.LaplaceNoise,
 			}},
 	} {
-		got, err := newCountFnTemp(PrivacySpec{preThreshold: tc.preThreshold, testMode: Disabled},
+		got, err := newCountFnTemp(PrivacySpec{preThreshold: tc.preThreshold, testMode: TestModeDisabled},
 			DistinctPrivacyIDParams{
 				AggregationEpsilon:       tc.aggregationEpsilon,
 				AggregationDelta:         tc.aggregationDelta,
@@ -682,7 +682,7 @@ func TestCountFnSetup(t *testing.T) {
 	}{
 		{"Laplace noise kind", noise.LaplaceNoise, noise.Laplace()},
 		{"Gaussian noise kind", noise.GaussianNoise, noise.Gaussian()}} {
-		got, err := newCountFn(1, 1e-5, 17, tc.noiseKind, false, Disabled)
+		got, err := newCountFn(1, 1e-5, 17, tc.noiseKind, false, TestModeDisabled)
 		if err != nil {
 			t.Fatalf("Couldn't get countFn: %v", err)
 		}

privacy-on-beam/pbeam/distinct_per_key.go

@@ -144,9 +144,11 @@ func DistinctPerKey(s beam.Scope, pcol PrivatePCollection, params DistinctPerKey
 		if err != nil {
 			log.Fatalf("Couldn't get aggregation budget for DistinctPerKey: %v", err)
 		}
-		params.PartitionSelectionParams.Epsilon, params.PartitionSelectionParams.Delta, err = spec.partitionSelectionBudget.get(params.PartitionSelectionParams.Epsilon, params.PartitionSelectionParams.Delta)
-		if err != nil {
-			log.Fatalf("Couldn't get partition selection budget for DistinctPerKey: %v", err)
+		if params.PublicPartitions == nil {
+			params.PartitionSelectionParams.Epsilon, params.PartitionSelectionParams.Delta, err = spec.partitionSelectionBudget.get(params.PartitionSelectionParams.Epsilon, params.PartitionSelectionParams.Delta)
+			if err != nil {
+				log.Fatalf("Couldn't get partition selection budget for DistinctPerKey: %v", err)
+			}
 		}
 	} else {
 		params.AggregationEpsilon, params.AggregationDelta, err = spec.budget.get(params.Epsilon, params.Delta)
@@ -157,7 +159,6 @@ func DistinctPerKey(s beam.Scope, pcol PrivatePCollection, params DistinctPerKey
 			params.AggregationEpsilon, params.AggregationDelta, params.PartitionSelectionParams.Epsilon, params.PartitionSelectionParams.Delta = splitBudget(params.AggregationEpsilon, params.AggregationDelta, noiseKind)
 		}
 	}
-
 	err = checkDistinctPerKeyParams(params, noiseKind, pcol.codec.KType.T)
 	if err != nil {
 		log.Fatalf("pbeam.DistinctPerKey: %v", err)
@@ -173,7 +174,7 @@ func DistinctPerKey(s beam.Scope, pcol PrivatePCollection, params DistinctPerKey
 	// This is not great in terms of utility, since dropping contributions randomly might
 	// mean that we keep duplicates instead of distinct values. However, this is necessary
 	// for the current algorithm to be DP.
-	if spec.testMode != NoNoiseWithoutContributionBounding {
+	if spec.testMode != TestModeWithoutContributionBounding {
 		// First, rekey by kv.Pair{ID,K} and do per-partition contribution bounding.
 		rekeyed := beam.ParDo(
 			s,
@@ -242,6 +243,8 @@ func DistinctPerKey(s beam.Scope, pcol PrivatePCollection, params DistinctPerKey
 		NoiseKind:                params.NoiseKind,
 		Epsilon:                  params.AggregationEpsilon,
 		Delta:                    params.AggregationDelta,
+		AggregationEpsilon:       params.AggregationEpsilon, // to also work with the new API
+		AggregationDelta:         params.AggregationDelta,   // to also work with the new API
 		MaxPartitionsContributed: params.MaxPartitionsContributed,
 		MaxValue:                 params.MaxContributionsPerPartition,
 		PublicPartitions:         params.PublicPartitions,

privacy-on-beam/pbeam/example_test.go

@@ -64,7 +64,14 @@ func Example() {
 	// provided by the pipeline.
 	const ε, δ = 1, 1e-3
 
-	privacySpec := pbeam.NewPrivacySpec(ε, δ)
+	privacySpec, err := pbeam.NewPrivacySpecTemp(pbeam.PrivacySpecParams{
+		AggregationEpsilon:        ε / 2,
+		PartitionSelectionEpsilon: ε / 2,
+		AggregationDelta:          δ,
+	})
+	if err != nil {
+		// Handle error.
+	}
 	pcol := pbeam.MakePrivateFromStruct(s, icol, privacySpec, "visitorID")
 	// pcol is now a PrivatePCollection<visit>.
 

privacy-on-beam/pbeam/mean.go

@@ -196,7 +196,7 @@ func MeanPerKey(s beam.Scope, pcol PrivatePCollection, params MeanParams) beam.P
 		beam.TypeDefinition{Var: beam.VType, T: pcol.codec.VType.T})
 
 	// Don't do per-partition contribution bounding if in test mode without contribution bounding.
-	if spec.testMode != NoNoiseWithoutContributionBounding {
+	if spec.testMode != TestModeWithoutContributionBounding {
 		decoded = boundContributions(s, decoded, params.MaxContributionsPerPartition)
 	}
 
@@ -218,7 +218,7 @@ func MeanPerKey(s beam.Scope, pcol PrivatePCollection, params MeanParams) beam.P
 	// Result is PCollection<ID, pairArrayFloat64>.
 	rekeyed := beam.ParDo(s, rekeyArrayFloat64, combined)
 	// Second, do cross-partition contribution bounding if not in test mode without contribution bounding.
-	if spec.testMode != NoNoiseWithoutContributionBounding {
+	if spec.testMode != TestModeWithoutContributionBounding {
 		rekeyed = boundContributions(s, rekeyed, params.MaxPartitionsContributed)
 	}
 
@@ -446,7 +446,7 @@ func (fn *boundedMeanFn) Setup() {
 }
 
 func (fn *boundedMeanFn) CreateAccumulator() (boundedMeanAccum, error) {
-	if fn.TestMode == NoNoiseWithoutContributionBounding && !fn.EmptyPartitions {
+	if fn.TestMode == TestModeWithoutContributionBounding && !fn.EmptyPartitions {
 		fn.Lower = math.Inf(-1)
 		fn.Upper = math.Inf(1)
 	}