Skip to content

Commit e81082f

Browse files
trxcllnttrxcllnt
committed
export vector views, allow cloning data as another type
1 parent 700a47c commit e81082f

File tree

7 files changed

+160
-34
lines changed

7 files changed

+160
-34
lines changed

‎js/src/Arrow.externs.js

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -374,6 +374,12 @@ ChunkedData.computeOffsets = function() {};
374374
var FlatVector = function() {};
375375
/** @type {?} */
376376
FlatVector.prototype.values;
377+
/** @type {?} */
378+
FlatVector.prototype.lows;
379+
/** @type {?} */
380+
FlatVector.prototype.highs;
381+
/** @type {?} */
382+
FlatVector.prototype.asInt32;
377383

378384
var ListVectorBase = function() {};
379385
/** @type {?} */

‎js/src/Arrow.ts

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import * as data_ from './data';
2020
import * as vector_ from './vector';
2121
import * as util_ from './util/int';
2222
import * as visitor_ from './visitor';
23+
import * as view_ from './vector/view';
2324
import { Vector } from './vector';
2425
import { RecordBatch } from './recordbatch';
2526
import { Schema, Field, Type } from './type';
@@ -122,6 +123,32 @@ export namespace visitor {
122123
export import VectorVisitor = visitor_.VectorVisitor;
123124
}
124125

126+
export namespace view {
127+
export import ChunkedView = view_.ChunkedView;
128+
export import DictionaryView = view_.DictionaryView;
129+
export import ListView = view_.ListView;
130+
export import FixedSizeListView = view_.FixedSizeListView;
131+
export import BinaryView = view_.BinaryView;
132+
export import Utf8View = view_.Utf8View;
133+
export import UnionView = view_.UnionView;
134+
export import DenseUnionView = view_.DenseUnionView;
135+
export import NestedView = view_.NestedView;
136+
export import StructView = view_.StructView;
137+
export import MapView = view_.MapView;
138+
export import FlatView = view_.FlatView;
139+
export import NullView = view_.NullView;
140+
export import BoolView = view_.BoolView;
141+
export import ValidityView = view_.ValidityView;
142+
export import FixedSizeView = view_.FixedSizeView;
143+
export import Float16View = view_.Float16View;
144+
export import DateDayView = view_.DateDayView;
145+
export import DateMillisecondView = view_.DateMillisecondView;
146+
export import IntervalYearMonthView = view_.IntervalYearMonthView;
147+
export import IntervalYearView = view_.IntervalYearView;
148+
export import IntervalMonthView = view_.IntervalMonthView;
149+
export import PrimitiveView = view_.PrimitiveView;
150+
}
151+
125152
/* These exports are needed for the closure and uglify umd targets */
126153
try {
127154
let Arrow: any = eval('exports');
@@ -130,6 +157,7 @@ try {
130157
Arrow['data'] = data;
131158
Arrow['type'] = type;
132159
Arrow['util'] = util;
160+
Arrow['view'] = view;
133161
Arrow['vector'] = vector;
134162
Arrow['visitor'] = visitor;
135163

‎js/src/data.ts

Lines changed: 46 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -90,13 +90,13 @@ export class BaseData<T extends DataType = DataType> implements VectorLike {
9090
}
9191
return nullCount;
9292
}
93-
public clone(length = this._length, offset = this._offset, nullCount = this._nullCount) {
94-
return new BaseData<T>(this._type, length, offset, nullCount) as this;
93+
public clone<R extends T>(type: R, length = this._length, offset = this._offset, nullCount = this._nullCount) {
94+
return new BaseData(type, length, offset, nullCount);
9595
}
9696
public slice(offset: number, length: number) {
9797
return length <= 0 ? this : this.sliceInternal(this.clone(
98-
length, this._offset + offset, +(this._nullCount === 0) - 1
99-
), offset, length);
98+
this._type, length, this._offset + offset, +(this._nullCount === 0) - 1
99+
) as any, offset, length);
100100
}
101101
protected sliceInternal(clone: this, offset: number, length: number) {
102102
let arr: any;
@@ -126,16 +126,13 @@ export class FlatData<T extends FlatType> extends BaseData<T> {
126126
this[VectorType.VALIDITY] = toTypedArray(Uint8Array, nullBitmap);
127127
}
128128
public get ArrayType(): T['ArrayType'] { return this._type.ArrayType; }
129-
public clone(length = this._length, offset = this._offset, nullCount = this._nullCount) {
130-
return new FlatData<T>(this._type, length, this[VectorType.VALIDITY], this[VectorType.DATA], offset, nullCount) as this;
129+
public clone<R extends T>(type: R, length = this._length, offset = this._offset, nullCount = this._nullCount) {
130+
return new (this.constructor as any)(type, length, this[VectorType.VALIDITY], this[VectorType.DATA], offset, nullCount) as FlatData<R>;
131131
}
132132
}
133133

134134
export class BoolData extends FlatData<Bool> {
135135
protected sliceData(data: Uint8Array) { return data; }
136-
public clone(length = this._length, offset = this._offset, nullCount = this._nullCount) {
137-
return new BoolData(this._type, length, this[VectorType.VALIDITY], this[VectorType.DATA], offset, nullCount) as this;
138-
}
139136
}
140137

141138
export class FlatListData<T extends FlatListType> extends FlatData<T> {
@@ -148,8 +145,8 @@ export class FlatListData<T extends FlatListType> extends FlatData<T> {
148145
super(type, length, nullBitmap, data, offset, nullCount);
149146
this[VectorType.OFFSET] = toTypedArray(Int32Array, valueOffsets);
150147
}
151-
public clone(length = this._length, offset = this._offset, nullCount = this._nullCount) {
152-
return new FlatListData<T>(this._type, length, this[VectorType.VALIDITY], this[VectorType.OFFSET], this[VectorType.DATA], offset, nullCount) as this;
148+
public clone<R extends T>(type: R, length = this._length, offset = this._offset, nullCount = this._nullCount) {
149+
return new FlatListData(type, length, this[VectorType.VALIDITY], this[VectorType.OFFSET], this[VectorType.DATA], offset, nullCount);
153150
}
154151
}
155152

@@ -165,8 +162,13 @@ export class DictionaryData<T extends DataType> extends BaseData<Dictionary<T>>
165162
}
166163
public get length() { return this._indicies.length; }
167164
public get nullCount() { return this._indicies.nullCount; }
168-
public clone(length = this._length, offset = this._offset) {
169-
return new DictionaryData<T>(this._type, this._dictionary, this._indicies.slice(offset - this._offset, length)) as this;
165+
public clone<R extends Dictionary<T>>(type: R, length = this._length, offset = this._offset) {
166+
const data = this._dictionary.data.clone(type.dictionary as any);
167+
return new DictionaryData<R>(
168+
this._type as any,
169+
this._dictionary.clone(data) as any,
170+
this._indicies.slice(offset - this._offset, length)
171+
) as any;
170172
}
171173
protected sliceInternal(clone: this, _offset: number, _length: number) {
172174
clone._length = clone._indicies.length;
@@ -182,8 +184,8 @@ export class NestedData<T extends NestedType = NestedType> extends BaseData<T> {
182184
this._childData = childData;
183185
this[VectorType.VALIDITY] = toTypedArray(Uint8Array, nullBitmap);
184186
}
185-
public clone(length = this._length, offset = this._offset, nullCount = this._nullCount) {
186-
return new NestedData<T>(this._type, length, this[VectorType.VALIDITY], this._childData, offset, nullCount) as this;
187+
public clone<R extends T>(type: R, length = this._length, offset = this._offset, nullCount = this._nullCount) {
188+
return new NestedData<R>(type, length, this[VectorType.VALIDITY], this._childData, offset, nullCount);
187189
}
188190
protected sliceInternal(clone: this, offset: number, length: number) {
189191
if (!this[VectorType.OFFSET]) {
@@ -204,8 +206,8 @@ export class ListData<T extends ListType> extends NestedData<T> {
204206
this._valuesData = valueChildData;
205207
this[VectorType.OFFSET] = toTypedArray(Int32Array, valueOffsets);
206208
}
207-
public clone(length = this._length, offset = this._offset, nullCount = this._nullCount) {
208-
return new ListData<T>(this._type, length, this[VectorType.VALIDITY], this[VectorType.OFFSET], this._valuesData, offset, nullCount) as this;
209+
public clone<R extends T>(type: R, length = this._length, offset = this._offset, nullCount = this._nullCount) {
210+
return new ListData<R>(type, length, this[VectorType.VALIDITY], this[VectorType.OFFSET], this._valuesData as any, offset, nullCount);
209211
}
210212
}
211213

@@ -216,17 +218,24 @@ export class UnionData<T extends (DenseUnion | SparseUnion) = any> extends Neste
216218
super(type, length, nullBitmap, childData, offset, nullCount);
217219
this[VectorType.TYPE] = toTypedArray(Int8Array, typeIds);
218220
}
219-
public clone(length = this._length, offset = this._offset, nullCount = this._nullCount) {
220-
return new UnionData<T>(this._type, length, this[VectorType.VALIDITY], this[VectorType.TYPE], this._childData, offset, nullCount) as this;
221+
public clone<R extends T>(type: R, length = this._length, offset = this._offset, nullCount = this._nullCount) {
222+
return new UnionData<R>(type, length, this[VectorType.VALIDITY], this[VectorType.TYPE], this._childData, offset, nullCount);
221223
}
222224
}
223225

224226
export class SparseUnionData extends UnionData<SparseUnion> {
225227
constructor(type: SparseUnion, length: number, nullBitmap: Uint8Array | null | undefined, typeIds: Iterable<number>, childData: Data<any>[], offset?: number, nullCount?: number) {
226228
super(type, length, nullBitmap, typeIds, childData, offset, nullCount);
227229
}
228-
public clone(length = this._length, offset = this._offset, nullCount = this._nullCount) {
229-
return new SparseUnionData(this._type, length, this[VectorType.VALIDITY], this[VectorType.TYPE], this._childData, offset, nullCount) as this;
230+
public clone<R extends SparseUnion>(type: R, length = this._length, offset = this._offset, nullCount = this._nullCount) {
231+
return new SparseUnionData(
232+
type,
233+
length,
234+
this[VectorType.VALIDITY],
235+
this[VectorType.TYPE],
236+
this._childData,
237+
offset, nullCount
238+
) as any as UnionData<R>;
230239
}
231240
}
232241

@@ -237,8 +246,16 @@ export class DenseUnionData extends UnionData<DenseUnion> {
237246
super(type, length, nullBitmap, typeIds, childData, offset, nullCount);
238247
this[VectorType.OFFSET] = toTypedArray(Int32Array, valueOffsets);
239248
}
240-
public clone(length = this._length, offset = this._offset, nullCount = this._nullCount) {
241-
return new DenseUnionData(this._type, length, this[VectorType.VALIDITY], this[VectorType.TYPE], this[VectorType.OFFSET], this._childData, offset, nullCount) as this;
249+
public clone<R extends DenseUnion>(type: R, length = this._length, offset = this._offset, nullCount = this._nullCount) {
250+
return new DenseUnionData(
251+
type,
252+
length,
253+
this[VectorType.VALIDITY],
254+
this[VectorType.TYPE],
255+
this[VectorType.OFFSET],
256+
this._childData,
257+
offset, nullCount
258+
) as any as UnionData<R>;
242259
}
243260
}
244261

@@ -263,8 +280,12 @@ export class ChunkedData<T extends DataType> extends BaseData<T> {
263280
}
264281
return nullCount;
265282
}
266-
public clone(length = this._length, offset = this._offset, nullCount = this._nullCount) {
267-
return new ChunkedData<T>(this._type, length, this._childVectors, offset, nullCount, this._childOffsets) as this;
283+
public clone<R extends T>(type: R, length = this._length, offset = this._offset, nullCount = this._nullCount) {
284+
return new ChunkedData<R>(
285+
type, length,
286+
this._childVectors.map((vec) => vec.clone(vec.data.clone(type))) as any,
287+
offset, nullCount, this._childOffsets
288+
);
268289
}
269290
protected sliceInternal(clone: this, offset: number, length: number) {
270291
const chunks = this._childVectors;

‎js/src/recordbatch.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -64,8 +64,8 @@ export class RecordBatch extends StructVector {
6464
this.numCols = schema.fields.length;
6565
}
6666
}
67-
public clone(data: Data<Struct>, view: View<Struct> = this.view.clone(data)): this {
68-
return new RecordBatch(this.schema, data, view) as this;
67+
public clone<R extends Struct>(data: Data<R>, view: View<R> = this.view.clone(data)): this {
68+
return new RecordBatch(this.schema, data as any, view) as any;
6969
}
7070
public select(...columnNames: string[]) {
7171
const fields = this.schema.fields;

‎js/src/vector.ts

Lines changed: 48 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18-
import { Data, ChunkedData, FlatData, BoolData } from './data';
18+
import { Data, ChunkedData, FlatData, BoolData, FlatListData, NestedData } from './data';
1919
import { VisitorNode, TypeVisitor, VectorVisitor } from './visitor';
2020
import { DataType, ListType, FlatType, NestedType, FlatListType } from './type';
2121
import { IterableArrayLike, Precision, DateUnit, IntervalUnit, UnionMode } from './type';
@@ -35,10 +35,14 @@ export class Vector<T extends DataType = any> implements VectorLike, View<T>, Vi
3535
public static create<T extends DataType>(data: Data<T>): Vector<T> {
3636
return createVector(data);
3737
}
38+
public type: T;
39+
public length: number;
3840
public readonly data: Data<T>;
3941
public readonly view: View<T>;
4042
constructor(data: Data<T>, view: View<T>) {
4143
this.data = data;
44+
this.type = data.type;
45+
this.length = data.length;
4246
let nulls: Uint8Array;
4347
if ((<any> data instanceof ChunkedData) && !(view instanceof ChunkedView)) {
4448
this.view = new ChunkedView(data);
@@ -49,15 +53,13 @@ export class Vector<T extends DataType = any> implements VectorLike, View<T>, Vi
4953
}
5054
}
5155

52-
public get type() { return this.data.type; }
53-
public get length() { return this.data.length; }
5456
public get nullCount() { return this.data.nullCount; }
5557
public get nullBitmap() { return this.data.nullBitmap; }
5658
public get [Symbol.toStringTag]() {
5759
return `Vector<${this.type[Symbol.toStringTag]}>`;
5860
}
5961
public toJSON(): any { return this.toArray(); }
60-
public clone(data: Data<T>, view: View<T> = this.view.clone(data)): this {
62+
public clone<R extends T>(data: Data<R>, view: View<R> = this.view.clone(data) as any): this {
6163
return new (this.constructor as any)(data, view);
6264
}
6365
public isValid(index: number): boolean {
@@ -111,6 +113,17 @@ export class Vector<T extends DataType = any> implements VectorLike, View<T>, Vi
111113

112114
export abstract class FlatVector<T extends FlatType> extends Vector<T> {
113115
public get values() { return this.data.values; }
116+
public lows(): IntVector<Int32> { return this.asInt32(0, 2); }
117+
public highs(): IntVector<Int32> { return this.asInt32(1, 2); }
118+
public asInt32(offset: number = 0, stride: number = 2): IntVector<Int32> {
119+
let data = (this.data as FlatData<any>).clone(new Int32());
120+
if (offset > 0) {
121+
data = data.slice(offset, this.length - offset);
122+
}
123+
const int32s = new IntVector(data, new PrimitiveView(data, stride));
124+
int32s.length = this.length / stride | 0;
125+
return int32s;
126+
}
114127
}
115128

116129
export abstract class ListVectorBase<T extends (ListType | FlatListType)> extends Vector<T> {
@@ -144,7 +157,7 @@ import { ChunkedView } from './vector/chunked';
144157
import { DictionaryView } from './vector/dictionary';
145158
import { ListView, FixedSizeListView, BinaryView, Utf8View } from './vector/list';
146159
import { UnionView, DenseUnionView, NestedView, StructView, MapView } from './vector/nested';
147-
import { FlatView, NullView, BoolView, ValidityView, FixedSizeView, Float16View, DateDayView, DateMillisecondView, IntervalYearMonthView } from './vector/flat';
160+
import { FlatView, NullView, BoolView, ValidityView, FixedSizeView, Float16View, DateDayView, DateMillisecondView, IntervalYearMonthView, PrimitiveView } from './vector/flat';
148161
import { packBools } from './util/bit';
149162

150163
export class NullVector extends Vector<Null> {
@@ -223,6 +236,12 @@ export class DateVector extends FlatVector<Date_> {
223236
constructor(data: Data<Date_>, view: View<Date_> = DateVector.defaultView(data)) {
224237
super(data, view);
225238
}
239+
public lows(): IntVector<Int32> {
240+
return this.type.unit === DateUnit.DAY ? this.asInt32(0, 1) : this.asInt32(0, 2);
241+
}
242+
public highs(): IntVector<Int32> {
243+
return this.type.unit === DateUnit.DAY ? this.asInt32(0, 1) : this.asInt32(1, 2);
244+
}
226245
}
227246

228247
export class DecimalVector extends FlatVector<Decimal> {
@@ -238,6 +257,12 @@ export class TimeVector extends FlatVector<Time> {
238257
constructor(data: Data<Time>, view: View<Time> = TimeVector.defaultView(data)) {
239258
super(data, view);
240259
}
260+
public lows(): IntVector<Int32> {
261+
return this.type.bitWidth <= 32 ? this.asInt32(0, 1) : this.asInt32(0, 2);
262+
}
263+
public highs(): IntVector<Int32> {
264+
return this.type.bitWidth <= 32 ? this.asInt32(0, 1) : this.asInt32(1, 2);
265+
}
241266
}
242267

243268
export class TimestampVector extends FlatVector<Timestamp> {
@@ -253,12 +278,21 @@ export class IntervalVector extends FlatVector<Interval> {
253278
constructor(data: Data<Interval>, view: View<Interval> = IntervalVector.defaultView(data)) {
254279
super(data, view);
255280
}
281+
public lows(): IntVector<Int32> {
282+
return this.type.unit === IntervalUnit.YEAR_MONTH ? this.asInt32(0, 1) : this.asInt32(0, 2);
283+
}
284+
public highs(): IntVector<Int32> {
285+
return this.type.unit === IntervalUnit.YEAR_MONTH ? this.asInt32(0, 1) : this.asInt32(1, 2);
286+
}
256287
}
257288

258289
export class BinaryVector extends ListVectorBase<Binary> {
259290
constructor(data: Data<Binary>, view: View<Binary> = new BinaryView(data)) {
260291
super(data, view);
261292
}
293+
public asUtf8() {
294+
return new Utf8Vector((this.data as FlatListData<any>).clone(new Utf8()));
295+
}
262296
}
263297

264298
export class FixedSizeBinaryVector extends FlatVector<FixedSizeBinary> {
@@ -271,6 +305,9 @@ export class Utf8Vector extends ListVectorBase<Utf8> {
271305
constructor(data: Data<Utf8>, view: View<Utf8> = new Utf8View(data)) {
272306
super(data, view);
273307
}
308+
public asBinary() {
309+
return new BinaryVector((this.data as FlatListData<any>).clone(new Binary()));
310+
}
274311
}
275312

276313
export class ListVector<T extends DataType = DataType> extends ListVectorBase<List<T>> {
@@ -289,12 +326,18 @@ export class MapVector extends NestedVector<Map_> {
289326
constructor(data: Data<Map_>, view: View<Map_> = new MapView(data)) {
290327
super(data, view);
291328
}
329+
public asStruct() {
330+
return new StructVector((this.data as NestedData<any>).clone(new Struct(this.type.children)));
331+
}
292332
}
293333

294334
export class StructVector extends NestedVector<Struct> {
295335
constructor(data: Data<Struct>, view: View<Struct> = new StructView(data)) {
296336
super(data, view);
297337
}
338+
public asMap(keysSorted: boolean = false) {
339+
return new MapVector((this.data as NestedData<any>).clone(new Map_(keysSorted, this.type.children)));
340+
}
298341
}
299342

300343
export class UnionVector<T extends (SparseUnion | DenseUnion) = any> extends NestedVector<T> {

0 commit comments

Comments
 (0)