1 files changed, 13 insertions, 3 deletions

diff --git a/tensorflow/core/kernels/data/stats_dataset_ops.cc b/tensorflow/core/kernels/data/stats_dataset_ops.cc
index 3e0a6ae049..58ec3d4495 100644
--- a/tensorflow/core/kernels/data/stats_dataset_ops.cc
+++ b/tensorflow/core/kernels/data/stats_dataset_ops.cc
@@ -310,16 +310,20 @@ class FeatureStatsDatasetOp : public UnaryDatasetOpKernel {
           for (const Tensor& t : *out_tensors) {
             auto record_t = t.flat<string>();
             Example example;
-            // TODO(shivaniagrawal): redundant parsing here, potential solutions
+            // TODO(b/111553342): redundant parsing here, potential solutions
             // to improve performance is to a) have a potential
             // ParseExampleDataset and collect stats from there and b) make
             // changes to parse_example() where it returns stats as well.
             for (int i = 0; i < record_t.size(); ++i) {
               if (example.ParseFromString(record_t(i))) {
+                stats_aggregator->IncrementCounter("examples_count", "trainer",
+                                                   1);
                 AddStatsFeatures(example, stats_aggregator);
               } else {
                 SequenceExample sequence_example;
                 if (sequence_example.ParseFromString(record_t(i))) {
+                  stats_aggregator->IncrementCounter("sequence_examples_count",
+                                                     "trainer", 1);
                   AddStatsFeatures(sequence_example, stats_aggregator);
                 }
               }
@@ -329,7 +333,6 @@ class FeatureStatsDatasetOp : public UnaryDatasetOpKernel {
         return s;
       }
 
-      // TODO(shivaniagrawal): Add features/feature-values to streamz metrics.
       int AddStatsFeatureValues(const Feature& feature) {
         int feature_values_list_size = 0;
         switch (feature.kind_case()) {
@@ -360,8 +363,11 @@ class FeatureStatsDatasetOp : public UnaryDatasetOpKernel {
 
         int feature_values_list_size_sum = 0;
         for (const auto& feature : example.features().feature()) {
+          stats_aggregator->IncrementCounter("features_count", "trainer", 1);
           feature_values_list_size_sum += AddStatsFeatureValues(feature.second);
         }
+        stats_aggregator->IncrementCounter("feature_values_count", "trainer",
+                                           feature_values_list_size_sum);
         stats_aggregator->AddToHistogram(
             strings::StrCat(dataset()->tag_, ":feature-values"),
             {static_cast<double>(feature_values_list_size_sum)});
@@ -378,16 +384,20 @@ class FeatureStatsDatasetOp : public UnaryDatasetOpKernel {
 
         int feature_values_list_size_sum = 0;
         for (const auto& feature : example.context().feature()) {
+          stats_aggregator->IncrementCounter("features_count", "trainer", 1);
           feature_values_list_size_sum += AddStatsFeatureValues(feature.second);
         }
 
         for (const auto& feature_list :
              example.feature_lists().feature_list()) {
+          stats_aggregator->IncrementCounter("feature_lists_count", "trainer",
+                                             1);
           for (const auto& feature : feature_list.second.feature()) {
             feature_values_list_size_sum += AddStatsFeatureValues(feature);
           }
         }
-
+        stats_aggregator->IncrementCounter("feature_values_count", "trainer",
+                                           feature_values_list_size_sum);
         stats_aggregator->AddToHistogram(
             strings::StrCat(dataset()->tag_, ":feature-values"),
             {static_cast<double>(feature_values_list_size_sum)});