Implement aggregated_values for DirectPipelineRunner

Create new class DirectPipelineResult, on which new method
aggregated_values() is defined.  This method lets you query the
accumulated value of an Aggregator after the pipeline has run.
----Release Notes----
[]
-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=120006979

Author: gildea

google/cloud/dataflow/examples/wordcount.py

@@ -24,6 +24,8 @@
 
 
 empty_line_aggregator = df.Aggregator('emptyLines')
+average_word_size_aggregator = df.Aggregator('averageWordLength',
+                                             df.combiners.Mean())
 
 
 class WordExtractingDoFn(df.DoFn):
@@ -43,7 +45,10 @@ def process(self, context):
     text_line = context.element.strip()
     if not text_line:
       context.aggregate_to(empty_line_aggregator, 1)
-    return re.findall(r'[A-Za-z\']+', text_line)
+    words = re.findall(r'[A-Za-z\']+', text_line)
+    for w in words:
+      context.aggregate_to(average_word_size_aggregator, float(len(w)))
+    return words
 
 
 def run(argv=None):
@@ -81,7 +86,11 @@ def run(argv=None):
   output | df.io.Write('write', df.io.TextFileSink(known_args.output))
 
   # Actually run the pipeline (all operations above are deferred).
-  p.run()
+  result = p.run()
+  empty_line_values = result.aggregated_values(empty_line_aggregator)
+  logging.info('number of empty lines: %d', sum(empty_line_values.values()))
+  word_length_values = result.aggregated_values(average_word_size_aggregator)
+  logging.info('average word lengths: %s', word_length_values.values())
 
 
 if __name__ == '__main__':

google/cloud/dataflow/runners/direct_runner.py

@@ -32,7 +32,9 @@
 from google.cloud.dataflow.pvalue import EmptySideInput
 from google.cloud.dataflow.runners.common import DoFnRunner
 from google.cloud.dataflow.runners.common import DoFnState
+from google.cloud.dataflow.runners.runner import PipelineResult
 from google.cloud.dataflow.runners.runner import PipelineRunner
+from google.cloud.dataflow.runners.runner import PipelineState
 from google.cloud.dataflow.runners.runner import PValueCache
 from google.cloud.dataflow.transforms import DoFnProcessContext
 from google.cloud.dataflow.transforms.window import GlobalWindows
@@ -89,6 +91,8 @@ def func_wrapper(self, pvalue, *args, **kwargs):
   def run(self, pipeline, node=None):
     super(DirectPipelineRunner, self).run(pipeline, node)
     logging.info('Final: Debug counters: %s', self.debug_counters)
+    return DirectPipelineResult(state=PipelineState.DONE,
+                                counter_factory=self._counter_factory)
 
   @skip_if_cached
   def run_ParDo(self, transform_node):
@@ -226,3 +230,14 @@ def run__NativeWrite(self, transform_node):
       for v in self._cache.get_pvalue(transform_node.inputs[0]):
         self.debug_counters['element_counts'][transform_node.full_label] += 1
         writer.Write(v.value)
+
+
+class DirectPipelineResult(PipelineResult):
+  """A DirectPipelineResult provides access to info about a pipeline."""
+
+  def __init__(self, state, counter_factory=None):
+    super(DirectPipelineResult, self).__init__(state)
+    self._counter_factory = counter_factory
+
+  def aggregated_values(self, aggregator_or_name):
+    return self._counter_factory.get_aggregator_values(aggregator_or_name)

google/cloud/dataflow/runners/runner.py

@@ -83,7 +83,6 @@ def visit_transform(self, transform_node):
           raise
 
     pipeline.visit(RunVisitor(self), node=node)
-    return PipelineResult(state=PipelineState.DONE)
 
   def clear(self, pipeline, node=None):
     """Clear all nodes or nodes reachable from node of materialized values.
@@ -267,3 +266,10 @@ def __init__(self, state):
   def current_state(self):
     """Return the current state of running the pipeline."""
     return self._state
+
+  # pylint: disable=unused-argument
+  def aggregated_values(self, aggregator_or_name):
+    """Return a dict of step names to values of the Aggregator."""
+    logging.warn('%s does not implement aggregated_values',
+                 self.__class__.__name__)
+    return {}

google/cloud/dataflow/utils/counters.py

@@ -89,6 +89,15 @@ def _update_small(self, delta):
   def total(self):
     return self.c_total + self.py_total
 
+  def value(self):
+    if self.aggregation_kind == self.SUM:
+      return self.total
+    elif self.aggregation_kind == self.MEAN:
+      return float(self.total)/self.elements
+    else:
+      # This can't happen, because we check in __init__
+      raise TypeError('%s.value(): unsupported aggregation_kind' % self)
+
   def __str__(self):
     return '<%s>' % self._str_internal()
 
@@ -125,6 +134,10 @@ def __init__(self, name='unnamed'):
                                       Counter.SUM)
 
 
+# Counters that represent Accumulators have names starting with this
+USER_COUNTER_PREFIX = 'user-'
+
+
 class CounterFactory(object):
   """Keeps track of unique counters."""
 
@@ -168,7 +181,7 @@ def get_aggregator_counter(self, step_name, aggregator):
       A new or existing counter.
     """
     with self._lock:
-      name = 'user-%s-%s' % (step_name, aggregator.name)
+      name = '%s%s-%s' % (USER_COUNTER_PREFIX, step_name, aggregator.name)
       aggregation_kind = aggregator.aggregation_kind
       counter = self.counters.get(name, None)
       if counter:
@@ -190,3 +203,32 @@ def get_counters(self):
     """
     with self._lock:
       return self.counters.values()
+
+  def get_aggregator_values(self, aggregator_or_name):
+    """Returns dict of step names to values of the aggregator."""
+    with self._lock:
+      return get_aggregator_values(
+          aggregator_or_name, self.counters, lambda counter: counter.value())
+
+
+def get_aggregator_values(aggregator_or_name, counter_dict,
+                          value_extractor=None):
+  """Extracts the named aggregator value from a set of counters.
+
+  Args:
+    aggregator_or_name: an Aggregator object or the name of one.
+    counter_dict: a dict object of {name: value_wrapper}
+    value_extractor: a function to convert the value_wrapper into a value.
+      If None, no extraction is done and the value is return unchanged.
+
+  Returns:
+    dict of step names to values of the aggregator.
+  """
+  name = aggregator_or_name
+  if value_extractor is None:
+    value_extractor = lambda x: x
+  if not isinstance(aggregator_or_name, basestring):
+    name = aggregator_or_name.name
+    return {n: value_extractor(c) for n, c in counter_dict.iteritems()
+            if n.startswith(USER_COUNTER_PREFIX)
+            and n.endswith('-%s' % name)}