protected RecordBatch CreateArrowObjectFromMessage(
Flatbuf.Message message, ByteBuffer bodyByteBuffer, IMemoryOwner <byte> memoryOwner)
{
switch (message.HeaderType)
{
case Flatbuf.MessageHeader.Schema:
// TODO: Read schema and verify equality?
break;
case Flatbuf.MessageHeader.DictionaryBatch:
// TODO: not supported currently
Debug.WriteLine("Dictionaries are not yet supported.");
break;
case Flatbuf.MessageHeader.RecordBatch:
Flatbuf.RecordBatch rb = message.Header <Flatbuf.RecordBatch>().Value;
List <IArrowArray> arrays = BuildArrays(Schema, bodyByteBuffer, rb);
return(new RecordBatch(Schema, memoryOwner, arrays, (int)rb.Length));
default:
// NOTE: Skip unsupported message type
Debug.WriteLine($"Skipping unsupported message type '{message.HeaderType}'");
break;
}
return(null);
}
protected RecordBatch ReadRecordBatch()
{
ReadSchema();
int messageLength = ReadMessageLength(throwOnFullRead: false);
if (messageLength == 0)
{
// reached end
return(null);
}
RecordBatch result = null;
ArrayPool <byte> .Shared.RentReturn(messageLength, messageBuff =>
{
int bytesRead = BaseStream.ReadFullBuffer(messageBuff);
EnsureFullRead(messageBuff, bytesRead);
Flatbuf.Message message = Flatbuf.Message.GetRootAsMessage(CreateByteBuffer(messageBuff));
int bodyLength = checked ((int)message.BodyLength);
IMemoryOwner <byte> bodyBuffOwner = _allocator.Allocate(bodyLength);
Memory <byte> bodyBuff = bodyBuffOwner.Memory.Slice(0, bodyLength);
bytesRead = BaseStream.ReadFullBuffer(bodyBuff);
EnsureFullRead(bodyBuff, bytesRead);
FlatBuffers.ByteBuffer bodybb = CreateByteBuffer(bodyBuff);
result = CreateArrowObjectFromMessage(message, bodybb, bodyBuffOwner);
});
return(result);
}
protected async ValueTask <RecordBatch> ReadRecordBatchAsync(CancellationToken cancellationToken = default)
{
await ReadSchemaAsync().ConfigureAwait(false);
int messageLength = 0;
await ArrayPool <byte> .Shared.RentReturnAsync(4, async (lengthBuffer) =>
{
// Get Length of record batch for message header.
int bytesRead = await BaseStream.ReadFullBufferAsync(lengthBuffer, cancellationToken)
.ConfigureAwait(false);
if (bytesRead == 4)
{
messageLength = BitUtility.ReadInt32(lengthBuffer);
}
}).ConfigureAwait(false);
if (messageLength == 0)
{
// reached end
return(null);
}
RecordBatch result = null;
await ArrayPool <byte> .Shared.RentReturnAsync(messageLength, async (messageBuff) =>
{
int bytesRead = await BaseStream.ReadFullBufferAsync(messageBuff, cancellationToken)
.ConfigureAwait(false);
EnsureFullRead(messageBuff, bytesRead);
Flatbuf.Message message = Flatbuf.Message.GetRootAsMessage(CreateByteBuffer(messageBuff));
int bodyLength = checked ((int)message.BodyLength);
IMemoryOwner <byte> bodyBuffOwner = _allocator.Allocate(bodyLength);
Memory <byte> bodyBuff = bodyBuffOwner?.Memory.Slice(0, bodyLength) ?? Memory <byte> .Empty;
bytesRead = await BaseStream.ReadFullBufferAsync(bodyBuff, cancellationToken)
.ConfigureAwait(false);
EnsureFullRead(bodyBuff, bytesRead);
FlatBuffers.ByteBuffer bodybb = CreateByteBuffer(bodyBuff);
result = CreateArrowObjectFromMessage(message, bodybb, bodyBuffOwner);
}).ConfigureAwait(false);
return(result);
}
protected static T ReadMessage <T>(ByteBuffer bb)
where T : struct, IFlatbufferObject
{
Type returnType = typeof(T);
Flatbuf.Message msg = Flatbuf.Message.GetRootAsMessage(bb);
if (MatchEnum(msg.HeaderType, returnType))
{
return(msg.Header <T>().Value);
}
else
{
throw new Exception($"Requested type '{returnType.Name}' " +
$"did not match type found at offset => '{msg.HeaderType}'");
}
}
protected async ValueTask <RecordBatch> ReadRecordBatchAsync(CancellationToken cancellationToken = default)
{
await ReadSchemaAsync().ConfigureAwait(false);
int messageLength = 0;
await Buffers.RentReturnAsync(4, async (lengthBuffer) =>
{
// Get Length of record batch for message header.
int bytesRead = await BaseStream.ReadFullBufferAsync(lengthBuffer, cancellationToken)
.ConfigureAwait(false);
if (bytesRead == 4)
{
messageLength = BitUtility.ReadInt32(lengthBuffer);
}
}).ConfigureAwait(false);
if (messageLength == 0)
{
// reached end
return(null);
}
RecordBatch result = null;
await Buffers.RentReturnAsync(messageLength, async (messageBuff) =>
{
int bytesRead = await BaseStream.ReadFullBufferAsync(messageBuff, cancellationToken)
.ConfigureAwait(false);
EnsureFullRead(messageBuff, bytesRead);
Flatbuf.Message message = Flatbuf.Message.GetRootAsMessage(CreateByteBuffer(messageBuff));
await Buffers.RentReturnAsync((int)message.BodyLength, async(bodyBuff) =>
{
int bodyBytesRead = await BaseStream.ReadFullBufferAsync(bodyBuff, cancellationToken)
.ConfigureAwait(false);
EnsureFullRead(bodyBuff, bodyBytesRead);
FlatBuffers.ByteBuffer bodybb = CreateByteBuffer(bodyBuff);
result = CreateArrowObjectFromMessage(message, bodybb);
}).ConfigureAwait(false);
}).ConfigureAwait(false);
return(result);
}
protected RecordBatch ReadRecordBatch()
{
ReadSchema();
int messageLength = 0;
Buffers.RentReturn(4, lengthBuffer =>
{
int bytesRead = BaseStream.ReadFullBuffer(lengthBuffer);
if (bytesRead == 4)
{
messageLength = BitUtility.ReadInt32(lengthBuffer);
}
});
if (messageLength == 0)
{
// reached end
return(null);
}
RecordBatch result = null;
Buffers.RentReturn(messageLength, messageBuff =>
{
int bytesRead = BaseStream.ReadFullBuffer(messageBuff);
EnsureFullRead(messageBuff, bytesRead);
Flatbuf.Message message = Flatbuf.Message.GetRootAsMessage(CreateByteBuffer(messageBuff));
Buffers.RentReturn((int)message.BodyLength, bodyBuff =>
{
int bodyBytesRead = BaseStream.ReadFullBuffer(bodyBuff);
EnsureFullRead(bodyBuff, bodyBytesRead);
FlatBuffers.ByteBuffer bodybb = CreateByteBuffer(bodyBuff);
result = CreateArrowObjectFromMessage(message, bodybb);
});
});
return(result);
}
protected async ValueTask <RecordBatch> ReadRecordBatchAsync(CancellationToken cancellationToken = default)
{
await ReadSchemaAsync().ConfigureAwait(false);
RecordBatch result = null;
while (result == null)
{
int messageLength = await ReadMessageLengthAsync(throwOnFullRead : false, cancellationToken)
.ConfigureAwait(false);
if (messageLength == 0)
{
// reached end
return(null);
}
await ArrayPool <byte> .Shared.RentReturnAsync(messageLength, async (messageBuff) =>
{
int bytesRead = await BaseStream.ReadFullBufferAsync(messageBuff, cancellationToken)
.ConfigureAwait(false);
EnsureFullRead(messageBuff, bytesRead);
Flatbuf.Message message = Flatbuf.Message.GetRootAsMessage(CreateByteBuffer(messageBuff));
int bodyLength = checked ((int)message.BodyLength);
IMemoryOwner <byte> bodyBuffOwner = _allocator.Allocate(bodyLength);
Memory <byte> bodyBuff = bodyBuffOwner.Memory.Slice(0, bodyLength);
bytesRead = await BaseStream.ReadFullBufferAsync(bodyBuff, cancellationToken)
.ConfigureAwait(false);
EnsureFullRead(bodyBuff, bytesRead);
FlatBuffers.ByteBuffer bodybb = CreateByteBuffer(bodyBuff);
result = CreateArrowObjectFromMessage(message, bodybb, bodyBuffOwner);
}).ConfigureAwait(false);
}
return(result);
}
