diff --git a/dpsynth/local_mode/beam_initializers.py b/dpsynth/local_mode/beam_initializers.py index 6fed006..4e48926 100644 --- a/dpsynth/local_mode/beam_initializers.py +++ b/dpsynth/local_mode/beam_initializers.py @@ -15,63 +15,100 @@ """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 + grid_size: int | None = None # numerical only + min_count: int = 1 # openset only + + +def _build_encode_specs( + init_specs: dict[str, InitSpec], +) -> tuple[list[tuple[str, str, dict[str, Any]]], dict[str, int]]: + """Derives encoding specs and openset min-counts from InitSpecs.""" + encode_specs = [] + openset_min_counts = {} + for column, spec in init_specs.items(): + if spec.column_type == ColumnType.NUMERICAL: + attr = spec.attribute + lower = attr.min_value + delta = (attr.exclusive_max_value - lower) / (spec.grid_size - 1) + encode_specs.append(( + column, + 'numerical', + { + 'attribute': attr, + 'lower': lower, + 'delta': delta, + }, + )) + elif spec.column_type == ColumnType.CATEGORICAL: + lookup = {str(v): i for i, v in enumerate(spec.attribute.possible_values)} + encode_specs.append(( + column, + 'categorical', + { + 'lookup': lookup, + 'default': spec.attribute.out_of_domain_index, + }, + )) + elif spec.column_type == ColumnType.OPENSET: + encode_specs.append((column, 'openset', {})) + openset_min_counts[column] = spec.min_count + return encode_specs, openset_min_counts 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, specs: list[tuple[str, str, dict[str, Any]]]): # 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 = attr.min_value - delta = (attr.exclusive_max_value - lower) / (init.grid_size - 1) - meta = {'attribute': attr, 'lower': lower, 'delta': delta} - self._specs.append((column, 'numerical', meta)) - - elif isinstance(init, initialization.CategoricalInitializer): - lookup = { - str(v): i for i, v in enumerate(init.attribute.possible_values) - } - meta = {'lookup': lookup, 'default': init.attribute.out_of_domain_index} - self._specs.append((column, 'categorical', meta)) - elif isinstance(init, initialization.OpenSetCategoricalInitializer): - self._specs.append((column, 'openset', {})) - else: - raise TypeError(f'Unsupported initializer type: {type(init)}') + self._specs = specs def process(self, row: Row): for column, kind, params in self._specs: @@ -109,31 +146,20 @@ 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. + """Computes per-column sufficient statistics in a single pass.""" - Attributes: - initializers: Calibrated initializers keyed by column name. - """ - - def __init__(self, initializers: dict[str, Initializer]): + def __init__(self, init_specs: dict[str, InitSpec]): super().__init__() - self._initializers = initializers - self._openset_min_counts = { - col: init.min_count - for col, init in initializers.items() - if isinstance(init, initialization.OpenSetCategoricalInitializer) - } + self._encode_specs, self._openset_min_counts = _build_encode_specs( + init_specs + ) def expand( self, rows: beam.PCollection[Row] ) -> beam.PCollection[tuple[str, list[tuple[Any, int]]]]: return ( rows - | 'Encode' >> beam.ParDo(_EncodeColumns(self._initializers)) + | 'Encode' >> beam.ParDo(_EncodeColumns(self._encode_specs)) | 'Count' >> beam.combiners.Count.PerElement() | 'Unpack' >> beam.Map(_unpack_count) # Aggregate data and materialize on the driver (see module header). @@ -168,75 +194,78 @@ 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 _edges_to_beam_result(raw_edges, attribute): + """Processes quantile edges into a BeamColumnResult (no MBI).""" + raw_edges = np.asarray(raw_edges, dtype=float) + if attribute.dtype == 'int': + raw_edges = np.floor(raw_edges) + bin_edges, _ = np.unique(raw_edges, return_counts=True) + # Remove edges at max_value to avoid degenerate tail bins. + if len(bin_edges) > 0 and bin_edges[-1] >= attribute.max_value: + bin_edges = bin_edges[:-1] + cat_attr = vtx.categorical_attribute_from_edges(bin_edges, attribute) + return BeamColumnResult(ColumnType.NUMERICAL, cat_attr, bin_edges=bin_edges) + + 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) + raw_edges = spec.mechanism(rng, counts) + results[column] = _edges_to_beam_result(raw_edges, spec.attribute) + elif spec.column_type == ColumnType.CATEGORICAL: + counts = _sparse_to_dense_categorical(sparse, spec.attribute.size) + result = spec.mechanism(rng, counts) + results[column] = BeamColumnResult( + ColumnType.CATEGORICAL, + spec.attribute, + noisy_counts=result.counts, + stddev=spec.mechanism.sigma, + ) + elif spec.column_type == ColumnType.OPENSET: unique_values, value_counts = _sparse_to_openset(sparse) - results[column] = init.from_summary(rng, unique_values, value_counts) + result = spec.mechanism.from_summary(rng, value_counts) + selected = list(unique_values[result.selected_partitions]) + possible = [spec.attribute.default_value] + selected + cat_attr = domain.CategoricalAttribute( + possible_values=possible, + out_of_domain_index=0, + ) + results[column] = BeamColumnResult( + ColumnType.OPENSET, + cat_attr, + noisy_counts=result.estimated_counts, + stddev=spec.mechanism.sigma, + ) 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 ae19f04..86e2115 100644 --- a/tests/local_mode/beam_initializers_test.py +++ b/tests/local_mode/beam_initializers_test.py @@ -29,18 +29,44 @@ def _store(x): _TEST_RESULTS.append(x) +def _init_spec_from_initializer(init): + """Helper: builds an InitSpec from a calibrated initializer.""" + if isinstance(init, initialization.NumericalInitializer): + return beam_initializers.InitSpec( + beam_initializers.ColumnType.NUMERICAL, + init.mechanism, + init.attribute, + 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, grid_size=101): init = initialization.NumericalInitializer( name='x', num_partitions=4, attribute=attr, grid_size=grid_size ).calibrate(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']) @@ -85,6 +111,7 @@ def test_basic_counts(self): init = initialization.CategoricalInitializer( name='col', attribute=attr ).calibrate(zcdp_rho=np.inf) + spec = _init_spec_from_initializer(init) rows = [ {'col': 'a'}, {'col': 'a'}, @@ -99,7 +126,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 +144,7 @@ def test_basic_counts(self): init = initialization.OpenSetCategoricalInitializer( name='col', attribute=attr, delta=0.01, min_count=1 ).calibrate(zcdp_rho=np.inf) + spec = _init_spec_from_initializer(init) rows = [ {'col': 'apple'}, {'col': 'apple'}, @@ -130,7 +158,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,49 +170,70 @@ def test_basic_counts(self): class BeamInitializeTest(absltest.TestCase): - def test_end_to_end_mixed(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=None) - - initializers = { + def _make_init_specs(self): + inits = { 'score': ( initialization.NumericalInitializer( - name='score', num_partitions=4, attribute=num_attr + name='score', + num_partitions=4, + attribute=domain.NumericalAttribute(min_value=0, max_value=100), ).calibrate(zcdp_rho=np.inf) ), 'grade': ( initialization.CategoricalInitializer( - name='grade', attribute=cat_attr + name='grade', + attribute=domain.CategoricalAttribute( + possible_values=['a', 'b'] + ), ).calibrate(zcdp_rho=np.inf) ), 'tag': ( initialization.OpenSetCategoricalInitializer( - name='tag', attribute=open_attr, delta=0.01, min_count=1 + name='tag', + attribute=domain.OpenSetCategoricalAttribute( + default_value=None + ), + delta=0.01, + min_count=1, ).calibrate(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__':