public void Write(Stream stream, CommandPayload commandPayload)
{
SerDe.Write(stream, (int)commandPayload.EvalType);
Write(stream, commandPayload.Commands);
if ((commandPayload.EvalType == PythonEvalType.SQL_SCALAR_PANDAS_UDF) ||
(commandPayload.EvalType == PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF))
{
SerDe.Write(stream, "unused timezone");
}
}
internal void Write(
Stream stream,
Payload payload,
CommandPayload commandPayload)
{
SerDe.Write(stream, payload.SplitIndex);
SerDe.Write(stream, payload.Version);
_taskContextWriter.Write(stream, payload.TaskContext);
SerDe.Write(stream, payload.SparkFilesDir);
Write(stream, payload.IncludeItems);
_broadcastVariableWriter.Write(stream, payload.BroadcastVariables);
_commandWriter.Write(stream, commandPayload);
}
public void TestClosedStreamWithSocket()
{
var commandPayload = new CommandPayload()
{
EvalType = UdfUtils.PythonEvalType.SQL_BATCHED_UDF,
Commands = new Command[] { }
};
PayloadWriter payloadWriter = new PayloadWriterFactory().Create();
Payload payload = TestData.GetDefaultPayload();
var serverListener = new DefaultSocketWrapper();
serverListener.Listen();
var port = (serverListener.LocalEndPoint as IPEndPoint).Port;
var clientSocket = new DefaultSocketWrapper();
clientSocket.Connect(IPAddress.Loopback, port, null);
using (ISocketWrapper serverSocket = serverListener.Accept())
{
Stream outStream = serverSocket.OutputStream;
payloadWriter.Write(outStream, payload, commandPayload);
outStream.Flush();
}
// At this point server socket is closed.
Stream inStream = clientSocket.InputStream;
// Consume bytes already written to the socket.
var payloadProcessor = new PayloadProcessor(payloadWriter.Version);
Payload actualPayload = payloadProcessor.Process(inStream);
Assert.Equal(payload.SplitIndex, actualPayload.SplitIndex);
Assert.Equal(payload.Version, actualPayload.Version);
Assert.Equal(payload.TaskContext, actualPayload.TaskContext);
Assert.Equal(payload.SparkFilesDir, actualPayload.SparkFilesDir);
Assert.Equal(payload.IncludeItems, actualPayload.IncludeItems);
Assert.Equal(payload.BroadcastVariables.Count, actualPayload.BroadcastVariables.Count);
ValidateCommandPayload(commandPayload, actualPayload.Command);
// Another read will detect that the socket is closed.
Assert.Null(payloadProcessor.Process(inStream));
}
public void TestPayloadProcessor(string version)
{
CommandPayload commandPayload = TestData.GetDefaultCommandPayload();
PayloadWriter payloadWriter = new PayloadWriterFactory().Create(new Version(version));
Payload payload = TestData.GetDefaultPayload();
Payload actualPayload = null;
using (var outStream = new MemoryStream())
{
payloadWriter.Write(outStream, payload, commandPayload);
using (var inputStream = new MemoryStream(outStream.ToArray()))
{
actualPayload =
new PayloadProcessor(payloadWriter.Version).Process(inputStream);
}
}
// Validate the read payload.
Assert.Equal(payload.SplitIndex, actualPayload.SplitIndex);
Assert.Equal(payload.Version, actualPayload.Version);
Assert.Equal(payload.TaskContext, actualPayload.TaskContext);
Assert.Equal(payload.SparkFilesDir, actualPayload.SparkFilesDir);
Assert.Equal(payload.IncludeItems, actualPayload.IncludeItems);
Assert.Equal(payload.BroadcastVariables.Count, actualPayload.BroadcastVariables.Count);
ValidateCommandPayload(commandPayload, actualPayload.Command);
// Validate the UDFs.
var actualCommand1 = (SqlCommand)actualPayload.Command.Commands[0];
var result1 = ((PicklingWorkerFunction)actualCommand1.WorkerFunction).Func(
0,
new object[] { "hello", 10, 20 },
actualCommand1.ArgOffsets);
Assert.Equal("udf2 udf1 hello", result1);
var actualCommand2 = (SqlCommand)actualPayload.Command.Commands[1];
var result2 = ((PicklingWorkerFunction)actualCommand2.WorkerFunction).Func(
0,
new object[] { "hello", 10, 20 },
actualCommand2.ArgOffsets);
Assert.Equal(30, result2);
}
public void Write(Stream stream, CommandPayload commandPayload)
{
if (commandPayload.EvalType == PythonEvalType.SQL_SCALAR_PANDAS_UDF ||
commandPayload.EvalType == PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF ||
commandPayload.EvalType == PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF ||
commandPayload.EvalType == PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF)
{
SerDe.Write(stream, 1);
for (int i = 0; i < 1; ++i)
{
SerDe.Write(stream, "unused key");
SerDe.Write(stream, "unused value");
}
}
SerDe.Write(stream, (int)commandPayload.EvalType);
Write(stream, commandPayload.Commands);
}
private void ValidateCommandPayload(
CommandPayload expected,
Worker.CommandPayload actual)
{
Assert.Equal(expected.EvalType, actual.EvalType);
Assert.Equal(expected.Commands.Length, actual.Commands.Count());
for (int i = 0; i < expected.Commands.Length; ++i)
{
Command expectedCommand = expected.Commands[i];
var actualCommand = (SqlCommand)actual.Commands[i];
Assert.Equal(expectedCommand.ArgOffsets, actualCommand.ArgOffsets);
Assert.Equal(
expectedCommand.ChainedUdfs.Length,
actualCommand.NumChainedFunctions);
Assert.Equal(
expectedCommand.SerializerMode,
actualCommand.SerializerMode);
Assert.Equal(
expectedCommand.DeserializerMode,
actualCommand.DeserializerMode);
}
}
public void WriteTestData(Stream stream)
{
Payload payload = TestData.GetDefaultPayload();
CommandPayload commandPayload = TestData.GetDefaultCommandPayload();
Write(stream, payload, commandPayload);
// Write 10 rows to the output stream.
var pickler = new Pickler();
for (int i = 0; i < 10; ++i)
{
byte[] pickled = pickler.dumps(
new[] { new object[] { i.ToString(), i, i } });
SerDe.Write(stream, pickled.Length);
SerDe.Write(stream, pickled);
}
// Signal the end of data and stream.
SerDe.Write(stream, (int)SpecialLengths.END_OF_DATA_SECTION);
SerDe.Write(stream, (int)SpecialLengths.END_OF_STREAM);
stream.Flush();
}
public void TestTaskRunner()
{
using var serverListener = new DefaultSocketWrapper();
serverListener.Listen();
var port = (serverListener.LocalEndPoint as IPEndPoint).Port;
var clientSocket = new DefaultSocketWrapper();
clientSocket.Connect(IPAddress.Loopback, port, null);
PayloadWriter payloadWriter = new PayloadWriterFactory().Create();
var taskRunner = new TaskRunner(0, clientSocket, false, payloadWriter.Version);
var clientTask = Task.Run(() => taskRunner.Run());
using (ISocketWrapper serverSocket = serverListener.Accept())
{
System.IO.Stream inputStream = serverSocket.InputStream;
System.IO.Stream outputStream = serverSocket.OutputStream;
Payload payload = TestData.GetDefaultPayload();
CommandPayload commandPayload = TestData.GetDefaultCommandPayload();
payloadWriter.Write(outputStream, payload, commandPayload);
// Write 10 rows to the output stream.
var pickler = new Pickler();
for (int i = 0; i < 10; ++i)
{
var pickled = pickler.dumps(
new[] { new object[] { i.ToString(), i, i } });
SerDe.Write(outputStream, pickled.Length);
SerDe.Write(outputStream, pickled);
}
// Signal the end of data and stream.
SerDe.Write(outputStream, (int)SpecialLengths.END_OF_DATA_SECTION);
SerDe.Write(outputStream, (int)SpecialLengths.END_OF_STREAM);
outputStream.Flush();
// Now process the bytes flowing in from the client.
var timingDataReceived = false;
var exceptionThrown = false;
var rowsReceived = new List <object[]>();
while (true)
{
var length = SerDe.ReadInt32(inputStream);
if (length > 0)
{
var pickledBytes = SerDe.ReadBytes(inputStream, length);
using var unpickler = new Unpickler();
var rows = unpickler.loads(pickledBytes) as ArrayList;
foreach (object row in rows)
{
rowsReceived.Add((object[])row);
}
}
else if (length == (int)SpecialLengths.TIMING_DATA)
{
var bootTime = SerDe.ReadInt64(inputStream);
var initTime = SerDe.ReadInt64(inputStream);
var finishTime = SerDe.ReadInt64(inputStream);
var memoryBytesSpilled = SerDe.ReadInt64(inputStream);
var diskBytesSpilled = SerDe.ReadInt64(inputStream);
timingDataReceived = true;
}
else if (length == (int)SpecialLengths.PYTHON_EXCEPTION_THROWN)
{
SerDe.ReadString(inputStream);
exceptionThrown = true;
break;
}
else if (length == (int)SpecialLengths.END_OF_DATA_SECTION)
{
var numAccumulatorUpdates = SerDe.ReadInt32(inputStream);
SerDe.ReadInt32(inputStream);
break;
}
}
Assert.True(timingDataReceived);
Assert.False(exceptionThrown);
// Validate rows received.
Assert.Equal(10, rowsReceived.Count);
for (int i = 0; i < 10; ++i)
{
// Two UDFs registered, thus expecting two columns.
// Refer to TestData.GetDefaultCommandPayload().
var row = rowsReceived[i];
Assert.Equal(2, rowsReceived[i].Length);
Assert.Equal($"udf2 udf1 {i}", row[0]);
Assert.Equal(i + i, row[1]);
}
}
Assert.True(clientTask.Wait(5000));
}
