fix writing empty tables

trxcllnt · trxcllnt · commit ea1c8db63191 · 2019-02-26T10:19:21.000-08:00
diff --git a/js/src/ipc/writer.ts b/js/src/ipc/writer.ts
@@ -33,7 +33,7 @@ import { JSONTypeAssembler } from '../visitor/jsontypeassembler';
 import { JSONVectorAssembler } from '../visitor/jsonvectorassembler';
 import { ArrayBufferViewInput, toUint8Array } from '../util/buffer';
 import { Writable, ReadableInterop, ReadableDOMStreamOptions } from '../io/interfaces';
-import { isPromise, isAsyncIterable, isWritableDOMStream, isWritableNodeStream } from '../util/compat';
+import { isPromise, isAsyncIterable, isWritableDOMStream, isWritableNodeStream, isIterable } from '../util/compat';
 
 export class RecordBatchWriter<T extends { [key: string]: DataType } = any> extends ReadableInterop<Uint8Array> implements Writable<RecordBatch<T>> {
 
@@ -140,21 +140,36 @@ export class RecordBatchWriter<T extends { [key: string]: DataType } = any> exte
         return this;
     }
 
-    public write(chunk?: Table<T> | RecordBatch<T> | null) {
-        let schema: Schema<T> | null;
+    public write(payload?: Table<T> | RecordBatch<T> | Iterable<RecordBatch<T>> | null) {
+
+        let schema: Schema<T> | null = null;
+
         if (!this._sink) {
             throw new Error(`RecordBatchWriter is closed`);
-        } else if (!chunk || !(schema = chunk.schema)) {
+        } else if (payload === null || payload === undefined) {
+            return this.finish() && undefined;
+        } else if (payload instanceof Table && !(schema = payload.schema)) {
             return this.finish() && undefined;
-        } else if (schema !== this._schema) {
+        } else if (payload instanceof RecordBatch && !(schema = payload.schema)) {
+            return this.finish() && undefined;
+        }
+
+        if (schema && !schema.compareTo(this._schema)) {
             if (this._started && this._autoDestroy) {
                 return this.close();
             }
             this.reset(this._sink, schema);
         }
-        (chunk instanceof Table)
-            ? this.writeAll(chunk.chunks)
-            : this._writeRecordBatch(chunk);
+
+        if (payload instanceof RecordBatch) {
+            if (payload.length > 0) {
+                this._writeRecordBatch(payload);
+            }
+        } else if (payload instanceof Table) {
+            this.writeAll(payload.chunks);
+        } else if (isIterable(payload)) {
+            this.writeAll(payload);
+        }
     }
 
     protected _writeMessage<T extends MessageHeader>(message: Message<T>, alignment = 8) {
@@ -363,7 +378,11 @@ export class RecordBatchJSONWriter<T extends { [key: string]: DataType } = any>
 
 /** @ignore */
 function writeAll<T extends { [key: string]: DataType } = any>(writer: RecordBatchWriter<T>, input: Table<T> | Iterable<RecordBatch<T>>) {
-    const chunks = (input instanceof Table) ? input.chunks : input;
+    let chunks = input as Iterable<RecordBatch<T>>;
+    if (input instanceof Table) {
+        chunks = input.chunks;
+        writer.reset(undefined, input.schema);
+    }
     for (const batch of chunks) {
         writer.write(batch);
     }
diff --git a/js/src/table.ts b/js/src/table.ts
@@ -179,9 +179,7 @@ export class Table<T extends { [key: string]: DataType } = any>
             throw new TypeError('Table must be initialized with a Schema or at least one RecordBatch');
         }
 
-        if (!chunks[0]) { chunks[0] = new RecordBatch(schema, 0, []); }
-
-        super(chunks[0].type, chunks);
+        super(new Struct<T>(schema.fields), chunks);
 
         this._schema = schema;
         this._chunks = chunks;
diff --git a/js/src/visitor/vectorassembler.ts b/js/src/visitor/vectorassembler.ts
@@ -59,7 +59,10 @@ export class VectorAssembler extends Visitor {
 
     /** @nocollapse */
     public static assemble<T extends Vector | RecordBatch>(...args: (T | T[])[]) {
-        return new VectorAssembler().visitMany(selectVectorChildrenArgs(RecordBatch, args))[0];
+        const assembler = new VectorAssembler();
+        const vectorChildren = selectVectorChildrenArgs(RecordBatch, args);
+        const [assembleResult = assembler] = assembler.visitMany(vectorChildren);
+        return assembleResult;
     }
 
     private constructor() { super(); }
diff --git a/js/test/unit/table/serialize-tests.ts b/js/test/unit/table/serialize-tests.ts
@@ -30,6 +30,15 @@ function createTable<T extends { [key: string]: DataType } = any>(schema: Schema
 }
 
 describe('Table#serialize()', () => {
+
+    test(`Table#empty round-trips through serialization`, () => {
+        const source = Table.empty();
+        expect(source.length).toBe(0);
+        expect(source.numCols).toBe(0);
+        const result = Table.from(source.serialize());
+        expect(result).toEqualTable(source);
+    });
+
     const chunkLengths = [] as number[];
     for (let i = -1; ++i < 3;) {
         chunkLengths[i] = (Math.random() * 100) | 0;

Original file line number	Diff line number	Diff line change
`@@ -179,9 +179,7 @@ export class Table<T extends { [key: string]: DataType } = any>`
`179`	`179`	`throw new TypeError('Table must be initialized with a Schema or at least one RecordBatch');`
`180`	`180`	`}`
`181`	`181`
`182`		`- if (!chunks[0]) { chunks[0] = new RecordBatch(schema, 0, []); }`
`183`		`-`
`184`		`- super(chunks[0].type, chunks);`
	`182`	`+ super(new Struct<T>(schema.fields), chunks);`
`185`	`183`
`186`	`184`	`this._schema = schema;`
`187`	`185`	`this._chunks = chunks;`