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.dictionary;
  
  import java.util.ArrayList;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  import org.apache.arrow.memory.BufferAllocator;
  import org.apache.arrow.memory.util.hash.ArrowBufHasher;
  import org.apache.arrow.memory.util.hash.SimpleHasher;
  import org.apache.arrow.util.AutoCloseables;
  import org.apache.arrow.util.Preconditions;
  import org.apache.arrow.vector.BaseIntVector;
  import org.apache.arrow.vector.FieldVector;
  import org.apache.arrow.vector.complex.StructVector;
  import org.apache.arrow.vector.ipc.message.ArrowFieldNode;
  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.util.TransferPair;
  
  /**
   * Sub fields encoder/decoder for Dictionary encoded {@link StructVector}. Notes that child vectors
   * within struct vector can either be dictionary encodable or not.
   */
  public class StructSubfieldEncoder {
  
    private final BufferAllocator allocator;
  
    private final DictionaryProvider.MapDictionaryProvider provider;
    private final Map<Long, DictionaryHashTable> dictionaryIdToHashTable;
  
    /** Construct an instance. */
    public StructSubfieldEncoder(
        BufferAllocator allocator, DictionaryProvider.MapDictionaryProvider provider) {
      this(allocator, provider, SimpleHasher.INSTANCE);
    }
  
    /** Construct an instance. */
    public StructSubfieldEncoder(
        BufferAllocator allocator,
        DictionaryProvider.MapDictionaryProvider provider,
        ArrowBufHasher hasher) {
  
      this.allocator = allocator;
      this.provider = provider;
  
      this.dictionaryIdToHashTable = new HashMap<>();
  
      provider
          .getDictionaryIds()
          .forEach(
              id ->
                  dictionaryIdToHashTable.put(
                      id, new DictionaryHashTable(provider.lookup(id).getVector(), hasher)));
    }
  
    private static FieldVector getChildVector(StructVector vector, int index) {
      return vector.getChildrenFromFields().get(index);
    }
  
    private static StructVector cloneVector(StructVector vector, BufferAllocator allocator) {
  
      final FieldType fieldType = vector.getField().getFieldType();
      StructVector cloned =
          (StructVector)
              fieldType.createNewSingleVector(
                  vector.getField().getName(), allocator, /* schemaCallback= */ null);
  
      final ArrowFieldNode fieldNode =
          new ArrowFieldNode(vector.getValueCount(), vector.getNullCount());
      cloned.loadFieldBuffers(fieldNode, vector.getFieldBuffers());
  
      return cloned;
    }
  
    /**
     * Dictionary encodes subfields for complex vector with a provided dictionary. The dictionary must
     * contain all values in the sub fields vector.
     *
     * @param vector vector to encode
     * @param columnToDictionaryId the mappings between child vector index and dictionary id. A null
     *     dictionary id indicates the child vector is not encodable.
     * @return dictionary encoded vector
     */
    public StructVector encode(StructVector vector, Map<Integer, Long> columnToDictionaryId) {
      final int valueCount = vector.getValueCount();
      final int childCount = vector.getChildrenFromFields().size();
  
      List<Field> childrenFields = new ArrayList<>();
  
      // initialize child fields.
      for (int i = 0; i < childCount; i++) {
        FieldVector childVector = getChildVector(vector, i);
        Long dictionaryId = columnToDictionaryId.get(i);
        // A null dictionaryId indicates the child vector shouldn't be encoded.
        if (dictionaryId == null) {
          childrenFields.add(childVector.getField());
        } else {
          Dictionary dictionary = provider.lookup(dictionaryId);
          Preconditions.checkNotNull(dictionary, "Dictionary not found with id:" + dictionaryId);
          FieldType indexFieldType =
              new FieldType(
                  childVector.getField().isNullable(),
                  dictionary.getEncoding().getIndexType(),
                  dictionary.getEncoding());
          childrenFields.add(
              new Field(childVector.getField().getName(), indexFieldType, /* children= */ null));
        }
      }
  
      // clone list vector and initialize data vector
      StructVector encoded = cloneVector(vector, allocator);
      try {
        encoded.initializeChildrenFromFields(childrenFields);
        encoded.setValueCount(valueCount);
  
        for (int index = 0; index < childCount; index++) {
          FieldVector childVector = getChildVector(vector, index);
          FieldVector encodedChildVector = getChildVector(encoded, index);
          Long dictionaryId = columnToDictionaryId.get(index);
          if (dictionaryId != null) {
            BaseIntVector indices = (BaseIntVector) encodedChildVector;
            DictionaryEncoder.buildIndexVector(
                childVector, indices, dictionaryIdToHashTable.get(dictionaryId), 0, valueCount);
          } else {
            childVector.makeTransferPair(encodedChildVector).splitAndTransfer(0, valueCount);
          }
        }
  
        return encoded;
      } catch (Exception e) {
        AutoCloseables.close(e, encoded);
        throw e;
      }
    }
  
    /**
     * Decodes a dictionary subfields encoded vector using the provided dictionary.
     *
     * <p>{@link StructSubfieldEncoder#decode(StructVector, DictionaryProvider.MapDictionaryProvider,
     * BufferAllocator)} should be used instead if only decoding is required as it can avoid building
     * the {@link DictionaryHashTable} which only makes sense when encoding.
     *
     * @param vector dictionary encoded vector, its child vector must be int type
     * @return vector with values restored from dictionary
     */
    public StructVector decode(StructVector vector) {
      return decode(vector, provider, allocator);
    }
  
    /**
     * Decodes a dictionary subfields encoded vector using the provided dictionary.
     *
     * @param vector dictionary encoded vector, its data vector must be int type
     * @param provider dictionary provider used to decode the values
     * @param allocator allocator the decoded values use
     * @return vector with values restored from dictionary
     */
    public static StructVector decode(
        StructVector vector,
        DictionaryProvider.MapDictionaryProvider provider,
        BufferAllocator allocator) {
      final int valueCount = vector.getValueCount();
      final int childCount = vector.getChildrenFromFields().size();
  
      // clone list vector and initialize child vectors
      StructVector decoded = cloneVector(vector, allocator);
      try {
        List<Field> childFields = new ArrayList<>();
        for (int i = 0; i < childCount; i++) {
          FieldVector childVector = getChildVector(vector, i);
          Dictionary dictionary = getChildVectorDictionary(childVector, provider);
          // childVector is not encoded.
          if (dictionary == null) {
            childFields.add(childVector.getField());
          } else {
            childFields.add(dictionary.getVector().getField());
          }
        }
        decoded.initializeChildrenFromFields(childFields);
        decoded.setValueCount(valueCount);
  
        for (int index = 0; index < childCount; index++) {
          // get child vector
          FieldVector childVector = getChildVector(vector, index);
          FieldVector decodedChildVector = getChildVector(decoded, index);
          Dictionary dictionary = getChildVectorDictionary(childVector, provider);
          if (dictionary == null) {
            childVector.makeTransferPair(decodedChildVector).splitAndTransfer(0, valueCount);
          } else {
            TransferPair transfer = dictionary.getVector().makeTransferPair(decodedChildVector);
            BaseIntVector indices = (BaseIntVector) childVector;
  
            DictionaryEncoder.retrieveIndexVector(indices, transfer, valueCount, 0, valueCount);
          }
        }
  
        return decoded;
      } catch (Exception e) {
        AutoCloseables.close(e, decoded);
        throw e;
      }
    }
  
    /** Get the child vector dictionary, return null if not dictionary encoded. */
    private static Dictionary getChildVectorDictionary(
        FieldVector childVector, DictionaryProvider.MapDictionaryProvider provider) {
      DictionaryEncoding dictionaryEncoding = childVector.getField().getDictionary();
      if (dictionaryEncoding != null) {
        Dictionary dictionary = provider.lookup(dictionaryEncoding.getId());
        Preconditions.checkNotNull(dictionary, "Dictionary not found with id:" + dictionary);
        return dictionary;
      }
      return null;
    }
  }
  