Fill in ExecutionContext

Author: IoannisPanagiotas

executor/src/integrationTest/java/org/neo4j/gds/executor/MemoryEstimationExecutorTest.java

@@ -74,6 +74,7 @@ void setup() throws Exception {
             .builder()
             .databaseId(DatabaseId.of(db.databaseName()))
             .dependencyResolver(GraphDatabaseApiProxy.dependencyResolver(db))
+            .memoryEstimationContext(new MemoryEstimationContext(true))
             .returnColumns(ProcedureReturnColumns.EMPTY)
             .log(Log.noOpLog())
             .taskRegistryFactory(EmptyTaskRegistryFactory.INSTANCE)

executor/src/integrationTest/java/org/neo4j/gds/executor/ProcedureExecutorTest.java

@@ -22,8 +22,6 @@
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
-import org.neo4j.common.DependencyResolver;
-import org.neo4j.configuration.Config;
 import org.neo4j.gds.api.CloseableResourceRegistry;
 import org.neo4j.gds.api.GraphStore;
 import org.neo4j.gds.api.NodeLookup;
@@ -48,7 +46,6 @@
 import java.time.Duration;
 import java.util.HashSet;
 import java.util.Map;
-import java.util.Optional;
 import java.util.Set;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -143,25 +140,7 @@ private ExecutionContext executionContext(TaskStore taskStore) {
             .nodeLookup(NodeLookup.EMPTY)
             .userLogRegistryFactory(EmptyUserLogRegistryFactory.INSTANCE)
             .metrics(Metrics.DISABLED)
-            .dependencyResolver(new DependencyResolver() {
-                @Override
-                public <T> T resolveDependency(Class<T> aClass, SelectionStrategy selectionStrategy) {
-                    if (aClass == Config.class){
-                        return (T) Config.defaults();
-                    }
-                    throw new IllegalStateException("What are you doing here, young one?");
-                }
-
-                @Override
-                public <T> Optional<T> resolveOptionalDependency(Class<T> aClass) {
-                    return Optional.empty();
-                }
-
-                @Override
-                public boolean containsDependency(Class<?> aClass) {
-                    return false;
-                }
-            })
+            .memoryEstimationContext(new MemoryEstimationContext(true))
             .build();
     }
 

executor/src/main/java/org/neo4j/gds/executor/ExecutionContext.java

@@ -45,6 +45,7 @@ public interface ExecutionContext {
 
     DatabaseId databaseId();
 
+    @Nullable
     DependencyResolver dependencyResolver();
 
     MemoryEstimationContext memoryEstimationContext();
@@ -114,8 +115,8 @@ public boolean containsDependency(Class<?> type) {
         }
     };
 
-
     MemoryEstimationContext EMPTY_MEMORY_CONTEXT = new MemoryEstimationContext(false);
+
     ExecutionContext EMPTY = new ExecutionContext() {
 
         @Override

pipeline/src/test/java/org/neo4j/gds/ml/pipeline/linkPipeline/LinkPredictionTrainingPipelineTest.java

@@ -185,6 +185,7 @@ void deriveRelationshipWeightProperty() {
         var executionContext = ImmutableExecutionContext.builder()
             .databaseId(DatabaseId.of(""))
             .dependencyResolver(ExecutionContext.EMPTY_DEPENDENCY_RESOLVER)
+            .memoryEstimationContext(ExecutionContext.EMPTY_MEMORY_CONTEXT)
             .username("")
             .terminationMonitor(TerminationMonitor.EMPTY)
             .closeableResourceRegistry(CloseableResourceRegistry.EMPTY)
@@ -229,6 +230,7 @@ void deriveRelationshipWeightPropertyFromTrainedModel() {
         var executionContext = ImmutableExecutionContext.builder()
             .databaseId(DatabaseId.of(""))
             .dependencyResolver(ExecutionContext.EMPTY_DEPENDENCY_RESOLVER)
+            .memoryEstimationContext(ExecutionContext.EMPTY_MEMORY_CONTEXT)
             .username("")
             .modelCatalog(modelCatalog)
             .terminationMonitor(TerminationMonitor.EMPTY)
@@ -273,6 +275,7 @@ void notDerivePropertyFromUnweightedTrainedModel() {
         var executionContext = ImmutableExecutionContext.builder()
             .databaseId(DatabaseId.of(""))
             .dependencyResolver(ExecutionContext.EMPTY_DEPENDENCY_RESOLVER)
+            .memoryEstimationContext(ExecutionContext.EMPTY_MEMORY_CONTEXT)
             .username("")
             .modelCatalog(modelCatalog)
             .terminationMonitor(TerminationMonitor.EMPTY)

proc/common/src/main/java/org/neo4j/gds/BaseProc.java

@@ -26,16 +26,17 @@
 import org.neo4j.gds.core.utils.warnings.UserLogRegistryFactory;
 import org.neo4j.gds.executor.ExecutionContext;
 import org.neo4j.gds.executor.ImmutableExecutionContext;
+import org.neo4j.gds.executor.MemoryEstimationContext;
 import org.neo4j.gds.logging.LogAdapter;
 import org.neo4j.gds.metrics.Metrics;
 import org.neo4j.gds.procedures.GraphDataScienceProcedures;
 import org.neo4j.gds.procedures.ProcedureCallContextReturnColumns;
-import org.neo4j.gds.transaction.TransactionCloseableResourceRegistry;
-import org.neo4j.gds.transaction.TransactionNodeLookup;
 import org.neo4j.gds.termination.TransactionTerminationMonitor;
 import org.neo4j.gds.transaction.DatabaseTransactionContext;
 import org.neo4j.gds.transaction.EmptyTransactionContext;
+import org.neo4j.gds.transaction.TransactionCloseableResourceRegistry;
 import org.neo4j.gds.transaction.TransactionContext;
+import org.neo4j.gds.transaction.TransactionNodeLookup;
 import org.neo4j.graphdb.GraphDatabaseService;
 import org.neo4j.graphdb.Transaction;
 import org.neo4j.internal.kernel.api.procs.ProcedureCallContext;
@@ -82,12 +83,15 @@ protected String username() {
     }
 
     public ExecutionContext executionContext() {
+
         return databaseService == null
             ? ExecutionContext.EMPTY
             : ImmutableExecutionContext
                 .builder()
                 .databaseId(databaseId())
                 .dependencyResolver(GraphDatabaseApiProxy.dependencyResolver(databaseService))
+                //Base Proc is only used by pregel  generated codes which have have same min-max expectation so fine
+                .memoryEstimationContext(new MemoryEstimationContext(true))
                 .log(new LogAdapter(log))
                 .returnColumns(new ProcedureCallContextReturnColumns(callContext))
                 .userLogRegistryFactory(userLogRegistryFactory)