re-enable histograms

Author: jrmccluskey

sdks/python/apache_beam/runners/dataflow/dataflow_runner.py

@@ -759,6 +759,8 @@ def api_jobstate_to_pipeline_state(api_jobstate):
             values_enum.JOB_STATE_CANCELLING: PipelineState.CANCELLING,
             values_enum.JOB_STATE_RESOURCE_CLEANING_UP: PipelineState.
             RESOURCE_CLEANING_UP,
+            # values_enum.JOB_STATE_PAUSING : PipelineState.PAUSING,
+            # values_enum.JOB_STATE_PAUSED : PipelineState.PAUSED,
         })
 
     return (

sdks/python/apache_beam/runners/dataflow/internal/apiclient.py

@@ -69,6 +69,7 @@
 from apache_beam.runners.internal import names as shared_names
 from apache_beam.runners.pipeline_utils import validate_pipeline_graph
 from apache_beam.runners.portability.stager import Stager
+from apache_beam.transforms import DataflowDistributionCounter
 from apache_beam.transforms import cy_combiners
 from apache_beam.transforms.display import DisplayData
 from apache_beam.transforms.environments import is_apache_beam_container
@@ -1109,6 +1110,12 @@ def get_sdk_package_name():
     return shared_names.BEAM_PACKAGE_NAME
 
 
+class Histogram(object):
+  def __init__(self):
+    self.firstBucketOffset = None
+    self.bucketCounts = None
+
+
 class DataflowJobAlreadyExistsError(retry.PermanentException):
   """A non-retryable exception that a job with the given name already exists."""
   # Inherits retry.PermanentException to avoid retry in
@@ -1133,9 +1140,11 @@ def translate_distribution(
   dist_update_proto.update({"count": distribution_update.count})
   dist_update_proto.update({"sum": distribution_update.sum})
   # DataflowDistributionCounter needs to translate histogram
-  # if isinstance(distribution_update, DataflowDistributionCounter):
-  #   dist_update_proto.histogram = dataflow.Histogram()
-  #   distribution_update.translate_to_histogram(dist_update_proto.histogram)
+  if isinstance(distribution_update, DataflowDistributionCounter):
+    histogram = Histogram()
+    distribution_update.translate_to_histogram(histogram)
+    dist_update_proto.update({"firstBucketOffset": histogram.firstBucketOffset})
+    dist_update_proto.update({"bucketCounts": histogram.bucketCounts})
   metric_update_proto.distribution = dist_update_proto
 
 

sdks/python/apache_beam/runners/dataflow/internal/apiclient_test.py

@@ -39,6 +39,7 @@
 from apache_beam.portability.api import beam_runner_api_pb2
 from apache_beam.runners.dataflow.internal import names
 from apache_beam.transforms import Create
+from apache_beam.transforms import DataflowDistributionCounter
 from apache_beam.transforms import DoFn
 from apache_beam.transforms import ParDo
 from apache_beam.transforms.environments import DockerEnvironment
@@ -492,25 +493,25 @@ def test_translate_distribution_using_distribution_data(self):
         metric_update.distribution['count'], distribution_update.count)
 
 
-#   def test_translate_distribution_using_dataflow_distribution_counter(self):
-#     counter_update = DataflowDistributionCounter()
-#     counter_update.add_input(1)
-#     counter_update.add_input(3)
-#     metric_proto = dataflow.MetricUpdate()
-#     apiclient.translate_distribution(counter_update, metric_proto)
-#     histogram = mock.Mock(firstBucketOffset=None, bucketCounts=None)
-#     counter_update.translate_to_histogram(histogram)
-#     self.assertEqual(metric_proto.distribution['min'], counter_update.min)
-#     self.assertEqual(metric_proto.distribution['max'], counter_update.max)
-#     self.assertEqual(metric_proto.distribution['sum'], counter_update.sum)
-#     self.assertEqual(
-#         metric_proto.distribution['count'], counter_update.count)
-#     self.assertEqual(
-#         metric_proto.distribution.histogram.bucketCounts,
-#         histogram.bucketCounts)
-#     self.assertEqual(
-#         metric_proto.distribution.histogram.firstBucketOffset,
-#         histogram.firstBucketOffset)
+  def test_translate_distribution_using_dataflow_distribution_counter(self):
+    counter_update = DataflowDistributionCounter()
+    counter_update.add_input(1)
+    counter_update.add_input(3)
+    metric_proto = dataflow.MetricUpdate()
+    apiclient.translate_distribution(counter_update, metric_proto)
+    histogram = mock.Mock(firstBucketOffset=None, bucketCounts=None)
+    counter_update.translate_to_histogram(histogram)
+    self.assertEqual(metric_proto.distribution['min'], counter_update.min)
+    self.assertEqual(metric_proto.distribution['max'], counter_update.max)
+    self.assertEqual(metric_proto.distribution['sum'], counter_update.sum)
+    self.assertEqual(
+        metric_proto.distribution['count'], counter_update.count)
+    self.assertEqual(
+        metric_proto.distribution['firstBucketOffset'],
+        histogram.firstBucketOffset)
+    self.assertEqual(
+        metric_proto.distribution['bucketCounts'],
+        histogram.bucketCounts)
 
   def test_translate_means(self):
     metric_update = dataflow.MetricUpdate()