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	"sync/atomic"
21
22	"github.com/apache/arrow/go/v6/arrow"
23	"github.com/apache/arrow/go/v6/arrow/bitutil"
24	"github.com/apache/arrow/go/v6/arrow/float16"
25	"github.com/apache/arrow/go/v6/arrow/internal/debug"
26	"github.com/apache/arrow/go/v6/arrow/memory"
27)
28
29type Float16Builder struct {
30	builder
31
32	data    *memory.Buffer
33	rawData []float16.Num
34}
35
36func NewFloat16Builder(mem memory.Allocator) *Float16Builder {
37	return &Float16Builder{builder: builder{refCount: 1, mem: mem}}
38}
39
40// Release decreases the reference count by 1.
41// When the reference count goes to zero, the memory is freed.
42func (b *Float16Builder) Release() {
43	debug.Assert(atomic.LoadInt64(&b.refCount) > 0, "too many releases")
44
45	if atomic.AddInt64(&b.refCount, -1) == 0 {
46		if b.nullBitmap != nil {
47			b.nullBitmap.Release()
48			b.nullBitmap = nil
49		}
50		if b.data != nil {
51			b.data.Release()
52			b.data = nil
53			b.rawData = nil
54		}
55	}
56}
57
58func (b *Float16Builder) Append(v float16.Num) {
59	b.Reserve(1)
60	b.UnsafeAppend(v)
61}
62
63func (b *Float16Builder) UnsafeAppend(v float16.Num) {
64	bitutil.SetBit(b.nullBitmap.Bytes(), b.length)
65	b.rawData[b.length] = v
66	b.length++
67}
68
69func (b *Float16Builder) AppendNull() {
70	b.Reserve(1)
71	b.UnsafeAppendBoolToBitmap(false)
72}
73
74func (b *Float16Builder) UnsafeAppendBoolToBitmap(isValid bool) {
75	if isValid {
76		bitutil.SetBit(b.nullBitmap.Bytes(), b.length)
77	} else {
78		b.nulls++
79	}
80	b.length++
81}
82
83// AppendValues will append the values in the v slice. The valid slice determines which values
84// in v are valid (not null). The valid slice must either be empty or be equal in length to v. If empty,
85// all values in v are appended and considered valid.
86func (b *Float16Builder) AppendValues(v []float16.Num, valid []bool) {
87	if len(v) != len(valid) && len(valid) != 0 {
88		panic("len(v) != len(valid) && len(valid) != 0")
89	}
90
91	if len(v) == 0 {
92		return
93	}
94
95	b.Reserve(len(v))
96	if len(v) > 0 {
97		arrow.Float16Traits.Copy(b.rawData[b.length:], v)
98	}
99	b.builder.unsafeAppendBoolsToBitmap(valid, len(v))
100}
101
102func (b *Float16Builder) init(capacity int) {
103	b.builder.init(capacity)
104
105	b.data = memory.NewResizableBuffer(b.mem)
106	bytesN := arrow.Uint16Traits.BytesRequired(capacity)
107	b.data.Resize(bytesN)
108	b.rawData = arrow.Float16Traits.CastFromBytes(b.data.Bytes())
109}
110
111// Reserve ensures there is enough space for appending n elements
112// by checking the capacity and calling Resize if necessary.
113func (b *Float16Builder) Reserve(n int) {
114	b.builder.reserve(n, b.Resize)
115}
116
117// Resize adjusts the space allocated by b to n elements. If n is greater than b.Cap(),
118// additional memory will be allocated. If n is smaller, the allocated memory may reduced.
119func (b *Float16Builder) Resize(n int) {
120	nBuilder := n
121	if n < minBuilderCapacity {
122		n = minBuilderCapacity
123	}
124
125	if b.capacity == 0 {
126		b.init(n)
127	} else {
128		b.builder.resize(nBuilder, b.init)
129		b.data.Resize(arrow.Float16Traits.BytesRequired(n))
130		b.rawData = arrow.Float16Traits.CastFromBytes(b.data.Bytes())
131	}
132}
133
134// NewArray creates a Float16 array from the memory buffers used by the builder and resets the Float16Builder
135// so it can be used to build a new array.
136func (b *Float16Builder) NewArray() Interface {
137	return b.NewFloat16Array()
138}
139
140// NewFloat16Array creates a Float16 array from the memory buffers used by the builder and resets the Float16Builder
141// so it can be used to build a new array.
142func (b *Float16Builder) NewFloat16Array() (a *Float16) {
143	data := b.newData()
144	a = NewFloat16Data(data)
145	data.Release()
146	return
147}
148
149func (b *Float16Builder) newData() (data *Data) {
150	bytesRequired := arrow.Float16Traits.BytesRequired(b.length)
151	if bytesRequired > 0 && bytesRequired < b.data.Len() {
152		// trim buffers
153		b.data.Resize(bytesRequired)
154	}
155	data = NewData(arrow.FixedWidthTypes.Float16, b.length, []*memory.Buffer{b.nullBitmap, b.data}, nil, b.nulls, 0)
156	b.reset()
157
158	if b.data != nil {
159		b.data.Release()
160		b.data = nil
161		b.rawData = nil
162	}
163
164	return
165}
166