content: /*
   * 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.ipc;
  
  import static org.apache.arrow.vector.testing.ValueVectorDataPopulator.setVector;
  import static org.junit.jupiter.api.Assertions.assertArrayEquals;
  import static org.junit.jupiter.api.Assertions.assertEquals;
  import static org.junit.jupiter.api.Assertions.assertNotNull;
  import static org.junit.jupiter.api.Assertions.assertTrue;
  
  import java.io.ByteArrayInputStream;
  import java.io.ByteArrayOutputStream;
  import java.io.IOException;
  import java.nio.channels.Channels;
  import java.nio.charset.StandardCharsets;
  import java.util.Arrays;
  import java.util.Map;
  import java.util.function.ToIntBiFunction;
  import java.util.stream.Stream;
  import org.apache.arrow.memory.BufferAllocator;
  import org.apache.arrow.memory.RootAllocator;
  import org.apache.arrow.vector.FieldVector;
  import org.apache.arrow.vector.UInt1Vector;
  import org.apache.arrow.vector.UInt2Vector;
  import org.apache.arrow.vector.UInt4Vector;
  import org.apache.arrow.vector.UInt8Vector;
  import org.apache.arrow.vector.ValueVector;
  import org.apache.arrow.vector.VarCharVector;
  import org.apache.arrow.vector.VectorSchemaRoot;
  import org.apache.arrow.vector.dictionary.Dictionary;
  import org.apache.arrow.vector.dictionary.DictionaryProvider;
  import org.apache.arrow.vector.types.pojo.ArrowType;
  import org.apache.arrow.vector.types.pojo.DictionaryEncoding;
  import org.apache.arrow.vector.types.pojo.Field;
  import org.apache.arrow.vector.types.pojo.FieldType;
  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 the round-trip of dictionary encoding, with unsigned integer as indices. */
  public class TestUIntDictionaryRoundTrip {
  
    private BufferAllocator allocator;
  
    private DictionaryProvider.MapDictionaryProvider dictionaryProvider;
  
    @BeforeEach
    public void init() {
      allocator = new RootAllocator(Long.MAX_VALUE);
      dictionaryProvider = new DictionaryProvider.MapDictionaryProvider();
    }
  
    @AfterEach
    public void terminate() throws Exception {
      allocator.close();
    }
  
    private byte[] writeData(boolean streamMode, FieldVector encodedVector) throws IOException {
      ByteArrayOutputStream out = new ByteArrayOutputStream();
      VectorSchemaRoot root =
          new VectorSchemaRoot(
              Arrays.asList(encodedVector.getField()),
              Arrays.asList(encodedVector),
              encodedVector.getValueCount());
      try (ArrowWriter writer =
          streamMode
              ? new ArrowStreamWriter(root, dictionaryProvider, out)
              : new ArrowFileWriter(root, dictionaryProvider, Channels.newChannel(out))) {
        writer.start();
        writer.writeBatch();
        writer.end();
  
        return out.toByteArray();
      }
    }
  
    private void readData(
        boolean streamMode,
        byte[] data,
        Field expectedField,
        ToIntBiFunction<ValueVector, Integer> valGetter,
        long dictionaryID,
        int[] expectedIndices,
        String[] expectedDictItems)
        throws IOException {
      try (ArrowReader reader =
          streamMode
              ? new ArrowStreamReader(new ByteArrayInputStream(data), allocator)
              : new ArrowFileReader(
                  new SeekableReadChannel(new ByteArrayReadableSeekableByteChannel(data)),
                  allocator)) {
  
        // verify schema
        Schema readSchema = reader.getVectorSchemaRoot().getSchema();
        assertEquals(1, readSchema.getFields().size());
        assertEquals(expectedField, readSchema.getFields().get(0));
  
        // verify vector schema root
        assertTrue(reader.loadNextBatch());
        VectorSchemaRoot root = reader.getVectorSchemaRoot();
  
        assertEquals(1, root.getFieldVectors().size());
        ValueVector encodedVector = root.getVector(0);
        assertEquals(expectedIndices.length, encodedVector.getValueCount());
  
        for (int i = 0; i < expectedIndices.length; i++) {
          assertEquals(expectedIndices[i], valGetter.applyAsInt(encodedVector, i));
        }
  
        // verify dictionary
        Map<Long, Dictionary> dictVectors = reader.getDictionaryVectors();
        assertEquals(1, dictVectors.size());
        Dictionary dictionary = dictVectors.get(dictionaryID);
        assertNotNull(dictionary);
  
        assertTrue(dictionary.getVector() instanceof VarCharVector);
        VarCharVector dictVector = (VarCharVector) dictionary.getVector();
        assertEquals(expectedDictItems.length, dictVector.getValueCount());
        for (int i = 0; i < dictVector.getValueCount(); i++) {
          assertArrayEquals(expectedDictItems[i].getBytes(StandardCharsets.UTF_8), dictVector.get(i));
        }
      }
    }
  
    private ValueVector createEncodedVector(int bitWidth, VarCharVector dictionaryVector) {
      final DictionaryEncoding dictionaryEncoding =
          new DictionaryEncoding(bitWidth, false, new ArrowType.Int(bitWidth, false));
      Dictionary dictionary = new Dictionary(dictionaryVector, dictionaryEncoding);
      dictionaryProvider.put(dictionary);
  
      final FieldType type =
          new FieldType(true, dictionaryEncoding.getIndexType(), dictionaryEncoding, null);
      final Field field = new Field("encoded", type, null);
      return field.createVector(allocator);
    }
  
    @ParameterizedTest(name = "stream mode = {0}")
    @MethodSource("getRepeat")
    public void testUInt1RoundTrip(boolean streamMode) throws IOException {
      final int vectorLength = UInt1Vector.MAX_UINT1 & UInt1Vector.PROMOTION_MASK;
      try (VarCharVector dictionaryVector = new VarCharVector("dictionary", allocator);
          UInt1Vector encodedVector1 = (UInt1Vector) createEncodedVector(8, dictionaryVector)) {
        int[] indices = new int[vectorLength];
        String[] dictionaryItems = new String[vectorLength];
        for (int i = 0; i < vectorLength; i++) {
          encodedVector1.setSafe(i, (byte) i);
          indices[i] = i;
          dictionaryItems[i] = String.valueOf(i);
        }
        encodedVector1.setValueCount(vectorLength);
        setVector(dictionaryVector, dictionaryItems);
        byte[] data = writeData(streamMode, encodedVector1);
        readData(
            streamMode,
            data,
            encodedVector1.getField(),
            (vector, index) -> (int) ((UInt1Vector) vector).getValueAsLong(index),
            8L,
            indices,
            dictionaryItems);
      }
    }
  
    @ParameterizedTest(name = "stream mode = {0}")
    @MethodSource("getRepeat")
    public void testUInt2RoundTrip(boolean streamMode) throws IOException {
      try (VarCharVector dictionaryVector = new VarCharVector("dictionary", allocator);
          UInt2Vector encodedVector2 = (UInt2Vector) createEncodedVector(16, dictionaryVector)) {
        int[] indices = new int[] {1, 3, 5, 7, 9, UInt2Vector.MAX_UINT2};
        String[] dictItems = new String[UInt2Vector.MAX_UINT2];
        for (int i = 0; i < UInt2Vector.MAX_UINT2; i++) {
          dictItems[i] = String.valueOf(i);
        }
  
        setVector(
            encodedVector2, (char) 1, (char) 3, (char) 5, (char) 7, (char) 9, UInt2Vector.MAX_UINT2);
        setVector(dictionaryVector, dictItems);
  
        byte[] data = writeData(streamMode, encodedVector2);
        readData(
            streamMode,
            data,
            encodedVector2.getField(),
            (vector, index) -> (int) ((UInt2Vector) vector).getValueAsLong(index),
            16L,
            indices,
            dictItems);
      }
    }
  
    @ParameterizedTest(name = "stream mode = {0}")
    @MethodSource("getRepeat")
    public void testUInt4RoundTrip(boolean streamMode) throws IOException {
      final int dictLength = 10;
      try (VarCharVector dictionaryVector = new VarCharVector("dictionary", allocator);
          UInt4Vector encodedVector4 = (UInt4Vector) createEncodedVector(32, dictionaryVector)) {
        int[] indices = new int[] {1, 3, 5, 7, 9};
        String[] dictItems = new String[dictLength];
        for (int i = 0; i < dictLength; i++) {
          dictItems[i] = String.valueOf(i);
        }
  
        setVector(encodedVector4, 1, 3, 5, 7, 9);
        setVector(dictionaryVector, dictItems);
  
        setVector(encodedVector4, 1, 3, 5, 7, 9);
        byte[] data = writeData(streamMode, encodedVector4);
        readData(
            streamMode,
            data,
            encodedVector4.getField(),
            (vector, index) -> (int) ((UInt4Vector) vector).getValueAsLong(index),
            32L,
            indices,
            dictItems);
      }
    }
  
    @ParameterizedTest(name = "stream mode = {0}")
    @MethodSource("getRepeat")
    public void testUInt8RoundTrip(boolean streamMode) throws IOException {
      final int dictLength = 10;
      try (VarCharVector dictionaryVector = new VarCharVector("dictionary", allocator);
          UInt8Vector encodedVector8 = (UInt8Vector) createEncodedVector(64, dictionaryVector)) {
        int[] indices = new int[] {1, 3, 5, 7, 9};
        String[] dictItems = new String[dictLength];
        for (int i = 0; i < dictLength; i++) {
          dictItems[i] = String.valueOf(i);
        }
  
        setVector(encodedVector8, 1L, 3L, 5L, 7L, 9L);
        setVector(dictionaryVector, dictItems);
  
        byte[] data = writeData(streamMode, encodedVector8);
        readData(
            streamMode,
            data,
            encodedVector8.getField(),
            (vector, index) -> (int) ((UInt8Vector) vector).getValueAsLong(index),
            64L,
            indices,
            dictItems);
      }
    }
  
    static Stream<Arguments> getRepeat() {
      return Stream.of(Arguments.of(true), Arguments.of(false));
    }
  }
  