Fix DataFrame incorrectly sets column value for index higher than Buffer.MaxCapacity

Author: asmirnov82

src/Microsoft.Data.Analysis/PrimitiveColumnContainer.cs

@@ -359,13 +359,13 @@ private byte SetBit(byte curBitMap, int index, bool value)
         }
 
         // private function. Faster to use when we already have a span since it avoids indexing
-        private void SetValidityBit(Span<byte> bitMapBufferSpan, int index, bool value)
+        private void SetValidityBit(Span<byte> validitySpan, int index, bool value)
         {
             int bitMapBufferIndex = (int)((uint)index / 8);
-            Debug.Assert(bitMapBufferSpan.Length >= bitMapBufferIndex);
-            byte curBitMap = bitMapBufferSpan[bitMapBufferIndex];
+            Debug.Assert(validitySpan.Length >= bitMapBufferIndex);
+            byte curBitMap = validitySpan[bitMapBufferIndex];
             byte newBitMap = SetBit(curBitMap, index, value);
-            bitMapBufferSpan[bitMapBufferIndex] = newBitMap;
+            validitySpan[bitMapBufferIndex] = newBitMap;
         }
 
         /// <summary>
@@ -457,25 +457,25 @@ public T? this[long rowIndex]
                     return null;
                 }
                 int arrayIndex = GetArrayContainingRowIndex(rowIndex);
-                rowIndex = rowIndex - arrayIndex * ReadOnlyDataFrameBuffer<T>.MaxCapacity;
-                return Buffers[arrayIndex][(int)rowIndex];
+                var bufferOffset = (int)(rowIndex % ReadOnlyDataFrameBuffer<T>.MaxCapacity);
+                return Buffers[arrayIndex][bufferOffset];
             }
             set
             {
                 int arrayIndex = GetArrayContainingRowIndex(rowIndex);
-                rowIndex = rowIndex - arrayIndex * ReadOnlyDataFrameBuffer<T>.MaxCapacity;
+                var bufferOffset = (int)(rowIndex % ReadOnlyDataFrameBuffer<T>.MaxCapacity);
 
                 Buffers.GetOrCreateMutable(arrayIndex);
                 NullBitMapBuffers.GetOrCreateMutable(arrayIndex);
 
                 if (value.HasValue)
                 {
-                    Buffers[arrayIndex][(int)rowIndex] = value.Value;
+                    Buffers[arrayIndex][bufferOffset] = value.Value;
                     SetValidityBit(rowIndex, true);
                 }
                 else
                 {
-                    Buffers[arrayIndex][(int)rowIndex] = default;
+                    Buffers[arrayIndex][bufferOffset] = default;
                     SetValidityBit(rowIndex, false);
                 }
             }

test/Microsoft.Data.Analysis.Tests/BufferTests.cs

@@ -8,12 +8,27 @@
 using System.Text;
 using Apache.Arrow;
 using Microsoft.ML.TestFramework.Attributes;
+using Newtonsoft.Json.Linq;
 using Xunit;
 
 namespace Microsoft.Data.Analysis.Tests
 {
     public class BufferTests
     {
+        [X64Fact("32-bit doesn't allow to allocate more than 2 Gb")]
+        public void TestGetterAndSetterForColumnsGreaterThanMaxCapacity()
+        {
+            const int MaxCapacityInBytes = 2147483591;
+
+            var length = MaxCapacityInBytes + 5;
+            var column = new PrimitiveDataFrameColumn<byte>("LargeColumn", length);
+            var index = length - 1;
+            column[index] = 33;
+
+            Assert.Equal((byte)33, column[index]);
+            Assert.Null(column[index % MaxCapacityInBytes]);
+        }
+
         [Fact]
         public void TestNullCounts()
         {
@@ -446,15 +461,15 @@ public void TestArrowStringColumnClone()
                 Assert.Null(clone[i]);
         }
 
-        [X64Fact("32-bit dosn't allow to allocate more than 2 Gb")]
+        [X64Fact("32-bit doesn't allow to allocate more than 2 Gb")]
         public void TestAppend_SizeMoreThanMaxBufferCapacity()
         {
             //Check appending value, than can increase buffer size over MaxCapacity (default strategy is to double buffer capacity)
             PrimitiveDataFrameColumn<byte> intColumn = new PrimitiveDataFrameColumn<byte>("Byte1", int.MaxValue / 2 - 1);
             intColumn.Append(10);
         }
 
-        [X64Fact("32-bit dosn't allow to allocate more than 2 Gb")]
+        [X64Fact("32-bit doesn't allow to allocate more than 2 Gb")]
         public void TestAppendMany_SizeMoreThanMaxBufferCapacity()
         {
             const int MaxCapacityInBytes = 2147483591;