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