schrodinger.seam.transforms.core module

schrodinger.seam.transforms.core.yields_elements(finish_bundle_method: Callable)

A decorator that wraps a DoFn’s finish_bundle method to allow it to yield elements.

Example usage:

>>> import apache_beam as beam
>>> from schrodinger.seam.transforms import core
>>>
>>> class Batcher(beam.DoFn):
...
...     def __init__(self):
...         self._batch = []
...
...     def start_bundle(self):
...         self._batch.clear()
...
...     def process(self, element):
...         self._batch.append(element)
...
...     @yields_elements
...     def finish_bundle(self):
...         yield self._batch
...         self._batch.clear()
...

class schrodinger.seam.transforms.core.GroupIntoBatches(batch_size: int)

Bases: PTransform

PTransform that batches the input into desired batch size after grouping by key.

Identical behavior to Beam’s GroupIntoBatches but doesn’t use Beam timer’s (which are unsupported by the SeamRunner).

__init__(batch_size: int)

classmethod with_hot_keys(max_values_per_shard: int)

PTransform that batches the input into desired batch size after grouping by key.

Has the added feature of lowering memory by adding additional data into the keys such that the intermediate GroupByKey transform does not have to store all the data associated with keys that have a lot of values (“hot keys”) in memory prior to batching.

Parameters:

max_values_per_shard – the expected maximum amount of values per shard (every shard is not guaranteed to have fewer than max_values_per_shard and can exceed it by a small percentage)

Returns:

a group into batches with hot keys PTransform with

expand(pcoll)

class schrodinger.seam.transforms.core.TeeIf(condition: bool, *consumers: Union[PTransform[PCollection[T], Any], Callable[[PCollection[T]], Any]])

Bases: PTransform

A PTransform which performs beam.Tee() if the condition is True, and does nothing if the condition is False.

Example usage:

>>> import apache_beam as beam
>>> from schrodinger import structure
>>> from schrodinger.seam.io import chemio
>>> from schrodinger.seam.transforms import core
>>>
>>> sts = [structure.create_new_structure(num_atoms=i) for i in range(1, 11)]
>>> debug_flag = True
>>>
>>> with beam.Pipeline() as p:
...     _ = (p
...         | beam.Create(sts)
...         | core.TeeIf(debug_flag, chemio.WriteStructuresToFile("output_file.maegz"))
...     )
>>> st_count = structure.count_structures("output_file.maegz")
>>> # st_count will be 10, and output_file.maegz will be created because debug_flag was set to True

__init__(condition: bool, *consumers: Union[PTransform[PCollection[T], Any], Callable[[PCollection[T]], Any]])

expand(pcoll)

class schrodinger.seam.transforms.core.MemSafeDistinct(label: Optional[str] = None)

Bases: PTransform

Remove duplicate elements from a PCollection. Elements are considered identical if their encoded byte values are equivalent.

This is a drop-in replacement for apache_beam.Distinct, which can load up to 100,000 unique, decoded elements from a single bundle due to the lift_combiners optimization phase.

The trade-off is that avoiding that memory pressure comes at the cost of increased PCollection materialization, and therefore more I/O. Consider which of the two matters more for a given pipeline and choose accordingly.

See https://github.com/apache/beam/issues/39754

expand(pcoll: PCollection[T]) → PCollection[T]