1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9// http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing, software
12// distributed under the License is distributed on an "AS IS" BASIS,
13// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14// See the License for the specific language governing permissions and
15// limitations under the License.
16
17package array_test
18
19import (
20	"fmt"
21	"math"
22	"testing"
23
24	"github.com/apache/arrow/go/v6/arrow"
25	"github.com/apache/arrow/go/v6/arrow/array"
26	"github.com/apache/arrow/go/v6/arrow/float16"
27	"github.com/apache/arrow/go/v6/arrow/internal/arrdata"
28	"github.com/apache/arrow/go/v6/arrow/memory"
29	"github.com/stretchr/testify/assert"
30)
31
32func TestArrayEqual(t *testing.T) {
33	for name, recs := range arrdata.Records {
34		t.Run(name, func(t *testing.T) {
35			rec := recs[0]
36			schema := rec.Schema()
37			for i, col := range rec.Columns() {
38				t.Run(schema.Field(i).Name, func(t *testing.T) {
39					arr := col
40					if !array.ArrayEqual(arr, arr) {
41						t.Fatalf("identical arrays should compare equal:\narray=%v", arr)
42					}
43					sub1 := array.NewSlice(arr, 1, int64(arr.Len()))
44					defer sub1.Release()
45
46					sub2 := array.NewSlice(arr, 0, int64(arr.Len()-1))
47					defer sub2.Release()
48
49					if array.ArrayEqual(sub1, sub2) && name != "nulls" {
50						t.Fatalf("non-identical arrays should not compare equal:\nsub1=%v\nsub2=%v\narrf=%v\n", sub1, sub2, arr)
51					}
52				})
53			}
54		})
55	}
56}
57
58func TestArraySliceEqual(t *testing.T) {
59	for name, recs := range arrdata.Records {
60		t.Run(name, func(t *testing.T) {
61			rec := recs[0]
62			schema := rec.Schema()
63			for i, col := range rec.Columns() {
64				t.Run(schema.Field(i).Name, func(t *testing.T) {
65					arr := col
66					if !array.ArraySliceEqual(
67						arr, 0, int64(arr.Len()),
68						arr, 0, int64(arr.Len()),
69					) {
70						t.Fatalf("identical slices should compare equal:\narray=%v", arr)
71					}
72					sub1 := array.NewSlice(arr, 1, int64(arr.Len()))
73					defer sub1.Release()
74
75					sub2 := array.NewSlice(arr, 0, int64(arr.Len()-1))
76					defer sub2.Release()
77
78					if array.ArraySliceEqual(sub1, 0, int64(sub1.Len()), sub2, 0, int64(sub2.Len())) && name != "nulls" {
79						t.Fatalf("non-identical slices should not compare equal:\nsub1=%v\nsub2=%v\narrf=%v\n", sub1, sub2, arr)
80					}
81				})
82			}
83		})
84	}
85}
86
87func TestArrayApproxEqual(t *testing.T) {
88	for name, recs := range arrdata.Records {
89		t.Run(name, func(t *testing.T) {
90			rec := recs[0]
91			schema := rec.Schema()
92			for i, col := range rec.Columns() {
93				t.Run(schema.Field(i).Name, func(t *testing.T) {
94					arr := col
95					if !array.ArrayApproxEqual(arr, arr) {
96						t.Fatalf("identical arrays should compare equal:\narray=%v", arr)
97					}
98					sub1 := array.NewSlice(arr, 1, int64(arr.Len()))
99					defer sub1.Release()
100
101					sub2 := array.NewSlice(arr, 0, int64(arr.Len()-1))
102					defer sub2.Release()
103
104					if array.ArrayApproxEqual(sub1, sub2) && name != "nulls" {
105						t.Fatalf("non-identical arrays should not compare equal:\nsub1=%v\nsub2=%v\narrf=%v\n", sub1, sub2, arr)
106					}
107				})
108			}
109		})
110	}
111}
112
113func TestArrayApproxEqualFloats(t *testing.T) {
114	f16sFrom := func(vs []float64) []float16.Num {
115		o := make([]float16.Num, len(vs))
116		for i, v := range vs {
117			o[i] = float16.New(float32(v))
118		}
119		return o
120	}
121
122	for _, tc := range []struct {
123		name string
124		a1   interface{}
125		a2   interface{}
126		opts []array.EqualOption
127		want bool
128	}{
129		{
130			name: "f16",
131			a1:   f16sFrom([]float64{1, 2, 3, 4, 5, 6}),
132			a2:   f16sFrom([]float64{1, 2, 3, 4, 5, 6}),
133			want: true,
134		},
135		{
136			name: "f16-no-tol",
137			a1:   f16sFrom([]float64{1, 2, 3, 4, 5, 6}),
138			a2:   f16sFrom([]float64{1, 2, 3, 4, 5, 7}),
139			want: false,
140		},
141		{
142			name: "f16-tol-ok",
143			a1:   f16sFrom([]float64{1, 2, 3, 4, 5, 6}),
144			a2:   f16sFrom([]float64{1, 2, 3, 4, 5, 7}),
145			opts: []array.EqualOption{array.WithAbsTolerance(1)},
146			want: true,
147		},
148		{
149			name: "f16-nan",
150			a1:   f16sFrom([]float64{1, 2, 3, 4, 5, 6}),
151			a2:   f16sFrom([]float64{1, 2, 3, 4, 5, math.NaN()}),
152			want: false,
153		},
154		{
155			name: "f16-nan-not",
156			a1:   f16sFrom([]float64{1, 2, 3, 4, 5, 6}),
157			a2:   f16sFrom([]float64{1, 2, 3, 4, 5, math.NaN()}),
158			opts: []array.EqualOption{array.WithNaNsEqual(true)},
159			want: false,
160		},
161		{
162			name: "f16-nan-ok",
163			a1:   f16sFrom([]float64{1, 2, 3, 4, 5, math.NaN()}),
164			a2:   f16sFrom([]float64{1, 2, 3, 4, 5, math.NaN()}),
165			opts: []array.EqualOption{array.WithNaNsEqual(true)},
166			want: true,
167		},
168		{
169			name: "f16-nan-no-tol",
170			a1:   f16sFrom([]float64{1, 2, 3, 4, 5, math.NaN()}),
171			a2:   f16sFrom([]float64{1, 2, 3, 4, 6, math.NaN()}),
172			opts: []array.EqualOption{array.WithNaNsEqual(true)},
173			want: false,
174		},
175		{
176			name: "f16-nan-tol",
177			a1:   f16sFrom([]float64{1, 2, 3, 4, 5, math.NaN()}),
178			a2:   f16sFrom([]float64{1, 2, 3, 4, 6, math.NaN()}),
179			opts: []array.EqualOption{array.WithNaNsEqual(true), array.WithAbsTolerance(1)},
180			want: true,
181		},
182		{
183			name: "f32",
184			a1:   []float32{1, 2, 3, 4, 5, 6},
185			a2:   []float32{1, 2, 3, 4, 5, 6},
186			want: true,
187		},
188		{
189			name: "f32-no-tol",
190			a1:   []float32{1, 2, 3, 4, 5, 6},
191			a2:   []float32{1, 2, 3, 4, 5, 7},
192			want: false,
193		},
194		{
195			name: "f32-tol-ok",
196			a1:   []float32{1, 2, 3, 4, 5, 6},
197			a2:   []float32{1, 2, 3, 4, 5, 7},
198			opts: []array.EqualOption{array.WithAbsTolerance(1)},
199			want: true,
200		},
201		{
202			name: "f32-nan",
203			a1:   []float32{1, 2, 3, 4, 5, 6},
204			a2:   []float32{1, 2, 3, 4, 5, float32(math.NaN())},
205			want: false,
206		},
207		{
208			name: "f32-nan-not",
209			a1:   []float32{1, 2, 3, 4, 5, 6},
210			a2:   []float32{1, 2, 3, 4, 5, float32(math.NaN())},
211			opts: []array.EqualOption{array.WithNaNsEqual(true)},
212			want: false,
213		},
214		{
215			name: "f32-nan-ok",
216			a1:   []float32{1, 2, 3, 4, 5, float32(math.NaN())},
217			a2:   []float32{1, 2, 3, 4, 5, float32(math.NaN())},
218			opts: []array.EqualOption{array.WithNaNsEqual(true)},
219			want: true,
220		},
221		{
222			name: "f32-nan-no-tol",
223			a1:   []float32{1, 2, 3, 4, 5, float32(math.NaN())},
224			a2:   []float32{1, 2, 3, 4, 6, float32(math.NaN())},
225			opts: []array.EqualOption{array.WithNaNsEqual(true)},
226			want: false,
227		},
228		{
229			name: "f32-nan-tol",
230			a1:   []float32{1, 2, 3, 4, 5, float32(math.NaN())},
231			a2:   []float32{1, 2, 3, 4, 6, float32(math.NaN())},
232			opts: []array.EqualOption{array.WithNaNsEqual(true), array.WithAbsTolerance(1)},
233			want: true,
234		},
235		{
236			name: "f64",
237			a1:   []float64{1, 2, 3, 4, 5, 6},
238			a2:   []float64{1, 2, 3, 4, 5, 6},
239			want: true,
240		},
241		{
242			name: "f64-no-tol",
243			a1:   []float64{1, 2, 3, 4, 5, 6},
244			a2:   []float64{1, 2, 3, 4, 5, 7},
245			want: false,
246		},
247		{
248			name: "f64-tol-ok",
249			a1:   []float64{1, 2, 3, 4, 5, 6},
250			a2:   []float64{1, 2, 3, 4, 5, 7},
251			opts: []array.EqualOption{array.WithAbsTolerance(1)},
252			want: true,
253		},
254		{
255			name: "f64-nan",
256			a1:   []float64{1, 2, 3, 4, 5, 6},
257			a2:   []float64{1, 2, 3, 4, 5, math.NaN()},
258			want: false,
259		},
260		{
261			name: "f64-nan-not",
262			a1:   []float64{1, 2, 3, 4, 5, 6},
263			a2:   []float64{1, 2, 3, 4, 5, math.NaN()},
264			opts: []array.EqualOption{array.WithNaNsEqual(true)},
265			want: false,
266		},
267		{
268			name: "f64-nan-ok",
269			a1:   []float64{1, 2, 3, 4, 5, math.NaN()},
270			a2:   []float64{1, 2, 3, 4, 5, math.NaN()},
271			opts: []array.EqualOption{array.WithNaNsEqual(true)},
272			want: true,
273		},
274		{
275			name: "f64-nan-no-tol",
276			a1:   []float64{1, 2, 3, 4, 5, math.NaN()},
277			a2:   []float64{1, 2, 3, 4, 6, math.NaN()},
278			opts: []array.EqualOption{array.WithNaNsEqual(true)},
279			want: false,
280		},
281		{
282			name: "f64-nan-tol",
283			a1:   []float64{1, 2, 3, 4, 5, math.NaN()},
284			a2:   []float64{1, 2, 3, 4, 6, math.NaN()},
285			opts: []array.EqualOption{array.WithNaNsEqual(true), array.WithAbsTolerance(1)},
286			want: true,
287		},
288	} {
289		t.Run(tc.name, func(t *testing.T) {
290			mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
291			defer mem.AssertSize(t, 0)
292
293			a1 := arrayOf(mem, tc.a1, nil)
294			defer a1.Release()
295			a2 := arrayOf(mem, tc.a2, nil)
296			defer a2.Release()
297
298			if got, want := array.ArrayApproxEqual(a1, a2, tc.opts...), tc.want; got != want {
299				t.Fatalf("invalid comparison: got=%v, want=%v\na1: %v\na2: %v\n", got, want, a1, a2)
300			}
301		})
302	}
303}
304
305func arrayOf(mem memory.Allocator, a interface{}, valids []bool) array.Interface {
306	if mem == nil {
307		mem = memory.NewGoAllocator()
308	}
309
310	switch a := a.(type) {
311	case []float16.Num:
312		bldr := array.NewFloat16Builder(mem)
313		defer bldr.Release()
314
315		bldr.AppendValues(a, valids)
316		return bldr.NewFloat16Array()
317
318	case []float32:
319		bldr := array.NewFloat32Builder(mem)
320		defer bldr.Release()
321
322		bldr.AppendValues(a, valids)
323		return bldr.NewFloat32Array()
324
325	case []float64:
326		bldr := array.NewFloat64Builder(mem)
327		defer bldr.Release()
328
329		bldr.AppendValues(a, valids)
330		return bldr.NewFloat64Array()
331
332	default:
333		panic(fmt.Errorf("arrdata: invalid data slice type %T", a))
334	}
335}
336
337func TestArrayEqualBaseArray(t *testing.T) {
338	mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
339	defer mem.AssertSize(t, 0)
340
341	b1 := array.NewBooleanBuilder(mem)
342	defer b1.Release()
343	b1.Append(true)
344	a1 := b1.NewBooleanArray()
345	defer a1.Release()
346
347	b2 := array.NewBooleanBuilder(mem)
348	defer b2.Release()
349	a2 := b2.NewBooleanArray()
350	defer a2.Release()
351
352	if array.ArrayEqual(a1, a2) {
353		t.Errorf("two arrays with different lengths must not be equal")
354	}
355
356	b3 := array.NewBooleanBuilder(mem)
357	defer b3.Release()
358	b3.AppendNull()
359	a3 := b3.NewBooleanArray()
360	defer a3.Release()
361
362	if array.ArrayEqual(a1, a3) {
363		t.Errorf("two arrays with different number of null values must not be equal")
364	}
365
366	b4 := array.NewInt32Builder(mem)
367	defer b4.Release()
368	b4.Append(0)
369	a4 := b4.NewInt32Array()
370	defer a4.Release()
371
372	if array.ArrayEqual(a1, a4) {
373		t.Errorf("two arrays with different types must not be equal")
374	}
375
376	b5 := array.NewBooleanBuilder(mem)
377	defer b5.Release()
378	b5.AppendNull()
379	b5.Append(true)
380	a5 := b5.NewBooleanArray()
381	defer a5.Release()
382	b1.AppendNull()
383
384	if array.ArrayEqual(a1, a5) {
385		t.Errorf("two arrays with different validity bitmaps must not be equal")
386	}
387}
388
389func TestArrayEqualNull(t *testing.T) {
390	mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
391	defer mem.AssertSize(t, 0)
392
393	null := array.NewNull(0)
394	defer null.Release()
395
396	if !array.ArrayEqual(null, null) {
397		t.Fatalf("identical arrays should compare equal")
398	}
399
400	n0 := array.NewNull(10)
401	defer n0.Release()
402
403	n1 := array.NewNull(10)
404	defer n1.Release()
405
406	if !array.ArrayEqual(n0, n0) {
407		t.Fatalf("identical arrays should compare equal")
408	}
409	if !array.ArrayEqual(n1, n1) {
410		t.Fatalf("identical arrays should compare equal")
411	}
412	if !array.ArrayEqual(n0, n1) || !array.ArrayEqual(n1, n0) {
413		t.Fatalf("n0 and n1 should compare equal")
414	}
415
416	sub07 := array.NewSlice(n0, 0, 7)
417	defer sub07.Release()
418	sub08 := array.NewSlice(n0, 0, 8)
419	defer sub08.Release()
420	sub19 := array.NewSlice(n0, 1, 9)
421	defer sub19.Release()
422
423	if !array.ArrayEqual(sub08, sub19) {
424		t.Fatalf("sub08 and sub19 should compare equal")
425	}
426
427	if array.ArrayEqual(sub08, sub07) {
428		t.Fatalf("sub08 and sub07 should not compare equal")
429	}
430}
431
432func TestArrayEqualMaskedArray(t *testing.T) {
433	mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
434	defer mem.AssertSize(t, 0)
435
436	ab := array.NewInt32Builder(mem)
437	defer ab.Release()
438
439	valids := []bool{false, false, false, false}
440	ab.AppendValues([]int32{1, 2, 0, 4}, valids)
441
442	a1 := ab.NewInt32Array()
443	defer a1.Release()
444
445	ab.AppendValues([]int32{1, 2, 3, 4}, valids)
446	a2 := ab.NewInt32Array()
447	defer a2.Release()
448
449	if !array.ArrayEqual(a1, a1) || !array.ArrayEqual(a2, a2) {
450		t.Errorf("an array must be equal to itself")
451	}
452
453	if !array.ArrayEqual(a1, a2) {
454		t.Errorf("%v must be equal to %v", a1, a2)
455	}
456}
457
458func TestArrayEqualDifferentMaskedValues(t *testing.T) {
459	// test 2 int32 arrays, with same nulls (but different masked values) compare equal.
460	mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
461	defer mem.AssertSize(t, 0)
462
463	ab := array.NewInt32Builder(mem)
464	defer ab.Release()
465
466	valids := []bool{true, true, false, true}
467	ab.AppendValues([]int32{1, 2, 0, 4}, valids)
468
469	a1 := ab.NewInt32Array()
470	defer a1.Release()
471
472	ab.AppendValues([]int32{1, 2, 3, 4}, valids)
473	a2 := ab.NewInt32Array()
474	defer a2.Release()
475
476	if !array.ArrayEqual(a1, a1) || !array.ArrayEqual(a2, a2) {
477		t.Errorf("an array must be equal to itself")
478	}
479
480	if !array.ArrayEqual(a1, a2) {
481		t.Errorf("%v must be equal to %v", a1, a2)
482	}
483}
484
485func TestRecordEqual(t *testing.T) {
486	for name, recs := range arrdata.Records {
487		t.Run(name, func(t *testing.T) {
488			rec0 := recs[0]
489			rec1 := recs[1]
490			if !array.RecordEqual(rec0, rec0) {
491				t.Fatalf("identical records should compare equal:\nrecord:\n%v", rec0)
492			}
493
494			if array.RecordEqual(rec0, rec1) && name != "nulls" {
495				t.Fatalf("non-identical records should not compare equal:\nrec0:\n%v\nrec1:\n%v", rec0, rec1)
496			}
497
498			sub00 := rec0.NewSlice(0, recs[0].NumRows()-1)
499			defer sub00.Release()
500			sub01 := rec0.NewSlice(1, recs[0].NumRows())
501			defer sub01.Release()
502
503			if array.RecordEqual(sub00, sub01) && name != "nulls" {
504				t.Fatalf("non-identical records should not compare equal:\nsub0:\n%v\nsub1:\n%v", sub00, sub01)
505			}
506		})
507	}
508}
509
510func TestRecordApproxEqual(t *testing.T) {
511	for name, recs := range arrdata.Records {
512		t.Run(name, func(t *testing.T) {
513			rec0 := recs[0]
514			rec1 := recs[1]
515			if !array.RecordApproxEqual(rec0, rec0) {
516				t.Fatalf("identical records should compare equal:\nrecord:\n%v", rec0)
517			}
518
519			if array.RecordApproxEqual(rec0, rec1) && name != "nulls" {
520				t.Fatalf("non-identical records should not compare equal:\nrec0:\n%v\nrec1:\n%v", rec0, rec1)
521			}
522
523			sub00 := rec0.NewSlice(0, recs[0].NumRows()-1)
524			defer sub00.Release()
525			sub01 := rec0.NewSlice(1, recs[0].NumRows())
526			defer sub01.Release()
527
528			if array.RecordApproxEqual(sub00, sub01) && name != "nulls" {
529				t.Fatalf("non-identical records should not compare equal:\nsub0:\n%v\nsub1:\n%v", sub00, sub01)
530			}
531		})
532	}
533}
534
535func TestChunkedEqual(t *testing.T) {
536	for name, recs := range arrdata.Records {
537		t.Run(name, func(t *testing.T) {
538			tbl := array.NewTableFromRecords(recs[0].Schema(), recs)
539			defer tbl.Release()
540
541			for i := 0; i < int(tbl.NumCols()); i++ {
542				if !array.ChunkedEqual(tbl.Column(i).Data(), tbl.Column(i).Data()) && name != "nulls" {
543					t.Fatalf("identical chunked arrays should compare as equal:\narr:%v\n", tbl.Column(i).Data())
544				}
545			}
546		})
547	}
548}
549
550func TestChunkedApproxEqual(t *testing.T) {
551	fb := array.NewFloat64Builder(memory.DefaultAllocator)
552	defer fb.Release()
553
554	fb.AppendValues([]float64{1, 2, 3, 4, 5}, nil)
555	f1 := fb.NewFloat64Array()
556	defer f1.Release()
557
558	fb.AppendValues([]float64{6, 7}, nil)
559	f2 := fb.NewFloat64Array()
560	defer f2.Release()
561
562	fb.AppendValues([]float64{8, 9, 10}, nil)
563	f3 := fb.NewFloat64Array()
564	defer f3.Release()
565
566	c1 := array.NewChunked(
567		arrow.PrimitiveTypes.Float64,
568		[]array.Interface{f1, f2, f3},
569	)
570	defer c1.Release()
571
572	fb.AppendValues([]float64{1, 2, 3}, nil)
573	f4 := fb.NewFloat64Array()
574	defer f4.Release()
575
576	fb.AppendValues([]float64{4, 5}, nil)
577	f5 := fb.NewFloat64Array()
578	defer f5.Release()
579
580	fb.AppendValues([]float64{6, 7, 8, 9}, nil)
581	f6 := fb.NewFloat64Array()
582	defer f6.Release()
583
584	fb.AppendValues([]float64{10}, nil)
585	f7 := fb.NewFloat64Array()
586	defer f7.Release()
587
588	c2 := array.NewChunked(
589		arrow.PrimitiveTypes.Float64,
590		[]array.Interface{f4, f5, f6, f7},
591	)
592	defer c2.Release()
593
594	assert.True(t, array.ChunkedEqual(c1, c2))
595	assert.True(t, array.ChunkedApproxEqual(c1, c2))
596}
597
598func TestTableEqual(t *testing.T) {
599	for name, recs := range arrdata.Records {
600		t.Run(name, func(t *testing.T) {
601			tbl := array.NewTableFromRecords(recs[0].Schema(), recs)
602			defer tbl.Release()
603
604			if !array.TableEqual(tbl, tbl) {
605				t.Fatalf("identical tables should compare as equal:\tbl:%v\n", tbl)
606			}
607			if !array.TableApproxEqual(tbl, tbl) {
608				t.Fatalf("identical tables should compare as approx equal:\tbl:%v\n", tbl)
609			}
610		})
611	}
612}
613