Skip to content

Commit ad88eee

Browse files
fangererfangerer
committed
Migrate Arrow callbacks to NativeAccess
1 parent fa4ee4d commit ad88eee

7 files changed

Lines changed: 126 additions & 219 deletions

File tree

‎graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/GraalPythonModuleBuiltins.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1537,11 +1537,11 @@ static PTuple doCreate(long arrowArrayAddr, long arrowSchemaAddr,
15371537
@Bind Node inliningTarget,
15381538
@Cached PythonCextCapsuleBuiltins.PyCapsuleNewNode pyCapsuleNewNode) {
15391539
PythonContext ctx = getContext(inliningTarget);
1540-
long arrayDestructor = ctx.arrowSupport.getArrowArrayDestructor(inliningTarget);
1540+
long arrayDestructor = ctx.arrowSupport.getArrowArrayDestructor();
15411541
long arrayCapsuleNamePointer = ctx.stringToNativeUtf8Bytes(ArrowArray.CAPSULE_NAME, true);
15421542
PyCapsule arrowArrayCapsule = pyCapsuleNewNode.execute(inliningTarget, arrowArrayAddr, arrayCapsuleNamePointer, arrayDestructor);
15431543

1544-
long schemaDestructor = ctx.arrowSupport.getArrowSchemaDestructor(inliningTarget);
1544+
long schemaDestructor = ctx.arrowSupport.getArrowSchemaDestructor();
15451545
long schemaCapsuleNamePointer = ctx.stringToNativeUtf8Bytes(ArrowSchema.CAPSULE_NAME, true);
15461546
PyCapsule arrowSchemaCapsule = pyCapsuleNewNode.execute(inliningTarget, arrowSchemaAddr, schemaCapsuleNamePointer, schemaDestructor);
15471547
return PFactory.createTuple(ctx.getLanguage(inliningTarget), new Object[]{arrowSchemaCapsule, arrowArrayCapsule});

‎graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/cext/common/NativePointer.java‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -47,10 +47,7 @@
4747
import com.oracle.truffle.api.library.ExportLibrary;
4848
import com.oracle.truffle.api.library.ExportMessage;
4949

50-
/**
51-
* Currently only used for wrapping pointers to call TruffleString, see GR-71311 and in
52-
* InvokeArrowReleaseCallbackNode which still uses original NFI.
53-
*/
50+
/** Currently only used for wrapping pointers to call TruffleString, see GR-71311. */
5451
@ExportLibrary(InteropLibrary.class)
5552
public final class NativePointer implements TruffleObject {
5653
private final long ptr;

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/runtime/arrow/ArrowUtil.java renamed to graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/arrow/ArrowReleaseCallback.java

Lines changed: 18 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright (c) 2024, 2026, Oracle and/or its affiliates. All rights reserved.
2+
* Copyright (c) 2025, 2026, Oracle and/or its affiliates. All rights reserved.
33
* DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
44
*
55
* The Universal Permissive License (UPL), Version 1.0
@@ -38,19 +38,26 @@
3838
* OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
3939
* SOFTWARE.
4040
*/
41-
package com.oracle.graal.python.runtime.arrow;
41+
package com.oracle.graal.python.nodes.arrow;
4242

43-
import static com.oracle.graal.python.nodes.StringLiterals.J_NFI_LANGUAGE;
44-
import static com.oracle.graal.python.util.PythonUtils.callCallTarget;
43+
import java.lang.invoke.MethodHandle;
44+
import java.lang.invoke.MethodType;
4545

46-
import com.oracle.graal.python.runtime.PythonContext;
47-
import com.oracle.truffle.api.nodes.Node;
48-
import com.oracle.truffle.api.source.Source;
46+
import com.oracle.graal.python.runtime.nativeaccess.NativeAccessSupport;
47+
import com.oracle.truffle.api.CompilerDirectives;
4948

50-
public class ArrowUtil {
49+
public final class ArrowReleaseCallback {
50+
private static final MethodHandle HANDLE = NativeAccessSupport.createDowncallHandle(
51+
MethodType.methodType(void.class, long.class, long.class), false);
5152

52-
public static Object createNfiSignature(Node location, String methodSignature, PythonContext ctx) {
53-
Source sigSource = Source.newBuilder(J_NFI_LANGUAGE, methodSignature, "python-nfi-signature").build();
54-
return callCallTarget(ctx.getEnv().parseInternal(sigSource), location);
53+
private ArrowReleaseCallback() {
54+
}
55+
56+
public static void execute(long releaseCallback, long baseStructure) {
57+
try {
58+
HANDLE.invokeExact(releaseCallback, baseStructure);
59+
} catch (Throwable e) {
60+
throw CompilerDirectives.shouldNotReachHere("Unable to call release callback. Error:", e);
61+
}
5562
}
5663
}

‎graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/arrow/InvokeArrowReleaseCallbackNode.java‎

Lines changed: 0 additions & 102 deletions
This file was deleted.

‎graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/arrow/capsule/ArrowArrayCapsuleDestructor.java‎

Lines changed: 46 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -40,60 +40,62 @@
4040
*/
4141
package com.oracle.graal.python.nodes.arrow.capsule;
4242

43+
import java.lang.invoke.MethodHandle;
44+
import java.lang.invoke.MethodHandles;
45+
import java.lang.invoke.MethodType;
46+
47+
import com.oracle.graal.python.annotations.CApiUpcallTarget;
4348
import com.oracle.graal.python.builtins.modules.cext.PythonCextCapsuleBuiltins.PyCapsuleGetPointerNode;
44-
import com.oracle.graal.python.builtins.objects.PNone;
4549
import com.oracle.graal.python.builtins.objects.cext.capi.transitions.CApiTransitions.NativeToPythonNode;
46-
import com.oracle.graal.python.runtime.nativeaccess.NativeMemory;
4750
import com.oracle.graal.python.nodes.arrow.ArrowArray;
48-
import com.oracle.graal.python.nodes.arrow.InvokeArrowReleaseCallbackNode;
51+
import com.oracle.graal.python.nodes.arrow.ArrowReleaseCallback;
4952
import com.oracle.graal.python.runtime.PythonContext;
50-
import com.oracle.truffle.api.CompilerDirectives;
51-
import com.oracle.truffle.api.dsl.Bind;
52-
import com.oracle.truffle.api.dsl.Cached;
53-
import com.oracle.truffle.api.interop.InteropLibrary;
54-
import com.oracle.truffle.api.interop.TruffleObject;
55-
import com.oracle.truffle.api.library.CachedLibrary;
56-
import com.oracle.truffle.api.library.ExportLibrary;
57-
import com.oracle.truffle.api.library.ExportMessage;
58-
import com.oracle.truffle.api.nodes.Node;
53+
import com.oracle.graal.python.runtime.nativeaccess.NativeMemory;
54+
import com.oracle.truffle.api.CompilerAsserts;
55+
56+
public final class ArrowArrayCapsuleDestructor {
57+
private static final MethodHandle HANDLE_EXECUTE;
5958

60-
@ExportLibrary(InteropLibrary.class)
61-
public class ArrowArrayCapsuleDestructor implements TruffleObject {
59+
static {
60+
try {
61+
HANDLE_EXECUTE = MethodHandles.lookup().findStatic(ArrowArrayCapsuleDestructor.class, "execute",
62+
MethodType.methodType(void.class, long.class));
63+
} catch (NoSuchMethodException | IllegalAccessException e) {
64+
throw new RuntimeException(e);
65+
}
66+
}
6267

63-
@ExportMessage
64-
boolean isExecutable() {
65-
return true;
68+
private ArrowArrayCapsuleDestructor() {
6669
}
6770

68-
@ExportMessage
69-
Object execute(Object[] args,
70-
@Bind Node inliningTarget,
71-
@CachedLibrary(limit = "1") InteropLibrary lib,
72-
@Cached NativeToPythonNode nativeToPythonNode,
73-
@Cached PyCapsuleGetPointerNode capsuleGetPointerNode,
74-
@Cached InvokeArrowReleaseCallbackNode.Lazy invokeReleaseCallbackNode) {
75-
if (args.length != 1 || !lib.isPointer(args[0])) {
76-
throw CompilerDirectives.shouldNotReachHere();
77-
}
71+
public static MethodHandle getMethodHandle() {
72+
return HANDLE_EXECUTE;
73+
}
7874

79-
Object capsule = nativeToPythonNode.execute(args[0]);
80-
PythonContext ctx = PythonContext.get(inliningTarget);
75+
@CApiUpcallTarget
76+
private static void execute(long capsulePointer) {
77+
CompilerAsserts.neverPartOfCompilation();
78+
PythonContext ctx = PythonContext.get(null);
8179
ctx.ensureNativeAccess();
82-
long capsuleNamePointer = ctx.stringToNativeUtf8Bytes(ArrowArray.CAPSULE_NAME, true);
83-
var arrowArray = ArrowArray.wrap(capsuleGetPointerNode.execute(inliningTarget, capsule, capsuleNamePointer));
84-
/*
85-
* The exported PyCapsules should have a destructor that calls the release callback of the
86-
* Arrow struct, if it is not already null. This prevents a memory leak in case the capsule
87-
* was never passed to another consumer.
88-
*
89-
* For more information see:
90-
* https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html#lifetime-
91-
* semantics
92-
*/
93-
if (!arrowArray.isReleased()) {
94-
invokeReleaseCallbackNode.get(inliningTarget).executeCached(arrowArray.releaseCallback(), arrowArray.memoryAddress());
80+
Object capsule = NativeToPythonNode.executeRawUncached(capsulePointer);
81+
long capsuleNamePointer = ctx.stringToNativeUtf8Bytes(ArrowArray.CAPSULE_NAME, false);
82+
try {
83+
var arrowArray = ArrowArray.wrap(PyCapsuleGetPointerNode.executeUncached(capsule, capsuleNamePointer));
84+
/*
85+
* The exported PyCapsules should have a destructor that calls the release callback of
86+
* the Arrow struct, if it is not already null. This prevents a memory leak in case the
87+
* capsule was never passed to another consumer.
88+
*
89+
* For more information see:
90+
* https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html#lifetime-
91+
* semantics
92+
*/
93+
if (!arrowArray.isReleased()) {
94+
ArrowReleaseCallback.execute(arrowArray.releaseCallback(), arrowArray.memoryAddress());
95+
}
96+
NativeMemory.free(arrowArray.memoryAddress());
97+
} finally {
98+
NativeMemory.free(capsuleNamePointer);
9599
}
96-
NativeMemory.free(arrowArray.memoryAddress());
97-
return PNone.NO_VALUE;
98100
}
99101
}

‎graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/arrow/capsule/ArrowSchemaCapsuleDestructor.java‎

Lines changed: 38 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -40,53 +40,55 @@
4040
*/
4141
package com.oracle.graal.python.nodes.arrow.capsule;
4242

43+
import java.lang.invoke.MethodHandle;
44+
import java.lang.invoke.MethodHandles;
45+
import java.lang.invoke.MethodType;
46+
47+
import com.oracle.graal.python.annotations.CApiUpcallTarget;
4348
import com.oracle.graal.python.builtins.modules.cext.PythonCextCapsuleBuiltins.PyCapsuleGetPointerNode;
44-
import com.oracle.graal.python.builtins.objects.PNone;
4549
import com.oracle.graal.python.builtins.objects.cext.capi.transitions.CApiTransitions.NativeToPythonNode;
46-
import com.oracle.graal.python.runtime.nativeaccess.NativeMemory;
50+
import com.oracle.graal.python.nodes.arrow.ArrowReleaseCallback;
4751
import com.oracle.graal.python.nodes.arrow.ArrowSchema;
48-
import com.oracle.graal.python.nodes.arrow.InvokeArrowReleaseCallbackNode;
4952
import com.oracle.graal.python.runtime.PythonContext;
50-
import com.oracle.truffle.api.CompilerDirectives;
51-
import com.oracle.truffle.api.dsl.Bind;
52-
import com.oracle.truffle.api.dsl.Cached;
53-
import com.oracle.truffle.api.interop.InteropLibrary;
54-
import com.oracle.truffle.api.interop.TruffleObject;
55-
import com.oracle.truffle.api.library.CachedLibrary;
56-
import com.oracle.truffle.api.library.ExportLibrary;
57-
import com.oracle.truffle.api.library.ExportMessage;
58-
import com.oracle.truffle.api.nodes.Node;
53+
import com.oracle.graal.python.runtime.nativeaccess.NativeMemory;
54+
import com.oracle.truffle.api.CompilerAsserts;
5955

60-
@ExportLibrary(InteropLibrary.class)
61-
public class ArrowSchemaCapsuleDestructor implements TruffleObject {
56+
public final class ArrowSchemaCapsuleDestructor {
57+
private static final MethodHandle HANDLE_EXECUTE;
6258

63-
@ExportMessage
64-
boolean isExecutable() {
65-
return true;
59+
static {
60+
try {
61+
HANDLE_EXECUTE = MethodHandles.lookup().findStatic(ArrowSchemaCapsuleDestructor.class, "execute",
62+
MethodType.methodType(void.class, long.class));
63+
} catch (NoSuchMethodException | IllegalAccessException e) {
64+
throw new RuntimeException(e);
65+
}
6666
}
6767

68-
@ExportMessage
69-
Object execute(Object[] args,
70-
@Bind Node inliningTarget,
71-
@CachedLibrary(limit = "1") InteropLibrary lib,
72-
@Cached NativeToPythonNode nativeToPythonNode,
73-
@Cached PyCapsuleGetPointerNode pyCapsuleGetPointerNode,
74-
@Cached InvokeArrowReleaseCallbackNode.Lazy invokeReleaseCallbackNode) {
75-
if (args.length != 1 || !lib.isPointer(args[0])) {
76-
throw CompilerDirectives.shouldNotReachHere();
77-
}
68+
private ArrowSchemaCapsuleDestructor() {
69+
}
70+
71+
public static MethodHandle getMethodHandle() {
72+
return HANDLE_EXECUTE;
73+
}
7874

79-
Object capsule = nativeToPythonNode.execute(args[0]);
80-
PythonContext ctx = PythonContext.get(inliningTarget);
75+
@CApiUpcallTarget
76+
private static void execute(long capsulePointer) {
77+
CompilerAsserts.neverPartOfCompilation();
78+
PythonContext ctx = PythonContext.get(null);
8179
ctx.ensureNativeAccess();
82-
long capsuleNamePointer = ctx.stringToNativeUtf8Bytes(ArrowSchema.CAPSULE_NAME, true);
83-
var arrowSchema = ArrowSchema.wrap(pyCapsuleGetPointerNode.execute(inliningTarget, capsule, capsuleNamePointer));
80+
Object capsule = NativeToPythonNode.executeRawUncached(capsulePointer);
81+
long capsuleNamePointer = ctx.stringToNativeUtf8Bytes(ArrowSchema.CAPSULE_NAME, false);
82+
try {
83+
var arrowSchema = ArrowSchema.wrap(PyCapsuleGetPointerNode.executeUncached(capsule, capsuleNamePointer));
8484

85-
if (!arrowSchema.isReleased()) {
86-
invokeReleaseCallbackNode.get(inliningTarget).executeCached(arrowSchema.releaseCallback(), arrowSchema.memoryAddress());
87-
}
85+
if (!arrowSchema.isReleased()) {
86+
ArrowReleaseCallback.execute(arrowSchema.releaseCallback(), arrowSchema.memoryAddress());
87+
}
8888

89-
NativeMemory.free(arrowSchema.memoryAddress());
90-
return PNone.NO_VALUE;
89+
NativeMemory.free(arrowSchema.memoryAddress());
90+
} finally {
91+
NativeMemory.free(capsuleNamePointer);
92+
}
9193
}
9294
}

0 commit comments

Comments
 (0)