address comments

apache · Apr 12, 2023 · 5fab041 · 5fab041
1 parent 77d8afb
commit 5fab041
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 7 deletions.
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
@@ -145,9 +145,6 @@ private int uncompress(CompressedShuffleBlock rawBlock, ByteBuffer rawData) {
 
     int uncompressedLen = rawBlock.getUncompressLength();
     if (codec != null) {
-      if (uncompressedData != null) {
-        LOG.warn("UncompressedData: " + uncompressedLen + " Capacity: " + uncompressedData.capacity());
-      }
       if (uncompressedData == null || uncompressedData.capacity() < uncompressedLen) {
         if (uncompressedData != null) {
           RssUtils.releaseByteBuffer(uncompressedData);
@@ -161,12 +158,12 @@ private int uncompress(CompressedShuffleBlock rawBlock, ByteBuffer rawData) {
       unCompressedBytesLength += uncompressedLen;
       long decompressDuration = System.currentTimeMillis() - startDecompress;
       decompressTime += decompressDuration;
+      // ByteBuffer's limit may not uncompressDataLength after using compress method.
+      // So we need set limit here.
+      uncompressedData.limit(uncompressedData.position() + uncompressedLen);
     } else {
       uncompressedData = rawData;
     }
-    // ByteBuffer's limit may not uncompressDataLength after using compress method.
-    // So we need set limit here.
-    uncompressedData.limit(uncompressedData.position() + uncompressedLen);
     return uncompressedLen;
   }
 

diff --git a/storage/src/main/java/org/apache/uniffle/storage/handler/impl/HdfsFileReader.java b/storage/src/main/java/org/apache/uniffle/storage/handler/impl/HdfsFileReader.java
@@ -98,7 +98,7 @@ public ByteBuffer readAsByteBuffer(long offset, int length) {
   public ByteBuffer readAsByteBuffer() {
     try {
       long length = getFileLen();
-      if (length > Integer.MAX_VALUE) {
+      if (length - fsDataInputStream.getPos() > Integer.MAX_VALUE) {
         LOG.warn("File " + path + "length is too long");
         return ByteBuffer.allocateDirect(0);
       }