| 594 | _ = generator | f'{split_name}_write' >> _encode_pcollection() # pylint: disable=no-value-for-parameter |
| 595 | |
| 596 | def _resolve_future(): |
| 597 | if self._in_contextmanager: |
| 598 | raise AssertionError( |
| 599 | '`future.result()` should be called after the ' |
| 600 | '`maybe_beam_pipeline` contextmanager.' |
| 601 | ) |
| 602 | logging.info('Retrieving split info for %s...', split_name) |
| 603 | shard_lengths, total_size = beam_writer.finalize() |
| 604 | return splits_lib.SplitInfo( |
| 605 | name=split_name, |
| 606 | shard_lengths=shard_lengths, |
| 607 | num_bytes=total_size, |
| 608 | filename_template=filename_template, |
| 609 | ) |
| 610 | |
| 611 | return _SplitInfoFuture(_resolve_future) |