HHH-15921 @BatchSize and @IdClass on join column throws exception

hibernate-core/src/main/java/org/hibernate/dialect/temptable/TemporaryTable.java

@@ -31,7 +31,6 @@ import org.hibernate.mapping.Selectable;
 import org.hibernate.mapping.SimpleValue;
 import org.hibernate.mapping.Value;
 import org.hibernate.metamodel.mapping.EntityDiscriminatorMapping;
-import org.hibernate.metamodel.mapping.EntityIdentifierMapping;
 import org.hibernate.metamodel.mapping.EntityMappingType;
 import org.hibernate.metamodel.mapping.ForeignKeyDescriptor;
 import org.hibernate.metamodel.mapping.JdbcMapping;
@@ -216,7 +215,7 @@ public class TemporaryTable implements Exportable, Contributable {
 											throw new IllegalStateException( "Not yet ready: " + pluralAttribute );
 										}
 										final ModelPart fkTarget = keyDescriptor.getTargetPart();
-										if ( !( fkTarget instanceof EntityIdentifierMapping ) ) {
+										if ( !fkTarget.isEntityIdentifierMapping() ) {
 											final Value value = entityBinding.getSubclassProperty( pluralAttribute.getAttributeName() )
 													.getValue();
 											final Iterator<Selectable> columnIterator =

hibernate-core/src/main/java/org/hibernate/metamodel/internal/EmbeddableInstantiatorPojoIndirecting.java

@@ -10,12 +10,13 @@ import java.lang.reflect.Constructor;
 
 import org.hibernate.InstantiationException;
 import org.hibernate.engine.spi.SessionFactoryImplementor;
+import org.hibernate.metamodel.spi.EmbeddableInstantiator;
 import org.hibernate.metamodel.spi.ValueAccess;
 
 /**
  * Support for instantiating embeddables as POJO representation through a constructor
  */
-public class EmbeddableInstantiatorPojoIndirecting extends AbstractPojoInstantiator implements StandardEmbeddableInstantiator {
+public class EmbeddableInstantiatorPojoIndirecting extends AbstractPojoInstantiator implements EmbeddableInstantiator {
 	protected final Constructor<?> constructor;
 	protected final int[] index;
 

hibernate-core/src/main/java/org/hibernate/metamodel/mapping/EntityIdentifierMapping.java

@@ -141,4 +141,9 @@ public interface EntityIdentifierMapping extends ValuedModelPart {
 		 */
 		VIRTUAL
 	}
+
+	@Override
+	default boolean isEntityIdentifierMapping() {
+		return true;
+	}
 }

hibernate-core/src/main/java/org/hibernate/metamodel/mapping/ModelPart.java

@@ -91,6 +91,10 @@ public interface ModelPart extends MappingModelExpressible {
 		return false;
 	}
 
+	default boolean isEntityIdentifierMapping() {
+		return false;
+	}
+
 	boolean hasPartitionedSelectionMapping();
 
 	/**

hibernate-core/src/main/java/org/hibernate/metamodel/mapping/internal/IdClassEmbeddable.java

@@ -198,7 +198,7 @@ public class IdClassEmbeddable extends AbstractEmbeddableMapping implements Iden
 				final ModelPart targetPart = toOneAttributeMapping.getForeignKeyDescriptor().getPart(
 						toOneAttributeMapping.getSideNature().inverse()
 				);
-				if ( targetPart instanceof EntityIdentifierMapping ) {
+				if ( targetPart.isEntityIdentifierMapping() ) {
 					propertyValues[i] = ( (EntityIdentifierMapping) targetPart ).getIdentifier( o );
 				}
 				else {

hibernate-core/src/main/java/org/hibernate/metamodel/mapping/internal/InverseNonAggregatedIdentifierMapping.java

@@ -206,7 +206,7 @@ public class InverseNonAggregatedIdentifierMapping extends EmbeddedAttributeMapp
 					final ModelPart targetPart = toOneAttributeMapping.getForeignKeyDescriptor().getPart(
 							toOneAttributeMapping.getSideNature().inverse()
 					);
-					if ( targetPart instanceof EntityIdentifierMapping ) {
+					if ( targetPart.isEntityIdentifierMapping() ) {
 						propertyValues[i] = ( (EntityIdentifierMapping) targetPart ).getIdentifier( o );
 					}
 					else {

hibernate-core/src/main/java/org/hibernate/metamodel/mapping/internal/ManyToManyCollectionPart.java

@@ -193,7 +193,7 @@ public class ManyToManyCollectionPart extends AbstractEntityCollectionPart imple
 
 	@Override
 	public boolean isReferenceToPrimaryKey() {
-		return getForeignKeyDescriptor().getTargetPart() instanceof EntityIdentifierMapping;
+		return getForeignKeyDescriptor().getTargetPart().isEntityIdentifierMapping();
 	}
 
 	@Override

hibernate-core/src/main/java/org/hibernate/metamodel/mapping/internal/NonAggregatedIdentifierMappingImpl.java

@@ -230,7 +230,7 @@ public class NonAggregatedIdentifierMappingImpl extends AbstractCompositeIdentif
 					final ModelPart targetPart = toOneAttributeMapping.getForeignKeyDescriptor().getPart(
 							toOneAttributeMapping.getSideNature().inverse()
 					);
-					if ( targetPart instanceof EntityIdentifierMapping ) {
+					if ( targetPart.isEntityIdentifierMapping() ) {
 						propertyValues[i] = ( (EntityIdentifierMapping) targetPart ).getIdentifier( o );
 					}
 					else {

hibernate-core/src/main/java/org/hibernate/metamodel/mapping/internal/SimpleForeignKeyDescriptor.java

@@ -442,7 +442,7 @@ public class SimpleForeignKeyDescriptor implements ForeignKeyDescriptor, BasicVa
 			}
 		}
 		final ModelPart modelPart = side.getModelPart();
-		if ( modelPart instanceof EntityIdentifierMapping ) {
+		if ( modelPart.isEntityIdentifierMapping() ) {
 			return ( (EntityIdentifierMapping) modelPart ).getIdentifierIfNotUnsaved( targetObject, session );
 		}
 		return ( (PropertyBasedMapping) modelPart ).getPropertyAccess().getGetter().get( targetObject );

hibernate-core/src/main/java/org/hibernate/metamodel/mapping/internal/ToOneAttributeMapping.java

@@ -770,7 +770,7 @@ public class ToOneAttributeMapping
 
 	@Override
 	public boolean isReferenceToPrimaryKey() {
-		return foreignKeyDescriptor.getSide( sideNature.inverse() ).getModelPart() instanceof EntityIdentifierMapping;
+		return foreignKeyDescriptor.getSide( sideNature.inverse() ).getModelPart().isEntityIdentifierMapping();
 	}
 
 	@Override

hibernate-core/src/main/java/org/hibernate/persister/entity/AbstractEntityPersister.java

@@ -5564,7 +5564,7 @@ public abstract class AbstractEntityPersister
 			final ModelPart superDefinedAttribute = superMappingType.findSubPart( name, superMappingType );
 			if ( superDefinedAttribute != null ) {
 				// Prefer the identifier mapping of the concrete class
-				if ( superDefinedAttribute instanceof EntityIdentifierMapping ) {
+				if ( superDefinedAttribute.isEntityIdentifierMapping() ) {
 					final ModelPart identifierModelPart = getIdentifierModelPart( name, treatTargetType );
 					if ( identifierModelPart != null ) {
 						return identifierModelPart;

hibernate-core/src/main/java/org/hibernate/query/results/ResultsHelper.java

@@ -113,7 +113,7 @@ public class ResultsHelper {
 	}
 
 	public static String attributeName(ModelPart identifierMapping) {
-		if ( identifierMapping instanceof EntityIdentifierMapping ) {
+		if ( identifierMapping.isEntityIdentifierMapping() ) {
 			return identifierMapping instanceof SingleAttributeIdentifierMapping
 					? ( (SingleAttributeIdentifierMapping) identifierMapping ).getAttributeName()
 					: null;

hibernate-core/src/main/java/org/hibernate/query/sqm/mutation/internal/MatchingIdSelectionHelper.java

@@ -251,8 +251,8 @@ public class MatchingIdSelectionHelper {
 
 							if ( pluralAttribute.getSeparateCollectionTable() != null ) {
 								// Ensure that the FK target columns are available
-								final boolean useFkTarget = !( pluralAttribute.getKeyDescriptor()
+								final boolean useFkTarget = !pluralAttribute.getKeyDescriptor()
-										.getTargetPart() instanceof EntityIdentifierMapping );
+										.getTargetPart().isEntityIdentifierMapping();
 								if ( useFkTarget ) {
 									final TableGroup mutatingTableGroup = sqmConverter.getMutatingTableGroup();
 									pluralAttribute.getKeyDescriptor().getTargetPart().applySqlSelections(

hibernate-core/src/main/java/org/hibernate/query/sqm/mutation/internal/cte/CteDeleteHandler.java

@@ -77,8 +77,8 @@ public class CteDeleteHandler extends AbstractCteMutationHandler implements Dele
 				pluralAttribute -> {
 					if ( pluralAttribute.getSeparateCollectionTable() != null ) {
 						// Ensure that the FK target columns are available
-						final boolean useFkTarget = !( pluralAttribute.getKeyDescriptor()
+						final boolean useFkTarget = !pluralAttribute.getKeyDescriptor()
-								.getTargetPart() instanceof EntityIdentifierMapping );
+								.getTargetPart().isEntityIdentifierMapping();
 						if ( useFkTarget ) {
 							final TableGroup mutatingTableGroup = sqmConverter.getMutatingTableGroup();
 							pluralAttribute.getKeyDescriptor().getTargetPart().applySqlSelections(

hibernate-core/src/main/java/org/hibernate/query/sqm/mutation/internal/inline/InlineDeleteHandler.java

@@ -103,7 +103,7 @@ public class InlineDeleteHandler implements DeleteHandler {
 						// collection table
 						final ModelPart fkTargetPart = pluralAttribute.getKeyDescriptor().getTargetPart();
 						final int valueIndex;
-						if ( fkTargetPart instanceof EntityIdentifierMapping ) {
+						if ( fkTargetPart.isEntityIdentifierMapping() ) {
 							valueIndex = 0;
 						}
 						else {

hibernate-core/src/main/java/org/hibernate/query/sqm/mutation/internal/temptable/RestrictedDeleteExecutionDelegate.java

@@ -258,7 +258,7 @@ public class RestrictedDeleteExecutionDelegate implements TableBasedDeleteHandle
 					final ForeignKeyDescriptor fkDescriptor = attributeMapping.getKeyDescriptor();
 					final QuerySpec idSelectFkSubQuery;
 					// todo (6.0): based on the location of the attribute mapping, we could prune the table group of the subquery
-					if ( fkDescriptor.getTargetPart() instanceof EntityIdentifierMapping ) {
+					if ( fkDescriptor.getTargetPart().isEntityIdentifierMapping() ) {
 						idSelectFkSubQuery = matchingIdSubQuerySpec;
 					}
 					else {
@@ -538,7 +538,7 @@ public class RestrictedDeleteExecutionDelegate implements TableBasedDeleteHandle
 				(tableReference, attributeMapping) -> {
 					final ForeignKeyDescriptor fkDescriptor = attributeMapping.getKeyDescriptor();
 					final QuerySpec idTableFkSubQuery;
-					if ( fkDescriptor.getTargetPart() instanceof EntityIdentifierMapping ) {
+					if ( fkDescriptor.getTargetPart().isEntityIdentifierMapping() ) {
 						idTableFkSubQuery = idTableIdentifierSubQuery;
 					}
 					else {

hibernate-core/src/main/java/org/hibernate/sql/results/graph/Initializer.java

@@ -9,6 +9,7 @@ package org.hibernate.sql.results.graph;
 import org.hibernate.Incubating;
 import org.hibernate.metamodel.mapping.AttributeMapping;
 import org.hibernate.persister.entity.EntityPersister;
+import org.hibernate.sql.results.graph.embeddable.EmbeddableInitializer;
 import org.hibernate.sql.results.graph.entity.EntityInitializer;
 import org.hibernate.sql.results.jdbc.spi.RowProcessingState;
 import org.hibernate.metamodel.mapping.ModelPart;
@@ -104,8 +105,22 @@ public interface Initializer {
 		return false;
 	}
 
+	/**
+	 * A utility method to avoid casting explicitly to EntityInitializer
+	 *
+	 * @return EntityInitializer if this is an instance of EntityInitializer otherwise {@code null}
+	 */
 	default EntityInitializer asEntityInitializer() {
 		return null;
 	}
 
+	/**
+	 * A utility method to avoid casting explicitly to EmbeddableInitializer
+	 *
+	 * @return EmbeddableInitializer if this is an instance of EmbeddableInitializer otherwise {@code null}
+	 */
+	default EmbeddableInitializer asEmbeddableInitializer() {
+		return null;
+	}
+
 }

hibernate-core/src/main/java/org/hibernate/sql/results/graph/embeddable/AbstractEmbeddableInitializer.java

@@ -219,6 +219,8 @@ public abstract class AbstractEmbeddableInitializer extends AbstractFetchParentA
 			notifyResolutionListeners( compositeInstance );
 
 			final LazyInitializer lazyInitializer = HibernateProxy.extractLazyInitializer( compositeInstance );
+			// If the composite instance has a lazy initializer attached, this means that the embeddable is actually virtual
+			// and the compositeInstance == entity, so we have to inject the row state into the entity when it finishes resolution
 			if ( lazyInitializer != null ) {
 				final Initializer parentInitializer = processingState.resolveInitializer( navigablePath.getParent() );
 				if ( parentInitializer != this ) {

hibernate-core/src/main/java/org/hibernate/sql/results/graph/embeddable/EmbeddableInitializer.java

@@ -27,8 +27,8 @@ public interface EmbeddableInitializer extends FetchParentAccess {
 	default RowProcessingState wrapProcessingState(RowProcessingState processingState) {
 		final FetchParentAccess fetchParentAccess = getFetchParentAccess();
 		if ( fetchParentAccess != null ) {
-			if ( fetchParentAccess instanceof EmbeddableInitializer ) {
+			if ( fetchParentAccess.isEmbeddableInitializer() ) {
-				return ( (EmbeddableInitializer) fetchParentAccess ).wrapProcessingState( processingState );
+				return ( fetchParentAccess.asEmbeddableInitializer() ).wrapProcessingState( processingState );
 			}
 		}
 		return processingState;
@@ -43,4 +43,9 @@ public interface EmbeddableInitializer extends FetchParentAccess {
 	default boolean isEmbeddableInitializer() {
 		return true;
 	}
+
+	@Override
+	default EmbeddableInitializer asEmbeddableInitializer() {
+		return this;
+	}
 }

hibernate-core/src/main/java/org/hibernate/sql/results/graph/embeddable/internal/AggregateEmbeddableFetchImpl.java

@@ -152,7 +152,7 @@ public class AggregateEmbeddableFetchImpl extends AbstractFetchParent implements
 	public DomainResultAssembler createAssembler(
 			FetchParentAccess parentAccess,
 			AssemblerCreationState creationState) {
-		final EmbeddableInitializer initializer = (EmbeddableInitializer) creationState.resolveInitializer(
+		final EmbeddableInitializer initializer = creationState.resolveInitializer(
 				getNavigablePath(),
 				getReferencedModePart(),
 				() -> new AggregateEmbeddableFetchInitializer(
@@ -161,7 +161,9 @@ public class AggregateEmbeddableFetchImpl extends AbstractFetchParent implements
 						creationState,
 						aggregateSelection
 				)
-		);
+		).asEmbeddableInitializer();
+
+		assert initializer != null;
 
 		return new EmbeddableAssembler( initializer );
 	}

hibernate-core/src/main/java/org/hibernate/sql/results/graph/embeddable/internal/AggregateEmbeddableResultImpl.java

@@ -151,7 +151,7 @@ public class AggregateEmbeddableResultImpl<T> extends AbstractFetchParent implem
 	public DomainResultAssembler<T> createResultAssembler(
 			FetchParentAccess parentAccess,
 			AssemblerCreationState creationState) {
-		final EmbeddableInitializer initializer = (EmbeddableInitializer) creationState.resolveInitializer(
+		final EmbeddableInitializer initializer = creationState.resolveInitializer(
 				initializerNavigablePath,
 				getReferencedModePart(),
 				() -> new AggregateEmbeddableResultInitializer(
@@ -160,7 +160,9 @@ public class AggregateEmbeddableResultImpl<T> extends AbstractFetchParent implem
 						creationState,
 						aggregateSelection
 				)
-		);
+		).asEmbeddableInitializer();
+
+		assert initializer != null;
 
 		return new EmbeddableAssembler( initializer );
 	}

hibernate-core/src/main/java/org/hibernate/sql/results/graph/embeddable/internal/EmbeddableExpressionResultImpl.java

@@ -122,7 +122,7 @@ public class EmbeddableExpressionResultImpl<T> extends AbstractFetchParent imple
 	public DomainResultAssembler<T> createResultAssembler(
 			FetchParentAccess parentAccess,
 			AssemblerCreationState creationState) {
-		final EmbeddableInitializer initializer = (EmbeddableInitializer) creationState.resolveInitializer(
+		final EmbeddableInitializer initializer = creationState.resolveInitializer(
 				getNavigablePath(),
 				getReferencedModePart(),
 				() -> new EmbeddableResultInitializer(
@@ -130,7 +130,9 @@ public class EmbeddableExpressionResultImpl<T> extends AbstractFetchParent imple
 						parentAccess,
 						creationState
 				)
-		);
+		).asEmbeddableInitializer();
+
+		assert initializer != null;
 
 		//noinspection unchecked
 		return new EmbeddableAssembler( initializer );

hibernate-core/src/main/java/org/hibernate/sql/results/graph/embeddable/internal/EmbeddableFetchImpl.java

@@ -129,7 +129,7 @@ public class EmbeddableFetchImpl extends AbstractFetchParent implements Embeddab
 	public DomainResultAssembler createAssembler(
 			FetchParentAccess parentAccess,
 			AssemblerCreationState creationState) {
-		final EmbeddableInitializer initializer = (EmbeddableInitializer) creationState.resolveInitializer(
+		final EmbeddableInitializer initializer = creationState.resolveInitializer(
 				getNavigablePath(),
 				getReferencedModePart(),
 				() -> new EmbeddableFetchInitializer(
@@ -137,7 +137,9 @@ public class EmbeddableFetchImpl extends AbstractFetchParent implements Embeddab
 						this,
 						creationState
 				)
-		);
+		).asEmbeddableInitializer();
+
+		assert initializer != null;
 
 		return new EmbeddableAssembler( initializer );
 	}

hibernate-core/src/main/java/org/hibernate/sql/results/graph/embeddable/internal/EmbeddableForeignKeyResultImpl.java

@@ -95,11 +95,13 @@ public class EmbeddableForeignKeyResultImpl<T>
 	public DomainResultAssembler<T> createResultAssembler(
 			FetchParentAccess parentAccess,
 			AssemblerCreationState creationState) {
-		final EmbeddableInitializer initializer = (EmbeddableInitializer) creationState.resolveInitializer(
+		final EmbeddableInitializer initializer = creationState.resolveInitializer(
 				getNavigablePath(),
 				getReferencedModePart(),
 				() -> new EmbeddableResultInitializer( this, parentAccess, creationState )
-		);
+		).asEmbeddableInitializer();
+
+		assert initializer != null;
 
 		//noinspection unchecked
 		return new EmbeddableAssembler( initializer );

hibernate-core/src/main/java/org/hibernate/sql/results/graph/embeddable/internal/EmbeddableResultImpl.java

@@ -121,7 +121,7 @@ public class EmbeddableResultImpl<T> extends AbstractFetchParent implements Embe
 	public DomainResultAssembler<T> createResultAssembler(
 			FetchParentAccess parentAccess,
 			AssemblerCreationState creationState) {
-		final EmbeddableInitializer initializer = (EmbeddableInitializer) creationState.resolveInitializer(
+		final EmbeddableInitializer initializer = creationState.resolveInitializer(
 				initializerNavigablePath,
 				getReferencedModePart(),
 				() -> new EmbeddableResultInitializer(
@@ -129,7 +129,9 @@ public class EmbeddableResultImpl<T> extends AbstractFetchParent implements Embe
 						parentAccess,
 						creationState
 				)
-		);
+		).asEmbeddableInitializer();
+
+		assert initializer != null;
 
 		//noinspection unchecked
 		return new EmbeddableAssembler( initializer );

hibernate-core/src/main/java/org/hibernate/sql/results/graph/entity/internal/AbstractNonJoinedEntityFetch.java

@@ -6,9 +6,6 @@
  */
 package org.hibernate.sql.results.graph.entity.internal;
 
-import java.util.Collections;
-import java.util.List;
-
 import org.hibernate.metamodel.mapping.EntityMappingType;
 import org.hibernate.spi.NavigablePath;
 import org.hibernate.sql.results.graph.Fetch;

hibernate-core/src/main/java/org/hibernate/sql/results/graph/entity/internal/BatchInitializeEntitySelectFetchInitializer.java

@@ -0,0 +1,121 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * License: GNU Lesser General Public License (LGPL), version 2.1 or later
+ * See the lgpl.txt file in the root directory or http://www.gnu.org/licenses/lgpl-2.1.html
+ */
+package org.hibernate.sql.results.graph.entity.internal;
+
+import java.util.HashSet;
+import java.util.Set;
+
+import org.hibernate.engine.spi.EntityKey;
+import org.hibernate.engine.spi.SharedSessionContractImplementor;
+import org.hibernate.internal.log.LoggingHelper;
+import org.hibernate.metamodel.mapping.internal.ToOneAttributeMapping;
+import org.hibernate.persister.entity.EntityPersister;
+import org.hibernate.spi.NavigablePath;
+import org.hibernate.sql.exec.spi.ExecutionContext;
+import org.hibernate.sql.results.graph.DomainResultAssembler;
+import org.hibernate.sql.results.graph.FetchParentAccess;
+import org.hibernate.sql.results.graph.entity.LoadingEntityEntry;
+import org.hibernate.sql.results.jdbc.spi.JdbcValuesSourceProcessingState;
+import org.hibernate.sql.results.jdbc.spi.RowProcessingState;
+
+/**
+ * Loads entities from the persistence context or creates proxies if not found there,
+ * and initializes all proxies in a batch.
+ */
+public class BatchInitializeEntitySelectFetchInitializer extends AbstractBatchEntitySelectFetchInitializer {
+
+	private final Set<EntityKey> toBatchLoad = new HashSet<>();
+	private State state = State.UNINITIALIZED;
+
+	public BatchInitializeEntitySelectFetchInitializer(
+			FetchParentAccess parentAccess,
+			ToOneAttributeMapping referencedModelPart,
+			NavigablePath fetchedNavigable,
+			EntityPersister concreteDescriptor,
+			DomainResultAssembler<?> identifierAssembler) {
+		super( parentAccess, referencedModelPart, fetchedNavigable, concreteDescriptor, identifierAssembler );
+	}
+
+	@Override
+	protected void registerResolutionListener() {
+		// No-op, because we resolve a proxy
+	}
+
+	@Override
+	public void resolveKey(RowProcessingState rowProcessingState) {
+		if ( state != State.UNINITIALIZED ) {
+			return;
+		}
+		super.resolveKey( rowProcessingState );
+		state = entityKey == null ? State.MISSING : State.KEY_RESOLVED;
+	}
+
+	@Override
+	public void resolveInstance(RowProcessingState rowProcessingState) {
+		if ( state != State.KEY_RESOLVED ) {
+			return;
+		}
+
+		state = State.INITIALIZED;
+		final SharedSessionContractImplementor session = rowProcessingState.getSession();
+		entityInstance = session.getPersistenceContext().getEntity( entityKey );
+		if ( entityInstance == null ) {
+			final LoadingEntityEntry loadingEntityEntry = rowProcessingState.getJdbcValuesSourceProcessingState()
+					.findLoadingEntityLocally( entityKey );
+			if ( loadingEntityEntry != null ) {
+				loadingEntityEntry.getEntityInitializer().resolveInstance( rowProcessingState );
+				entityInstance = loadingEntityEntry.getEntityInstance();
+			}
+			else {
+				if ( entityInstance == null ) {
+					// Force creating a proxy
+					entityInstance = session.internalLoad(
+							entityKey.getEntityName(),
+							entityKey.getIdentifier(),
+							false,
+							false
+					);
+					toBatchLoad.add( entityKey );
+				}
+			}
+		}
+	}
+
+	@Override
+	public boolean isEntityInitialized() {
+		return state == State.INITIALIZED;
+	}
+
+	@Override
+	public void finishUpRow(RowProcessingState rowProcessingState) {
+		super.finishUpRow( rowProcessingState );
+		state = State.UNINITIALIZED;
+	}
+
+	@Override
+	public void endLoading(ExecutionContext context) {
+		final SharedSessionContractImplementor session = context.getSession();
+		for ( EntityKey key : toBatchLoad ) {
+			loadInstance( key, referencedModelPart, session );
+		}
+		toBatchLoad.clear();
+		parentAccess = null;
+	}
+
+	@Override
+	public String toString() {
+		return "BatchInitializeEntitySelectFetchInitializer(" + LoggingHelper.toLoggableString( getNavigablePath() ) + ")";
+	}
+
+	enum State {
+		UNINITIALIZED,
+		MISSING,
+		KEY_RESOLVED,
+		INITIALIZED
+	}
+
+}

hibernate-core/src/main/java/org/hibernate/sql/results/graph/entity/internal/EntityFetchSelectImpl.java

@@ -8,7 +8,6 @@ package org.hibernate.sql.results.graph.entity.internal;
 
 import org.hibernate.engine.FetchTiming;
 import org.hibernate.metamodel.mapping.internal.ToOneAttributeMapping;
-import org.hibernate.persister.entity.EntityPersister;
 import org.hibernate.spi.NavigablePath;
 import org.hibernate.sql.results.graph.AssemblerCreationState;
 import org.hibernate.sql.results.graph.DomainResult;
@@ -17,7 +16,6 @@ import org.hibernate.sql.results.graph.DomainResultCreationState;
 import org.hibernate.sql.results.graph.FetchParent;
 import org.hibernate.sql.results.graph.FetchParentAccess;
 import org.hibernate.sql.results.graph.Initializer;
-import org.hibernate.sql.results.graph.entity.EntityInitializer;
 
 /**
  * An eager entity fetch performed as a subsequent (n+1) select
@@ -52,52 +50,22 @@ public class EntityFetchSelectImpl extends AbstractNonJoinedEntityFetch {
 	}
 
 	@Override
-	public DomainResultAssembler<?> createAssembler(FetchParentAccess parentAccess, AssemblerCreationState creationState) {
+	public DomainResultAssembler<?> createAssembler(
+			FetchParentAccess parentAccess,
+			AssemblerCreationState creationState) {
 		final Initializer initializer = creationState.resolveInitializer(
 				getNavigablePath(),
 				getFetchedMapping(),
-				() -> {
+				() ->
-
+						EntitySelectFetchInitializerBuilder.createInitializer(
-					EntityPersister entityPersister = getReferencedMappingContainer().getEntityPersister();
-
-					final ToOneAttributeMapping fetchedAttribute = (ToOneAttributeMapping) getFetchedMapping();
-					if ( selectByUniqueKey ) {
-						return new EntitySelectFetchByUniqueKeyInitializer(
 								parentAccess,
-								fetchedAttribute,
+								(ToOneAttributeMapping) getFetchedMapping(),
+								getReferencedMappingContainer().getEntityPersister(),
+								keyResult,
 								getNavigablePath(),
-								entityPersister,
+								selectByUniqueKey,
-								keyResult.createResultAssembler( parentAccess, creationState )
+								creationState
-						);
+						)
-					}
-					if ( entityPersister.isBatchLoadable() && !creationState.isScrollResult() ) {
-						if ( parentAccess.isEmbeddableInitializer() ) {
-							return new BatchEntityInsideEmbeddableSelectFetchInitializer(
-									parentAccess,
-									fetchedAttribute,
-									getNavigablePath(),
-									entityPersister,
-									keyResult.createResultAssembler( parentAccess, creationState )
-							);
-						}
-						return new BatchEntitySelectFetchInitializer(
-								parentAccess,
-								fetchedAttribute,
-								getNavigablePath(),
-								entityPersister,
-								keyResult.createResultAssembler( parentAccess, creationState )
-						);
-					}
-					else {
-						return new EntitySelectFetchInitializer(
-								parentAccess,
-								fetchedAttribute,
-								getNavigablePath(),
-								entityPersister,
-								keyResult.createResultAssembler( parentAccess, creationState )
-						);
-					}
-				}
 		);
 
 		return new EntityAssembler( getResultJavaType(), initializer.asEntityInitializer() );

hibernate-core/src/main/java/org/hibernate/sql/results/graph/entity/internal/EntitySelectFetchInitializerBuilder.java

@@ -0,0 +1,111 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * License: GNU Lesser General Public License (LGPL), version 2.1 or later
+ * See the lgpl.txt file in the root directory or http://www.gnu.org/licenses/lgpl-2.1.html
+ */
+package org.hibernate.sql.results.graph.entity.internal;
+
+import org.hibernate.metamodel.internal.StandardEmbeddableInstantiator;
+import org.hibernate.metamodel.mapping.EmbeddableValuedModelPart;
+import org.hibernate.metamodel.mapping.internal.ToOneAttributeMapping;
+import org.hibernate.persister.entity.EntityPersister;
+import org.hibernate.spi.NavigablePath;
+import org.hibernate.sql.results.graph.AbstractFetchParentAccess;
+import org.hibernate.sql.results.graph.AssemblerCreationState;
+import org.hibernate.sql.results.graph.DomainResult;
+import org.hibernate.sql.results.graph.FetchParentAccess;
+import org.hibernate.sql.results.graph.embeddable.EmbeddableInitializer;
+
+public class EntitySelectFetchInitializerBuilder {
+
+	public static AbstractFetchParentAccess createInitializer(
+			FetchParentAccess parentAccess,
+			ToOneAttributeMapping fetchedAttribute,
+			EntityPersister entityPersister,
+			DomainResult<?> keyResult,
+			NavigablePath navigablePath,
+			boolean selectByUniqueKey,
+			AssemblerCreationState creationState) {
+		if ( selectByUniqueKey ) {
+			return new EntitySelectFetchByUniqueKeyInitializer(
+					parentAccess,
+					fetchedAttribute,
+					navigablePath,
+					entityPersister,
+					keyResult.createResultAssembler( parentAccess, creationState )
+			);
+		}
+		final BatchMode batchMode = determineBatchMode( entityPersister, parentAccess, creationState );
+		switch ( batchMode ) {
+			case NONE:
+				return new EntitySelectFetchInitializer(
+						parentAccess,
+						fetchedAttribute,
+						navigablePath,
+						entityPersister,
+						keyResult.createResultAssembler( parentAccess, creationState )
+				);
+			case BATCH_LOAD:
+				if ( parentAccess.isEmbeddableInitializer() ) {
+					return new BatchEntityInsideEmbeddableSelectFetchInitializer(
+							parentAccess,
+							fetchedAttribute,
+							navigablePath,
+							entityPersister,
+							keyResult.createResultAssembler( parentAccess, creationState )
+					);
+				}
+				else {
+					return new BatchEntitySelectFetchInitializer(
+							parentAccess,
+							fetchedAttribute,
+							navigablePath,
+							entityPersister,
+							keyResult.createResultAssembler( parentAccess, creationState )
+					);
+				}
+			case BATCH_INITIALIZE:
+				return new BatchInitializeEntitySelectFetchInitializer(
+						parentAccess,
+						fetchedAttribute,
+						navigablePath,
+						entityPersister,
+						keyResult.createResultAssembler( parentAccess, creationState )
+				);
+		}
+		throw new IllegalStateException( "Should be unreachable" );
+	}
+
+	private static BatchMode determineBatchMode(EntityPersister entityPersister, FetchParentAccess parentAccess, AssemblerCreationState creationState) {
+		if ( !entityPersister.isBatchLoadable() || creationState.isScrollResult() ) {
+			return BatchMode.NONE;
+		}
+		while ( parentAccess.isEmbeddableInitializer() ) {
+			final EmbeddableInitializer embeddableInitializer = parentAccess.asEmbeddableInitializer();
+			final EmbeddableValuedModelPart initializedPart = embeddableInitializer.getInitializedPart();
+			// For entity identifier mappings we can't batch load,
+			// because the entity identifier needs the instance in the resolveKey phase,
+			// but batch loading is inherently executed out of order
+			if ( initializedPart.isEntityIdentifierMapping()
+					// todo: check if the virtual check is necessary
+					|| initializedPart.isVirtual()
+					// If the parent embeddable has a custom instantiator, we can't inject entities later through setValues()
+					|| !( initializedPart.getMappedType().getRepresentationStrategy().getInstantiator() instanceof StandardEmbeddableInstantiator ) ) {
+				return entityPersister.hasSubclasses() ? BatchMode.NONE : BatchMode.BATCH_INITIALIZE;
+			}
+			parentAccess = parentAccess.getFetchParentAccess();
+			if ( parentAccess == null ) {
+				break;
+			}
+		}
+		return BatchMode.BATCH_LOAD;
+	}
+
+	enum BatchMode {
+		NONE,
+		BATCH_LOAD,
+		BATCH_INITIALIZE
+	}
+
+}

hibernate-core/src/main/java/org/hibernate/sql/results/internal/domain/CircularFetchImpl.java

@@ -9,7 +9,6 @@ package org.hibernate.sql.results.internal.domain;
 import org.hibernate.engine.FetchTiming;
 import org.hibernate.metamodel.mapping.EntityMappingType;
 import org.hibernate.metamodel.mapping.internal.ToOneAttributeMapping;
-import org.hibernate.persister.entity.EntityPersister;
 import org.hibernate.spi.NavigablePath;
 import org.hibernate.sql.results.graph.AssemblerCreationState;
 import org.hibernate.sql.results.graph.BiDirectionalFetch;
@@ -20,11 +19,8 @@ import org.hibernate.sql.results.graph.FetchParentAccess;
 import org.hibernate.sql.results.graph.Fetchable;
 import org.hibernate.sql.results.graph.Initializer;
 import org.hibernate.sql.results.graph.entity.EntityInitializer;
-import org.hibernate.sql.results.graph.entity.internal.BatchEntityInsideEmbeddableSelectFetchInitializer;
-import org.hibernate.sql.results.graph.entity.internal.BatchEntitySelectFetchInitializer;
 import org.hibernate.sql.results.graph.entity.internal.EntityDelayedFetchInitializer;
-import org.hibernate.sql.results.graph.entity.internal.EntitySelectFetchByUniqueKeyInitializer;
+import org.hibernate.sql.results.graph.entity.internal.EntitySelectFetchInitializerBuilder;
-import org.hibernate.sql.results.graph.entity.internal.EntitySelectFetchInitializer;
 import org.hibernate.sql.results.jdbc.spi.JdbcValuesSourceProcessingOptions;
 import org.hibernate.sql.results.jdbc.spi.RowProcessingState;
 import org.hibernate.type.descriptor.java.JavaType;
@@ -95,50 +91,21 @@ public class CircularFetchImpl implements BiDirectionalFetch {
 	public DomainResultAssembler<?> createAssembler(
 			FetchParentAccess parentAccess,
 			AssemblerCreationState creationState) {
-		final DomainResultAssembler<?> keyAssembler = keyResult.createResultAssembler( parentAccess, creationState );
 
 		final Initializer initializer = creationState.resolveInitializer(
 				getNavigablePath(),
 				referencedModelPart,
 				() -> {
 					if ( timing == FetchTiming.IMMEDIATE ) {
-						if ( selectByUniqueKey ) {
+						return EntitySelectFetchInitializerBuilder.createInitializer(
-							return new EntitySelectFetchByUniqueKeyInitializer(
+								parentAccess,
-									parentAccess,
+								fetchable,
-									fetchable,
+								entityMappingType.getEntityPersister(),
-									getNavigablePath(),
+								keyResult,
-									entityMappingType.getEntityPersister(),
+								getNavigablePath(),
-									keyAssembler
+								selectByUniqueKey,
-							);
+								creationState
-						}
+						);
-						final EntityPersister entityPersister = entityMappingType.getEntityPersister();
-						if ( entityPersister.isBatchLoadable() ) {
-							if ( parentAccess.isEmbeddableInitializer() ) {
-								return new BatchEntityInsideEmbeddableSelectFetchInitializer(
-										parentAccess,
-										referencedModelPart,
-										getNavigablePath(),
-										entityPersister,
-										keyResult.createResultAssembler( parentAccess, creationState )
-								);
-							}
-							return new BatchEntitySelectFetchInitializer(
-									parentAccess,
-									referencedModelPart,
-									getReferencedPath(),
-									entityPersister,
-									keyAssembler
-							);
-						}
-						else {
-							return new EntitySelectFetchInitializer(
-									parentAccess,
-									(ToOneAttributeMapping) referencedModelPart,
-									getReferencedPath(),
-									entityPersister,
-									keyAssembler
-							);
-						}
 					}
 					else {
 						return new EntityDelayedFetchInitializer(
@@ -146,7 +113,7 @@ public class CircularFetchImpl implements BiDirectionalFetch {
 								getReferencedPath(),
 								fetchable,
 								selectByUniqueKey,
-								keyAssembler
+								keyResult.createResultAssembler( parentAccess, creationState )
 						);
 					}
 				}
@@ -158,6 +125,8 @@ public class CircularFetchImpl implements BiDirectionalFetch {
 		);
 	}
 
+
+
 	@Override
 	public FetchTiming getTiming() {
 		return timing;