From 9346c28c8ea4b6b8515eff79c0355d889fb8394f Mon Sep 17 00:00:00 2001
From: Marco Belladelli <marcobladel@gmail.com>
Date: Tue, 23 Apr 2024 15:54:44 +0200
Subject: [PATCH] HHH-17983 Remove entity keys from batch queue before loading

---
 .../ast/internal/EntityBatchLoaderArrayParam.java  | 14 +++++++-------
 .../ast/internal/EntityBatchLoaderInPredicate.java |  9 +++++----
 2 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/hibernate-core/src/main/java/org/hibernate/loader/ast/internal/EntityBatchLoaderArrayParam.java b/hibernate-core/src/main/java/org/hibernate/loader/ast/internal/EntityBatchLoaderArrayParam.java
index 0c1e561acdc1..a36590b4a793 100644
--- a/hibernate-core/src/main/java/org/hibernate/loader/ast/internal/EntityBatchLoaderArrayParam.java
+++ b/hibernate-core/src/main/java/org/hibernate/loader/ast/internal/EntityBatchLoaderArrayParam.java
@@ -134,6 +134,13 @@ protected void initializeEntities(
 					getLoadable().getEntityName(), id, Arrays.toString(idsToInitialize) );
 		}
 
+		for ( Object initializedId : idsToInitialize ) {
+			if ( initializedId != null ) {
+				// found or not, remove the key from the batch-fetch queue
+				removeBatchLoadableEntityKey( initializedId, getLoadable(), session );
+			}
+		}
+
 		LoaderHelper.loadByArrayParameter(
 				idsToInitialize,
 				sqlAst,
@@ -147,13 +154,6 @@ protected void initializeEntities(
 				readOnly,
 				session
 		);
-
-		for ( Object initializedId : idsToInitialize ) {
-			if ( initializedId != null ) {
-				// found or not, remove the key from the batch-fetch queue
-				removeBatchLoadableEntityKey( initializedId, getLoadable(), session );
-			}
-		}
 	}
 
 	@Override
diff --git a/hibernate-core/src/main/java/org/hibernate/loader/ast/internal/EntityBatchLoaderInPredicate.java b/hibernate-core/src/main/java/org/hibernate/loader/ast/internal/EntityBatchLoaderInPredicate.java
index 6f2ea798cbd1..04d27c5f5479 100644
--- a/hibernate-core/src/main/java/org/hibernate/loader/ast/internal/EntityBatchLoaderInPredicate.java
+++ b/hibernate-core/src/main/java/org/hibernate/loader/ast/internal/EntityBatchLoaderInPredicate.java
@@ -141,7 +141,6 @@ protected void initializeEntities(
 		);
 
 		final BatchFetchQueue batchFetchQueue = session.getPersistenceContextInternal().getBatchFetchQueue();
-		final List<EntityKey> entityKeys = arrayList( sqlBatchSize );
 
 		chunker.processChunks(
 				idsToInitialize,
@@ -166,7 +165,11 @@ protected void initializeEntities(
 				},
 				(key, relativePosition, absolutePosition) -> {
 					if ( key != null ) {
-						entityKeys.add( session.generateEntityKey( key, getLoadable().getEntityPersister() ) );
+						final EntityKey entityKey = session.generateEntityKey(
+								key,
+								getLoadable().getEntityPersister()
+						);
+						batchFetchQueue.removeBatchLoadableEntityKey( entityKey );
 					}
 				},
 				(startIndex) -> {
@@ -181,8 +184,6 @@ protected void initializeEntities(
 					}
 				},
 				(startIndex, nonNullElementCount) -> {
-					entityKeys.forEach( batchFetchQueue::removeBatchLoadableEntityKey );
-					entityKeys.clear();
 				},
 				session
 		);