HHH-17983 Remove entity keys from batch queue before loading

hibernate-core/src/main/java/org/hibernate/loader/ast/internal/EntityBatchLoaderArrayParam.java

@@ -134,6 +134,13 @@ public class EntityBatchLoaderArrayParam<T>
 					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 @@ public class EntityBatchLoaderArrayParam<T>
 				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

hibernate-core/src/main/java/org/hibernate/loader/ast/internal/EntityBatchLoaderInPredicate.java

@@ -141,7 +141,6 @@ public class EntityBatchLoaderInPredicate<T>
 		);
 
 		final BatchFetchQueue batchFetchQueue = session.getPersistenceContextInternal().getBatchFetchQueue();
-		final List<EntityKey> entityKeys = arrayList( sqlBatchSize );
 
 		chunker.processChunks(
 				idsToInitialize,
@@ -166,7 +165,11 @@ public class EntityBatchLoaderInPredicate<T>
 				},
 				(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 @@ public class EntityBatchLoaderInPredicate<T>
 					}
 				},
 				(startIndex, nonNullElementCount) -> {
-					entityKeys.forEach( batchFetchQueue::removeBatchLoadableEntityKey );
-					entityKeys.clear();
 				},
 				session
 		);