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