| // Licensed to the Apache Software Foundation (ASF) under one |
| // or more contributor license agreements. See the NOTICE file |
| // distributed with this work for additional information |
| // regarding copyright ownership. The ASF licenses this file |
| // to you under the Apache License, Version 2.0 (the |
| // "License"); you may not use this file except in compliance |
| // with the License. You may obtain a copy of the License at |
| // |
| // http://www.apache.org/licenses/LICENSE-2.0 |
| // |
| // Unless required by applicable law or agreed to in writing, software |
| // distributed under the License is distributed on an "AS IS" BASIS, |
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| // See the License for the specific language governing permissions and |
| // limitations under the License. |
| |
| package array |
| |
| import ( |
| "sync/atomic" |
| |
| "github.com/apache/arrow-go/v18/arrow" |
| "github.com/apache/arrow-go/v18/arrow/bitutil" |
| "github.com/apache/arrow-go/v18/arrow/internal/debug" |
| ) |
| |
| const ( |
| // UnknownNullCount specifies the NullN should be calculated from the null bitmap buffer. |
| UnknownNullCount = -1 |
| |
| // NullValueStr represents a null value in arrow.Array.ValueStr and in Builder.AppendValueFromString. |
| // It should be returned from the arrow.Array.ValueStr implementations. |
| // Using it as the value in Builder.AppendValueFromString should be equivalent to Builder.AppendNull. |
| NullValueStr = "(null)" |
| ) |
| |
| type array struct { |
| refCount atomic.Int64 |
| data *Data |
| nullBitmapBytes []byte |
| } |
| |
| // Retain increases the reference count by 1. |
| // Retain may be called simultaneously from multiple goroutines. |
| func (a *array) Retain() { |
| a.refCount.Add(1) |
| } |
| |
| // Release decreases the reference count by 1. |
| // Release may be called simultaneously from multiple goroutines. |
| // When the reference count goes to zero, the memory is freed. |
| func (a *array) Release() { |
| debug.Assert(a.refCount.Load() > 0, "too many releases") |
| |
| if a.refCount.Add(-1) == 0 { |
| a.data.Release() |
| a.data, a.nullBitmapBytes = nil, nil |
| } |
| } |
| |
| // DataType returns the type metadata for this instance. |
| func (a *array) DataType() arrow.DataType { return a.data.dtype } |
| |
| // NullN returns the number of null values in the array. |
| func (a *array) NullN() int { |
| if a.data.nulls < 0 { |
| a.data.nulls = a.data.length - bitutil.CountSetBits(a.nullBitmapBytes, a.data.offset, a.data.length) |
| } |
| return a.data.nulls |
| } |
| |
| // NullBitmapBytes returns a byte slice of the validity bitmap. |
| func (a *array) NullBitmapBytes() []byte { return a.nullBitmapBytes } |
| |
| func (a *array) Data() arrow.ArrayData { return a.data } |
| |
| // Len returns the number of elements in the array. |
| func (a *array) Len() int { return a.data.length } |
| |
| // IsNull returns true if value at index is null. |
| // NOTE: IsNull will panic if NullBitmapBytes is not empty and 0 > i ≥ Len. |
| func (a *array) IsNull(i int) bool { |
| return len(a.nullBitmapBytes) != 0 && bitutil.BitIsNotSet(a.nullBitmapBytes, a.data.offset+i) |
| } |
| |
| // IsValid returns true if value at index is not null. |
| // NOTE: IsValid will panic if NullBitmapBytes is not empty and 0 > i ≥ Len. |
| func (a *array) IsValid(i int) bool { |
| return len(a.nullBitmapBytes) == 0 || bitutil.BitIsSet(a.nullBitmapBytes, a.data.offset+i) |
| } |
| |
| func (a *array) setData(data *Data) { |
| // Retain before releasing in case a.data is the same as data. |
| data.Retain() |
| |
| if a.data != nil { |
| a.data.Release() |
| } |
| |
| if len(data.buffers) > 0 && data.buffers[0] != nil { |
| a.nullBitmapBytes = data.buffers[0].Bytes() |
| } |
| a.data = data |
| } |
| |
| func (a *array) Offset() int { |
| return a.data.Offset() |
| } |
| |
| type arrayConstructorFn func(arrow.ArrayData) arrow.Array |
| |
| var makeArrayFn [64]arrayConstructorFn |
| |
| func invalidDataType(data arrow.ArrayData) arrow.Array { |
| panic("invalid data type: " + data.DataType().ID().String()) |
| } |
| |
| // MakeFromData constructs a strongly-typed array instance from generic Data. |
| func MakeFromData(data arrow.ArrayData) arrow.Array { |
| id := data.DataType().ID() |
| if id < 0 || int(id) >= len(makeArrayFn) || makeArrayFn[id] == nil { |
| return invalidDataType(data) |
| } |
| return makeArrayFn[id](data) |
| } |
| |
| // NewSlice constructs a zero-copy slice of the array with the indicated |
| // indices i and j, corresponding to array[i:j]. |
| // The returned array must be Release()'d after use. |
| // |
| // NewSlice panics if the slice is outside the valid range of the input array. |
| // NewSlice panics if j < i. |
| func NewSlice(arr arrow.Array, i, j int64) arrow.Array { |
| data := NewSliceData(arr.Data(), i, j) |
| slice := MakeFromData(data) |
| data.Release() |
| return slice |
| } |
| |
| func init() { |
| makeArrayFn = [...]arrayConstructorFn{ |
| arrow.NULL: func(data arrow.ArrayData) arrow.Array { return NewNullData(data) }, |
| arrow.BOOL: func(data arrow.ArrayData) arrow.Array { return NewBooleanData(data) }, |
| arrow.UINT8: func(data arrow.ArrayData) arrow.Array { return NewUint8Data(data) }, |
| arrow.INT8: func(data arrow.ArrayData) arrow.Array { return NewInt8Data(data) }, |
| arrow.UINT16: func(data arrow.ArrayData) arrow.Array { return NewUint16Data(data) }, |
| arrow.INT16: func(data arrow.ArrayData) arrow.Array { return NewInt16Data(data) }, |
| arrow.UINT32: func(data arrow.ArrayData) arrow.Array { return NewUint32Data(data) }, |
| arrow.INT32: func(data arrow.ArrayData) arrow.Array { return NewInt32Data(data) }, |
| arrow.UINT64: func(data arrow.ArrayData) arrow.Array { return NewUint64Data(data) }, |
| arrow.INT64: func(data arrow.ArrayData) arrow.Array { return NewInt64Data(data) }, |
| arrow.FLOAT16: func(data arrow.ArrayData) arrow.Array { return NewFloat16Data(data) }, |
| arrow.FLOAT32: func(data arrow.ArrayData) arrow.Array { return NewFloat32Data(data) }, |
| arrow.FLOAT64: func(data arrow.ArrayData) arrow.Array { return NewFloat64Data(data) }, |
| arrow.STRING: func(data arrow.ArrayData) arrow.Array { return NewStringData(data) }, |
| arrow.BINARY: func(data arrow.ArrayData) arrow.Array { return NewBinaryData(data) }, |
| arrow.FIXED_SIZE_BINARY: func(data arrow.ArrayData) arrow.Array { return NewFixedSizeBinaryData(data) }, |
| arrow.DATE32: func(data arrow.ArrayData) arrow.Array { return NewDate32Data(data) }, |
| arrow.DATE64: func(data arrow.ArrayData) arrow.Array { return NewDate64Data(data) }, |
| arrow.TIMESTAMP: func(data arrow.ArrayData) arrow.Array { return NewTimestampData(data) }, |
| arrow.TIME32: func(data arrow.ArrayData) arrow.Array { return NewTime32Data(data) }, |
| arrow.TIME64: func(data arrow.ArrayData) arrow.Array { return NewTime64Data(data) }, |
| arrow.INTERVAL_MONTHS: func(data arrow.ArrayData) arrow.Array { return NewMonthIntervalData(data) }, |
| arrow.INTERVAL_DAY_TIME: func(data arrow.ArrayData) arrow.Array { return NewDayTimeIntervalData(data) }, |
| arrow.DECIMAL32: func(data arrow.ArrayData) arrow.Array { return NewDecimal32Data(data) }, |
| arrow.DECIMAL64: func(data arrow.ArrayData) arrow.Array { return NewDecimal64Data(data) }, |
| arrow.DECIMAL128: func(data arrow.ArrayData) arrow.Array { return NewDecimal128Data(data) }, |
| arrow.DECIMAL256: func(data arrow.ArrayData) arrow.Array { return NewDecimal256Data(data) }, |
| arrow.LIST: func(data arrow.ArrayData) arrow.Array { return NewListData(data) }, |
| arrow.STRUCT: func(data arrow.ArrayData) arrow.Array { return NewStructData(data) }, |
| arrow.SPARSE_UNION: func(data arrow.ArrayData) arrow.Array { return NewSparseUnionData(data) }, |
| arrow.DENSE_UNION: func(data arrow.ArrayData) arrow.Array { return NewDenseUnionData(data) }, |
| arrow.DICTIONARY: func(data arrow.ArrayData) arrow.Array { return NewDictionaryData(data) }, |
| arrow.MAP: func(data arrow.ArrayData) arrow.Array { return NewMapData(data) }, |
| arrow.EXTENSION: func(data arrow.ArrayData) arrow.Array { return NewExtensionData(data) }, |
| arrow.FIXED_SIZE_LIST: func(data arrow.ArrayData) arrow.Array { return NewFixedSizeListData(data) }, |
| arrow.DURATION: func(data arrow.ArrayData) arrow.Array { return NewDurationData(data) }, |
| arrow.LARGE_STRING: func(data arrow.ArrayData) arrow.Array { return NewLargeStringData(data) }, |
| arrow.LARGE_BINARY: func(data arrow.ArrayData) arrow.Array { return NewLargeBinaryData(data) }, |
| arrow.LARGE_LIST: func(data arrow.ArrayData) arrow.Array { return NewLargeListData(data) }, |
| arrow.INTERVAL_MONTH_DAY_NANO: func(data arrow.ArrayData) arrow.Array { return NewMonthDayNanoIntervalData(data) }, |
| arrow.RUN_END_ENCODED: func(data arrow.ArrayData) arrow.Array { return NewRunEndEncodedData(data) }, |
| arrow.LIST_VIEW: func(data arrow.ArrayData) arrow.Array { return NewListViewData(data) }, |
| arrow.LARGE_LIST_VIEW: func(data arrow.ArrayData) arrow.Array { return NewLargeListViewData(data) }, |
| arrow.BINARY_VIEW: func(data arrow.ArrayData) arrow.Array { return NewBinaryViewData(data) }, |
| arrow.STRING_VIEW: func(data arrow.ArrayData) arrow.Array { return NewStringViewData(data) }, |
| // invalid data types to fill out array to size 2^6 - 1 |
| 63: invalidDataType, |
| } |
| } |