From 5abbd4b2e1cb3625469d12ec146d373be3b9343d Mon Sep 17 00:00:00 2001 From: Ryan McKenna Date: Thu, 2 Jul 2026 14:01:36 -0700 Subject: [PATCH] Decouple beam_initializers from MBI/JAX dependency. PiperOrigin-RevId: 941815983 --- dpsynth/local_mode/beam_initializers.py | 208 +++++++++++---------- tests/local_mode/beam_initializers_test.py | 72 +++++-- 2 files changed, 172 insertions(+), 108 deletions(-) diff --git a/dpsynth/local_mode/beam_initializers.py b/dpsynth/local_mode/beam_initializers.py index be00197..537b60f 100644 --- a/dpsynth/local_mode/beam_initializers.py +++ b/dpsynth/local_mode/beam_initializers.py @@ -15,63 +15,82 @@ """Beam-backed column initializers for DP Synth. Computes per-column sufficient statistics via Apache Beam PTransforms, -then delegates to the existing initializers' ``from_summary()`` methods -for DP mechanism execution on the driver. The central assumption in this file -is that the data is too large to feasibly materialize in memory on the driver, -but the per-column sufficient statistics can easily fit. The intention is to -use Beam where it is absolutely needed, but quickly delegate to local-mode -implementations as soon as the sufficient statistics are available, creating -a clear separation of concerns. All beam-related logic necessary to use the -local mode variant of DPSynth is contained in this file. +then runs DP mechanisms from ``primitives.py`` directly on the driver. +No dependency on MBI or JAX — only ``numpy``, ``domain``, ``primitives``, +and ``vectorized_transformations`` are imported. All outputs are pure +NumPy arrays in lightweight dataclasses. """ from __future__ import annotations +import dataclasses +import enum import math from typing import Any import apache_beam as beam -from dpsynth.local_mode import initialization +from dpsynth import domain +from dpsynth.local_mode import vectorized_transformations as vtx import numpy as np -# A single row of tabular data: column name -> raw value. -# representation for large pipelines. Consider supporting named tuples or -# a schema-aware format (e.g. Beam Rows, protos) to reduce per-element overhead. +# representation for large pipelines. Consider named tuples or Beam Rows. Row = dict[str, Any] -Initializer = ( - initialization.NumericalInitializer - | initialization.CategoricalInitializer - | initialization.OpenSetCategoricalInitializer -) + +class ColumnType(enum.Enum): + NUMERICAL = 'numerical' + CATEGORICAL = 'categorical' + OPENSET = 'openset' + + +@dataclasses.dataclass +class BeamColumnResult: + """Lightweight column result without MBI/JAX dependency.""" + + column_type: ColumnType + categorical_attribute: domain.CategoricalAttribute + bin_edges: np.ndarray | None = None + noisy_counts: np.ndarray | None = None + stddev: float | None = None + + +@dataclasses.dataclass +class InitSpec: + """Per-column mechanism + attribute specification (MBI-free).""" + + column_type: ColumnType + mechanism: Any # primitives.DPMechanism subclass + attribute: Any # domain.*Attribute + # Numerical quantile grid, from the calibrated DPQuantiles mechanism. + grid_lower: float | None = None + grid_upper: float | None = None + grid_size: int | None = None + min_count: int = 1 # openset only class _EncodeColumns(beam.DoFn): """Encodes each row into (column, key) pairs for all columns at once.""" - def __init__(self, initializers: dict[str, Initializer]): + def __init__(self, init_specs: dict[str, InitSpec]): # Do all setup in __init__ so that process below is cheaper. # We handle all columns at once here to reduce the size of the DAG in Beam. super().__init__() self._specs: list[tuple[str, str, dict[str, Any]]] = [] - for column, init in initializers.items(): - if isinstance(init, initialization.NumericalInitializer): - attr = init.attribute - lower, upper, gs = init._grid_spec - delta = (upper - lower) / (gs - 1) + for column, spec in init_specs.items(): + if spec.column_type == ColumnType.NUMERICAL: + attr = spec.attribute + lower = spec.grid_lower + delta = (spec.grid_upper - lower) / (spec.grid_size - 1) meta = {'attribute': attr, 'lower': lower, 'delta': delta} self._specs.append((column, 'numerical', meta)) - - elif isinstance(init, initialization.CategoricalInitializer): + elif spec.column_type == ColumnType.CATEGORICAL: lookup = { - str(v): i for i, v in enumerate(init.attribute.possible_values) + str(v): i for i, v in enumerate(spec.attribute.possible_values) } - meta = {'lookup': lookup, 'default': init.attribute.out_of_domain_index} + meta = {'lookup': lookup, 'default': spec.attribute.out_of_domain_index} self._specs.append((column, 'categorical', meta)) - elif isinstance(init, initialization.OpenSetCategoricalInitializer): + elif spec.column_type == ColumnType.OPENSET: self._specs.append((column, 'openset', {})) - else: - raise TypeError(f'Unsupported initializer type: {type(init)}') def process(self, row: Row): for column, kind, params in self._specs: @@ -109,23 +128,15 @@ def _materialize_pairs(col, pairs): class ComputeSufficientStats(beam.PTransform): - """Computes per-column sufficient statistics in a single pass. - - Encodes all columns in one ``DoFn``, then counts via a single - ``Count.PerElement`` and groups by column. The output is a ``PCollection`` - of ``(column_name, sparse_counts_list)`` pairs. - - Attributes: - initializers: Calibrated initializers keyed by column name. - """ + """Computes per-column sufficient statistics in a single pass.""" - def __init__(self, initializers: dict[str, Initializer]): + def __init__(self, init_specs: dict[str, InitSpec]): super().__init__() - self._initializers = initializers + self._init_specs = init_specs self._openset_min_counts = { - col: init.min_count - for col, init in initializers.items() - if isinstance(init, initialization.OpenSetCategoricalInitializer) + col: spec.min_count + for col, spec in init_specs.items() + if spec.column_type == ColumnType.OPENSET } def expand( @@ -133,7 +144,7 @@ def expand( ) -> beam.PCollection[tuple[str, list[tuple[Any, int]]]]: return ( rows - | 'Encode' >> beam.ParDo(_EncodeColumns(self._initializers)) + | 'Encode' >> beam.ParDo(_EncodeColumns(self._init_specs)) | 'Count' >> beam.combiners.Count.PerElement() | 'Unpack' >> beam.Map(_unpack_count) # Aggregate data and materialize on the driver (see module header). @@ -168,75 +179,84 @@ def _sparse_to_openset(sparse): return np.array(keys), np.array(vals, dtype=np.float64) -# into the Beam pipeline, which can increase setup time for each worker. +def _numerical_result(spec, rng, counts): + """Runs the quantile mechanism and builds a numerical BeamColumnResult.""" + raw_edges = np.asarray(spec.mechanism(rng, counts), dtype=float) + bin_edges, _ = np.unique(raw_edges, return_counts=True) + # Edges at or above max_value produce a degenerate empty tail bin. + if len(bin_edges) > 0 and bin_edges[-1] >= spec.attribute.max_value: + bin_edges = bin_edges[:-1] + cat_attr = vtx.categorical_attribute_from_edges(bin_edges, spec.attribute) + return BeamColumnResult(ColumnType.NUMERICAL, cat_attr, bin_edges=bin_edges) + + +def _categorical_result(spec, rng, counts): + """Runs the count mechanism and builds a categorical BeamColumnResult.""" + result = spec.mechanism(rng, counts) + return BeamColumnResult( + ColumnType.CATEGORICAL, + spec.attribute, + noisy_counts=result.counts, + stddev=spec.mechanism.sigma, + ) + + +def _openset_result(spec, rng, unique_values, value_counts): + """Runs partition selection and builds an open-set BeamColumnResult.""" + result = spec.mechanism.from_summary(rng, value_counts) + selected = [str(v) for v in unique_values[result.selected_partitions]] + possible = [spec.attribute.default_value] + selected + cat_attr = domain.CategoricalAttribute( + possible_values=possible, out_of_domain_index=0 + ) + return BeamColumnResult( + ColumnType.OPENSET, + cat_attr, + noisy_counts=result.estimated_counts, + stddev=spec.mechanism.sigma, + ) + + def run_from_summary( sparse_stats: dict[str, list[tuple[Any, int]]], - initializers: dict[str, Initializer], + init_specs: dict[str, InitSpec], rng: np.random.Generator, -) -> dict[str, initialization.ColumnMeasurement]: - """Converts materialized sparse stats to ColumnMeasurements on the driver. - - Meant to be called after ``ComputeSufficientStats`` results have been - materialized (e.g. via ``beam.combiners.ToDict()``). - - Args: - sparse_stats: Column-keyed dict of sparse (key, count) pair lists, as - produced by ``ComputeSufficientStats``. - initializers: Calibrated initializers keyed by column name. - rng: NumPy random generator for DP noise. - - Returns: - Per-column ``ColumnMeasurement`` results. - """ - results: dict[str, initialization.ColumnMeasurement] = {} - for column, init in initializers.items(): +) -> dict[str, BeamColumnResult]: + """Runs DP mechanisms via primitives and returns pure NumPy results.""" + results: dict[str, BeamColumnResult] = {} + for column, spec in init_specs.items(): sparse = sparse_stats[column] - if isinstance(init, initialization.NumericalInitializer): - counts = _sparse_to_dense_numerical(sparse, init.grid_size) - results[column] = init.from_summary(rng, counts) - elif isinstance(init, initialization.CategoricalInitializer): - counts = _sparse_to_dense_categorical(sparse, init.attribute.size) - results[column] = init.from_summary(rng, counts) - elif isinstance(init, initialization.OpenSetCategoricalInitializer): + if spec.column_type == ColumnType.NUMERICAL: + counts = _sparse_to_dense_numerical(sparse, spec.grid_size) + results[column] = _numerical_result(spec, rng, counts) + elif spec.column_type == ColumnType.CATEGORICAL: + counts = _sparse_to_dense_categorical(sparse, spec.attribute.size) + results[column] = _categorical_result(spec, rng, counts) + elif spec.column_type == ColumnType.OPENSET: unique_values, value_counts = _sparse_to_openset(sparse) - results[column] = init.from_summary(rng, unique_values, value_counts) + results[column] = _openset_result(spec, rng, unique_values, value_counts) return results class BeamInitialize(beam.PTransform): - """End-to-end: computes sufficient stats and runs DP initialization. - - Composes ``ComputeSufficientStats`` with sparse-to-dense conversion and - ``from_summary()`` calls. Produces a singleton ``PCollection`` containing - one ``dict[str, ColumnMeasurement]`` with all results. - - Attributes: - initializers: Calibrated initializers keyed by column name. - rng: NumPy random generator for DP noise. - """ - - def __init__( - self, - initializers: dict[str, Initializer], - rng: np.random.Generator, - ): + """Computes sufficient stats and runs DP initialization.""" + + def __init__(self, init_specs: dict[str, InitSpec], rng: np.random.Generator): super().__init__() - self._initializers = initializers + self._init_specs = init_specs self._rng = rng def expand( self, rows: beam.PCollection[Row] - ) -> beam.PCollection[dict[str, initialization.ColumnMeasurement]]: + ) -> beam.PCollection[dict[str, BeamColumnResult]]: return ( rows - | 'Stats' >> ComputeSufficientStats(self._initializers) + | 'Stats' >> ComputeSufficientStats(self._init_specs) | 'ToDict' >> beam.combiners.ToDict() | 'Initialize' - # Since all sufficient stats have been computed and materialized on the - # driver, passing a single rng is fine here. >> beam.Map( run_from_summary, - initializers=self._initializers, + init_specs=self._init_specs, rng=self._rng, ) ) diff --git a/tests/local_mode/beam_initializers_test.py b/tests/local_mode/beam_initializers_test.py index ffb5cee..693d4d0 100644 --- a/tests/local_mode/beam_initializers_test.py +++ b/tests/local_mode/beam_initializers_test.py @@ -29,18 +29,46 @@ def _store(x): _TEST_RESULTS.append(x) +def _init_spec_from_initializer(init): + """Builds an InitSpec from a configured initializer (as a caller would).""" + if isinstance(init, initialization.NumericalInitializer): + return beam_initializers.InitSpec( + beam_initializers.ColumnType.NUMERICAL, + init.mechanism, + init.attribute, + grid_lower=init.mechanism.lower, + grid_upper=init.mechanism.upper, + grid_size=init.grid_size, + ) + elif isinstance(init, initialization.CategoricalInitializer): + return beam_initializers.InitSpec( + beam_initializers.ColumnType.CATEGORICAL, + init.mechanism, + init.attribute, + ) + elif isinstance(init, initialization.OpenSetCategoricalInitializer): + return beam_initializers.InitSpec( + beam_initializers.ColumnType.OPENSET, + init.mechanism, + init.attribute, + min_count=init.min_count, + ) + raise TypeError(type(init)) + + class NumericalHistogramTest(absltest.TestCase): def _run(self, rows, attr, max_grid_size=101): init = initialization.NumericalInitializer( name='x', num_partitions=4, attribute=attr, max_grid_size=max_grid_size - ).calibrate(zcdp_rho=np.inf) + ).configure(zcdp_rho=np.inf) + spec = _init_spec_from_initializer(init) _TEST_RESULTS.clear() with beam.Pipeline() as p: stats = ( p | beam.Create(rows) - | beam_initializers.ComputeSufficientStats({'x': init}) + | beam_initializers.ComputeSufficientStats({'x': spec}) ) _ = stats | beam.combiners.ToDict() | beam.Map(_store) return dict(_TEST_RESULTS[0]['x']) @@ -84,7 +112,8 @@ def test_basic_counts(self): ) init = initialization.CategoricalInitializer( name='col', attribute=attr - ).calibrate(zcdp_rho=np.inf) + ).configure(zcdp_rho=np.inf) + spec = _init_spec_from_initializer(init) rows = [ {'col': 'a'}, {'col': 'a'}, @@ -99,7 +128,7 @@ def test_basic_counts(self): stats = ( p | beam.Create(rows) - | beam_initializers.ComputeSufficientStats({'col': init}) + | beam_initializers.ComputeSufficientStats({'col': spec}) ) _ = stats | beam.combiners.ToDict() | beam.Map(_store) counts = dict(_TEST_RESULTS[0]['col']) @@ -117,6 +146,7 @@ def test_basic_counts(self): init = initialization.OpenSetCategoricalInitializer( name='col', attribute=attr, delta=0.01, min_count=1 ).configure(zcdp_rho=np.inf) + spec = _init_spec_from_initializer(init) rows = [ {'col': 'apple'}, {'col': 'apple'}, @@ -130,7 +160,7 @@ def test_basic_counts(self): stats = ( p | beam.Create(rows) - | beam_initializers.ComputeSufficientStats({'col': init}) + | beam_initializers.ComputeSufficientStats({'col': spec}) ) _ = stats | beam.combiners.ToDict() | beam.Map(_store) counts = dict(_TEST_RESULTS[0]['col']) @@ -142,12 +172,11 @@ def test_basic_counts(self): class BeamInitializeTest(absltest.TestCase): - def test_end_to_end_mixed(self): + def _make_init_specs(self): num_attr = domain.NumericalAttribute(min_value=0, max_value=100) cat_attr = domain.CategoricalAttribute(possible_values=['a', 'b']) open_attr = domain.OpenSetCategoricalAttribute(default_value='') - - initializers = { + inits = { 'score': ( initialization.NumericalInitializer( name='score', num_partitions=4, attribute=num_attr @@ -164,27 +193,42 @@ def test_end_to_end_mixed(self): ).configure(zcdp_rho=np.inf) ), } + return {k: _init_spec_from_initializer(v) for k, v in inits.items()} + def _run_pipeline(self, init_specs): rows = [ {'score': 25.0, 'grade': 'a', 'tag': 'p'}, {'score': 50.0, 'grade': 'b', 'tag': 'q'}, {'score': 75.0, 'grade': 'a', 'tag': 'p'}, ] rng = np.random.default_rng(42) - _TEST_RESULTS.clear() with beam.Pipeline() as p: result = ( p | beam.Create(rows) - | beam_initializers.BeamInitialize(initializers, rng) + | beam_initializers.BeamInitialize(init_specs, rng) ) _ = result | beam.Map(_store) - measurements = _TEST_RESULTS[0] + return _TEST_RESULTS[0] - self.assertLen(measurements, 3) - for cm in measurements.values(): - self.assertIsInstance(cm, initialization.ColumnMeasurement) + def test_end_to_end_mixed(self): + results = self._run_pipeline(self._make_init_specs()) + self.assertLen(results, 3) + for br in results.values(): + self.assertIsInstance(br, beam_initializers.BeamColumnResult) + self.assertEqual( + results['score'].column_type, + beam_initializers.ColumnType.NUMERICAL, + ) + self.assertEqual( + results['grade'].column_type, + beam_initializers.ColumnType.CATEGORICAL, + ) + self.assertEqual( + results['tag'].column_type, + beam_initializers.ColumnType.OPENSET, + ) if __name__ == '__main__':