Skip to content

Commit

Permalink
GH-39255: [JS] Allow customization of schema when passing vectors to …
Browse files Browse the repository at this point in the history
…table constructor (#39256)

Merge after #39254.

* Closes: #39255
  • Loading branch information
domoritz committed Jan 3, 2024
1 parent 984eb38 commit 3d1324e
Show file tree
Hide file tree
Showing 3 changed files with 22 additions and 3 deletions.
2 changes: 1 addition & 1 deletion js/src/builder/largebinary.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 LargeBinaryBuilder<TNull = any> extends VariableWidthBuilder<LargeBinary, TNull> {
constructor(opts: BuilderOptions<LargeBinary, 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
6 changes: 4 additions & 2 deletions js/src/table.ts
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,8 @@ export class Table<T extends TypeMap = any> {
constructor(...batches: readonly RecordBatch<T>[]);
constructor(...columns: { [P in keyof T]: Vector<T[P]> }[]);
constructor(...columns: { [P in keyof T]: Data<T[P]> | DataProps<T[P]> }[]);
constructor(schema: Schema<T>, ...columns: { [P in keyof T]: Vector<T[P]> }[]);
constructor(schema: Schema<T>, ...columns: { [P in keyof T]: Data<T[P]> | DataProps<T[P]> }[]);
constructor(schema: Schema<T>, data?: RecordBatch<T> | RecordBatch<T>[]);
constructor(schema: Schema<T>, data?: RecordBatch<T> | RecordBatch<T>[], offsets?: Uint32Array);
constructor(...args: any[]) {
Expand Down Expand Up @@ -112,8 +114,8 @@ export class Table<T extends TypeMap = any> {
} else if (typeof x === 'object') {
const keys = Object.keys(x) as (keyof T)[];
const vecs = keys.map((k) => new Vector([x[k]]));
const schema = new Schema(keys.map((k, i) => new Field(String(k), vecs[i].type, vecs[i].nullCount > 0)));
const [, batches] = distributeVectorsIntoRecordBatches(schema, vecs);
const batchSchema = schema ?? new Schema(keys.map((k, i) => new Field(String(k), vecs[i].type, vecs[i].nullCount > 0)));
const [, batches] = distributeVectorsIntoRecordBatches(batchSchema, vecs);
return batches.length === 0 ? [new RecordBatch(x)] : batches;
}
}
Expand Down
17 changes: 17 additions & 0 deletions js/test/unit/table-tests.ts
Original file line number Diff line number Diff line change
Expand Up @@ -151,6 +151,23 @@ describe(`Table`, () => {
expect(i32).toEqualVector(makeVector(i32s));
});

test(`creates a new Table from a Typed Array and force nullable`, () => {
const i32s = new Int32Array(arange(new Array<number>(10)));
const i32 = makeVector([i32s]);
expect(i32).toHaveLength(i32s.length);
expect(i32.nullCount).toBe(0);

const table = new Table(new Schema([new Field('i32', new Int32, true)]), { i32 });
const i32Field = table.schema.fields[0];

expect(i32Field.name).toBe('i32');
expect(i32).toHaveLength(i32s.length);
expect(i32Field.nullable).toBe(true);
expect(i32.nullCount).toBe(0);

expect(i32).toEqualVector(makeVector(i32s));
});

test(`creates a new Table from Typed Arrays`, () => {
const i32s = new Int32Array(arange(new Array<number>(10)));
const f32s = new Float32Array(arange(new Array<number>(10)));
Expand Down

0 comments on commit 3d1324e

Please sign in to comment.