]>
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, | |
12 | // software distributed under the License is distributed on an | |
13 | // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | |
14 | // KIND, either express or implied. See the License for the | |
15 | // specific language governing permissions and limitations | |
16 | // under the License. | |
17 | ||
18 | import { Data } from '../data'; | |
19 | import { Field } from '../schema'; | |
20 | import { Column } from '../column'; | |
21 | import { Vector } from '../vector'; | |
22 | import { DataType, Float32, Float64, FloatArray, IntArray, Int16, Int32, Int64, Int8, Uint16, Uint32, Uint64, Uint8 } from '../type'; | |
23 | import { Chunked } from '../vector/chunked'; | |
24 | import { BigIntArray, TypedArray as TypedArray_ } from '../interfaces'; | |
25 | import { FloatArrayCtor } from '../vector/float'; | |
26 | import { IntArrayCtor } from '../vector/int'; | |
27 | ||
28 | type RecordBatchCtor = typeof import('../recordbatch').RecordBatch; | |
29 | ||
30 | const isArray = Array.isArray; | |
31 | ||
32 | type TypedArray = Exclude<TypedArray_ | BigIntArray, Uint8ClampedArray>; | |
33 | ||
34 | /** @ignore */ | |
35 | export function isTypedArray(arr: any): arr is TypedArray { | |
36 | return ArrayBuffer.isView(arr) && 'BYTES_PER_ELEMENT' in arr; | |
37 | } | |
38 | ||
39 | ||
40 | /** @ignore */ | |
41 | type ArrayCtor = FloatArrayCtor | IntArrayCtor; | |
42 | ||
43 | /** @ignore */ | |
44 | export function arrayTypeToDataType(ctor: ArrayCtor) { | |
45 | switch (ctor) { | |
46 | case Int8Array: return Int8; | |
47 | case Int16Array: return Int16; | |
48 | case Int32Array: return Int32; | |
49 | case BigInt64Array: return Int64; | |
50 | case Uint8Array: return Uint8; | |
51 | case Uint16Array: return Uint16; | |
52 | case Uint32Array: return Uint32; | |
53 | case BigUint64Array: return Uint64; | |
54 | case Float32Array: return Float32; | |
55 | case Float64Array: return Float64; | |
56 | default: return null; | |
57 | } | |
58 | } | |
59 | ||
60 | /** @ignore */ | |
61 | function vectorFromTypedArray(array: TypedArray): Vector { | |
62 | const ArrowType = arrayTypeToDataType(array.constructor as ArrayCtor); | |
63 | if (!ArrowType) { | |
64 | throw new TypeError('Unrecognized Array input'); | |
65 | } | |
66 | const type = new ArrowType(); | |
67 | const data = Data.new(type, 0, array.length, 0, [undefined, array as IntArray | FloatArray]); | |
68 | return Vector.new(data); | |
69 | } | |
70 | ||
71 | /** @ignore */ | |
72 | export const selectArgs = <T>(Ctor: any, vals: any[]) => _selectArgs(Ctor, vals, [], 0) as T[]; | |
73 | /** @ignore */ | |
74 | export const selectColumnArgs = <T extends { [key: string]: DataType }>(args: any[]) => { | |
75 | const [fields, values] = _selectFieldArgs<T>(args, [[], []]); | |
76 | return values.map((x, i) => | |
77 | x instanceof Column ? Column.new(x.field.clone(fields[i]), x) : | |
78 | x instanceof Vector ? Column.new(fields[i], x) as Column<T[keyof T]> : | |
79 | isTypedArray(x) ? Column.new(fields[i], vectorFromTypedArray(x)) as Column<T[keyof T]> : | |
80 | Column.new(fields[i], [] as Vector<T[keyof T]>[])); | |
81 | }; | |
82 | ||
83 | /** @ignore */ | |
84 | export const selectFieldArgs = <T extends { [key: string]: DataType }>(args: any[]) => _selectFieldArgs<T>(args, [[], []]); | |
85 | /** @ignore */ | |
86 | export const selectChunkArgs = <T>(Ctor: any, vals: any[]) => _selectChunkArgs(Ctor, vals, [], 0) as T[]; | |
87 | /** @ignore */ | |
88 | export const selectVectorChildrenArgs = <T extends Vector>(Ctor: RecordBatchCtor, vals: any[]) => _selectVectorChildrenArgs(Ctor, vals, [], 0) as T[]; | |
89 | /** @ignore */ | |
90 | export const selectColumnChildrenArgs = <T extends Column>(Ctor: RecordBatchCtor, vals: any[]) => _selectColumnChildrenArgs(Ctor, vals, [], 0) as T[]; | |
91 | ||
92 | /** @ignore */ | |
93 | function _selectArgs<T>(Ctor: any, vals: any[], res: T[], idx: number) { | |
94 | let value: any, j = idx; | |
95 | let i = -1; | |
96 | const n = vals.length; | |
97 | while (++i < n) { | |
98 | if (isArray(value = vals[i])) { | |
99 | j = _selectArgs(Ctor, value, res, j).length; | |
100 | } else if (value instanceof Ctor) { res[j++] = value; } | |
101 | } | |
102 | return res; | |
103 | } | |
104 | ||
105 | /** @ignore */ | |
106 | function _selectChunkArgs<T>(Ctor: any, vals: any[], res: T[], idx: number) { | |
107 | let value: any, j = idx; | |
108 | let i = -1; | |
109 | const n = vals.length; | |
110 | while (++i < n) { | |
111 | if (isArray(value = vals[i])) { | |
112 | j = _selectChunkArgs(Ctor, value, res, j).length; | |
113 | } else if (value instanceof Chunked) { | |
114 | j = _selectChunkArgs(Ctor, value.chunks, res, j).length; | |
115 | } else if (value instanceof Ctor) { res[j++] = value; } | |
116 | } | |
117 | return res; | |
118 | } | |
119 | ||
120 | /** @ignore */ | |
121 | function _selectVectorChildrenArgs<T extends Vector>(Ctor: RecordBatchCtor, vals: any[], res: T[], idx: number) { | |
122 | let value: any, j = idx; | |
123 | let i = -1; | |
124 | const n = vals.length; | |
125 | while (++i < n) { | |
126 | if (isArray(value = vals[i])) { | |
127 | j = _selectVectorChildrenArgs(Ctor, value, res, j).length; | |
128 | } else if (value instanceof Ctor) { | |
129 | j = _selectArgs(Vector, value.schema.fields.map((_, i) => value.getChildAt(i)!), res, j).length; | |
130 | } else if (value instanceof Vector) { res[j++] = value as T; } | |
131 | } | |
132 | return res; | |
133 | } | |
134 | ||
135 | /** @ignore */ | |
136 | function _selectColumnChildrenArgs<T extends Column>(Ctor: RecordBatchCtor, vals: any[], res: T[], idx: number) { | |
137 | let value: any, j = idx; | |
138 | let i = -1; | |
139 | const n = vals.length; | |
140 | while (++i < n) { | |
141 | if (isArray(value = vals[i])) { | |
142 | j = _selectColumnChildrenArgs(Ctor, value, res, j).length; | |
143 | } else if (value instanceof Ctor) { | |
144 | j = _selectArgs(Column, value.schema.fields.map((f, i) => Column.new(f, value.getChildAt(i)!)), res, j).length; | |
145 | } else if (value instanceof Column) { res[j++] = value as T; } | |
146 | } | |
147 | return res; | |
148 | } | |
149 | ||
150 | /** @ignore */ | |
151 | const toKeysAndValues = (xs: [any[], any[]], [k, v]: [any, any], i: number) => (xs[0][i] = k, xs[1][i] = v, xs); | |
152 | ||
153 | /** @ignore */ | |
154 | function _selectFieldArgs<T extends { [key: string]: DataType }>(vals: any[], ret: [Field<T[keyof T]>[], (Vector<T[keyof T]> | TypedArray)[]]): [Field<T[keyof T]>[], (T[keyof T] | Vector<T[keyof T]> | TypedArray)[]] { | |
155 | let keys: any[]; | |
156 | let n: number; | |
157 | switch (n = vals.length) { | |
158 | case 0: return ret; | |
159 | case 1: | |
160 | keys = ret[0]; | |
161 | if (!(vals[0])) { return ret; } | |
162 | if (isArray(vals[0])) { return _selectFieldArgs(vals[0], ret); } | |
163 | if (!(vals[0] instanceof Data || vals[0] instanceof Vector || isTypedArray(vals[0]) || vals[0] instanceof DataType)) { | |
164 | [keys, vals] = Object.entries(vals[0]).reduce(toKeysAndValues, ret); | |
165 | } | |
166 | break; | |
167 | default: | |
168 | !isArray(keys = vals[n - 1]) | |
169 | ? (vals = isArray(vals[0]) ? vals[0] : vals, keys = []) | |
170 | : (vals = isArray(vals[0]) ? vals[0] : vals.slice(0, n - 1)); | |
171 | } | |
172 | ||
173 | let fieldIndex = -1; | |
174 | let valueIndex = -1; | |
175 | let idx = -1; | |
176 | const len = vals.length; | |
177 | let field: number | string | Field<T[keyof T]>; | |
178 | let val: Vector<T[keyof T]> | Data<T[keyof T]>; | |
179 | const [fields, values] = ret as [Field<T[keyof T]>[], any[]]; | |
180 | ||
181 | while (++idx < len) { | |
182 | val = vals[idx]; | |
183 | if (val instanceof Column && (values[++valueIndex] = val)) { | |
184 | fields[++fieldIndex] = val.field.clone(keys[idx], val.type, true); | |
185 | } else { | |
186 | ({ [idx]: field = idx } = keys); | |
187 | if (val instanceof DataType && (values[++valueIndex] = val)) { | |
188 | fields[++fieldIndex] = Field.new(field, val as DataType, true) as Field<T[keyof T]>; | |
189 | } else if (val?.type && (values[++valueIndex] = val)) { | |
190 | val instanceof Data && (values[valueIndex] = val = Vector.new(val) as Vector); | |
191 | fields[++fieldIndex] = Field.new(field, val.type, true) as Field<T[keyof T]>; | |
192 | } | |
193 | } | |
194 | } | |
195 | return ret; | |
196 | } |