blob: ca4b5ef904d5304f12817946a79a26842ceb2797 [file] [log] [blame]
// 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.
package array_test
import (
"reflect"
"testing"
"github.com/apache/arrow/go/v9/arrow"
"github.com/apache/arrow/go/v9/arrow/array"
"github.com/apache/arrow/go/v9/arrow/memory"
)
func TestStructArray(t *testing.T) {
pool := memory.NewCheckedAllocator(memory.NewGoAllocator())
defer pool.AssertSize(t, 0)
var (
f1s = []byte{'j', 'o', 'e', 'b', 'o', 'b', 'm', 'a', 'r', 'k'}
f2s = []int32{1, 2, 3, 4}
f1Lengths = []int{3, 0, 3, 4}
f1Offsets = []int32{0, 3, 3, 6, 10}
f1Valids = []bool{true, false, true, true}
isValid = []bool{true, true, true, true}
fields = []arrow.Field{
{Name: "f1", Type: arrow.ListOf(arrow.PrimitiveTypes.Uint8)},
{Name: "f2", Type: arrow.PrimitiveTypes.Int32},
}
dtype = arrow.StructOf(fields...)
)
sb := array.NewStructBuilder(pool, dtype)
defer sb.Release()
for i := 0; i < 10; i++ {
f1b := sb.FieldBuilder(0).(*array.ListBuilder)
f1vb := f1b.ValueBuilder().(*array.Uint8Builder)
f2b := sb.FieldBuilder(1).(*array.Int32Builder)
if got, want := sb.NumField(), 2; got != want {
t.Fatalf("got=%d, want=%d", got, want)
}
sb.Resize(len(f1Lengths))
f1vb.Resize(len(f1s))
f2b.Resize(len(f2s))
pos := 0
for i, length := range f1Lengths {
f1b.Append(f1Valids[i])
for j := 0; j < length; j++ {
f1vb.Append(f1s[pos])
pos++
}
f2b.Append(f2s[i])
}
for _, valid := range isValid {
sb.Append(valid)
}
arr := sb.NewArray().(*array.Struct)
defer arr.Release()
arr.Retain()
arr.Release()
if got, want := arr.DataType().ID(), arrow.STRUCT; got != want {
t.Fatalf("got=%v, want=%v", got, want)
}
if got, want := arr.Len(), len(isValid); got != want {
t.Fatalf("got=%d, want=%d", got, want)
}
for i, valid := range isValid {
if got, want := arr.IsValid(i), valid; got != want {
t.Fatalf("got[%d]=%v, want[%d]=%v", i, got, i, want)
}
}
{
f1arr := arr.Field(0).(*array.List)
if got, want := f1arr.Len(), len(f1Lengths); got != want {
t.Fatalf("got=%d, want=%d", got, want)
}
for i := range f1Lengths {
if got, want := f1arr.IsValid(i), f1Valids[i]; got != want {
t.Fatalf("got[%d]=%v, want[%d]=%v", i, got, i, want)
}
if got, want := f1arr.IsNull(i), f1Lengths[i] == 0; got != want {
t.Fatalf("got[%d]=%v, want[%d]=%v", i, got, i, want)
}
}
if got, want := f1arr.Offsets(), f1Offsets; !reflect.DeepEqual(got, want) {
t.Fatalf("got=%v, want=%v", got, want)
}
varr := f1arr.ListValues().(*array.Uint8)
if got, want := varr.Uint8Values(), f1s; !reflect.DeepEqual(got, want) {
t.Fatalf("got=%v, want=%v", got, want)
}
}
{
f2arr := arr.Field(1).(*array.Int32)
if got, want := f2arr.Len(), len(f2s); got != want {
t.Fatalf("got=%d, want=%d", got, want)
}
if got, want := f2arr.Int32Values(), f2s; !reflect.DeepEqual(got, want) {
t.Fatalf("got=%d, want=%d", got, want)
}
}
}
}
func TestStructArrayEmpty(t *testing.T) {
pool := memory.NewCheckedAllocator(memory.NewGoAllocator())
defer pool.AssertSize(t, 0)
sb := array.NewStructBuilder(pool, arrow.StructOf())
defer sb.Release()
if got, want := sb.NumField(), 0; got != want {
t.Fatalf("got=%d, want=%d", got, want)
}
arr := sb.NewArray().(*array.Struct)
if got, want := arr.Len(), 0; got != want {
t.Fatalf("got=%d, want=%d", got, want)
}
if got, want := arr.NumField(), 0; got != want {
t.Fatalf("got=%d, want=%d", got, want)
}
}
func TestStructArrayBulkAppend(t *testing.T) {
pool := memory.NewCheckedAllocator(memory.NewGoAllocator())
defer pool.AssertSize(t, 0)
var (
f1s = []byte{'j', 'o', 'e', 'b', 'o', 'b', 'm', 'a', 'r', 'k'}
f2s = []int32{1, 2, 3, 4}
f1Lengths = []int{3, 0, 3, 4}
f1Offsets = []int32{0, 3, 3, 6, 10}
f1Valids = []bool{true, false, true, true}
isValid = []bool{true, true, true, true}
fields = []arrow.Field{
{Name: "f1", Type: arrow.ListOf(arrow.PrimitiveTypes.Uint8)},
{Name: "f2", Type: arrow.PrimitiveTypes.Int32},
}
dtype = arrow.StructOf(fields...)
)
sb := array.NewStructBuilder(pool, dtype)
defer sb.Release()
for i := 0; i < 10; i++ {
f1b := sb.FieldBuilder(0).(*array.ListBuilder)
f1vb := f1b.ValueBuilder().(*array.Uint8Builder)
f2b := sb.FieldBuilder(1).(*array.Int32Builder)
if got, want := sb.NumField(), 2; got != want {
t.Fatalf("got=%d, want=%d", got, want)
}
sb.Resize(len(f1Lengths))
f1vb.Resize(len(f1s))
f2b.Resize(len(f2s))
sb.AppendValues(isValid)
f1b.AppendValues(f1Offsets, f1Valids)
f1vb.AppendValues(f1s, nil)
f2b.AppendValues(f2s, nil)
arr := sb.NewArray().(*array.Struct)
defer arr.Release()
if got, want := arr.DataType().ID(), arrow.STRUCT; got != want {
t.Fatalf("got=%v, want=%v", got, want)
}
if got, want := arr.Len(), len(isValid); got != want {
t.Fatalf("got=%d, want=%d", got, want)
}
for i, valid := range isValid {
if got, want := arr.IsValid(i), valid; got != want {
t.Fatalf("got[%d]=%v, want[%d]=%v", i, got, i, want)
}
}
{
f1arr := arr.Field(0).(*array.List)
if got, want := f1arr.Len(), len(f1Lengths); got != want {
t.Fatalf("got=%d, want=%d", got, want)
}
for i := range f1Lengths {
if got, want := f1arr.IsValid(i), f1Valids[i]; got != want {
t.Fatalf("got[%d]=%v, want[%d]=%v", i, got, i, want)
}
if got, want := f1arr.IsNull(i), f1Lengths[i] == 0; got != want {
t.Fatalf("got[%d]=%v, want[%d]=%v", i, got, i, want)
}
}
if got, want := f1arr.Offsets(), f1Offsets; !reflect.DeepEqual(got, want) {
t.Fatalf("got=%v, want=%v", got, want)
}
varr := f1arr.ListValues().(*array.Uint8)
if got, want := varr.Uint8Values(), f1s; !reflect.DeepEqual(got, want) {
t.Fatalf("got=%v, want=%v", got, want)
}
}
{
f2arr := arr.Field(1).(*array.Int32)
if got, want := f2arr.Len(), len(f2s); got != want {
t.Fatalf("got=%d, want=%d", got, want)
}
if got, want := f2arr.Int32Values(), f2s; !reflect.DeepEqual(got, want) {
t.Fatalf("got=%d, want=%d", got, want)
}
}
}
}
func TestStructArrayStringer(t *testing.T) {
pool := memory.NewCheckedAllocator(memory.NewGoAllocator())
defer pool.AssertSize(t, 0)
var (
f1s = []float64{1.1, 1.2, 1.3, 1.4}
f2s = []int32{1, 2, 3, 4}
fields = []arrow.Field{
{Name: "f1", Type: arrow.PrimitiveTypes.Float64},
{Name: "f2", Type: arrow.PrimitiveTypes.Int32},
}
dtype = arrow.StructOf(fields...)
)
sb := array.NewStructBuilder(pool, dtype)
defer sb.Release()
f1b := sb.FieldBuilder(0).(*array.Float64Builder)
f2b := sb.FieldBuilder(1).(*array.Int32Builder)
if got, want := sb.NumField(), 2; got != want {
t.Fatalf("got=%d, want=%d", got, want)
}
for i := range f1s {
sb.Append(true)
switch i {
case 1:
f1b.AppendNull()
f2b.Append(f2s[i])
case 2:
f1b.Append(f1s[i])
f2b.AppendNull()
default:
f1b.Append(f1s[i])
f2b.Append(f2s[i])
}
}
arr := sb.NewArray().(*array.Struct)
defer arr.Release()
want := "{[1.1 (null) 1.3 1.4] [1 2 (null) 4]}"
got := arr.String()
if got != want {
t.Fatalf("invalid string representation:\ngot = %q\nwant= %q", got, want)
}
}
func TestStructArraySlice(t *testing.T) {
pool := memory.NewCheckedAllocator(memory.NewGoAllocator())
defer pool.AssertSize(t, 0)
var (
f1s = []float64{1.1, 1.2, 1.3, 1.4}
f2s = []int32{1, 2, 3, 4}
valids = []bool{true, true, true, true}
fields = []arrow.Field{
{Name: "f1", Type: arrow.PrimitiveTypes.Float64},
{Name: "f2", Type: arrow.PrimitiveTypes.Int32},
}
dtype = arrow.StructOf(fields...)
)
sb := array.NewStructBuilder(pool, dtype)
defer sb.Release()
f1b := sb.FieldBuilder(0).(*array.Float64Builder)
f2b := sb.FieldBuilder(1).(*array.Int32Builder)
if got, want := sb.NumField(), 2; got != want {
t.Fatalf("got=%d, want=%d", got, want)
}
for i := range f1s {
sb.Append(valids[i])
switch i {
case 1:
f1b.AppendNull()
f2b.Append(f2s[i])
case 2:
f1b.Append(f1s[i])
f2b.AppendNull()
default:
f1b.Append(f1s[i])
f2b.Append(f2s[i])
}
}
arr := sb.NewArray().(*array.Struct)
defer arr.Release()
// Slice
arrSlice := array.NewSlice(arr, 2, 4).(*array.Struct)
defer arrSlice.Release()
want := "{[1.3 1.4] [(null) 4]}"
got := arrSlice.String()
if got != want {
t.Fatalf("invalid string representation:\ngot = %q\nwant= %q", got, want)
}
}
func TestStructArrayNullBitmap(t *testing.T) {
pool := memory.NewCheckedAllocator(memory.NewGoAllocator())
defer pool.AssertSize(t, 0)
var (
f1s = []float64{1.1, 1.2, 1.3, 1.4}
f2s = []int32{1, 2, 3, 4}
valids = []bool{true, true, true, false}
fields = []arrow.Field{
{Name: "f1", Type: arrow.PrimitiveTypes.Float64},
{Name: "f2", Type: arrow.PrimitiveTypes.Int32},
}
dtype = arrow.StructOf(fields...)
)
sb := array.NewStructBuilder(pool, dtype)
defer sb.Release()
f1b := sb.FieldBuilder(0).(*array.Float64Builder)
f2b := sb.FieldBuilder(1).(*array.Int32Builder)
if got, want := sb.NumField(), 2; got != want {
t.Fatalf("got=%d, want=%d", got, want)
}
sb.AppendValues(valids)
for i := range f1s {
f1b.Append(f1s[i])
switch i {
case 1:
f2b.AppendNull()
default:
f2b.Append(f2s[i])
}
}
arr := sb.NewArray().(*array.Struct)
defer arr.Release()
want := "{[1.1 1.2 1.3 (null)] [1 (null) 3 (null)]}"
got := arr.String()
if got != want {
t.Fatalf("invalid string representation:\ngot = %q\nwant= %q", got, want)
}
}
func TestStructArrayUnmarshalJSONMissingFields(t *testing.T) {
pool := memory.NewGoAllocator()
var (
fields = []arrow.Field{
{Name: "f1", Type: arrow.PrimitiveTypes.Float64, Nullable: true},
{Name: "f2", Type: arrow.PrimitiveTypes.Int32},
{
Name: "f3", Type: arrow.StructOf(
[]arrow.Field{
{Name: "f3_1", Type: arrow.BinaryTypes.String, Nullable: true},
{Name: "f3_2", Type: arrow.BinaryTypes.String, Nullable: true},
{Name: "f3_3", Type: arrow.BinaryTypes.String, Nullable: false},
}...,
),
},
}
dtype = arrow.StructOf(fields...)
)
tests := []struct {
name string
jsonInput string
want string
panic bool
}{
{
name: "missing required field",
jsonInput: `[{"f2": 3, "f3": {"f3_1": "test"}}]`,
panic: true,
want: "",
},
{
name: "missing optional fields",
jsonInput: `[{"f2": 3, "f3": {"f3_3": "test"}}]`,
panic: false,
want: `{[(null)] [3] {[(null)] [(null)] ["test"]}}`,
},
}
for _, tc := range tests {
t.Run(
tc.name, func(t *testing.T) {
var val bool
sb := array.NewStructBuilder(pool, dtype)
defer sb.Release()
if tc.panic {
defer func() {
e := recover()
if e == nil {
t.Fatalf("this should have panicked, but did not; slice value %v", val)
}
if got, want := e.(string), "arrow/array: index out of range"; got != want {
t.Fatalf("invalid error. got=%q, want=%q", got, want)
}
}()
} else {
defer func() {
if e := recover(); e != nil {
t.Fatalf("unexpected panic: %v", e)
}
}()
}
err := sb.UnmarshalJSON([]byte(tc.jsonInput))
if err != nil {
t.Fatal(err)
}
arr := sb.NewArray().(*array.Struct)
defer arr.Release()
got := arr.String()
if got != tc.want {
t.Fatalf("invalid string representation:\ngot = %q\nwant= %q", got, tc.want)
}
},
)
}
}