Add support for string vectors

dotnet · Apr 21, 2023 · f7b8346 · f7b8346
1 parent fc67fd1
commit f7b8346
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 1 deletion.
diff --git a/src/Microsoft.Data.Analysis/IDataView.Extension.cs b/src/Microsoft.Data.Analysis/IDataView.Extension.cs
@@ -204,6 +204,10 @@ private static DataFrameColumn GetVectorDataFrame(VectorDataViewType vectorType,
             {
                 return new VBufferDataFrameColumn<decimal>(name);
             }
+            else if (itemType.RawType == typeof(ReadOnlyMemory<char>))
+            {
+                return new VBufferDataFrameColumn<ReadOnlyMemory<char>>(name);
+            }
 
             throw new NotSupportedException(String.Format(Microsoft.Data.Strings.VectorSubTypeNotSupported, itemType.ToString()));
         }

diff --git a/src/Microsoft.Data.Analysis/VBufferDataFrameColumn.cs b/src/Microsoft.Data.Analysis/VBufferDataFrameColumn.cs
@@ -390,6 +390,10 @@ private static VectorDataViewType GetDataViewType()
             {
                 return new VectorDataViewType(NumberDataViewType.Double);
             }
+            else if (typeof(T) == typeof(ReadOnlyMemory<char>))
+            {
+                return new VectorDataViewType(TextDataViewType.Instance);
+            }
 
             throw new NotSupportedException();
         }

diff --git a/test/Microsoft.Data.Analysis.Tests/DataFrameIDataViewTests.cs b/test/Microsoft.Data.Analysis.Tests/DataFrameIDataViewTests.cs
@@ -461,6 +461,7 @@ public void TestDataFrameFromIDataView_VBufferType()
                     ushortFeatures = new ushort[] {0, 0},
                     uintFeatures = new uint[] {0, 0},
                     ulongFeatures = new ulong[] {0, 0},
+                    stringFeatures = new string[]{ "A", "B"},
                 },
                 new {
                     boolFeature = new bool[] {false, false},
@@ -474,13 +475,14 @@ public void TestDataFrameFromIDataView_VBufferType()
                     ushortFeatures = new ushort[] {0, 0},
                     uintFeatures = new uint[] {0, 0},
                     ulongFeatures = new ulong[] {0, 0},
+                    stringFeatures = new string[]{ "A", "B"},
                 }
             };
 
             var data = mlContext.Data.LoadFromEnumerable(inputData);
             var df = data.ToDataFrame();
 
-            Assert.Equal(11, df.Columns.Count);
+            Assert.Equal(12, df.Columns.Count);
             Assert.Equal(2, df.Rows.Count);
         }
     }