From 149efd9441b46523c549802ef73856f1b88b4538 Mon Sep 17 00:00:00 2001 From: liyafan82 Date: Tue, 17 Sep 2019 21:44:33 -0700 Subject: [PATCH] ARROW-5917: [Java] Redesign the dictionary encoder The current dictionary encoder implementation (org.apache.arrow.vector.dictionary.DictionaryEncoder) has heavy performance overhead, which prevents it from being useful in practice: * There are repeated conversions between Java objects and bytes (e.g. vector.getObject). * Unnecessary memory copy (the vector data must be copied to the hash table). * The hash table cannot be reused for encoding multiple vectors (other data structure & results cannot be reused either). * The output vector should not be created/managed by the encoder (just like in the out-of-place sorter) * The hash table requires that the hashCode & equals methods be implemented appropriately, but this is not guaranteed. We plan to implement a new one in the algorithm module, and gradually deprecate the current one. Closes #4994 from liyafan82/fly_0712_encode and squashes the following commits: 8b699a8f7 Redesign the dictionary encoder Authored-by: liyafan82 Signed-off-by: Micah Kornfield --- .../dictionary/SearchDictionaryEncoder.java | 105 +++++ .../TestSearchDictionaryEncoder.java | 358 ++++++++++++++++++ 2 files changed, 463 insertions(+) create mode 100644 java/algorithm/src/main/java/org/apache/arrow/algorithm/dictionary/SearchDictionaryEncoder.java create mode 100644 java/algorithm/src/test/java/org/apache/arrow/algorithm/dictionary/TestSearchDictionaryEncoder.java diff --git a/java/algorithm/src/main/java/org/apache/arrow/algorithm/dictionary/SearchDictionaryEncoder.java b/java/algorithm/src/main/java/org/apache/arrow/algorithm/dictionary/SearchDictionaryEncoder.java new file mode 100644 index 0000000000000..091a6f44dee68 --- /dev/null +++ b/java/algorithm/src/main/java/org/apache/arrow/algorithm/dictionary/SearchDictionaryEncoder.java @@ -0,0 +1,105 @@ +/* + * 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.algorithm.dictionary; + +import org.apache.arrow.algorithm.search.VectorSearcher; +import org.apache.arrow.algorithm.sort.VectorValueComparator; +import org.apache.arrow.vector.BaseIntVector; +import org.apache.arrow.vector.ValueVector; + +/** + * Dictionary encoder based on searching. + * @param encoded vector type. + * @param decoded vector type, which is also the dictionary type. + */ +public class SearchDictionaryEncoder { + + /** + * The dictionary for encoding/decoding. + * It must be sorted. + */ + private final D dictionary; + + /** + * The criteria by which the dictionary is sorted. + */ + private final VectorValueComparator comparator; + + /** + * A flag indicating if null should be encoded. + */ + private final boolean encodeNull; + + /** + * Constructs a dictionary encoder. + * @param dictionary the dictionary. It must be in sorted order. + * @param comparator the criteria for sorting. + */ + public SearchDictionaryEncoder(D dictionary, VectorValueComparator comparator) { + this(dictionary, comparator, false); + } + + /** + * Constructs a dictionary encoder. + * @param dictionary the dictionary. It must be in sorted order. + * @param comparator the criteria for sorting. + * @param encodeNull a flag indicating if null should be encoded. + * It determines the behaviors for processing null values in the input during encoding/decoding. + *
  • + * For encoding, when a null is encountered in the input, + * 1) If the flag is set to true, the encoder searches for the value in the dictionary, + * and outputs the index in the dictionary. + * 2) If the flag is set to false, the encoder simply produces a null in the output. + *
  • + *
  • + * For decoding, when a null is encountered in the input, + * 1) If the flag is set to true, the decoder should never expect a null in the input. + * 2) If set to false, the decoder simply produces a null in the output. + *
  • + */ + public SearchDictionaryEncoder(D dictionary, VectorValueComparator comparator, boolean encodeNull) { + this.dictionary = dictionary; + this.comparator = comparator; + this.encodeNull = encodeNull; + } + + /** + * Encodes an input vector by binary search. + * So the algorithm takes O(n * log(m)) time, where n is the length of the input vector, + * and m is the length of the dictionary. + * @param input the input vector. + * @param output the output vector. Note that it must be in a fresh state. At least, + * all its validity bits should be clear. + */ + public void encode(D input, E output) { + for (int i = 0; i < input.getValueCount(); i++) { + if (!encodeNull && input.isNull(i)) { + // for this case, we should simply output a null in the output. + // by assuming the output vector is fresh, we do nothing here. + continue; + } + + int index = VectorSearcher.binarySearch(dictionary, comparator, input, i); + if (index == -1) { + throw new IllegalArgumentException("The data element is not found in the dictionary: " + i); + } + output.setWithPossibleTruncate(i, index); + } + output.setValueCount(input.getValueCount()); + } +} diff --git a/java/algorithm/src/test/java/org/apache/arrow/algorithm/dictionary/TestSearchDictionaryEncoder.java b/java/algorithm/src/test/java/org/apache/arrow/algorithm/dictionary/TestSearchDictionaryEncoder.java new file mode 100644 index 0000000000000..22aaf183d8ac9 --- /dev/null +++ b/java/algorithm/src/test/java/org/apache/arrow/algorithm/dictionary/TestSearchDictionaryEncoder.java @@ -0,0 +1,358 @@ +/* + * 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.algorithm.dictionary; + +import static junit.framework.TestCase.assertTrue; +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.Random; + +import org.apache.arrow.algorithm.sort.DefaultVectorComparators; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; + +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VarBinaryVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.dictionary.Dictionary; +import org.apache.arrow.vector.dictionary.DictionaryEncoder; +import org.apache.arrow.vector.types.pojo.DictionaryEncoding; +import org.junit.After; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; + +/** + * Test cases for {@link SearchDictionaryEncoder}. + */ +public class TestSearchDictionaryEncoder { + + private final int VECTOR_LENGTH = 50; + + private final int DICTIONARY_LENGTH = 10; + + private BufferAllocator allocator; + + byte[] zero = "000".getBytes(StandardCharsets.UTF_8); + byte[] one = "111".getBytes(StandardCharsets.UTF_8); + byte[] two = "222".getBytes(StandardCharsets.UTF_8); + + byte[][] data = new byte[][]{zero, one, two}; + + @Before + public void prepare() { + allocator = new RootAllocator(1024 * 1024); + } + + @After + public void shutdown() { + allocator.close(); + } + + @Test + public void testEncodeAndDecode() { + Random random = new Random(); + try (VarCharVector rawVector = new VarCharVector("original vector", allocator); + IntVector encodedVector = new IntVector("encoded vector", allocator); + VarCharVector dictionary = new VarCharVector("dictionary", allocator)) { + + // set up dictionary + dictionary.allocateNew(); + for (int i = 0; i < DICTIONARY_LENGTH; i++) { + // encode "i" as i + dictionary.setSafe(i, String.valueOf(i).getBytes()); + } + dictionary.setValueCount(DICTIONARY_LENGTH); + + // set up raw vector + rawVector.allocateNew(10 * VECTOR_LENGTH, VECTOR_LENGTH); + for (int i = 0; i < VECTOR_LENGTH; i++) { + int val = (random.nextInt() & Integer.MAX_VALUE) % DICTIONARY_LENGTH; + rawVector.set(i, String.valueOf(val).getBytes()); + } + rawVector.setValueCount(VECTOR_LENGTH); + + SearchDictionaryEncoder encoder = + new SearchDictionaryEncoder<>( + dictionary, DefaultVectorComparators.createDefaultComparator(rawVector), false); + + // perform encoding + encodedVector.allocateNew(); + encoder.encode(rawVector, encodedVector); + + // verify encoding results + assertEquals(rawVector.getValueCount(), encodedVector.getValueCount()); + for (int i = 0; i < VECTOR_LENGTH; i++) { + assertArrayEquals(rawVector.get(i), String.valueOf(encodedVector.get(i)).getBytes()); + } + + // perform decoding + Dictionary dict = new Dictionary(dictionary, new DictionaryEncoding(1L, false, null)); + try (VarCharVector decodedVector = (VarCharVector) DictionaryEncoder.decode(encodedVector, dict)) { + + // verify decoding results + assertEquals(encodedVector.getValueCount(), decodedVector.getValueCount()); + for (int i = 0; i < VECTOR_LENGTH; i++) { + assertArrayEquals(String.valueOf(encodedVector.get(i)).getBytes(), decodedVector.get(i)); + } + } + } + } + + @Test + public void testEncodeAndDecodeWithNull() { + Random random = new Random(); + try (VarCharVector rawVector = new VarCharVector("original vector", allocator); + IntVector encodedVector = new IntVector("encoded vector", allocator); + VarCharVector dictionary = new VarCharVector("dictionary", allocator)) { + + // set up dictionary + dictionary.allocateNew(); + dictionary.setNull(0); + for (int i = 1; i < DICTIONARY_LENGTH; i++) { + // encode "i" as i + dictionary.setSafe(i, String.valueOf(i).getBytes()); + } + dictionary.setValueCount(DICTIONARY_LENGTH); + + // set up raw vector + rawVector.allocateNew(10 * VECTOR_LENGTH, VECTOR_LENGTH); + for (int i = 0; i < VECTOR_LENGTH; i++) { + if (i % 10 == 0) { + rawVector.setNull(i); + } else { + int val = (random.nextInt() & Integer.MAX_VALUE) % (DICTIONARY_LENGTH - 1) + 1; + rawVector.set(i, String.valueOf(val).getBytes()); + } + } + rawVector.setValueCount(VECTOR_LENGTH); + + SearchDictionaryEncoder encoder = + new SearchDictionaryEncoder<>( + dictionary, DefaultVectorComparators.createDefaultComparator(rawVector), true); + + // perform encoding + encodedVector.allocateNew(); + encoder.encode(rawVector, encodedVector); + + // verify encoding results + assertEquals(rawVector.getValueCount(), encodedVector.getValueCount()); + for (int i = 0; i < VECTOR_LENGTH; i++) { + if (i % 10 == 0) { + assertEquals(0, encodedVector.get(i)); + } else { + assertArrayEquals(rawVector.get(i), String.valueOf(encodedVector.get(i)).getBytes()); + } + } + + // perform decoding + Dictionary dict = new Dictionary(dictionary, new DictionaryEncoding(1L, false, null)); + try (VarCharVector decodedVector = (VarCharVector) DictionaryEncoder.decode(encodedVector, dict)) { + + // verify decoding results + assertEquals(encodedVector.getValueCount(), decodedVector.getValueCount()); + for (int i = 0; i < VECTOR_LENGTH; i++) { + if (i % 10 == 0) { + assertTrue(decodedVector.isNull(i)); + } else { + assertArrayEquals(String.valueOf(encodedVector.get(i)).getBytes(), decodedVector.get(i)); + } + } + } + } + } + + @Test + public void testEncodeNoNullInDictionary() { + try (VarCharVector rawVector = new VarCharVector("original vector", allocator); + IntVector encodedVector = new IntVector("encoded vector", allocator); + VarCharVector dictionary = new VarCharVector("dictionary", allocator)) { + + // set up dictionary, with no null in it. + dictionary.allocateNew(); + for (int i = 0; i < DICTIONARY_LENGTH; i++) { + // encode "i" as i + dictionary.setSafe(i, String.valueOf(i).getBytes()); + } + dictionary.setValueCount(DICTIONARY_LENGTH); + + // the vector to encode has a null inside. + rawVector.allocateNew(1); + rawVector.setNull(0); + rawVector.setValueCount(1); + + encodedVector.allocateNew(); + + SearchDictionaryEncoder encoder = + new SearchDictionaryEncoder<>( + dictionary, DefaultVectorComparators.createDefaultComparator(rawVector), true); + + // the encoder should encode null, but no null in the dictionary, + // so an exception should be thrown. + assertThrows(IllegalArgumentException.class, () -> { + encoder.encode(rawVector, encodedVector); + }); + } + } + + @Test + public void testEncodeStrings() { + // Create a new value vector + try (final VarCharVector vector = new VarCharVector("foo", allocator); + final IntVector encoded = new IntVector("encoded", allocator); + final VarCharVector dictionaryVector = new VarCharVector("dict", allocator)) { + + vector.allocateNew(512, 5); + encoded.allocateNew(); + + // set some values + vector.setSafe(0, zero, 0, zero.length); + vector.setSafe(1, one, 0, one.length); + vector.setSafe(2, one, 0, one.length); + vector.setSafe(3, two, 0, two.length); + vector.setSafe(4, zero, 0, zero.length); + vector.setValueCount(5); + + // set some dictionary values + dictionaryVector.allocateNew(512, 3); + dictionaryVector.setSafe(0, zero, 0, one.length); + dictionaryVector.setSafe(1, one, 0, two.length); + dictionaryVector.setSafe(2, two, 0, zero.length); + dictionaryVector.setValueCount(3); + + SearchDictionaryEncoder encoder = + new SearchDictionaryEncoder<>( + dictionaryVector, DefaultVectorComparators.createDefaultComparator(vector)); + encoder.encode(vector, encoded); + + // verify indices + assertEquals(5, encoded.getValueCount()); + assertEquals(0, encoded.get(0)); + assertEquals(1, encoded.get(1)); + assertEquals(1, encoded.get(2)); + assertEquals(2, encoded.get(3)); + assertEquals(0, encoded.get(4)); + + // now run through the decoder and verify we get the original back + Dictionary dict = new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null)); + try (VarCharVector decoded = (VarCharVector) DictionaryEncoder.decode(encoded, dict)) { + assertEquals(vector.getValueCount(), decoded.getValueCount()); + for (int i = 0; i < 5; i++) { + assertEquals(vector.getObject(i), decoded.getObject(i)); + } + } + } + } + + @Test + public void testEncodeLargeVector() { + // Create a new value vector + try (final VarCharVector vector = new VarCharVector("foo", allocator); + final IntVector encoded = new IntVector("encoded", allocator); + final VarCharVector dictionaryVector = new VarCharVector("dict", allocator)) { + vector.allocateNew(); + encoded.allocateNew(); + + int count = 10000; + + for (int i = 0; i < 10000; ++i) { + vector.setSafe(i, data[i % 3], 0, data[i % 3].length); + } + vector.setValueCount(count); + + dictionaryVector.allocateNew(512, 3); + dictionaryVector.setSafe(0, zero, 0, one.length); + dictionaryVector.setSafe(1, one, 0, two.length); + dictionaryVector.setSafe(2, two, 0, zero.length); + dictionaryVector.setValueCount(3); + + SearchDictionaryEncoder encoder = + new SearchDictionaryEncoder<>( + dictionaryVector, DefaultVectorComparators.createDefaultComparator(vector)); + encoder.encode(vector, encoded); + + assertEquals(count, encoded.getValueCount()); + for (int i = 0; i < count; ++i) { + assertEquals(i % 3, encoded.get(i)); + } + + // now run through the decoder and verify we get the original back + Dictionary dict = new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null)); + try (VarCharVector decoded = (VarCharVector) DictionaryEncoder.decode(encoded, dict)) { + assertEquals(vector.getClass(), decoded.getClass()); + assertEquals(vector.getValueCount(), decoded.getValueCount()); + for (int i = 0; i < count; ++i) { + assertEquals(vector.getObject(i), decoded.getObject(i)); + } + } + } + } + + @Test + public void testEncodeBinaryVector() { + // Create a new value vector + try (final VarBinaryVector vector = new VarBinaryVector("foo", allocator); + final VarBinaryVector dictionaryVector = new VarBinaryVector("dict", allocator); + final IntVector encoded = new IntVector("encoded", allocator)) { + vector.allocateNew(512, 5); + vector.allocateNew(); + encoded.allocateNew(); + + // set some values + vector.setSafe(0, zero, 0, zero.length); + vector.setSafe(1, one, 0, one.length); + vector.setSafe(2, one, 0, one.length); + vector.setSafe(3, two, 0, two.length); + vector.setSafe(4, zero, 0, zero.length); + vector.setValueCount(5); + + // set some dictionary values + dictionaryVector.allocateNew(512, 3); + dictionaryVector.setSafe(0, zero, 0, one.length); + dictionaryVector.setSafe(1, one, 0, two.length); + dictionaryVector.setSafe(2, two, 0, zero.length); + dictionaryVector.setValueCount(3); + + SearchDictionaryEncoder encoder = + new SearchDictionaryEncoder<>( + dictionaryVector, DefaultVectorComparators.createDefaultComparator(vector)); + encoder.encode(vector, encoded); + + assertEquals(5, encoded.getValueCount()); + assertEquals(0, encoded.get(0)); + assertEquals(1, encoded.get(1)); + assertEquals(1, encoded.get(2)); + assertEquals(2, encoded.get(3)); + assertEquals(0, encoded.get(4)); + + // now run through the decoder and verify we get the original back + Dictionary dict = new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null)); + try (VarBinaryVector decoded = (VarBinaryVector) DictionaryEncoder.decode(encoded, dict)) { + assertEquals(vector.getClass(), decoded.getClass()); + assertEquals(vector.getValueCount(), decoded.getValueCount()); + for (int i = 0; i < 5; i++) { + Assert.assertTrue(Arrays.equals(vector.getObject(i), decoded.getObject(i))); + } + } + } + } +}