1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
|
// 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.
//go:build go1.18
package exec_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/compute/exec"
"github.com/apache/arrow-go/v18/arrow/memory"
"github.com/stretchr/testify/assert"
)
func TestRechunkConsistentArraysTrivial(t *testing.T) {
var groups [][]arrow.Array
rechunked := exec.RechunkArraysConsistently(groups)
assert.Zero(t, rechunked)
mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
defer mem.AssertSize(t, 0)
a1 := exec.ArrayFromSlice(mem, []int16{})
defer a1.Release()
a2 := exec.ArrayFromSlice(mem, []int16{})
defer a2.Release()
b1 := exec.ArrayFromSlice(mem, []int32{})
defer b1.Release()
groups = [][]arrow.Array{{a1, a2}, {}, {b1}}
rechunked = exec.RechunkArraysConsistently(groups)
assert.Len(t, rechunked, 3)
for _, arrvec := range rechunked {
for _, arr := range arrvec {
assert.Zero(t, arr.Len())
}
}
}
func assertEqual[T arrow.NumericType](t *testing.T, mem memory.Allocator, arr arrow.Array, data []T) {
exp := exec.ArrayFromSlice(mem, data)
defer exp.Release()
assert.Truef(t, array.Equal(exp, arr), "expected: %s\ngot: %s", exp, arr)
}
func TestRechunkArraysConsistentlyPlain(t *testing.T) {
mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
defer mem.AssertSize(t, 0)
a1 := exec.ArrayFromSlice(mem, []int16{1, 2, 3})
defer a1.Release()
a2 := exec.ArrayFromSlice(mem, []int16{4, 5})
defer a2.Release()
a3 := exec.ArrayFromSlice(mem, []int16{6, 7, 8, 9})
defer a3.Release()
b1 := exec.ArrayFromSlice(mem, []int32{41, 42})
defer b1.Release()
b2 := exec.ArrayFromSlice(mem, []int32{43, 44, 45})
defer b2.Release()
b3 := exec.ArrayFromSlice(mem, []int32{46, 47})
defer b3.Release()
b4 := exec.ArrayFromSlice(mem, []int32{48, 49})
defer b4.Release()
groups := [][]arrow.Array{{a1, a2, a3}, {b1, b2, b3, b4}}
rechunked := exec.RechunkArraysConsistently(groups)
assert.Len(t, rechunked, 2)
ra := rechunked[0]
rb := rechunked[1]
assert.Len(t, ra, 5)
assertEqual(t, mem, ra[0], []int16{1, 2})
ra[0].Release()
assertEqual(t, mem, ra[1], []int16{3})
ra[1].Release()
assertEqual(t, mem, ra[2], []int16{4, 5})
ra[2].Release()
assertEqual(t, mem, ra[3], []int16{6, 7})
ra[3].Release()
assertEqual(t, mem, ra[4], []int16{8, 9})
ra[4].Release()
assert.Len(t, rb, 5)
assertEqual(t, mem, rb[0], []int32{41, 42})
rb[0].Release()
assertEqual(t, mem, rb[1], []int32{43})
rb[1].Release()
assertEqual(t, mem, rb[2], []int32{44, 45})
rb[2].Release()
assertEqual(t, mem, rb[3], []int32{46, 47})
rb[3].Release()
assertEqual(t, mem, rb[4], []int32{48, 49})
rb[4].Release()
}
|
