Check for null Label, Feature, and GroupId columns (dotnet#4265)

najeeb-kazmi · web-flow · commit 4944be7e970e · 2019-10-01T16:16:15.000-07:00
diff --git a/src/Microsoft.ML.EntryPoints/PermutationFeatureImportance.cs b/src/Microsoft.ML.EntryPoints/PermutationFeatureImportance.cs
@@ -62,6 +62,8 @@ internal static IDataView GetMetrics(
             RoleMappedData roleMappedData,
             PermutationFeatureImportanceArguments input)
         {
+            Contracts.Check(roleMappedData.Schema.Feature != null, "Feature column not found.");
+            Contracts.Check(roleMappedData.Schema.Label != null, "Label column not found.");
             IDataView result;
             if (predictor.PredictionKind == PredictionKind.BinaryClassification)
                 result = GetBinaryMetrics(env, predictor, roleMappedData, input);
@@ -252,6 +254,7 @@ private static IDataView GetRankingMetrics(
             RoleMappedData roleMappedData,
             PermutationFeatureImportanceArguments input)
         {
+            Contracts.Check(roleMappedData.Schema.Group != null, "Group ID column not found.");
             var roles = roleMappedData.Schema.GetColumnRoleNames();
             var featureColumnName = roles.Where(x => x.Key.Value == RoleMappedSchema.ColumnRole.Feature.Value).First().Value;
             var labelColumnName = roles.Where(x => x.Key.Value == RoleMappedSchema.ColumnRole.Label.Value).First().Value;
