From 47138ad06333c3e22dd34bd77d233f07ea429cc5 Mon Sep 17 00:00:00 2001
From: Minh Vu <vuhoangminh97@gmail.com>
Date: Sat, 4 Jul 2026 21:41:18 +0200
Subject: [PATCH] BUG: fix builder SetNull null count tracking

---
 arrow/array/builder.go      |  3 +++
 arrow/array/builder_test.go | 13 +++++++++++++
 arrow/array/map_test.go     |  4 ++++
 3 files changed, 20 insertions(+)

diff --git a/arrow/array/builder.go b/arrow/array/builder.go
index 0b3a4e9af..ed4f08e72 100644
--- a/arrow/array/builder.go
+++ b/arrow/array/builder.go
@@ -133,6 +133,9 @@ func (b *builder) SetNull(i int) {
 	if i < 0 || i >= b.length {
 		panic("arrow/array: index out of range")
 	}
+	if bitutil.BitIsSet(b.nullBitmap.Bytes(), i) {
+		b.nulls++
+	}
 	bitutil.ClearBit(b.nullBitmap.Bytes(), i)
 }
 
diff --git a/arrow/array/builder_test.go b/arrow/array/builder_test.go
index 2fb47b5cb..99fdd04aa 100644
--- a/arrow/array/builder_test.go
+++ b/arrow/array/builder_test.go
@@ -112,6 +112,18 @@ func TestBuilder_SetNull(t *testing.T) {
 			b.SetNull(i)
 		}
 	}
+	assert.Equal(t, n/2, b.nulls)
+	assert.Equal(t, n/2, b.NullN())
+
+	b.SetNull(0)
+	assert.Equal(t, n/2, b.nulls)
+	assert.Equal(t, n/2, b.NullN())
+	b.SetNull(1)
+	assert.Equal(t, n/2+1, b.nulls)
+	assert.Equal(t, n/2+1, b.NullN())
+	b.SetNull(1)
+	assert.Equal(t, n/2+1, b.nulls)
+	assert.Equal(t, n/2+1, b.NullN())
 
 	for i := 0; i < n; i++ {
 		if i%2 == 0 {
@@ -120,4 +132,5 @@ func TestBuilder_SetNull(t *testing.T) {
 			assert.False(t, b.IsNull(i))
 		}
 	}
+	assert.True(t, b.IsNull(1))
 }
diff --git a/arrow/array/map_test.go b/arrow/array/map_test.go
index 13b0c46ea..086ba7a26 100644
--- a/arrow/array/map_test.go
+++ b/arrow/array/map_test.go
@@ -244,6 +244,9 @@ func TestMapBuilder_SetNull(t *testing.T) {
 
 	bldr.SetNull(0)
 	bldr.SetNull(3)
+	assert.Equal(t, 2, bldr.NullN())
+	bldr.SetNull(3)
+	assert.Equal(t, 2, bldr.NullN())
 
 	arr = bldr.NewMapArray()
 	defer arr.Release()
@@ -251,4 +254,5 @@ func TestMapBuilder_SetNull(t *testing.T) {
 	assert.True(t, arr.IsNull(0))
 	assert.True(t, arr.IsValid(1))
 	assert.True(t, arr.IsNull(3))
+	assert.Equal(t, 2, arr.NullN())
 }