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.c;
  
  import static org.junit.jupiter.api.Assertions.assertEquals;
  import static org.junit.jupiter.api.Assertions.assertFalse;
  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.Collections;
  import org.apache.arrow.memory.BufferAllocator;
  import org.apache.arrow.memory.RootAllocator;
  import org.apache.arrow.util.AutoCloseables;
  import org.apache.arrow.vector.FieldVector;
  import org.apache.arrow.vector.IntVector;
  import org.apache.arrow.vector.ValueVector;
  import org.apache.arrow.vector.VarCharVector;
  import org.apache.arrow.vector.VectorSchemaRoot;
  import org.apache.arrow.vector.ViewVarCharVector;
  import org.apache.arrow.vector.compare.VectorEqualsVisitor;
  import org.apache.arrow.vector.complex.StructVector;
  import org.apache.arrow.vector.dictionary.Dictionary;
  import org.apache.arrow.vector.dictionary.DictionaryEncoder;
  import org.apache.arrow.vector.dictionary.DictionaryProvider;
  import org.apache.arrow.vector.ipc.ArrowStreamReader;
  import org.apache.arrow.vector.ipc.ArrowStreamWriter;
  import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
  import org.apache.arrow.vector.types.Types.MinorType;
  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.junit.jupiter.api.AfterEach;
  import org.junit.jupiter.api.BeforeEach;
  import org.junit.jupiter.api.Test;
  
  public class DictionaryTest {
    private RootAllocator allocator = null;
  
    @BeforeEach
    public void setUp() {
      allocator = new RootAllocator(Long.MAX_VALUE);
    }
  
    @AfterEach
    public void tearDown() {
      allocator.close();
    }
  
    void roundtrip(FieldVector vector, DictionaryProvider provider, Class<?> clazz) {
      // Consumer allocates empty structures
      try (ArrowSchema consumerArrowSchema = ArrowSchema.allocateNew(allocator);
          ArrowArray consumerArrowArray = ArrowArray.allocateNew(allocator)) {
  
        // Producer creates structures from existing memory pointers
        try (ArrowSchema arrowSchema = ArrowSchema.wrap(consumerArrowSchema.memoryAddress());
            ArrowArray arrowArray = ArrowArray.wrap(consumerArrowArray.memoryAddress())) {
          // Producer exports vector into the C Data Interface structures
          Data.exportVector(allocator, vector, provider, arrowArray, arrowSchema);
        }
  
        // Consumer imports vector
        try (CDataDictionaryProvider cDictionaryProvider = new CDataDictionaryProvider();
            FieldVector imported =
                Data.importVector(
                    allocator, consumerArrowArray, consumerArrowSchema, cDictionaryProvider); ) {
          assertTrue(
              clazz.isInstance(imported),
              String.format("expected %s but was %s", clazz, imported.getClass()));
          assertTrue(
              VectorEqualsVisitor.vectorEquals(vector, imported), "vectors are not equivalent");
          for (long id : cDictionaryProvider.getDictionaryIds()) {
            ValueVector exportedDictionaryVector = provider.lookup(id).getVector();
            ValueVector importedDictionaryVector = cDictionaryProvider.lookup(id).getVector();
            assertTrue(
                VectorEqualsVisitor.vectorEquals(exportedDictionaryVector, importedDictionaryVector),
                String.format("Dictionary vectors for ID %d are not equivalent", id));
          }
        }
      }
    }
  
    @Test
    public void testWithDictionary() throws Exception {
      DictionaryProvider.MapDictionaryProvider provider =
          new DictionaryProvider.MapDictionaryProvider();
      // create dictionary and provider
      final VarCharVector dictVector = new VarCharVector("dict", allocator);
      dictVector.allocateNewSafe();
      dictVector.setSafe(0, "aa".getBytes(StandardCharsets.UTF_8));
      dictVector.setSafe(1, "bb".getBytes(StandardCharsets.UTF_8));
      dictVector.setSafe(2, "cc".getBytes(StandardCharsets.UTF_8));
      dictVector.setValueCount(3);
  
      Dictionary dictionary =
          new Dictionary(dictVector, new DictionaryEncoding(0L, false, /* indexType= */ null));
      provider.put(dictionary);
  
      // create vector and encode it
      final VarCharVector vector = new VarCharVector("vector", allocator);
      vector.allocateNewSafe();
      vector.setSafe(0, "bb".getBytes(StandardCharsets.UTF_8));
      vector.setSafe(1, "bb".getBytes(StandardCharsets.UTF_8));
      vector.setSafe(2, "cc".getBytes(StandardCharsets.UTF_8));
      vector.setSafe(3, "aa".getBytes(StandardCharsets.UTF_8));
      vector.setValueCount(4);
  
      // get the encoded vector
      IntVector encodedVector = (IntVector) DictionaryEncoder.encode(vector, dictionary);
  
      // Perform roundtrip using C Data Interface
      roundtrip(encodedVector, provider, IntVector.class);
  
      // Close all
      AutoCloseables.close((AutoCloseable) vector, encodedVector, dictVector);
    }
  
    @Test
    public void testRoundtripMultipleBatches() throws IOException {
      try (ArrowStreamReader reader = createMultiBatchReader();
          ArrowSchema consumerArrowSchema = ArrowSchema.allocateNew(allocator)) {
        // Load first batch
        reader.loadNextBatch();
        // Producer fills consumer schema structure
        Data.exportSchema(
            allocator, reader.getVectorSchemaRoot().getSchema(), reader, consumerArrowSchema);
        // Consumer loads it as an empty vector schema root
        try (CDataDictionaryProvider consumerDictionaryProvider = new CDataDictionaryProvider();
            VectorSchemaRoot consumerRoot =
                Data.importVectorSchemaRoot(
                    allocator, consumerArrowSchema, consumerDictionaryProvider)) {
          do {
            try (ArrowArray consumerArray = ArrowArray.allocateNew(allocator)) {
              // Producer exports next data
              Data.exportVectorSchemaRoot(
                  allocator, reader.getVectorSchemaRoot(), reader, consumerArray);
              // Consumer loads next data
              Data.importIntoVectorSchemaRoot(
                  allocator, consumerArray, consumerRoot, consumerDictionaryProvider);
  
              // Roundtrip validation
              assertTrue(
                  consumerRoot.equals(reader.getVectorSchemaRoot()),
                  "vector schema roots are not equivalent");
              for (long id : consumerDictionaryProvider.getDictionaryIds()) {
                ValueVector exportedDictionaryVector = reader.lookup(id).getVector();
                ValueVector importedDictionaryVector =
                    consumerDictionaryProvider.lookup(id).getVector();
                assertTrue(
                    VectorEqualsVisitor.vectorEquals(
                        exportedDictionaryVector, importedDictionaryVector),
                    String.format("Dictionary vectors for ID %d are not equivalent", id));
              }
            }
          } while (reader.loadNextBatch());
        }
      }
    }
  
    private ArrowStreamReader createMultiBatchReader() throws IOException {
      ByteArrayOutputStream os = new ByteArrayOutputStream();
      try (final VarCharVector dictVector = new VarCharVector("dict", allocator);
          IntVector vector = new IntVector("foo", allocator)) {
        // create dictionary and provider
        DictionaryProvider.MapDictionaryProvider provider =
            new DictionaryProvider.MapDictionaryProvider();
        dictVector.allocateNewSafe();
        dictVector.setSafe(0, "aa".getBytes(StandardCharsets.UTF_8));
        dictVector.setSafe(1, "bb".getBytes(StandardCharsets.UTF_8));
        dictVector.setSafe(2, "cc".getBytes(StandardCharsets.UTF_8));
        dictVector.setSafe(3, "dd".getBytes(StandardCharsets.UTF_8));
        dictVector.setSafe(4, "ee".getBytes(StandardCharsets.UTF_8));
        dictVector.setValueCount(5);
        Dictionary dictionary =
            new Dictionary(dictVector, new DictionaryEncoding(0L, false, /* indexType= */ null));
        provider.put(dictionary);
  
        Schema schema = new Schema(Collections.singletonList(vector.getField()));
        try (VectorSchemaRoot root =
                new VectorSchemaRoot(
                    schema, Collections.singletonList(vector), vector.getValueCount());
            ArrowStreamWriter writer =
                new ArrowStreamWriter(root, provider, Channels.newChannel(os)); ) {
  
          writer.start();
  
          // Batch 1
          vector.setNull(0);
          vector.setSafe(1, 1);
          vector.setSafe(2, 2);
          vector.setNull(3);
          vector.setSafe(4, 1);
          vector.setValueCount(5);
          root.setRowCount(5);
          writer.writeBatch();
  
          // Batch 2
          vector.setNull(0);
          vector.setSafe(1, 1);
          vector.setSafe(2, 2);
          vector.setValueCount(3);
          root.setRowCount(3);
          writer.writeBatch();
  
          // Batch 3
          vector.setSafe(0, 0);
          vector.setSafe(1, 1);
          vector.setSafe(2, 2);
          vector.setSafe(3, 3);
          vector.setSafe(4, 4);
          vector.setValueCount(5);
          root.setRowCount(5);
          writer.writeBatch();
  
          writer.end();
        }
      }
  
      ByteArrayInputStream in = new ByteArrayInputStream(os.toByteArray());
      return new ArrowStreamReader(in, allocator);
    }
  
    private void createStructVector(StructVector vector) {
      final ViewVarCharVector child1 =
          vector.addOrGet(
              "f0", FieldType.nullable(MinorType.VIEWVARCHAR.getType()), ViewVarCharVector.class);
      final IntVector child2 =
          vector.addOrGet("f1", FieldType.nullable(MinorType.INT.getType()), IntVector.class);
  
      // Write the values to child 1
      child1.allocateNew();
      child1.set(0, "01234567890".getBytes(StandardCharsets.UTF_8));
      child1.set(1, "012345678901234567".getBytes(StandardCharsets.UTF_8));
      vector.setIndexDefined(0);
  
      // Write the values to child 2
      child2.allocateNew();
      child2.set(0, 10);
      child2.set(1, 11);
      vector.setIndexDefined(1);
  
      vector.setValueCount(2);
    }
  
    private void createStructVectorInline(StructVector vector) {
      final ViewVarCharVector child1 =
          vector.addOrGet(
              "f0", FieldType.nullable(MinorType.VIEWVARCHAR.getType()), ViewVarCharVector.class);
      final IntVector child2 =
          vector.addOrGet("f1", FieldType.nullable(MinorType.INT.getType()), IntVector.class);
  
      // Write the values to child 1
      child1.allocateNew();
      child1.set(0, "012345678".getBytes(StandardCharsets.UTF_8));
      child1.set(1, "01234".getBytes(StandardCharsets.UTF_8));
      vector.setIndexDefined(0);
  
      // Write the values to child 2
      child2.allocateNew();
      child2.set(0, 10);
      child2.set(1, 11);
      vector.setIndexDefined(1);
  
      vector.setValueCount(2);
    }
  
    @Test
    public void testVectorLoadUnloadOnStructVector() {
      try (final StructVector structVector1 = StructVector.empty("struct", allocator)) {
        createStructVector(structVector1);
        Field field1 = structVector1.getField();
        Schema schema = new Schema(field1.getChildren());
        StructVectorUnloader vectorUnloader = new StructVectorUnloader(structVector1);
  
        try (ArrowRecordBatch recordBatch = vectorUnloader.getRecordBatch();
            BufferAllocator finalVectorsAllocator =
                allocator.newChildAllocator("struct", 0, Long.MAX_VALUE); ) {
          // validating recordBatch contains an output for variadicBufferCounts
          assertFalse(recordBatch.getVariadicBufferCounts().isEmpty());
          assertEquals(1, recordBatch.getVariadicBufferCounts().size());
          assertEquals(1, recordBatch.getVariadicBufferCounts().get(0));
  
          StructVectorLoader vectorLoader = new StructVectorLoader(schema);
          try (StructVector structVector2 = vectorLoader.load(finalVectorsAllocator, recordBatch)) {
            // Improve this after fixing https://github.com/apache/arrow/issues/41933
            // assertTrue(VectorEqualsVisitor.vectorEquals(structVector1, structVector2), "vectors are
            // not equivalent");
            assertTrue(
                VectorEqualsVisitor.vectorEquals(
                    structVector1.getChild("f0"), structVector2.getChild("f0")),
                "vectors are not equivalent");
            assertTrue(
                VectorEqualsVisitor.vectorEquals(
                    structVector1.getChild("f1"), structVector2.getChild("f1")),
                "vectors are not equivalent");
          }
        }
      }
    }
  
    @Test
    public void testVectorLoadUnloadOnStructVectorWithInline() {
      try (final StructVector structVector1 = StructVector.empty("struct", allocator)) {
        createStructVectorInline(structVector1);
        Field field1 = structVector1.getField();
        Schema schema = new Schema(field1.getChildren());
        StructVectorUnloader vectorUnloader = new StructVectorUnloader(structVector1);
  
        try (ArrowRecordBatch recordBatch = vectorUnloader.getRecordBatch();
            BufferAllocator finalVectorsAllocator =
                allocator.newChildAllocator("struct", 0, Long.MAX_VALUE); ) {
          // validating recordBatch contains an output for variadicBufferCounts
          assertFalse(recordBatch.getVariadicBufferCounts().isEmpty());
          assertEquals(1, recordBatch.getVariadicBufferCounts().size());
          assertEquals(0, recordBatch.getVariadicBufferCounts().get(0));
  
          StructVectorLoader vectorLoader = new StructVectorLoader(schema);
          try (StructVector structVector2 = vectorLoader.load(finalVectorsAllocator, recordBatch)) {
            // Improve this after fixing https://github.com/apache/arrow/issues/41933
            // assertTrue(VectorEqualsVisitor.vectorEquals(structVector1, structVector2), "vectors are
            // not equivalent");
            assertTrue(
                VectorEqualsVisitor.vectorEquals(
                    structVector1.getChild("f0"), structVector2.getChild("f0")),
                "vectors are not equivalent");
            assertTrue(
                VectorEqualsVisitor.vectorEquals(
                    structVector1.getChild("f1"), structVector2.getChild("f1")),
                "vectors are not equivalent");
          }
        }
      }
    }
  }
  