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