]>
Commit | Line | Data |
---|---|---|
1d09f67e TL |
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 | ||
17 | package array | |
18 | ||
19 | import ( | |
20 | "fmt" | |
21 | "strings" | |
22 | "sync/atomic" | |
23 | ||
24 | "github.com/apache/arrow/go/v6/arrow" | |
25 | "github.com/apache/arrow/go/v6/arrow/bitutil" | |
26 | "github.com/apache/arrow/go/v6/arrow/internal/debug" | |
27 | "github.com/apache/arrow/go/v6/arrow/memory" | |
28 | ) | |
29 | ||
30 | // FixedSizeList represents an immutable sequence of N array values. | |
31 | type FixedSizeList struct { | |
32 | array | |
33 | n int32 | |
34 | values Interface | |
35 | } | |
36 | ||
37 | // NewFixedSizeListData returns a new List array value, from data. | |
38 | func NewFixedSizeListData(data *Data) *FixedSizeList { | |
39 | a := &FixedSizeList{} | |
40 | a.refCount = 1 | |
41 | a.setData(data) | |
42 | return a | |
43 | } | |
44 | ||
45 | func (a *FixedSizeList) ListValues() Interface { return a.values } | |
46 | ||
47 | func (a *FixedSizeList) String() string { | |
48 | o := new(strings.Builder) | |
49 | o.WriteString("[") | |
50 | for i := 0; i < a.Len(); i++ { | |
51 | if i > 0 { | |
52 | o.WriteString(" ") | |
53 | } | |
54 | if !a.IsValid(i) { | |
55 | o.WriteString("(null)") | |
56 | continue | |
57 | } | |
58 | sub := a.newListValue(i) | |
59 | fmt.Fprintf(o, "%v", sub) | |
60 | sub.Release() | |
61 | } | |
62 | o.WriteString("]") | |
63 | return o.String() | |
64 | } | |
65 | ||
66 | func (a *FixedSizeList) newListValue(i int) Interface { | |
67 | n := int64(a.n) | |
68 | off := int64(a.array.data.offset) | |
69 | beg := (off + int64(i)) * n | |
70 | end := (off + int64(i+1)) * n | |
71 | sli := NewSlice(a.values, beg, end) | |
72 | return sli | |
73 | } | |
74 | ||
75 | func (a *FixedSizeList) setData(data *Data) { | |
76 | a.array.setData(data) | |
77 | a.n = a.DataType().(*arrow.FixedSizeListType).Len() | |
78 | a.values = MakeFromData(data.childData[0]) | |
79 | } | |
80 | ||
81 | func arrayEqualFixedSizeList(left, right *FixedSizeList) bool { | |
82 | for i := 0; i < left.Len(); i++ { | |
83 | if left.IsNull(i) { | |
84 | continue | |
85 | } | |
86 | o := func() bool { | |
87 | l := left.newListValue(i) | |
88 | defer l.Release() | |
89 | r := right.newListValue(i) | |
90 | defer r.Release() | |
91 | return ArrayEqual(l, r) | |
92 | }() | |
93 | if !o { | |
94 | return false | |
95 | } | |
96 | } | |
97 | return true | |
98 | } | |
99 | ||
100 | // Len returns the number of elements in the array. | |
101 | func (a *FixedSizeList) Len() int { return a.array.Len() } | |
102 | ||
103 | func (a *FixedSizeList) Retain() { | |
104 | a.array.Retain() | |
105 | a.values.Retain() | |
106 | } | |
107 | ||
108 | func (a *FixedSizeList) Release() { | |
109 | a.array.Release() | |
110 | a.values.Release() | |
111 | } | |
112 | ||
113 | type FixedSizeListBuilder struct { | |
114 | builder | |
115 | ||
116 | etype arrow.DataType // data type of the list's elements. | |
117 | n int32 // number of elements in the fixed-size list. | |
118 | values Builder // value builder for the list's elements. | |
119 | } | |
120 | ||
121 | // NewFixedSizeListBuilder returns a builder, using the provided memory allocator. | |
122 | // The created list builder will create a list whose elements will be of type etype. | |
123 | func NewFixedSizeListBuilder(mem memory.Allocator, n int32, etype arrow.DataType) *FixedSizeListBuilder { | |
124 | return &FixedSizeListBuilder{ | |
125 | builder: builder{refCount: 1, mem: mem}, | |
126 | etype: etype, | |
127 | n: n, | |
128 | values: NewBuilder(mem, etype), | |
129 | } | |
130 | } | |
131 | ||
132 | // Release decreases the reference count by 1. | |
133 | // When the reference count goes to zero, the memory is freed. | |
134 | func (b *FixedSizeListBuilder) Release() { | |
135 | debug.Assert(atomic.LoadInt64(&b.refCount) > 0, "too many releases") | |
136 | ||
137 | if atomic.AddInt64(&b.refCount, -1) == 0 { | |
138 | if b.nullBitmap != nil { | |
139 | b.nullBitmap.Release() | |
140 | b.nullBitmap = nil | |
141 | } | |
142 | if b.values != nil { | |
143 | b.values.Release() | |
144 | b.values = nil | |
145 | } | |
146 | } | |
147 | } | |
148 | ||
149 | func (b *FixedSizeListBuilder) Append(v bool) { | |
150 | b.Reserve(1) | |
151 | b.unsafeAppendBoolToBitmap(v) | |
152 | } | |
153 | ||
154 | func (b *FixedSizeListBuilder) AppendNull() { | |
155 | b.Reserve(1) | |
156 | b.unsafeAppendBoolToBitmap(false) | |
157 | } | |
158 | ||
159 | func (b *FixedSizeListBuilder) AppendValues(valid []bool) { | |
160 | b.Reserve(len(valid)) | |
161 | b.builder.unsafeAppendBoolsToBitmap(valid, len(valid)) | |
162 | } | |
163 | ||
164 | func (b *FixedSizeListBuilder) unsafeAppend(v bool) { | |
165 | bitutil.SetBit(b.nullBitmap.Bytes(), b.length) | |
166 | b.length++ | |
167 | } | |
168 | ||
169 | func (b *FixedSizeListBuilder) unsafeAppendBoolToBitmap(isValid bool) { | |
170 | if isValid { | |
171 | bitutil.SetBit(b.nullBitmap.Bytes(), b.length) | |
172 | } else { | |
173 | b.nulls++ | |
174 | } | |
175 | b.length++ | |
176 | } | |
177 | ||
178 | func (b *FixedSizeListBuilder) init(capacity int) { | |
179 | b.builder.init(capacity) | |
180 | } | |
181 | ||
182 | // Reserve ensures there is enough space for appending n elements | |
183 | // by checking the capacity and calling Resize if necessary. | |
184 | func (b *FixedSizeListBuilder) Reserve(n int) { | |
185 | b.builder.reserve(n, b.Resize) | |
186 | } | |
187 | ||
188 | // Resize adjusts the space allocated by b to n elements. If n is greater than b.Cap(), | |
189 | // additional memory will be allocated. If n is smaller, the allocated memory may reduced. | |
190 | func (b *FixedSizeListBuilder) Resize(n int) { | |
191 | if n < minBuilderCapacity { | |
192 | n = minBuilderCapacity | |
193 | } | |
194 | ||
195 | if b.capacity == 0 { | |
196 | b.init(n) | |
197 | } else { | |
198 | b.builder.resize(n, b.builder.init) | |
199 | } | |
200 | } | |
201 | ||
202 | func (b *FixedSizeListBuilder) ValueBuilder() Builder { | |
203 | return b.values | |
204 | } | |
205 | ||
206 | // NewArray creates a List array from the memory buffers used by the builder and resets the FixedSizeListBuilder | |
207 | // so it can be used to build a new array. | |
208 | func (b *FixedSizeListBuilder) NewArray() Interface { | |
209 | return b.NewListArray() | |
210 | } | |
211 | ||
212 | // NewListArray creates a List array from the memory buffers used by the builder and resets the FixedSizeListBuilder | |
213 | // so it can be used to build a new array. | |
214 | func (b *FixedSizeListBuilder) NewListArray() (a *FixedSizeList) { | |
215 | data := b.newData() | |
216 | a = NewFixedSizeListData(data) | |
217 | data.Release() | |
218 | return | |
219 | } | |
220 | ||
221 | func (b *FixedSizeListBuilder) newData() (data *Data) { | |
222 | values := b.values.NewArray() | |
223 | defer values.Release() | |
224 | ||
225 | data = NewData( | |
226 | arrow.FixedSizeListOf(b.n, b.etype), b.length, | |
227 | []*memory.Buffer{b.nullBitmap}, | |
228 | []*Data{values.Data()}, | |
229 | b.nulls, | |
230 | 0, | |
231 | ) | |
232 | b.reset() | |
233 | ||
234 | return | |
235 | } | |
236 | ||
237 | var ( | |
238 | _ Interface = (*FixedSizeList)(nil) | |
239 | _ Builder = (*FixedSizeListBuilder)(nil) | |
240 | ) |