diff --git a/src/Microsoft.ML.AutoML/CodeGen/map_value_to_key_search_space.json b/src/Microsoft.ML.AutoML/CodeGen/map_value_to_key_search_space.json
index 1663a95c18..6e9390f141 100644
--- a/src/Microsoft.ML.AutoML/CodeGen/map_value_to_key_search_space.json
+++ b/src/Microsoft.ML.AutoML/CodeGen/map_value_to_key_search_space.json
@@ -14,6 +14,10 @@
       "name": "AddKeyValueAnnotationsAsText",
       "type": "boolean",
       "default": false
+    },
+    {
+      "name": "KeyData",
+      "type": "dataView"
     }
   ]
 }
diff --git a/src/Microsoft.ML.AutoML/CodeGen/search-space-schema.json b/src/Microsoft.ML.AutoML/CodeGen/search-space-schema.json
index b780854c87..c0512afad1 100644
--- a/src/Microsoft.ML.AutoML/CodeGen/search-space-schema.json
+++ b/src/Microsoft.ML.AutoML/CodeGen/search-space-schema.json
@@ -240,7 +240,8 @@
         "predictedAnswerColumnName",
         "TopKAnswers",
         "TargetType",
-        "PredictionColumnName"
+        "PredictionColumnName",
+        "KeyData"
       ]
     },
     "option_type": {
@@ -259,7 +260,8 @@
         "dnnModelFactory",
         "bertArchitecture",
         "imageClassificationArchType",
-        "dataKind"
+        "dataKind",
+        "dataView"
       ]
     }
   },
diff --git a/src/Microsoft.ML.AutoML/CodeGen/type_converter_search_space.json b/src/Microsoft.ML.AutoML/CodeGen/type_converter_search_space.json
index 38d860e9c4..b512e32fc0 100644
--- a/src/Microsoft.ML.AutoML/CodeGen/type_converter_search_space.json
+++ b/src/Microsoft.ML.AutoML/CodeGen/type_converter_search_space.json
@@ -14,6 +14,10 @@
       "name": "TargetType",
       "type": "dataKind",
       "default": "DataKind.Single"
+    },
+    {
+      "name": "KeyData",
+      "type": "dataView"
     }
   ]
 }
diff --git a/src/Microsoft.ML.AutoML/SweepableEstimator/Estimators/MapValueToKey.cs b/src/Microsoft.ML.AutoML/SweepableEstimator/Estimators/MapValueToKey.cs
index a48827478f..decbd39109 100644
--- a/src/Microsoft.ML.AutoML/SweepableEstimator/Estimators/MapValueToKey.cs
+++ b/src/Microsoft.ML.AutoML/SweepableEstimator/Estimators/MapValueToKey.cs
@@ -8,7 +8,7 @@ internal partial class MapValueToKey
     {
         public override IEstimator<ITransformer> BuildFromOption(MLContext context, MapValueToKeyOption param)
         {
-            return context.Transforms.Conversion.MapValueToKey(param.OutputColumnName, param.InputColumnName, addKeyValueAnnotationsAsText: param.AddKeyValueAnnotationsAsText);
+            return context.Transforms.Conversion.MapValueToKey(param.OutputColumnName, param.InputColumnName, addKeyValueAnnotationsAsText: param.AddKeyValueAnnotationsAsText, keyData: param.KeyData);
         }
     }
 
diff --git a/tools-local/Microsoft.ML.AutoML.SourceGenerator/SearchSpaceGenerator.cs b/tools-local/Microsoft.ML.AutoML.SourceGenerator/SearchSpaceGenerator.cs
index 379d9e618e..85ed140019 100644
--- a/tools-local/Microsoft.ML.AutoML.SourceGenerator/SearchSpaceGenerator.cs
+++ b/tools-local/Microsoft.ML.AutoML.SourceGenerator/SearchSpaceGenerator.cs
@@ -57,6 +57,7 @@ public void Execute(GeneratorExecutionContext context)
                         "bertArchitecture" => "BertArchitecture",
                         "imageClassificationArchType" => "Microsoft.ML.Vision.ImageClassificationTrainer.Architecture",
                         "dataKind" => "Microsoft.ML.Data.DataKind",
+                        "dataView" => "Microsoft.ML.IDataView",
                         _ => throw new ArgumentException("unknown type"),
                     };
 
@@ -76,6 +77,7 @@ public void Execute(GeneratorExecutionContext context)
                         (_, "BertArchitecture") => defaultToken.GetValue<string>(),
                         (_, "Microsoft.ML.Vision.ImageClassificationTrainer.Architecture") => defaultToken.GetValue<string>(),
                         (_, "Microsoft.ML.Data.DataKind") => defaultToken.GetValue<string>(),
+                        (_, "Microsoft.ML.IDataView") => defaultToken.GetValue<string>(),
                         (_, _) => throw new ArgumentException("unknown"),
                     };