| // Licensed to the Apache Software Foundation (ASF) under one |
| // or more contributor license agreements. See the NOTICE file |
| // distributed with this work for additional information |
| // regarding copyright ownership. The ASF licenses this file |
| // to you under the Apache License, Version 2.0 (the |
| // "License"); you may not use this file except in compliance |
| // with the License. You may obtain a copy of the License at |
| // |
| // http://www.apache.org/licenses/LICENSE-2.0 |
| // |
| // Unless required by applicable law or agreed to in writing, |
| // software distributed under the License is distributed on an |
| // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| // KIND, either express or implied. See the License for the |
| // specific language governing permissions and limitations |
| // under the License. |
| |
| import * as flatbuffers from 'flatbuffers'; |
| |
| import { Schema as _Schema } from '../../fb/schema.js'; |
| import { Int as _Int } from '../../fb/int.js'; |
| import { RecordBatch as _RecordBatch } from '../../fb/record-batch.js'; |
| import { DictionaryBatch as _DictionaryBatch } from '../../fb/dictionary-batch.js'; |
| import { Buffer as _Buffer } from '../../fb/buffer.js'; |
| import { Field as _Field } from '../../fb/field.js'; |
| import { FieldNode as _FieldNode } from '../../fb/field-node.js'; |
| import { DictionaryEncoding as _DictionaryEncoding } from '../../fb/dictionary-encoding.js'; |
| import { Type } from '../../fb/type.js'; |
| import { KeyValue as _KeyValue } from '../../fb/key-value.js'; |
| import { Endianness as _Endianness } from '../../fb/endianness.js'; |
| import { FloatingPoint as _FloatingPoint } from '../../fb/floating-point.js'; |
| import { Decimal as _Decimal } from '../../fb/decimal.js'; |
| import { Date as _Date } from '../../fb/date.js'; |
| import { Time as _Time } from '../../fb/time.js'; |
| import { Timestamp as _Timestamp } from '../../fb/timestamp.js'; |
| import { Interval as _Interval } from '../../fb/interval.js'; |
| import { Duration as _Duration } from '../../fb/duration.js'; |
| import { Union as _Union } from '../../fb/union.js'; |
| import { FixedSizeBinary as _FixedSizeBinary } from '../../fb/fixed-size-binary.js'; |
| import { FixedSizeList as _FixedSizeList } from '../../fb/fixed-size-list.js'; |
| import { Map as _Map } from '../../fb/map.js'; |
| import { Message as _Message } from '../../fb/message.js'; |
| import { CompressionType as _CompressionType } from '../../fb/compression-type.js'; |
| import { BodyCompression as _BodyCompression } from '../../fb/body-compression.js'; |
| import { BodyCompressionMethod as _BodyCompressionMethod } from '../../fb/body-compression-method.js'; |
| |
| import { Schema, Field } from '../../schema.js'; |
| import { toUint8Array } from '../../util/buffer.js'; |
| import { ArrayBufferViewInput } from '../../util/buffer.js'; |
| import { bigIntToNumber } from '../../util/bigint.js'; |
| import { MessageHeader, MetadataVersion } from '../../enum.js'; |
| import { instance as typeAssembler } from '../../visitor/typeassembler.js'; |
| import { fieldFromJSON, schemaFromJSON, recordBatchFromJSON, dictionaryBatchFromJSON } from './json.js'; |
| |
| import Builder = flatbuffers.Builder; |
| import ByteBuffer = flatbuffers.ByteBuffer; |
| |
| import { |
| DataType, Dictionary, TimeBitWidth, |
| Utf8, LargeUtf8, Binary, LargeBinary, BinaryView, Utf8View, Decimal, FixedSizeBinary, |
| List, LargeList, FixedSizeList, Map_, Struct, Union, |
| Bool, Null, Int, Float, Date_, Time, Interval, Timestamp, IntBitWidth, Int32, TKeys, Duration, |
| } from '../../type.js'; |
| |
| /** |
| * @ignore |
| * @private |
| **/ |
| export class Message<T extends MessageHeader = any> { |
| |
| /** @nocollapse */ |
| public static fromJSON<T extends MessageHeader>(msg: any, headerType: T): Message<T> { |
| const message = new Message(0, MetadataVersion.V5, headerType); |
| message._createHeader = messageHeaderFromJSON(msg, headerType); |
| return message; |
| } |
| |
| /** @nocollapse */ |
| public static decode(buf: ArrayBufferViewInput) { |
| buf = new ByteBuffer(toUint8Array(buf)); |
| const _message = _Message.getRootAsMessage(buf); |
| const bodyLength: bigint = _message.bodyLength()!; |
| const version: MetadataVersion = _message.version(); |
| const headerType: MessageHeader = _message.headerType(); |
| const metadata = decodeMessageCustomMetadata(_message); |
| const message = new Message(bodyLength, version, headerType, undefined, metadata); |
| message._createHeader = decodeMessageHeader(_message, headerType); |
| return message; |
| } |
| |
| /** @nocollapse */ |
| public static encode<T extends MessageHeader>(message: Message<T>) { |
| const b = new Builder(); |
| let headerOffset = -1; |
| if (message.isSchema()) { |
| headerOffset = Schema.encode(b, message.header() as Schema); |
| } else if (message.isRecordBatch()) { |
| headerOffset = RecordBatch.encode(b, message.header() as RecordBatch); |
| } else if (message.isDictionaryBatch()) { |
| headerOffset = DictionaryBatch.encode(b, message.header() as DictionaryBatch); |
| } |
| |
| // Encode custom metadata if present (must be done before startMessage) |
| const customMetadataOffset = !(message.metadata && message.metadata.size > 0) ? -1 : |
| _Message.createCustomMetadataVector(b, [...message.metadata].map(([k, v]) => { |
| const key = b.createString(`${k}`); |
| const val = b.createString(`${v}`); |
| _KeyValue.startKeyValue(b); |
| _KeyValue.addKey(b, key); |
| _KeyValue.addValue(b, val); |
| return _KeyValue.endKeyValue(b); |
| })); |
| |
| _Message.startMessage(b); |
| _Message.addVersion(b, MetadataVersion.V5); |
| _Message.addHeader(b, headerOffset); |
| _Message.addHeaderType(b, message.headerType); |
| _Message.addBodyLength(b, BigInt(message.bodyLength)); |
| if (customMetadataOffset !== -1) { _Message.addCustomMetadata(b, customMetadataOffset); } |
| _Message.finishMessageBuffer(b, _Message.endMessage(b)); |
| return b.asUint8Array(); |
| } |
| |
| /** @nocollapse */ |
| public static from(header: Schema | RecordBatch | DictionaryBatch, bodyLength = 0) { |
| if (header instanceof Schema) { |
| return new Message(0, MetadataVersion.V5, MessageHeader.Schema, header); |
| } |
| if (header instanceof RecordBatch) { |
| return new Message(bodyLength, MetadataVersion.V5, MessageHeader.RecordBatch, header, header.metadata); |
| } |
| if (header instanceof DictionaryBatch) { |
| return new Message(bodyLength, MetadataVersion.V5, MessageHeader.DictionaryBatch, header); |
| } |
| throw new Error(`Unrecognized Message header: ${header}`); |
| } |
| |
| public body: Uint8Array; |
| protected _headerType: T; |
| protected _bodyLength: number; |
| protected _version: MetadataVersion; |
| protected _compression: BodyCompression | null; |
| protected _metadata: Map<string, string>; |
| public get type() { return this.headerType; } |
| public get version() { return this._version; } |
| public get headerType() { return this._headerType; } |
| public get compression() { return this._compression; } |
| public get bodyLength() { return this._bodyLength; } |
| public get metadata() { return this._metadata; } |
| declare protected _createHeader: MessageHeaderDecoder; |
| public header() { return this._createHeader<T>(); } |
| public isSchema(): this is Message<MessageHeader.Schema> { return this.headerType === MessageHeader.Schema; } |
| public isRecordBatch(): this is Message<MessageHeader.RecordBatch> { return this.headerType === MessageHeader.RecordBatch; } |
| public isDictionaryBatch(): this is Message<MessageHeader.DictionaryBatch> { return this.headerType === MessageHeader.DictionaryBatch; } |
| |
| constructor(bodyLength: bigint | number, version: MetadataVersion, headerType: T, header?: any, metadata?: Map<string, string>) { |
| this._version = version; |
| this._headerType = headerType; |
| this.body = new Uint8Array(0); |
| this._compression = header?.compression; |
| header && (this._createHeader = () => header); |
| this._bodyLength = bigIntToNumber(bodyLength); |
| this._metadata = metadata || new Map(); |
| } |
| } |
| |
| /** |
| * @ignore |
| * @private |
| **/ |
| export class RecordBatch { |
| protected _length: number; |
| protected _nodes: FieldNode[]; |
| protected _buffers: BufferRegion[]; |
| protected _compression: BodyCompression | null; |
| protected _variadicBufferCounts: number[]; |
| protected _metadata: Map<string, string>; |
| public get nodes() { return this._nodes; } |
| public get length() { return this._length; } |
| public get buffers() { return this._buffers; } |
| public get compression() { return this._compression; } |
| public get variadicBufferCounts() { return this._variadicBufferCounts; } |
| public get metadata() { return this._metadata; } |
| constructor( |
| length: bigint | number, |
| nodes: FieldNode[], |
| buffers: BufferRegion[], |
| compression: BodyCompression | null, |
| variadicBufferCounts: number[] = [], |
| metadata?: Map<string, string> |
| ) { |
| this._nodes = nodes; |
| this._buffers = buffers; |
| this._length = bigIntToNumber(length); |
| this._compression = compression; |
| this._variadicBufferCounts = variadicBufferCounts; |
| this._metadata = metadata || new Map(); |
| } |
| } |
| |
| /** |
| * @ignore |
| * @private |
| **/ |
| export class DictionaryBatch { |
| |
| protected _id: number; |
| protected _isDelta: boolean; |
| protected _data: RecordBatch; |
| public get id() { return this._id; } |
| public get data() { return this._data; } |
| public get isDelta() { return this._isDelta; } |
| public get length(): number { return this.data.length; } |
| public get nodes(): FieldNode[] { return this.data.nodes; } |
| public get buffers(): BufferRegion[] { return this.data.buffers; } |
| |
| constructor(data: RecordBatch, id: bigint | number, isDelta = false) { |
| this._data = data; |
| this._isDelta = isDelta; |
| this._id = bigIntToNumber(id); |
| } |
| } |
| |
| /** |
| * @ignore |
| * @private |
| **/ |
| export class BufferRegion { |
| public offset: number; |
| public length: number; |
| constructor(offset: bigint | number, length: bigint | number) { |
| this.offset = bigIntToNumber(offset); |
| this.length = bigIntToNumber(length); |
| } |
| } |
| |
| /** |
| * @ignore |
| * @private |
| **/ |
| export class FieldNode { |
| public length: number; |
| public nullCount: number; |
| constructor(length: bigint | number, nullCount: bigint | number) { |
| this.length = bigIntToNumber(length); |
| this.nullCount = bigIntToNumber(nullCount); |
| } |
| } |
| |
| /** |
| * @ignore |
| * @private |
| **/ |
| export class BodyCompression { |
| public type: _CompressionType; |
| public method: _BodyCompressionMethod; |
| constructor(type: _CompressionType, method: _BodyCompressionMethod = _BodyCompressionMethod.BUFFER) { |
| this.type = type; |
| this.method = method; |
| } |
| } |
| |
| /** @ignore */ |
| function messageHeaderFromJSON(message: any, type: MessageHeader) { |
| return (() => { |
| switch (type) { |
| case MessageHeader.Schema: return Schema.fromJSON(message); |
| case MessageHeader.RecordBatch: return RecordBatch.fromJSON(message); |
| case MessageHeader.DictionaryBatch: return DictionaryBatch.fromJSON(message); |
| } |
| throw new Error(`Unrecognized Message type: { name: ${MessageHeader[type]}, type: ${type} }`); |
| }) as MessageHeaderDecoder; |
| } |
| |
| /** @ignore */ |
| function decodeMessageHeader(message: _Message, type: MessageHeader) { |
| return (() => { |
| switch (type) { |
| case MessageHeader.Schema: return Schema.decode(message.header(new _Schema())!, new Map(), message.version()); |
| case MessageHeader.RecordBatch: return RecordBatch.decode(message.header(new _RecordBatch())!, message.version()); |
| case MessageHeader.DictionaryBatch: return DictionaryBatch.decode(message.header(new _DictionaryBatch())!, message.version()); |
| } |
| throw new Error(`Unrecognized Message type: { name: ${MessageHeader[type]}, type: ${type} }`); |
| }) as MessageHeaderDecoder; |
| } |
| |
| Field['encode'] = encodeField; |
| Field['decode'] = decodeField; |
| Field['fromJSON'] = fieldFromJSON; |
| |
| Schema['encode'] = encodeSchema; |
| Schema['decode'] = decodeSchema; |
| Schema['fromJSON'] = schemaFromJSON; |
| |
| RecordBatch['encode'] = encodeRecordBatch; |
| RecordBatch['decode'] = decodeRecordBatch; |
| RecordBatch['fromJSON'] = recordBatchFromJSON; |
| |
| DictionaryBatch['encode'] = encodeDictionaryBatch; |
| DictionaryBatch['decode'] = decodeDictionaryBatch; |
| DictionaryBatch['fromJSON'] = dictionaryBatchFromJSON; |
| |
| FieldNode['encode'] = encodeFieldNode; |
| FieldNode['decode'] = decodeFieldNode; |
| |
| BufferRegion['encode'] = encodeBufferRegion; |
| BufferRegion['decode'] = decodeBufferRegion; |
| |
| BodyCompression['encode'] = encodeBodyCompression; |
| BodyCompression['decode'] = decodeBodyCompression; |
| |
| declare module '../../schema' { |
| namespace Field { |
| export { encodeField as encode }; |
| export { decodeField as decode }; |
| export { fieldFromJSON as fromJSON }; |
| } |
| namespace Schema { |
| export { encodeSchema as encode }; |
| export { decodeSchema as decode }; |
| export { schemaFromJSON as fromJSON }; |
| } |
| } |
| |
| declare module './message' { |
| namespace RecordBatch { |
| export { encodeRecordBatch as encode }; |
| export { decodeRecordBatch as decode }; |
| export { recordBatchFromJSON as fromJSON }; |
| } |
| namespace DictionaryBatch { |
| export { encodeDictionaryBatch as encode }; |
| export { decodeDictionaryBatch as decode }; |
| export { dictionaryBatchFromJSON as fromJSON }; |
| } |
| namespace FieldNode { |
| export { encodeFieldNode as encode }; |
| export { decodeFieldNode as decode }; |
| } |
| namespace BufferRegion { |
| export { encodeBufferRegion as encode }; |
| export { decodeBufferRegion as decode }; |
| } |
| namespace BodyCompression { |
| export { encodeBodyCompression as encode }; |
| export { decodeBodyCompression as decode }; |
| } |
| } |
| |
| /** @ignore */ |
| function decodeSchema(_schema: _Schema, dictionaries: Map<number, DataType> = new Map(), version = MetadataVersion.V5) { |
| const fields = decodeSchemaFields(_schema, dictionaries); |
| return new Schema(fields, decodeCustomMetadata(_schema), dictionaries, version); |
| } |
| |
| /** @ignore */ |
| function decodeRecordBatch(batch: _RecordBatch, version = MetadataVersion.V5) { |
| const recordBatch = new RecordBatch( |
| batch.length(), |
| decodeFieldNodes(batch), |
| decodeBuffers(batch, version), |
| decodeBodyCompression(batch.compression()), |
| decodeVariadicBufferCounts(batch) |
| ); |
| return recordBatch; |
| } |
| |
| /** @ignore */ |
| function decodeDictionaryBatch(batch: _DictionaryBatch, version = MetadataVersion.V5) { |
| return new DictionaryBatch(RecordBatch.decode(batch.data()!, version), batch.id(), batch.isDelta()); |
| } |
| |
| /** @ignore */ |
| function decodeBufferRegion(b: _Buffer) { |
| return new BufferRegion(b.offset(), b.length()); |
| } |
| |
| /** @ignore */ |
| function decodeFieldNode(f: _FieldNode) { |
| return new FieldNode(f.length(), f.nullCount()); |
| } |
| |
| /** @ignore */ |
| function decodeFieldNodes(batch: _RecordBatch) { |
| const nodes = [] as FieldNode[]; |
| for (let f, i = -1, j = -1, n = batch.nodesLength(); ++i < n;) { |
| if (f = batch.nodes(i)) { |
| nodes[++j] = FieldNode.decode(f); |
| } |
| } |
| return nodes; |
| } |
| |
| /** @ignore */ |
| function decodeBuffers(batch: _RecordBatch, version: MetadataVersion) { |
| const bufferRegions = [] as BufferRegion[]; |
| for (let b, i = -1, j = -1, n = batch.buffersLength(); ++i < n;) { |
| if (b = batch.buffers(i)) { |
| // If this Arrow buffer was written before version 4, |
| // advance the buffer's bb_pos 8 bytes to skip past |
| // the now-removed page_id field |
| if (version < MetadataVersion.V4) { |
| b.bb_pos += (8 * (i + 1)); |
| } |
| bufferRegions[++j] = BufferRegion.decode(b); |
| } |
| } |
| return bufferRegions; |
| } |
| |
| /** @ignore */ |
| function decodeVariadicBufferCounts(batch: _RecordBatch) { |
| const counts = [] as number[]; |
| const length = Math.trunc(batch.variadicBufferCountsLength()); |
| for (let i = 0; i < length; ++i) { |
| counts.push(bigIntToNumber(batch.variadicBufferCounts(i)!)); |
| } |
| return counts; |
| } |
| |
| /** @ignore */ |
| function decodeSchemaFields(schema: _Schema, dictionaries?: Map<number, DataType>) { |
| const fields = [] as Field[]; |
| for (let f, i = -1, j = -1, n = schema.fieldsLength(); ++i < n;) { |
| if (f = schema.fields(i)) { |
| fields[++j] = Field.decode(f, dictionaries); |
| } |
| } |
| return fields; |
| } |
| |
| /** @ignore */ |
| function decodeFieldChildren(field: _Field, dictionaries?: Map<number, DataType>): Field[] { |
| const children = [] as Field[]; |
| for (let f, i = -1, j = -1, n = field.childrenLength(); ++i < n;) { |
| if (f = field.children(i)) { |
| children[++j] = Field.decode(f, dictionaries); |
| } |
| } |
| return children; |
| } |
| |
| /** @ignore */ |
| function decodeField(f: _Field, dictionaries?: Map<number, DataType>) { |
| |
| let id: number; |
| let field: Field | void; |
| let type: DataType<any>; |
| let keys: _Int | TKeys | null; |
| let dictType: Dictionary; |
| let dictMeta: _DictionaryEncoding | null; |
| |
| // If no dictionary encoding |
| if (!dictionaries || !(dictMeta = f.dictionary())) { |
| type = decodeFieldType(f, decodeFieldChildren(f, dictionaries)); |
| field = new Field(f.name()!, type, f.nullable(), decodeCustomMetadata(f)); |
| } |
| // If dictionary encoded and the first time we've seen this dictionary id, decode |
| // the data type and child fields, then wrap in a Dictionary type and insert the |
| // data type into the dictionary types map. |
| else if (!dictionaries.has(id = bigIntToNumber(dictMeta.id()))) { |
| // a dictionary index defaults to signed 32 bit int if unspecified |
| keys = (keys = dictMeta.indexType()) ? decodeIndexType(keys) as TKeys : new Int32(); |
| dictionaries.set(id, type = decodeFieldType(f, decodeFieldChildren(f, dictionaries))); |
| dictType = new Dictionary(type, keys, id, dictMeta.isOrdered()); |
| field = new Field(f.name()!, dictType, f.nullable(), decodeCustomMetadata(f)); |
| } |
| // If dictionary encoded, and have already seen this dictionary Id in the schema, then reuse the |
| // data type and wrap in a new Dictionary type and field. |
| else { |
| // a dictionary index defaults to signed 32 bit int if unspecified |
| keys = (keys = dictMeta.indexType()) ? decodeIndexType(keys) as TKeys : new Int32(); |
| dictType = new Dictionary(dictionaries.get(id)!, keys, id, dictMeta.isOrdered()); |
| field = new Field(f.name()!, dictType, f.nullable(), decodeCustomMetadata(f)); |
| } |
| return field || null; |
| } |
| |
| /** @ignore */ |
| function decodeCustomMetadata(parent?: _Schema | _Field | null) { |
| const data = new Map<string, string>(); |
| if (parent) { |
| for (let entry, key, i = -1, n = Math.trunc(parent.customMetadataLength()); ++i < n;) { |
| if ((entry = parent.customMetadata(i)) && (key = entry.key()) != null) { |
| data.set(key, entry.value()!); |
| } |
| } |
| } |
| return data; |
| } |
| |
| /** @ignore */ |
| function decodeMessageCustomMetadata(message: _Message) { |
| const data = new Map<string, string>(); |
| for (let entry, key, i = -1, n = Math.trunc(message.customMetadataLength()); ++i < n;) { |
| if ((entry = message.customMetadata(i)) && (key = entry.key()) != null) { |
| data.set(key, entry.value()!); |
| } |
| } |
| return data; |
| } |
| |
| /** @ignore */ |
| function decodeIndexType(_type: _Int) { |
| return new Int(_type.isSigned(), _type.bitWidth() as IntBitWidth); |
| } |
| |
| /** @ignore */ |
| function decodeFieldType(f: _Field, children?: Field[]): DataType<any> { |
| |
| const typeId = f.typeType(); |
| |
| switch (typeId) { |
| case Type['NONE']: return new Null(); |
| case Type['Null']: return new Null(); |
| case Type['Binary']: return new Binary(); |
| case Type['LargeBinary']: return new LargeBinary(); |
| case Type['BinaryView']: return new BinaryView(); |
| case Type['Utf8']: return new Utf8(); |
| case Type['LargeUtf8']: return new LargeUtf8(); |
| case Type['Utf8View']: return new Utf8View(); |
| case Type['Bool']: return new Bool(); |
| case Type['List']: return new List((children || [])[0]); |
| case Type['LargeList']: return new LargeList((children || [])[0]); |
| case Type['Struct_']: return new Struct(children || []); |
| } |
| |
| switch (typeId) { |
| case Type['Int']: { |
| const t = f.type(new _Int())!; |
| return new Int(t.isSigned(), t.bitWidth()); |
| } |
| case Type['FloatingPoint']: { |
| const t = f.type(new _FloatingPoint())!; |
| return new Float(t.precision()); |
| } |
| case Type['Decimal']: { |
| const t = f.type(new _Decimal())!; |
| return new Decimal(t.scale(), t.precision(), t.bitWidth()); |
| } |
| case Type['Date']: { |
| const t = f.type(new _Date())!; |
| return new Date_(t.unit()); |
| } |
| case Type['Time']: { |
| const t = f.type(new _Time())!; |
| return new Time(t.unit(), t.bitWidth() as TimeBitWidth); |
| } |
| case Type['Timestamp']: { |
| const t = f.type(new _Timestamp())!; |
| return new Timestamp(t.unit(), t.timezone()); |
| } |
| case Type['Interval']: { |
| const t = f.type(new _Interval())!; |
| return new Interval(t.unit()); |
| } |
| case Type['Duration']: { |
| const t = f.type(new _Duration())!; |
| return new Duration(t.unit()); |
| } |
| case Type['Union']: { |
| const t = f.type(new _Union())!; |
| return new Union(t.mode(), t.typeIdsArray() || [], children || []); |
| } |
| case Type['FixedSizeBinary']: { |
| const t = f.type(new _FixedSizeBinary())!; |
| return new FixedSizeBinary(t.byteWidth()); |
| } |
| case Type['FixedSizeList']: { |
| const t = f.type(new _FixedSizeList())!; |
| return new FixedSizeList(t.listSize(), (children || [])[0]); |
| } |
| case Type['Map']: { |
| const t = f.type(new _Map())!; |
| return new Map_((children || [])[0], t.keysSorted()); |
| } |
| } |
| throw new Error(`Unrecognized type: "${Type[typeId]}" (${typeId})`); |
| } |
| |
| /** @ignore */ |
| function decodeBodyCompression(b: _BodyCompression | null) { |
| return b ? new BodyCompression(b.codec(), b.method()) : null; |
| } |
| |
| /** @ignore */ |
| function encodeSchema(b: Builder, schema: Schema) { |
| |
| const fieldOffsets = schema.fields.map((f) => Field.encode(b, f)); |
| |
| _Schema.startFieldsVector(b, fieldOffsets.length); |
| |
| const fieldsVectorOffset = _Schema.createFieldsVector(b, fieldOffsets); |
| |
| const metadataOffset = !(schema.metadata && schema.metadata.size > 0) ? -1 : |
| _Schema.createCustomMetadataVector(b, [...schema.metadata].map(([k, v]) => { |
| const key = b.createString(`${k}`); |
| const val = b.createString(`${v}`); |
| _KeyValue.startKeyValue(b); |
| _KeyValue.addKey(b, key); |
| _KeyValue.addValue(b, val); |
| return _KeyValue.endKeyValue(b); |
| })); |
| |
| _Schema.startSchema(b); |
| _Schema.addFields(b, fieldsVectorOffset); |
| _Schema.addEndianness(b, platformIsLittleEndian ? _Endianness.Little : _Endianness.Big); |
| |
| if (metadataOffset !== -1) { _Schema.addCustomMetadata(b, metadataOffset); } |
| |
| return _Schema.endSchema(b); |
| } |
| |
| /** @ignore */ |
| function encodeField(b: Builder, field: Field) { |
| |
| let nameOffset = -1; |
| let typeOffset = -1; |
| let dictionaryOffset = -1; |
| |
| const type = field.type; |
| let typeId: Type = <any>field.typeId; |
| |
| if (!DataType.isDictionary(type)) { |
| typeOffset = typeAssembler.visit(type, b)!; |
| } else { |
| typeId = type.dictionary.typeId; |
| dictionaryOffset = typeAssembler.visit(type, b)!; |
| typeOffset = typeAssembler.visit(type.dictionary, b)!; |
| } |
| |
| const childOffsets = (type.children || []).map((f: Field) => Field.encode(b, f)); |
| const childrenVectorOffset = _Field.createChildrenVector(b, childOffsets); |
| |
| const metadataOffset = !(field.metadata && field.metadata.size > 0) ? -1 : |
| _Field.createCustomMetadataVector(b, [...field.metadata].map(([k, v]) => { |
| const key = b.createString(`${k}`); |
| const val = b.createString(`${v}`); |
| _KeyValue.startKeyValue(b); |
| _KeyValue.addKey(b, key); |
| _KeyValue.addValue(b, val); |
| return _KeyValue.endKeyValue(b); |
| })); |
| |
| if (field.name) { |
| nameOffset = b.createString(field.name); |
| } |
| |
| _Field.startField(b); |
| _Field.addType(b, typeOffset); |
| _Field.addTypeType(b, typeId); |
| _Field.addChildren(b, childrenVectorOffset); |
| _Field.addNullable(b, !!field.nullable); |
| |
| if (nameOffset !== -1) { _Field.addName(b, nameOffset); } |
| if (dictionaryOffset !== -1) { _Field.addDictionary(b, dictionaryOffset); } |
| if (metadataOffset !== -1) { _Field.addCustomMetadata(b, metadataOffset); } |
| |
| return _Field.endField(b); |
| } |
| |
| /** @ignore */ |
| function encodeRecordBatch(b: Builder, recordBatch: RecordBatch) { |
| |
| const nodes = recordBatch.nodes || []; |
| const buffers = recordBatch.buffers || []; |
| const variadicBufferCounts = recordBatch.variadicBufferCounts || []; |
| |
| _RecordBatch.startNodesVector(b, nodes.length); |
| for (const n of nodes.slice().reverse()) FieldNode.encode(b, n); |
| |
| const nodesVectorOffset = b.endVector(); |
| |
| _RecordBatch.startBuffersVector(b, buffers.length); |
| for (const b_ of buffers.slice().reverse()) BufferRegion.encode(b, b_); |
| |
| const buffersVectorOffset = b.endVector(); |
| |
| let bodyCompressionOffset = null; |
| if (recordBatch.compression !== null) { |
| bodyCompressionOffset = encodeBodyCompression(b, recordBatch.compression); |
| } |
| |
| let variadicBufferCountsOffset = -1; |
| if (variadicBufferCounts.length > 0) { |
| variadicBufferCountsOffset = _RecordBatch.createVariadicBufferCountsVector(b, variadicBufferCounts.map(BigInt)); |
| } |
| |
| _RecordBatch.startRecordBatch(b); |
| _RecordBatch.addLength(b, BigInt(recordBatch.length)); |
| _RecordBatch.addNodes(b, nodesVectorOffset); |
| _RecordBatch.addBuffers(b, buffersVectorOffset); |
| if (recordBatch.compression !== null && bodyCompressionOffset) { |
| _RecordBatch.addCompression(b, bodyCompressionOffset); |
| } |
| if (variadicBufferCountsOffset !== -1) { |
| _RecordBatch.addVariadicBufferCounts(b, variadicBufferCountsOffset); |
| } |
| return _RecordBatch.endRecordBatch(b); |
| } |
| |
| /** @ignore */ |
| function encodeBodyCompression(b: Builder, node: BodyCompression) { |
| _BodyCompression.startBodyCompression(b); |
| _BodyCompression.addCodec(b, node.type); |
| _BodyCompression.addMethod(b, node.method); |
| return _BodyCompression.endBodyCompression(b); |
| } |
| |
| /** @ignore */ |
| function encodeDictionaryBatch(b: Builder, dictionaryBatch: DictionaryBatch) { |
| const dataOffset = RecordBatch.encode(b, dictionaryBatch.data); |
| _DictionaryBatch.startDictionaryBatch(b); |
| _DictionaryBatch.addId(b, BigInt(dictionaryBatch.id)); |
| _DictionaryBatch.addIsDelta(b, dictionaryBatch.isDelta); |
| _DictionaryBatch.addData(b, dataOffset); |
| return _DictionaryBatch.endDictionaryBatch(b); |
| } |
| |
| /** @ignore */ |
| function encodeFieldNode(b: Builder, node: FieldNode) { |
| return _FieldNode.createFieldNode(b, BigInt(node.length), BigInt(node.nullCount)); |
| } |
| |
| /** @ignore */ |
| function encodeBufferRegion(b: Builder, node: BufferRegion) { |
| return _Buffer.createBuffer(b, BigInt(node.offset), BigInt(node.length)); |
| } |
| |
| /** @ignore */ |
| const platformIsLittleEndian = (() => { |
| const buffer = new ArrayBuffer(2); |
| new DataView(buffer).setInt16(0, 256, true /* littleEndian */); |
| // Int16Array uses the platform's endianness. |
| return new Int16Array(buffer)[0] === 256; |
| })(); |
| |
| /** @ignore */ |
| type MessageHeaderDecoder = <T extends MessageHeader>() => T extends MessageHeader.Schema ? Schema |
| : T extends MessageHeader.RecordBatch ? RecordBatch |
| : T extends MessageHeader.DictionaryBatch ? DictionaryBatch : never; |