Merge pull request #9709 from soerenreichardt/cypher-aggregation-regression-fix

Use adaptive batch size for CypherAggregation progress logging

Author: soerenreichardt

cypher-aggregation/src/main/java/org/neo4j/gds/projection/GraphAggregator.java

@@ -39,6 +39,7 @@
 import org.neo4j.gds.core.loading.construction.NodeLabelTokens;
 import org.neo4j.gds.core.loading.construction.PropertyValues;
 import org.neo4j.gds.core.utils.ProgressTimer;
+import org.neo4j.gds.core.utils.progress.BatchingTaskProgressTracker;
 import org.neo4j.gds.core.utils.progress.TaskRegistryFactory;
 import org.neo4j.gds.core.utils.progress.TaskStore;
 import org.neo4j.gds.core.utils.progress.tasks.TaskProgressTracker;
@@ -204,14 +205,15 @@ private GraphImporter createGraphImporter(
         var idMapBuilder = idMapBuilder(config.readConcurrency());
 
         var taskVolume = queryEstimator.estimateRows(query);
-        var progressTracker = new TaskProgressTracker(
+        var internalProgressTracker = new TaskProgressTracker(
             GraphImporter.graphImporterTask(taskVolume),
             log,
             config.readConcurrency(),
             config.jobId(),
             TaskRegistryFactory.local(username, taskStore),
             EmptyUserLogRegistryFactory.INSTANCE
         );
+        var progressTracker = BatchingTaskProgressTracker.create(internalProgressTracker, taskVolume, config.readConcurrency());
 
         return new GraphImporter(
             config,

progress-tracking/src/main/java/org/neo4j/gds/core/utils/progress/BatchingProgressLogger.java

@@ -22,9 +22,9 @@
 import org.apache.commons.lang3.mutable.MutableLong;
 import org.neo4j.gds.core.concurrency.Concurrency;
 import org.neo4j.gds.core.utils.progress.tasks.Task;
+import org.neo4j.gds.logging.Log;
 import org.neo4j.gds.mem.BitUtil;
 import org.neo4j.gds.utils.CloseableThreadLocal;
-import org.neo4j.gds.logging.Log;
 
 import java.util.Objects;
 import java.util.concurrent.atomic.LongAdder;
@@ -49,7 +49,7 @@ private static long calculateBatchSize(Task task, Concurrency concurrency) {
         return calculateBatchSize(Math.max(1L, task.getProgress().volume()), concurrency);
     }
 
-    private static long calculateBatchSize(long taskVolume, Concurrency concurrency) {
+    static long calculateBatchSize(long taskVolume, Concurrency concurrency) {
         // target 100 logs per full run (every 1 percent)
         var batchSize = taskVolume / 100;
         // split batchSize into thread-local chunks

progress-tracking/src/main/java/org/neo4j/gds/core/utils/progress/BatchingTaskProgressTracker.java

@@ -0,0 +1,75 @@
+/*
+ * Copyright (c) "Neo4j"
+ * Neo4j Sweden AB [http://neo4j.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * Neo4j is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+package org.neo4j.gds.core.utils.progress;
+
+import org.neo4j.gds.core.concurrency.Concurrency;
+import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
+import org.neo4j.gds.core.utils.progress.tasks.ProgressTrackerAdapter;
+import org.neo4j.gds.core.utils.progress.tasks.Task;
+
+public final class BatchingTaskProgressTracker {
+
+    private BatchingTaskProgressTracker() {}
+
+    public static ProgressTracker create(ProgressTracker delegate, long volume, Concurrency concurrency) {
+        return volume == Task.UNKNOWN_VOLUME
+            ? new WithoutLogging(delegate)
+            : new WithLogging(delegate, volume, concurrency);
+    }
+
+    static class WithLogging extends ProgressTrackerAdapter {
+
+        private final long batchSize;
+        private long rowCounter;
+
+        WithLogging(ProgressTracker delegate, long volume, Concurrency concurrency) {
+            super(delegate);
+            this.batchSize = BatchingProgressLogger.calculateBatchSize(volume, concurrency);
+            this.rowCounter = 0;
+        }
+
+        @Override
+        public void logProgress() {
+            if (++rowCounter == batchSize) {
+                super.logProgress(batchSize);
+                rowCounter = 0;
+            }
+        }
+    }
+
+    static class WithoutLogging extends ProgressTrackerAdapter {
+
+        WithoutLogging(ProgressTracker delegate) {
+            super(delegate);
+        }
+
+        @Override
+        public void logProgress() {
+        }
+
+        @Override
+        public void logProgress(long value) {
+        }
+
+        @Override
+        public void logProgress(long value, String messageTemplate) {
+        }
+    }
+}

triplet-graph-builder/src/main/java/org/neo4j/gds/projection/BatchingTaskProgressTracker.java

@@ -63,7 +63,6 @@
 
 public final class GraphImporter {
 
-    private static final int PROGRESS_TRACKING_THRESHOLD = 10_000;
     public static final int NO_TARGET_NODE = -1;
 
     private final GraphProjectConfig config;
@@ -74,7 +73,7 @@ public final class GraphImporter {
     private final WriteMode writeMode;
     private final String query;
 
-    private final BatchingTaskProgressTracker progressTracker;
+    private final ProgressTracker progressTracker;
 
     private final Map<RelationshipType, RelationshipsBuilder> relImporters;
     private final ImmutableMutableGraphSchema.Builder graphSchemaBuilder;
@@ -95,36 +94,14 @@ public GraphImporter(
         WriteMode writeMode,
         String query,
         ProgressTracker progressTracker
-    ) {
-        this(
-            config,
-            undirectedRelationshipTypes,
-            inverseIndexedRelationshipTypes,
-            idMapBuilder,
-            writeMode,
-            query,
-            progressTracker,
-            PROGRESS_TRACKING_THRESHOLD
-        );
-    }
-
-    GraphImporter(
-        GraphProjectConfig config,
-        List<String> undirectedRelationshipTypes,
-        List<String> inverseIndexedRelationshipTypes,
-        LazyIdMapBuilder idMapBuilder,
-        WriteMode writeMode,
-        String query,
-        ProgressTracker progressTracker,
-        int progressTrackingThreshold
     ) {
         this.config = config;
         this.undirectedRelationshipTypes = undirectedRelationshipTypes;
         this.inverseIndexedRelationshipTypes = inverseIndexedRelationshipTypes;
         this.idMapBuilder = idMapBuilder;
         this.writeMode = writeMode;
         this.query = query;
-        this.progressTracker = new BatchingTaskProgressTracker(progressTracker, progressTrackingThreshold);
+        this.progressTracker = progressTracker;
         this.relImporters = new ConcurrentHashMap<>();
         this.graphSchemaBuilder = MutableGraphSchema.builder();
 
@@ -188,7 +165,6 @@ public AggregationResult result(
         ProgressTimer timer,
         boolean hasSeenArbitraryId
     ) {
-        progressTracker.logRemainingProgress();
         progressTracker.endSubTask("Update aggregation");
         progressTracker.beginSubTask("Build graph store");
         progressTracker.beginSubTask("Nodes");

triplet-graph-builder/src/main/java/org/neo4j/gds/projection/GraphImporter.java

@@ -413,8 +413,7 @@ void shouldRegisterTaskAndLogProgress() {
                 .build(),
             Capabilities.WriteMode.REMOTE,
             "",
-            progressTracker,
-            1
+            progressTracker
         );
 
         for (int i = 0; i < 2; i++) {