How to send packed array from math kernel to 3rd party app via WSTP

Question

Background

I'm developing a small python wrapper PyWSTP. The wrapper is still in very early stage and lacking features. One important feature I'm considering is transferring dense array between math kernel and python. (Yep, the all awesome NumPy) I do realize LibraryLink would be better suited for this, but WSTP has the capability of sending stuff over TCP.

Problem

Whenever kernel sends packed array to a 3rd party app via LinkWrite[_LinkObject, _], it sends unpacked array. On the receiving side, tokens such WSTKPACKED are never seen. Thus WSTP C APIs such as WSGetReal64Array cannot be used.

Here's some simple code to illustrate the problem:

server.m

(* run via "math -script" *)
Needs["Developer`"];

link = LinkCreate["testlink", LinkMode->Listen];

obj1 = Image[RandomInteger[{0, 255}, {4, 4}]];
obj2 = Developer`ToPackedArray[RandomReal[1., {4,4}]];
LinkWrite[link, #]& /@ {obj1, obj2};
InputString["Done, ENTER to exit ..."];

LinkClose[link];

client.py

import wstp

ctx = wstp.WSTPContext()
link = ctx.open_link('testlink', mode='connect')

obj1 = link.recv()
print('got object1:')
print(obj1)
obj2 = link.recv()
print('got object2:')
print(obj2)
input('Done, ENTER to exit ...')

ctx.close_link('testlink')

On PyWSTP, I have not implemented handlers for WSTKPACKED yet. Only basic data types (WSTKSYM, WSTKFUNC, WSTKINT, WSTKREAL, WSTKSTR) are implemented. Yet client side is able to correctly receive the object:

got object1:
Image[RawArray["Real64", List[List[207.0, 142.0, 166.0, 205.0], List[137.0, 115.0, 60.0, 39.0], List[37.0, 110.0, 9.0, 89.0], List[56.0, 162.0, 116.0, 225.0]]], "Real", Rule[ColorSpace, Automatic], Rule[Interleaving, None]]
got object2:
List[List[0.8681530207587866, 0.0734728851659936, 0.9233463976844913, 0.1346824422121633], List[0.7711652733626824, 0.88933420636346, 0.9202474236790081, 0.5288719560153152], List[0.18096422954894353, 0.04400711306110572, 0.2330108388845873, 0.7479540642843265], List[0.4382030298759463, 0.42470372650663424, 0.9718400392443247, 0.6100416353164417]]
Done, ENTER to exit ...

This indicates math kernel unpacks dense array before sending over WSTP.

Question summary

1. How to make sure kernel sends packed array to 3rd party app over WSTP?

2. Is it possible to achieve 1. without writing C code?

Answer 1

Following a dedicated discussion with @Szabolcs, very helpful comments from @ihojnicki, and information from an old thread, I successfully implemented NumPy array sender/receiver in PyWSTP.

Assume we are trying to receive a Real64 array in C++.

(* server.m *)
LinkOpen["someLink", LinkMode->Listen];
LinkWrite[link, RandomReal[1., {4,4}]];
InputString[];
LinkClose[link];

The kernel does send packed array in packed form. However, in order to successfully receive packed array, one must call WSGetRawType to retrieve low level type tokens:

#include "wstp.h"

int err;
auto env_p = WSInitialize((void*)0);
auto link_p = WSOpenString(env_p, "someLink -linkmode connect", &err);
if (err != WSEOK) { /* error ... */ }

auto token = WSGetType(link_p);
auto raw_token = WSGetRawType(link_p);

Here token will be WSTKFUNC, however raw_token is WSTKARRAY. The problem here is how to determine data type and size of the multidimensional array. From this point on, two methods can be used (as of WSTPv4):

Method 1 - Create a temporary loopback link to store array metadata

auto temp_link_p = WSLoopbackOpen(env_p, &err);
if (err != WSEOK) { /* error ... */ }
array_meterp am;
long depth;
mlapi__token leaf_token; // this token seems unused

if (!WSGetArrayTypeWithDepthAndLeafType(link_p, temp_link_p, &am, &depth, &leaf_token)) {
    /* error ... */
}

size_t total_size = 1;
const char *head_p;
int argcount;
vector<int> shape;
vector<string> heads;

for (int i=0; i<depth; ++i) {
    WSGetFunction(temp_link_p, &head_p, &argcount);
    total_size *= argcount;
    shape.push_back(argcount);
    heads.push_back(string(head_p));
    WSReleaseSymbol(link_p, head_p);
}

auto elem_token = WSGetRawType(link_p);

switch (elem_token) {
case WSTK_CDOUBLE:
    total_size *= sizeof(double);
    break;
case WSTK_CFLOAT:
    total_size *= sizeof(float);
    break;
/* case ... */
default:
    /* error ... */
    break;
}

void *array_data = malloc(total_size);
WSGetBinaryNumberArrayData(link_p, am, array_data, total_size, elem_token);
WSReleaseGetArrayState(link_p, temp_link_p, am);
WSClose(temp_link_p);

for (int i=0; i<16; ++i) {
printf("%f\n", ((double*)array_data)[i]);
}
/* do stuff with array */

free(array_data);

Method 2 - Set mark, read metadata manually, then seek back and read data

vector<int> shape;
vector<string> heads;
size_t total_size = 1;
const char *head_p;
int argcount;
int elem_token;
auto mark = WSCreateMark(link_p);

do {
    WSGetFunction(link_p, &head_p, &argcount);
    total_size *= argcount;
    shape.push_back(argcount);
    heads.push_back(string(head_p));
    WSReleaseSymbol(link_p, head_p);
    elem_token = WSGetRawType(link_p);
} while (elem_token == WSTKFUNC);

switch (elem_token) {
case WSTK_CDOUBLE:
    total_size *= sizeof(double);
    break;
case WSTK_CFLOAT:
    total_size *= sizeof(float);
    break;
/* case ... */
default:
    /* error ... */
    break;
}

WSSeekToMark(link_p, mark, 0);
WSDestroyMark(link_p, mark);

void *array_data;
long *received_dims;
char **received_heads;
long received_depth;

WSGetBinaryNumberArray(
    link_p, &array_data, &received_dims, &received_heads, &received_depth, elem_token);

/* do stuff with array */

// there is a int/long mismatch on received_dims, I don't understand why
WSReleaseBinaryNumberArray(
    link_p, array_data, (int*)received_dims, received_heads, total_size, elem_token);

Observations

• array_meterp is pointer to an opaque internal struct, likely holding array metadata.

• WSGetBinaryNumberArray* is using platform specific type tokens. They can be found in wstp.h.

• As of WSTP v4, WSGet*Array (e.g WSGetInteger64Array) are wrappers to WSGetBinaryNumberArray.

• Temporary link in method 1 just store function heads, attempts to retrieve full exrepssion may fail.

• Generally the heads in each level are "List".

• The last head can be "Complex" in case of complex data type. E.g a complex unitary matrix would have List, List, and Complex.

The full code can be found here. Under linux, compile via:

WSTP_PATH=$MATHEMATICA_HOME/SystemFiles/Links/WSTP/DeveloperKit/Linux-x86-64/CompilerAdditions
g++ a.cpp -g -I$WSTP_PATH -L$WSTP_PATH -lWSTP64i4 -luuid -lrt -ldl

Credits to @ihojnicki, and @Szabolcs.