Support the decompress feature in shuffle component. (#3)

JiaKe · liyafan82 · web-flow · commit 23939324f906 · 2021-03-03T13:46:05.000+08:00
* ARROW-10880: [Java] Support compressing RecordBatch IPC buffers by LZ4

* ARROW-10880: [Java] Support reading/writing big-endian message size

* ARROW-10880: [Java] Adjust variable names

* ARROW-10880: [Java] Support empty buffers

* ARROW-10880: [Java] Support passing raw data

* ARROW-10880: [Java] Switch to commons-compress library

* bug fix and support the fastpfor codec in the IPC framework

* update the access permission from private to protected

* disable the decompress function when loading the buffer

Co-authored-by: liyafan82 &lt;fan_li_ya@foxmail.com&gt;
diff --git a/cpp/src/arrow/ipc/metadata_internal.cc b/cpp/src/arrow/ipc/metadata_internal.cc
@@ -919,6 +919,8 @@ static Status GetBodyCompression(FBB& fbb, const IpcWriteOptions& options,
       codec = flatbuf::CompressionType::LZ4_FRAME;
     } else if (options.codec->compression_type() == Compression::ZSTD) {
       codec = flatbuf::CompressionType::ZSTD;
+    } else if (options.codec->compression_type() == Compression::FASTPFOR) {
+      codec = flatbuf::CompressionType::FASTPFOR;
     } else {
       return Status::Invalid("Unsupported IPC compression codec: ",
                              options.codec->name());
diff --git a/cpp/src/arrow/ipc/reader.cc b/cpp/src/arrow/ipc/reader.cc
@@ -609,6 +609,8 @@ Status GetCompression(const flatbuf::RecordBatch* batch, Compression::type* out)
       *out = Compression::LZ4_FRAME;
     } else if (compression->codec() == flatbuf::CompressionType::ZSTD) {
       *out = Compression::ZSTD;
+    } else if (compression->codec() == flatbuf::CompressionType::FASTPFOR) {
+      *out = Compression::FASTPFOR;
     } else {
       return Status::Invalid("Unsupported codec in RecordBatch::compression metadata");
     }
diff --git a/cpp/src/arrow/util/compression.cc b/cpp/src/arrow/util/compression.cc
@@ -88,7 +88,7 @@ Result<Compression::type> Codec::GetCompressionType(const std::string& name) {
     return Compression::ZSTD;
   } else if (name == "bz2") {
     return Compression::BZ2;
-  } else if (name == "FASTPFOR") {
+  } else if (name == "fastpfor") {
     return Compression::FASTPFOR;
   } else {
     return Status::Invalid("Unrecognized compression type: ", name);
diff --git a/cpp/src/generated/Message_generated.h b/cpp/src/generated/Message_generated.h
@@ -32,22 +32,25 @@ struct MessageBuilder;
 enum class CompressionType : int8_t {
   LZ4_FRAME = 0,
   ZSTD = 1,
+  FASTPFOR = 2,
   MIN = LZ4_FRAME,
   MAX = ZSTD
 };
 
-inline const CompressionType (&EnumValuesCompressionType())[2] {
+inline const CompressionType (&EnumValuesCompressionType())[3] {
   static const CompressionType values[] = {
     CompressionType::LZ4_FRAME,
-    CompressionType::ZSTD
+    CompressionType::ZSTD,
+    CompressionType::FASTPFOR
   };
   return values;
 }
 
 inline const char * const *EnumNamesCompressionType() {
-  static const char * const names[3] = {
+  static const char * const names[4] = {
     "LZ4_FRAME",
     "ZSTD",
+    "FASTPFOR",
     nullptr
   };
   return names;
diff --git a/format/Message.fbs b/format/Message.fbs
@@ -49,7 +49,8 @@ enum CompressionType:byte {
   LZ4_FRAME,
 
   // Zstandard
-  ZSTD
+  ZSTD,
+  FASTPFOR
 }
 
 /// Provided for forward compatibility in case we need to support different
diff --git a/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/MemoryUtil.java b/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/MemoryUtil.java
@@ -21,6 +21,7 @@
 import java.lang.reflect.Field;
 import java.lang.reflect.InvocationTargetException;
 import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
 import java.security.AccessController;
 import java.security.PrivilegedAction;
 
@@ -48,6 +49,11 @@ public class MemoryUtil {
    */
   static final long BYTE_BUFFER_ADDRESS_OFFSET;
 
+  /**
+   * If the native byte order is little-endian.
+   */
+  public static final boolean LITTLE_ENDIAN = ByteOrder.nativeOrder() == ByteOrder.LITTLE_ENDIAN;
+
   static {
     try {
       // try to get the unsafe object
@@ -132,7 +138,7 @@ public Object run() {
   }
 
   /**
-   * Given a {@link ByteBuf}, gets the address the underlying memory space.
+   * Given a {@link ByteBuffer}, gets the address the underlying memory space.
    *
    * @param buf the byte buffer.
    * @return address of the underlying memory.
diff --git a/java/vector/pom.xml b/java/vector/pom.xml
@@ -74,6 +74,11 @@
       <groupId>org.slf4j</groupId>
       <artifactId>slf4j-api</artifactId>
     </dependency>
+    <dependency>
+      <groupId>org.apache.commons</groupId>
+      <artifactId>commons-compress</artifactId>
+      <version>1.20</version>
+    </dependency>
   </dependencies>
 
   <pluginRepositories>
diff --git a/java/vector/src/main/java/org/apache/arrow/vector/VectorLoader.java b/java/vector/src/main/java/org/apache/arrow/vector/VectorLoader.java
@@ -56,9 +56,9 @@ public VectorLoader(VectorSchemaRoot root) {
   public void load(ArrowRecordBatch recordBatch) {
     Iterator<ArrowBuf> buffers = recordBatch.getBuffers().iterator();
     Iterator<ArrowFieldNode> nodes = recordBatch.getNodes().iterator();
-    CompressionCodec codec = CompressionUtil.createCodec(recordBatch.getBodyCompression().getCodec());
+
     for (FieldVector fieldVector : root.getFieldVectors()) {
-      loadBuffers(fieldVector, fieldVector.getField(), buffers, nodes, codec);
+      loadBuffers(fieldVector, fieldVector.getField(), buffers, nodes);
     }
     root.setRowCount(recordBatch.getLength());
     if (nodes.hasNext() || buffers.hasNext()) {
@@ -71,15 +71,14 @@ protected void loadBuffers(
       FieldVector vector,
       Field field,
       Iterator<ArrowBuf> buffers,
-      Iterator<ArrowFieldNode> nodes,
-      CompressionCodec codec) {
+      Iterator<ArrowFieldNode> nodes) {
     checkArgument(nodes.hasNext(), "no more field nodes for for field %s and vector %s", field, vector);
     ArrowFieldNode fieldNode = nodes.next();
     int bufferLayoutCount = TypeLayout.getTypeBufferCount(field.getType());
     List<ArrowBuf> ownBuffers = new ArrayList<>(bufferLayoutCount);
     for (int j = 0; j < bufferLayoutCount; j++) {
       ArrowBuf nextBuf = buffers.next();
-      ownBuffers.add(codec.decompress(vector.getAllocator(), nextBuf));
+      ownBuffers.add(nextBuf);
     }
     try {
       vector.loadFieldBuffers(fieldNode, ownBuffers);
@@ -96,7 +95,7 @@ protected void loadBuffers(
       for (int i = 0; i < childrenFromFields.size(); i++) {
         Field child = children.get(i);
         FieldVector fieldVector = childrenFromFields.get(i);
-        loadBuffers(fieldVector, child, buffers, nodes, codec);
+        loadBuffers(fieldVector, child, buffers, nodes);
       }
     }
   }
diff --git a/java/vector/src/main/java/org/apache/arrow/vector/compression/CompressionCodec.java b/java/vector/src/main/java/org/apache/arrow/vector/compression/CompressionCodec.java
@@ -28,11 +28,11 @@ public interface CompressionCodec {
   /**
    * Compress a buffer.
    * @param allocator the allocator for allocating memory for compressed buffer.
-   * @param unCompressedBuffer the buffer to compress.
+   * @param uncompressedBuffer the buffer to compress.
    *                           Implementation of this method should take care of releasing this buffer.
    * @return the compressed buffer.
    */
-  ArrowBuf compress(BufferAllocator allocator, ArrowBuf unCompressedBuffer);
+  ArrowBuf compress(BufferAllocator allocator, ArrowBuf uncompressedBuffer);
 
   /**
    * Decompress a buffer.
diff --git a/java/vector/src/main/java/org/apache/arrow/vector/compression/CompressionUtil.java b/java/vector/src/main/java/org/apache/arrow/vector/compression/CompressionUtil.java
@@ -19,13 +19,23 @@
 
 import org.apache.arrow.flatbuf.BodyCompressionMethod;
 import org.apache.arrow.flatbuf.CompressionType;
+import org.apache.arrow.memory.ArrowBuf;
+import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.vector.ipc.message.ArrowBodyCompression;
 
 /**
  * Utilities for data compression/decompression.
  */
 public class CompressionUtil {
 
+  static final long SIZE_OF_UNCOMPRESSED_LENGTH = 8L;
+
+  /**
+   * Special flag to indicate no compression.
+   * (e.g. when the compressed buffer has a larger size.)
+   */
+  static final long NO_COMPRESSION_LENGTH = -1L;
+
   private CompressionUtil() {
   }
 
@@ -53,8 +63,29 @@ public static CompressionCodec createCodec(byte compressionType) {
     switch (compressionType) {
       case NoCompressionCodec.COMPRESSION_TYPE:
         return NoCompressionCodec.INSTANCE;
+      case CompressionType.LZ4_FRAME:
+        return new Lz4CompressionCodec();
       default:
         throw new IllegalArgumentException("Compression type not supported: " + compressionType);
     }
   }
+
+  /**
+   * Process compression by compressing the buffer as is.
+   */
+  public static ArrowBuf compressRawBuffer(BufferAllocator allocator, ArrowBuf inputBuffer) {
+    ArrowBuf compressedBuffer = allocator.buffer(SIZE_OF_UNCOMPRESSED_LENGTH + inputBuffer.writerIndex());
+    compressedBuffer.setLong(0, NO_COMPRESSION_LENGTH);
+    compressedBuffer.setBytes(SIZE_OF_UNCOMPRESSED_LENGTH, inputBuffer, 0, inputBuffer.writerIndex());
+    compressedBuffer.writerIndex(SIZE_OF_UNCOMPRESSED_LENGTH + inputBuffer.writerIndex());
+    return compressedBuffer;
+  }
+
+  /**
+   * Process decompression by decompressing the buffer as is.
+   */
+  public static ArrowBuf decompressRawBuffer(ArrowBuf inputBuffer) {
+    return inputBuffer.slice(SIZE_OF_UNCOMPRESSED_LENGTH,
+        inputBuffer.writerIndex() - SIZE_OF_UNCOMPRESSED_LENGTH);
+  }
 }
diff --git a/java/vector/src/main/java/org/apache/arrow/vector/compression/Lz4CompressionCodec.java b/java/vector/src/main/java/org/apache/arrow/vector/compression/Lz4CompressionCodec.java
@@ -0,0 +1,160 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.arrow.vector.compression;
+
+import static org.apache.arrow.memory.util.MemoryUtil.LITTLE_ENDIAN;
+import static org.apache.arrow.vector.compression.CompressionUtil.NO_COMPRESSION_LENGTH;
+import static org.apache.arrow.vector.compression.CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+
+import org.apache.arrow.flatbuf.CompressionType;
+import org.apache.arrow.memory.ArrowBuf;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.util.Preconditions;
+import org.apache.commons.compress.compressors.lz4.FramedLZ4CompressorInputStream;
+import org.apache.commons.compress.compressors.lz4.FramedLZ4CompressorOutputStream;
+import org.apache.commons.compress.utils.IOUtils;
+
+import io.netty.util.internal.PlatformDependent;
+
+/**
+ * Compression codec for the LZ4 algorithm.
+ */
+public class Lz4CompressionCodec implements CompressionCodec {
+
+  @Override
+  public ArrowBuf compress(BufferAllocator allocator, ArrowBuf uncompressedBuffer) {
+    Preconditions.checkArgument(uncompressedBuffer.writerIndex() <= Integer.MAX_VALUE,
+        "The uncompressed buffer size exceeds the integer limit");
+
+    if (uncompressedBuffer.writerIndex() == 0L) {
+      // shortcut for empty buffer
+      ArrowBuf compressedBuffer = allocator.buffer(SIZE_OF_UNCOMPRESSED_LENGTH);
+      compressedBuffer.setLong(0, 0);
+      compressedBuffer.writerIndex(SIZE_OF_UNCOMPRESSED_LENGTH);
+      uncompressedBuffer.close();
+      return compressedBuffer;
+    }
+
+    try {
+      ArrowBuf compressedBuffer = doCompress(allocator, uncompressedBuffer);
+      long compressedLength = compressedBuffer.writerIndex() - SIZE_OF_UNCOMPRESSED_LENGTH;
+      if (compressedLength > uncompressedBuffer.writerIndex()) {
+        // compressed buffer is larger, send the raw buffer
+        compressedBuffer.close();
+        compressedBuffer = CompressionUtil.compressRawBuffer(allocator, uncompressedBuffer);
+      }
+
+      uncompressedBuffer.close();
+      return compressedBuffer;
+    } catch (IOException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  private ArrowBuf doCompress(BufferAllocator allocator, ArrowBuf uncompressedBuffer) throws IOException {
+    byte[] inBytes = new byte[(int) uncompressedBuffer.writerIndex()];
+    PlatformDependent.copyMemory(uncompressedBuffer.memoryAddress(), inBytes, 0, uncompressedBuffer.writerIndex());
+    ByteArrayOutputStream baos = new ByteArrayOutputStream();
+    try (InputStream in = new ByteArrayInputStream(inBytes);
+         OutputStream out = new FramedLZ4CompressorOutputStream(baos)) {
+      IOUtils.copy(in, out);
+    }
+
+    byte[] outBytes = baos.toByteArray();
+
+    ArrowBuf compressedBuffer = allocator.buffer(SIZE_OF_UNCOMPRESSED_LENGTH + outBytes.length);
+
+    long uncompressedLength = uncompressedBuffer.writerIndex();
+    if (!LITTLE_ENDIAN) {
+      uncompressedLength = Long.reverseBytes(uncompressedLength);
+    }
+    // first 8 bytes reserved for uncompressed length, to be consistent with the
+    // C++ implementation.
+    compressedBuffer.setLong(0, uncompressedLength);
+
+    PlatformDependent.copyMemory(
+        outBytes, 0, compressedBuffer.memoryAddress() + SIZE_OF_UNCOMPRESSED_LENGTH, outBytes.length);
+    compressedBuffer.writerIndex(SIZE_OF_UNCOMPRESSED_LENGTH + outBytes.length);
+    return compressedBuffer;
+  }
+
+  @Override
+  public ArrowBuf decompress(BufferAllocator allocator, ArrowBuf compressedBuffer) {
+    Preconditions.checkArgument(compressedBuffer.writerIndex() <= Integer.MAX_VALUE,
+        "The compressed buffer size exceeds the integer limit");
+
+    Preconditions.checkArgument(compressedBuffer.writerIndex() >= SIZE_OF_UNCOMPRESSED_LENGTH,
+        "Not enough data to decompress.");
+
+    long decompressedLength = compressedBuffer.getLong(0);
+    if (!LITTLE_ENDIAN) {
+      decompressedLength = Long.reverseBytes(decompressedLength);
+    }
+
+    if (decompressedLength == 0L) {
+      // shortcut for empty buffer
+      compressedBuffer.close();
+      return allocator.getEmpty();
+    }
+
+    if (decompressedLength == NO_COMPRESSION_LENGTH) {
+      // no compression
+      return CompressionUtil.decompressRawBuffer(compressedBuffer);
+    }
+
+    try {
+      ArrowBuf decompressedBuffer = doDecompress(allocator, compressedBuffer);
+      compressedBuffer.close();
+      return decompressedBuffer;
+    } catch (IOException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  private ArrowBuf doDecompress(BufferAllocator allocator, ArrowBuf compressedBuffer) throws IOException {
+    long decompressedLength = compressedBuffer.getLong(0);
+    if (!LITTLE_ENDIAN) {
+      decompressedLength = Long.reverseBytes(decompressedLength);
+    }
+
+    byte[] inBytes = new byte[(int) (compressedBuffer.writerIndex() - SIZE_OF_UNCOMPRESSED_LENGTH)];
+    PlatformDependent.copyMemory(
+        compressedBuffer.memoryAddress() + SIZE_OF_UNCOMPRESSED_LENGTH, inBytes, 0, inBytes.length);
+    ByteArrayOutputStream out = new ByteArrayOutputStream((int) decompressedLength);
+    try (InputStream in = new FramedLZ4CompressorInputStream(new ByteArrayInputStream(inBytes))) {
+      IOUtils.copy(in, out);
+    }
+
+    byte[] outBytes = out.toByteArray();
+    ArrowBuf decompressedBuffer = allocator.buffer(outBytes.length);
+    PlatformDependent.copyMemory(outBytes, 0, decompressedBuffer.memoryAddress(), outBytes.length);
+    decompressedBuffer.writerIndex(decompressedLength);
+    return decompressedBuffer;
+  }
+
+  @Override
+  public String getCodecName() {
+    return CompressionType.name(CompressionType.LZ4_FRAME);
+  }
+}
diff --git a/java/vector/src/main/java/org/apache/arrow/vector/compression/NoCompressionCodec.java b/java/vector/src/main/java/org/apache/arrow/vector/compression/NoCompressionCodec.java
@@ -38,8 +38,8 @@ private NoCompressionCodec() {
   }
 
   @Override
-  public ArrowBuf compress(BufferAllocator allocator, ArrowBuf unCompressedBuffer) {
-    return unCompressedBuffer;
+  public ArrowBuf compress(BufferAllocator allocator, ArrowBuf uncompressedBuffer) {
+    return uncompressedBuffer;
   }
 
   @Override
diff --git a/java/vector/src/main/java/org/apache/arrow/vector/ipc/ArrowStreamReader.java b/java/vector/src/main/java/org/apache/arrow/vector/ipc/ArrowStreamReader.java
diff --git a/java/vector/src/test/java/org/apache/arrow/vector/compression/TestCompressionCodec.java b/java/vector/src/test/java/org/apache/arrow/vector/compression/TestCompressionCodec.java

Original file line number	Diff line number	Diff line change
`@@ -609,6 +609,8 @@ Status GetCompression(const flatbuf::RecordBatch* batch, Compression::type* out)`
`609`	`609`	`*out = Compression::LZ4_FRAME;`
`610`	`610`	`} else if (compression->codec() == flatbuf::CompressionType::ZSTD) {`
`611`	`611`	`*out = Compression::ZSTD;`
	`612`	`+ } else if (compression->codec() == flatbuf::CompressionType::FASTPFOR) {`
	`613`	`+ *out = Compression::FASTPFOR;`
`612`	`614`	`} else {`
`613`	`615`	`return Status::Invalid("Unsupported codec in RecordBatch::compression metadata");`
`614`	`616`	`}`
Original file line number	Diff line number	Diff line change
`@@ -49,7 +49,8 @@ enum CompressionType:byte {`
`49`	`49`	`LZ4_FRAME,`
`50`	`50`
`51`	`51`	`// Zstandard`
`52`		`- ZSTD`
	`52`	`+ ZSTD,`
	`53`	`+ FASTPFOR`
`53`	`54`	`}`
`54`	`55`
`55`	`56`	`/// Provided for forward compatibility in case we need to support different`