DoFn 的进程函数未执行

Question

>我正在尝试编写一个光束变换，例如

util.py

class GroupIntoBatches(PTransform):
  def __init__(self, batch_size):
    self.batch_size = batch_size
  @staticmethod
  def of_size(batch_size):
    return GroupIntoBatches(batch_size)
  def expand(self, pcoll):
    input_coder = coders.registry.get_coder(pcoll)
    if not input_coder.is_kv_coder():
          raise ValueError(
            'coder specified in the input PCollection is not a KvCoder')
    key_coder = input_coder.key_coder()
    value_coder = input_coder.value_coder()
    return pcoll | ParDo(_GroupIntoBatchesDoFn(self.batch_size, key_coder, value_coder))

class _GroupIntoBatchesDoFn(DoFn):
    def __init__(self, batch_size, input_key_coder, input_value_coder):
      self.batch_size = batch_size
      self.batch_spec = BagStateSpec("GroupIntoBatches", input_value_coder)
    def process(self, element):
      raise Exception("Not getting to this point") # This is not working
      print element

尝试通过测试用例执行此转换

util_test.py

class GroupIntoBatchesTest(unittest.TestCase):
  NUM_ELEMENTS = 10
  BATCH_SIZE = 5
  @staticmethod
  def _create_test_data():
    scientists = [
      "Einstein",
      "Darwin",
      "Copernicus",
      "Pasteur",
      "Curie",
      "Faraday",
      "Newton",
      "Bohr",
      "Galilei",
      "Maxwell"
    ]
    data = []
    for i in range(GroupIntoBatchesTest.NUM_ELEMENTS):
      index = i % len(scientists)
      data.append(("key", scientists[index]))
    return data
  def test_in_global_window(self):
    pipeline = TestPipeline()
    collection = pipeline | beam.Create(GroupIntoBatchesTest._create_test_data()) | util.GroupIntoBatches.of_size(GroupIntoBatchesTest.BATCH_SIZE)

我的问题是process函数没有在我的_GroupIntoBatchesDoFn上调用的原因是什么

我在运行测试用例时得到这个结果

test_in_global_window (apache_beam.transforms.util_test.组进批次测试( ...还行

Answer 1

测试正在构造管道，但实际上并未执行它。你需要写

pipeline = TestPipeline()
collection = pipeline | ...
pipeline.run()

或者，

或者，或者

with TestPipeline() as pipeline:
    collection = pipeline | ...
# run is implicitly called on exit of the with block

(您可能还对批处理元素转换感兴趣。