Use KeyTypeAttribute from Schema in CreateTextLoader (#5082)

* Read KeyTypeAttribute during CreateTextLoader

* Changed to using Assert.Equal

Author: aslotte

src/Microsoft.ML.Data/DataLoadSave/Text/TextLoader.cs

@@ -1489,6 +1489,13 @@ internal static TextLoader CreateTextLoader<TInput>(IHostEnvironment host,
                 var column = new Column();
                 column.Name = mappingAttrName?.Name ?? memberInfo.Name;
                 column.Source = mappingAttr.Sources.ToArray();
+
+                var keyTypeAttr = memberInfo.GetCustomAttribute<KeyTypeAttribute>();
+                if (keyTypeAttr != null)
+                {
+                    column.KeyCount = keyTypeAttr.KeyCount;
+                }
+
                 InternalDataKind dk;
                 switch (memberInfo)
                 {

test/Microsoft.ML.Tests/TextLoaderTests.cs

@@ -877,5 +877,48 @@ public void TestTextLoaderNoFields()
                 Assert.StartsWith("Should define at least one public, readable field or property in TInput.", ex.Message);
             }
         }
+
+        public class BreastCancerInputModelWithKeyType
+        {
+            [LoadColumn(0)]
+            public bool IsMalignant { get; set; }
+
+            [LoadColumn(1), KeyType(10)]
+            public uint Thickness { get; set; }
+        }
+
+        public class BreastCancerInputModelWithoutKeyType
+        {
+            [LoadColumn(0)]
+            public bool IsMalignant { get; set; }
+
+            [LoadColumn(1)]
+            public uint Thickness { get; set; }
+        }
+
+        [Fact]
+        public void TestLoadTextWithKeyTypeAttribute()
+        {
+            ulong expectedCount = 10;
+
+            var mlContext = new MLContext(seed: 1);
+            string breastCancerPath = GetDataPath(TestDatasets.breastCancer.trainFilename);
+
+            var data = mlContext.Data.CreateTextLoader<BreastCancerInputModelWithKeyType>(separatorChar: ',').Load(breastCancerPath);
+
+            Assert.Equal(expectedCount, data.Schema[1].Type.GetKeyCount());
+        }
+
+        [Fact]
+        public void TestLoadTextWithoutKeyTypeAttribute()
+        {
+            ulong expectedCount = 0;
+            var mlContext = new MLContext(seed: 1);
+            string breastCancerPath = GetDataPath(TestDatasets.breastCancer.trainFilename);
+
+            var data = mlContext.Data.CreateTextLoader<BreastCancerInputModelWithoutKeyType>(separatorChar: ',').Load(breastCancerPath);
+
+            Assert.Equal(expectedCount, data.Schema[1].Type.GetKeyCount());
+        }
     }
 }