Skip to content

Decouple beam_initializers from MBI/JAX dependency. #77

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Decouple beam_initializers from MBI/JAX dependency. #77

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
208 changes: 114 additions & 94 deletions dpsynth/local_mode/beam_initializers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -109,31 +128,23 @@ 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(
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._init_specs))
| 'Count' >> beam.combiners.Count.PerElement()
| 'Unpack' >> beam.Map(_unpack_count)
# Aggregate data and materialize on the driver (see module header).
Expand Down Expand Up @@ -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,
)
)
Loading
Loading