org.apache.flink.runtime.io.network.buffer.BufferDecompressor Maven / Gradle / Ivy
/*
* 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.flink.runtime.io.network.buffer;
import org.apache.flink.annotation.VisibleForTesting;
import org.apache.flink.core.memory.MemorySegment;
import org.apache.flink.core.memory.MemorySegmentFactory;
import org.apache.flink.runtime.io.compression.BlockCompressionFactory;
import org.apache.flink.runtime.io.compression.BlockDecompressor;
import static org.apache.flink.util.Preconditions.checkArgument;
import static org.apache.flink.util.Preconditions.checkNotNull;
import static org.apache.flink.util.Preconditions.checkState;
/**
* Decompressor for compressed {@link Buffer}.
*/
public class BufferDecompressor {
/** The backing block decompressor for data decompression. */
private final BlockDecompressor blockDecompressor;
/** The intermediate buffer for the decompressed data. */
private final NetworkBuffer internalBuffer;
public BufferDecompressor(int bufferSize, String factoryName) {
checkArgument(bufferSize > 0);
checkNotNull(factoryName);
// the decompressed data size should be never larger than the configured buffer size
final byte[] heapBuffer = new byte[bufferSize];
this.internalBuffer = new NetworkBuffer(MemorySegmentFactory.wrap(heapBuffer), FreeingBufferRecycler.INSTANCE);
this.blockDecompressor = BlockCompressionFactory.createBlockCompressionFactory(factoryName).getDecompressor();
}
/**
* Decompresses the given {@link Buffer} using {@link BlockDecompressor}. The decompressed data will be stored
* in the intermediate buffer of this {@link BufferDecompressor} and returned to the caller. The caller must
* guarantee that the returned {@link Buffer} has been freed when calling the method next time.
*
* Notes that the decompression will always start from offset 0 to the size of the input {@link Buffer}.
*/
public Buffer decompressToIntermediateBuffer(Buffer buffer) {
int decompressedLen = decompress(buffer);
internalBuffer.setSize(decompressedLen);
return internalBuffer.retainBuffer();
}
/**
* The difference between this method and {@link #decompressToIntermediateBuffer(Buffer)} is that this method
* copies the decompressed data to the input {@link Buffer} starting from offset 0.
*
*
The caller must guarantee that the input {@link Buffer} is writable and there's enough space left.
*/
@VisibleForTesting
public Buffer decompressToOriginalBuffer(Buffer buffer) {
int decompressedLen = decompress(buffer);
// copy the decompressed data back
int memorySegmentOffset = buffer.getMemorySegmentOffset();
MemorySegment segment = buffer.getMemorySegment();
segment.put(memorySegmentOffset, internalBuffer.array(), 0, decompressedLen);
return new ReadOnlySlicedNetworkBuffer(buffer.asByteBuf(), 0, decompressedLen, memorySegmentOffset, false);
}
/**
* Decompresses the input {@link Buffer} into the intermediate buffer and returns the decompressed data size.
*/
private int decompress(Buffer buffer) {
checkArgument(buffer != null, "The input buffer must not be null.");
checkArgument(buffer.isBuffer(), "Event can not be decompressed.");
checkArgument(buffer.isCompressed(), "Buffer not compressed.");
checkArgument(buffer.getReaderIndex() == 0, "Reader index of the input buffer must be 0.");
checkArgument(buffer.readableBytes() > 0, "No data to be decompressed.");
checkState(internalBuffer.refCnt() == 1, "Illegal reference count, buffer need to be released.");
int length = buffer.getSize();
// decompress the given buffer into the internal heap buffer
return blockDecompressor.decompress(
buffer.getNioBuffer(0, length),
0,
length,
internalBuffer.getNioBuffer(0, internalBuffer.capacity()),
0);
}
}