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 18import { Data } from '../data'; 19import { Field } from '../schema'; 20import { Column } from '../column'; 21import { Vector } from '../vector'; 22import { DataType } from '../type'; 23import { Chunked } from '../vector/chunked'; 24 25type RecordBatchCtor = typeof import('../recordbatch').RecordBatch; 26 27const isArray = Array.isArray; 28 29/** @ignore */ 30export const selectArgs = <T>(Ctor: any, vals: any[]) => _selectArgs(Ctor, vals, [], 0) as T[]; 31/** @ignore */ 32export const selectColumnArgs = <T extends { [key: string]: DataType }>(args: any[]) => { 33 const [fields, values] = _selectFieldArgs<T>(args, [[], []]); 34 return values.map((x, i) => 35 x instanceof Column ? Column.new(x.field.clone(fields[i]), x) : 36 x instanceof Vector ? Column.new(fields[i], x) as Column<T[keyof T]> : 37 Column.new(fields[i], [] as Vector<T[keyof T]>[])); 38}; 39 40/** @ignore */ 41export const selectFieldArgs = <T extends { [key: string]: DataType }>(args: any[]) => _selectFieldArgs<T>(args, [[], []]); 42/** @ignore */ 43export const selectChunkArgs = <T>(Ctor: any, vals: any[]) => _selectChunkArgs(Ctor, vals, [], 0) as T[]; 44/** @ignore */ 45export const selectVectorChildrenArgs = <T extends Vector>(Ctor: RecordBatchCtor, vals: any[]) => _selectVectorChildrenArgs(Ctor, vals, [], 0) as T[]; 46/** @ignore */ 47export const selectColumnChildrenArgs = <T extends Column>(Ctor: RecordBatchCtor, vals: any[]) => _selectColumnChildrenArgs(Ctor, vals, [], 0) as T[]; 48 49/** @ignore */ 50function _selectArgs<T>(Ctor: any, vals: any[], res: T[], idx: number) { 51 let value: any, j = idx; 52 let i = -1, n = vals.length; 53 while (++i < n) { 54 if (isArray(value = vals[i])) { 55 j = _selectArgs(Ctor, value, res, j).length; 56 } else if (value instanceof Ctor) { res[j++] = value; } 57 } 58 return res; 59} 60 61/** @ignore */ 62function _selectChunkArgs<T>(Ctor: any, vals: any[], res: T[], idx: number) { 63 let value: any, j = idx; 64 let i = -1, n = vals.length; 65 while (++i < n) { 66 if (isArray(value = vals[i])) { 67 j = _selectChunkArgs(Ctor, value, res, j).length; 68 } else if (value instanceof Chunked) { 69 j = _selectChunkArgs(Ctor, value.chunks, res, j).length; 70 } else if (value instanceof Ctor) { res[j++] = value; } 71 } 72 return res; 73} 74 75/** @ignore */ 76function _selectVectorChildrenArgs<T extends Vector>(Ctor: RecordBatchCtor, vals: any[], res: T[], idx: number) { 77 let value: any, j = idx; 78 let i = -1, n = vals.length; 79 while (++i < n) { 80 if (isArray(value = vals[i])) { 81 j = _selectVectorChildrenArgs(Ctor, value, res, j).length; 82 } else if (value instanceof Ctor) { 83 j = _selectArgs(Vector, value.schema.fields.map((_, i) => value.getChildAt(i)!), res, j).length; 84 } else if (value instanceof Vector) { res[j++] = value as T; } 85 } 86 return res; 87} 88 89/** @ignore */ 90function _selectColumnChildrenArgs<T extends Column>(Ctor: RecordBatchCtor, vals: any[], res: T[], idx: number) { 91 let value: any, j = idx; 92 let i = -1, n = vals.length; 93 while (++i < n) { 94 if (isArray(value = vals[i])) { 95 j = _selectColumnChildrenArgs(Ctor, value, res, j).length; 96 } else if (value instanceof Ctor) { 97 j = _selectArgs(Column, value.schema.fields.map((f, i) => Column.new(f, value.getChildAt(i)!)), res, j).length; 98 } else if (value instanceof Column) { res[j++] = value as T; } 99 } 100 return res; 101} 102 103/** @ignore */ 104const toKeysAndValues = (xs: [any[], any[]], [k, v]: [any, any], i: number) => (xs[0][i] = k, xs[1][i] = v, xs); 105 106/** @ignore */ 107function _selectFieldArgs<T extends { [key: string]: DataType }>(vals: any[], ret: [Field<T[keyof T]>[], Vector<T[keyof T]>[]]): [Field<T[keyof T]>[], (T[keyof T] | Vector<T[keyof T]>)[]] { 108 let keys: any[], n: number; 109 switch (n = vals.length) { 110 case 0: return ret; 111 case 1: 112 keys = ret[0]; 113 if (!(vals[0])) { return ret; } 114 if (isArray(vals[0])) { return _selectFieldArgs(vals[0], ret); } 115 if (!(vals[0] instanceof Data || vals[0] instanceof Vector || vals[0] instanceof DataType)) { 116 [keys, vals] = Object.entries(vals[0]).reduce(toKeysAndValues, ret); 117 } 118 break; 119 default: 120 !isArray(keys = vals[n - 1]) 121 ? (vals = isArray(vals[0]) ? vals[0] : vals, keys = []) 122 : (vals = isArray(vals[0]) ? vals[0] : vals.slice(0, n - 1)); 123 } 124 125 let fieldIndex = -1; 126 let valueIndex = -1; 127 let idx = -1, len = vals.length; 128 let field: number | string | Field<T[keyof T]>; 129 let val: Vector<T[keyof T]> | Data<T[keyof T]>; 130 let [fields, values] = ret as [Field<T[keyof T]>[], any[]]; 131 132 while (++idx < len) { 133 val = vals[idx]; 134 if (val instanceof Column && (values[++valueIndex] = val)) { 135 fields[++fieldIndex] = val.field.clone(keys[idx], val.type, true); 136 } else { 137 ({ [idx]: field = idx } = keys); 138 if (val instanceof DataType && (values[++valueIndex] = val)) { 139 fields[++fieldIndex] = Field.new(field, val as DataType, true) as Field<T[keyof T]>; 140 } else if (val && val.type && (values[++valueIndex] = val)) { 141 val instanceof Data && (values[valueIndex] = val = Vector.new(val) as Vector); 142 fields[++fieldIndex] = Field.new(field, val.type, true) as Field<T[keyof T]>; 143 } 144 } 145 } 146 return ret; 147} 148