Fix bug around handling WeightedFeatures in validation.

PiperOrigin-RevId: 290172331

Author: tfx-copybara

tensorflow_data_validation/anomalies/schema.cc

@@ -41,6 +41,7 @@ limitations under the License.
 #include "tensorflow/core/platform/protobuf.h"
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow_metadata/proto/v0/anomalies.pb.h"
+#include "tensorflow_metadata/proto/v0/schema.pb.h"
 #include "tensorflow_metadata/proto/v0/statistics.pb.h"
 
 namespace tensorflow {
@@ -55,6 +56,7 @@ using ::tensorflow::metadata::v0::Feature;
 using ::tensorflow::metadata::v0::FeatureNameStatistics;
 using ::tensorflow::metadata::v0::SparseFeature;
 using ::tensorflow::metadata::v0::StringDomain;
+using ::tensorflow::metadata::v0::WeightedFeature;
 using PathProto = ::tensorflow::metadata::v0::Path;
 
 constexpr char kTrainingServingSkew[] = "Training/Serving skew";
@@ -205,6 +207,13 @@ tensorflow::Status Schema::UpdateFeature(
   Feature* feature = GetExistingFeature(feature_stats_view.GetPath());
   SparseFeature* sparse_feature =
       GetExistingSparseFeature(feature_stats_view.GetPath());
+  const WeightedFeature* weighted_feature =
+      GetExistingWeightedFeature(feature_stats_view.GetPath());
+  if (weighted_feature != nullptr) {
+    // TODO(b/141961105): Add validation logic for weighted features.
+    return Status::OK();
+  }
+
   if (sparse_feature != nullptr &&
       !::tensorflow::data_validation::SparseFeatureIsDeprecated(
           *sparse_feature)) {
@@ -554,7 +563,8 @@ tensorflow::metadata::v0::Schema Schema::GetSchema() const { return schema_; }
 
 bool Schema::FeatureExists(const Path& path) {
   return GetExistingFeature(path) != nullptr ||
-         GetExistingSparseFeature(path) != nullptr;
+         GetExistingSparseFeature(path) != nullptr ||
+         GetExistingWeightedFeature(path) != nullptr;
 }
 
 Feature* Schema::GetExistingFeature(const Path& path) {
@@ -595,6 +605,23 @@ SparseFeature* Schema::GetExistingSparseFeature(const Path& path) {
         parent_feature->mutable_struct_domain()->mutable_sparse_feature());
   }
 }
+
+const WeightedFeature* Schema::GetExistingWeightedFeature(
+    const Path& path) const {
+  CHECK(!path.empty());
+  if (path.size() != 1) {
+    // Weighted features are always top-level features with single-step paths.
+    return nullptr;
+  }
+  auto name = path.last_step();
+  for (const WeightedFeature& weighted_feature : schema_.weighted_feature()) {
+    if (weighted_feature.name() == name) {
+      return &weighted_feature;
+    }
+  }
+  return nullptr;
+}
+
 Feature* Schema::GetNewFeature(const Path& path) {
   CHECK(!path.empty());
   if (path.size() > 1) {

tensorflow_data_validation/anomalies/schema.h

@@ -168,6 +168,7 @@ class Schema {
   using Feature = tensorflow::metadata::v0::Feature;
   using SparseFeature = tensorflow::metadata::v0::SparseFeature;
   using StringDomain = tensorflow::metadata::v0::StringDomain;
+  using WeightedFeature = tensorflow::metadata::v0::WeightedFeature;
   // Updates Schema given new data, but only on the columns specified.
   // If you have a new, previously unseen column on the list of columns to
   // consider, then config is used to create it.
@@ -226,12 +227,15 @@ class Schema {
   // values if necessary.
   StringDomain* GetStringDomain(const string& name);
 
-  // Gets an existing feature, and returns null if it doesn't exist.
+  // Gets an existing feature, or returns null if it doesn't exist.
   Feature* GetExistingFeature(const Path& path);
 
-  // Gets an existing sparse feature, and returns null if it doesn't exist.
+  // Gets an existing sparse feature, or returns null if it doesn't exist.
   SparseFeature* GetExistingSparseFeature(const Path& path);
 
+  // Gets an existing weighted feature, or returns null if it doesn't exist.
+  const WeightedFeature* GetExistingWeightedFeature(const Path& path) const;
+
   // Gets a new feature. Assumes that the feature does not already exist.
   Feature* GetNewFeature(const Path& path);
 

tensorflow_data_validation/api/validation_api_test.py

@@ -857,6 +857,59 @@ def test_validate_stats(self):
     anomalies = validation_api.validate_statistics(statistics, schema)
     self._assert_equal_anomalies(anomalies, expected_anomalies)
 
+  def test_validate_stats_weighted_feature(self):
+    # This test is not intended to verify the anomaly detection logic, but just
+    # ensure that the validation API will not fail on custom stats for weighted
+    # features.
+    schema = text_format.Parse(
+        """
+        feature {
+          name: "f"
+        }
+        feature {
+          name: "w"
+        }
+        weighted_feature {
+          name: "weighted_feature"
+          feature {
+            step: "f"
+          }
+          weight_feature {
+            step: "w"
+          }
+        }
+        """, schema_pb2.Schema())
+    statistics = text_format.Parse(
+        """
+        datasets {
+          num_examples: 10
+          features {
+            path { step: 'weighted_feature' }
+            custom_stats {
+              name: 'missing_weight'
+              num: 0.0
+            }
+            custom_stats {
+              name: 'missing_value'
+              num: 0.0
+            }
+            custom_stats {
+              name: 'min_weight_length_diff'
+              num: 0.0
+            }
+            custom_stats {
+              name: 'max_weight_length_diff'
+              num: 0.0
+            }
+          }
+        }
+        """, statistics_pb2.DatasetFeatureStatisticsList())
+    expected_anomalies = {}
+
+    # Validate the stats.
+    anomalies = validation_api.validate_statistics(statistics, schema)
+    self._assert_equal_anomalies(anomalies, expected_anomalies)
+
   # pylint: disable=line-too-long
   _annotated_enum_anomaly_info = """
             path {