github.com/apache/arrow/go/v15@v15.0.1/arrow/compute/exec/span_test.go (about) 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 17 //go:build go1.18 18 19 package exec_test 20 21 import ( 22 "reflect" 23 "strings" 24 "testing" 25 "unsafe" 26 27 "github.com/apache/arrow/go/v15/arrow" 28 "github.com/apache/arrow/go/v15/arrow/array" 29 "github.com/apache/arrow/go/v15/arrow/compute/exec" 30 "github.com/apache/arrow/go/v15/arrow/decimal128" 31 "github.com/apache/arrow/go/v15/arrow/endian" 32 "github.com/apache/arrow/go/v15/arrow/memory" 33 "github.com/apache/arrow/go/v15/arrow/scalar" 34 "github.com/apache/arrow/go/v15/internal/types" 35 "github.com/stretchr/testify/assert" 36 ) 37 38 func TestBufferSpan_SetBuffer(t *testing.T) { 39 type fields struct { 40 Buf []byte 41 Owner *memory.Buffer 42 SelfAlloc bool 43 } 44 type args struct { 45 buf *memory.Buffer 46 } 47 foo := []byte{0xde, 0xad, 0xbe, 0xef} 48 own := memory.NewBufferBytes(foo) 49 tests := []struct { 50 name string 51 fields fields 52 args args 53 }{ 54 {"simple", fields{SelfAlloc: true}, args{own}}, 55 } 56 for _, tt := range tests { 57 t.Run(tt.name, func(t *testing.T) { 58 b := &exec.BufferSpan{ 59 Buf: tt.fields.Buf, 60 Owner: tt.fields.Owner, 61 SelfAlloc: tt.fields.SelfAlloc, 62 } 63 b.SetBuffer(tt.args.buf) 64 assert.Same(t, &foo[0], &b.Buf[0]) 65 assert.Same(t, own, b.Owner) 66 assert.False(t, b.SelfAlloc) 67 }) 68 } 69 } 70 71 func TestBufferSpan_WrapBuffer(t *testing.T) { 72 type fields struct { 73 Buf []byte 74 Owner *memory.Buffer 75 SelfAlloc bool 76 } 77 type args struct { 78 buf *memory.Buffer 79 } 80 foo := []byte{0xde, 0xad, 0xbe, 0xef} 81 own := memory.NewBufferBytes(foo) 82 tests := []struct { 83 name string 84 fields fields 85 args args 86 }{ 87 {"simple", fields{SelfAlloc: false}, args{own}}, 88 } 89 for _, tt := range tests { 90 t.Run(tt.name, func(t *testing.T) { 91 b := &exec.BufferSpan{ 92 Buf: tt.fields.Buf, 93 Owner: tt.fields.Owner, 94 SelfAlloc: tt.fields.SelfAlloc, 95 } 96 b.WrapBuffer(tt.args.buf) 97 assert.Same(t, &foo[0], &b.Buf[0]) 98 assert.Same(t, own, b.Owner) 99 assert.True(t, b.SelfAlloc) 100 }) 101 } 102 } 103 104 func TestArraySpan_UpdateNullCount(t *testing.T) { 105 type fields struct { 106 Type arrow.DataType 107 Len int64 108 Nulls int64 109 Offset int64 110 Buffers [3]exec.BufferSpan 111 Scratch [2]uint64 112 Children []exec.ArraySpan 113 } 114 tests := []struct { 115 name string 116 fields fields 117 want int64 118 }{ 119 {"known", fields{Nulls: 25}, 25}, 120 {"unknown", fields{ 121 Nulls: array.UnknownNullCount, 122 Len: 8, // 0b01101101 123 Buffers: [3]exec.BufferSpan{{Buf: []byte{109}}, {}, {}}}, 3}, 124 {"unknown with offset", fields{ 125 Nulls: array.UnknownNullCount, 126 Len: 4, 127 Offset: 2, // 0b01101101 128 Buffers: [3]exec.BufferSpan{{Buf: []byte{109}}, {}, {}}}, 1}, 129 } 130 for _, tt := range tests { 131 t.Run(tt.name, func(t *testing.T) { 132 a := &exec.ArraySpan{ 133 Type: tt.fields.Type, 134 Len: tt.fields.Len, 135 Nulls: tt.fields.Nulls, 136 Offset: tt.fields.Offset, 137 Buffers: tt.fields.Buffers, 138 Scratch: tt.fields.Scratch, 139 Children: tt.fields.Children, 140 } 141 if got := a.UpdateNullCount(); got != tt.want { 142 t.Errorf("ArraySpan.UpdateNullCount() = %v, want %v", got, tt.want) 143 } 144 }) 145 } 146 } 147 148 func TestArraySpan_Dictionary(t *testing.T) { 149 type fields struct { 150 Type arrow.DataType 151 Len int64 152 Nulls int64 153 Offset int64 154 Buffers [3]exec.BufferSpan 155 Scratch [2]uint64 156 Children []exec.ArraySpan 157 } 158 children := []exec.ArraySpan{{}} 159 tests := []struct { 160 name string 161 fields fields 162 want *exec.ArraySpan 163 }{ 164 {"basic", fields{Children: children}, &children[0]}, 165 } 166 for _, tt := range tests { 167 t.Run(tt.name, func(t *testing.T) { 168 a := &exec.ArraySpan{ 169 Type: tt.fields.Type, 170 Len: tt.fields.Len, 171 Nulls: tt.fields.Nulls, 172 Offset: tt.fields.Offset, 173 Buffers: tt.fields.Buffers, 174 Scratch: tt.fields.Scratch, 175 Children: tt.fields.Children, 176 } 177 if got := a.Dictionary(); !reflect.DeepEqual(got, tt.want) { 178 t.Errorf("ArraySpan.Dictionary() = %v, want %v", got, tt.want) 179 } 180 }) 181 } 182 } 183 184 func TestArraySpan_NumBuffers(t *testing.T) { 185 type fields struct { 186 Type arrow.DataType 187 Len int64 188 Nulls int64 189 Offset int64 190 Buffers [3]exec.BufferSpan 191 Scratch [2]uint64 192 Children []exec.ArraySpan 193 } 194 195 arrow.RegisterExtensionType(types.NewUUIDType()) 196 defer arrow.UnregisterExtensionType("uuid") 197 198 tests := []struct { 199 name string 200 fields fields 201 want int 202 }{ 203 {"null", fields{Type: arrow.Null}, 1}, 204 {"struct", fields{Type: arrow.StructOf()}, 1}, 205 {"fixed size list", fields{Type: arrow.FixedSizeListOf(4, arrow.PrimitiveTypes.Int32)}, 1}, 206 {"binary", fields{Type: arrow.BinaryTypes.Binary}, 3}, 207 {"large binary", fields{Type: arrow.BinaryTypes.LargeBinary}, 3}, 208 {"string", fields{Type: arrow.BinaryTypes.String}, 3}, 209 {"large string", fields{Type: arrow.BinaryTypes.LargeString}, 3}, 210 {"extension", fields{Type: types.NewUUIDType()}, 2}, 211 {"int32", fields{Type: arrow.PrimitiveTypes.Int32}, 2}, 212 } 213 for _, tt := range tests { 214 t.Run(tt.name, func(t *testing.T) { 215 a := &exec.ArraySpan{ 216 Type: tt.fields.Type, 217 Len: tt.fields.Len, 218 Nulls: tt.fields.Nulls, 219 Offset: tt.fields.Offset, 220 Buffers: tt.fields.Buffers, 221 Scratch: tt.fields.Scratch, 222 Children: tt.fields.Children, 223 } 224 if got := a.NumBuffers(); got != tt.want { 225 t.Errorf("ArraySpan.NumBuffers() = %v, want %v", got, tt.want) 226 } 227 }) 228 } 229 } 230 231 func TestArraySpan_MakeData(t *testing.T) { 232 type fields struct { 233 Type arrow.DataType 234 Len int64 235 Nulls int64 236 Offset int64 237 Buffers [3]exec.BufferSpan 238 Scratch [2]uint64 239 Children []exec.ArraySpan 240 } 241 242 var ( 243 buf1 *memory.Buffer 244 ) 245 arrow.RegisterExtensionType(types.NewDictExtensionType()) 246 defer arrow.UnregisterExtensionType("dict-extension") 247 248 tests := []struct { 249 name string 250 fields func(mem memory.Allocator) fields 251 want func(mem memory.Allocator) arrow.ArrayData 252 }{ 253 {"null type", func(mem memory.Allocator) fields { 254 return fields{ 255 Type: arrow.Null, 256 Len: 5, 257 Nulls: array.UnknownNullCount, 258 } 259 }, func(mem memory.Allocator) arrow.ArrayData { 260 return array.NewData(arrow.Null, 5, []*memory.Buffer{nil}, nil, 5, 0) 261 }}, 262 {"zero len", func(mem memory.Allocator) fields { 263 return fields{Type: arrow.PrimitiveTypes.Int32} 264 }, func(mem memory.Allocator) arrow.ArrayData { 265 return array.NewData(arrow.PrimitiveTypes.Int32, 0, []*memory.Buffer{nil, nil}, nil, 0, 0) 266 }}, 267 {"non-owning offset", func(mem memory.Allocator) fields { 268 ret := fields{ 269 Type: arrow.PrimitiveTypes.Int8, 270 Len: 4, 271 Nulls: 1, 272 Offset: 1, 273 } 274 buf1 = memory.NewResizableBuffer(mem) 275 buf1.Resize(1) 276 buf1.Bytes()[0] = 109 277 ret.Buffers[0].SetBuffer(buf1) 278 ret.Buffers[1].SetBuffer(memory.NewBufferBytes([]byte{5, 5, 5, 5, 5})) 279 return ret 280 }, func(mem memory.Allocator) arrow.ArrayData { 281 // created in the above func, we release after constructing 282 // the NewData so the refcount is as expected 283 defer buf1.Release() 284 return array.NewData(arrow.PrimitiveTypes.Int8, 4, 285 []*memory.Buffer{buf1, memory.NewBufferBytes([]byte{5, 5, 5, 5, 5})}, nil, 1, 1) 286 }}, 287 {"self-alloc", func(mem memory.Allocator) fields { 288 ret := fields{ 289 Type: arrow.PrimitiveTypes.Int8, 290 Len: 4, 291 } 292 buf := memory.NewResizableBuffer(mem) 293 buf.Resize(1) 294 ret.Buffers[0].WrapBuffer(buf) 295 buf2 := memory.NewResizableBuffer(mem) 296 buf2.Resize(4) 297 ret.Buffers[1].WrapBuffer(buf2) 298 return ret 299 }, func(mem memory.Allocator) arrow.ArrayData { 300 buf := memory.NewResizableBuffer(mem) 301 buf.Resize(1) 302 defer buf.Release() 303 buf2 := memory.NewResizableBuffer(mem) 304 buf2.Resize(4) 305 defer buf2.Release() 306 return array.NewData(arrow.PrimitiveTypes.Int8, 4, []*memory.Buffer{buf, buf2}, nil, 0, 0) 307 }}, 308 {"with children", func(mem memory.Allocator) fields { 309 ret := fields{ 310 Type: arrow.ListOf(arrow.PrimitiveTypes.Int8), 311 Len: 1, 312 Children: []exec.ArraySpan{{ 313 Type: arrow.PrimitiveTypes.Int8, 314 Len: 4, 315 }}, 316 } 317 var offsets [8]byte 318 endian.Native.PutUint32(offsets[4:], 4) 319 ret.Buffers[1].SetBuffer(memory.NewBufferBytes(offsets[:])) 320 buf := memory.NewResizableBuffer(mem) 321 buf.Resize(4) 322 buf.Bytes()[0] = 1 323 buf.Bytes()[1] = 2 324 buf.Bytes()[2] = 3 325 buf.Bytes()[3] = 4 326 327 ret.Children[0].Buffers[1].WrapBuffer(buf) 328 return ret 329 }, func(mem memory.Allocator) arrow.ArrayData { 330 buf := memory.NewResizableBuffer(mem) 331 buf.Resize(4) 332 buf.Bytes()[0] = 1 333 buf.Bytes()[1] = 2 334 buf.Bytes()[2] = 3 335 buf.Bytes()[3] = 4 336 defer buf.Release() 337 child := array.NewData(arrow.PrimitiveTypes.Int8, 4, []*memory.Buffer{nil, buf}, nil, 0, 0) 338 defer child.Release() 339 340 var offsets [8]byte 341 endian.Native.PutUint32(offsets[4:], 4) 342 343 return array.NewData(arrow.ListOf(arrow.PrimitiveTypes.Int8), 1, 344 []*memory.Buffer{nil, memory.NewBufferBytes(offsets[:])}, 345 []arrow.ArrayData{child}, 0, 0) 346 }}, 347 {"dict-extension-type", func(mem memory.Allocator) fields { 348 // dict-extension-type is dict(Index: int8, Value: string) 349 // so there should be an int8 in the arrayspan and 350 // a child of a string arrayspan in the first index of 351 // Children 352 ret := fields{ 353 Type: types.NewDictExtensionType(), 354 Len: 1, 355 Children: []exec.ArraySpan{{ 356 Type: arrow.BinaryTypes.String, 357 Len: 2, 358 }}, 359 } 360 361 indices := memory.NewResizableBuffer(mem) 362 indices.Resize(1) 363 indices.Bytes()[0] = 1 364 ret.Buffers[1].WrapBuffer(indices) 365 366 offsets := memory.NewResizableBuffer(mem) 367 offsets.Resize(3 * arrow.Int32SizeBytes) 368 copy(offsets.Bytes(), arrow.Int32Traits.CastToBytes([]int32{0, 5, 10})) 369 370 values := memory.NewResizableBuffer(mem) 371 values.Resize(len("HelloWorld")) 372 copy(values.Bytes(), []byte("HelloWorld")) 373 374 nulls := memory.NewResizableBuffer(mem) 375 nulls.Resize(1) 376 nulls.Bytes()[0] = 3 377 ret.Children[0].Buffers[0].WrapBuffer(nulls) 378 ret.Children[0].Buffers[1].WrapBuffer(offsets) 379 ret.Children[0].Buffers[2].WrapBuffer(values) 380 381 return ret 382 }, func(mem memory.Allocator) arrow.ArrayData { 383 dict, _, _ := array.FromJSON(mem, arrow.BinaryTypes.String, strings.NewReader(`["Hello", "World"]`)) 384 defer dict.Release() 385 index, _, _ := array.FromJSON(mem, arrow.PrimitiveTypes.Int8, strings.NewReader(`[1]`)) 386 defer index.Release() 387 388 out := array.NewData(types.NewDictExtensionType(), 1, []*memory.Buffer{nil, index.Data().Buffers()[1]}, nil, 0, 0) 389 out.SetDictionary(dict.Data()) 390 return out 391 }}, 392 } 393 for _, tt := range tests { 394 t.Run(tt.name, func(t *testing.T) { 395 mem := memory.NewCheckedAllocator(memory.DefaultAllocator) 396 defer mem.AssertSize(t, 0) 397 398 t.Run("MakeData", func(t *testing.T) { 399 f := tt.fields(mem) 400 a := &exec.ArraySpan{ 401 Type: f.Type, 402 Len: f.Len, 403 Nulls: f.Nulls, 404 Offset: f.Offset, 405 Buffers: f.Buffers, 406 Scratch: f.Scratch, 407 Children: f.Children, 408 } 409 got := a.MakeData() 410 want := tt.want(mem) 411 if !reflect.DeepEqual(got, want) { 412 t.Errorf("ArraySpan.MakeData() = %v, want %v", got, want) 413 } 414 want.Release() 415 got.Release() 416 }) 417 418 t.Run("MakeArray", func(t *testing.T) { 419 f := tt.fields(mem) 420 a := &exec.ArraySpan{ 421 Type: f.Type, 422 Len: f.Len, 423 Nulls: f.Nulls, 424 Offset: f.Offset, 425 Buffers: f.Buffers, 426 Scratch: f.Scratch, 427 Children: f.Children, 428 } 429 arr := a.MakeArray() 430 want := tt.want(mem) 431 defer want.Release() 432 exp := array.MakeFromData(want) 433 434 assert.Truef(t, array.Equal(arr, exp), "expected: %s\ngot: %s", exp, arr) 435 436 exp.Release() 437 arr.Release() 438 }) 439 }) 440 } 441 } 442 443 func TestArraySpan_SetSlice(t *testing.T) { 444 type fields struct { 445 Type arrow.DataType 446 Len int64 447 Nulls int64 448 Offset int64 449 Buffers [3]exec.BufferSpan 450 Scratch [2]uint64 451 Children []exec.ArraySpan 452 } 453 type args struct { 454 off int64 455 length int64 456 } 457 tests := []struct { 458 name string 459 fields fields 460 args args 461 wantNulls int64 462 }{ 463 {"null type", fields{Type: arrow.Null}, args{5, 10}, 10}, 464 {"not-null type", fields{Type: arrow.PrimitiveTypes.Int8}, args{5, 10}, 0}, 465 {"not-null type with nulls", fields{Type: arrow.PrimitiveTypes.Int8, Nulls: -1}, args{5, 10}, array.UnknownNullCount}, 466 } 467 for _, tt := range tests { 468 t.Run(tt.name, func(t *testing.T) { 469 a := &exec.ArraySpan{ 470 Type: tt.fields.Type, 471 Len: tt.fields.Len, 472 Nulls: tt.fields.Nulls, 473 Offset: tt.fields.Offset, 474 Buffers: tt.fields.Buffers, 475 Scratch: tt.fields.Scratch, 476 Children: tt.fields.Children, 477 } 478 a.SetSlice(tt.args.off, tt.args.length) 479 assert.Equal(t, tt.args.off, a.Offset) 480 assert.Equal(t, tt.args.length, a.Len) 481 assert.Equal(t, tt.wantNulls, a.Nulls) 482 }) 483 } 484 } 485 486 func TestArraySpan_FillFromScalar(t *testing.T) { 487 var ( 488 expDecimalBuf [arrow.Decimal128SizeBytes]byte 489 expScratch [2]uint64 490 ) 491 492 endian.Native.PutUint64(expDecimalBuf[:], 1234) 493 endian.Native.PutUint32(arrow.Uint64Traits.CastToBytes(expScratch[:])[4:], 10) 494 495 dict, _, _ := array.FromJSON(memory.DefaultAllocator, arrow.BinaryTypes.String, strings.NewReader(`["Hello", "World"]`)) 496 defer dict.Release() 497 498 tests := []struct { 499 name string 500 args scalar.Scalar 501 exp exec.ArraySpan 502 }{ 503 {"null-type", 504 scalar.MakeNullScalar(arrow.Null), 505 exec.ArraySpan{Type: arrow.Null, Len: 1, Nulls: 1}}, 506 {"bool valid", 507 scalar.MakeScalar(true), 508 exec.ArraySpan{ 509 Type: arrow.FixedWidthTypes.Boolean, 510 Len: 1, 511 Nulls: 0, 512 Buffers: [3]exec.BufferSpan{{Buf: []byte{0x01}}, {Buf: []byte{0x01}}, {}}, 513 }}, 514 {"bool valid false", 515 scalar.MakeScalar(false), 516 exec.ArraySpan{ 517 Type: arrow.FixedWidthTypes.Boolean, 518 Len: 1, 519 Nulls: 0, 520 Buffers: [3]exec.BufferSpan{{Buf: []byte{0x01}}, {Buf: []byte{0x00}}, {}}, 521 }}, 522 {"primitive null", 523 scalar.MakeNullScalar(arrow.PrimitiveTypes.Int32), 524 exec.ArraySpan{ 525 Type: arrow.PrimitiveTypes.Int32, 526 Len: 1, 527 Nulls: 1, 528 Buffers: [3]exec.BufferSpan{{Buf: []byte{0x00}}, {Buf: []byte{0, 0, 0, 0}}, {}}, 529 }}, 530 {"decimal valid", 531 scalar.NewDecimal128Scalar(decimal128.FromU64(1234), &arrow.Decimal128Type{Precision: 12, Scale: 2}), 532 exec.ArraySpan{ 533 Type: &arrow.Decimal128Type{Precision: 12, Scale: 2}, 534 Len: 1, 535 Nulls: 0, 536 Buffers: [3]exec.BufferSpan{{Buf: []byte{0x01}}, {Buf: expDecimalBuf[:]}, {}}, 537 }}, 538 {"dictionary scalar", 539 scalar.NewDictScalar(scalar.NewInt8Scalar(1), dict), 540 exec.ArraySpan{ 541 Type: &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int8, ValueType: arrow.BinaryTypes.String}, 542 Len: 1, 543 Nulls: 0, 544 Buffers: [3]exec.BufferSpan{{Buf: []byte{0x01}}, 545 {Buf: []byte{1}}, {}, 546 }, 547 Children: []exec.ArraySpan{{ 548 Type: arrow.BinaryTypes.String, 549 Len: 2, 550 Buffers: [3]exec.BufferSpan{ 551 {Buf: dict.NullBitmapBytes(), Owner: dict.Data().Buffers()[0]}, 552 {Buf: dict.Data().Buffers()[1].Bytes(), Owner: dict.Data().Buffers()[1]}, 553 {Buf: dict.Data().Buffers()[2].Bytes(), Owner: dict.Data().Buffers()[2]}, 554 }, 555 }}, 556 }, 557 }, 558 {"binary scalar", 559 scalar.NewBinaryScalar(dict.Data().Buffers()[2], arrow.BinaryTypes.String), 560 exec.ArraySpan{ 561 Type: arrow.BinaryTypes.String, 562 Len: 1, 563 Nulls: 0, 564 Scratch: expScratch, 565 Buffers: [3]exec.BufferSpan{ 566 {Buf: []byte{0x01}}, 567 {Buf: arrow.Uint64Traits.CastToBytes(expScratch[:1])}, 568 {Buf: dict.Data().Buffers()[2].Bytes(), Owner: dict.Data().Buffers()[2]}}, 569 }, 570 }, 571 {"large binary", 572 scalar.NewLargeStringScalarFromBuffer(dict.Data().Buffers()[2]), 573 exec.ArraySpan{ 574 Type: arrow.BinaryTypes.LargeString, 575 Len: 1, 576 Nulls: 0, 577 Scratch: [2]uint64{0, 10}, 578 Buffers: [3]exec.BufferSpan{ 579 {Buf: []byte{0x01}}, 580 {Buf: arrow.Uint64Traits.CastToBytes([]uint64{0, 10})}, 581 {Buf: dict.Data().Buffers()[2].Bytes(), Owner: dict.Data().Buffers()[2]}}, 582 }}, 583 {"fixed size binary", 584 scalar.NewFixedSizeBinaryScalar(dict.Data().Buffers()[2], &arrow.FixedSizeBinaryType{ByteWidth: 10}), 585 exec.ArraySpan{ 586 Type: &arrow.FixedSizeBinaryType{ByteWidth: 10}, 587 Len: 1, 588 Buffers: [3]exec.BufferSpan{ 589 {Buf: []byte{0x01}}, 590 {Buf: dict.Data().Buffers()[2].Bytes(), Owner: dict.Data().Buffers()[2]}, {}, 591 }, 592 }}, 593 {"map scalar null value", 594 scalar.MakeNullScalar(arrow.MapOf(arrow.PrimitiveTypes.Int8, arrow.BinaryTypes.String)), 595 exec.ArraySpan{ 596 Type: arrow.MapOf(arrow.PrimitiveTypes.Int8, arrow.BinaryTypes.String), 597 Len: 1, 598 Nulls: 1, 599 Buffers: [3]exec.BufferSpan{ 600 {Buf: []byte{0}}, 601 {Buf: []byte{0, 0, 0, 0, 0, 0, 0, 0}}, 602 {}, 603 }, 604 Children: []exec.ArraySpan{{ 605 Type: arrow.StructOf(arrow.Field{Name: "key", Type: arrow.PrimitiveTypes.Int8}, 606 arrow.Field{Name: "value", Type: arrow.BinaryTypes.String, Nullable: true}), 607 Len: 0, 608 Nulls: 0, 609 Buffers: [3]exec.BufferSpan{ 610 {Buf: []byte{}}, {}, {}, 611 }, 612 Children: []exec.ArraySpan{ 613 { 614 Type: arrow.PrimitiveTypes.Int8, 615 Buffers: [3]exec.BufferSpan{ 616 {Buf: []byte{}}, {Buf: []byte{}}, {}, 617 }, 618 }, 619 { 620 Type: arrow.BinaryTypes.String, 621 Buffers: [3]exec.BufferSpan{ 622 {Buf: []byte{}}, {Buf: []byte{}}, {Buf: []byte{}}, 623 }, 624 }, 625 }, 626 }}, 627 }}, 628 {"list scalar", 629 scalar.NewListScalarData(dict.Data()), 630 exec.ArraySpan{ 631 Type: arrow.ListOf(arrow.BinaryTypes.String), 632 Len: 1, 633 Scratch: [2]uint64{ 634 *(*uint64)(unsafe.Pointer(&[]int32{0, 2}[0])), 635 0, 636 }, 637 Buffers: [3]exec.BufferSpan{ 638 {Buf: []byte{0x1}}, 639 {Buf: arrow.Int32Traits.CastToBytes([]int32{0, 2})}, 640 }, 641 Children: []exec.ArraySpan{{ 642 Type: arrow.BinaryTypes.String, 643 Len: 2, 644 Buffers: [3]exec.BufferSpan{ 645 {Buf: dict.NullBitmapBytes(), Owner: dict.Data().Buffers()[0]}, 646 {Buf: dict.Data().Buffers()[1].Bytes(), Owner: dict.Data().Buffers()[1]}, 647 {Buf: dict.Data().Buffers()[2].Bytes(), Owner: dict.Data().Buffers()[2]}, 648 }, 649 }}, 650 }, 651 }, 652 {"large list scalar", 653 scalar.NewLargeListScalarData(dict.Data()), 654 exec.ArraySpan{ 655 Type: arrow.LargeListOf(arrow.BinaryTypes.String), 656 Len: 1, 657 Scratch: [2]uint64{0, 2}, 658 Buffers: [3]exec.BufferSpan{ 659 {Buf: []byte{0x1}}, 660 {Buf: arrow.Int64Traits.CastToBytes([]int64{0, 2})}, 661 }, 662 Children: []exec.ArraySpan{{ 663 Type: arrow.BinaryTypes.String, 664 Len: 2, 665 Buffers: [3]exec.BufferSpan{ 666 {Buf: dict.NullBitmapBytes(), Owner: dict.Data().Buffers()[0]}, 667 {Buf: dict.Data().Buffers()[1].Bytes(), Owner: dict.Data().Buffers()[1]}, 668 {Buf: dict.Data().Buffers()[2].Bytes(), Owner: dict.Data().Buffers()[2]}, 669 }, 670 }}, 671 }, 672 }, 673 {"fixed size list", 674 scalar.NewFixedSizeListScalar(dict), 675 exec.ArraySpan{ 676 Type: arrow.FixedSizeListOf(2, arrow.BinaryTypes.String), 677 Len: 1, 678 Buffers: [3]exec.BufferSpan{ 679 {Buf: []byte{0x1}}, 680 {}, {}, 681 }, 682 Children: []exec.ArraySpan{{ 683 Type: arrow.BinaryTypes.String, 684 Len: 2, 685 Buffers: [3]exec.BufferSpan{ 686 {Buf: dict.NullBitmapBytes(), Owner: dict.Data().Buffers()[0]}, 687 {Buf: dict.Data().Buffers()[1].Bytes(), Owner: dict.Data().Buffers()[1]}, 688 {Buf: dict.Data().Buffers()[2].Bytes(), Owner: dict.Data().Buffers()[2]}, 689 }, 690 }}, 691 }, 692 }, 693 {"struct scalar", 694 func() scalar.Scalar { 695 s, _ := scalar.NewStructScalarWithNames([]scalar.Scalar{ 696 scalar.MakeScalar(int32(5)), scalar.MakeScalar(uint8(10)), 697 }, []string{"int32", "uint8"}) 698 return s 699 }(), 700 exec.ArraySpan{ 701 Type: arrow.StructOf( 702 arrow.Field{Name: "int32", Type: arrow.PrimitiveTypes.Int32, Nullable: true}, 703 arrow.Field{Name: "uint8", Type: arrow.PrimitiveTypes.Uint8, Nullable: true}), 704 Buffers: [3]exec.BufferSpan{ 705 {Buf: []byte{0x1}}, {}, {}, 706 }, 707 Len: 1, 708 Children: []exec.ArraySpan{ 709 { 710 Type: arrow.PrimitiveTypes.Int32, 711 Len: 1, 712 Buffers: [3]exec.BufferSpan{ 713 {Buf: []byte{0x1}}, 714 {Buf: arrow.Int32Traits.CastToBytes([]int32{5})}, 715 {}, 716 }, 717 }, 718 { 719 Type: arrow.PrimitiveTypes.Uint8, 720 Len: 1, 721 Buffers: [3]exec.BufferSpan{ 722 {Buf: []byte{0x1}}, 723 {Buf: []byte{10}}, 724 {}, 725 }, 726 }, 727 }, 728 }, 729 }, 730 {"dense union scalar", 731 func() scalar.Scalar { 732 dt := arrow.UnionOf(arrow.DenseMode, []arrow.Field{ 733 {Name: "string", Type: arrow.BinaryTypes.String, Nullable: true}, 734 {Name: "number", Type: arrow.PrimitiveTypes.Uint64, Nullable: true}, 735 {Name: "other_number", Type: arrow.PrimitiveTypes.Uint64, Nullable: true}, 736 }, []arrow.UnionTypeCode{3, 42, 43}) 737 return scalar.NewDenseUnionScalar(scalar.MakeScalar(uint64(25)), 42, dt.(*arrow.DenseUnionType)) 738 }(), 739 exec.ArraySpan{ 740 Type: arrow.UnionOf(arrow.DenseMode, []arrow.Field{ 741 {Name: "string", Type: arrow.BinaryTypes.String, Nullable: true}, 742 {Name: "number", Type: arrow.PrimitiveTypes.Uint64, Nullable: true}, 743 {Name: "other_number", Type: arrow.PrimitiveTypes.Uint64, Nullable: true}, 744 }, []arrow.UnionTypeCode{3, 42, 43}), 745 Len: 1, 746 Scratch: [2]uint64{42, 1}, 747 Buffers: [3]exec.BufferSpan{{}, 748 {Buf: []byte{42}}, {Buf: arrow.Int32Traits.CastToBytes([]int32{0, 1})}, 749 }, 750 Children: []exec.ArraySpan{ 751 { 752 Type: arrow.BinaryTypes.String, 753 Buffers: [3]exec.BufferSpan{ 754 {Buf: []byte{}}, {Buf: []byte{}}, {Buf: []byte{}}, 755 }, 756 }, 757 { 758 Type: arrow.PrimitiveTypes.Uint64, 759 Len: 1, 760 Buffers: [3]exec.BufferSpan{ 761 {Buf: []byte{0x1}}, 762 {Buf: arrow.Uint64Traits.CastToBytes([]uint64{25})}, 763 {}, 764 }, 765 }, 766 { 767 Type: arrow.PrimitiveTypes.Uint64, 768 Buffers: [3]exec.BufferSpan{ 769 {Buf: []byte{}}, {Buf: []byte{}}, {}, 770 }, 771 }, 772 }, 773 }, 774 }, 775 {"sparse union", 776 func() scalar.Scalar { 777 dt := arrow.UnionOf(arrow.SparseMode, []arrow.Field{ 778 {Name: "string", Type: arrow.BinaryTypes.String, Nullable: true}, 779 {Name: "number", Type: arrow.PrimitiveTypes.Uint64, Nullable: true}, 780 {Name: "other_number", Type: arrow.PrimitiveTypes.Uint64, Nullable: true}, 781 }, []arrow.UnionTypeCode{3, 42, 43}) 782 return scalar.NewSparseUnionScalarFromValue(scalar.MakeScalar(uint64(25)), 1, dt.(*arrow.SparseUnionType)) 783 }(), 784 exec.ArraySpan{ 785 Type: arrow.UnionOf(arrow.SparseMode, []arrow.Field{ 786 {Name: "string", Type: arrow.BinaryTypes.String, Nullable: true}, 787 {Name: "number", Type: arrow.PrimitiveTypes.Uint64, Nullable: true}, 788 {Name: "other_number", Type: arrow.PrimitiveTypes.Uint64, Nullable: true}, 789 }, []arrow.UnionTypeCode{3, 42, 43}), 790 Len: 1, 791 Scratch: [2]uint64{42, 0}, 792 Buffers: [3]exec.BufferSpan{{}, 793 {Buf: []byte{42}}, {}, 794 }, 795 Children: []exec.ArraySpan{ 796 { 797 Type: arrow.BinaryTypes.String, 798 Len: 1, 799 Nulls: 1, 800 Buffers: [3]exec.BufferSpan{ 801 {Buf: []byte{0x0}}, 802 {Buf: []byte{0, 0, 0, 0, 0, 0, 0, 0}}, 803 {}, 804 }, 805 }, 806 { 807 Type: arrow.PrimitiveTypes.Uint64, 808 Len: 1, 809 Buffers: [3]exec.BufferSpan{ 810 {Buf: []byte{0x1}}, 811 {Buf: arrow.Uint64Traits.CastToBytes([]uint64{25})}, 812 {}, 813 }, 814 }, 815 { 816 Type: arrow.PrimitiveTypes.Uint64, 817 Len: 1, 818 Nulls: 1, 819 Buffers: [3]exec.BufferSpan{ 820 {Buf: []byte{0x0}}, {Buf: []byte{0, 0, 0, 0, 0, 0, 0, 0}}, {}, 821 }, 822 }, 823 }, 824 }, 825 }, 826 } 827 for _, tt := range tests { 828 t.Run(tt.name, func(t *testing.T) { 829 a := &exec.ArraySpan{ 830 Nulls: array.UnknownNullCount, 831 Buffers: [3]exec.BufferSpan{{SelfAlloc: true, Owner: &memory.Buffer{}}, {SelfAlloc: true, Owner: &memory.Buffer{}}, {}}, 832 } 833 a.FillFromScalar(tt.args) 834 assert.Equal(t, tt.exp, *a) 835 }) 836 } 837 }