[#1358] fix(spark): pre-check bytebuffer whether is direct before unc…

…ompress
apache · Dec 8, 2023 · e15a1c6 · e15a1c6
1 parent c5be58f
commit e15a1c6
Showing 1 changed file with 15 additions and 1 deletion.
diff --git a/...nt-spark/common/src/main/java/org/apache/spark/shuffle/reader/RssShuffleDataIterator.java b/...nt-spark/common/src/main/java/org/apache/spark/shuffle/reader/RssShuffleDataIterator.java
@@ -150,14 +150,28 @@ public boolean hasNext() {
     return recordsIterator.hasNext();
   }
 
+  private boolean isSameMemoryType(ByteBuffer left, ByteBuffer right) {
+    if (left.isDirect() && right.isDirect()) {
+      return true;
+    }
+
+    if (!left.isDirect() && !right.isDirect()) {
+      return true;
+    }
+
+    return false;
+  }
+
   private int uncompress(CompressedShuffleBlock rawBlock, ByteBuffer rawData) {
     long rawDataLength = rawData.limit() - rawData.position();
     totalRawBytesLength += rawDataLength;
     shuffleReadMetrics.incRemoteBytesRead(rawDataLength);
 
     int uncompressedLen = rawBlock.getUncompressLength();
     if (codec != null) {
-      if (uncompressedData == null || uncompressedData.capacity() < uncompressedLen) {
+      if (uncompressedData == null
+          || uncompressedData.capacity() < uncompressedLen
+          || !isSameMemoryType(uncompressedData, rawData)) {
         if (uncompressedData != null) {
           RssUtils.releaseByteBuffer(uncompressedData);
         }