schrodinger.seam.transforms.samplers module

class schrodinger.seam.transforms.samplers.RandomSample(n: int, seed: Optional[int] = None, distinct=False)

Bases: PTransform

A PTransform that returns approximately n random elements.

On average, the number of elements sampled will be at most 0.3% off from n. For small numbers of n (less than or equal to 100,000), it will be exactly n.

The seed value is only used if n is larger than 100,000.

Example usage:

>>> with beam.Pipeline() as p:
...     sample = (p | beam.Create(range(10))
...                 | RandomSample(3))
...     # sample will contain three randomly selected elements

If distinct is True, then the input pcollection is first deduplicated before sampling. The transform’s display data exposes that unique collection as the unique count source for finished-run metrics.

N_CUTOFF = 100000

UNIQUE_COUNT_SOURCE = 'unique'

__init__(n: int, seed: Optional[int] = None, distinct=False)

display_data() → dict[str, object]

Returns the display data associated to a pipeline component.

It should be reimplemented in pipeline components that wish to have static display data.

Returns:

Dict[str, Any]: A dictionary containing key:value pairs. The value might be an integer, float or string value; a DisplayDataItem for values that have more data (e.g. short value, label, url); or a HasDisplayData instance that has more display data that should be picked up. For example:

{
  'key1': 'string_value',
  'key2': 1234,
  'key3': 3.14159265,
  'key4': DisplayDataItem('apache.org', url='http://apache.org'),
  'key5': subComponent
}

class schrodinger.seam.transforms.samplers.ThrottleConfig(*, max_cpu_time_sec: float, elements_per_sec: float, max_sample_size: int, distinct: bool = False, seed: int | None = None)

Bases: BaseModel

Configuration for Throttle.

Can also be used standalone to compute the sample size without creating a full Throttle PTransform.

Parameters:

• max_cpu_time_sec – total CPU time budget in seconds (typically num_workers * avg_worker_time).

• elements_per_sec – historical throughput of the inner transform.

• max_sample_size – upper bound on the number of sampled elements.

• distinct – if True, deduplicate elements before sampling.

• seed – random seed for reproducible sampling.

max_cpu_time_sec: float

elements_per_sec: float

max_sample_size: int

distinct: bool

seed: int | None

property sample_size: int

The number of elements to sample, at least 1.

model_config: ClassVar[ConfigDict] = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class schrodinger.seam.transforms.samplers.Throttle(inner: PTransform, *, max_cpu_time_sec: float, elements_per_sec: float, max_sample_size: int, distinct: bool = False, seed: int | None = None)

Bases: PTransform

Wraps a PTransform with a RandomSample pre-filter to limit the number of elements processed, based on a compute budget.

The sample size is computed as:

sample_size = min(max_cpu_time_sec * elements_per_sec, max_sample_size)

where elements_per_sec is the historical throughput of the inner transform and max_cpu_time_sec is the total CPU time budget the caller is willing to spend (typically num_workers * avg_worker_time).

The output type is determined by the inner transform and is not annotated on this class.

Example usage:

result = pcoll | Throttle(
    MyExpensiveTransform(),
    max_cpu_time_sec=num_workers * 600,
    elements_per_sec=2.0,
    max_sample_size=100_000,
)

__init__(inner: PTransform, *, max_cpu_time_sec: float, elements_per_sec: float, max_sample_size: int, distinct: bool = False, seed: int | None = None)