<commit_before>/**
 * Copyright (C) 2014-2015 LinkedIn Corp. (<EMAIL>)
 *
 * Licensed 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 com.linkedin.pinot.core.operator.aggregation.groupby;

import com.linkedin.pinot.core.common.Block;
import com.linkedin.pinot.core.common.BlockMultiValIterator;
import com.linkedin.pinot.core.common.BlockSingleValIterator;
import com.linkedin.pinot.core.common.DataSource;
import com.linkedin.pinot.core.indexsegment.IndexSegment;
import com.linkedin.pinot.core.operator.aggregation.ResultHolderFactory;
import com.linkedin.pinot.core.query.aggregation.groupby.GroupByConstants;
import com.linkedin.pinot.core.segment.index.readers.Dictionary;
import it.unimi.dsi.fastutil.longs.Long2IntOpenHashMap;
import it.unimi.dsi.fastutil.objects.ObjectIterator;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Map;


/**
 * Class for generating group-by keys for a given list of group-by columns.
 * Uses dictionary id's to generate group-by keys.
 *
 * If maximum number of possible group keys are less than a threshold (10K),
 * returns the group key generated based on dictionary id's for the group by columns.
 *
 * For larger values, maps the actual generated group by keys onto contiguous
 * indices, and returns the latter. Keeps a mapping between the two internally.
 *
 */
public class DefaultGroupKeyGenerator implements GroupKeyGenerator {
  private static final int INVALID_ID = -1;

  public enum STORAGE_TYPE {
    ARRAY_BASED,
    MAP_BASED
  }

  private final String[] _groupByColumns;

  private final BlockSingleValIterator[] _singleValIterators;
  private final BlockMultiValIterator[] _multiValIterators;

  private final Dictionary[] _dictionaries;
  private final int[] _cardinalities;

  // A reusable array of number of group by columns size,
  // to avoid creating int[] objects with each call.
  private final int[] _reusableGroupByValuesArray;
  private final int[][] _multiValueResulableArray;

  private final boolean[] _isSingleValueGroupByColumn;
  private STORAGE_TYPE _storageType;
  private int _numUniqueGroupKeys;

  // For ARRAY_BASED storage type.
  private boolean[] _uniqueGroupKeysFlag;

  // For MAP_BASED storage type.
  private Long2IntOpenHashMap _groupKeyToId;

  /**
   * Constructor for the class. Initializes data members (reusable arrays).
   *
   * @param indexSegment
   * @param groupByColumns
   * @param maxNumGroupKeys
   */
  DefaultGroupKeyGenerator(IndexSegment indexSegment, String[] groupByColumns, int maxNumGroupKeys) {
    _groupByColumns = groupByColumns;

    int numGroupByColumns = groupByColumns.length;
    _multiValueResulableArray = new int[numGroupByColumns][];

    _singleValIterators = new BlockSingleValIterator[numGroupByColumns];
    _multiValIterators = new BlockMultiValIterator[numGroupByColumns];

    _dictionaries = new Dictionary[numGroupByColumns];
    _cardinalities = new int[numGroupByColumns];
    _isSingleValueGroupByColumn = new boolean[numGroupByColumns];

    for (int i = 0; i < numGroupByColumns; i++) {
      DataSource dataSource = indexSegment.getDataSource(_groupByColumns[i]);

      _dictionaries[i] = dataSource.getDictionary();
      _cardinalities[i] = dataSource.getDataSourceMetadata().cardinality();
      _isSingleValueGroupByColumn[i] = (dataSource.getDataSourceMetadata().isSingleValue()) ? true : false;

      Block block = dataSource.nextBlock();

      if (_isSingleValueGroupByColumn[i]) {
        _singleValIterators[i] = (BlockSingleValIterator) block.getBlockValueSet().iterator();
      } else {
        int maxNumberOfMultiValues = block.getMetadata().getMaxNumberOfMultiValues();
        _multiValueResulableArray[i] = new int[maxNumberOfMultiValues];
        _multiValIterators[i] = (BlockMultiValIterator) block.getBlockValueSet().iterator();
      }
    }

    _numUniqueGroupKeys = 0;
    if (maxNumGroupKeys <= ResultHolderFactory.MAX_INITIAL_RESULT_HOLDER_CAPACITY) {
      _storageType = STORAGE_TYPE.ARRAY_BASED;
      _uniqueGroupKeysFlag = new boolean[maxNumGroupKeys];
    } else {
      _storageType = STORAGE_TYPE.MAP_BASED;
      _groupKeyToId = new Long2IntOpenHashMap();
      _groupKeyToId.defaultReturnValue(INVALID_ID);
    }

    _reusableGroupByValuesArray = new int[numGroupByColumns];
  }

  /**
   * Generate group-by key for a given array of column values, and corresponding
   * column cardinality.
   *
   *
   * @param values
   * @param cardinalities
   * @return
   */
  public static long generateRawKey(int[] values, int[] cardinalities) {
    long groupKey = 0;
    for (int i = 0; i < values.length; i++) {
      groupKey = groupKey * cardinalities[i] + values[i];
    }
    return groupKey;
  }

  /**
   * Decode the individual column values from a given group by key and the individual
   * column cardinalities.
   *
   * @param rawGroupKey
   * @param cardinalities
   * @param decoded
   */
  public static void decodeRawGroupKey(long rawGroupKey, int[] cardinalities, int[] decoded) {
    int length = cardinalities.length;

    for (int i = length - 1; i >= 0; --i) {
      decoded[i] = (int) (rawGroupKey % cardinalities[i]);
      rawGroupKey = rawGroupKey / cardinalities[i];
    }
  }

  /**
   * {@inheritDoc}
   *
   * The group-by key generated is always 'int'. For array-based storage, the rawKey is
   * guaranteed to be < 10K. And for map-based storage, we map the rawKey to an index,
   * which is also guaranteed to fit in 'int'.
   *
   * @param docId
   * @return
   */
  @Override
  public int generateKeyForDocId(int docId) {
    for (int i = 0; i < _groupByColumns.length; i++) {
      if (_isSingleValueGroupByColumn[i]) {
        _singleValIterators[i].skipTo(docId);
        _reusableGroupByValuesArray[i] = _singleValIterators[i].nextIntVal();
      } else {
        throw new RuntimeException("GroupBy on multi-valued columns not supported.");
      }
    }

    long rawKey = generateRawKey(_reusableGroupByValuesArray, _cardinalities);
    return updateRawKeyToGroupKeyMapping(rawKey);
  }

  /**
   * {@inheritDoc}
   * @param docIdSet
   * @param startIndex
   * @param length
   * @param docIdToGroupKey
   */
  @Override
  public void generateKeysForDocIdSet(int[] docIdSet, int startIndex, int length, int[] docIdToGroupKey) {
    int endIndex = startIndex + length - 1;
    for (int i = startIndex; i <= endIndex; ++i) {
      int docId = docIdSet[i];
      docIdToGroupKey[i] = generateKeyForDocId(docId);
    }
  }

  /**
   * Given a docId, generate and return an array of unique group by keys, based on the
   * values of group by columns of the record of given docId.
   *
   * @param docId
   * @return
   */
  public int[] generateKeysForDocId(int docId) {
    List<Long> groupByKeys = new ArrayList<>();<commit_msg>Optimize for doc level function in group key generator (#150)<commit_after>/**
 * Copyright (C) 2014-2015 LinkedIn Corp. (<EMAIL>)
 *
 * Licensed 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 com.linkedin.pinot.core.operator.aggregation.groupby;

import com.linkedin.pinot.core.common.Block;
import com.linkedin.pinot.core.common.BlockMultiValIterator;
import com.linkedin.pinot.core.common.BlockSingleValIterator;
import com.linkedin.pinot.core.common.DataSource;
import com.linkedin.pinot.core.indexsegment.IndexSegment;
import com.linkedin.pinot.core.operator.aggregation.ResultHolderFactory;
import com.linkedin.pinot.core.query.aggregation.groupby.GroupByConstants;
import com.linkedin.pinot.core.segment.index.readers.Dictionary;
import it.unimi.dsi.fastutil.longs.Long2IntOpenHashMap;
import it.unimi.dsi.fastutil.objects.ObjectIterator;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Map;


/**
 * Class for generating group-by keys for a given list of group-by columns.
 * Uses dictionary id's to generate group-by keys.
 *
 * If maximum number of possible group keys are less than a threshold (10K),
 * returns the group key generated based on dictionary id's for the group by columns.
 *
 * For larger values, maps the actual generated group by keys onto contiguous
 * indices, and returns the latter. Keeps a mapping between the two internally.
 *
 */
public class DefaultGroupKeyGenerator implements GroupKeyGenerator {
  private static final int INVALID_ID = -1;

  public enum STORAGE_TYPE {
    ARRAY_BASED,
    MAP_BASED
  }

  private final String[] _groupByColumns;

  private final BlockSingleValIterator[] _singleValIterators;
  private final BlockMultiValIterator[] _multiValIterators;

  private final Dictionary[] _dictionaries;
  private final int[] _cardinalities;

  // A reusable array of number of group by columns size,
  // to avoid creating int[] objects with each call.
  private final int[] _reusableGroupByValuesArray;
  private final int[][] _multiValueResulableArray;

  private final boolean[] _isSingleValueGroupByColumn;
  private STORAGE_TYPE _storageType;
  private int _numUniqueGroupKeys;

  // For ARRAY_BASED storage type.
  private boolean[] _uniqueGroupKeysFlag;

  // For MAP_BASED storage type.
  private Long2IntOpenHashMap _groupKeyToId;

  /**
   * Constructor for the class. Initializes data members (reusable arrays).
   *
   * @param indexSegment
   * @param groupByColumns
   * @param maxNumGroupKeys
   */
  DefaultGroupKeyGenerator(IndexSegment indexSegment, String[] groupByColumns, int maxNumGroupKeys) {
    _groupByColumns = groupByColumns;

    int numGroupByColumns = groupByColumns.length;
    _multiValueResulableArray = new int[numGroupByColumns][];

    _singleValIterators = new BlockSingleValIterator[numGroupByColumns];
    _multiValIterators = new BlockMultiValIterator[numGroupByColumns];

    _dictionaries = new Dictionary[numGroupByColumns];
    _cardinalities = new int[numGroupByColumns];
    _isSingleValueGroupByColumn = new boolean[numGroupByColumns];

    for (int i = 0; i < numGroupByColumns; i++) {
      DataSource dataSource = indexSegment.getDataSource(_groupByColumns[i]);

      _dictionaries[i] = dataSource.getDictionary();
      _cardinalities[i] = dataSource.getDataSourceMetadata().cardinality();
      _isSingleValueGroupByColumn[i] = (dataSource.getDataSourceMetadata().isSingleValue()) ? true : false;

      Block block = dataSource.nextBlock();

      if (_isSingleValueGroupByColumn[i]) {
        _singleValIterators[i] = (BlockSingleValIterator) block.getBlockValueSet().iterator();
      } else {
        int maxNumberOfMultiValues = block.getMetadata().getMaxNumberOfMultiValues();
        _multiValueResulableArray[i] = new int[maxNumberOfMultiValues];
        _multiValIterators[i] = (BlockMultiValIterator) block.getBlockValueSet().iterator();
      }
    }

    _numUniqueGroupKeys = 0;
    if (maxNumGroupKeys <= ResultHolderFactory.MAX_INITIAL_RESULT_HOLDER_CAPACITY) {
      _storageType = STORAGE_TYPE.ARRAY_BASED;
      _uniqueGroupKeysFlag = new boolean[maxNumGroupKeys];
    } else {
      _storageType = STORAGE_TYPE.MAP_BASED;
      _groupKeyToId = new Long2IntOpenHashMap();
      _groupKeyToId.defaultReturnValue(INVALID_ID);
    }

    _reusableGroupByValuesArray = new int[numGroupByColumns];
  }

  /**
   * Generate group-by key for a given array of column values, and corresponding
   * column cardinality.
   *
   *
   * @param values
   * @param cardinalities
   * @return
   */
  public static long generateRawKey(int[] values, int[] cardinalities) {
    long groupKey = 0;
    for (int i = 0; i < values.length; i++) {
      groupKey = groupKey * cardinalities[i] + values[i];
    }
    return groupKey;
  }

  /**
   * Decode the individual column values from a given group by key and the individual
   * column cardinalities.
   *
   * @param rawGroupKey
   * @param cardinalities
   * @param decoded
   */
  public static void decodeRawGroupKey(long rawGroupKey, int[] cardinalities, int[] decoded) {
    int length = cardinalities.length;

    for (int i = length - 1; i >= 0; --i) {
      decoded[i] = (int) (rawGroupKey % cardinalities[i]);
      rawGroupKey = rawGroupKey / cardinalities[i];
    }
  }

  /**
   * {@inheritDoc}
   *
   * The group-by key generated is always 'int'. For array-based storage, the rawKey is
   * guaranteed to be < 10K. And for map-based storage, we map the rawKey to an index,
   * which is also guaranteed to fit in 'int'.
   *
   * @param docId
   * @return
   */
  @Override
  public int generateKeyForDocId(int docId) {
    long rawKey = 0;

    for (int i = 0; i < _groupByColumns.length; i++) {
      BlockSingleValIterator singleValIterator = _singleValIterators[i];
      singleValIterator.skipTo(docId);
      rawKey = rawKey * _cardinalities[i] + singleValIterator.nextIntVal();
    }

    return updateRawKeyToGroupKeyMapping(rawKey);
  }

  /**
   * {@inheritDoc}
   * @param docIdSet
   * @param startIndex
   * @param length
   * @param docIdToGroupKey
   */
  @Override
  public void generateKeysForDocIdSet(int[] docIdSet, int startIndex, int length, int[] docIdToGroupKey) {
    int endIndex = startIndex + length - 1;
    for (int i = startIndex; i <= endIndex; ++i) {
      int docId = docIdSet[i];
      docIdToGroupKey[i] = generateKeyForDocId(docId);
    }
  }

  /**
   * Given a docId, generate and return an array of unique group by keys, based on the
   * values of group by columns of the record of given docId.
   *
   * @param docId
   * @return
   */
  public int[] generateKeysForDocId(int docId) {
    List<Long> groupByKeys = new ArrayList<>();
{'fraction_non_alphanumeric': 0.09284532671629446, 'fraction_numerical': 0.0041356492969396195, 'mean_word_length': 3.8411745078411745, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3146858', 'n_tokens_mistral': 4398, 'n_tokens_neox': 4249, 'n_words': 1219}