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
18
19import (
20	"bytes"
21	"fmt"
22	"strings"
23	"sync/atomic"
24
25	"github.com/apache/arrow/go/arrow"
26	"github.com/apache/arrow/go/arrow/bitutil"
27	"github.com/apache/arrow/go/arrow/internal/debug"
28	"github.com/apache/arrow/go/arrow/memory"
29)
30
31// Struct represents an ordered sequence of relative types.
32type Struct struct {
33	array
34	fields []Interface
35}
36
37// NewStructData returns a new Struct array value from data.
38func NewStructData(data *Data) *Struct {
39	a := &Struct{}
40	a.refCount = 1
41	a.setData(data)
42	return a
43}
44
45func (a *Struct) NumField() int         { return len(a.fields) }
46func (a *Struct) Field(i int) Interface { return a.fields[i] }
47
48func (a *Struct) String() string {
49	o := new(strings.Builder)
50	o.WriteString("{")
51
52	structBitmap := a.NullBitmapBytes()
53	for i, v := range a.fields {
54		if i > 0 {
55			o.WriteString(" ")
56		}
57		if !bytes.Equal(structBitmap, v.NullBitmapBytes()) {
58			masked := a.newStructFieldWithParentValidityMask(i)
59			fmt.Fprintf(o, "%v", masked)
60			masked.Release()
61			continue
62		}
63		fmt.Fprintf(o, "%v", v)
64	}
65	o.WriteString("}")
66	return o.String()
67}
68
69// newStructFieldWithParentValidityMask returns the Interface at fieldIndex
70// with a nullBitmapBytes adjusted according on the parent struct nullBitmapBytes.
71// From the docs:
72//   "When reading the struct array the parent validity bitmap takes priority."
73func (a *Struct) newStructFieldWithParentValidityMask(fieldIndex int) Interface {
74	field := a.Field(fieldIndex)
75	nullBitmapBytes := field.NullBitmapBytes()
76	maskedNullBitmapBytes := make([]byte, len(nullBitmapBytes))
77	copy(maskedNullBitmapBytes, nullBitmapBytes)
78	for i := 0; i < field.Len(); i++ {
79		if !a.IsValid(i) {
80			bitutil.ClearBit(maskedNullBitmapBytes, i)
81		}
82	}
83	data := NewSliceData(field.Data(), 0, int64(field.Len()))
84	defer data.Release()
85	bufs := make([]*memory.Buffer, len(data.buffers))
86	copy(bufs, data.buffers)
87	bufs[0].Release()
88	bufs[0] = memory.NewBufferBytes(maskedNullBitmapBytes)
89	data.buffers = bufs
90	maskedField := MakeFromData(data)
91	return maskedField
92}
93
94func (a *Struct) setData(data *Data) {
95	a.array.setData(data)
96	a.fields = make([]Interface, len(data.childData))
97	for i, child := range data.childData {
98		if data.offset != 0 || child.length != data.length {
99			sub := NewSliceData(child, int64(data.offset), int64(data.offset+data.length))
100			a.fields[i] = MakeFromData(sub)
101			sub.Release()
102		} else {
103			a.fields[i] = MakeFromData(child)
104		}
105	}
106}
107
108func arrayEqualStruct(left, right *Struct) bool {
109	for i, lf := range left.fields {
110		rf := right.fields[i]
111		if !ArrayEqual(lf, rf) {
112			return false
113		}
114	}
115	return true
116}
117
118func (a *Struct) Retain() {
119	a.array.Retain()
120	for _, f := range a.fields {
121		f.Retain()
122	}
123}
124
125func (a *Struct) Release() {
126	a.array.Release()
127	for _, f := range a.fields {
128		f.Release()
129	}
130}
131
132type StructBuilder struct {
133	builder
134
135	dtype  arrow.DataType
136	fields []Builder
137}
138
139// NewStructBuilder returns a builder, using the provided memory allocator.
140func NewStructBuilder(mem memory.Allocator, dtype *arrow.StructType) *StructBuilder {
141	b := &StructBuilder{
142		builder: builder{refCount: 1, mem: mem},
143		dtype:   dtype,
144		fields:  make([]Builder, len(dtype.Fields())),
145	}
146	for i, f := range dtype.Fields() {
147		b.fields[i] = newBuilder(b.mem, f.Type)
148	}
149	return b
150}
151
152// Release decreases the reference count by 1.
153// When the reference count goes to zero, the memory is freed.
154func (b *StructBuilder) Release() {
155	debug.Assert(atomic.LoadInt64(&b.refCount) > 0, "too many releases")
156
157	if atomic.AddInt64(&b.refCount, -1) == 0 {
158		if b.nullBitmap != nil {
159			b.nullBitmap.Release()
160			b.nullBitmap = nil
161		}
162	}
163
164	for _, f := range b.fields {
165		f.Release()
166	}
167}
168
169func (b *StructBuilder) Append(v bool) {
170	b.Reserve(1)
171	b.unsafeAppendBoolToBitmap(v)
172	if !v {
173		for _, f := range b.fields {
174			f.AppendNull()
175		}
176	}
177}
178
179func (b *StructBuilder) AppendValues(valids []bool) {
180	b.Reserve(len(valids))
181	b.builder.unsafeAppendBoolsToBitmap(valids, len(valids))
182}
183
184func (b *StructBuilder) AppendNull() { b.Append(false) }
185
186func (b *StructBuilder) unsafeAppend(v bool) {
187	bitutil.SetBit(b.nullBitmap.Bytes(), b.length)
188	b.length++
189}
190
191func (b *StructBuilder) unsafeAppendBoolToBitmap(isValid bool) {
192	if isValid {
193		bitutil.SetBit(b.nullBitmap.Bytes(), b.length)
194	} else {
195		b.nulls++
196	}
197	b.length++
198}
199
200func (b *StructBuilder) init(capacity int) {
201	b.builder.init(capacity)
202}
203
204// Reserve ensures there is enough space for appending n elements
205// by checking the capacity and calling Resize if necessary.
206func (b *StructBuilder) Reserve(n int) {
207	b.builder.reserve(n, b.resizeHelper)
208	for _, f := range b.fields {
209		f.Reserve(n)
210	}
211}
212
213// Resize adjusts the space allocated by b to n elements. If n is greater than b.Cap(),
214// additional memory will be allocated. If n is smaller, the allocated memory may reduced.
215func (b *StructBuilder) Resize(n int) {
216	b.resizeHelper(n)
217	for _, f := range b.fields {
218		f.Resize(n)
219	}
220}
221
222func (b *StructBuilder) resizeHelper(n int) {
223	if n < minBuilderCapacity {
224		n = minBuilderCapacity
225	}
226
227	if b.capacity == 0 {
228		b.init(n)
229	} else {
230		b.builder.resize(n, b.builder.init)
231	}
232}
233
234func (b *StructBuilder) NumField() int              { return len(b.fields) }
235func (b *StructBuilder) FieldBuilder(i int) Builder { return b.fields[i] }
236
237// NewArray creates a Struct array from the memory buffers used by the builder and resets the StructBuilder
238// so it can be used to build a new array.
239func (b *StructBuilder) NewArray() Interface {
240	return b.NewStructArray()
241}
242
243// NewStructArray creates a Struct array from the memory buffers used by the builder and resets the StructBuilder
244// so it can be used to build a new array.
245func (b *StructBuilder) NewStructArray() (a *Struct) {
246	data := b.newData()
247	a = NewStructData(data)
248	data.Release()
249	return
250}
251
252func (b *StructBuilder) newData() (data *Data) {
253	fields := make([]*Data, len(b.fields))
254	for i, f := range b.fields {
255		arr := f.NewArray()
256		defer arr.Release()
257		fields[i] = arr.Data()
258	}
259
260	data = NewData(
261		b.dtype, b.length,
262		[]*memory.Buffer{
263			b.nullBitmap,
264			nil, // FIXME(sbinet)
265		},
266		fields,
267		b.nulls,
268		0,
269	)
270	b.reset()
271
272	return
273}
274
275var (
276	_ Interface = (*Struct)(nil)
277	_ Builder   = (*StructBuilder)(nil)
278)
279