From ddebbdca1f420747c464af2c75067383c2b3e2dc Mon Sep 17 00:00:00 2001
From: ShryukGrandhi <85529990+ShryukGrandhi@users.noreply.github.com>
Date: Tue, 20 Oct 2026 02:50:48 -0700
Subject: [PATCH] Split buffers by target type capacity when cloning primitive
 columns

The CloneAs*Container methods mapped each source buffer to one target
buffer, which fails when the target element type is larger than the
source (e.g. a full int buffer cannot fit in a single double buffer) and
leaves non-full buffers when it is smaller, breaking row index lookups.

Fixes #7381
---
 .../PrimitiveColumnContainer.cs               | 248 +++++-------------
 .../PrimitiveDataFrameColumnTests.cs          |  35 +++
 2 files changed, 101 insertions(+), 182 deletions(-)

diff --git a/src/Microsoft.Data.Analysis/PrimitiveColumnContainer.cs b/src/Microsoft.Data.Analysis/PrimitiveColumnContainer.cs
index 7c52032785..6e83420205 100644
--- a/src/Microsoft.Data.Analysis/PrimitiveColumnContainer.cs
+++ b/src/Microsoft.Data.Analysis/PrimitiveColumnContainer.cs
@@ -543,239 +543,123 @@ public PrimitiveColumnContainer<T> Clone()
             return ret;
         }
 
-        internal PrimitiveColumnContainer<bool> CloneAsBoolContainer()
+        // The values of the source buffers are redistributed so that every buffer of the result (except the last one)
+        // holds exactly ReadOnlyDataFrameBuffer<U>.MaxCapacity values, as sizeof(U) may differ from sizeof(T)
+        private PrimitiveColumnContainer<U> CloneAs<U>(Func<T, U> convert)
+            where U : unmanaged
         {
-            var ret = new PrimitiveColumnContainer<bool>();
-            foreach (ReadOnlyDataFrameBuffer<T> buffer in Buffers)
-            {
-                DataFrameBuffer<bool> newBuffer = new DataFrameBuffer<bool>(buffer.Length);
-                ret.Buffers.Add(newBuffer);
-                newBuffer.IncreaseSize(buffer.Length);
+            var ret = new PrimitiveColumnContainer<U>();
+            int maxCapacity = ReadOnlyDataFrameBuffer<U>.MaxCapacity;
+            long remaining = Length;
+            DataFrameBuffer<U> newBuffer = null;
+            DataFrameBuffer<byte> newNullBitMapBuffer = null;
 
-                if (typeof(T) == typeof(bool))
-                {
-                    var localBuffer = buffer;
-                    ReadOnlyDataFrameBuffer<bool> boolLocalBuffer = Unsafe.As<ReadOnlyDataFrameBuffer<T>, ReadOnlyDataFrameBuffer<bool>>(ref localBuffer);
-                    boolLocalBuffer.ReadOnlySpan.TryCopyTo(newBuffer.RawSpan);
-                }
-                else
+            for (int b = 0; b < Buffers.Count; b++)
+            {
+                ReadOnlySpan<T> span = Buffers[b].ReadOnlySpan;
+                ReadOnlySpan<byte> nullBitMapSpan = NullBitMapBuffers[b].ReadOnlySpan;
+                int index = 0;
+                while (index < span.Length)
                 {
-                    newBuffer.Span.Fill(false);
+                    if (newBuffer == null || newBuffer.Length == maxCapacity)
+                    {
+                        int capacity = (int)Math.Min(remaining, maxCapacity);
+                        newBuffer = new DataFrameBuffer<U>(capacity);
+                        newNullBitMapBuffer = new DataFrameBuffer<byte>((capacity + 7) / 8);
+                        newNullBitMapBuffer.IncreaseSize((capacity + 7) / 8);
+                        ret.Buffers.Add(newBuffer);
+                        ret.NullBitMapBuffers.Add(newNullBitMapBuffer);
+                    }
+
+                    int newIndex = newBuffer.Length;
+                    int count = Math.Min(span.Length - index, maxCapacity - newIndex);
+                    newBuffer.IncreaseSize(count);
+                    Span<U> newSpan = newBuffer.Span;
+                    Span<byte> newNullBitMapSpan = newNullBitMapBuffer.Span;
+                    for (int i = 0; i < count; i++)
+                    {
+                        newSpan[newIndex + i] = convert(span[index + i]);
+                    }
+
+                    if (NullCount == 0)
+                    {
+                        BitUtility.SetBits(newNullBitMapSpan, newIndex, count, true);
+                    }
+                    else
+                    {
+                        for (int i = 0; i < count; i++)
+                        {
+                            BitUtility.SetBit(newNullBitMapSpan, newIndex + i, BitUtility.IsValid(nullBitMapSpan, index + i));
+                        }
+                    }
+
+                    index += count;
+                    remaining -= count;
                 }
-                ret.Length += buffer.Length;
             }
-            ret.NullBitMapBuffers = CloneNullBitMapBuffers();
+            ret.Length = Length;
             ret.NullCount = NullCount;
             return ret;
         }
 
+        internal PrimitiveColumnContainer<bool> CloneAsBoolContainer()
+        {
+            return CloneAs<bool>(value => typeof(T) == typeof(bool) && Unsafe.As<T, bool>(ref value));
+        }
+
         internal PrimitiveColumnContainer<byte> CloneAsByteContainer()
         {
-            var ret = new PrimitiveColumnContainer<byte>();
-            foreach (ReadOnlyDataFrameBuffer<T> buffer in Buffers)
-            {
-                ret.Length += buffer.Length;
-                DataFrameBuffer<byte> newBuffer = new DataFrameBuffer<byte>(buffer.Length);
-                ret.Buffers.Add(newBuffer);
-                ReadOnlySpan<T> span = buffer.ReadOnlySpan;
-                for (int i = 0; i < span.Length; i++)
-                {
-                    newBuffer.Append(ByteConverter<T>.Instance.GetByte(span[i]));
-                }
-            }
-            ret.NullBitMapBuffers = CloneNullBitMapBuffers();
-            ret.NullCount = NullCount;
-            return ret;
+            return CloneAs<byte>(value => ByteConverter<T>.Instance.GetByte(value));
         }
 
         internal PrimitiveColumnContainer<sbyte> CloneAsSByteContainer()
         {
-            var ret = new PrimitiveColumnContainer<sbyte>();
-            foreach (ReadOnlyDataFrameBuffer<T> buffer in Buffers)
-            {
-                ret.Length += buffer.Length;
-                DataFrameBuffer<sbyte> newBuffer = new DataFrameBuffer<sbyte>(buffer.Length);
-                ret.Buffers.Add(newBuffer);
-                ReadOnlySpan<T> span = buffer.ReadOnlySpan;
-                for (int i = 0; i < span.Length; i++)
-                {
-                    newBuffer.Append(SByteConverter<T>.Instance.GetSByte(span[i]));
-                }
-            }
-            ret.NullBitMapBuffers = CloneNullBitMapBuffers();
-            ret.NullCount = NullCount;
-            return ret;
+            return CloneAs<sbyte>(value => SByteConverter<T>.Instance.GetSByte(value));
         }
 
         internal PrimitiveColumnContainer<double> CloneAsDoubleContainer()
         {
-            var ret = new PrimitiveColumnContainer<double>();
-            foreach (ReadOnlyDataFrameBuffer<T> buffer in Buffers)
-            {
-                ret.Length += buffer.Length;
-                DataFrameBuffer<double> newBuffer = new DataFrameBuffer<double>(buffer.Length);
-                ret.Buffers.Add(newBuffer);
-                ReadOnlySpan<T> span = buffer.ReadOnlySpan;
-                for (int i = 0; i < span.Length; i++)
-                {
-                    newBuffer.Append(DoubleConverter<T>.Instance.GetDouble(span[i]));
-                }
-            }
-            ret.NullBitMapBuffers = CloneNullBitMapBuffers();
-            ret.NullCount = NullCount;
-            return ret;
+            return CloneAs<double>(value => DoubleConverter<T>.Instance.GetDouble(value));
         }
 
         internal PrimitiveColumnContainer<decimal> CloneAsDecimalContainer()
         {
-            var ret = new PrimitiveColumnContainer<decimal>();
-            foreach (ReadOnlyDataFrameBuffer<T> buffer in Buffers)
-            {
-                ret.Length += buffer.Length;
-                DataFrameBuffer<decimal> newBuffer = new DataFrameBuffer<decimal>(buffer.Length);
-                ret.Buffers.Add(newBuffer);
-                ReadOnlySpan<T> span = buffer.ReadOnlySpan;
-                for (int i = 0; i < span.Length; i++)
-                {
-                    newBuffer.Append(DecimalConverter<T>.Instance.GetDecimal(span[i]));
-                }
-            }
-            ret.NullBitMapBuffers = CloneNullBitMapBuffers();
-            ret.NullCount = NullCount;
-            return ret;
+            return CloneAs<decimal>(value => DecimalConverter<T>.Instance.GetDecimal(value));
         }
 
         internal PrimitiveColumnContainer<short> CloneAsShortContainer()
         {
-            var ret = new PrimitiveColumnContainer<short>();
-            foreach (ReadOnlyDataFrameBuffer<T> buffer in Buffers)
-            {
-                ret.Length += buffer.Length;
-                DataFrameBuffer<short> newBuffer = new DataFrameBuffer<short>(buffer.Length);
-                ret.Buffers.Add(newBuffer);
-                ReadOnlySpan<T> span = buffer.ReadOnlySpan;
-                for (int i = 0; i < span.Length; i++)
-                {
-                    newBuffer.Append(Int16Converter<T>.Instance.GetInt16(span[i]));
-                }
-            }
-            ret.NullBitMapBuffers = CloneNullBitMapBuffers();
-            ret.NullCount = NullCount;
-            return ret;
+            return CloneAs<short>(value => Int16Converter<T>.Instance.GetInt16(value));
         }
 
         internal PrimitiveColumnContainer<ushort> CloneAsUShortContainer()
         {
-            var ret = new PrimitiveColumnContainer<ushort>();
-            foreach (ReadOnlyDataFrameBuffer<T> buffer in Buffers)
-            {
-                ret.Length += buffer.Length;
-                DataFrameBuffer<ushort> newBuffer = new DataFrameBuffer<ushort>(buffer.Length);
-                ret.Buffers.Add(newBuffer);
-                ReadOnlySpan<T> span = buffer.ReadOnlySpan;
-                for (int i = 0; i < span.Length; i++)
-                {
-                    newBuffer.Append(UInt16Converter<T>.Instance.GetUInt16(span[i]));
-                }
-            }
-            ret.NullBitMapBuffers = CloneNullBitMapBuffers();
-            ret.NullCount = NullCount;
-            return ret;
+            return CloneAs<ushort>(value => UInt16Converter<T>.Instance.GetUInt16(value));
         }
 
         internal PrimitiveColumnContainer<int> CloneAsIntContainer()
         {
-            var ret = new PrimitiveColumnContainer<int>();
-            foreach (ReadOnlyDataFrameBuffer<T> buffer in Buffers)
-            {
-                ret.Length += buffer.Length;
-                DataFrameBuffer<int> newBuffer = new DataFrameBuffer<int>(buffer.Length);
-                ret.Buffers.Add(newBuffer);
-                ReadOnlySpan<T> span = buffer.ReadOnlySpan;
-                for (int i = 0; i < span.Length; i++)
-                {
-                    newBuffer.Append(Int32Converter<T>.Instance.GetInt32(span[i]));
-                }
-            }
-            ret.NullBitMapBuffers = CloneNullBitMapBuffers();
-            ret.NullCount = NullCount;
-            return ret;
+            return CloneAs<int>(value => Int32Converter<T>.Instance.GetInt32(value));
         }
 
         internal PrimitiveColumnContainer<uint> CloneAsUIntContainer()
         {
-            var ret = new PrimitiveColumnContainer<uint>();
-            foreach (ReadOnlyDataFrameBuffer<T> buffer in Buffers)
-            {
-                ret.Length += buffer.Length;
-                DataFrameBuffer<uint> newBuffer = new DataFrameBuffer<uint>(buffer.Length);
-                ret.Buffers.Add(newBuffer);
-                ReadOnlySpan<T> span = buffer.ReadOnlySpan;
-                for (int i = 0; i < span.Length; i++)
-                {
-                    newBuffer.Append(UInt32Converter<T>.Instance.GetUInt32(span[i]));
-                }
-            }
-            ret.NullBitMapBuffers = CloneNullBitMapBuffers();
-            ret.NullCount = NullCount;
-            return ret;
+            return CloneAs<uint>(value => UInt32Converter<T>.Instance.GetUInt32(value));
         }
 
         internal PrimitiveColumnContainer<long> CloneAsLongContainer()
         {
-            var ret = new PrimitiveColumnContainer<long>();
-            foreach (ReadOnlyDataFrameBuffer<T> buffer in Buffers)
-            {
-                ret.Length += buffer.Length;
-                DataFrameBuffer<long> newBuffer = new DataFrameBuffer<long>(buffer.Length);
-                ret.Buffers.Add(newBuffer);
-                ReadOnlySpan<T> span = buffer.ReadOnlySpan;
-                for (int i = 0; i < span.Length; i++)
-                {
-                    newBuffer.Append(Int64Converter<T>.Instance.GetInt64(span[i]));
-                }
-            }
-            ret.NullBitMapBuffers = CloneNullBitMapBuffers();
-            ret.NullCount = NullCount;
-            return ret;
+            return CloneAs<long>(value => Int64Converter<T>.Instance.GetInt64(value));
         }
 
         internal PrimitiveColumnContainer<ulong> CloneAsULongContainer()
         {
-            var ret = new PrimitiveColumnContainer<ulong>();
-            foreach (ReadOnlyDataFrameBuffer<T> buffer in Buffers)
-            {
-                ret.Length += buffer.Length;
-                DataFrameBuffer<ulong> newBuffer = new DataFrameBuffer<ulong>(buffer.Length);
-                ret.Buffers.Add(newBuffer);
-                ReadOnlySpan<T> span = buffer.ReadOnlySpan;
-                for (int i = 0; i < span.Length; i++)
-                {
-                    newBuffer.Append(UInt64Converter<T>.Instance.GetUInt64(span[i]));
-                }
-            }
-            ret.NullBitMapBuffers = CloneNullBitMapBuffers();
-            ret.NullCount = NullCount;
-            return ret;
+            return CloneAs<ulong>(value => UInt64Converter<T>.Instance.GetUInt64(value));
         }
 
         internal PrimitiveColumnContainer<float> CloneAsFloatContainer()
         {
-            var ret = new PrimitiveColumnContainer<float>();
-            foreach (ReadOnlyDataFrameBuffer<T> buffer in Buffers)
-            {
-                ret.Length += buffer.Length;
-                DataFrameBuffer<float> newBuffer = new DataFrameBuffer<float>(buffer.Length);
-                ret.Buffers.Add(newBuffer);
-                ReadOnlySpan<T> span = buffer.ReadOnlySpan;
-                for (int i = 0; i < span.Length; i++)
-                {
-                    newBuffer.Append(SingleConverter<T>.Instance.GetSingle(span[i]));
-                }
-            }
-            ret.NullBitMapBuffers = CloneNullBitMapBuffers();
-            ret.NullCount = NullCount;
-            return ret;
+            return CloneAs<float>(value => SingleConverter<T>.Instance.GetSingle(value));
         }
     }
 }
diff --git a/test/Microsoft.Data.Analysis.Tests/PrimitiveDataFrameColumnTests.cs b/test/Microsoft.Data.Analysis.Tests/PrimitiveDataFrameColumnTests.cs
index 0b58d3d3e8..7a106c767b 100644
--- a/test/Microsoft.Data.Analysis.Tests/PrimitiveDataFrameColumnTests.cs
+++ b/test/Microsoft.Data.Analysis.Tests/PrimitiveDataFrameColumnTests.cs
@@ -54,6 +54,41 @@ public void PrimitiveDataFrameColumn_Int32_ElementwiseEquals_AcrossMultipleBuffe
             Assert.False(results[count - 1]);           // null != 0
         }
 
+        [X64Fact("32-bit doesn't allow to allocate more than 2 Gb")]
+        public void PrimitiveDataFrameColumn_Int32_CloneAsDouble_AcrossMultipleBuffers()
+        {
+            // Mirrors ReadOnlyDataFrameBuffer<int>.MaxCapacity, the number of values held by a single buffer.
+            // A full int buffer holds twice as many values as a double buffer can.
+            const int MaxCapacity = 0X7FEFFFFF / sizeof(int);
+            const int DoubleMaxCapacity = 0X7FEFFFFF / sizeof(double);
+
+            const long count = MaxCapacity + 16L;
+            var column = new Int32DataFrameColumn("Int", count);
+            Assert.Equal(2, column.GetReadOnlyDataBuffers().Count());
+
+            column[0] = 1;
+            column[DoubleMaxCapacity - 1] = 2;
+            column[DoubleMaxCapacity] = 3;
+            column[MaxCapacity - 1] = 4;
+            column[MaxCapacity] = 5;
+            column[count - 1] = 6;
+
+            // Adding a double value clones the int column as a double column
+            DoubleDataFrameColumn result = column.Add(0.5);
+
+            Assert.Equal(count, result.Length);
+            Assert.Equal(count - 6, result.NullCount);
+            Assert.Equal(1.5, result[0]);
+            Assert.Null(result[1]);
+            Assert.Equal(2.5, result[DoubleMaxCapacity - 1]);
+            Assert.Equal(3.5, result[DoubleMaxCapacity]);
+            Assert.Null(result[DoubleMaxCapacity + 1]);
+            Assert.Equal(4.5, result[MaxCapacity - 1]);
+            Assert.Equal(5.5, result[MaxCapacity]);
+            Assert.Null(result[MaxCapacity + 1]);
+            Assert.Equal(6.5, result[count - 1]);
+        }
+
         [X64Fact("32-bit doesn't allow to allocate more than 2 Gb")]
         public void TestGetterAndSetterForColumnsGreaterThanMaxCapacity()
         {