From 684a79d1615e86f29ecbedcbdf539a205d0267ff Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Fri, 28 Aug 2026 12:31:22 +0200 Subject: [PATCH 1/3] perf(compute): pre-size SetLookupState memo tables --- .../internal/kernels/scalar_set_lookup.go | 2 +- arrow/compute/internal/kernels/vector_hash.go | 18 ++-- arrow/compute/scalar_set_lookup_bench_test.go | 97 +++++++++++++++++++ 3 files changed, 107 insertions(+), 10 deletions(-) create mode 100644 arrow/compute/scalar_set_lookup_bench_test.go diff --git a/arrow/compute/internal/kernels/scalar_set_lookup.go b/arrow/compute/internal/kernels/scalar_set_lookup.go index 7b1caeb70..03dbfe4eb 100644 --- a/arrow/compute/internal/kernels/scalar_set_lookup.go +++ b/arrow/compute/internal/kernels/scalar_set_lookup.go @@ -188,7 +188,7 @@ func (s *SetLookupState[T]) Init(opts SetLookupOptions) error { } } } - lookup, err := newMemoTable(s.Alloc, memoType) + lookup, err := newMemoTable(s.Alloc, memoType, opts.TotalLen*2) if err != nil { return err } diff --git a/arrow/compute/internal/kernels/vector_hash.go b/arrow/compute/internal/kernels/vector_hash.go index 00f481734..c8244d960 100644 --- a/arrow/compute/internal/kernels/vector_hash.go +++ b/arrow/compute/internal/kernels/vector_hash.go @@ -256,7 +256,7 @@ func (rhs *regularHashState) ValueType() arrow.DataType { return rhs.typ } func (rhs *regularHashState) Reset() error { if rhs.memoReleased { - memoTable, err := newMemoTable(rhs.mem, rhs.typ.ID()) + memoTable, err := newMemoTable(rhs.mem, rhs.typ.ID(), 0) if err != nil { return err } @@ -592,25 +592,25 @@ func nullHashInit(actionInit initAction) exec.KernelInitFn { } } -func newMemoTable(mem memory.Allocator, dt arrow.Type) (hashing.MemoTable, error) { +func newMemoTable(mem memory.Allocator, dt arrow.Type, initial int64) (hashing.MemoTable, error) { switch dt { case arrow.BOOL, arrow.INT8, arrow.UINT8: - return hashing.NewMemoTable[uint8](0), nil + return hashing.NewMemoTable[uint8](initial), nil case arrow.INT16, arrow.UINT16: - return hashing.NewMemoTable[uint16](0), nil + return hashing.NewMemoTable[uint16](initial), nil case arrow.INT32, arrow.UINT32, arrow.FLOAT32, arrow.DECIMAL32, arrow.DATE32, arrow.TIME32, arrow.INTERVAL_MONTHS: - return hashing.NewMemoTable[uint32](0), nil + return hashing.NewMemoTable[uint32](initial), nil case arrow.INT64, arrow.UINT64, arrow.FLOAT64, arrow.DECIMAL64, arrow.DATE64, arrow.TIME64, arrow.TIMESTAMP, arrow.DURATION, arrow.INTERVAL_DAY_TIME: - return hashing.NewMemoTable[uint64](0), nil + return hashing.NewMemoTable[uint64](initial), nil case arrow.BINARY, arrow.STRING, arrow.FIXED_SIZE_BINARY, arrow.DECIMAL128, arrow.DECIMAL256, arrow.INTERVAL_MONTH_DAY_NANO: - return hashing.NewBinaryMemoTable(0, 0, + return hashing.NewBinaryMemoTable(int(initial), 0, array.NewBinaryBuilder(mem, arrow.BinaryTypes.Binary)), nil case arrow.LARGE_BINARY, arrow.LARGE_STRING: - return hashing.NewBinaryMemoTable(0, 0, + return hashing.NewBinaryMemoTable(int(initial), 0, array.NewBinaryBuilder(mem, arrow.BinaryTypes.LargeBinary)), nil default: return nil, fmt.Errorf("%w: unsupported type %s", arrow.ErrNotImplemented, dt) @@ -624,7 +624,7 @@ func regularHashInit(dt arrow.DataType, actionInit initAction, appendFn func(Act if err != nil { return nil, err } - memoTable, err := newMemoTable(mem, dt.ID()) + memoTable, err := newMemoTable(mem, dt.ID(), 0) if err != nil { return nil, err } diff --git a/arrow/compute/scalar_set_lookup_bench_test.go b/arrow/compute/scalar_set_lookup_bench_test.go new file mode 100644 index 000000000..57af76aaa --- /dev/null +++ b/arrow/compute/scalar_set_lookup_bench_test.go @@ -0,0 +1,97 @@ +// 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" + "fmt" + "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" +) + +func BenchmarkIsInMemoPreSizing(b *testing.B) { + cases := []struct { + name string + typ arrow.DataType + valueSetSize int + cardinality int + }{ + {name: "int64/unique-1k", typ: arrow.PrimitiveTypes.Int64, valueSetSize: 1_000, cardinality: 1_000}, + {name: "int64/unique-64k", typ: arrow.PrimitiveTypes.Int64, valueSetSize: 64_000, cardinality: 64_000}, + {name: "int64/repeated-64k", typ: arrow.PrimitiveTypes.Int64, valueSetSize: 64_000, cardinality: 64}, + {name: "string/unique-1k", typ: arrow.BinaryTypes.String, valueSetSize: 1_000, cardinality: 1_000}, + {name: "string/unique-64k", typ: arrow.BinaryTypes.String, valueSetSize: 64_000, cardinality: 64_000}, + {name: "string/repeated-64k", typ: arrow.BinaryTypes.String, valueSetSize: 64_000, cardinality: 64}, + } + + for _, tc := range cases { + b.Run(tc.name, func(b *testing.B) { + mem := memory.DefaultAllocator + ctx := compute.WithAllocator(context.Background(), mem) + valueSet := newMemoBenchmarkArray(b, tc.typ, tc.valueSetSize, tc.cardinality) + defer valueSet.Release() + input := newMemoBenchmarkArray(b, tc.typ, 4_096, tc.cardinality*2) + defer input.Release() + + opts := compute.SetOptions{ + ValueSet: compute.NewDatumWithoutOwning(valueSet), + } + inputDatum := compute.NewDatumWithoutOwning(input) + + b.ReportAllocs() + b.ResetTimer() + for range b.N { + result, err := compute.IsIn(ctx, opts, inputDatum) + if err != nil { + b.Fatal(err) + } + result.Release() + } + }) + } +} + +func newMemoBenchmarkArray(b *testing.B, typ arrow.DataType, length, cardinality int) arrow.Array { + b.Helper() + switch typ.ID() { + case arrow.INT64: + builder := array.NewInt64Builder(memory.DefaultAllocator) + builder.Reserve(length) + for i := 0; i < length; i++ { + builder.Append(int64(i % cardinality)) + } + result := builder.NewArray() + builder.Release() + return result + case arrow.STRING: + builder := array.NewStringBuilder(memory.DefaultAllocator) + builder.Reserve(length) + for i := 0; i < length; i++ { + builder.Append(fmt.Sprintf("value-%08d", i%cardinality)) + } + result := builder.NewArray() + builder.Release() + return result + default: + b.Fatalf("unsupported benchmark type: %s", typ) + return nil + } +} From d844424a7984dcde9204bcd05d19bbf18b466d5f Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Fri, 28 Aug 2026 15:46:21 +0200 Subject: [PATCH 2/3] fix(compute): handle fixed binary extensions in set lookup --- .../internal/kernels/scalar_set_lookup.go | 8 ++- arrow/compute/scalar_set_lookup_test.go | 59 +++++++++++++++++++ 2 files changed, 64 insertions(+), 3 deletions(-) diff --git a/arrow/compute/internal/kernels/scalar_set_lookup.go b/arrow/compute/internal/kernels/scalar_set_lookup.go index 03dbfe4eb..cd4c525a6 100644 --- a/arrow/compute/internal/kernels/scalar_set_lookup.go +++ b/arrow/compute/internal/kernels/scalar_set_lookup.go @@ -165,9 +165,11 @@ func (s *SetLookupState[T]) Init(opts SetLookupOptions) error { s.NullBehavior = opts.NullBehavior s.MemoIndexToValueIndex = make([]int32, 0, opts.TotalLen) s.NullIndex = -1 - memoType := s.ValueSetType.ID() + valueSetType := s.ValueSetType + memoType := valueSetType.ID() if memoType == arrow.EXTENSION { - memoType = s.ValueSetType.(arrow.ExtensionType).StorageType().ID() + valueSetType = s.ValueSetType.(arrow.ExtensionType).StorageType() + memoType = valueSetType.ID() } // FixedSizeBinary with byte-widths 1/2/4/8 takes the numeric fast-path // in CreateSetLookupState (SetLookupState[uintN] + visitNumeric[uintN]), @@ -175,7 +177,7 @@ func (s *SetLookupState[T]) Init(opts SetLookupOptions) error { // the BinaryMemoTable that newMemoTable would otherwise return for // FIXED_SIZE_BINARY. if memoType == arrow.FIXED_SIZE_BINARY { - if fsb, ok := s.ValueSetType.(*arrow.FixedSizeBinaryType); ok { + if fsb, ok := valueSetType.(*arrow.FixedSizeBinaryType); ok { switch fsb.ByteWidth { case 1: memoType = arrow.UINT8 diff --git a/arrow/compute/scalar_set_lookup_test.go b/arrow/compute/scalar_set_lookup_test.go index 18b9a78e9..0b6de8fbc 100644 --- a/arrow/compute/scalar_set_lookup_test.go +++ b/arrow/compute/scalar_set_lookup_test.go @@ -20,6 +20,7 @@ import ( "context" "encoding/base64" "fmt" + "reflect" "strings" "sync" "testing" @@ -27,10 +28,34 @@ import ( "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/compute/exec" + "github.com/apache/arrow-go/v18/arrow/compute/internal/kernels" "github.com/apache/arrow-go/v18/arrow/memory" "github.com/stretchr/testify/suite" ) +type fixedSizeBinaryExtensionType struct { + arrow.ExtensionBase +} + +func (t *fixedSizeBinaryExtensionType) ArrayType() reflect.Type { + return reflect.TypeOf(fixedSizeBinaryExtensionArray{}) +} + +func (t *fixedSizeBinaryExtensionType) ExtensionEquals(other arrow.ExtensionType) bool { + return t.ExtensionName() == other.ExtensionName() && arrow.TypeEqual(t.StorageType(), other.StorageType()) +} + +func (*fixedSizeBinaryExtensionType) ExtensionName() string { return "compute.test.fixed_size_binary" } +func (*fixedSizeBinaryExtensionType) Serialize() string { return "" } +func (t *fixedSizeBinaryExtensionType) Deserialize(storageType arrow.DataType, _ string) (arrow.ExtensionType, error) { + return &fixedSizeBinaryExtensionType{ExtensionBase: arrow.ExtensionBase{Storage: storageType}}, nil +} + +type fixedSizeBinaryExtensionArray struct { + array.ExtensionArrayBase +} + type ScalarSetLookupSuite struct { suite.Suite @@ -297,6 +322,40 @@ func (ss *ScalarSetLookupSuite) TestIsInFixedSizeBinaryFastPaths() { } } +func (ss *ScalarSetLookupSuite) TestIsInFixedSizeBinaryExtensionFastPaths() { + for _, width := range []int{1, 2, 4, 8} { + width := width + ss.Run(fmt.Sprintf("ByteWidth=%d", width), func() { + typ := &fixedSizeBinaryExtensionType{ + ExtensionBase: arrow.ExtensionBase{ + Storage: &arrow.FixedSizeBinaryType{ByteWidth: width}, + }, + } + builder := array.NewFixedSizeBinaryBuilder(ss.mem, typ.StorageType().(*arrow.FixedSizeBinaryType)) + values := make([][]byte, 2) + for i := range values { + values[i] = make([]byte, width) + values[i][0] = byte(i + 1) + } + builder.AppendValues(values, nil) + storage := builder.NewFixedSizeBinaryArray() + builder.Release() + defer storage.Release() + + var span exec.ArraySpan + span.SetMembers(storage.Data()) + state, err := kernels.CreateSetLookupState(kernels.SetLookupOptions{ + ValueSetType: typ, + TotalLen: int64(storage.Len()), + ValueSet: []exec.ArraySpan{span}, + NullBehavior: kernels.NullMatchingMatch, + }, ss.mem) + ss.Require().NoError(err) + ss.Require().Equal(typ, state.(interface{ ValueType() arrow.DataType }).ValueType()) + }) + } +} + func (ss *ScalarSetLookupSuite) TestIsInDecimal() { type testCase struct { expected string From 9aee7b711556108ebdc2d39a225bbeb58b9b7309 Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Fri, 28 Aug 2026 15:57:56 +0200 Subject: [PATCH 3/3] fix(compute): read boolean values from bitmaps --- .../internal/kernels/scalar_set_lookup.go | 21 ++++++++ arrow/compute/scalar_set_lookup_test.go | 48 +++++++++++++++++++ 2 files changed, 69 insertions(+) diff --git a/arrow/compute/internal/kernels/scalar_set_lookup.go b/arrow/compute/internal/kernels/scalar_set_lookup.go index cd4c525a6..020b66180 100644 --- a/arrow/compute/internal/kernels/scalar_set_lookup.go +++ b/arrow/compute/internal/kernels/scalar_set_lookup.go @@ -50,6 +50,22 @@ func visitBinary[OffsetT int32 | int64](data *exec.ArraySpan, valid func([]byte) }, null) } +func visitBool(data *exec.ArraySpan, valid func(uint8) error, null func() error) error { + if data.Len == 0 { + return nil + } + + values := data.Buffers[1].Buf + return bitutils.VisitBitBlocksShort(data.Buffers[0].Buf, data.Offset, data.Len, + func(pos int64) error { + var value uint8 + if bitutil.BitIsSet(values, int(data.Offset+pos)) { + value = 1 + } + return valid(value) + }, null) +} + func visitNumeric[T arrow.FixedWidthType](data *exec.ArraySpan, valid func(T) error, null func() error) error { if data.Len == 0 { return nil @@ -108,6 +124,11 @@ func CreateSetLookupState(opts SetLookupOptions, alloc memory.Allocator) (exec.K visitFn: visitBinary[int64], } } + case *arrow.BooleanType: + state = &SetLookupState[uint8]{ + Alloc: alloc, + visitFn: visitBool, + } case arrow.FixedWidthDataType: switch ty.Bytes() { case 1: diff --git a/arrow/compute/scalar_set_lookup_test.go b/arrow/compute/scalar_set_lookup_test.go index 0b6de8fbc..835e5dd36 100644 --- a/arrow/compute/scalar_set_lookup_test.go +++ b/arrow/compute/scalar_set_lookup_test.go @@ -192,6 +192,54 @@ func (ss *ScalarSetLookupSuite) TestIsInPrimitive() { } } +func (ss *ScalarSetLookupSuite) TestIsInBoolean() { + for _, tc := range []struct { + name string + input string + valueset string + expected string + matching compute.NullMatchingBehavior + }{ + { + name: "no nulls", + input: `[false, true, false, true]`, + valueset: `[true]`, + expected: `[false, true, false, true]`, + matching: compute.NullMatchingMatch, + }, + { + name: "nulls in both", + input: `[false, true, null, false]`, + valueset: `[true, null]`, + expected: `[false, true, true, false]`, + matching: compute.NullMatchingMatch, + }, + { + name: "inconclusive nulls", + input: `[false, true, null, false]`, + valueset: `[true, null]`, + expected: `[null, true, null, null]`, + matching: compute.NullMatchingInconclusive, + }, + } { + ss.Run(tc.name, func() { + ss.checkIsInFromJSON(arrow.FixedWidthTypes.Boolean, + tc.input, tc.valueset, tc.expected, tc.matching) + }) + } + + input := ss.getArr(arrow.FixedWidthTypes.Boolean, `[false, true, false, true, false]`) + defer input.Release() + valueSet := ss.getArr(arrow.FixedWidthTypes.Boolean, `[false, true, false]`) + defer valueSet.Release() + + inputSlice := array.NewSlice(input, 1, 4) + defer inputSlice.Release() + valueSetSlice := array.NewSlice(valueSet, 1, 3) + defer valueSetSlice.Release() + ss.checkIsIn(inputSlice, valueSetSlice, `[true, true, true]`, compute.NullMatchingMatch) +} + func (ss *ScalarSetLookupSuite) TestDurationCasts() { vals := ss.getArr(arrow.FixedWidthTypes.Duration_s, `[0, 1, 2]`) defer vals.Release()