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 { Type } from '../enum';
20import { Visitor } from '../visitor';
21import { VectorType } from '../interfaces';
22import { BitIterator } from '../util/bit';
23import { instance as getVisitor } from './get';
24import {
25    DataType, Dictionary,
26    Bool, Null, Utf8, Binary, Decimal, FixedSizeBinary, List, FixedSizeList, Map_, Struct,
27    Float, Float16, Float32, Float64,
28    Int, Uint8, Uint16, Uint32, Uint64, Int8, Int16, Int32, Int64,
29    Date_, DateDay, DateMillisecond,
30    Interval, IntervalDayTime, IntervalYearMonth,
31    Time, TimeSecond, TimeMillisecond, TimeMicrosecond, TimeNanosecond,
32    Timestamp, TimestampSecond, TimestampMillisecond, TimestampMicrosecond, TimestampNanosecond,
33    Union, DenseUnion, SparseUnion,
34} from '../type';
35
36/** @ignore */
37export interface IteratorVisitor extends Visitor {
38    visit<T extends VectorType>(node: T): IterableIterator<T['TValue'] | null>;
39    visitMany <T extends VectorType>(nodes: T[]): IterableIterator<T['TValue'] | null>[];
40    getVisitFn<T extends Type>(node: T): (vector: VectorType<T>) => IterableIterator<VectorType<T>['TValue'] | null>;
41    getVisitFn<T extends DataType>(node: VectorType<T> | Data<T> | T): (vector: VectorType<T>) => IterableIterator<VectorType<T>['TValue'] | null>;
42    visitNull                 <T extends Null>                 (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
43    visitBool                 <T extends Bool>                 (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
44    visitInt                  <T extends Int>                  (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
45    visitInt8                 <T extends Int8>                 (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
46    visitInt16                <T extends Int16>                (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
47    visitInt32                <T extends Int32>                (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
48    visitInt64                <T extends Int64>                (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
49    visitUint8                <T extends Uint8>                (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
50    visitUint16               <T extends Uint16>               (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
51    visitUint32               <T extends Uint32>               (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
52    visitUint64               <T extends Uint64>               (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
53    visitFloat                <T extends Float>                (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
54    visitFloat16              <T extends Float16>              (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
55    visitFloat32              <T extends Float32>              (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
56    visitFloat64              <T extends Float64>              (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
57    visitUtf8                 <T extends Utf8>                 (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
58    visitBinary               <T extends Binary>               (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
59    visitFixedSizeBinary      <T extends FixedSizeBinary>      (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
60    visitDate                 <T extends Date_>                (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
61    visitDateDay              <T extends DateDay>              (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
62    visitDateMillisecond      <T extends DateMillisecond>      (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
63    visitTimestamp            <T extends Timestamp>            (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
64    visitTimestampSecond      <T extends TimestampSecond>      (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
65    visitTimestampMillisecond <T extends TimestampMillisecond> (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
66    visitTimestampMicrosecond <T extends TimestampMicrosecond> (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
67    visitTimestampNanosecond  <T extends TimestampNanosecond>  (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
68    visitTime                 <T extends Time>                 (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
69    visitTimeSecond           <T extends TimeSecond>           (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
70    visitTimeMillisecond      <T extends TimeMillisecond>      (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
71    visitTimeMicrosecond      <T extends TimeMicrosecond>      (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
72    visitTimeNanosecond       <T extends TimeNanosecond>       (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
73    visitDecimal              <T extends Decimal>              (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
74    visitList                 <T extends List>                 (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
75    visitStruct               <T extends Struct>               (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
76    visitUnion                <T extends Union>                (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
77    visitDenseUnion           <T extends DenseUnion>           (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
78    visitSparseUnion          <T extends SparseUnion>          (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
79    visitDictionary           <T extends Dictionary>           (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
80    visitInterval             <T extends Interval>             (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
81    visitIntervalDayTime      <T extends IntervalDayTime>      (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
82    visitIntervalYearMonth    <T extends IntervalYearMonth>    (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
83    visitFixedSizeList        <T extends FixedSizeList>        (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
84    visitMap                  <T extends Map_>                 (vector: VectorType<T>): IterableIterator<T['TValue'] | null>;
85}
86
87/** @ignore */
88export class IteratorVisitor extends Visitor {}
89
90/** @ignore */
91function nullableIterator<T extends DataType>(vector: VectorType<T>): IterableIterator<T['TValue'] | null> {
92    const getFn = getVisitor.getVisitFn(vector);
93    return new BitIterator<T['TValue'] | null>(
94        vector.data.nullBitmap, vector.data.offset, vector.length, vector,
95        (vec: VectorType<T>, idx: number, nullByte: number, nullBit: number) =>
96            ((nullByte & 1 << nullBit) !== 0) ? getFn(vec, idx) : null
97    );
98}
99
100/** @ignore */
101class VectorIterator<T extends DataType> implements IterableIterator<T['TValue'] | null> {
102    private index = 0;
103
104    constructor(
105        private vector: VectorType<T>,
106        private getFn: (vector: VectorType<T>, index: number) => VectorType<T>['TValue']
107    ) {}
108
109    next(): IteratorResult<T['TValue'] | null> {
110        if (this.index < this.vector.length) {
111            return {
112                value: this.getFn(this.vector, this.index++)
113            };
114        }
115
116        return {done: true, value: null};
117    }
118
119    [Symbol.iterator]() {
120        return this;
121    }
122}
123
124/** @ignore */
125function vectorIterator<T extends DataType>(vector: VectorType<T>): IterableIterator<T['TValue'] | null> {
126
127    // If nullable, iterate manually
128    if (vector.nullCount > 0) {
129        return nullableIterator<T>(vector);
130    }
131
132    const { type, typeId, length } = vector;
133
134    // Fast case, defer to native iterators if possible
135    if (vector.stride === 1 && (
136        (typeId === Type.Timestamp) ||
137        (typeId === Type.Int && (type as Int).bitWidth !== 64) ||
138        (typeId === Type.Time && (type as Time).bitWidth !== 64) ||
139        (typeId === Type.Float && (type as Float).precision > 0 /* Precision.HALF */)
140    )) {
141        return vector.data.values.subarray(0, length)[Symbol.iterator]();
142    }
143
144    // Otherwise, iterate manually
145    return new VectorIterator(vector, getVisitor.getVisitFn(vector));
146}
147
148IteratorVisitor.prototype.visitNull                 = vectorIterator;
149IteratorVisitor.prototype.visitBool                 = vectorIterator;
150IteratorVisitor.prototype.visitInt                  = vectorIterator;
151IteratorVisitor.prototype.visitInt8                 = vectorIterator;
152IteratorVisitor.prototype.visitInt16                = vectorIterator;
153IteratorVisitor.prototype.visitInt32                = vectorIterator;
154IteratorVisitor.prototype.visitInt64                = vectorIterator;
155IteratorVisitor.prototype.visitUint8                = vectorIterator;
156IteratorVisitor.prototype.visitUint16               = vectorIterator;
157IteratorVisitor.prototype.visitUint32               = vectorIterator;
158IteratorVisitor.prototype.visitUint64               = vectorIterator;
159IteratorVisitor.prototype.visitFloat                = vectorIterator;
160IteratorVisitor.prototype.visitFloat16              = vectorIterator;
161IteratorVisitor.prototype.visitFloat32              = vectorIterator;
162IteratorVisitor.prototype.visitFloat64              = vectorIterator;
163IteratorVisitor.prototype.visitUtf8                 = vectorIterator;
164IteratorVisitor.prototype.visitBinary               = vectorIterator;
165IteratorVisitor.prototype.visitFixedSizeBinary      = vectorIterator;
166IteratorVisitor.prototype.visitDate                 = vectorIterator;
167IteratorVisitor.prototype.visitDateDay              = vectorIterator;
168IteratorVisitor.prototype.visitDateMillisecond      = vectorIterator;
169IteratorVisitor.prototype.visitTimestamp            = vectorIterator;
170IteratorVisitor.prototype.visitTimestampSecond      = vectorIterator;
171IteratorVisitor.prototype.visitTimestampMillisecond = vectorIterator;
172IteratorVisitor.prototype.visitTimestampMicrosecond = vectorIterator;
173IteratorVisitor.prototype.visitTimestampNanosecond  = vectorIterator;
174IteratorVisitor.prototype.visitTime                 = vectorIterator;
175IteratorVisitor.prototype.visitTimeSecond           = vectorIterator;
176IteratorVisitor.prototype.visitTimeMillisecond      = vectorIterator;
177IteratorVisitor.prototype.visitTimeMicrosecond      = vectorIterator;
178IteratorVisitor.prototype.visitTimeNanosecond       = vectorIterator;
179IteratorVisitor.prototype.visitDecimal              = vectorIterator;
180IteratorVisitor.prototype.visitList                 = vectorIterator;
181IteratorVisitor.prototype.visitStruct               = vectorIterator;
182IteratorVisitor.prototype.visitUnion                = vectorIterator;
183IteratorVisitor.prototype.visitDenseUnion           = vectorIterator;
184IteratorVisitor.prototype.visitSparseUnion          = vectorIterator;
185IteratorVisitor.prototype.visitDictionary           = vectorIterator;
186IteratorVisitor.prototype.visitInterval             = vectorIterator;
187IteratorVisitor.prototype.visitIntervalDayTime      = vectorIterator;
188IteratorVisitor.prototype.visitIntervalYearMonth    = vectorIterator;
189IteratorVisitor.prototype.visitFixedSizeList        = vectorIterator;
190IteratorVisitor.prototype.visitMap                  = vectorIterator;
191
192/** @ignore */
193export const instance = new IteratorVisitor();
194