From 39235a76d98666053ebe16e2b98abdf46d849e9a Mon Sep 17 00:00:00 2001 From: Aleksei Smirnov Date: Sat, 26 Aug 2023 00:59:41 +0300 Subject: [PATCH] Fix issue with addIndexColumn in DataFrame.LoadCsv (#6769) * Fix issue with addIndexColumn in DataFrame.LoadCsv * Fix tests --- src/Microsoft.Data.Analysis/DataFrame.IO.cs | 7 +++---- .../DataFrame.IOTests.cs | 13 +++++++++++++ 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/src/Microsoft.Data.Analysis/DataFrame.IO.cs b/src/Microsoft.Data.Analysis/DataFrame.IO.cs index 44448da008..247df0a57e 100644 --- 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 indexColumn = new PrimitiveDataFrameColumn("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 index 05565673b0..f2a1182b59 100644 --- 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() {