ColumnName indexer on DataFrame (dotnet#2959)

* ColumnName indexer on DataFrame Fixes dotnet/corefxlab#2934 * Unit tests * Null column name
pgovind · Sep 2, 2020 · 7ebe8bc · 7ebe8bc
1 parent 6c2d800
commit 7ebe8bc
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 5 deletions.
diff --git a/src/Microsoft.Data.Analysis/DataFrame.cs b/src/Microsoft.Data.Analysis/DataFrame.cs
@@ -147,6 +147,18 @@ public DataFrame this[IEnumerable<bool> rowFilter]
             }
         }
 
+        /// <summary>
+        /// An indexer based on <see cref="DataFrameColumn.Name"/>
+        /// </summary>
+        /// <param name="columnName">The name of a <see cref="DataFrameColumn"/></param>
+        /// <returns>A <see cref="DataFrameColumn"/> if it exists.</returns>
+        /// <exception cref="ArgumentException">Throws if <paramref name="columnName"/> is not present in this <see cref="DataFrame"/></exception>
+        public DataFrameColumn this[string columnName]
+        {
+            get => Columns[columnName];
+            set => Columns[columnName] = value;
+        }
+
         /// <summary>
         /// Returns the first <paramref name="numberOfRows"/> rows
         /// </summary>

diff --git a/src/Microsoft.Data.Analysis/DataFrameColumnCollection.cs b/src/Microsoft.Data.Analysis/DataFrameColumnCollection.cs
@@ -133,7 +133,7 @@ public void Remove(string columnName)
         /// <param name="columnName"></param>
         public int IndexOf(string columnName)
         {
-            if (_columnNameToIndexDictionary.TryGetValue(columnName, out int columnIndex))
+            if (columnName != null && _columnNameToIndexDictionary.TryGetValue(columnName, out int columnIndex))
             {
                 return columnIndex;
             }

diff --git a/tests/Microsoft.Data.Analysis.Tests/DataFrameTests.cs b/tests/Microsoft.Data.Analysis.Tests/DataFrameTests.cs
@@ -200,10 +200,11 @@ public void TestIndexer()
             var row = dataFrame.Rows[4];
             Assert.Equal(14, (int)row[1]);
 
-            var column = dataFrame.Columns["Int2"] as Int32DataFrameColumn;
+            var column = dataFrame["Int2"] as Int32DataFrameColumn;
             Assert.Equal(1000, (int)column[2]);
 
-            Assert.Throws<ArgumentException>(() => dataFrame.Columns["Int5"]);
+            Assert.Throws<ArgumentException>(() => dataFrame["Int5"]);
+            Assert.Throws<ArgumentException>(() => dataFrame[(string)null]);
         }
 
         [Fact]
@@ -645,15 +646,15 @@ public void TestProjectionAndAppend()
         {
             DataFrame df = MakeDataFrameWithTwoColumns(10);
 
-            df.Columns["Int3"] = df.Columns["Int1"] * 2 + df.Columns["Int2"];
+            df["Int3"] = df.Columns["Int1"] * 2 + df.Columns["Int2"];
             Assert.Equal(16, df.Columns["Int3"][2]);
         }
 
         [Fact]
         public void TestComputations()
         {
             DataFrame df = MakeDataFrameWithAllMutableColumnTypes(10);
-            df.Columns["Int"][0] = -10;
+            df["Int"][0] = -10;
             Assert.Equal(-10, df.Columns["Int"][0]);
 
             DataFrameColumn absColumn = df.Columns["Int"].Abs();