NAReplace estimator

src/Microsoft.ML.Data/Transforms/KeyToVectorTransform.cs

@@ -280,6 +280,7 @@ private ColInfo[] CreateInfos(ISchema inputSchema)
                     if (!inputSchema.TryGetColumnIndex(_parent.ColumnPairs[i].input, out int colSrc))
                         throw Host.ExceptSchemaMismatch(nameof(inputSchema), "input", _parent.ColumnPairs[i].input);
                     var type = inputSchema.GetColumnType(colSrc);
+                    _parent.CheckInputColumn(inputSchema, i, colSrc);
                     infos[i] = new ColInfo(_parent.ColumnPairs[i].output, _parent.ColumnPairs[i].input, type);
                 }
                 return infos;

src/Microsoft.ML.Data/Transforms/TermEstimator.cs

@@ -15,6 +15,7 @@ public sealed class TermEstimator : IEstimator<TermTransform>
     {
         private readonly IHost _host;
         private readonly TermTransform.ColumnInfo[] _columns;
+
         public TermEstimator(IHostEnvironment env, string name, string source = null, int maxNumTerms = TermTransform.Defaults.MaxNumTerms, TermTransform.SortOrder sort = TermTransform.Defaults.Sort) :
            this(env, new TermTransform.ColumnInfo(name, source ?? name, maxNumTerms, sort))
         {

src/Microsoft.ML.Transforms/KeyToBinaryVectorTransform.cs

@@ -69,6 +69,7 @@ private static (string input, string output)[] GetColumnPairs(ColumnInfo[] colum
             Contracts.CheckValue(columns, nameof(columns));
             return columns.Select(x => (x.Input, x.Output)).ToArray();
         }
+
         public IReadOnlyCollection<ColumnInfo> Columns => _columns.AsReadOnly();
         private readonly ColumnInfo[] _columns;
 
@@ -209,7 +210,7 @@ private ColInfo[] CreateInfos(ISchema inputSchema)
                     if (!inputSchema.TryGetColumnIndex(_parent.ColumnPairs[i].input, out int colSrc))
                         throw Host.ExceptSchemaMismatch(nameof(inputSchema), "input", _parent.ColumnPairs[i].input);
                     var type = inputSchema.GetColumnType(colSrc);
-
+                    _parent.CheckInputColumn(inputSchema, i, colSrc);
                     infos[i] = new ColInfo(_parent.ColumnPairs[i].output, _parent.ColumnPairs[i].input, type);
                 }
                 return infos;

src/Microsoft.ML.Transforms/NAHandleTransform.cs

@@ -135,7 +135,7 @@ public static IDataTransform Create(IHostEnvironment env, Arguments args, IDataV
             h.CheckValue(input, nameof(input));
             h.CheckUserArg(Utils.Size(args.Column) > 0, nameof(args.Column));
 
-            var replaceCols = new List<NAReplaceTransform.Column>();
+            var replaceCols = new List<NAReplaceTransform.ColumnInfo>();
             var naIndicatorCols = new List<NAIndicatorTransform.Column>();
             var naConvCols = new List<ConvertTransform.Column>();
             var concatCols = new List<ConcatTransform.TaggedColumn>();
@@ -149,14 +149,7 @@ public static IDataTransform Create(IHostEnvironment env, Arguments args, IDataV
                 var addInd = column.ConcatIndicator ?? args.Concat;
                 if (!addInd)
                 {
-                    replaceCols.Add(
-                        new NAReplaceTransform.Column()
-                        {
-                            Kind = (NAReplaceTransform.ReplacementKind?)column.Kind,
-                            Name = column.Name,
-                            Source = column.Source,
-                            Slot = column.ImputeBySlot
-                        });
+                    replaceCols.Add(new NAReplaceTransform.ColumnInfo(column.Source, column.Name, (NAReplaceTransform.ReplacementKind)(column.Kind ?? args.ReplaceWith), column.ImputeBySlot ?? args.ImputeBySlot));
                     continue;
                 }
 
@@ -186,14 +179,7 @@ public static IDataTransform Create(IHostEnvironment env, Arguments args, IDataV
                     naConvCols.Add(new ConvertTransform.Column() { Name = tmpIsMissingColName, Source = tmpIsMissingColName, ResultType = replaceType.ItemType.RawKind });
 
                 // Add the NAReplaceTransform column.
-                replaceCols.Add(
-                    new NAReplaceTransform.Column()
-                    {
-                        Kind = (NAReplaceTransform.ReplacementKind?)column.Kind,
-                        Name = tmpReplacementColName,
-                        Source = column.Source,
-                        Slot = column.ImputeBySlot
-                    });
+                replaceCols.Add(new NAReplaceTransform.ColumnInfo(column.Source, tmpReplacementColName, (NAReplaceTransform.ReplacementKind)(column.Kind ?? args.ReplaceWith), column.ImputeBySlot ?? args.ImputeBySlot));
 
                 // Add the ConcatTransform column.
                 if (replaceType.IsVector)
@@ -237,15 +223,8 @@ public static IDataTransform Create(IHostEnvironment env, Arguments args, IDataV
                 h.AssertValue(output);
                 output = new ConvertTransform(h, new ConvertTransform.Arguments() { Column = naConvCols.ToArray() }, output);
             }
-
             // Create the NAReplace transform.
-            output = new NAReplaceTransform(h,
-                new NAReplaceTransform.Arguments()
-                {
-                    Column = replaceCols.ToArray(),
-                    ReplacementKind = (NAReplaceTransform.ReplacementKind)args.ReplaceWith,
-                    ImputeBySlot = args.ImputeBySlot
-                }, output ?? input);
+            output = NAReplaceTransform.Create(env, output ?? input, replaceCols.ToArray());
 
             // Concat the NAReplaceTransform output and the NAIndicatorTransform output.
             if (naIndicatorCols.Count > 0)

src/Microsoft.ML.Transforms/NAHandling.cs

@@ -88,7 +88,7 @@ public static CommonOutputs.TransformOutput Indicator(IHostEnvironment env, NAIn
         public static CommonOutputs.TransformOutput Replace(IHostEnvironment env, NAReplaceTransform.Arguments input)
         {
             var h = EntryPointUtils.CheckArgsAndCreateHost(env, "NAReplace", input);
-            var xf = new NAReplaceTransform(h, input, input.Data);
+            var xf = NAReplaceTransform.Create(h, input, input.Data);
             return new CommonOutputs.TransformOutput()
             {
                 Model = new TransformModel(h, xf, input.Data),