/
type_traits.go
162 lines (142 loc) · 6.05 KB
/
type_traits.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
// 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 arrow
import (
"reflect"
"unsafe"
"github.com/apache/arrow/go/v16/arrow/decimal128"
"github.com/apache/arrow/go/v16/arrow/decimal256"
"github.com/apache/arrow/go/v16/arrow/float16"
"golang.org/x/exp/constraints"
)
// IntType is a type constraint for raw values represented as signed
// integer types by We aren't just using constraints.Signed
// because we don't want to include the raw `int` type here whose size
// changes based on the architecture (int32 on 32-bit architectures and
// int64 on 64-bit architectures).
//
// This will also cover types like MonthInterval or the time types
// as their underlying types are int32 and int64 which will get covered
// by using the ~
type IntType interface {
~int8 | ~int16 | ~int32 | ~int64
}
// UintType is a type constraint for raw values represented as unsigned
// integer types by We aren't just using constraints.Unsigned
// because we don't want to include the raw `uint` type here whose size
// changes based on the architecture (uint32 on 32-bit architectures and
// uint64 on 64-bit architectures). We also don't want to include uintptr
type UintType interface {
~uint8 | ~uint16 | ~uint32 | ~uint64
}
// FloatType is a type constraint for raw values for representing
// floating point values in This consists of constraints.Float and
// float16.Num
type FloatType interface {
float16.Num | constraints.Float
}
// NumericType is a type constraint for just signed/unsigned integers
// and float32/float64.
type NumericType interface {
IntType | UintType | constraints.Float
}
// FixedWidthType is a type constraint for raw values in Arrow that
// can be represented as FixedWidth byte slices. Specifically this is for
// using Go generics to easily re-type a byte slice to a properly-typed
// slice. Booleans are excluded here since they are represented by Arrow
// as a bitmap and thus the buffer can't be just reinterpreted as a []bool
type FixedWidthType interface {
IntType | UintType |
FloatType | decimal128.Num | decimal256.Num |
DayTimeInterval | MonthDayNanoInterval
}
type TemporalType interface {
Date32 | Date64 | Time32 | Time64 |
Timestamp | Duration | DayTimeInterval |
MonthInterval | MonthDayNanoInterval
}
func reinterpretSlice[Out, T any](b []T) []Out {
if cap(b) == 0 {
return nil
}
out := (*Out)(unsafe.Pointer(&b[:1][0]))
lenBytes := len(b) * int(unsafe.Sizeof(b[0]))
capBytes := cap(b) * int(unsafe.Sizeof(b[0]))
lenOut := lenBytes / int(unsafe.Sizeof(*out))
capOut := capBytes / int(unsafe.Sizeof(*out))
return unsafe.Slice(out, capOut)[:lenOut]
}
// GetValues reinterprets the data.Buffers()[i] to a slice of T with len=data.Len().
//
// If the buffer is nil, nil will be returned.
//
// NOTE: the buffer's length must be a multiple of Sizeof(T).
func GetValues[T FixedWidthType](data ArrayData, i int) []T {
if data.Buffers()[i] == nil || data.Buffers()[i].Len() == 0 {
return nil
}
return reinterpretSlice[T](data.Buffers()[i].Bytes())[data.Offset() : data.Offset()+data.Len()]
}
// GetOffsets reinterprets the data.Buffers()[i] to a slice of T with len=data.Len()+1.
//
// NOTE: the buffer's length must be a multiple of Sizeof(T).
func GetOffsets[T int32 | int64](data ArrayData, i int) []T {
return reinterpretSlice[T](data.Buffers()[i].Bytes())[data.Offset() : data.Offset()+data.Len()+1]
}
// GetBytes reinterprets a slice of T to a slice of bytes.
func GetBytes[T FixedWidthType | ViewHeader](in []T) []byte {
return reinterpretSlice[byte](in)
}
// GetData reinterprets a slice of bytes to a slice of T.
//
// NOTE: the buffer's length must be a multiple of Sizeof(T).
func GetData[T FixedWidthType | ViewHeader](in []byte) []T {
return reinterpretSlice[T](in)
}
var typMap = map[reflect.Type]DataType{
reflect.TypeOf(false): FixedWidthTypes.Boolean,
reflect.TypeOf(int8(0)): PrimitiveTypes.Int8,
reflect.TypeOf(int16(0)): PrimitiveTypes.Int16,
reflect.TypeOf(int32(0)): PrimitiveTypes.Int32,
reflect.TypeOf(int64(0)): PrimitiveTypes.Int64,
reflect.TypeOf(uint8(0)): PrimitiveTypes.Uint8,
reflect.TypeOf(uint16(0)): PrimitiveTypes.Uint16,
reflect.TypeOf(uint32(0)): PrimitiveTypes.Uint32,
reflect.TypeOf(uint64(0)): PrimitiveTypes.Uint64,
reflect.TypeOf(float32(0)): PrimitiveTypes.Float32,
reflect.TypeOf(float64(0)): PrimitiveTypes.Float64,
reflect.TypeOf(string("")): BinaryTypes.String,
reflect.TypeOf(Date32(0)): FixedWidthTypes.Date32,
reflect.TypeOf(Date64(0)): FixedWidthTypes.Date64,
reflect.TypeOf(true): FixedWidthTypes.Boolean,
reflect.TypeOf(float16.Num{}): FixedWidthTypes.Float16,
reflect.TypeOf([]byte{}): BinaryTypes.Binary,
}
// GetDataType returns the appropriate DataType for the given type T
// only for non-parametric types. This uses a map and reflection internally
// so don't call this in a tight loop, instead call this once and then use
// a closure with the result.
func GetDataType[T NumericType | bool | string | []byte | float16.Num]() DataType {
var z T
return typMap[reflect.TypeOf(z)]
}
// GetType returns the appropriate Type type T, only for non-parametric
// types. This uses a map and reflection internally so don't call this in
// a tight loop, instead call it once and then use a closure with the result.
func GetType[T NumericType | bool | string]() Type {
var z T
return typMap[reflect.TypeOf(z)].ID()
}