diff --git a/arrow/array/compare.go b/arrow/array/compare.go index 37ef0594..1f45b11d 100644 --- a/arrow/array/compare.go +++ b/arrow/array/compare.go @@ -354,6 +354,12 @@ func Equal(left, right arrow.Array) bool { // SliceEqual reports whether slices left[lbeg:lend] and right[rbeg:rend] are equal. func SliceEqual(left arrow.Array, lbeg, lend int64, right arrow.Array, rbeg, rend int64) bool { + if lbeg == 0 && lend == int64(left.Len()) && + rbeg == 0 && rend == int64(right.Len()) && + canEqualDirectly(left) && canEqualDirectly(right) { + return Equal(left, right) + } + l := NewSlice(left, lbeg, lend) defer l.Release() r := NewSlice(right, rbeg, rend) @@ -362,6 +368,27 @@ func SliceEqual(left arrow.Array, lbeg, lend int64, right arrow.Array, rbeg, ren return Equal(l, r) } +// canEqualDirectly reports whether Equal can handle an array without first +// normalizing it through NewSlice. Equal uses concrete type assertions, so +// generic arrow.Array implementations must keep the normalization path. +func canEqualDirectly(arr arrow.Array) bool { + switch arr.(type) { + case *Null, *Boolean, *FixedSizeBinary, *Binary, *String, + *LargeBinary, *LargeString, *BinaryView, *StringView, + *Int8, *Int16, *Int32, *Int64, *Uint8, *Uint16, *Uint32, *Uint64, + *Float16, *Float32, *Float64, + *Decimal32, *Decimal64, *Decimal128, *Decimal256, + *Date32, *Date64, *Time32, *Time64, *Timestamp, + *List, *LargeList, *ListView, *LargeListView, *FixedSizeList, + *Struct, *MonthInterval, *DayTimeInterval, *MonthDayNanoInterval, + *Duration, *Map, ExtensionArray, *Dictionary, *SparseUnion, + *DenseUnion, *RunEndEncoded: + return true + default: + return false + } +} + type listOffset interface { int32 | int64 } diff --git a/arrow/array/compare_test.go b/arrow/array/compare_test.go index 4c5d0373..b52fd888 100644 --- a/arrow/array/compare_test.go +++ b/arrow/array/compare_test.go @@ -85,6 +85,51 @@ func TestArraySliceEqual(t *testing.T) { } } +func TestArraySliceEqualFullRange(t *testing.T) { + builder := array.NewInt64Builder(memory.DefaultAllocator) + builder.AppendValues([]int64{1, 2, 3}, nil) + arr := builder.NewInt64Array() + builder.Release() + defer arr.Release() + + short := array.NewSlice(arr, 0, 2) + defer short.Release() + shifted := array.NewSlice(arr, 1, 3) + defer shifted.Release() + + assert.True(t, array.SliceEqual(arr, 0, int64(arr.Len()), arr, 0, int64(arr.Len()))) + assert.True(t, array.SliceEqual(shifted, 0, int64(shifted.Len()), shifted, 0, int64(shifted.Len()))) + assert.False(t, array.SliceEqual(arr, 0, int64(arr.Len()), short, 0, int64(short.Len()))) +} + +type arrayWrapper struct { + arrow.Array +} + +func TestArraySliceEqualFullRangeGenericArray(t *testing.T) { + builder := array.NewInt64Builder(memory.DefaultAllocator) + builder.AppendValues([]int64{1, 2, 3}, nil) + arr := builder.NewInt64Array() + builder.Release() + defer arr.Release() + + wrapped := arrayWrapper{Array: arr} + assert.True(t, array.SliceEqual( + wrapped, 0, int64(wrapped.Len()), + wrapped, 0, int64(wrapped.Len()), + )) + assert.True(t, array.SliceEqual( + arr, 0, int64(arr.Len()), + wrapped, 0, int64(wrapped.Len()), + )) + + left := arrow.NewChunked(arrow.PrimitiveTypes.Int64, []arrow.Array{wrapped}) + right := arrow.NewChunked(arrow.PrimitiveTypes.Int64, []arrow.Array{wrapped}) + assert.True(t, array.ChunkedEqual(left, right)) + left.Release() + right.Release() +} + func TestListEqualByValidRuns(t *testing.T) { for _, dt := range []arrow.DataType{ arrow.ListOf(arrow.PrimitiveTypes.Int32), diff --git a/arrow/array/slice_equal_bench_test.go b/arrow/array/slice_equal_bench_test.go new file mode 100644 index 00000000..24a99012 --- /dev/null +++ b/arrow/array/slice_equal_bench_test.go @@ -0,0 +1,124 @@ +// 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 array_test + +import ( + "testing" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/memory" +) + +func BenchmarkSliceEqualFullRange(b *testing.B) { + tests := []struct { + name string + newArray func() arrow.Array + }{ + {name: "int64_64", newArray: func() arrow.Array { + return makeSliceEqualInt64Array(64) + }}, + {name: "string_64", newArray: func() arrow.Array { + return makeSliceEqualStringArray(64) + }}, + } + + for _, test := range tests { + arr := test.newArray() + b.Run(test.name, func(b *testing.B) { + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + if !array.SliceEqual(arr, 0, int64(arr.Len()), arr, 0, int64(arr.Len())) { + b.Fatal("array should equal itself") + } + } + }) + arr.Release() + } +} + +func BenchmarkChunkedEqualFullChunks(b *testing.B) { + tests := []struct { + name string + numChunks int + chunkLength int + }{ + {name: "64chunks_1024values", numChunks: 64, chunkLength: 1024}, + {name: "1024chunks_64values", numChunks: 1024, chunkLength: 64}, + } + + for _, test := range tests { + left, right := makeSliceEqualChunkedArrays(test.numChunks, test.chunkLength) + b.Run(test.name, func(b *testing.B) { + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + if !array.ChunkedEqual(left, right) { + b.Fatal("chunked arrays should be equal") + } + } + }) + left.Release() + right.Release() + } +} + +func makeSliceEqualInt64Array(length int) arrow.Array { + builder := array.NewInt64Builder(memory.DefaultAllocator) + defer builder.Release() + + values := make([]int64, length) + for i := range values { + values[i] = int64(i) + } + builder.AppendValues(values, nil) + return builder.NewInt64Array() +} + +func makeSliceEqualStringArray(length int) arrow.Array { + builder := array.NewStringBuilder(memory.DefaultAllocator) + defer builder.Release() + + values := make([]string, length) + for i := range values { + values[i] = "value" + } + builder.AppendValues(values, nil) + return builder.NewStringArray() +} + +func makeSliceEqualChunkedArrays(numChunks, chunkLength int) (*arrow.Chunked, *arrow.Chunked) { + chunks := make([]arrow.Array, numChunks) + values := make([]int64, chunkLength) + for i := 0; i < numChunks; i++ { + builder := array.NewInt64Builder(memory.DefaultAllocator) + for j := range values { + values[j] = int64(i*chunkLength + j) + } + builder.AppendValues(values, nil) + chunks[i] = builder.NewInt64Array() + builder.Release() + } + + left := arrow.NewChunked(arrow.PrimitiveTypes.Int64, chunks) + right := arrow.NewChunked(arrow.PrimitiveTypes.Int64, chunks) + for _, chunk := range chunks { + chunk.Release() + } + return left, right +}