diff --git a/arrow/compute/exec/span.go b/arrow/compute/exec/span.go index 49350eeb5..d071c0c6b 100644 --- a/arrow/compute/exec/span.go +++ b/arrow/compute/exec/span.go @@ -214,13 +214,13 @@ func (a *ArraySpan) SetSlice(off, length int64) { return } + // The cached null count describes the span as it was, not the slice. + // As in the C++ ArraySpan::SetSlice, a slice of anything with a validity + // bitmap has an unknown count until someone asks for it; only the null + // type and a span without a bitmap can be counted without looking. if a.Type.ID() != arrow.NULL { - if a.Nulls != 0 { - if a.Nulls == a.Len { - a.Nulls = length - } else { - a.Nulls = array.UnknownNullCount - } + if a.Nulls != 0 || len(a.Buffers[0].Buf) != 0 { + a.Nulls = array.UnknownNullCount } } else { a.Nulls = length diff --git a/arrow/compute/exec/span_test.go b/arrow/compute/exec/span_test.go index 90d4c93eb..3a19b21c0 100644 --- a/arrow/compute/exec/span_test.go +++ b/arrow/compute/exec/span_test.go @@ -523,6 +523,8 @@ func TestArraySpan_SetSlice(t *testing.T) { {"null type", fields{Type: arrow.Null}, args{5, 10}, 10}, {"not-null type", fields{Type: arrow.PrimitiveTypes.Int8}, args{5, 10}, 0}, {"not-null type with nulls", fields{Type: arrow.PrimitiveTypes.Int8, Nulls: -1}, args{5, 10}, array.UnknownNullCount}, + {"no nulls but a validity bitmap", fields{Type: arrow.PrimitiveTypes.Int8, Len: 16, Buffers: [3]exec.BufferSpan{{Buf: []byte{0xff, 0xff}}}}, args{5, 10}, array.UnknownNullCount}, + {"all null before the slice", fields{Type: arrow.PrimitiveTypes.Int8, Len: 16, Nulls: 16, Buffers: [3]exec.BufferSpan{{Buf: []byte{0, 0}}}}, args{5, 10}, array.UnknownNullCount}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { diff --git a/arrow/compute/scalar_boolean_chunksize_test.go b/arrow/compute/scalar_boolean_chunksize_test.go new file mode 100644 index 000000000..61c63a1ce --- /dev/null +++ b/arrow/compute/scalar_boolean_chunksize_test.go @@ -0,0 +1,76 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package compute_test + +import ( + "context" + "testing" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/compute" + "github.com/apache/arrow-go/v18/arrow/memory" +) + +// and_kleene and or_kleene over two boolean arrays with nulls, executed with +// different ExecCtx.ChunkSize values; the result must not depend on ChunkSize. +func TestKleeneChunkSizeIndependent(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + // left: [true, null, true, null, false, true, null, true] + // right: [null, true, true, false, null, null, true, true] + lb := array.NewBooleanBuilder(mem) + lb.AppendValues([]bool{true, false, true, false, false, true, false, true}, []bool{true, false, true, false, true, true, false, true}) + left := lb.NewArray() + lb.Release() + defer left.Release() + rb := array.NewBooleanBuilder(mem) + rb.AppendValues([]bool{false, true, true, false, false, false, true, true}, []bool{false, true, true, true, false, false, true, true}) + right := rb.NewArray() + rb.Release() + defer right.Release() + + for _, fn := range []string{"and_kleene", "or_kleene"} { + t.Run(fn, func(t *testing.T) { checkChunkSizeIndependent(t, mem, fn, left, right) }) + } +} + +func checkChunkSizeIndependent(t *testing.T, mem memory.Allocator, fn string, left, right arrow.Array) { + run := func(chunk int64) string { + ectx := compute.DefaultExecCtx() + ectx.ChunkSize = chunk + ctx := compute.SetExecCtx(compute.WithAllocator(context.Background(), mem), ectx) + out, err := compute.CallFunction(ctx, fn, nil, &compute.ArrayDatum{Value: left.Data()}, &compute.ArrayDatum{Value: right.Data()}) + if err != nil { + t.Fatal(err) + } + defer out.Release() + arr := out.(*compute.ArrayDatum).MakeArray() + defer arr.Release() + return arr.String() + } + ref := run(compute.DefaultMaxChunkSize) + t.Logf("ChunkSize default: %s", ref) + for _, n := range []int64{1, 2, 3, 4, 8} { + got := run(n) + t.Logf("ChunkSize %d: %s", n, got) + if got != ref { + t.Errorf("ChunkSize %d: got %s, want %s", n, got, ref) + } + } +}