public void CanWriteToNetworkStream()
{
RecordBatch originalBatch = TestData.CreateSampleRecordBatch(length: 100);
const int port = 32153;
TcpListener listener = new TcpListener(IPAddress.Loopback, port);
listener.Start();
using (TcpClient sender = new TcpClient())
{
sender.Connect(IPAddress.Loopback, port);
NetworkStream stream = sender.GetStream();
using (var writer = new ArrowStreamWriter(stream, originalBatch.Schema))
{
writer.WriteRecordBatch(originalBatch);
writer.WriteEnd();
stream.Flush();
}
}
using (TcpClient receiver = listener.AcceptTcpClient())
{
NetworkStream stream = receiver.GetStream();
using (var reader = new ArrowStreamReader(stream))
{
RecordBatch newBatch = reader.ReadNextRecordBatch();
ArrowReaderVerifier.CompareBatches(originalBatch, newBatch);
}
}
}
private static void TestRoundTripRecordBatches(List <RecordBatch> originalBatches, IpcOptions options = null)
{
using (MemoryStream stream = new MemoryStream())
{
using (var writer = new ArrowStreamWriter(stream, originalBatches[0].Schema, leaveOpen: true, options))
{
foreach (RecordBatch originalBatch in originalBatches)
{
writer.WriteRecordBatch(originalBatch);
}
writer.WriteEnd();
}
stream.Position = 0;
using (var reader = new ArrowStreamReader(stream))
{
foreach (RecordBatch originalBatch in originalBatches)
{
RecordBatch newBatch = reader.ReadNextRecordBatch();
ArrowReaderVerifier.CompareBatches(originalBatch, newBatch);
}
}
}
}
public void WritesEmptyFile()
{
RecordBatch originalBatch = TestData.CreateSampleRecordBatch(length: 1);
var stream = new MemoryStream();
var writer = new ArrowStreamWriter(stream, originalBatch.Schema);
writer.WriteStart();
writer.WriteEnd();
stream.Position = 0;
var reader = new ArrowStreamReader(stream);
RecordBatch readBatch = reader.ReadNextRecordBatch();
Assert.Null(readBatch);
SchemaComparer.Compare(originalBatch.Schema, reader.Schema);
}
private static void TestRoundTripRecordBatch(RecordBatch originalBatch)
{
using (MemoryStream stream = new MemoryStream())
{
using (var writer = new ArrowStreamWriter(stream, originalBatch.Schema, leaveOpen: true))
{
writer.WriteRecordBatch(originalBatch);
writer.WriteEnd();
}
stream.Position = 0;
using (var reader = new ArrowStreamReader(stream))
{
RecordBatch newBatch = reader.ReadNextRecordBatch();
ArrowReaderVerifier.CompareBatches(originalBatch, newBatch);
}
}
}
public void WriteLegacyIpcFormat(bool writeLegacyIpcFormat)
{
RecordBatch originalBatch = TestData.CreateSampleRecordBatch(length: 100);
var options = new IpcOptions()
{
WriteLegacyIpcFormat = writeLegacyIpcFormat
};
using (MemoryStream stream = new MemoryStream())
{
using (var writer = new ArrowStreamWriter(stream, originalBatch.Schema, leaveOpen: true, options))
{
writer.WriteRecordBatch(originalBatch);
writer.WriteEnd();
}
stream.Position = 0;
// ensure the continuation is written correctly
byte[] buffer = stream.ToArray();
int messageLength = BinaryPrimitives.ReadInt32LittleEndian(buffer);
int endOfBuffer1 = BinaryPrimitives.ReadInt32LittleEndian(buffer.AsSpan(buffer.Length - 8));
int endOfBuffer2 = BinaryPrimitives.ReadInt32LittleEndian(buffer.AsSpan(buffer.Length - 4));
if (writeLegacyIpcFormat)
{
// the legacy IPC format doesn't have a continuation token at the start
Assert.NotEqual(-1, messageLength);
Assert.NotEqual(-1, endOfBuffer1);
}
else
{
// the latest IPC format has a continuation token at the start
Assert.Equal(-1, messageLength);
Assert.Equal(-1, endOfBuffer1);
}
Assert.Equal(0, endOfBuffer2);
}
}
public void WriteBatchWithCorrectPadding()
{
byte value1 = 0x04;
byte value2 = 0x14;
var batch = new RecordBatch(
new Schema.Builder()
.Field(f => f.Name("age").DataType(Int32Type.Default))
.Field(f => f.Name("characterCount").DataType(Int32Type.Default))
.Build(),
new IArrowArray[]
{
new Int32Array(
new ArrowBuffer(new byte[] { value1, value1, 0x00, 0x00 }),
ArrowBuffer.Empty,
length: 1,
nullCount: 0,
offset: 0),
new Int32Array(
new ArrowBuffer(new byte[] { value2, value2, 0x00, 0x00 }),
ArrowBuffer.Empty,
length: 1,
nullCount: 0,
offset: 0)
},
length: 1);
TestRoundTripRecordBatch(batch);
using (MemoryStream stream = new MemoryStream())
{
using (var writer = new ArrowStreamWriter(stream, batch.Schema, leaveOpen: true))
{
writer.WriteRecordBatch(batch);
writer.WriteEnd();
}
byte[] writtenBytes = stream.ToArray();
// ensure that the data buffers at the end are 8-byte aligned
Assert.Equal(value1, writtenBytes[writtenBytes.Length - 24]);
Assert.Equal(value1, writtenBytes[writtenBytes.Length - 23]);
for (int i = 22; i > 16; i--)
{
Assert.Equal(0, writtenBytes[writtenBytes.Length - i]);
}
Assert.Equal(value2, writtenBytes[writtenBytes.Length - 16]);
Assert.Equal(value2, writtenBytes[writtenBytes.Length - 15]);
for (int i = 14; i > 8; i--)
{
Assert.Equal(0, writtenBytes[writtenBytes.Length - i]);
}
// verify the EOS is written correctly
for (int i = 8; i > 4; i--)
{
Assert.Equal(0xFF, writtenBytes[writtenBytes.Length - i]);
}
for (int i = 4; i > 0; i--)
{
Assert.Equal(0x00, writtenBytes[writtenBytes.Length - i]);
}
}
}
