perf(compute): pre-size SetLookupState memo tables

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:
@@ -165,17 +186,19 @@ 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]),
 	// so the lookup table must be the matching TypedMemoTable[uintN], not
 	// 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
@@ -188,7 +211,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
 	}

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
 		}

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
+	}
+}

arrow/compute/scalar_set_lookup_test.go

@@ -20,17 +20,42 @@ import (
 	"context"
 	"encoding/base64"
 	"fmt"
+	"reflect"
 	"strings"
 	"sync"
 	"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/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
 
@@ -167,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()
@@ -297,6 +370,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