Make initialSampler optional for LP predict

Co-Authored-By: Jacob Sznajdman <[email protected]>
Co-Authored-By: Florentin Dörre <[email protected]>

Author: 3 people

algo/src/main/java/org/neo4j/gds/similarity/knn/KnnSampler.java

@@ -22,6 +22,7 @@
 import java.util.Arrays;
 import java.util.List;
 import java.util.Locale;
+import java.util.Optional;
 import java.util.function.LongPredicate;
 import java.util.stream.Collectors;
 
@@ -67,6 +68,10 @@ else if (input instanceof SamplerType) {
             ));
         }
 
+        public static Optional<SamplerType> parseToOptional(String input) {
+            return Optional.of(parse(input));
+        }
+
         public static String toString(SamplerType samplerType) {
             return samplerType.toString();
         }

alpha/alpha-proc/src/main/java/org/neo4j/gds/ml/linkmodels/pipeline/predict/LinkPredictionPredictPipelineBaseConfig.java

@@ -44,6 +44,7 @@
 public interface LinkPredictionPredictPipelineBaseConfig extends AlgoBaseConfig, SingleThreadedRandomSeedConfig, ModelConfig {
 
     double DEFAULT_THRESHOLD = 0.0;
+    String MISSING_INITIAL_SAMPLER = "MISSING_VALUE";
 
     //TODO make this a parameter
     String graphName();
@@ -75,11 +76,19 @@ default double sampleRate() {
     @Configuration.IntegerRange(min = 0)
     Optional<Integer> randomJoins();
 
-    @Value.Default
-    @Configuration.ConvertWith("org.neo4j.gds.similarity.knn.KnnSampler.SamplerType#parse")
-    @Configuration.ToMapValue("org.neo4j.gds.similarity.knn.KnnSampler.SamplerType#toString")
-    default KnnSampler.SamplerType initialSampler() {
-        return KnnSampler.SamplerType.UNIFORM;
+    default String initialSampler() {
+        return MISSING_INITIAL_SAMPLER;
+    }
+
+    @Value.Derived
+    @Configuration.Ignore
+    default Optional<KnnSampler.SamplerType> derivedSampler() {
+        String sampler = initialSampler();
+        if (sampler.equals(MISSING_INITIAL_SAMPLER)) {
+            return Optional.empty();
+        }
+
+        return KnnSampler.SamplerType.parseToOptional(sampler);
     }
 
     @Value.Check
@@ -96,7 +105,7 @@ default void validateParameterCombinations() {
                 "deltaThreshold", deltaThreshold().isPresent(),
                 "maxIterations", maxIterations().isPresent(),
                 "randomJoins", randomJoins().isPresent(),
-                "initialSampler", randomJoins().isPresent());
+                "initialSampler", derivedSampler().isPresent());
             validateStrategySpecificParameters(approximateStrategyParameters, "less than 1");
 
             topN().orElseThrow(()-> MissingParameterExceptions.missingValueFor("topN", Collections.emptyList()));
@@ -132,13 +141,13 @@ default KnnBaseConfig approximateConfig() {
             .sampleRate(sampleRate())
             .nodeProperties(List.of(new KnnNodePropertySpec("NotUsedInLP")))
             .minBatchSize(LinkPrediction.MIN_NODE_BATCH_SIZE)
-            .initialSampler(initialSampler())
             .concurrency(concurrency());
 
         topK().ifPresent(knnBuilder::topK);
         deltaThreshold().ifPresent(knnBuilder::deltaThreshold);
         maxIterations().ifPresent(knnBuilder::maxIterations);
         randomJoins().ifPresent(knnBuilder::randomJoins);
+        derivedSampler().ifPresent(knnBuilder::initialSampler);
         randomSeed().ifPresent(knnBuilder::randomSeed);
 
         return knnBuilder.build();

alpha/alpha-proc/src/test/java/org/neo4j/gds/ml/linkmodels/pipeline/predict/LinkPredictionPipelineStreamProcTest.java

@@ -50,8 +50,7 @@ void shouldPredictWithTopN(int concurrency, String nodeLabel) {
             " modelName: 'model'," +
             " threshold: 0," +
             " topN: $topN," +
-            " concurrency:" +
-            " $concurrency" +
+            " concurrency: $concurrency" +
             "})" +
             "YIELD node1, node2, probability" +
             " RETURN node1, node2, probability" +
@@ -111,8 +110,7 @@ void shouldPredictWithInitialSamplerSet() {
             " randomSeed: 42," +
             " topK: $topK," +
             " initialSampler: 'randomWalk'," +
-            " concurrency:" +
-            " $concurrency" +
+            " concurrency: $concurrency" +
             "})" +
             "YIELD node1, node2, probability" +
             " RETURN node1, node2, probability" +