Python pipeline.Pipeline类代码示例

 def test_timestamped_with_combiners(self):
   p = Pipeline('DirectPipelineRunner')
   result = (p
             # Create some initial test values.
             | Create('start', [(k, k) for k in range(10)])
             # The purpose of the WindowInto transform is to establish a
             # FixedWindows windowing function for the PCollection.
             # It does not bucket elements into windows since the timestamps
             # from Create are not spaced 5 ms apart and very likely they all
             # fall into the same window.
             | WindowInto('w', FixedWindows(5))
             # Generate timestamped values using the values as timestamps.
             # Now there are values 5 ms apart and since Map propagates the
             # windowing function from input to output the output PCollection
             # will have elements falling into different 5ms windows.
             | Map(lambda (x, t): TimestampedValue(x, t))
             # We add a 'key' to each value representing the index of the
             # window. This is important since there is no guarantee of
             # order for the elements of a PCollection.
             | Map(lambda v: (v / 5, v)))
   # Sum all elements associated with a key and window. Although it
   # is called CombinePerKey it is really CombinePerKeyAndWindow the
   # same way GroupByKey is really GroupByKeyAndWindow.
   sum_per_window = result | CombinePerKey(sum)
   # Compute mean per key and window.
   mean_per_window = result | combiners.Mean.PerKey()
   assert_that(sum_per_window, equal_to([(0, 10), (1, 35)]),
               label='assert:sum')
   assert_that(mean_per_window, equal_to([(0, 2.0), (1, 7.0)]),
               label='assert:mean')
   p.run()

  def _run_write_test(self,
                      data,
                      return_init_result=True,
                      return_write_results=True):
    write_to_test_sink = WriteToTestSink(return_init_result,
                                         return_write_results)
    p = Pipeline(options=PipelineOptions([]))
    result = p | df.Create('start', data) | write_to_test_sink

    assert_that(result, is_empty())
    p.run()

    sink = write_to_test_sink.last_sink
    self.assertIsNotNone(sink)

    self.assertEqual(sink.state, _TestSink.STATE_FINALIZED)
    if data:
      self.assertIsNotNone(sink.last_writer)
      self.assertEqual(sink.last_writer.state, _TestWriter.STATE_CLOSED)
      self.assertEqual(sink.last_writer.write_output, data)
      if return_init_result:
        self.assertEqual(sink.last_writer.init_result,
                         _TestSink.TEST_INIT_RESULT)
        self.assertEqual(sink.init_result_at_finalize,
                         _TestSink.TEST_INIT_RESULT)
      self.assertIsNotNone(sink.last_writer.uid)
      if return_write_results:
        self.assertEqual(sink.write_results_at_finalize,
                         [_TestWriter.TEST_WRITE_RESULT])
    else:
      self.assertIsNone(sink.last_writer)

 def test_multi_valued_singleton_side_input(self):
   pipeline = Pipeline('DirectPipelineRunner')
   pcol = pipeline | Create('start', [1, 2])
   side = pipeline | Create('side', [3, 4])  # 2 values in side input.
   pcol | FlatMap('compute', lambda x, s: [x * s], AsSingleton(side))
   with self.assertRaises(ValueError) as e:
     pipeline.run()

 def test_word_count_using_get(self):
   pipeline = Pipeline('DirectPipelineRunner')
   lines = pipeline | Create('SomeWords', [DataflowTest.SAMPLE_DATA])
   result = (
       (lines | FlatMap('GetWords', lambda x: re.findall(r'\w+', x)))
       .apply('CountWords', DataflowTest.Count))
   assert_that(result, equal_to(DataflowTest.SAMPLE_RESULT))
   pipeline.run()

 def test_iterable_side_input(self):
   pipeline = Pipeline('DirectPipelineRunner')
   pcol = pipeline | Create('start', [1, 2])
   side = pipeline | Create('side', [3, 4])  # 2 values in side input.
   result = pcol | FlatMap('compute',
                           lambda x, s: [x * y for y in s], AllOf(side))
   assert_that(result, equal_to([3, 4, 6, 8]))
   pipeline.run()

 def test_default_value_singleton_side_input(self):
   pipeline = Pipeline('DirectPipelineRunner')
   pcol = pipeline | Create('start', [1, 2])
   side = pipeline | Create('side', [])  # 0 values in side input.
   result = (
       pcol | FlatMap('compute', lambda x, s: [x * s], AsSingleton(side, 10)))
   assert_that(result, equal_to([10, 20]))
   pipeline.run()

 def test_map(self):
   pipeline = Pipeline('DirectPipelineRunner')
   lines = pipeline | Create('input', ['a', 'b', 'c'])
   result = (lines
             | Map('upper', str.upper)
             | Map('prefix', lambda x, prefix: prefix + x, 'foo-'))
   assert_that(result, equal_to(['foo-A', 'foo-B', 'foo-C']))
   pipeline.run()

 def test_tuple_combine_fn_without_defaults(self):
   p = Pipeline('DirectPipelineRunner')
   result = (
       p
       | Create([1, 1, 2, 3])
       | df.CombineGlobally(
           combine.TupleCombineFn(min, combine.MeanCombineFn(), max)
           .with_common_input()).without_defaults())
   assert_that(result, equal_to([(1, 7.0 / 4, 3)]))
   p.run()

 def test_tuple_combine_fn(self):
   p = Pipeline('DirectPipelineRunner')
   result = (
       p
       | Create([('a', 100, 0.0), ('b', 10, -1), ('c', 1, 100)])
       | df.CombineGlobally(combine.TupleCombineFn(max,
                                                   combine.MeanCombineFn(),
                                                   sum)).without_defaults())
   assert_that(result, equal_to([('c', 111.0 / 3, 99.0)]))
   p.run()

  def test_create(self):
    pipeline = Pipeline('DirectPipelineRunner')
    pcoll = pipeline | Create('label1', [1, 2, 3])
    assert_that(pcoll, equal_to([1, 2, 3]))

    # Test if initial value is an iterator object.
    pcoll2 = pipeline | Create('label2', iter((4, 5, 6)))
    pcoll3 = pcoll2 | FlatMap('do', lambda x: [x + 10])
    assert_that(pcoll3, equal_to([14, 15, 16]), label='pcoll3')
    pipeline.run()

 def test_reuse_cloned_custom_transform_instance(self):
   pipeline = Pipeline(DirectPipelineRunner())
   pcoll1 = pipeline | Create('pcoll1', [1, 2, 3])
   pcoll2 = pipeline | Create('pcoll2', [4, 5, 6])
   transform = PipelineTest.CustomTransform()
   result1 = pcoll1 | transform
   result2 = pcoll2 | transform.clone('new label')
   assert_that(result1, equal_to([2, 3, 4]), label='r1')
   assert_that(result2, equal_to([5, 6, 7]), label='r2')
   pipeline.run()

  def test_cached_pvalues_are_refcounted(self):
    """Test that cached PValues are refcounted and deleted.

    The intermediary PValues computed by the workflow below contain
    one million elements so if the refcounting does not work the number of
    objects tracked by the garbage collector will increase by a few millions
    by the time we execute the final Map checking the objects tracked.
    Anything that is much larger than what we started with will fail the test.
    """
    def check_memory(value, count_threshold):
      gc.collect()
      objects_count = len(gc.get_objects())
      if objects_count > count_threshold:
        raise RuntimeError(
            'PValues are not refcounted: %s, %s' % (
                objects_count, count_threshold))
      return value

    def create_dupes(o, _):
      yield o
      yield SideOutputValue('side', o)

    pipeline = Pipeline('DirectPipelineRunner')

    gc.collect()
    count_threshold = len(gc.get_objects()) + 10000
    biglist = pipeline | Create('oom:create', ['x'] * 1000000)
    dupes = (
        biglist
        | Map('oom:addone', lambda x: (x, 1))
        | FlatMap('oom:dupes', create_dupes,
                  AsIter(biglist)).with_outputs('side', main='main'))
    result = (
        (dupes.side, dupes.main, dupes.side)
        | Flatten('oom:flatten')
        | CombinePerKey('oom:combine', sum)
        | Map('oom:check', check_memory, count_threshold))

    assert_that(result, equal_to([('x', 3000000)]))
    pipeline.run()
    self.assertEqual(
        pipeline.runner.debug_counters['element_counts'],
        {
            'oom:flatten': 3000000,
            ('oom:combine/GroupByKey/reify_windows', None): 3000000,
            ('oom:dupes/oom:dupes', 'side'): 1000000,
            ('oom:dupes/oom:dupes', None): 1000000,
            'oom:create': 1000000,
            ('oom:addone', None): 1000000,
            'oom:combine/GroupByKey/group_by_key': 1,
            ('oom:check', None): 1,
            'assert_that/singleton': 1,
            ('assert_that/Map(match)', None): 1,
            ('oom:combine/GroupByKey/group_by_window', None): 1,
            ('oom:combine/Combine/ParDo(CombineValuesDoFn)', None): 1})

 def test_empty_side_outputs(self):
   pipeline = Pipeline('DirectPipelineRunner')
   nums = pipeline | Create('Some Numbers', [1, 3, 5])
   results = nums | FlatMap(
       'ClassifyNumbers',
       lambda x: [x, SideOutputValue('even' if x % 2 == 0 else 'odd', x)]
   ).with_outputs('odd', 'even', main='main')
   assert_that(results.main, equal_to([1, 3, 5]))
   assert_that(results.odd, equal_to([1, 3, 5]), label='assert:odd')
   assert_that(results.even, equal_to([]), label='assert:even')
   pipeline.run()

  def test_empty_singleton_side_input(self):
    pipeline = Pipeline('DirectPipelineRunner')
    pcol = pipeline | Create('start', [1, 2])
    side = pipeline | Create('side', [])  # Empty side input.

    def my_fn(k, s):
      v = ('empty' if isinstance(s, EmptySideInput) else 'full')
      return [(k, v)]
    result = pcol | FlatMap('compute', my_fn, AsSingleton(side))
    assert_that(result, equal_to([(1, 'empty'), (2, 'empty')]))
    pipeline.run()

 def test_timestamped_value(self):
   p = Pipeline('DirectPipelineRunner')
   result = (p
             | Create('start', [(k, k) for k in range(10)])
             | Map(lambda (x, t): TimestampedValue(x, t))
             | WindowInto('w', FixedWindows(5))
             | Map(lambda v: ('key', v))
             | GroupByKey())
   assert_that(result, equal_to([('key', [0, 1, 2, 3, 4]),
                                 ('key', [5, 6, 7, 8, 9])]))
   p.run()