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 // import "github.com/apache/arrow/go/v6/arrow/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/internal/debug"
25)
26
27// A type which satisfies array.Interface represents an immutable sequence of values.
28type Interface interface {
29	// DataType returns the type metadata for this instance.
30	DataType() arrow.DataType
31
32	// NullN returns the number of null values in the array.
33	NullN() int
34
35	// NullBitmapBytes returns a byte slice of the validity bitmap.
36	NullBitmapBytes() []byte
37
38	// IsNull returns true if value at index is null.
39	// NOTE: IsNull will panic if NullBitmapBytes is not empty and 0 > i ≥ Len.
40	IsNull(i int) bool
41
42	// IsValid returns true if value at index is not null.
43	// NOTE: IsValid will panic if NullBitmapBytes is not empty and 0 > i ≥ Len.
44	IsValid(i int) bool
45
46	Data() *Data
47
48	// Len returns the number of elements in the array.
49	Len() int
50
51	// Retain increases the reference count by 1.
52	// Retain may be called simultaneously from multiple goroutines.
53	Retain()
54
55	// Release decreases the reference count by 1.
56	// Release may be called simultaneously from multiple goroutines.
57	// When the reference count goes to zero, the memory is freed.
58	Release()
59}
60
61const (
62	// UnknownNullCount specifies the NullN should be calculated from the null bitmap buffer.
63	UnknownNullCount = -1
64)
65
66type array struct {
67	refCount        int64
68	data            *Data
69	nullBitmapBytes []byte
70}
71
72// Retain increases the reference count by 1.
73// Retain may be called simultaneously from multiple goroutines.
74func (a *array) Retain() {
75	atomic.AddInt64(&a.refCount, 1)
76}
77
78// Release decreases the reference count by 1.
79// Release may be called simultaneously from multiple goroutines.
80// When the reference count goes to zero, the memory is freed.
81func (a *array) Release() {
82	debug.Assert(atomic.LoadInt64(&a.refCount) > 0, "too many releases")
83
84	if atomic.AddInt64(&a.refCount, -1) == 0 {
85		a.data.Release()
86		a.data, a.nullBitmapBytes = nil, nil
87	}
88}
89
90// DataType returns the type metadata for this instance.
91func (a *array) DataType() arrow.DataType { return a.data.dtype }
92
93// NullN returns the number of null values in the array.
94func (a *array) NullN() int {
95	if a.data.nulls < 0 {
96		a.data.nulls = a.data.length - bitutil.CountSetBits(a.nullBitmapBytes, a.data.offset, a.data.length)
97	}
98	return a.data.nulls
99}
100
101// NullBitmapBytes returns a byte slice of the validity bitmap.
102func (a *array) NullBitmapBytes() []byte { return a.nullBitmapBytes }
103
104func (a *array) Data() *Data { return a.data }
105
106// Len returns the number of elements in the array.
107func (a *array) Len() int { return a.data.length }
108
109// IsNull returns true if value at index is null.
110// NOTE: IsNull will panic if NullBitmapBytes is not empty and 0 > i ≥ Len.
111func (a *array) IsNull(i int) bool {
112	return len(a.nullBitmapBytes) != 0 && bitutil.BitIsNotSet(a.nullBitmapBytes, a.data.offset+i)
113}
114
115// IsValid returns true if value at index is not null.
116// NOTE: IsValid will panic if NullBitmapBytes is not empty and 0 > i ≥ Len.
117func (a *array) IsValid(i int) bool {
118	return len(a.nullBitmapBytes) == 0 || bitutil.BitIsSet(a.nullBitmapBytes, a.data.offset+i)
119}
120
121func (a *array) setData(data *Data) {
122	// Retain before releasing in case a.data is the same as data.
123	data.Retain()
124
125	if a.data != nil {
126		a.data.Release()
127	}
128
129	if len(data.buffers) > 0 && data.buffers[0] != nil {
130		a.nullBitmapBytes = data.buffers[0].Bytes()
131	}
132	a.data = data
133}
134
135func (a *array) Offset() int {
136	return a.data.Offset()
137}
138
139type arrayConstructorFn func(*Data) Interface
140
141var (
142	makeArrayFn [64]arrayConstructorFn
143)
144
145func unsupportedArrayType(data *Data) Interface {
146	panic("unsupported data type: " + data.dtype.ID().String())
147}
148
149func invalidDataType(data *Data) Interface {
150	panic("invalid data type: " + data.dtype.ID().String())
151}
152
153// MakeFromData constructs a strongly-typed array instance from generic Data.
154func MakeFromData(data *Data) Interface {
155	return makeArrayFn[byte(data.dtype.ID()&0x3f)](data)
156}
157
158// NewSlice constructs a zero-copy slice of the array with the indicated
159// indices i and j, corresponding to array[i:j].
160// The returned array must be Release()'d after use.
161//
162// NewSlice panics if the slice is outside the valid range of the input array.
163// NewSlice panics if j < i.
164func NewSlice(arr Interface, i, j int64) Interface {
165	data := NewSliceData(arr.Data(), i, j)
166	slice := MakeFromData(data)
167	data.Release()
168	return slice
169}
170
171func init() {
172	makeArrayFn = [...]arrayConstructorFn{
173		arrow.NULL:                    func(data *Data) Interface { return NewNullData(data) },
174		arrow.BOOL:                    func(data *Data) Interface { return NewBooleanData(data) },
175		arrow.UINT8:                   func(data *Data) Interface { return NewUint8Data(data) },
176		arrow.INT8:                    func(data *Data) Interface { return NewInt8Data(data) },
177		arrow.UINT16:                  func(data *Data) Interface { return NewUint16Data(data) },
178		arrow.INT16:                   func(data *Data) Interface { return NewInt16Data(data) },
179		arrow.UINT32:                  func(data *Data) Interface { return NewUint32Data(data) },
180		arrow.INT32:                   func(data *Data) Interface { return NewInt32Data(data) },
181		arrow.UINT64:                  func(data *Data) Interface { return NewUint64Data(data) },
182		arrow.INT64:                   func(data *Data) Interface { return NewInt64Data(data) },
183		arrow.FLOAT16:                 func(data *Data) Interface { return NewFloat16Data(data) },
184		arrow.FLOAT32:                 func(data *Data) Interface { return NewFloat32Data(data) },
185		arrow.FLOAT64:                 func(data *Data) Interface { return NewFloat64Data(data) },
186		arrow.STRING:                  func(data *Data) Interface { return NewStringData(data) },
187		arrow.BINARY:                  func(data *Data) Interface { return NewBinaryData(data) },
188		arrow.FIXED_SIZE_BINARY:       func(data *Data) Interface { return NewFixedSizeBinaryData(data) },
189		arrow.DATE32:                  func(data *Data) Interface { return NewDate32Data(data) },
190		arrow.DATE64:                  func(data *Data) Interface { return NewDate64Data(data) },
191		arrow.TIMESTAMP:               func(data *Data) Interface { return NewTimestampData(data) },
192		arrow.TIME32:                  func(data *Data) Interface { return NewTime32Data(data) },
193		arrow.TIME64:                  func(data *Data) Interface { return NewTime64Data(data) },
194		arrow.INTERVAL_MONTHS:         func(data *Data) Interface { return NewMonthIntervalData(data) },
195		arrow.INTERVAL_DAY_TIME:       func(data *Data) Interface { return NewDayTimeIntervalData(data) },
196		arrow.DECIMAL128:              func(data *Data) Interface { return NewDecimal128Data(data) },
197		arrow.DECIMAL256:              unsupportedArrayType,
198		arrow.LIST:                    func(data *Data) Interface { return NewListData(data) },
199		arrow.STRUCT:                  func(data *Data) Interface { return NewStructData(data) },
200		arrow.SPARSE_UNION:            unsupportedArrayType,
201		arrow.DENSE_UNION:             unsupportedArrayType,
202		arrow.DICTIONARY:              unsupportedArrayType,
203		arrow.MAP:                     func(data *Data) Interface { return NewMapData(data) },
204		arrow.EXTENSION:               func(data *Data) Interface { return NewExtensionData(data) },
205		arrow.FIXED_SIZE_LIST:         func(data *Data) Interface { return NewFixedSizeListData(data) },
206		arrow.DURATION:                func(data *Data) Interface { return NewDurationData(data) },
207		arrow.LARGE_STRING:            unsupportedArrayType,
208		arrow.LARGE_BINARY:            unsupportedArrayType,
209		arrow.LARGE_LIST:              unsupportedArrayType,
210		arrow.INTERVAL:                func(data *Data) Interface { return NewIntervalData(data) },
211		arrow.INTERVAL_MONTH_DAY_NANO: func(data *Data) Interface { return NewMonthDayNanoIntervalData(data) },
212
213		// invalid data types to fill out array to size 2^6 - 1
214		63: invalidDataType,
215	}
216}
217