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 * as type from '../type';
20import { Field } from '../schema';
21import { Vector } from '../vector';
22import { DataType } from '../type';
23import { Visitor } from '../visitor';
24import { packBools } from '../util/bit';
25import { encodeUtf8 } from '../util/utf8';
26import { Int64, Int128 } from '../util/int';
27import { UnionMode, DateUnit } from '../enum';
28import { toArrayBufferView } from '../util/buffer';
29import { BufferRegion, FieldNode } from '../ipc/metadata/message';
30
31/** @ignore */
32export interface VectorLoader extends Visitor {
33    visit<T extends DataType>(node: Field<T> | T): Data<T>;
34    visitMany<T extends DataType>(nodes: (Field<T> | T)[]): Data<T>[];
35}
36
37/** @ignore */
38export class VectorLoader extends Visitor {
39    private bytes: Uint8Array;
40    private nodes: FieldNode[];
41    private nodesIndex = -1;
42    private buffers: BufferRegion[];
43    private buffersIndex = -1;
44    private dictionaries: Map<number, Vector<any>>;
45    constructor(bytes: Uint8Array, nodes: FieldNode[], buffers: BufferRegion[], dictionaries: Map<number, Vector<any>>) {
46        super();
47        this.bytes = bytes;
48        this.nodes = nodes;
49        this.buffers = buffers;
50        this.dictionaries = dictionaries;
51    }
52
53    public visit<T extends DataType>(node: Field<T> | T): Data<T> {
54        return super.visit(node instanceof Field ? node.type : node);
55    }
56
57    public visitNull            <T extends type.Null>            (type: T, { length,           } = this.nextFieldNode()) { return            Data.Null(type, 0, length);                                                                                                                                 }
58    public visitBool            <T extends type.Bool>            (type: T, { length, nullCount } = this.nextFieldNode()) { return            Data.Bool(type, 0, length, nullCount, this.readNullBitmap(type, nullCount), this.readData(type));                                                           }
59    public visitInt             <T extends type.Int>             (type: T, { length, nullCount } = this.nextFieldNode()) { return             Data.Int(type, 0, length, nullCount, this.readNullBitmap(type, nullCount), this.readData(type));                                                           }
60    public visitFloat           <T extends type.Float>           (type: T, { length, nullCount } = this.nextFieldNode()) { return           Data.Float(type, 0, length, nullCount, this.readNullBitmap(type, nullCount), this.readData(type));                                                           }
61    public visitUtf8            <T extends type.Utf8>            (type: T, { length, nullCount } = this.nextFieldNode()) { return            Data.Utf8(type, 0, length, nullCount, this.readNullBitmap(type, nullCount), this.readOffsets(type), this.readData(type));                                   }
62    public visitBinary          <T extends type.Binary>          (type: T, { length, nullCount } = this.nextFieldNode()) { return          Data.Binary(type, 0, length, nullCount, this.readNullBitmap(type, nullCount), this.readOffsets(type), this.readData(type));                                   }
63    public visitFixedSizeBinary <T extends type.FixedSizeBinary> (type: T, { length, nullCount } = this.nextFieldNode()) { return Data.FixedSizeBinary(type, 0, length, nullCount, this.readNullBitmap(type, nullCount), this.readData(type));                                                           }
64    public visitDate            <T extends type.Date_>           (type: T, { length, nullCount } = this.nextFieldNode()) { return            Data.Date(type, 0, length, nullCount, this.readNullBitmap(type, nullCount), this.readData(type));                                                           }
65    public visitTimestamp       <T extends type.Timestamp>       (type: T, { length, nullCount } = this.nextFieldNode()) { return       Data.Timestamp(type, 0, length, nullCount, this.readNullBitmap(type, nullCount), this.readData(type));                                                           }
66    public visitTime            <T extends type.Time>            (type: T, { length, nullCount } = this.nextFieldNode()) { return            Data.Time(type, 0, length, nullCount, this.readNullBitmap(type, nullCount), this.readData(type));                                                           }
67    public visitDecimal         <T extends type.Decimal>         (type: T, { length, nullCount } = this.nextFieldNode()) { return         Data.Decimal(type, 0, length, nullCount, this.readNullBitmap(type, nullCount), this.readData(type));                                                           }
68    public visitList            <T extends type.List>            (type: T, { length, nullCount } = this.nextFieldNode()) { return            Data.List(type, 0, length, nullCount, this.readNullBitmap(type, nullCount), this.readOffsets(type), this.visit(type.children[0]));                          }
69    public visitStruct          <T extends type.Struct>          (type: T, { length, nullCount } = this.nextFieldNode()) { return          Data.Struct(type, 0, length, nullCount, this.readNullBitmap(type, nullCount), this.visitMany(type.children));                                                 }
70    public visitUnion           <T extends type.Union>           (type: T                                              ) { return type.mode === UnionMode.Sparse ? this.visitSparseUnion(type as type.SparseUnion) : this.visitDenseUnion(type as type.DenseUnion);                                      }
71    public visitDenseUnion      <T extends type.DenseUnion>      (type: T, { length, nullCount } = this.nextFieldNode()) { return           Data.Union(type, 0, length, nullCount, this.readNullBitmap(type, nullCount), this.readTypeIds(type), this.readOffsets(type), this.visitMany(type.children)); }
72    public visitSparseUnion     <T extends type.SparseUnion>     (type: T, { length, nullCount } = this.nextFieldNode()) { return           Data.Union(type, 0, length, nullCount, this.readNullBitmap(type, nullCount), this.readTypeIds(type), this.visitMany(type.children));                         }
73    public visitDictionary      <T extends type.Dictionary>      (type: T, { length, nullCount } = this.nextFieldNode()) { return      Data.Dictionary(type, 0, length, nullCount, this.readNullBitmap(type, nullCount), this.readData(type.indices), this.readDictionary(type));                        }
74    public visitInterval        <T extends type.Interval>        (type: T, { length, nullCount } = this.nextFieldNode()) { return        Data.Interval(type, 0, length, nullCount, this.readNullBitmap(type, nullCount), this.readData(type));                                                           }
75    public visitFixedSizeList   <T extends type.FixedSizeList>   (type: T, { length, nullCount } = this.nextFieldNode()) { return   Data.FixedSizeList(type, 0, length, nullCount, this.readNullBitmap(type, nullCount), this.visit(type.children[0]));                                                  }
76    public visitMap             <T extends type.Map_>            (type: T, { length, nullCount } = this.nextFieldNode()) { return             Data.Map(type, 0, length, nullCount, this.readNullBitmap(type, nullCount), this.readOffsets(type), this.visit(type.children[0]));                          }
77
78    protected nextFieldNode() { return this.nodes[++this.nodesIndex]; }
79    protected nextBufferRange() { return this.buffers[++this.buffersIndex]; }
80    protected readNullBitmap<T extends DataType>(type: T, nullCount: number, buffer = this.nextBufferRange()) {
81        return nullCount > 0 && this.readData(type, buffer) || new Uint8Array(0);
82    }
83    protected readOffsets<T extends DataType>(type: T, buffer?: BufferRegion) { return this.readData(type, buffer); }
84    protected readTypeIds<T extends DataType>(type: T, buffer?: BufferRegion) { return this.readData(type, buffer); }
85    protected readData<T extends DataType>(_type: T, { length, offset } = this.nextBufferRange()) {
86        return this.bytes.subarray(offset, offset + length);
87    }
88    protected readDictionary<T extends type.Dictionary>(type: T): Vector<T['dictionary']> {
89        return this.dictionaries.get(type.id)!;
90    }
91}
92
93/** @ignore */
94export class JSONVectorLoader extends VectorLoader {
95    private sources: any[][];
96    constructor(sources: any[][], nodes: FieldNode[], buffers: BufferRegion[], dictionaries: Map<number, Vector<any>>) {
97        super(new Uint8Array(0), nodes, buffers, dictionaries);
98        this.sources = sources;
99    }
100    protected readNullBitmap<T extends DataType>(_type: T, nullCount: number, { offset } = this.nextBufferRange()) {
101        return nullCount <= 0 ? new Uint8Array(0) : packBools(this.sources[offset]);
102    }
103    protected readOffsets<T extends DataType>(_type: T, { offset } = this.nextBufferRange()) {
104        return toArrayBufferView(Uint8Array, toArrayBufferView(Int32Array, this.sources[offset]));
105    }
106    protected readTypeIds<T extends DataType>(type: T, { offset } = this.nextBufferRange()) {
107        return toArrayBufferView(Uint8Array, toArrayBufferView(type.ArrayType, this.sources[offset]));
108    }
109    protected readData<T extends DataType>(type: T, { offset } = this.nextBufferRange()) {
110        const { sources } = this;
111        if (DataType.isTimestamp(type)) {
112            return toArrayBufferView(Uint8Array, Int64.convertArray(sources[offset] as string[]));
113        } else if ((DataType.isInt(type) || DataType.isTime(type)) && type.bitWidth === 64) {
114            return toArrayBufferView(Uint8Array, Int64.convertArray(sources[offset] as string[]));
115        } else if (DataType.isDate(type) && type.unit === DateUnit.MILLISECOND) {
116            return toArrayBufferView(Uint8Array, Int64.convertArray(sources[offset] as string[]));
117        } else if (DataType.isDecimal(type)) {
118            return toArrayBufferView(Uint8Array, Int128.convertArray(sources[offset] as string[]));
119        } else if (DataType.isBinary(type) || DataType.isFixedSizeBinary(type)) {
120            return binaryDataFromJSON(sources[offset] as string[]);
121        } else if (DataType.isBool(type)) {
122            return packBools(sources[offset] as number[]);
123        } else if (DataType.isUtf8(type)) {
124            return encodeUtf8((sources[offset] as string[]).join(''));
125        }
126        return toArrayBufferView(Uint8Array, toArrayBufferView(type.ArrayType, sources[offset].map((x) => +x)));
127    }
128}
129
130/** @ignore */
131function binaryDataFromJSON(values: string[]) {
132    // "DATA": ["49BC7D5B6C47D2","3F5FB6D9322026"]
133    // There are definitely more efficient ways to do this... but it gets the
134    // job done.
135    const joined = values.join('');
136    const data = new Uint8Array(joined.length / 2);
137    for (let i = 0; i < joined.length; i += 2) {
138        data[i >> 1] = parseInt(joined.substr(i, 2), 16);
139    }
140    return data;
141}
142