1   /**
2    * Licensed to the Apache Software Foundation (ASF) under one or more
3    * contributor license agreements. See the NOTICE file distributed with this
4    * work for additional information regarding copyright ownership. The ASF
5    * licenses this file to you under the Apache License, Version 2.0 (the
6    * "License"); you may not use this file except in compliance with the License.
7    * You may obtain a copy of the License at
8    *
9    * http://www.apache.org/licenses/LICENSE-2.0
10   *
11   * Unless required by applicable law or agreed to in writing, software
12   * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
13   * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
14   * License for the specific language governing permissions and limitations under
15   * the License.
16   */
17  package org.apache.hadoop.hbase.io.hfile;
18  
19  import java.util.Random;
20  
21  import org.apache.hadoop.io.BytesWritable;
22  import org.apache.hadoop.io.WritableComparator;
23  
24  /**
25   * Generate random <key, value> pairs.
26   * <p>
27   * Copied from
28   * <a href="https://issues.apache.org/jira/browse/HADOOP-3315">hadoop-3315 tfile</a>.
29   * Remove after tfile is committed and use the tfile version of this class
30   * instead.</p>
31   */
32  class KVGenerator {
33    private final Random random;
34    private final byte[][] dict;
35    private final boolean sorted;
36    private final RandomDistribution.DiscreteRNG keyLenRNG, valLenRNG;
37    private BytesWritable lastKey;
38    private static final int MIN_KEY_LEN = 4;
39    private final byte prefix[] = new byte[MIN_KEY_LEN];
40  
41    public KVGenerator(Random random, boolean sorted,
42        RandomDistribution.DiscreteRNG keyLenRNG,
43        RandomDistribution.DiscreteRNG valLenRNG,
44        RandomDistribution.DiscreteRNG wordLenRNG, int dictSize) {
45      this.random = random;
46      dict = new byte[dictSize][];
47      this.sorted = sorted;
48      this.keyLenRNG = keyLenRNG;
49      this.valLenRNG = valLenRNG;
50      for (int i = 0; i < dictSize; ++i) {
51        int wordLen = wordLenRNG.nextInt();
52        dict[i] = new byte[wordLen];
53        random.nextBytes(dict[i]);
54      }
55      lastKey = new BytesWritable();
56      fillKey(lastKey);
57    }
58  
59    private void fillKey(BytesWritable o) {
60      int len = keyLenRNG.nextInt();
61      if (len < MIN_KEY_LEN) len = MIN_KEY_LEN;
62      o.setSize(len);
63      int n = MIN_KEY_LEN;
64      while (n < len) {
65        byte[] word = dict[random.nextInt(dict.length)];
66        int l = Math.min(word.length, len - n);
67        System.arraycopy(word, 0, o.get(), n, l);
68        n += l;
69      }
70      if (sorted
71          && WritableComparator.compareBytes(lastKey.get(), MIN_KEY_LEN, lastKey
72              .getSize()
73              - MIN_KEY_LEN, o.get(), MIN_KEY_LEN, o.getSize() - MIN_KEY_LEN) > 0) {
74        incrementPrefix();
75      }
76  
77      System.arraycopy(prefix, 0, o.get(), 0, MIN_KEY_LEN);
78      lastKey.set(o);
79    }
80  
81    private void fillValue(BytesWritable o) {
82      int len = valLenRNG.nextInt();
83      o.setSize(len);
84      int n = 0;
85      while (n < len) {
86        byte[] word = dict[random.nextInt(dict.length)];
87        int l = Math.min(word.length, len - n);
88        System.arraycopy(word, 0, o.get(), n, l);
89        n += l;
90      }
91    }
92  
93    private void incrementPrefix() {
94      for (int i = MIN_KEY_LEN - 1; i >= 0; --i) {
95        ++prefix[i];
96        if (prefix[i] != 0) return;
97      }
98  
99      throw new RuntimeException("Prefix overflown");
100   }
101 
102   public void next(BytesWritable key, BytesWritable value, boolean dupKey) {
103     if (dupKey) {
104       key.set(lastKey);
105     }
106     else {
107       fillKey(key);
108     }
109     fillValue(value);
110   }
111 }