Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

GH-39251: [JS] Use resizable buffer in builder #39252

Merged
merged 5 commits into from
Dec 27, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion js/src/builder.ts
Original file line number Diff line number Diff line change
Expand Up @@ -342,7 +342,7 @@ export abstract class Builder<T extends DataType = any, TNull = any> {
export abstract class FixedWidthBuilder<T extends Int | Float | FixedSizeBinary | Date_ | Timestamp | Time | Decimal | Interval | Duration = any, TNull = any> extends Builder<T, TNull> {
constructor(opts: BuilderOptions<T, TNull>) {
super(opts);
this._values = new DataBufferBuilder(new this.ArrayType(0), this.stride);
this._values = new DataBufferBuilder(this.ArrayType, 0, this.stride);
}
public setValue(index: number, value: T['TValue']) {
const values = this._values;
Expand Down
4 changes: 2 additions & 2 deletions js/src/builder/binary.ts
Original file line number Diff line number Diff line change
Expand Up @@ -16,15 +16,15 @@
// under the License.

import { Binary } from '../type.js';
import { toUint8Array } from '../util/buffer.js';
import { BufferBuilder } from './buffer.js';
import { VariableWidthBuilder, BuilderOptions } from '../builder.js';
import { toUint8Array } from '../util/buffer.js';

/** @ignore */
export class BinaryBuilder<TNull = any> extends VariableWidthBuilder<Binary, TNull> {
constructor(opts: BuilderOptions<Binary, TNull>) {
super(opts);
this._values = new BufferBuilder(new Uint8Array(0));
this._values = new BufferBuilder(Uint8Array);
}
public get byteLength(): number {
let size = this._pendingLength + (this.length * 4);
Expand Down
44 changes: 30 additions & 14 deletions js/src/builder/buffer.ts
Original file line number Diff line number Diff line change
Expand Up @@ -24,20 +24,36 @@ function roundLengthUpToNearest64Bytes(len: number, BPE: number) {
const bytesMinus1 = Math.ceil(len) * BPE - 1;
return ((bytesMinus1 - bytesMinus1 % 64 + 64) || 64) / BPE;
}

/** @ignore */
const sliceOrExtendArray = <T extends TypedArray | BigIntArray>(arr: T, len = 0) => (
arr.length >= len ? arr.subarray(0, len) : memcpy(new (arr.constructor as any)(len), arr, 0)
) as T;
function resizeArray<T extends TypedArray | BigIntArray>(arr: T, len = 0): T {
// TODO: remove when https://github.com/microsoft/TypeScript/issues/54636 is fixed
const buffer = arr.buffer as ArrayBufferLike & { resizable: boolean; resize: (byteLength: number) => void; maxByteLength: number };
const byteLength = len * arr.BYTES_PER_ELEMENT;
if (buffer.resizable && byteLength <= buffer.maxByteLength) {
buffer.resize(byteLength);
return arr;
}

// Fallback for non-resizable buffers
return arr.length >= len ?
arr.subarray(0, len) as T :
memcpy(new (arr.constructor as any)(len), arr, 0);
}

/** @ignore */
export const SAFE_ARRAY_SIZE = 2 ** 32 - 1;

/** @ignore */
export class BufferBuilder<T extends TypedArray | BigIntArray> {

constructor(buffer: T, stride = 1) {
this.buffer = buffer;
constructor(bufferType: ArrayCtor<T>, initialSize = 0, stride = 1) {
this.length = Math.ceil(initialSize / stride);
// TODO: remove as any when https://github.com/microsoft/TypeScript/issues/54636 is fixed
this.buffer = new bufferType(new (ArrayBuffer as any)(this.length * bufferType.BYTES_PER_ELEMENT, { maxByteLength: SAFE_ARRAY_SIZE })) as T;
this.stride = stride;
this.BYTES_PER_ELEMENT = buffer.BYTES_PER_ELEMENT;
this.ArrayType = buffer.constructor as ArrayCtor<T>;
this._resize(this.length = Math.ceil(buffer.length / stride));
this.BYTES_PER_ELEMENT = bufferType.BYTES_PER_ELEMENT;
this.ArrayType = bufferType;
}

public buffer: T;
Expand Down Expand Up @@ -72,17 +88,18 @@ export class BufferBuilder<T extends TypedArray | BigIntArray> {
}
public flush(length = this.length) {
length = roundLengthUpToNearest64Bytes(length * this.stride, this.BYTES_PER_ELEMENT);
const array = sliceOrExtendArray<T>(this.buffer, length);
const array = resizeArray<T>(this.buffer, length);
this.clear();
return array;
}
public clear() {
this.length = 0;
this._resize(0);
// TODO: remove as any when https://github.com/microsoft/TypeScript/issues/54636 is fixed
this.buffer = new this.ArrayType(new (ArrayBuffer as any)(0, { maxByteLength: SAFE_ARRAY_SIZE })) as T;
return this;
}
protected _resize(newLength: number) {
return this.buffer = <T>memcpy(new this.ArrayType(newLength), this.buffer);
return this.buffer = resizeArray<T>(this.buffer, newLength);
}
}

Expand All @@ -100,7 +117,7 @@ export class DataBufferBuilder<T extends TypedArray | BigIntArray> extends Buffe
/** @ignore */
export class BitmapBufferBuilder extends DataBufferBuilder<Uint8Array> {

constructor(data = new Uint8Array(0)) { super(data, 1 / 8); }
constructor() { super(Uint8Array, 0, 1 / 8); }

public numValid = 0;
public get numInvalid() { return this.length - this.numValid; }
Expand All @@ -123,9 +140,8 @@ export class BitmapBufferBuilder extends DataBufferBuilder<Uint8Array> {
/** @ignore */
export class OffsetsBufferBuilder<T extends DataType> extends DataBufferBuilder<T['TOffsetArray']> {
constructor(type: T) {
super(new type.OffsetArrayType(1), 1);
super(type.OffsetArrayType as ArrayCtor<T['TOffsetArray']>, 1, 1);
}

public append(value: T['TOffsetArray'][0]) {
return this.set(this.length - 1, value);
}
Expand Down
2 changes: 1 addition & 1 deletion js/src/builder/largeutf8.ts
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ import { VariableWidthBuilder, BuilderOptions } from '../builder.js';
export class LargeUtf8Builder<TNull = any> extends VariableWidthBuilder<LargeUtf8, TNull> {
constructor(opts: BuilderOptions<LargeUtf8, TNull>) {
super(opts);
this._values = new BufferBuilder(new Uint8Array(0));
this._values = new BufferBuilder(Uint8Array);
}
public get byteLength(): number {
let size = this._pendingLength + (this.length * 4);
Expand Down
4 changes: 2 additions & 2 deletions js/src/builder/union.ts
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ export abstract class UnionBuilder<T extends Union, TNull = any> extends Builder

constructor(options: UnionBuilderOptions<T, TNull>) {
super(options);
this._typeIds = new DataBufferBuilder(new Int8Array(0), 1);
this._typeIds = new DataBufferBuilder(Int8Array, 0, 1);
if (typeof options['valueToChildTypeId'] === 'function') {
this._valueToChildTypeId = options['valueToChildTypeId'];
}
Expand Down Expand Up @@ -84,7 +84,7 @@ export class DenseUnionBuilder<T extends DenseUnion, TNull = any> extends UnionB

constructor(options: UnionBuilderOptions<T, TNull>) {
super(options);
this._offsets = new DataBufferBuilder(new Int32Array(0));
this._offsets = new DataBufferBuilder(Int32Array);
}

/** @ignore */
Expand Down
2 changes: 1 addition & 1 deletion js/src/builder/utf8.ts
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ import { VariableWidthBuilder, BuilderOptions } from '../builder.js';
export class Utf8Builder<TNull = any> extends VariableWidthBuilder<Utf8, TNull> {
constructor(opts: BuilderOptions<Utf8, TNull>) {
super(opts);
this._values = new BufferBuilder(new Uint8Array(0));
this._values = new BufferBuilder(Uint8Array);
}
public get byteLength(): number {
let size = this._pendingLength + (this.length * 4);
Expand Down
Loading