hbase-common/src/main/java/org/apache/hadoop/hbase/util/ReservoirSample.java

   1 /**
   2  * Licensed to the Apache Software Foundation (ASF) under one
   3  * or more contributor license agreements.  See the NOTICE file
   4  * distributed with this work for additional information
   5  * regarding copyright ownership.  The ASF licenses this file
   6  * to you under the Apache License, Version 2.0 (the
   7  * "License"); you may not use this file except in compliance
   8  * with the License.  You may obtain a copy of the License at
   9  *
  10  *     http://www.apache.org/licenses/LICENSE-2.0
  11  *
  12  * Unless required by applicable law or agreed to in writing, software
  13  * distributed under the License is distributed on an "AS IS" BASIS,
  14  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  15  * See the License for the specific language governing permissions and
  16  * limitations under the License.
  17  */
  18 package org.apache.hadoop.hbase.util;
  19
  20 import java.util.ArrayList;
  21 import java.util.Iterator;
  22 import java.util.List;
  23 import java.util.concurrent.ThreadLocalRandom;
  24 import java.util.stream.Stream;
  25 import org.apache.yetus.audience.InterfaceAudience;
  26
  27 import org.apache.hbase.thirdparty.com.google.common.base.Preconditions;
  28
  29 /**
  30  * The simple version of reservoir sampling implementation. It is enough for the usage in HBase.
  31  * <p/>
  32  * See https://en.wikipedia.org/wiki/Reservoir_sampling.
  33  */
  34 @InterfaceAudience.Private
  35 public class ReservoirSample<T> {
  36
  37   private final List<T> r;
  38
  39   private final int k;
  40
  41   private int n;
  42
  43   public ReservoirSample(int k) {
  44     Preconditions.checkArgument(k > 0, "negative sampling number(%d) is not allowed");
  45     r = new ArrayList<>(k);
  46     this.k = k;
  47   }
  48
  49   public void add(T t) {
  50     if (n < k) {
  51       r.add(t);
  52     } else {
  53       int j = ThreadLocalRandom.current().nextInt(n + 1);
  54       if (j < k) {
  55         r.set(j, t);
  56       }
  57     }
  58     n++;
  59   }
  60
  61   public void add(Iterator<T> iter) {
  62     iter.forEachRemaining(this::add);
  63   }
  64
  65   public void add(Stream<T> s) {
  66     s.forEachOrdered(this::add);
  67   }
  68
  69   public List<T> getSamplingResult() {
  70     return r;
  71   }
  72 }