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.

N_CUTOFF = 100000

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

display_data() → dict

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
}

WithCount()

Returns a tuple of the sampled pcollection and a pcollection containing the number of inputs that were sampled from.

Example usage:

>>> with beam.Pipeline() as p:
...     sample, count = (p | beam.Create([1, 1, 2, 2])
...                         | RandomSample(3, distinct=True).WithCount())
...     # sample will contain three randomly selected elements
...     # count will contain the number of elements in the input pcollection (2)

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)