public static void CompareBatches(RecordBatch expectedBatch, RecordBatch actualBatch)
{
SchemaComparer.Compare(expectedBatch.Schema, actualBatch.Schema);
Assert.Equal(expectedBatch.Length, actualBatch.Length);
Assert.Equal(expectedBatch.ColumnCount, actualBatch.ColumnCount);
for (int i = 0; i < expectedBatch.ColumnCount; i++)
{
IArrowArray expectedArray = expectedBatch.Arrays.ElementAt(i);
IArrowArray actualArray = actualBatch.Arrays.ElementAt(i);
actualArray.Accept(new ArrayComparer(expectedArray));
}
}
private Tuple <ArrowRecordBatchFlatBufferBuilder, VectorOffset> PreparingWritingRecordBatch(IReadOnlyDictionary <string, Field> fields, IReadOnlyList <IArrowArray> arrays)
{
Builder.Clear();
// Serialize field nodes
int fieldCount = fields.Count;
Flatbuf.RecordBatch.StartNodesVector(Builder, CountAllNodes(fields));
// flatbuffer struct vectors have to be created in reverse order
for (int i = fieldCount - 1; i >= 0; i--)
{
CreateSelfAndChildrenFieldNodes(arrays[i].Data);
}
VectorOffset fieldNodesVectorOffset = Builder.EndVector();
// Serialize buffers
var recordBatchBuilder = new ArrowRecordBatchFlatBufferBuilder();
for (int i = 0; i < fieldCount; i++)
{
IArrowArray fieldArray = arrays[i];
fieldArray.Accept(recordBatchBuilder);
}
IReadOnlyList <ArrowRecordBatchFlatBufferBuilder.Buffer> buffers = recordBatchBuilder.Buffers;
Flatbuf.RecordBatch.StartBuffersVector(Builder, buffers.Count);
// flatbuffer struct vectors have to be created in reverse order
for (int i = buffers.Count - 1; i >= 0; i--)
{
Flatbuf.Buffer.CreateBuffer(Builder,
buffers[i].Offset, buffers[i].DataBuffer.Length);
}
return(Tuple.Create(recordBatchBuilder, fieldNodesVectorOffset));
}
private protected async Task WriteRecordBatchInternalAsync(RecordBatch recordBatch,
CancellationToken cancellationToken = default)
{
// TODO: Truncate buffers with extraneous padding / unused capacity
if (!HasWrittenSchema)
{
await WriteSchemaAsync(Schema, cancellationToken).ConfigureAwait(false);
HasWrittenSchema = true;
}
Builder.Clear();
// Serialize field nodes
int fieldCount = Schema.Fields.Count;
Flatbuf.RecordBatch.StartNodesVector(Builder, CountAllNodes());
// flatbuffer struct vectors have to be created in reverse order
for (int i = fieldCount - 1; i >= 0; i--)
{
CreateSelfAndChildrenFieldNodes(recordBatch.Column(i).Data);
}
VectorOffset fieldNodesVectorOffset = Builder.EndVector();
// Serialize buffers
var recordBatchBuilder = new ArrowRecordBatchFlatBufferBuilder();
for (int i = 0; i < fieldCount; i++)
{
IArrowArray fieldArray = recordBatch.Column(i);
fieldArray.Accept(recordBatchBuilder);
}
IReadOnlyList <ArrowRecordBatchFlatBufferBuilder.Buffer> buffers = recordBatchBuilder.Buffers;
Flatbuf.RecordBatch.StartBuffersVector(Builder, buffers.Count);
// flatbuffer struct vectors have to be created in reverse order
for (int i = buffers.Count - 1; i >= 0; i--)
{
Flatbuf.Buffer.CreateBuffer(Builder,
buffers[i].Offset, buffers[i].DataBuffer.Length);
}
VectorOffset buffersVectorOffset = Builder.EndVector();
// Serialize record batch
StartingWritingRecordBatch();
Offset <Flatbuf.RecordBatch> recordBatchOffset = Flatbuf.RecordBatch.CreateRecordBatch(Builder, recordBatch.Length,
fieldNodesVectorOffset,
buffersVectorOffset);
long metadataLength = await WriteMessageAsync(Flatbuf.MessageHeader.RecordBatch,
recordBatchOffset, recordBatchBuilder.TotalLength,
cancellationToken).ConfigureAwait(false);
// Write buffer data
long bodyLength = 0;
for (int i = 0; i < buffers.Count; i++)
{
ArrowBuffer buffer = buffers[i].DataBuffer;
if (buffer.IsEmpty)
{
continue;
}
await WriteBufferAsync(buffer, cancellationToken).ConfigureAwait(false);
int paddedLength = checked ((int)BitUtility.RoundUpToMultipleOf8(buffer.Length));
int padding = paddedLength - buffer.Length;
if (padding > 0)
{
await WritePaddingAsync(padding).ConfigureAwait(false);
}
bodyLength += paddedLength;
}
// Write padding so the record batch message body length is a multiple of 8 bytes
int bodyPaddingLength = CalculatePadding(bodyLength);
await WritePaddingAsync(bodyPaddingLength).ConfigureAwait(false);
FinishedWritingRecordBatch(bodyLength + bodyPaddingLength, metadataLength);
}
public static void CompareArrays(IArrowArray expectedArray, IArrowArray actualArray, bool strictCompare = true)
{
actualArray.Accept(new ArrayComparer(expectedArray, strictCompare));
}
