@useairfoil/flight 0.3.0 → 0.4.1

package/src/record-batch-decoder.ts

@@ -45,12 +45,11 @@ import {
 import * as metadata from "apache-arrow/ipc/metadata/message";
 import { _InternalEmptyPlaceholderRecordBatch } from "apache-arrow/recordbatch";
 import { bigIntToNumber } from "apache-arrow/util/bigint";
-import {
-  CompressedVectorLoader,
-  VectorLoader,
-} from "apache-arrow/visitor/vectorloader";
+import { CompressedVectorLoader, VectorLoader } from "apache-arrow/visitor/vectorloader";
 import * as flatbuffers from "flatbuffers";
+
 import type { FlightData } from "./proto/Flight";
+import type { RecordBatchWithMetadata } from "./record-batch-with-metadata";
 
 const invalidMessageType = (type: MessageHeader) =>
   `Expected ${MessageHeader[type]} Message in stream, but was null or length 0.`;
@@ -61,10 +60,8 @@ const invalidMessageType = (type: MessageHeader) =>
 // const invalidMessageBodyLength = (expected: number, actual: number) =>
 //   `Expected to read ${expected} bytes for message body, but only read ${actual}.`;
 
-abstract class RecordBatchReaderImpl<T extends TypeMap = any>
-  implements RecordBatchReaderImpl<T>
-{
-  public declare schema: Schema<T>;
+abstract class RecordBatchReaderImpl<T extends TypeMap = any> implements RecordBatchReaderImpl<T> {
+  declare public schema: Schema<T>;
   public closed = false;
   public autoDestroy = true;
   public dictionaries: Map<number, Vector>;
@@ -105,10 +102,7 @@ abstract class RecordBatchReaderImpl<T extends TypeMap = any>
     return this;
   }
 
-  protected _loadRecordBatch(
-    header: metadata.RecordBatch,
-    body: Uint8Array,
-  ): RecordBatch<T> {
+  protected _loadRecordBatch(header: metadata.RecordBatch, body: Uint8Array): RecordBatch<T> {
     // Ensure the body buffer is properly aligned for Apache Arrow
     // Apache Arrow requires 8-byte aligned buffers for typed arrays
     const alignedBody = this._ensureAlignedBuffer(body);
@@ -117,22 +111,9 @@ abstract class RecordBatchReaderImpl<T extends TypeMap = any>
     if (header.compression != null) {
       const codec = compressionRegistry.get(header.compression.type);
       if (codec?.decode && typeof codec.decode === "function") {
-        const { decommpressedBody, buffers } = this._decompressBuffers(
-          header,
-          alignedBody,
-          codec,
-        );
-        children = this._loadCompressedVectors(
-          header,
-          decommpressedBody,
-          this.schema.fields,
-        );
-        header = new metadata.RecordBatch(
-          header.length,
-          header.nodes,
-          buffers,
-          null,
-        );
+        const { decommpressedBody, buffers } = this._decompressBuffers(header, alignedBody, codec);
+        children = this._loadCompressedVectors(header, decommpressedBody, this.schema.fields);
+        header = new metadata.RecordBatch(header.length, header.nodes, buffers, null);
       } else {
         throw new Error("Record batch is compressed but codec not found");
       }
@@ -164,10 +145,7 @@ abstract class RecordBatchReaderImpl<T extends TypeMap = any>
     return alignedBuffer;
   }
 
-  protected _loadDictionaryBatch(
-    header: metadata.DictionaryBatch,
-    body: Uint8Array,
-  ) {
+  protected _loadDictionaryBatch(header: metadata.DictionaryBatch, body: Uint8Array) {
     const { id, isDelta } = header;
     const { dictionaries, schema } = this;
     const dictionary = dictionaries.get(id);
@@ -176,21 +154,10 @@ abstract class RecordBatchReaderImpl<T extends TypeMap = any>
     if (header.data.compression != null) {
       const codec = compressionRegistry.get(header.data.compression.type);
       if (codec?.decode && typeof codec.decode === "function") {
-        const { decommpressedBody, buffers } = this._decompressBuffers(
-          header.data,
-          body,
-          codec,
-        );
-        data = this._loadCompressedVectors(header.data, decommpressedBody, [
-          type,
-        ]);
+        const { decommpressedBody, buffers } = this._decompressBuffers(header.data, body, codec);
+        data = this._loadCompressedVectors(header.data, decommpressedBody, [type]);
         header = new metadata.DictionaryBatch(
-          new metadata.RecordBatch(
-            header.data.length,
-            header.data.nodes,
-            buffers,
-            null,
-          ),
+          new metadata.RecordBatch(header.data.length, header.data.nodes, buffers, null),
           id,
           isDelta,
         );
@@ -202,9 +169,7 @@ abstract class RecordBatchReaderImpl<T extends TypeMap = any>
     }
     // const data = this._loadVectors(header.data, body, [type]);
     return (
-      dictionary && isDelta
-        ? dictionary.concat(new Vector(data))
-        : new Vector(data)
+      dictionary && isDelta ? dictionary.concat(new Vector(data)) : new Vector(data)
     ).memoize() as Vector;
   }
 
@@ -251,25 +216,19 @@ abstract class RecordBatchReaderImpl<T extends TypeMap = any>
         newBufferRegions.push(new metadata.BufferRegion(currentOffset, 0));
         continue;
       }
-      const byteBuf = new flatbuffers.ByteBuffer(
-        body.subarray(offset, offset + length),
-      );
+      const byteBuf = new flatbuffers.ByteBuffer(body.subarray(offset, offset + length));
       const uncompressedLenth = bigIntToNumber(byteBuf.readInt64(0));
 
       const bytes = byteBuf.bytes().subarray(COMPRESS_LENGTH_PREFIX);
 
       const decompressed =
-        uncompressedLenth === LENGTH_NO_COMPRESSED_DATA
-          ? bytes
-          : codec.decode!(bytes);
+        uncompressedLenth === LENGTH_NO_COMPRESSED_DATA ? bytes : codec.decode!(bytes);
 
       decompressedBuffers.push(decompressed);
 
       const padding = ((currentOffset + 7) & ~7) - currentOffset;
       currentOffset += padding;
-      newBufferRegions.push(
-        new metadata.BufferRegion(currentOffset, decompressed.length),
-      );
+      newBufferRegions.push(new metadata.BufferRegion(currentOffset, decompressed.length));
       currentOffset += decompressed.length;
     }
 
@@ -282,7 +241,7 @@ abstract class RecordBatchReaderImpl<T extends TypeMap = any>
 
 export class RecordBatchStreamReaderFromFlightData<T extends TypeMap = any>
   extends RecordBatchReaderImpl<T>
-  implements AsyncIterableIterator<RecordBatch<T>>
+  implements AsyncIterableIterator<RecordBatchWithMetadata<T>>
 {
   protected _reader: AsyncIterator<FlightData>;
 
@@ -303,8 +262,8 @@ export class RecordBatchStreamReaderFromFlightData<T extends TypeMap = any>
   //   return true;
   // }
 
-  public [Symbol.asyncIterator](): AsyncIterableIterator<RecordBatch<T>> {
-    return this as AsyncIterableIterator<RecordBatch<T>>;
+  public [Symbol.asyncIterator](): AsyncIterableIterator<RecordBatchWithMetadata<T>> {
+    return this as AsyncIterableIterator<RecordBatchWithMetadata<T>>;
   }
 
   /*
@@ -362,7 +321,13 @@ export class RecordBatchStreamReaderFromFlightData<T extends TypeMap = any>
         this._recordBatchIndex++;
         const header = message.header();
         const recordBatch = this._loadRecordBatch(header, flight.dataBody);
-        return { done: false, value: recordBatch };
+        return {
+          done: false,
+          value: {
+            batch: recordBatch,
+            appMetadata: flight.appMetadata,
+          },
+        };
       } else if (message.isDictionaryBatch()) {
         this._dictionaryIndex++;
         const header = message.header();
@@ -375,7 +340,10 @@ export class RecordBatchStreamReaderFromFlightData<T extends TypeMap = any>
       this._recordBatchIndex++;
       return {
         done: false,
-        value: new _InternalEmptyPlaceholderRecordBatch<T>(this.schema),
+        value: {
+          batch: new _InternalEmptyPlaceholderRecordBatch<T>(this.schema),
+          appMetadata: new Uint8Array(0),
+        },
       };
     }
 
@@ -383,9 +351,7 @@ export class RecordBatchStreamReaderFromFlightData<T extends TypeMap = any>
     return ITERATOR_DONE;
   }
 
-  protected async _readNextMessageAndValidate<T extends MessageHeader>(
-    type?: T | null,
-  ) {
+  protected async _readNextMessageAndValidate<T extends MessageHeader>(type?: T | null) {
     const { done, value } = await this._reader.next();
     if (done) {
       return null;

package/src/record-batch-with-metadata.ts

@@ -0,0 +1,6 @@
+import type { RecordBatch, TypeMap } from "apache-arrow";
+
+export type RecordBatchWithMetadata<T extends TypeMap = any> = {
+  readonly batch: RecordBatch<T>;
+  readonly appMetadata: Uint8Array;
+};