move import inside function
this way it works in cloud mode, avoids ErrorName error
@@ -96,6 +96,7 @@ def parse_arguments(argv):
# TODO: Perhaps use Reshuffle (https://issues.apache.org/jira/browse/BEAM-1872)?
@beam.ptransform_fn
def _Shuffle(pcoll): # pylint: disable=invalid-name
+ import random
return (pcoll
| 'PairWithRandom' >> beam.Map(lambda x: (random.random(), x))
| 'GroupByRandom' >> beam.GroupByKey()
