content: package org.apache.cassandra.stress.generate;
  /*
   * 
   * 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.
   * 
   */
  
  
  import java.nio.ByteBuffer;
  import java.util.*;
  import java.util.stream.Collectors;
  
  import com.google.common.collect.Iterables;
  
  import org.apache.cassandra.stress.generate.values.Generator;
  
  public class PartitionGenerator
  {
  
      public static enum Order
      {
          ARBITRARY, SHUFFLED, SORTED
      }
  
      public final double maxRowCount;
      public final double minRowCount;
      final List<Generator> partitionKey;
      final List<Generator> clusteringComponents;
      final List<Generator> valueComponents;
      final int[] clusteringDescendantAverages;
      final int[] clusteringComponentAverages;
  
      private final Map<String, Integer> indexMap;
      final Order order;
  
      public PartitionGenerator(List<Generator> partitionKey, List<Generator> clusteringComponents, List<Generator> valueComponents, Order order)
      {
          this.partitionKey = partitionKey;
          this.clusteringComponents = clusteringComponents;
          this.valueComponents = valueComponents;
          this.order = order;
          this.clusteringDescendantAverages = new int[clusteringComponents.size()];
          this.clusteringComponentAverages = new int[clusteringComponents.size()];
          for (int i = 0 ; i < clusteringComponentAverages.length ; i++)
              clusteringComponentAverages[i] = (int) clusteringComponents.get(i).clusteringDistribution.average();
          for (int i = clusteringDescendantAverages.length - 1 ; i >= 0 ; i--)
              clusteringDescendantAverages[i] = (int) (i < (clusteringDescendantAverages.length - 1) ? clusteringComponentAverages[i + 1] * clusteringDescendantAverages[i + 1] : 1);
          double maxRowCount = 1d;
          double minRowCount = 1d;
          for (Generator component : clusteringComponents)
          {
              maxRowCount *= component.clusteringDistribution.maxValue();
              minRowCount *= component.clusteringDistribution.minValue();
          }
          this.maxRowCount = maxRowCount;
          this.minRowCount = minRowCount;
          this.indexMap = new LinkedHashMap<>();
          int i = 0;
          for (Generator generator : partitionKey)
              indexMap.put(generator.name, --i);
          i = 0;
          for (Generator generator : Iterables.concat(clusteringComponents, valueComponents))
              indexMap.put(generator.name, i++);
      }
  
      public boolean permitNulls(int index)
      {
          return !(index < 0 || index < clusteringComponents.size());
      }
  
      public List<Generator> getPartitionKey()
      {
          return Collections.unmodifiableList(partitionKey);
      }
  
      public List<Generator> getClusteringComponents()
      {
          return Collections.unmodifiableList(clusteringComponents);
      }
  
      public int indexOf(String name)
      {
          Integer i = indexMap.get(name);
          if (i == null)
              throw new NoSuchElementException();
          return i;
      }
  
      public ByteBuffer convert(int c, Object v)
      {
          if (c < 0)
              return partitionKey.get(-1-c).type.decomposeUntyped(v);
          if (c < clusteringComponents.size())
              return clusteringComponents.get(c).type.decomposeUntyped(v);
          return valueComponents.get(c - clusteringComponents.size()).type.decomposeUntyped(v);
      }
  
      public Object convert(int c, ByteBuffer v)
      {
          if (c < 0)
              return partitionKey.get(-1-c).type.compose(v);
          if (c < clusteringComponents.size())
              return clusteringComponents.get(c).type.compose(v);
          return valueComponents.get(c - clusteringComponents.size()).type.compose(v);
      }
  
      public List<String> getColumnNames()
      {
          return indexMap.keySet().stream().collect(Collectors.toList());
      }
  }
  