Support wrapping a Java-side block of memory allocated off-heap to Python

[ctrueden]

With ByteBuffer.allocateDirect you can allocate memory off-heap in Java, which can then be shared with other processes. We want to easy manufacturing of numpy arrays and xarrays that wrap this sort of off-heap memory, so that you can directly change data in Python that originated in Java (by some definition of "originated"—since it's off-heap).

In #73, @hanslovsky wrote:

should be possible already, albeit I don't think there is a convenience method for that. Things may have changed since I was last involved with imglyb (it was still pyjnius back then), but there is/was a way to generate ImgLib2 ArrayImgs backed by native memory and you can then simply pass that pointer into a numpy array.

See also https://github.com/imglib/imglib2-cache-python

[gselzer]

See https://github.com/imglib/imglyb/pull/14

[ctrueden]

This is already doable easily with currently released versions of scyjava + numpy + jpype, along with @mkitti's ByteBufferAccess:

import scyjava
import numpy
import jpype
import random

scyjava.config.endpoints.append("net.imglib2:imglib2:6.0.0")
print("Populating 5 x 7 x 11 direct byte buffer")
ByteBuffer = scyjava.jimport('java.nio.ByteBuffer')
jbuf = ByteBuffer.allocateDirect(5 * 7 * 11)
for _ in range(jbuf.limit()):
    jbuf.put(random.randint(-128, 127))
print(f"Buffer limit = {jbuf.limit()}")

print("Wrapping it to an ndarray")
view = memoryview(jbuf)
narr = numpy.frombuffer(view, count=jbuf.limit(), dtype=numpy.int8)
narr = narr.reshape([11, 7, 5])
print(f"Shape = {narr.shape}")

print()
print("Wrapping to ImgLib2 image")
ArrayImgs = scyjava.jimport('net.imglib2.img.array.ArrayImgs')
ByteBufferAccess = scyjava.jimport('net.imglib2.img.basictypeaccess.nio.ByteBufferAccess')
access = ByteBufferAccess(jbuf, True)
dims = scyjava.jarray('j', 3)
dims[0] = 5; dims[1] = 7; dims[2] = 11
img = ArrayImgs.bytes(access, dims);
print(img)

def print_value(x, y, z):
    print(f"- narr[{z}, {y}, {x}] = {narr[z, y, x]}")
    print(f"- jbuf.get({z}*5*7 + {y}*5 + {x}) = {jbuf.get(z*5*7 + y*5 + x)}")
    pos = scyjava.jarray('j', 3)
    pos[0] = x; pos[1] = y; pos[2] = z
    print(f"- img.getAt({x}, {y}, {z}).get() = {img.getAt(pos).get()}")

print()
print("Initial values:")
print_value(0, 0, 0)
print_value(2, 4, 6)

print()
print("Changing values:")
print("- narr[6, 4, 2] -> 17")
narr[6, 4, 2] = 17
print("- narr[0, 0, 0] -> 23")
narr[0, 0, 0] = 23

print()
print("Values after:")
print_value(0, 0, 0)
print_value(2, 4, 6)

produces:

Populating 5 x 7 x 11 direct byte buffer
Buffer limit = 385
Wrapping it to an ndarray
Shape = (11, 7, 5)

Wrapping to ImgLib2 image
ArrayImg [5x7x11]

Initial values:
- narr[0, 0, 0] = 4
- jbuf.get(0*5*7 + 0*5 + 0) = 4
- img.getAt(0, 0, 0).get() = 4
- narr[6, 4, 2] = 125
- jbuf.get(6*5*7 + 4*5 + 2) = 125
- img.getAt(2, 4, 6).get() = 125

Changing values:
- narr[6, 4, 2] -> 17
- narr[0, 0, 0] -> 23

Values after:
- narr[0, 0, 0] = 23
- jbuf.get(0*5*7 + 0*5 + 0) = 23
- img.getAt(0, 0, 0).get() = 23
- narr[6, 4, 2] = 17
- jbuf.get(6*5*7 + 4*5 + 2) = 17
- img.getAt(2, 4, 6).get() = 17

Note: There is a bug in the above code relating to the memoryview, resulting in an error sometimes:

Traceback (most recent call last):
  File "...py", line 16, in <module>
    narr = numpy.frombuffer(view, count=jbuf.limit(), dtype=numpy.int8)
ValueError: buffer is smaller than requested size

On my system it works maybe 10-20% of the time, producing the above error the other 80-90%. But I don't have time to debug right now—I just wanted to post the code as a starting point should anyone feel like working on this issue. This issue becomes mostly just: A) fixing said bug; and B) deciding how to slot in this logic most conveniently to PyImageJ's API; and C) generalizing it to use buffer views for other types besides just ByteType.

Note also that ByteBuffer.allocateDirect is limited to 2GB in size. For larger images we will likely want to use the foreign memory API (which is not yet a permanent feature of Java, but is available as a preview feature in Java 19).

[Thrameos]

Could the error be related to the nbytes field incorrect in currently released versions of JPype? I fixed it recently but it after the release.

[ctrueden]

@Thrameos Yep, works every time with JPype installed from the current master branch (jpype-project/JPype@4bacf4c9ca79d7744b4bdf74c826046f8ffefd45). Thanks!

[Thrameos]

Drat. That means I will need another micro release. Perhaps I can finish these requests for random access file conversions and then release it together in Nov/Dec timeframe?

[ctrueden]

Sure, no rush on my side! JPype works great for how we're using it now. Supporting ByteBuffer.allocateDirect is not urgent for PyImageJ.