Fix DataFrame Merge issue

asmirnov82 · asmirnov82 · commit 4ebddeb1b5ae · 2023-05-21T12:19:56.000+03:00
# Conflicts:
#	src/Microsoft.Data.Analysis/BooleanDataFrameColumn.cs
#	src/Microsoft.Data.Analysis/ByteDataFrameColumn.cs
#	src/Microsoft.Data.Analysis/CharDataFrameColumn.cs
#	src/Microsoft.Data.Analysis/DateTimeDataFrameColumn.cs
#	src/Microsoft.Data.Analysis/DecimalDataFrameColumn.cs
#	src/Microsoft.Data.Analysis/DoubleDataFrameColumn.cs
#	src/Microsoft.Data.Analysis/Int16DataFrameColumn.cs
#	src/Microsoft.Data.Analysis/Int32DataFrameColumn.cs
#	src/Microsoft.Data.Analysis/Int64DataFrameColumn.cs
#	src/Microsoft.Data.Analysis/SByteDataFrameColumn.cs
#	src/Microsoft.Data.Analysis/SingleDataFrameColumn.cs
#	src/Microsoft.Data.Analysis/UInt16DataFrameColumn.cs
#	src/Microsoft.Data.Analysis/UInt32DataFrameColumn.cs
#	src/Microsoft.Data.Analysis/UInt64DataFrameColumn.cs
#	test/Microsoft.Data.Analysis.Tests/DataFrame.IOTests.cs
diff --git a/src/Microsoft.Data.Analysis/DataFrameColumnCollection.cs b/src/Microsoft.Data.Analysis/DataFrameColumnCollection.cs
@@ -220,6 +220,23 @@ public PrimitiveDataFrameColumn<T> GetPrimitiveColumn<T>(string name)
             throw new ArgumentException(string.Format(Strings.BadColumnCast, column.DataType, typeof(T)), nameof(T));
         }
 
+        /// <summary>
+        /// Gets the <see cref="PrimitiveDataFrameColumn{T}"/> with the specified <paramref name="name"/>.
+        /// </summary>
+        /// <param name="name">The name of the column</param>
+        /// <returns><see cref="PrimitiveDataFrameColumn{T}"/>.</returns>
+        /// <exception cref="ArgumentException">A column named <paramref name="name"/> cannot be found, or if the column's type doesn't match.</exception>
+        public PrimitiveDataFrameColumn<DateTime> GetDateTimeColumn(string name)
+        {
+            DataFrameColumn column = this[name];
+            if (column is PrimitiveDataFrameColumn<DateTime> ret)
+            {
+                return ret;
+            }
+
+            throw new ArgumentException(string.Format(Strings.BadColumnCast, column.DataType, typeof(DateTime)));
+        }
+
         /// <summary>
         /// Gets the <see cref="ArrowStringDataFrameColumn"/> with the specified <paramref name="name"/>.
         /// </summary>
diff --git a/src/Microsoft.Data.Analysis/PrimitiveDataFrameColumn.cs b/src/Microsoft.Data.Analysis/PrimitiveDataFrameColumn.cs
@@ -220,6 +220,16 @@ protected override IReadOnlyList<object> GetValues(long startIndex, int length)
             return ret;
         }
 
+        internal virtual PrimitiveDataFrameColumn<T> CreateNewColumn(string name, PrimitiveColumnContainer<T> container)
+        {
+            return new PrimitiveDataFrameColumn<T>(name, container);
+        }
+
+        protected virtual PrimitiveDataFrameColumn<T> CreateNewColumn(string name, long length = 0)
+        {
+            return new PrimitiveDataFrameColumn<T>(name, length);
+        }
+
         internal T? GetTypedValue(long rowIndex) => _columnContainer[rowIndex];
 
         protected override object GetValue(long rowIndex) => GetTypedValue(rowIndex);
@@ -411,7 +421,7 @@ private PrimitiveDataFrameColumn<T> Clone(PrimitiveDataFrameColumn<bool> boolCol
         {
             if (boolColumn.Length > Length)
                 throw new ArgumentException(Strings.MapIndicesExceedsColumnLenth, nameof(boolColumn));
-            PrimitiveDataFrameColumn<T> ret = new PrimitiveDataFrameColumn<T>(Name);
+            PrimitiveDataFrameColumn<T> ret = CreateNewColumn(Name);
             for (long i = 0; i < boolColumn.Length; i++)
             {
                 bool? value = boolColumn[i];
@@ -438,7 +448,8 @@ private PrimitiveDataFrameColumn<T> CloneImplementation<U>(PrimitiveDataFrameCol
             }
             else
                 throw new NotImplementedException();
-            PrimitiveDataFrameColumn<T> ret = new PrimitiveDataFrameColumn<T>(Name, retContainer);
+
+            PrimitiveDataFrameColumn<T> ret = CreateNewColumn(Name, retContainer);
             return ret;
         }
 
@@ -447,7 +458,7 @@ public PrimitiveDataFrameColumn<T> Clone(PrimitiveDataFrameColumn<long> mapIndic
             if (mapIndices is null)
             {
                 PrimitiveColumnContainer<T> newColumnContainer = _columnContainer.Clone();
-                return new PrimitiveDataFrameColumn<T>(Name, newColumnContainer);
+                return CreateNewColumn(Name, newColumnContainer);
             }
             else
             {
diff --git a/test/Microsoft.Data.Analysis.Tests/DataFrameTests.cs b/test/Microsoft.Data.Analysis.Tests/DataFrameTests.cs
@@ -2766,6 +2766,56 @@ public void TestMerge_Issue5778()
             MatchRowsOnMergedDataFrame(merge, left, right, 1, 1, 0);
         }
 
+        [Fact]
+        //Issue 6127
+        public void TestMerge_CorrectColumnTypes()
+        {
+            DataFrame left = MakeDataFrameWithAllMutableColumnTypes(2, false);
+            DataFrame right = MakeDataFrameWithAllMutableColumnTypes(1);
+
+            DataFrame merge = left.Merge<int>(right, "Int", "Int");
+
+            Assert.NotNull(merge.Columns.GetBooleanColumn("Bool_left"));
+            Assert.NotNull(merge.Columns.GetBooleanColumn("Bool_right"));
+
+            Assert.NotNull(merge.Columns.GetDecimalColumn("Decimal_left"));
+            Assert.NotNull(merge.Columns.GetDecimalColumn("Decimal_right"));
+
+            Assert.NotNull(merge.Columns.GetSingleColumn("Float_left"));
+            Assert.NotNull(merge.Columns.GetSingleColumn("Float_right"));
+
+            Assert.NotNull(merge.Columns.GetDoubleColumn("Double_left"));
+            Assert.NotNull(merge.Columns.GetDoubleColumn("Double_right"));
+
+            Assert.NotNull(merge.Columns.GetByteColumn("Byte_left"));
+            Assert.NotNull(merge.Columns.GetByteColumn("Byte_right"));
+
+            Assert.NotNull(merge.Columns.GetCharColumn("Char_left"));
+            Assert.NotNull(merge.Columns.GetCharColumn("Char_right"));
+
+            Assert.NotNull(merge.Columns.GetInt16Column("Short_left"));
+            Assert.NotNull(merge.Columns.GetInt16Column("Short_right"));
+
+            Assert.NotNull(merge.Columns.GetUInt16Column("Ushort_left"));
+            Assert.NotNull(merge.Columns.GetUInt16Column("Ushort_right"));
+
+            Assert.NotNull(merge.Columns.GetInt32Column("Int_left"));
+            Assert.NotNull(merge.Columns.GetInt32Column("Int_right"));
+
+            Assert.NotNull(merge.Columns.GetUInt32Column("Uint_left"));
+            Assert.NotNull(merge.Columns.GetUInt32Column("Uint_right"));
+
+            Assert.NotNull(merge.Columns.GetInt64Column("Long_left"));
+            Assert.NotNull(merge.Columns.GetInt64Column("Long_right"));
+
+            Assert.NotNull(merge.Columns.GetUInt64Column("Ulong_left"));
+            Assert.NotNull(merge.Columns.GetUInt64Column("Ulong_right"));
+
+            Assert.NotNull(merge.Columns.GetDateTimeColumn("DateTime_left"));
+            Assert.NotNull(merge.Columns.GetDateTimeColumn("DateTime_right"));
+
+        }
+
         [Fact]
         public void TestDescription()
         {

Original file line number	Diff line number	Diff line change
`@@ -220,6 +220,16 @@ protected override IReadOnlyList<object> GetValues(long startIndex, int length)`
`220`	`220`	`return ret;`
`221`	`221`	`}`
`222`	`222`
	`223`	`+ internal virtual PrimitiveDataFrameColumn<T> CreateNewColumn(string name, PrimitiveColumnContainer<T> container)`
	`224`	`+ {`
	`225`	`+ return new PrimitiveDataFrameColumn<T>(name, container);`
	`226`	`+ }`
	`227`	`+`
	`228`	`+ protected virtual PrimitiveDataFrameColumn<T> CreateNewColumn(string name, long length = 0)`
	`229`	`+ {`
	`230`	`+ return new PrimitiveDataFrameColumn<T>(name, length);`
	`231`	`+ }`
	`232`	`+`
`223`	`233`	`internal T? GetTypedValue(long rowIndex) => _columnContainer[rowIndex];`
`224`	`234`
`225`	`235`	`protected override object GetValue(long rowIndex) => GetTypedValue(rowIndex);`
`@@ -411,7 +421,7 @@ private PrimitiveDataFrameColumn<T> Clone(PrimitiveDataFrameColumn<bool> boolCol`
`411`	`421`	`{`
`412`	`422`	`if (boolColumn.Length > Length)`
`413`	`423`	`throw new ArgumentException(Strings.MapIndicesExceedsColumnLenth, nameof(boolColumn));`
`414`		`- PrimitiveDataFrameColumn<T> ret = new PrimitiveDataFrameColumn<T>(Name);`
	`424`	`+ PrimitiveDataFrameColumn<T> ret = CreateNewColumn(Name);`
`415`	`425`	`for (long i = 0; i < boolColumn.Length; i++)`
`416`	`426`	`{`
`417`	`427`	`bool? value = boolColumn[i];`
`@@ -438,7 +448,8 @@ private PrimitiveDataFrameColumn<T> CloneImplementation<U>(PrimitiveDataFrameCol`
`438`	`448`	`}`
`439`	`449`	`else`
`440`	`450`	`throw new NotImplementedException();`
`441`		`- PrimitiveDataFrameColumn<T> ret = new PrimitiveDataFrameColumn<T>(Name, retContainer);`
	`451`	`+`
	`452`	`+ PrimitiveDataFrameColumn<T> ret = CreateNewColumn(Name, retContainer);`
`442`	`453`	`return ret;`
`443`	`454`	`}`
`444`	`455`
`@@ -447,7 +458,7 @@ public PrimitiveDataFrameColumn<T> Clone(PrimitiveDataFrameColumn<long> mapIndic`
`447`	`458`	`if (mapIndices is null)`
`448`	`459`	`{`
`449`	`460`	`PrimitiveColumnContainer<T> newColumnContainer = _columnContainer.Clone();`
`450`		`- return new PrimitiveDataFrameColumn<T>(Name, newColumnContainer);`
	`461`	`+ return CreateNewColumn(Name, newColumnContainer);`
`451`	`462`	`}`
`452`	`463`	`else`
`453`	`464`	`{`