Add SetNull

apache · Jun 17, 2023 · cccf233 · cccf233
1 parent 9f6d825
commit cccf233
Show file tree

Hide file tree

Showing 3 changed files with 80 additions and 11 deletions.
diff --git a/csharp/src/Apache.Arrow/Arrays/BinaryArray.cs b/csharp/src/Apache.Arrow/Arrays/BinaryArray.cs
@@ -263,24 +263,29 @@ public TBuilder Set(int index, byte value)
                 return Instance;
             }
 
+            public TBuilder SetNull(int offset)
+            {
+                int index = ValueOffsets.Span[offset];
+                int length = GetOffsetValueLength(offset);
+
+                for (int i = 0; i < length; i++)
+                {
+                    ValueBuffer.Span[index + i] = 0;
+                }
+
+                ValidityBuffer.Set(offset, false);
+
+                return Instance;
+            }
+
             public TBuilder Set(int offset, ReadOnlySpan<byte> values)
             {
 #if NETSTANDARD2_1_OR_GREATER || NET6_0_OR_GREATER
                 byte[] newValues = values.ToArray();
                 int index = ValueOffsets.Span[offset];
-                int existingValueLength;
+                int existingValueLength = GetOffsetValueLength(offset);
                 int newValueLength = newValues.Length;
 
-                if (offset + 1 < ValueOffsets.Length)
-                {
-                    int nextIndex = ValueOffsets.Span[offset + 1];
-                    existingValueLength = nextIndex - index;
-                }
-                else
-                {
-                    existingValueLength = ValueBuffer.Length - ValueOffsets.Span[offset];
-                }
-
                 // Resize and shift the value and offset buffers
                 if (existingValueLength != newValueLength)
                 {
@@ -320,6 +325,13 @@ public TBuilder Set(int offset, ReadOnlySpan<byte> values)
 #endif
             }
 
+            private int GetOffsetValueLength(int offset)
+            {
+                int index = ValueOffsets.Span[offset];
+                int nextIndex = offset + 1 < ValueOffsets.Length ? ValueOffsets.Span[offset + 1] : ValueBuffer.Length;
+                return nextIndex - index;
+            }
+
             /// <summary>
             /// Clear all contents appended so far.
             /// </summary>

diff --git a/csharp/test/Apache.Arrow.Tests/BinaryArrayBuilderTests.cs b/csharp/test/Apache.Arrow.Tests/BinaryArrayBuilderTests.cs
@@ -508,6 +508,35 @@ public void EnsureSetsUpdatesTheCorrectBytes(int offset, string newValue)
                     Assert.Equal(expectedValues[i], actual);
                 }
             }
+
+            [Theory]
+            [InlineData(0)]
+            [InlineData(1)]
+            [InlineData(2)]
+            public void EnsureSetNullUpdatesTheCorrectBytes(int offset)
+            {
+                // Arrange
+                var builder = new BinaryArray.Builder();
+                var defaultValues = new[] { "Test",  "Test1", "Test2"};
+                var expectedValues = (string[])defaultValues.Clone();
+                expectedValues[offset] = null;
+                foreach (string defaultValue in defaultValues)
+                {
+                    builder.Append(StringArray.DefaultEncoding.GetBytes(defaultValue).AsSpan());
+                }
+
+                // Act
+                builder.SetNull(offset);
+                var array = builder.Build();
+
+                // Assert
+                for (int i = 0; i <expectedValues.Length; i++)
+                {
+                    ReadOnlySpan<byte> byteValue = array.GetBytes(i);
+                    string actual = byteValue == null ? null : StringArray.DefaultEncoding.GetString(byteValue);
+                    Assert.Equal(expectedValues[i], actual);
+                }
+            }
         }
 
         private static void AssertArrayContents(IEnumerable<byte[]> expectedContents, BinaryArray array)

diff --git a/csharp/test/Apache.Arrow.Tests/StringArrayBuilderTests.cs b/csharp/test/Apache.Arrow.Tests/StringArrayBuilderTests.cs
@@ -53,5 +53,33 @@ public void EnsureSetsUpdatesTheCorrectStrings(int offset, string newValue)
                 Assert.Equal(expectedValues[i], actual);
             }
         }
+
+        [Theory]
+        [InlineData(0)]
+        [InlineData(1)]
+        [InlineData(2)]
+        public void EnsureSetNullUpdatesTheCorrectStrings(int offset)
+        {
+            // Arrange
+            var builder = new StringArray.Builder();
+            var defaultValues = new[] { "Test",  "Test1", "Test2"};
+            var expectedValues = (string[])defaultValues.Clone();
+            expectedValues[offset] = null;
+            foreach (string defaultValue in defaultValues)
+            {
+                builder.Append(defaultValue);
+            }
+
+            // Act
+            builder.SetNull(offset);
+            var array = builder.Build();
+
+            // Assert
+            for (int i = 0; i <expectedValues.Length; i++)
+            {
+                string actual = array.GetString(i);
+                Assert.Equal(expectedValues[i], actual);
+            }
+        }
     }
 }