Skip to content

Commit 3695a54

Browse files
committed
fix arraydata equal func
1 parent 0666d21 commit 3695a54

3 files changed

Lines changed: 58 additions & 1 deletion

File tree

arrow/array.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,8 @@ type ArrayData interface {
8383
Dictionary() ArrayData
8484
// SizeInBytes returns the size of the ArrayData buffers and any children and/or dictionary in bytes.
8585
SizeInBytes() uint64
86+
87+
Equal(ArrayData) bool
8688
}
8789

8890
// Array represents an immutable sequence of values using the Arrow in-memory format.

arrow/array/data.go

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
package array
1818

1919
import (
20+
"bytes"
2021
"hash/maphash"
2122
"math/bits"
2223
"sync/atomic"
@@ -77,6 +78,59 @@ func NewDataWithDictionary(dtype arrow.DataType, length int, buffers []*memory.B
7778
return data
7879
}
7980

81+
func (d *Data) Equal(other arrow.ArrayData) bool {
82+
rhs, ok := other.(*Data)
83+
if !ok {
84+
return false
85+
}
86+
87+
if d == rhs {
88+
return true
89+
}
90+
91+
switch {
92+
case !arrow.TypeEqual(d.dtype, rhs.dtype):
93+
return false
94+
case d.length != rhs.length || d.nulls != rhs.nulls || d.offset != rhs.offset:
95+
return false
96+
case len(d.buffers) != len(rhs.buffers):
97+
return false
98+
case len(d.childData) != len(rhs.childData):
99+
return false
100+
case d.dictionary != nil && rhs.dictionary == nil:
101+
return false
102+
case d.dictionary == nil && rhs.dictionary != nil:
103+
return false
104+
}
105+
106+
if d.dictionary != nil {
107+
if !d.dictionary.Equal(rhs.dictionary) {
108+
return false
109+
}
110+
}
111+
112+
for i := range d.childData {
113+
if !d.childData[i].Equal(rhs.childData[i]) {
114+
return false
115+
}
116+
}
117+
118+
for i, b := range d.buffers {
119+
switch {
120+
case b == nil:
121+
if rhs.buffers[i] != nil {
122+
return false
123+
}
124+
case rhs.buffers[i] == nil:
125+
return false
126+
case !bytes.Equal(b.Bytes(), rhs.buffers[i].Bytes()):
127+
return false
128+
}
129+
}
130+
131+
return true
132+
}
133+
80134
func (d *Data) Copy() *Data {
81135
// don't pass the slices directly, otherwise it retains the connection
82136
// we need to make new slices and populate them with the same pointers

arrow/compute/exec/span_test.go

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -406,7 +406,8 @@ func TestArraySpan_MakeData(t *testing.T) {
406406
}
407407
got := a.MakeData()
408408
want := tt.want(mem)
409-
if !reflect.DeepEqual(got, want) {
409+
410+
if !got.Equal(want) {
410411
t.Errorf("ArraySpan.MakeData() = %v, want %v", got, want)
411412
}
412413
want.Release()

0 commit comments

Comments
 (0)