diff --git a/arrow/array/builder.go b/arrow/array/builder.go index 0b3a4e9a..ed4f08e7 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 2fb47b5c..99fdd04a 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 13b0c46e..086ba7a2 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()) }