Fixes MAHOUT-1180

git-svn-id: https://svn.apache.org/repos/asf/mahout/trunk@1478723 13f79535-47bb-0310-9956-ffa450edef68
apache · May 3, 2013 · bc1b16d · bc1b16d
1 parent 7e0978e
commit bc1b16d
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 10 deletions.
diff --git a/math/src/main/java/org/apache/mahout/math/random/Multinomial.java b/math/src/main/java/org/apache/mahout/math/random/Multinomial.java
@@ -17,18 +17,20 @@
 
 package org.apache.mahout.math.random;
 
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.Random;
+
 import com.google.common.base.Preconditions;
+import com.google.common.collect.AbstractIterator;
+import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Multiset;
 import org.apache.mahout.common.RandomUtils;
 import org.apache.mahout.math.list.DoubleArrayList;
 
-import java.util.Iterator;
-import java.util.List;
-import java.util.Map;
-import java.util.Random;
-
 /**
  * Multinomial sampler that allows updates to element probabilities.  The basic idea is that sampling is
  * done by using a simple balanced tree.  Probabilities are kept in the tree so that we can navigate to
@@ -66,6 +68,7 @@ public Multinomial(Iterable<WeightedThing<T>> things) {
   }
 
   public void add(T value, double w) {
+    Preconditions.checkNotNull(value);
     Preconditions.checkArgument(!items.containsKey(value));
 
     int n = this.weight.size();
@@ -182,6 +185,18 @@ List<Double> getWeights() {
 
   @Override
   public Iterator<T> iterator() {
-    return items.keySet().iterator();
+    return new AbstractIterator<T>() {
+      Iterator<T> valuesIterator = Iterables.skip(values, 1).iterator();
+      @Override
+      protected T computeNext() {
+        while (valuesIterator.hasNext()) {
+          T next = valuesIterator.next();
+          if (items.containsKey(next)) {
+            return next;
+          }
+        }
+        return endOfData();
+      }
+    };
   }
 }
diff --git a/math/src/test/java/org/apache/mahout/math/random/MultinomialTest.java b/math/src/test/java/org/apache/mahout/math/random/MultinomialTest.java
@@ -17,6 +17,10 @@
 
 package org.apache.mahout.math.random;
 
+import java.util.List;
+import java.util.Map;
+import java.util.Random;
+
 import com.google.common.collect.HashMultiset;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Lists;
@@ -26,10 +30,6 @@
 import org.junit.Before;
 import org.junit.Test;
 
-import java.util.List;
-import java.util.Map;
-import java.util.Random;
-
 public class MultinomialTest extends MahoutTestCase {
     @Override
     @Before
@@ -166,6 +166,27 @@ public void testInsert() {
     }
 
     @Test
+  public void testSetZeroWhileIterating() {
+    Multinomial<Integer> table = new Multinomial<Integer>();
+    for (int i = 0; i < 10000; ++i) {
+      table.add(i, i);
+    }
+    // Setting a sample's weight to 0 removes from the items map.
+    // If that map is used when iterating (it used to be), it will
+    // trigger a ConcurrentModificationException.
+    for (Integer sample : table) {
+      table.set(sample, 0);
+    }
+  }
+
+  @Test(expected=NullPointerException.class)
+  public void testNoNullValuesAllowed() {
+    Multinomial<Integer> table = new Multinomial<Integer>();
+    // No null values should be allowed.
+    table.add(null, 1);
+  }
+
+  @Test
     public void testDeleteAndUpdate() {
         Random rand = RandomUtils.getRandom();
         Multinomial<Integer> table = new Multinomial<Integer>();