diff --git a/arrow/compute/internal/kernels/scalar_comparisons.go b/arrow/compute/internal/kernels/scalar_comparisons.go index aed44da6..4dfe00cc 100644 --- a/arrow/compute/internal/kernels/scalar_comparisons.go +++ b/arrow/compute/internal/kernels/scalar_comparisons.go @@ -55,24 +55,28 @@ func comparePrimitiveArrayArray[T arrow.FixedWidthType](op cmpFn[T, T]) binaryKe left = arrow.GetData[T](leftBytes) right = arrow.GetData[T](rightBytes) nvals = len(left) - nbatches = nvals / batchSize tmpOutput [batchSize]uint32 ) tmpOutSlice := tmpOutput[:] if prefix := offset % 8; prefix != 0 { - vals := 8 - prefix + vals := min(nvals, 8-prefix) op(left[:vals], right[:vals], tmpOutSlice[:vals]) left, right = left[vals:], right[vals:] + nvals -= vals for i, v := range tmpOutSlice[:vals] { bitutil.SetBitTo(out, prefix+i, v != 0) } + if nvals == 0 { + return + } out = out[1:] } + nbatches := nvals / batchSize for j := 0; j < nbatches; j++ { - op(left, right, tmpOutSlice) + op(left[:batchSize], right[:batchSize], tmpOutSlice) left, right = left[batchSize:], right[batchSize:] packBits(tmpOutput, out) out = out[batchSize/8:] @@ -93,22 +97,26 @@ func comparePrimitiveArrayScalar[T arrow.FixedWidthType](op cmpScalarRight[T, T] left = arrow.GetData[T](leftBytes) rightVal = *(*T)(unsafe.Pointer(&rightBytes[0])) nvals = len(left) - nbatches = nvals / batchSize tmpOutput [batchSize]uint32 ) tmpOutSlice := tmpOutput[:] if prefix := offset % 8; prefix != 0 { - vals := 8 - prefix + vals := min(nvals, 8-prefix) op(left[:vals], rightVal, tmpOutSlice[:vals]) left = left[vals:] + nvals -= vals for i, v := range tmpOutSlice[:vals] { bitutil.SetBitTo(out, prefix+i, v != 0) } + if nvals == 0 { + return + } out = out[1:] } + nbatches := nvals / batchSize for j := 0; j < nbatches; j++ { op(left[:batchSize], rightVal, tmpOutSlice) left = left[batchSize:] @@ -132,22 +140,26 @@ func comparePrimitiveScalarArray[T arrow.FixedWidthType](op cmpScalarLeft[T, T]) right = arrow.GetData[T](rightBytes) nvals = len(right) - nbatches = nvals / batchSize tmpOutput [batchSize]uint32 ) tmpOutSlice := tmpOutput[:] if prefix := offset % 8; prefix != 0 { - vals := 8 - prefix + vals := min(nvals, 8-prefix) op(leftVal, right[:vals], tmpOutSlice[:vals]) right = right[vals:] + nvals -= vals for i, v := range tmpOutSlice[:vals] { bitutil.SetBitTo(out, prefix+i, v != 0) } + if nvals == 0 { + return + } out = out[1:] } + nbatches := nvals / batchSize for j := 0; j < nbatches; j++ { op(leftVal, right[:batchSize], tmpOutSlice) right = right[batchSize:] diff --git a/arrow/compute/internal/kernels/scalar_comparisons_test.go b/arrow/compute/internal/kernels/scalar_comparisons_test.go new file mode 100644 index 00000000..e99bb01d --- /dev/null +++ b/arrow/compute/internal/kernels/scalar_comparisons_test.go @@ -0,0 +1,102 @@ +// 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. + +//go:build go1.18 + +package kernels + +import ( + "bytes" + "fmt" + "testing" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/bitutil" +) + +func TestPrimitiveComparisonsWithOutputOffset(t *testing.T) { + const maxLength = 65 + + left := make([]int32, maxLength) + right := make([]int32, maxLength) + for i := range left { + left[i] = int32((i*7 + 1) % 11) + right[i] = int32((i*5 + 3) % 11) + } + + type operation struct { + op CompareOperator + eval func(int32, int32) bool + } + operations := []operation{ + {CmpEQ, func(l, r int32) bool { return l == r }}, + {CmpNE, func(l, r int32) bool { return l != r }}, + {CmpGT, func(l, r int32) bool { return l > r }}, + {CmpGE, func(l, r int32) bool { return l >= r }}, + } + + leftBytes := arrow.Int32Traits.CastToBytes(left) + rightBytes := arrow.Int32Traits.CastToBytes(right) + leftScalar := int32(6) + rightScalar := int32(4) + leftScalarBytes := arrow.Int32Traits.CastToBytes([]int32{leftScalar}) + rightScalarBytes := arrow.Int32Traits.CastToBytes([]int32{rightScalar}) + + for _, operation := range operations { + t.Run(operation.op.String(), func(t *testing.T) { + cmp := getCmpOp[int32](operation.op) + for _, shape := range []string{"array_array", "array_scalar", "scalar_array"} { + t.Run(shape, func(t *testing.T) { + for offset := 0; offset < 8; offset++ { + t.Run(fmt.Sprintf("offset_%d", offset), func(t *testing.T) { + for length := 0; length <= maxLength; length++ { + out := bytes.Repeat([]byte{0xa5}, int(bitutil.BytesForBits(int64(offset+length)))) + expected := append([]byte(nil), out...) + + for i := 0; i < length; i++ { + var result bool + switch shape { + case "array_array": + result = operation.eval(left[i], right[i]) + case "array_scalar": + result = operation.eval(left[i], rightScalar) + case "scalar_array": + result = operation.eval(leftScalar, right[i]) + } + bitutil.SetBitTo(expected, offset+i, result) + } + + switch shape { + case "array_array": + comparePrimitiveArrayArray(cmp.arrArr)(leftBytes[:length*4], rightBytes[:length*4], out, offset) + case "array_scalar": + comparePrimitiveArrayScalar(cmp.arrScalar)(leftBytes[:length*4], rightScalarBytes, out, offset) + case "scalar_array": + comparePrimitiveScalarArray(cmp.scalarArr)(leftScalarBytes, rightBytes[:length*4], out, offset) + } + + if !bytes.Equal(expected, out) { + t.Fatalf("length %d: expected %08b, got %08b", length, expected, out) + } + } + }) + } + }) + } + }) + } +}