[GR-48910] Introduce enable_java_integration method.

horakivo · steve-s · commit 201f7c85721a · 2025-01-23T06:25:04.000Z
PullRequest: graalpython/3656
diff --git a/graalpython/lib-graalpython/modules/_polyglot_arrow.py b/graalpython/lib-graalpython/modules/_polyglot_arrow.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2024, Oracle and/or its affiliates. All rights reserved.
+# Copyright (c) 2024, 2025, Oracle and/or its affiliates. All rights reserved.
 # DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 #
 # The Universal Permissive License (UPL), Version 1.0
@@ -43,102 +43,124 @@
 try:
     java.type("org.apache.arrow.vector.BaseFixedWidthVector")
 except KeyError:
-    raise ImportError("It is not possible to import Apache Arrow Vector classes because arrow-vector package is not on the class path. Please add this library to your project.")
+    raise ImportError(
+        "It is not possible to import Apache Arrow Vector classes because arrow-vector package is not on the class path. Please add this library to your project.")
 
 if not __graalpython__.host_import_enabled:
     raise NotImplementedError("Host lookup is not allowed. You can allow it while building python context.")
-else:
-    class TinyIntVector:
 
-        def __len__(self):
-            return self.getValueCount()
 
-        def __arrow_c_array__(self, requested_schema=None):
-            return __graalpython__.export_arrow_vector(self)
+class TinyIntVector:
 
+    def __len__(self):
+        return self.getValueCount()
 
-    class SmallIntVector:
+    def __arrow_c_array__(self, requested_schema=None):
+        return __graalpython__.export_arrow_vector(self)
 
-        def __len__(self):
-            return self.getValueCount()
 
-        def __arrow_c_array__(self, requested_schema=None):
-            return __graalpython__.export_arrow_vector(self)
+class SmallIntVector:
 
+    def __len__(self):
+        return self.getValueCount()
 
-    class IntVector:
+    def __arrow_c_array__(self, requested_schema=None):
+        return __graalpython__.export_arrow_vector(self)
 
-        def __len__(self):
-            return self.getValueCount()
 
-        def __arrow_c_array__(self, requested_schema=None):
-            return __graalpython__.export_arrow_vector(self)
+class IntVector:
 
+    def __len__(self):
+        return self.getValueCount()
 
-    class BigIntVector:
+    def __arrow_c_array__(self, requested_schema=None):
+        return __graalpython__.export_arrow_vector(self)
 
-        def __len__(self):
-            return self.getValueCount()
 
-        def __arrow_c_array__(self, requested_schema=None):
-            return __graalpython__.export_arrow_vector(self)
+class BigIntVector:
 
+    def __len__(self):
+        return self.getValueCount()
 
-    class BitVector:
+    def __arrow_c_array__(self, requested_schema=None):
+        return __graalpython__.export_arrow_vector(self)
 
-        def __len__(self):
-            return self.getValueCount()
 
-        def __arrow_c_array__(self, requested_schema=None):
-            return __graalpython__.export_arrow_vector(self)
+class BitVector:
 
+    def __len__(self):
+        return self.getValueCount()
 
-    class Float2Vector:
+    def __arrow_c_array__(self, requested_schema=None):
+        return __graalpython__.export_arrow_vector(self)
 
-        def __len__(self):
-            return self.getValueCount()
 
-        def __arrow_c_array__(self, requested_schema=None):
-            return __graalpython__.export_arrow_vector(self)
+class Float2Vector:
 
+    def __len__(self):
+        return self.getValueCount()
 
-    class Float4Vector:
+    def __arrow_c_array__(self, requested_schema=None):
+        return __graalpython__.export_arrow_vector(self)
 
-        def __len__(self):
-            return self.getValueCount()
 
-        def __arrow_c_array__(self, requested_schema=None):
-            return __graalpython__.export_arrow_vector(self)
+class Float4Vector:
 
+    def __len__(self):
+        return self.getValueCount()
 
-    class Float8Vector:
+    def __arrow_c_array__(self, requested_schema=None):
+        return __graalpython__.export_arrow_vector(self)
 
-        def __len__(self):
-            return self.getValueCount()
 
-        def __arrow_c_array__(self, requested_schema=None):
-            return __graalpython__.export_arrow_vector(self)
+class Float8Vector:
 
+    def __len__(self):
+        return self.getValueCount()
 
-    # Ints
-    int8_vector_class_path = java.type("org.apache.arrow.vector.TinyIntVector")
-    int16_vector_class_path = java.type("org.apache.arrow.vector.SmallIntVector")
-    int32_vector_class_path = java.type("org.apache.arrow.vector.IntVector")
-    int64_vector_class_path = java.type("org.apache.arrow.vector.BigIntVector")
-    # Boolean
-    boolean_vector_class_path = java.type("org.apache.arrow.vector.BitVector")
-    # Floats
-    float2_vector_class_path = java.type("org.apache.arrow.vector.Float2Vector")
-    float4_vector_class_path = java.type("org.apache.arrow.vector.Float4Vector")
-    float8_vector_class_path = java.type("org.apache.arrow.vector.Float8Vector")
+    def __arrow_c_array__(self, requested_schema=None):
+        return __graalpython__.export_arrow_vector(self)
 
-    polyglot.register_interop_type(int8_vector_class_path, TinyIntVector)
-    polyglot.register_interop_type(int16_vector_class_path, SmallIntVector)
-    polyglot.register_interop_type(int32_vector_class_path, IntVector)
-    polyglot.register_interop_type(int64_vector_class_path, BigIntVector)
 
-    polyglot.register_interop_type(boolean_vector_class_path, BitVector)
+__enabled_java_integration = False
 
-    polyglot.register_interop_type(float2_vector_class_path, Float2Vector)
-    polyglot.register_interop_type(float4_vector_class_path, Float4Vector)
-    polyglot.register_interop_type(float8_vector_class_path, Float8Vector)
+
+def enable_java_integration(allow_method_overwrites: bool = False):
+    """
+    This method enables passing Java Apache Arrow vector classes directly to Python code without copying any memory.
+    It basically calls `polyglot.register_interop_type` on every vector class defined in the library.
+    Calling the method more than once has no effect.
+
+    If allow_method_overwrites=True, defining the same method is explicitly allowed.
+    """
+    global __enabled_java_integration
+    if not __enabled_java_integration:
+        __enabled_java_integration = True
+        # Ints
+        int8_vector_class_path = java.type("org.apache.arrow.vector.TinyIntVector")
+        int16_vector_class_path = java.type("org.apache.arrow.vector.SmallIntVector")
+        int32_vector_class_path = java.type("org.apache.arrow.vector.IntVector")
+        int64_vector_class_path = java.type("org.apache.arrow.vector.BigIntVector")
+        # Boolean
+        boolean_vector_class_path = java.type("org.apache.arrow.vector.BitVector")
+        # Floats
+        float2_vector_class_path = java.type("org.apache.arrow.vector.Float2Vector")
+        float4_vector_class_path = java.type("org.apache.arrow.vector.Float4Vector")
+        float8_vector_class_path = java.type("org.apache.arrow.vector.Float8Vector")
+
+        polyglot.register_interop_type(int8_vector_class_path, TinyIntVector,
+                                       allow_method_overwrites=allow_method_overwrites)
+        polyglot.register_interop_type(int16_vector_class_path, SmallIntVector,
+                                       allow_method_overwrites=allow_method_overwrites)
+        polyglot.register_interop_type(int32_vector_class_path, IntVector,
+                                       allow_method_overwrites=allow_method_overwrites)
+        polyglot.register_interop_type(int64_vector_class_path, BigIntVector,
+                                       allow_method_overwrites=allow_method_overwrites)
+        polyglot.register_interop_type(boolean_vector_class_path, BitVector,
+                                       allow_method_overwrites=allow_method_overwrites)
+        polyglot.register_interop_type(float2_vector_class_path, Float2Vector,
+                                       allow_method_overwrites=allow_method_overwrites)
+        polyglot.register_interop_type(float4_vector_class_path, Float4Vector,
+                                       allow_method_overwrites=allow_method_overwrites)
+        polyglot.register_interop_type(float8_vector_class_path, Float8Vector,
+                                       allow_method_overwrites=allow_method_overwrites)
