--- /dev/null
+/*
+ * 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.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 LinearDictionaryEncoder}.
+ */
+public class TestLinearDictionaryEncoder {
+
+ 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);
+
+ LinearDictionaryEncoder<IntVector, VarCharVector> encoder =
+ new LinearDictionaryEncoder<>(dictionary, 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);
+
+ LinearDictionaryEncoder<IntVector, VarCharVector> encoder =
+ new LinearDictionaryEncoder<>(dictionary, 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 testEncodeNullWithoutNullInDictionary() {
+ 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();
+
+ LinearDictionaryEncoder<IntVector, VarCharVector> encoder =
+ new LinearDictionaryEncoder<>(dictionary, 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);
+
+ LinearDictionaryEncoder<IntVector, VarCharVector> encoder =
+ new LinearDictionaryEncoder<>(dictionaryVector);
+ 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);
+
+ LinearDictionaryEncoder<IntVector, VarCharVector> encoder =
+ new LinearDictionaryEncoder<>(dictionaryVector);
+ 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);
+
+ LinearDictionaryEncoder<IntVector, VarBinaryVector> encoder =
+ new LinearDictionaryEncoder<>(dictionaryVector);
+ 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)));
+ }
+ }
+ }
+ }
+}
projects / ceph.git / blobdiff