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)); }