diff --git a/hibernate-reactive-core/src/main/java/org/hibernate/reactive/loader/ast/internal/ReactiveEntityBatchLoaderArrayParam.java b/hibernate-reactive-core/src/main/java/org/hibernate/reactive/loader/ast/internal/ReactiveEntityBatchLoaderArrayParam.java
index ba69fd407..04f971ffd 100644
--- a/hibernate-reactive-core/src/main/java/org/hibernate/reactive/loader/ast/internal/ReactiveEntityBatchLoaderArrayParam.java
+++ b/hibernate-reactive-core/src/main/java/org/hibernate/reactive/loader/ast/internal/ReactiveEntityBatchLoaderArrayParam.java
@@ -10,7 +10,6 @@
 import java.util.concurrent.CompletionStage;
 
 import org.hibernate.LockOptions;
-import org.hibernate.engine.internal.BatchFetchQueueHelper;
 import org.hibernate.engine.spi.EntityKey;
 import org.hibernate.engine.spi.LoadQueryInfluencers;
 import org.hibernate.engine.spi.SessionFactoryImplementor;
@@ -160,7 +159,8 @@ private CompletionStage<Void> initializeEntities(
 							continue;
 						}
 						// found or not, remove the key from the batch-fetch queue
-						BatchFetchQueueHelper.removeBatchLoadableEntityKey( id, getLoadable(), session );
+						session.getPersistenceContextInternal().getBatchFetchQueue()
+								.removeBatchLoadableEntityKey( session.generateEntityKey( id, getLoadable().getEntityPersister() ) );
 					}
 				} );
 	}
diff --git a/hibernate-reactive-core/src/main/java/org/hibernate/reactive/loader/ast/internal/ReactiveMultiIdEntityLoaderArrayParam.java b/hibernate-reactive-core/src/main/java/org/hibernate/reactive/loader/ast/internal/ReactiveMultiIdEntityLoaderArrayParam.java
index 8495d115b..8239a6e15 100644
--- a/hibernate-reactive-core/src/main/java/org/hibernate/reactive/loader/ast/internal/ReactiveMultiIdEntityLoaderArrayParam.java
+++ b/hibernate-reactive-core/src/main/java/org/hibernate/reactive/loader/ast/internal/ReactiveMultiIdEntityLoaderArrayParam.java
@@ -13,7 +13,6 @@
 
 import org.hibernate.LockMode;
 import org.hibernate.LockOptions;
-import org.hibernate.engine.internal.BatchFetchQueueHelper;
 import org.hibernate.engine.spi.BatchFetchQueue;
 import org.hibernate.engine.spi.EntityEntry;
 import org.hibernate.engine.spi.EntityKey;
@@ -206,7 +205,7 @@ protected <K> CompletionStage<List<E>> performOrderedMultiLoad(
 						// the element value at this position in the result List should be
 						// the EntityKey for that entity - reuse it
 						final EntityKey entityKey = (EntityKey) result.get( resultIndex );
-						BatchFetchQueueHelper.removeBatchLoadableEntityKey( entityKey, session );
+						session.getPersistenceContextInternal().getBatchFetchQueue().removeBatchLoadableEntityKey( entityKey );
 						Object entity = persistenceContext.getEntity( entityKey );
 						if ( entity != null && !loadOptions.isReturnOfDeletedEntitiesEnabled() ) {
 							// make sure it is not DELETED
@@ -293,7 +292,8 @@ protected <K> CompletionStage<List<E>> performUnorderedMultiLoad(
 					continue;
 				}
 				// found or not, remove the key from the batch-fetch queue
-				BatchFetchQueueHelper.removeBatchLoadableEntityKey( id, getLoadable(), session );
+				session.getPersistenceContextInternal().getBatchFetchQueue()
+						.removeBatchLoadableEntityKey( session.generateEntityKey( id, getLoadable().getEntityPersister() ) );
 			}
 
 			return result;