datafaker-net · kingthorin · Nov 24, 2024 · Oct 21, 2024 · Oct 30, 2024 · Oct 30, 2024
diff --git a/docs/documentation/custom-providers.md b/docs/documentation/custom-providers.md
@@ -21,8 +21,15 @@ Create a custom provider of data:
 === "Java"
 
     ``` java
-    public class Insect extends AbstractProvider<BaseProviders> {
-        private static final String[] INSECT_NAMES = new String[]{"Ant", "Beetle", "Butterfly", "Wasp"};
+    public static class Insect extends AbstractProvider<BaseProviders> {
+        private static final WeightedRandomSelector selector = new WeightedRandomSelector(new Random());
+
+        private static final String[] INSECT_NAMES = { "Ant", "Beetle", "Butterfly", "Wasp" };
+        private static final List<Map<String, Object>> WEIGHTED_INSECTS = List.of(
+            Map.of("value", "Driver ant", "weight", 6.0),
+            Map.of("value", "Fire ant", "weight", 3.0),
+            Map.of("value", "Harvester ant", "weight", 1.0)
+        );
 
         public Insect(BaseProviders faker) {
             super(faker);
@@ -31,10 +38,13 @@ Create a custom provider of data:
         public String nextInsectName() {
             return INSECT_NAMES[faker.random().nextInt(INSECT_NAMES.length)];
         }
+
+        public String weightedInsectName() {
+            return selector.select(WEIGHTED_INSECTS);
+        }
     }
     ```
-
-
+
 ### Register provider
 
 Create your own custom faker, which extends `Faker`, and register the custom provider:
@@ -66,6 +76,22 @@ This will print something like the following:
 Wasp
 ```
 
+**Usafe of weigted random selector is in the POC stage and is currently available only for custom hardcoded providers.**
+
+To use a random selector based on weights, you can do the following:
+
+=== "Java"
+
+    ``` java
+    MyCustomFaker myFaker = new MyCustomFaker();
+    System.out.println(myFaker.insect().weightedInsectName());
+    ```
+
+This will return a random insect name but based on the provided weights
+```
+Driver ant
+```
+
 ## Custom provider using Yaml file
 
 In case you have a large set of data to load, it might be better to use a Yaml file.

diff --git a/src/main/java/net/datafaker/service/WeightedRandomSelector.java b/src/main/java/net/datafaker/service/WeightedRandomSelector.java
@@ -0,0 +1,136 @@
+package net.datafaker.service;
+
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Random;
+import java.util.Set;
+
+/**
+ * A utility class for selecting a random element from a list based on assigned weights.
+ **/
+public record WeightedRandomSelector(Random random) {
+    private static final String WEIGHT_KEY = "weight";
+    private static final String VALUE_KEY = "value";
+
+    public WeightedRandomSelector(Random random) {
+        this.random = random != null ? random : new Random();
+    }
+
+    /**
+     * Returns a weighted random element from the given list, where each element is represented as a Map
+     * containing a weight and the corresponding value.
+     * <p>
+     *
+     * @param items A list of maps, where each map contains:
+     *              - weight: A Double representing the weight of the element, influencing its selection probability.
+     *              - value: The actual element of type T to be randomly selected based on its weight.
+     * @param <T>   The type of the element to be selected from the list. The value associated with the weight can be of any type.
+     * @return A randomly selected element based on its weight.
+     * @throws IllegalArgumentException if:
+     *                                  - the list is null or empty,
+     *                                  - any item in the list is null or empty,
+     *                                  - the item does not contain 'weight' or 'value' keys,
+     *                                  - any weight is null, non-positive, NaN or infinite,
+     *                                  - any values in the list are not unique or null,
+     *                                  - the sum of weights exceeds Double.MAX_VALUE.
+     */
+    public <T> T select(List<Map<String, Object>> items) {
-    public <T> T select(List<Map<String, Object>> items) {
+    public static <T> T select(List<Map<String, Object>> items) {
-    public <T> T select(List<Map<String, Object>> items) {
+    public static <T> T select(List<Map<String, Object>> items) {
+        validateItemsList(items);
+
+        Object[] values = new Object[items.size()];
+        double[] cumulativeWeights = preprocessItems(items, values);
+
+        double randomValue = random.nextDouble() * cumulativeWeights[cumulativeWeights.length - 1];
+        return selectWeightedElement(randomValue, cumulativeWeights, values);
+    }
+
+    private static void validateItemsList(List<Map<String, Object>> items) {
+        if (items == null) {
+            throw new IllegalArgumentException("Input list cannot be null");
+        }
+        if (items.isEmpty()) {
+            throw new IllegalArgumentException("Input list cannot be empty");
+        }
+
+        Set<Object> uniqueValues = new HashSet<>();
+
+        for (var item : items) {
+            validateItem(item);
+            assertUniqueValues(item, uniqueValues);
+        }
+    }
+
+    private static void assertUniqueValues(Map<String, Object> item, Set<Object> values) {
+        Object value = item.get(VALUE_KEY);
+        if (!values.add(value)) {
+            throw new IllegalArgumentException("Duplicate value found: " + value + ". Values must be unique.");
+        }
+    }
+
+    private static void validateItem(Map<String, Object> item) {
+        if (item == null) {
+            throw new IllegalArgumentException("Item cannot be null");
+        }
+        if (item.isEmpty()) {
+            throw new IllegalArgumentException("Item cannot be empty");
+        }
+        if (!item.containsKey(WEIGHT_KEY) || !item.containsKey(VALUE_KEY)) {
+            throw new IllegalArgumentException("Each item must contain 'weight' and 'value' keys");
+        }
+        validateValue(item.get(VALUE_KEY));
+        validateWeight(item.get(WEIGHT_KEY));
+    }
+
+    private static void validateValue(Object valueObj) {
+        if (valueObj == null) {
+            throw new IllegalArgumentException("Value cannot be null");
+        }
+    }
+
+    private static void validateWeight(Object weightObj) {
+        if (!(weightObj instanceof Double weight)) {
+            throw new IllegalArgumentException("Weight must be a non-null Double");
+        }
+        if (weight < 0 || Double.isNaN(weight) || Double.isInfinite(weight)) {
+            throw new IllegalArgumentException("Weight must be a non-negative number and cannot be NaN or infinite");
+        }
+    }
+
+    private static void validateTotalWeight(double totalWeight) {
+        if (totalWeight <= 0) {
+            throw new IllegalArgumentException("The total weight must be greater than 0. At least one item must have a positive weight");
+        }
+    }
+
+    static double[] preprocessItems(List<Map<String, Object>> items, Object[] values) {
+        double[] cumulativeWeights = new double[items.size()];
+
+        double totalWeight = 0.0;
+        for (int i = 0; i < items.size(); i++) {
+            double weight = (Double) items.get(i).get(WEIGHT_KEY);
+            if (Double.MAX_VALUE - totalWeight < weight) {
+                throw new IllegalArgumentException("Sum of the weights exceeds Double.MAX_VALUE");
+            }
+            totalWeight += weight;
+            cumulativeWeights[i] = totalWeight;
+            values[i] = items.get(i).get(VALUE_KEY);
+        }
+
+        validateTotalWeight(totalWeight);
+
+        return cumulativeWeights;
+    }
+
+    static <T> T selectWeightedElement(double randomValue, double[] cumulativeWeights, Object[] values) {
+        int index = Arrays.binarySearch(cumulativeWeights, randomValue);
+        index = (index < 0) ? -index - 1 : index;
+
+        if (index >= cumulativeWeights.length) {
+            index = cumulativeWeights.length - 1;
+        }
-        if (index >= cumulativeWeights.length) {
-            index = cumulativeWeights.length - 1;
-        }
+        index = Math.max(index, cumulativeWeights.length - 1);
-        if (index >= cumulativeWeights.length) {
-            index = cumulativeWeights.length - 1;
-        }
+        index = Math.max(index, cumulativeWeights.length - 1);
+
+        return (T) values[index];
+    }
+}
diff --git a/src/test/java/net/datafaker/providers/base/CustomFakerTest.java b/src/test/java/net/datafaker/providers/base/CustomFakerTest.java
@@ -4,13 +4,18 @@
 import org.junit.jupiter.api.Test;
 
 import java.nio.file.Paths;
+import java.util.HashMap;
+import java.util.List;
 import java.util.Locale;
+import java.util.Map;
+import java.util.Random;
 
+import net.datafaker.service.WeightedRandomSelector;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 /**
- * This is a demo of how to create a custom faker and register a custom faker in it.
+ * This is a demo of how to create a custom data provider and register a custom faker to use it.
  */
 class CustomFakerTest {
     public static class MyCustomFaker extends BaseFaker {
@@ -24,7 +29,14 @@ public InsectFromFile insectFromFile() {
     }
 
     public static class Insect extends AbstractProvider<BaseProviders> {
-        private static final String[] INSECT_NAMES = {"Ant", "Beetle", "Butterfly", "Wasp"};
+        private static final WeightedRandomSelector selector = new WeightedRandomSelector(new Random());
+
+        private static final String[] INSECT_NAMES = { "Ant", "Beetle", "Butterfly", "Wasp" };
+        private static final List<Map<String, Object>> WEIGHTED_INSECTS = List.of(
+            Map.of("value", "Driver ant", "weight", 6.0),
+            Map.of("value", "Fire ant", "weight", 3.0),
+            Map.of("value", "Harvester ant", "weight", 1.0)
+        );
 
         public Insect(BaseProviders faker) {
             super(faker);
@@ -33,6 +45,10 @@ public Insect(BaseProviders faker) {
         public String nextInsectName() {
             return INSECT_NAMES[faker.random().nextInt(INSECT_NAMES.length)];
         }
+
+        public String weightedInsectName() {
+            return selector.select(WEIGHTED_INSECTS);
+        }
     }
 
     public static class InsectFromFile extends AbstractProvider<BaseProviders> {
@@ -107,18 +123,31 @@ class InsectFaker extends BaseFaker {
             public InsectFaker(Locale locale) {
                 super(locale);
             }
-
-            public Insect insect() {
-                return getProvider(Insect.class, Insect::new);
-            }
         }
         BaseFaker faker1 = new InsectFaker(Locale.ENGLISH);
         BaseFaker faker2 = new InsectFaker(Locale.GERMAN);
 
-        Insect insect1 = faker1.getProvider("Insect");
-        Insect insect2 = faker2.getProvider("Insect");
+        Insect insect1 = faker1.getProvider(Insect.class, Insect::new);
+        Insect insect2 = faker2.getProvider(Insect.class, Insect::new);
         assertThat(insect1).isNotNull();
         assertThat(insect2).isNotNull();
         assertThat(insect1).isNotSameAs(insect2);
     }
+
+    @Test
+    void weightedInsectNameTest() {
+        MyCustomFaker myFaker = new MyCustomFaker();
+        Map<String, Integer> insectCounts = new HashMap<>();
+        insectCounts.put("Driver ant", 0);
+        insectCounts.put("Fire ant", 0);
+        insectCounts.put("Harvester ant", 0);
+
+        for (int i = 0; i < 100; i++) {
+            String selectedInsect = myFaker.insect().weightedInsectName();
+            insectCounts.put(selectedInsect, insectCounts.get(selectedInsect) + 1);
+        }
+
+        assertThat(insectCounts.get("Driver ant")).isGreaterThan(insectCounts.get("Fire ant"));
+        assertThat(insectCounts.get("Fire ant")).isGreaterThan(insectCounts.get("Harvester ant"));
+    }
 }