fix(#1552): use user provided weights instead of uniform distribution

finos · Jul 15, 2020 · 954273e · 954273e
1 parent da3646a
commit 954273e
Show file tree

Hide file tree

Showing 6 changed files with 27 additions and 4 deletions.
diff --git a/...on/src/main/java/com/scottlogic/datahelix/generator/common/whitelist/DistributedList.java b/...on/src/main/java/com/scottlogic/datahelix/generator/common/whitelist/DistributedList.java
@@ -54,6 +54,16 @@ public static <T> DistributedList<T> singleton(final T element) {
         return DistributedList.uniform(Collections.singleton(element));
     }
 
+    public static <T> DistributedList<T> weightedOrDefault(final Collection<T> underlyingSet) {
+        return new DistributedList<>(
+            underlyingSet.stream()
+                .map(element -> element instanceof WeightedElement
+                        ? (WeightedElement<T>) element
+                        : WeightedElement.withDefaultWeight(element)
+                )
+                .collect(Collectors.toList()));
+    }
+
     public static <T> DistributedList<T> uniform(final Collection<T> underlyingSet) {
         return new DistributedList<>(
             underlyingSet.stream()

diff --git a/...c/datahelix/generator/profile/factories/constraint_factories/AtomicConstraintFactory.java b/...c/datahelix/generator/profile/factories/constraint_factories/AtomicConstraintFactory.java
@@ -107,7 +107,7 @@ public InMapRelation createInMapRelation(InMapConstraintDTO dto, Fields fields)
 
     private InSetConstraint createInSetConstraint(InSetConstraintDTO dto, Field field)
     {
-        DistributedList<Object> values = DistributedList.uniform(dto.values.stream()
+        DistributedList<Object> values = DistributedList.weightedOrDefault(dto.values.stream()
             .distinct()
             .map(this::parseValue)
             .collect(Collectors.toList()));

diff --git a/.../datahelix/generator/profile/factories/constraint_factories/NumericConstraintFactory.java b/.../datahelix/generator/profile/factories/constraint_factories/NumericConstraintFactory.java
@@ -20,6 +20,7 @@
 import com.scottlogic.datahelix.generator.common.profile.Field;
 import com.scottlogic.datahelix.generator.common.profile.NumericGranularity;
 import com.scottlogic.datahelix.generator.common.util.NumberUtils;
+import com.scottlogic.datahelix.generator.common.whitelist.WeightedElement;
 import com.scottlogic.datahelix.generator.core.profile.constraints.atomic.*;
 import com.scottlogic.datahelix.generator.profile.dtos.constraints.atomic.GranularToConstraintDTO;
 import com.scottlogic.datahelix.generator.profile.dtos.constraints.atomic.integer.LongerThanConstraintDTO;
@@ -38,6 +39,13 @@ public class NumericConstraintFactory extends AtomicConstraintFactory
     @Override
     Object parseValue(Object value)
     {
+        if (value instanceof WeightedElement) {
+            WeightedElement weightedElement = (WeightedElement) value;
+            return new WeightedElement<>(
+                NumberUtils.coerceToBigDecimal(weightedElement.element()),
+                weightedElement.weight()
+            );
+        }
         return NumberUtils.coerceToBigDecimal(value);
     }
 

diff --git a/...java/com/scottlogic/datahelix/generator/profile/serialisation/ConstraintDeserializer.java b/...java/com/scottlogic/datahelix/generator/profile/serialisation/ConstraintDeserializer.java
@@ -35,6 +35,7 @@
 import java.io.IOException;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.stream.Collectors;
@@ -84,7 +85,7 @@ private InMapConstraintDTO map(InMapFromFileConstraintDTO dto)
 
     private InSetConstraintDTO map(InSetFromFileConstraintDTO dto)
     {
-        List<Object> values = fileReader.setFromFile(getFile(dto.file)).stream().collect(Collectors.toList());
+        List<Object> values = new ArrayList<>(fileReader.setFromFile(getFile(dto.file)).distributedList());
         InSetConstraintDTO inSetConstraintDTO = new InSetConstraintDTO();
         inSetConstraintDTO.field = dto.field;
         inSetConstraintDTO.values = values;

diff --git a/...ix/generator/profile/validators/profile/constraints/atomic/AtomicConstraintValidator.java b/...ix/generator/profile/validators/profile/constraints/atomic/AtomicConstraintValidator.java
@@ -18,6 +18,7 @@
 
 import com.scottlogic.datahelix.generator.common.profile.FieldType;
 import com.scottlogic.datahelix.generator.common.validators.ValidationResult;
+import com.scottlogic.datahelix.generator.common.whitelist.WeightedElement;
 import com.scottlogic.datahelix.generator.profile.dtos.FieldDTO;
 import com.scottlogic.datahelix.generator.profile.dtos.constraints.atomic.AtomicConstraintDTO;
 import com.scottlogic.datahelix.generator.profile.validators.profile.ConstraintValidator;
@@ -59,7 +60,9 @@ ValidationResult fieldTypeMustMatchValueType(T dto, Object value)
         {
             return ValidationResult.failure("Value " + value + " must be a boolean" + getErrorInfo(dto));
         }
-        if (!(value instanceof Number || value instanceof String && isNumber((String)value)) && fieldType == FieldType.NUMERIC)
+        if (!(value instanceof Number || value instanceof String && isNumber((String)value) ||
+            value instanceof WeightedElement && !(((WeightedElement) value).element() instanceof Number)) &&
+            fieldType == FieldType.NUMERIC)
         {
             return ValidationResult.failure("Value " + value + " must be a number" + getErrorInfo(dto));
         }

diff --git a/...st/java/com/scottlogic/datahelix/generator/profile/AtomicConstraintDeserialiserTests.java b/...st/java/com/scottlogic/datahelix/generator/profile/AtomicConstraintDeserialiserTests.java
@@ -19,6 +19,7 @@
 import com.fasterxml.jackson.core.JsonParseException;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.module.SimpleModule;
+import com.scottlogic.datahelix.generator.common.whitelist.WeightedElement;
 import com.scottlogic.datahelix.generator.profile.dtos.constraints.ConstraintDTO;
 import com.scottlogic.datahelix.generator.profile.dtos.constraints.atomic.EqualToConstraintDTO;
 import com.scottlogic.datahelix.generator.profile.dtos.constraints.atomic.GranularToConstraintDTO;
@@ -110,7 +111,7 @@ public void shouldDeserialiseInSetCsvFileWithoutException() throws IOException {
         // Assert
         InSetConstraintDTO expected = new InSetConstraintDTO();
         expected.field = "country";
-        expected.values = Collections.singletonList("test");
+        expected.values = Collections.singletonList(new WeightedElement<>("test", 1.0));
 
         assertThat(actual, sameBeanAs(expected));
     }