diff --git a/src/main/java/org/springframework/data/neo4j/core/Neo4jTemplate.java b/src/main/java/org/springframework/data/neo4j/core/Neo4jTemplate.java index 9549817cf..5d3505bb0 100644 --- a/src/main/java/org/springframework/data/neo4j/core/Neo4jTemplate.java +++ b/src/main/java/org/springframework/data/neo4j/core/Neo4jTemplate.java @@ -1202,8 +1202,8 @@ public final class Neo4jTemplate implements return NodesAndRelationshipsByIdStatementProvider.EMPTY; } // load first level relationships - final Set relationshipIds = new HashSet<>(); - final Set relatedNodeIds = new HashSet<>(); +// final Set relationshipIds = new HashSet<>(); + final Map> relationshipsToRelatedNodeIds = new HashMap<>(); for (RelationshipDescription relationshipDescription : entityMetaData.getRelationshipsInHierarchy(queryFragments::includeField)) { @@ -1217,14 +1217,14 @@ public final class Neo4jTemplate implements .bindAll(usedParameters) .fetch() .one() - .ifPresent(iterateAndMapNextLevel(relationshipIds, relatedNodeIds, relationshipDescription, PropertyPathWalkStep.empty())); + .ifPresent(iterateAndMapNextLevel(relationshipsToRelatedNodeIds, relationshipDescription, PropertyPathWalkStep.empty())); } - return new NodesAndRelationshipsByIdStatementProvider(rootNodeIds, relationshipIds, relatedNodeIds, queryFragments, elementIdOrIdFunction); + return new NodesAndRelationshipsByIdStatementProvider(rootNodeIds, relationshipsToRelatedNodeIds.keySet(), relationshipsToRelatedNodeIds.values().stream().flatMap(Collection::stream).toList(), queryFragments, elementIdOrIdFunction); } - private void iterateNextLevel(Collection nodeIds, RelationshipDescription sourceRelationshipDescription, Set relationshipIds, - Set relatedNodeIds, PropertyPathWalkStep currentPathStep) { + private void iterateNextLevel(Collection nodeIds, RelationshipDescription sourceRelationshipDescription, + Map> relationshipsToRelatedNodes, PropertyPathWalkStep currentPathStep) { Neo4jPersistentEntity target = (Neo4jPersistentEntity) sourceRelationshipDescription.getTarget(); @@ -1258,32 +1258,42 @@ public final class Neo4jTemplate implements .bindAll(Collections.singletonMap(Constants.NAME_OF_IDS, TemplateSupport.convertToLongIdOrStringElementId(nodeIds))) .fetch() .one() - .ifPresent(iterateAndMapNextLevel(relationshipIds, relatedNodeIds, relationshipDescription, nextPathStep)); + .ifPresent(iterateAndMapNextLevel(relationshipsToRelatedNodes, relationshipDescription, nextPathStep)); } } @NonNull - private Consumer> iterateAndMapNextLevel(Set relationshipIds, - Set relatedNodeIds, + private Consumer> iterateAndMapNextLevel(Map> relationshipsToRelatedNodes, RelationshipDescription relationshipDescription, PropertyPathWalkStep currentPathStep) { return record -> { + + Map> relatedNodesVisited = new HashMap<>(relationshipsToRelatedNodes); @SuppressWarnings("unchecked") List newRelationshipIds = ((List) record.get(Constants.NAME_OF_SYNTHESIZED_RELATIONS)).stream().map(TemplateSupport::convertIdOrElementIdToString).toList(); - relationshipIds.addAll(newRelationshipIds); - @SuppressWarnings("unchecked") - List newRelatedNodeIds = ((List) record.get(Constants.NAME_OF_SYNTHESIZED_RELATED_NODES)).stream().map(TemplateSupport::convertIdOrElementIdToString).toList(); + Set relatedIds = new HashSet<>(((List) record.get(Constants.NAME_OF_SYNTHESIZED_RELATED_NODES)).stream().map(TemplateSupport::convertIdOrElementIdToString).toList()); - Set relatedIds = new HashSet<>(newRelatedNodeIds); // use this list to get down the road // 1. remove already visited ones; - relatedIds.removeAll(relatedNodeIds); - relatedNodeIds.addAll(relatedIds); + // we don't know which id came with which node, so we need to assume that a relationshipId connects to all related nodes + for (String newRelationshipId : newRelationshipIds) { + relatedNodesVisited.put(newRelationshipId, relatedIds); + Set knownRelatedNodesBefore = relationshipsToRelatedNodes.get(newRelationshipId); + if (knownRelatedNodesBefore != null) { + Set mergedKnownRelatedNodes = new HashSet<>(knownRelatedNodesBefore); + // there are already existing nodes in there for this relationship + mergedKnownRelatedNodes.addAll(relatedIds); + relatedNodesVisited.put(newRelationshipId, mergedKnownRelatedNodes); + relatedIds.removeAll(knownRelatedNodesBefore); + } + } + + relationshipsToRelatedNodes.putAll(relatedNodesVisited); // 2. for the rest start the exploration if (!relatedIds.isEmpty()) { - iterateNextLevel(relatedIds, relationshipDescription, relationshipIds, relatedNodeIds, currentPathStep); + iterateNextLevel(relatedIds, relationshipDescription, relationshipsToRelatedNodes, currentPathStep); } }; } diff --git a/src/main/java/org/springframework/data/neo4j/core/ReactiveNeo4jTemplate.java b/src/main/java/org/springframework/data/neo4j/core/ReactiveNeo4jTemplate.java index b0f3a45d6..195d1ad48 100644 --- a/src/main/java/org/springframework/data/neo4j/core/ReactiveNeo4jTemplate.java +++ b/src/main/java/org/springframework/data/neo4j/core/ReactiveNeo4jTemplate.java @@ -720,8 +720,7 @@ public final class ReactiveNeo4jTemplate implements Class rootClass = entityMetaData.getUnderlyingClass(); Set rootNodeIds = ctx.get("rootNodes"); - Set processedRelationshipIds = ctx.get("processedRelationships"); - Set processedNodeIds = ctx.get("processedNodes"); + Map> relationshipsToRelatedNodeIds = ctx.get("relationshipsToRelatedNodeIds"); return Flux.fromIterable(entityMetaData.getRelationshipsInHierarchy(queryFragments::includeField)) .concatMap(relationshipDescription -> { @@ -748,12 +747,11 @@ public final class ReactiveNeo4jTemplate implements }) .expand(iterateAndMapNextLevel(relationshipDescription, queryFragments, rootClass, PropertyPathWalkStep.empty())); }) - .then(Mono.fromSupplier(() -> new NodesAndRelationshipsByIdStatementProvider(rootNodeIds, processedRelationshipIds, processedNodeIds, queryFragments, elementIdOrIdFunction))); + .then(Mono.fromSupplier(() -> new NodesAndRelationshipsByIdStatementProvider(rootNodeIds, relationshipsToRelatedNodeIds.keySet(), relationshipsToRelatedNodeIds.values().stream().flatMap(Collection::stream).toList(), queryFragments, elementIdOrIdFunction))); }) .contextWrite(ctx -> ctx .put("rootNodes", ConcurrentHashMap.newKeySet()) - .put("processedNodes", ConcurrentHashMap.newKeySet()) - .put("processedRelationships", ConcurrentHashMap.newKeySet())); + .put("relationshipsToRelatedNodeIds", new ConcurrentHashMap<>())); } @@ -812,22 +810,29 @@ public final class ReactiveNeo4jTemplate implements return newRelationshipAndRelatedNodeIds -> Flux.deferContextual(ctx -> { - Set relationshipIds = ctx.get("processedRelationships"); - Set processedNodeIds = ctx.get("processedNodes"); + Map> relationshipsToRelatedNodeIds = ctx.get("relationshipsToRelatedNodeIds"); + Map> relatedNodesVisited = new HashMap<>(relationshipsToRelatedNodeIds); Collection newRelationshipIds = newRelationshipAndRelatedNodeIds.getT1(); - Set tmpProcessedRels = ConcurrentHashMap.newKeySet(newRelationshipIds.size()); - tmpProcessedRels.addAll(newRelationshipIds); - tmpProcessedRels.removeAll(relationshipIds); - relationshipIds.addAll(newRelationshipIds); Collection newRelatedNodeIds = newRelationshipAndRelatedNodeIds.getT2(); - Set tmpProcessedNodes = ConcurrentHashMap.newKeySet(newRelatedNodeIds.size()); - tmpProcessedNodes.addAll(newRelatedNodeIds); - tmpProcessedNodes.removeAll(processedNodeIds); - processedNodeIds.addAll(newRelatedNodeIds); + Set relatedIds = ConcurrentHashMap.newKeySet(newRelatedNodeIds.size()); + relatedIds.addAll(newRelatedNodeIds); - if (tmpProcessedRels.isEmpty() && tmpProcessedNodes.isEmpty()) { + for (String newRelationshipId : newRelationshipIds) { + relatedNodesVisited.put(newRelationshipId, relatedIds); + Set knownRelatedNodesBefore = relationshipsToRelatedNodeIds.get(newRelationshipId); + if (knownRelatedNodesBefore != null) { + Set mergedKnownRelatedNodes = new HashSet<>(knownRelatedNodesBefore); + // there are already existing nodes in there for this relationship + mergedKnownRelatedNodes.addAll(relatedIds); + relatedNodesVisited.put(newRelationshipId, mergedKnownRelatedNodes); + relatedIds.removeAll(knownRelatedNodesBefore); + } + } + relationshipsToRelatedNodeIds.putAll(relatedNodesVisited); + + if (relatedIds.isEmpty()) { return Mono.empty(); }