Fix DataFrame Merge issue (#6677)

asmirnov82 · web-flow · commit c259a33de003 · 2023-06-13T16:43:37.000-06:00
* Fix DataFrame Merge issue

* Return commented test
diff --git a/src/Microsoft.Data.Analysis/BooleanDataFrameColumn.cs b/src/Microsoft.Data.Analysis/BooleanDataFrameColumn.cs
@@ -19,5 +19,15 @@ public BooleanDataFrameColumn(string name, long length = 0) : base(name, length)
         public BooleanDataFrameColumn(string name, ReadOnlyMemory<byte> buffer, ReadOnlyMemory<byte> nullBitMap, int length = 0, int nullCount = 0) : base(name, buffer, nullBitMap, length, nullCount) { }
 
         internal BooleanDataFrameColumn(string name, PrimitiveColumnContainer<bool> values) : base(name, values) { }
+
+        protected override PrimitiveDataFrameColumn<bool> CreateNewColumn(string name, long length = 0)
+        {
+            return new BooleanDataFrameColumn(name, length);
+        }
+
+        internal override PrimitiveDataFrameColumn<bool> CreateNewColumn(string name, PrimitiveColumnContainer<bool> container)
+        {
+            return new BooleanDataFrameColumn(name, container);
+        }
     }
 }
diff --git a/src/Microsoft.Data.Analysis/ByteDataFrameColumn.cs b/src/Microsoft.Data.Analysis/ByteDataFrameColumn.cs
@@ -19,5 +19,15 @@ public ByteDataFrameColumn(string name, long length = 0) : base(name, length) {
         public ByteDataFrameColumn(string name, ReadOnlyMemory<byte> buffer, ReadOnlyMemory<byte> nullBitMap, int length = 0, int nullCount = 0) : base(name, buffer, nullBitMap, length, nullCount) { }
 
         internal ByteDataFrameColumn(string name, PrimitiveColumnContainer<byte> values) : base(name, values) { }
+
+        protected override PrimitiveDataFrameColumn<byte> CreateNewColumn(string name, long length = 0)
+        {
+            return new ByteDataFrameColumn(name, length);
+        }
+
+        internal override PrimitiveDataFrameColumn<byte> CreateNewColumn(string name, PrimitiveColumnContainer<byte> container)
+        {
+            return new ByteDataFrameColumn(name, container);
+        }
     }
 }
diff --git a/src/Microsoft.Data.Analysis/CharDataFrameColumn.cs b/src/Microsoft.Data.Analysis/CharDataFrameColumn.cs
@@ -19,5 +19,15 @@ public CharDataFrameColumn(string name, long length = 0) : base(name, length) {
         public CharDataFrameColumn(string name, ReadOnlyMemory<byte> buffer, ReadOnlyMemory<byte> nullBitMap, int length = 0, int nullCount = 0) : base(name, buffer, nullBitMap, length, nullCount) { }
 
         internal CharDataFrameColumn(string name, PrimitiveColumnContainer<char> values) : base(name, values) { }
+
+        protected override PrimitiveDataFrameColumn<char> CreateNewColumn(string name, long length = 0)
+        {
+            return new CharDataFrameColumn(name, length);
+        }
+
+        internal override PrimitiveDataFrameColumn<char> CreateNewColumn(string name, PrimitiveColumnContainer<char> container)
+        {
+            return new CharDataFrameColumn(name, container);
+        }
     }
 }
diff --git a/src/Microsoft.Data.Analysis/DataFrameColumnCollection.cs b/src/Microsoft.Data.Analysis/DataFrameColumnCollection.cs
@@ -191,6 +191,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="DateTimeDataFrameColumn"/> with the specified <paramref name="name"/>.
+        /// </summary>
+        /// <param name="name">The name of the column</param>
+        /// <returns><see cref="DateTimeDataFrameColumn"/>.</returns>
+        /// <exception cref="ArgumentException">A column named <paramref name="name"/> cannot be found, or if the column's type doesn't match.</exception>
+        public DateTimeDataFrameColumn GetDateTimeColumn(string name)
+        {
+            DataFrameColumn column = this[name];
+            if (column is DateTimeDataFrameColumn 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/DateTimeDataFrameColumn.cs b/src/Microsoft.Data.Analysis/DateTimeDataFrameColumn.cs
@@ -19,5 +19,15 @@ public DateTimeDataFrameColumn(string name, long length = 0) : base(name, length
         public DateTimeDataFrameColumn(string name, ReadOnlyMemory<byte> buffer, ReadOnlyMemory<byte> nullBitMap, int length = 0, int nullCount = 0) : base(name, buffer, nullBitMap, length, nullCount) { }
 
         internal DateTimeDataFrameColumn(string name, PrimitiveColumnContainer<DateTime> values) : base(name, values) { }
+
+        protected override PrimitiveDataFrameColumn<DateTime> CreateNewColumn(string name, long length = 0)
+        {
+            return new DateTimeDataFrameColumn(name, length);
+        }
+
+        internal override PrimitiveDataFrameColumn<DateTime> CreateNewColumn(string name, PrimitiveColumnContainer<DateTime> container)
+        {
+            return new DateTimeDataFrameColumn(name, container);
+        }
     }
 }
diff --git a/src/Microsoft.Data.Analysis/DecimalDataFrameColumn.cs b/src/Microsoft.Data.Analysis/DecimalDataFrameColumn.cs
@@ -19,5 +19,15 @@ public DecimalDataFrameColumn(string name, long length = 0) : base(name, length)
         public DecimalDataFrameColumn(string name, ReadOnlyMemory<byte> buffer, ReadOnlyMemory<byte> nullBitMap, int length = 0, int nullCount = 0) : base(name, buffer, nullBitMap, length, nullCount) { }
 
         internal DecimalDataFrameColumn(string name, PrimitiveColumnContainer<decimal> values) : base(name, values) { }
+
+        protected override PrimitiveDataFrameColumn<decimal> CreateNewColumn(string name, long length = 0)
+        {
+            return new DecimalDataFrameColumn(name, length);
+        }
+
+        internal override PrimitiveDataFrameColumn<decimal> CreateNewColumn(string name, PrimitiveColumnContainer<decimal> container)
+        {
+            return new DecimalDataFrameColumn(name, container);
+        }
     }
 }
diff --git a/src/Microsoft.Data.Analysis/DoubleDataFrameColumn.cs b/src/Microsoft.Data.Analysis/DoubleDataFrameColumn.cs
@@ -19,5 +19,15 @@ public DoubleDataFrameColumn(string name, long length = 0) : base(name, length)
         public DoubleDataFrameColumn(string name, ReadOnlyMemory<byte> buffer, ReadOnlyMemory<byte> nullBitMap, int length = 0, int nullCount = 0) : base(name, buffer, nullBitMap, length, nullCount) { }
 
         internal DoubleDataFrameColumn(string name, PrimitiveColumnContainer<double> values) : base(name, values) { }
+
+        protected override PrimitiveDataFrameColumn<double> CreateNewColumn(string name, long length = 0)
+        {
+            return new DoubleDataFrameColumn(name, length);
+        }
+
+        internal override PrimitiveDataFrameColumn<double> CreateNewColumn(string name, PrimitiveColumnContainer<double> container)
+        {
+            return new DoubleDataFrameColumn(name, container);
+        }
     }
 }
diff --git a/src/Microsoft.Data.Analysis/Int16DataFrameColumn.cs b/src/Microsoft.Data.Analysis/Int16DataFrameColumn.cs
@@ -19,5 +19,15 @@ public Int16DataFrameColumn(string name, long length = 0) : base(name, length) {
         public Int16DataFrameColumn(string name, ReadOnlyMemory<byte> buffer, ReadOnlyMemory<byte> nullBitMap, int length = 0, int nullCount = 0) : base(name, buffer, nullBitMap, length, nullCount) { }
 
         internal Int16DataFrameColumn(string name, PrimitiveColumnContainer<short> values) : base(name, values) { }
+
+        protected override PrimitiveDataFrameColumn<short> CreateNewColumn(string name, long length = 0)
+        {
+            return new Int16DataFrameColumn(name, length);
+        }
+
+        internal override PrimitiveDataFrameColumn<short> CreateNewColumn(string name, PrimitiveColumnContainer<short> container)
+        {
+            return new Int16DataFrameColumn(name, container);
+        }
     }
 }
diff --git a/src/Microsoft.Data.Analysis/Int32DataFrameColumn.cs b/src/Microsoft.Data.Analysis/Int32DataFrameColumn.cs
@@ -19,5 +19,15 @@ public Int32DataFrameColumn(string name, long length = 0) : base(name, length) {
         public Int32DataFrameColumn(string name, ReadOnlyMemory<byte> buffer, ReadOnlyMemory<byte> nullBitMap, int length = 0, int nullCount = 0) : base(name, buffer, nullBitMap, length, nullCount) { }
 
         internal Int32DataFrameColumn(string name, PrimitiveColumnContainer<int> values) : base(name, values) { }
+
+        protected override PrimitiveDataFrameColumn<int> CreateNewColumn(string name, long length = 0)
+        {
+            return new Int32DataFrameColumn(name, length);
+        }
+
+        internal override PrimitiveDataFrameColumn<int> CreateNewColumn(string name, PrimitiveColumnContainer<int> container)
+        {
+            return new Int32DataFrameColumn(name, container);
+        }
     }
 }
diff --git a/src/Microsoft.Data.Analysis/Int64DataFrameColumn.cs b/src/Microsoft.Data.Analysis/Int64DataFrameColumn.cs
@@ -19,5 +19,15 @@ public Int64DataFrameColumn(string name, long length = 0) : base(name, length) {
         public Int64DataFrameColumn(string name, ReadOnlyMemory<byte> buffer, ReadOnlyMemory<byte> nullBitMap, int length = 0, int nullCount = 0) : base(name, buffer, nullBitMap, length, nullCount) { }
 
         internal Int64DataFrameColumn(string name, PrimitiveColumnContainer<long> values) : base(name, values) { }
+
+        protected override PrimitiveDataFrameColumn<long> CreateNewColumn(string name, long length = 0)
+        {
+            return new Int64DataFrameColumn(name, length);
+        }
+
+        internal override PrimitiveDataFrameColumn<long> CreateNewColumn(string name, PrimitiveColumnContainer<long> container)
+        {
+            return new Int64DataFrameColumn(name, container);
+        }
     }
 }
diff --git a/src/Microsoft.Data.Analysis/PrimitiveDataFrameColumn.cs b/src/Microsoft.Data.Analysis/PrimitiveDataFrameColumn.cs
@@ -188,6 +188,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);
@@ -379,7 +389,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];
@@ -406,7 +416,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;
         }
 
@@ -415,7 +426,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/src/Microsoft.Data.Analysis/SByteDataFrameColumn.cs b/src/Microsoft.Data.Analysis/SByteDataFrameColumn.cs
@@ -19,5 +19,15 @@ public SByteDataFrameColumn(string name, long length = 0) : base(name, length) {
         public SByteDataFrameColumn(string name, ReadOnlyMemory<byte> buffer, ReadOnlyMemory<byte> nullBitMap, int length = 0, int nullCount = 0) : base(name, buffer, nullBitMap, length, nullCount) { }
 
         internal SByteDataFrameColumn(string name, PrimitiveColumnContainer<sbyte> values) : base(name, values) { }
+
+        protected override PrimitiveDataFrameColumn<sbyte> CreateNewColumn(string name, long length = 0)
+        {
+            return new SByteDataFrameColumn(name, length);
+        }
+
+        internal override PrimitiveDataFrameColumn<sbyte> CreateNewColumn(string name, PrimitiveColumnContainer<sbyte> container)
+        {
+            return new SByteDataFrameColumn(name, container);
+        }
     }
 }
diff --git a/src/Microsoft.Data.Analysis/SingleDataFrameColumn.cs b/src/Microsoft.Data.Analysis/SingleDataFrameColumn.cs
@@ -19,5 +19,15 @@ public SingleDataFrameColumn(string name, long length = 0) : base(name, length)
         public SingleDataFrameColumn(string name, ReadOnlyMemory<byte> buffer, ReadOnlyMemory<byte> nullBitMap, int length = 0, int nullCount = 0) : base(name, buffer, nullBitMap, length, nullCount) { }
 
         internal SingleDataFrameColumn(string name, PrimitiveColumnContainer<float> values) : base(name, values) { }
+
+        protected override PrimitiveDataFrameColumn<float> CreateNewColumn(string name, long length = 0)
+        {
+            return new SingleDataFrameColumn(name, length);
+        }
+
+        internal override PrimitiveDataFrameColumn<float> CreateNewColumn(string name, PrimitiveColumnContainer<float> container)
+        {
+            return new SingleDataFrameColumn(name, container);
+        }
     }
 }
diff --git a/src/Microsoft.Data.Analysis/UInt16DataFrameColumn.cs b/src/Microsoft.Data.Analysis/UInt16DataFrameColumn.cs
@@ -19,5 +19,15 @@ public UInt16DataFrameColumn(string name, long length = 0) : base(name, length)
         public UInt16DataFrameColumn(string name, ReadOnlyMemory<byte> buffer, ReadOnlyMemory<byte> nullBitMap, int length = 0, int nullCount = 0) : base(name, buffer, nullBitMap, length, nullCount) { }
 
         internal UInt16DataFrameColumn(string name, PrimitiveColumnContainer<ushort> values) : base(name, values) { }
+
+        protected override PrimitiveDataFrameColumn<ushort> CreateNewColumn(string name, long length = 0)
+        {
+            return new UInt16DataFrameColumn(name, length);
+        }
+
+        internal override PrimitiveDataFrameColumn<ushort> CreateNewColumn(string name, PrimitiveColumnContainer<ushort> container)
+        {
+            return new UInt16DataFrameColumn(name, container);
+        }
     }
 }
diff --git a/src/Microsoft.Data.Analysis/UInt32DataFrameColumn.cs b/src/Microsoft.Data.Analysis/UInt32DataFrameColumn.cs
@@ -19,5 +19,15 @@ public UInt32DataFrameColumn(string name, long length = 0) : base(name, length)
         public UInt32DataFrameColumn(string name, ReadOnlyMemory<byte> buffer, ReadOnlyMemory<byte> nullBitMap, int length = 0, int nullCount = 0) : base(name, buffer, nullBitMap, length, nullCount) { }
 
         internal UInt32DataFrameColumn(string name, PrimitiveColumnContainer<uint> values) : base(name, values) { }
+
+        protected override PrimitiveDataFrameColumn<uint> CreateNewColumn(string name, long length = 0)
+        {
+            return new UInt32DataFrameColumn(name, length);
+        }
+
+        internal override PrimitiveDataFrameColumn<uint> CreateNewColumn(string name, PrimitiveColumnContainer<uint> container)
+        {
+            return new UInt32DataFrameColumn(name, container);
+        }
     }
 }
diff --git a/src/Microsoft.Data.Analysis/UInt64DataFrameColumn.cs b/src/Microsoft.Data.Analysis/UInt64DataFrameColumn.cs
@@ -19,5 +19,15 @@ public UInt64DataFrameColumn(string name, long length = 0) : base(name, length)
         public UInt64DataFrameColumn(string name, ReadOnlyMemory<byte> buffer, ReadOnlyMemory<byte> nullBitMap, int length = 0, int nullCount = 0) : base(name, buffer, nullBitMap, length, nullCount) { }
 
         internal UInt64DataFrameColumn(string name, PrimitiveColumnContainer<ulong> values) : base(name, values) { }
+
+        protected override PrimitiveDataFrameColumn<ulong> CreateNewColumn(string name, long length = 0)
+        {
+            return new UInt64DataFrameColumn(name, length);
+        }
+
+        internal override PrimitiveDataFrameColumn<ulong> CreateNewColumn(string name, PrimitiveColumnContainer<ulong> container)
+        {
+            return new UInt64DataFrameColumn(name, container);
+        }
     }
 }
diff --git a/test/Microsoft.Data.Analysis.Tests/DataFrameTests.cs b/test/Microsoft.Data.Analysis.Tests/DataFrameTests.cs
@@ -2659,6 +2659,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()
         {
