[xgb] Fix xgb intern to close replaced array (#3558)

* [xgb] Fix xgb intern to close replaced array * [xgb] Fix xgb intern to close replaced array * Review changes * Review changes --------- Co-authored-by: Xin Yang <xyang19@gmail.com>
deepjavalibrary · Jan 22, 2025 · 0ea9091 · 0ea9091
1 parent a4e61ea
commit 0ea9091
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 2 deletions.
diff --git a/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java b/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
@@ -40,7 +40,7 @@ public abstract class NDArrayAdapter implements NDArray {
 
     protected NDManager manager;
     protected NDManager alternativeManager;
-    private NDArray alternativeArray;
+    protected NDArray alternativeArray;
 
     protected Shape shape;
     protected DataType dataType;

diff --git a/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbNDArray.java b/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbNDArray.java
@@ -88,14 +88,33 @@ public ByteBuffer toByteBuffer(boolean tryDirect) {
     /** {@inheritDoc} */
     @Override
     public void intern(NDArray replaced) {
+        if (!(replaced instanceof XgbNDArray)) {
+            throw new IllegalArgumentException(
+                    "The replaced NDArray must be an instance of XgbNDArray.");
+        }
+        XgbNDArray array = (XgbNDArray) replaced;
+        if (isReleased()) {
+            throw new IllegalArgumentException("This array is already closed");
+        }
+        if (replaced.isReleased()) {
+            throw new IllegalArgumentException("This target array is already closed");
+        }
+
         if (handle != null && handle.get() != 0L) {
             long pointer = handle.getAndSet(0L);
             JniUtils.deleteDMatrix(pointer);
         }
-        XgbNDArray array = (XgbNDArray) replaced;
+        if (alternativeArray != null) {
+            alternativeArray.close();
+        }
+
         data = array.data;
         handle = array.handle;
         format = array.format;
+        alternativeArray = array.alternativeArray;
+        array.handle = null;
+        array.alternativeArray = null;
+        array.close();
     }
 
     /** {@inheritDoc} */