Fix issue with addIndexColumn in DataFrame.LoadCsv (#6769)

* Fix issue with addIndexColumn in DataFrame.LoadCsv * Fix tests
dotnet · Aug 25, 2023 · 39235a7 · 39235a7
1 parent 43a6a81
commit 39235a7
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 4 deletions.
diff --git a/src/Microsoft.Data.Analysis/DataFrame.IO.cs b/src/Microsoft.Data.Analysis/DataFrame.IO.cs
@@ -439,14 +439,13 @@ private static DataFrame ReadCsvLinesIntoDataFrame(WrappedStreamReaderOrStringRe
 
                 if (addIndexColumn)
                 {
-                    PrimitiveDataFrameColumn<int> indexColumn = new PrimitiveDataFrameColumn<int>("IndexColumn", columns[0].Length);
-                    for (int i = 0; i < columns[0].Length; i++)
+                    Int64DataFrameColumn indexColumn = new Int64DataFrameColumn("IndexColumn", columns[0].Length);
+                    for (long i = 0; i < columns[0].Length; i++)
                     {
                         indexColumn[i] = i;
                     }
-                    columns.Insert(0, indexColumn);
+                    ret.Columns.Insert(0, indexColumn);
                 }
-
             }
 
             return ret;

diff --git a/test/Microsoft.Data.Analysis.Tests/DataFrame.IOTests.cs b/test/Microsoft.Data.Analysis.Tests/DataFrame.IOTests.cs
@@ -622,6 +622,19 @@ public void TestReadCsvWithRepeatColumnNameInHeader()
             Assert.Matches(@"DataFrame already contains a column called Column( \(Parameter 'column'\)|\r\nParameter name: column)", exp.Message);
         }
 
+        [Fact]
+        public void TestLoadCsvWithAddIndexColumn()
+        {
+            var dataFrame = DataFrame.LoadCsvFromString("11\r\n22\r\n33", header: false, addIndexColumn: true);
+
+            Assert.Equal(2, dataFrame.Columns.Count);
+            Assert.Equal("IndexColumn", dataFrame.Columns[0].Name);
+            Assert.Equal(3, dataFrame.Columns[0].Length);
+
+            for (long i = 0; i < dataFrame.Columns[0].Length; i++)
+                Assert.Equal(i, dataFrame.Columns[0][i]);
+        }
+
         [Fact]
         public void TestReadCsvWithExtraColumnInRow()
         {