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
112
113
114
|
// 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 exprs
import (
"context"
"strings"
"testing"
"github.com/apache/arrow-go/v18/arrow"
"github.com/apache/arrow-go/v18/arrow/array"
"github.com/apache/arrow-go/v18/arrow/compute"
"github.com/apache/arrow-go/v18/arrow/memory"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
var (
boringArrowSchema = arrow.NewSchema([]arrow.Field{
{Name: "bool", Type: arrow.FixedWidthTypes.Boolean, Nullable: true},
{Name: "i8", Type: arrow.PrimitiveTypes.Int8, Nullable: true},
{Name: "i32", Type: arrow.PrimitiveTypes.Int32, Nullable: true},
{Name: "i32_req", Type: arrow.PrimitiveTypes.Int32},
{Name: "u32", Type: arrow.PrimitiveTypes.Uint32, Nullable: true},
{Name: "i64", Type: arrow.PrimitiveTypes.Int64, Nullable: true},
{Name: "f32", Type: arrow.PrimitiveTypes.Float32, Nullable: true},
{Name: "f32_req", Type: arrow.PrimitiveTypes.Float32},
{Name: "f64", Type: arrow.PrimitiveTypes.Float64, Nullable: true},
{Name: "date32", Type: arrow.FixedWidthTypes.Date32, Nullable: true},
{Name: "str", Type: arrow.BinaryTypes.String, Nullable: true},
{Name: "bin", Type: arrow.BinaryTypes.Binary, Nullable: true},
}, nil)
)
func TestMakeExecBatch(t *testing.T) {
mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
defer mem.AssertSize(t, 0)
const numRows = 3
var (
ctx = compute.WithAllocator(context.Background(), mem)
i32, _, _ = array.FromJSON(mem, arrow.PrimitiveTypes.Int32, strings.NewReader(`[1, 2, 3]`))
f32, _, _ = array.FromJSON(mem, arrow.PrimitiveTypes.Float32, strings.NewReader(`[1.5, 2.25, 3.125]`))
empty, _, _ = array.RecordFromJSON(mem, boringArrowSchema, strings.NewReader(`[]`))
)
defer i32.Release()
defer f32.Release()
getField := func(n string) arrow.Field {
f, _ := boringArrowSchema.FieldsByName(n)
return f[0]
}
tests := []struct {
name string
batch arrow.Record
}{
{"empty", empty},
{"subset", array.NewRecord(arrow.NewSchema([]arrow.Field{getField("i32"), getField("f32")}, nil),
[]arrow.Array{i32, f32}, numRows)},
{"flipped subset", array.NewRecord(arrow.NewSchema([]arrow.Field{getField("f32"), getField("i32")}, nil),
[]arrow.Array{f32, i32}, numRows)},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
defer tt.batch.Release()
batch, err := makeExecBatch(ctx, boringArrowSchema, compute.NewDatumWithoutOwning(tt.batch))
require.NoError(t, err)
require.Equal(t, tt.batch.NumRows(), batch.Len)
defer func() {
for _, v := range batch.Values {
v.Release()
}
}()
for i, field := range boringArrowSchema.Fields() {
typ := batch.Values[i].(compute.ArrayLikeDatum).Type()
assert.Truef(t, arrow.TypeEqual(typ, field.Type),
"expected: %s\ngot: %s", field.Type, typ)
idxes := tt.batch.Schema().FieldIndices(field.Name)
if batch.Values[i].Kind() == compute.KindScalar {
assert.False(t, batch.Values[i].(*compute.ScalarDatum).Value.IsValid(),
"null placeholder should be injected")
assert.Len(t, idxes, 0, "should only happen when column isn't found")
} else {
col := tt.batch.Column(idxes[0])
val := batch.Values[i].(*compute.ArrayDatum).MakeArray()
defer val.Release()
assert.Truef(t, array.Equal(col, val), "expected: %s\ngot: %s", col, val)
}
}
})
}
}
|
