| /* |
| * 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.compression; |
| |
| import static org.junit.jupiter.api.Assertions.assertArrayEquals; |
| import static org.junit.jupiter.api.Assertions.assertEquals; |
| import static org.junit.jupiter.api.Assertions.assertFalse; |
| import static org.junit.jupiter.api.Assertions.assertNotEquals; |
| import static org.junit.jupiter.api.Assertions.assertTrue; |
| |
| import java.io.ByteArrayOutputStream; |
| import java.io.IOException; |
| import java.nio.channels.Channels; |
| import java.nio.charset.StandardCharsets; |
| import java.util.ArrayList; |
| import java.util.Arrays; |
| import java.util.Collection; |
| import java.util.HashMap; |
| import java.util.List; |
| import java.util.Optional; |
| import java.util.function.BiConsumer; |
| import java.util.stream.Stream; |
| import org.apache.arrow.memory.ArrowBuf; |
| import org.apache.arrow.memory.BufferAllocator; |
| import org.apache.arrow.memory.RootAllocator; |
| import org.apache.arrow.util.AutoCloseables; |
| import org.apache.arrow.vector.IntVector; |
| import org.apache.arrow.vector.VarBinaryVector; |
| import org.apache.arrow.vector.VarCharVector; |
| import org.apache.arrow.vector.VectorSchemaRoot; |
| import org.apache.arrow.vector.VectorUnloader; |
| import org.apache.arrow.vector.compression.CompressionCodec; |
| import org.apache.arrow.vector.compression.CompressionUtil; |
| import org.apache.arrow.vector.compression.NoCompressionCodec; |
| import org.apache.arrow.vector.dictionary.DictionaryProvider; |
| import org.apache.arrow.vector.ipc.ArrowFileReader; |
| import org.apache.arrow.vector.ipc.ArrowFileWriter; |
| import org.apache.arrow.vector.ipc.ArrowStreamReader; |
| import org.apache.arrow.vector.ipc.ArrowStreamWriter; |
| import org.apache.arrow.vector.ipc.message.ArrowFieldNode; |
| import org.apache.arrow.vector.ipc.message.IpcOption; |
| import org.apache.arrow.vector.types.pojo.ArrowType; |
| import org.apache.arrow.vector.types.pojo.Field; |
| import org.apache.arrow.vector.types.pojo.Schema; |
| import org.apache.arrow.vector.util.ByteArrayReadableSeekableByteChannel; |
| import org.junit.jupiter.api.AfterEach; |
| import org.junit.jupiter.api.BeforeEach; |
| import org.junit.jupiter.params.ParameterizedTest; |
| import org.junit.jupiter.params.provider.Arguments; |
| import org.junit.jupiter.params.provider.MethodSource; |
| |
| /** Test cases for {@link CompressionCodec}s. */ |
| class TestCompressionCodec { |
| private BufferAllocator allocator; |
| |
| @BeforeEach |
| void init() { |
| allocator = new RootAllocator(Integer.MAX_VALUE); |
| } |
| |
| @AfterEach |
| void terminate() { |
| allocator.close(); |
| } |
| |
| static Collection<Arguments> codecs() { |
| List<Arguments> params = new ArrayList<>(); |
| |
| int[] lengths = new int[] {10, 100, 1000}; |
| for (int len : lengths) { |
| CompressionCodec dumbCodec = NoCompressionCodec.INSTANCE; |
| params.add(Arguments.arguments(len, dumbCodec)); |
| |
| CompressionCodec lz4Codec = new Lz4CompressionCodec(); |
| params.add(Arguments.arguments(len, lz4Codec)); |
| |
| CompressionCodec zstdCodec = new ZstdCompressionCodec(); |
| params.add(Arguments.arguments(len, zstdCodec)); |
| |
| CompressionCodec zstdCodecAndCompressionLevel = new ZstdCompressionCodec(7); |
| params.add(Arguments.arguments(len, zstdCodecAndCompressionLevel)); |
| } |
| return params; |
| } |
| |
| private List<ArrowBuf> compressBuffers(CompressionCodec codec, List<ArrowBuf> inputBuffers) { |
| List<ArrowBuf> outputBuffers = new ArrayList<>(inputBuffers.size()); |
| for (ArrowBuf buf : inputBuffers) { |
| outputBuffers.add(codec.compress(allocator, buf)); |
| } |
| return outputBuffers; |
| } |
| |
| private List<ArrowBuf> deCompressBuffers(CompressionCodec codec, List<ArrowBuf> inputBuffers) { |
| List<ArrowBuf> outputBuffers = new ArrayList<>(inputBuffers.size()); |
| for (ArrowBuf buf : inputBuffers) { |
| outputBuffers.add(codec.decompress(allocator, buf)); |
| } |
| return outputBuffers; |
| } |
| |
| private void assertWriterIndex(List<ArrowBuf> decompressedBuffers) { |
| for (ArrowBuf decompressedBuf : decompressedBuffers) { |
| assertTrue(decompressedBuf.writerIndex() > 0); |
| } |
| } |
| |
| @ParameterizedTest |
| @MethodSource("codecs") |
| void testCompressFixedWidthBuffers(int vectorLength, CompressionCodec codec) throws Exception { |
| // prepare vector to compress |
| IntVector origVec = new IntVector("vec", allocator); |
| origVec.allocateNew(vectorLength); |
| for (int i = 0; i < vectorLength; i++) { |
| if (i % 10 == 0) { |
| origVec.setNull(i); |
| } else { |
| origVec.set(i, i); |
| } |
| } |
| origVec.setValueCount(vectorLength); |
| int nullCount = origVec.getNullCount(); |
| |
| // compress & decompress |
| List<ArrowBuf> origBuffers = origVec.getFieldBuffers(); |
| List<ArrowBuf> compressedBuffers = compressBuffers(codec, origBuffers); |
| List<ArrowBuf> decompressedBuffers = deCompressBuffers(codec, compressedBuffers); |
| |
| assertEquals(2, decompressedBuffers.size()); |
| assertWriterIndex(decompressedBuffers); |
| |
| // orchestrate new vector |
| IntVector newVec = new IntVector("new vec", allocator); |
| newVec.loadFieldBuffers(new ArrowFieldNode(vectorLength, nullCount), decompressedBuffers); |
| |
| // verify new vector |
| assertEquals(vectorLength, newVec.getValueCount()); |
| for (int i = 0; i < vectorLength; i++) { |
| if (i % 10 == 0) { |
| assertTrue(newVec.isNull(i)); |
| } else { |
| assertEquals(i, newVec.get(i)); |
| } |
| } |
| |
| newVec.close(); |
| AutoCloseables.close(decompressedBuffers); |
| } |
| |
| @ParameterizedTest |
| @MethodSource("codecs") |
| void testCompressVariableWidthBuffers(int vectorLength, CompressionCodec codec) throws Exception { |
| // prepare vector to compress |
| VarCharVector origVec = new VarCharVector("vec", allocator); |
| origVec.allocateNew(); |
| for (int i = 0; i < vectorLength; i++) { |
| if (i % 10 == 0) { |
| origVec.setNull(i); |
| } else { |
| origVec.setSafe(i, String.valueOf(i).getBytes(StandardCharsets.UTF_8)); |
| } |
| } |
| origVec.setValueCount(vectorLength); |
| int nullCount = origVec.getNullCount(); |
| |
| // compress & decompress |
| List<ArrowBuf> origBuffers = origVec.getFieldBuffers(); |
| List<ArrowBuf> compressedBuffers = compressBuffers(codec, origBuffers); |
| List<ArrowBuf> decompressedBuffers = deCompressBuffers(codec, compressedBuffers); |
| |
| assertEquals(3, decompressedBuffers.size()); |
| assertWriterIndex(decompressedBuffers); |
| |
| // orchestrate new vector |
| VarCharVector newVec = new VarCharVector("new vec", allocator); |
| newVec.loadFieldBuffers(new ArrowFieldNode(vectorLength, nullCount), decompressedBuffers); |
| |
| // verify new vector |
| assertEquals(vectorLength, newVec.getValueCount()); |
| for (int i = 0; i < vectorLength; i++) { |
| if (i % 10 == 0) { |
| assertTrue(newVec.isNull(i)); |
| } else { |
| assertArrayEquals(String.valueOf(i).getBytes(StandardCharsets.UTF_8), newVec.get(i)); |
| } |
| } |
| |
| newVec.close(); |
| AutoCloseables.close(decompressedBuffers); |
| } |
| |
| @ParameterizedTest |
| @MethodSource("codecs") |
| void testEmptyBuffer(int vectorLength, CompressionCodec codec) throws Exception { |
| final VarBinaryVector origVec = new VarBinaryVector("vec", allocator); |
| |
| origVec.allocateNew(vectorLength); |
| |
| // Do not set any values (all missing) |
| origVec.setValueCount(vectorLength); |
| |
| final List<ArrowBuf> origBuffers = origVec.getFieldBuffers(); |
| final List<ArrowBuf> compressedBuffers = compressBuffers(codec, origBuffers); |
| final List<ArrowBuf> decompressedBuffers = deCompressBuffers(codec, compressedBuffers); |
| |
| // orchestrate new vector |
| VarBinaryVector newVec = new VarBinaryVector("new vec", allocator); |
| newVec.loadFieldBuffers(new ArrowFieldNode(vectorLength, vectorLength), decompressedBuffers); |
| |
| // verify new vector |
| assertEquals(vectorLength, newVec.getValueCount()); |
| for (int i = 0; i < vectorLength; i++) { |
| assertTrue(newVec.isNull(i)); |
| } |
| |
| newVec.close(); |
| AutoCloseables.close(decompressedBuffers); |
| } |
| |
| private static Stream<CompressionUtil.CodecType> codecTypes() { |
| return Arrays.stream(CompressionUtil.CodecType.values()); |
| } |
| |
| @ParameterizedTest |
| @MethodSource("codecTypes") |
| void testReadWriteStream(CompressionUtil.CodecType codec) throws Exception { |
| withRoot( |
| codec, |
| (factory, root) -> { |
| ByteArrayOutputStream compressedStream = new ByteArrayOutputStream(); |
| try (final ArrowStreamWriter writer = |
| new ArrowStreamWriter( |
| root, |
| new DictionaryProvider.MapDictionaryProvider(), |
| Channels.newChannel(compressedStream), |
| IpcOption.DEFAULT, |
| factory, |
| codec, |
| Optional.of(7))) { |
| writer.start(); |
| writer.writeBatch(); |
| writer.end(); |
| } catch (IOException e) { |
| throw new RuntimeException(e); |
| } |
| |
| try (ArrowStreamReader reader = |
| new ArrowStreamReader( |
| new ByteArrayReadableSeekableByteChannel(compressedStream.toByteArray()), |
| allocator, |
| factory)) { |
| assertTrue(reader.loadNextBatch()); |
| assertTrue(root.equals(reader.getVectorSchemaRoot())); |
| assertFalse(reader.loadNextBatch()); |
| } catch (IOException e) { |
| throw new RuntimeException(e); |
| } |
| }); |
| } |
| |
| @ParameterizedTest |
| @MethodSource("codecTypes") |
| void testReadWriteFile(CompressionUtil.CodecType codec) throws Exception { |
| withRoot( |
| codec, |
| (factory, root) -> { |
| ByteArrayOutputStream compressedStream = new ByteArrayOutputStream(); |
| try (final ArrowFileWriter writer = |
| new ArrowFileWriter( |
| root, |
| new DictionaryProvider.MapDictionaryProvider(), |
| Channels.newChannel(compressedStream), |
| new HashMap<>(), |
| IpcOption.DEFAULT, |
| factory, |
| codec, |
| Optional.of(7))) { |
| writer.start(); |
| writer.writeBatch(); |
| writer.end(); |
| } catch (IOException e) { |
| throw new RuntimeException(e); |
| } |
| |
| try (ArrowFileReader reader = |
| new ArrowFileReader( |
| new ByteArrayReadableSeekableByteChannel(compressedStream.toByteArray()), |
| allocator, |
| factory)) { |
| assertTrue(reader.loadNextBatch()); |
| assertTrue(root.equals(reader.getVectorSchemaRoot())); |
| assertFalse(reader.loadNextBatch()); |
| } catch (IOException e) { |
| throw new RuntimeException(e); |
| } |
| }); |
| } |
| |
| /** Unloading a vector should not free source buffers. */ |
| @ParameterizedTest |
| @MethodSource("codecTypes") |
| void testUnloadCompressed(CompressionUtil.CodecType codec) { |
| withRoot( |
| codec, |
| (factory, root) -> { |
| root.getFieldVectors() |
| .forEach( |
| (vector) -> { |
| Arrays.stream(vector.getBuffers(/*clear*/ false)) |
| .forEach( |
| (buf) -> { |
| assertNotEquals(0, buf.getReferenceManager().getRefCount()); |
| }); |
| }); |
| |
| final VectorUnloader unloader = |
| new VectorUnloader( |
| root, /*includeNullCount*/ |
| true, |
| factory.createCodec(codec), /*alignBuffers*/ |
| true); |
| unloader.getRecordBatch().close(); |
| |
| root.getFieldVectors() |
| .forEach( |
| (vector) -> { |
| Arrays.stream(vector.getBuffers(/*clear*/ false)) |
| .forEach( |
| (buf) -> { |
| assertNotEquals(0, buf.getReferenceManager().getRefCount()); |
| }); |
| }); |
| }); |
| } |
| |
| void withRoot( |
| CompressionUtil.CodecType codec, |
| BiConsumer<CompressionCodec.Factory, VectorSchemaRoot> testBody) { |
| final Schema schema = |
| new Schema( |
| Arrays.asList( |
| Field.nullable("ints", new ArrowType.Int(32, true)), |
| Field.nullable("strings", ArrowType.Utf8.INSTANCE))); |
| CompressionCodec.Factory factory = |
| codec == CompressionUtil.CodecType.NO_COMPRESSION |
| ? NoCompressionCodec.Factory.INSTANCE |
| : CommonsCompressionFactory.INSTANCE; |
| try (final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { |
| final IntVector ints = (IntVector) root.getVector(0); |
| final VarCharVector strings = (VarCharVector) root.getVector(1); |
| // Doesn't get compressed |
| ints.setSafe(0, 0x4a3e); |
| ints.setSafe(1, 0x8aba); |
| ints.setSafe(2, 0x4362); |
| ints.setSafe(3, 0x383f); |
| // Gets compressed |
| String compressibleString = " "; // 16 bytes |
| compressibleString = compressibleString + compressibleString; |
| compressibleString = compressibleString + compressibleString; |
| compressibleString = compressibleString + compressibleString; |
| compressibleString = compressibleString + compressibleString; |
| compressibleString = compressibleString + compressibleString; // 512 bytes |
| byte[] compressibleData = compressibleString.getBytes(StandardCharsets.UTF_8); |
| strings.setSafe(0, compressibleData); |
| strings.setSafe(1, compressibleData); |
| strings.setSafe(2, compressibleData); |
| strings.setSafe(3, compressibleData); |
| root.setRowCount(4); |
| |
| testBody.accept(factory, root); |
| } |
| } |
| } |
