share more code between the batch loaders

hibernate-core/src/main/java/org/hibernate/engine/spi/BatchFetchQueue.java

@@ -267,7 +267,8 @@ public class BatchFetchQueue {
 
 		// TODO: this needn't exclude subclasses...
 
-		LinkedHashSet<EntityKey> set =  batchLoadableEntityKeys.get( entityDescriptor.getEntityName() );
+		final LinkedHashSet<EntityKey> set =
+				batchLoadableEntityKeys.get( entityDescriptor.getEntityName() );
 		if ( set != null ) {
 			for ( EntityKey key : set ) {
 				if ( checkForEnd && i == end ) {

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

@@ -17,6 +17,10 @@ import org.hibernate.loader.ast.spi.CollectionBatchLoader;
 import org.hibernate.metamodel.mapping.PluralAttributeMapping;
 import org.hibernate.sql.results.internal.ResultsHelper;
 
+import java.lang.reflect.Array;
+
+import static org.hibernate.loader.ast.internal.MultiKeyLoadHelper.hasSingleId;
+import static org.hibernate.loader.ast.internal.MultiKeyLoadHelper.trimIdBatch;
 import static org.hibernate.loader.ast.internal.MultiKeyLoadLogging.MULTI_KEY_LOAD_DEBUG_ENABLED;
 import static org.hibernate.loader.ast.internal.MultiKeyLoadLogging.MULTI_KEY_LOAD_LOGGER;
 
@@ -31,6 +35,8 @@ public abstract class AbstractCollectionBatchLoader implements CollectionBatchLo
 
 	private final int keyJdbcCount;
 
+	private final CollectionLoaderSingleKey singleKeyLoader;
+
 	public AbstractCollectionBatchLoader(
 			int domainBatchSize,
 			LoadQueryInfluencers influencers,
@@ -42,6 +48,8 @@ public abstract class AbstractCollectionBatchLoader implements CollectionBatchLo
 		this.keyJdbcCount = attributeMapping.getJdbcTypeCount();
 		this.sessionFactory = sessionFactory;
 		this.influencers = influencers;
+
+		singleKeyLoader = new CollectionLoaderSingleKey( getLoadable(), getInfluencers(), getSessionFactory() );
 	}
 
 	@Override
@@ -66,15 +74,35 @@ public abstract class AbstractCollectionBatchLoader implements CollectionBatchLo
 		return keyJdbcCount;
 	}
 
-	protected void finishInitializingKey(
+	abstract void initializeKeys(Object key, Object[] keysToInitialize, SharedSessionContractImplementor session);
-			Object key,
+
-			SharedSessionContractImplementor session) {
+	@Override
+	public PersistentCollection<?> load(Object key, SharedSessionContractImplementor session) {
+		if ( MULTI_KEY_LOAD_DEBUG_ENABLED ) {
+			MULTI_KEY_LOAD_LOGGER.debugf( "Batch fetching collection: %s.%s",
+					getLoadable().getNavigableRole().getFullPath(), key );
+		}
+
+		final Object[] keys = resolveKeysToInitialize( key, session );
+
+		if ( hasSingleId( keys ) ) {
+			return singleKeyLoader.load( key, session );
+		}
+
+		initializeKeys( key, keys, session );
+
+		final CollectionKey collectionKey = new CollectionKey( getLoadable().getCollectionDescriptor(), key );
+		return session.getPersistenceContext().getCollection( collectionKey );
+	}
+
+	protected void finishInitializingKey(Object key, SharedSessionContractImplementor session) {
 		if ( key == null ) {
 			return;
 		}
 
 		if ( MULTI_KEY_LOAD_DEBUG_ENABLED ) {
-			MULTI_KEY_LOAD_LOGGER.debugf( "Finishing initializing batch-fetched collection : %s.%s", attributeMapping.getNavigableRole().getFullPath(), key );
+			MULTI_KEY_LOAD_LOGGER.debugf( "Finishing initializing batch-fetched collection: %s.%s",
+					attributeMapping.getNavigableRole().getFullPath(), key );
 		}
 
 		final PersistenceContext persistenceContext = session.getPersistenceContext();
@@ -93,4 +121,20 @@ public abstract class AbstractCollectionBatchLoader implements CollectionBatchLo
 		}
 
 	}
+
+	Object[] resolveKeysToInitialize(Object keyBeingLoaded, SharedSessionContractImplementor session) {
+		final int length = getDomainBatchSize();
+		final Class<?> keyType = getLoadable().getKeyDescriptor().getJavaType().getJavaTypeClass();
+		final Object[] keysToInitialize = (Object[]) Array.newInstance( keyType, length );
+		session.getPersistenceContextInternal().getBatchFetchQueue()
+				.collectBatchLoadableCollectionKeys(
+						length,
+						(index, key) -> keysToInitialize[index] = key,
+						keyBeingLoaded,
+						getLoadable()
+				);
+		// now trim down the array to the number of keys we found
+		return trimIdBatch( length, keysToInitialize );
+	}
+
 }

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

@@ -0,0 +1,65 @@
+/*
+ * 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.loader.ast.internal;
+
+import org.hibernate.LockOptions;
+import org.hibernate.engine.spi.EntityKey;
+import org.hibernate.engine.spi.SessionFactoryImplementor;
+import org.hibernate.engine.spi.SharedSessionContractImplementor;
+import org.hibernate.loader.ast.spi.EntityBatchLoader;
+import org.hibernate.metamodel.mapping.EntityMappingType;
+
+import static org.hibernate.loader.ast.internal.MultiKeyLoadHelper.hasSingleId;
+import static org.hibernate.loader.ast.internal.MultiKeyLoadLogging.MULTI_KEY_LOAD_DEBUG_ENABLED;
+import static org.hibernate.loader.ast.internal.MultiKeyLoadLogging.MULTI_KEY_LOAD_LOGGER;
+
+public abstract class AbstractEntityBatchLoader<T>
+		extends SingleIdEntityLoaderSupport<T>
+		implements EntityBatchLoader<T> {
+
+	private final SingleIdEntityLoaderStandardImpl<T> singleIdLoader;
+
+	public AbstractEntityBatchLoader(EntityMappingType entityDescriptor, SessionFactoryImplementor sessionFactory) {
+		super( entityDescriptor, sessionFactory );
+		singleIdLoader = new SingleIdEntityLoaderStandardImpl<>( entityDescriptor, sessionFactory );
+	}
+
+	protected abstract void initializeEntities(
+			Object[] idsToInitialize,
+			Object pkValue,
+			Object entityInstance,
+			LockOptions lockOptions,
+			Boolean readOnly,
+			SharedSessionContractImplementor session);
+
+	protected abstract Object[] resolveIdsToInitialize(Object id, SharedSessionContractImplementor session);
+
+	@Override
+	public final T load(
+			Object id,
+			Object entityInstance,
+			LockOptions lockOptions,
+			Boolean readOnly,
+			SharedSessionContractImplementor session) {
+		if ( MULTI_KEY_LOAD_DEBUG_ENABLED ) {
+			MULTI_KEY_LOAD_LOGGER.debugf( "Batch fetching entity `%s#%s`", getLoadable().getEntityName(), id );
+		}
+
+		final Object[] ids = resolveIdsToInitialize( id, session );
+
+		if ( hasSingleId( ids ) ) {
+			return singleIdLoader.load( id, entityInstance, lockOptions, readOnly, session );
+		}
+
+		initializeEntities( ids, id, entityInstance, lockOptions, readOnly, session );
+
+		final EntityKey entityKey = session.generateEntityKey( id, getLoadable().getEntityPersister() );
+		//noinspection unchecked
+		return (T) session.getPersistenceContext().getEntity( entityKey );
+	}
+
+}

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

@@ -7,17 +7,15 @@
 package org.hibernate.loader.ast.internal;
 
 import java.lang.reflect.Array;
-import java.util.Arrays;
 
 import org.hibernate.LockOptions;
-import org.hibernate.collection.spi.PersistentCollection;
-import org.hibernate.engine.spi.CollectionKey;
 import org.hibernate.engine.spi.LoadQueryInfluencers;
 import org.hibernate.engine.spi.SessionFactoryImplementor;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.engine.spi.SubselectFetch;
 import org.hibernate.loader.ast.spi.CollectionBatchLoader;
 import org.hibernate.loader.ast.spi.SqlArrayMultiKeyLoader;
+import org.hibernate.metamodel.mapping.ForeignKeyDescriptor;
 import org.hibernate.metamodel.mapping.JdbcMapping;
 import org.hibernate.metamodel.mapping.PluralAttributeMapping;
 import org.hibernate.metamodel.mapping.internal.SimpleForeignKeyDescriptor;
@@ -34,23 +32,21 @@ import org.hibernate.sql.results.spi.ListResultsConsumer;
 import org.hibernate.type.BasicType;
 
 import static java.util.Collections.singletonList;
-import static org.hibernate.loader.ast.internal.MultiKeyLoadHelper.hasSingleId;
 import static org.hibernate.loader.ast.internal.MultiKeyLoadLogging.MULTI_KEY_LOAD_DEBUG_ENABLED;
 import static org.hibernate.loader.ast.internal.MultiKeyLoadLogging.MULTI_KEY_LOAD_LOGGER;
 
 /**
- * CollectionBatchLoader using a SQL ARRAY parameter to pass the key values
+ * {@link CollectionBatchLoader} using a SQL {@code ARRAY} parameter to pass the key values.
  *
  * @author Steve Ebersole
  */
 public class CollectionBatchLoaderArrayParam
 		extends AbstractCollectionBatchLoader
-		implements CollectionBatchLoader, SqlArrayMultiKeyLoader {
+		implements SqlArrayMultiKeyLoader {
 	private final JdbcMapping arrayJdbcMapping;
 	private final JdbcParameter jdbcParameter;
 	private final SelectStatement sqlSelect;
 	private final JdbcOperationQuerySelect jdbcSelectOperation;
-	private final CollectionLoaderSingleKey singleKeyLoader;
 
 	public CollectionBatchLoaderArrayParam(
 			int domainBatchSize,
@@ -67,17 +63,19 @@ public class CollectionBatchLoaderArrayParam
 			);
 		}
 
-		final SimpleForeignKeyDescriptor keyDescriptor = getKeyDescriptor();
+		final ForeignKeyDescriptor keyDescriptor = getLoadable().getKeyDescriptor();
-
 		final Class<?> keyType = keyDescriptor.getJavaType().getJavaTypeClass();
 		final Class<?> arrayClass = Array.newInstance( keyType, 0 ).getClass();
 
+		// this typecast is always safe because we don't instantiate this class unless the FK is "simple"
+		final SimpleForeignKeyDescriptor simpleKeyDescriptor = (SimpleForeignKeyDescriptor) keyDescriptor;
+
 		final BasicType<?> arrayBasicType = getSessionFactory().getTypeConfiguration()
 				.getBasicTypeRegistry()
 				.getRegisteredType( arrayClass );
 		arrayJdbcMapping = MultiKeyLoadHelper.resolveArrayJdbcMapping(
 				arrayBasicType,
-				keyDescriptor.getJdbcMapping(),
+				simpleKeyDescriptor.getJdbcMapping(),
 				arrayClass,
 				getSessionFactory()
 		);
@@ -85,7 +83,7 @@ public class CollectionBatchLoaderArrayParam
 		jdbcParameter = new JdbcParameterImpl( arrayJdbcMapping );
 		sqlSelect = LoaderSelectBuilder.createSelectBySingleArrayParameter(
 				getLoadable(),
-				keyDescriptor.getKeyPart(),
+				simpleKeyDescriptor.getKeyPart(),
 				getInfluencers(),
 				LockOptions.NONE,
 				jdbcParameter,
@@ -97,51 +95,19 @@ public class CollectionBatchLoaderArrayParam
 				.getSqlAstTranslatorFactory()
 				.buildSelectTranslator( getSessionFactory(), sqlSelect )
 				.translate( JdbcParameterBindings.NO_BINDINGS, QueryOptions.NONE );
-
-		singleKeyLoader = new CollectionLoaderSingleKey( attributeMapping, loadQueryInfluencers, sessionFactory );
-	}
-
-	private SimpleForeignKeyDescriptor getKeyDescriptor() {
-		return (SimpleForeignKeyDescriptor) getLoadable().getKeyDescriptor();
 	}
 
 	@Override
-	public PersistentCollection<?> load(Object key, SharedSessionContractImplementor session) {
+	void initializeKeys(Object key, Object[] keysToInitialize, SharedSessionContractImplementor session) {
 		if ( MULTI_KEY_LOAD_DEBUG_ENABLED ) {
-			MULTI_KEY_LOAD_LOGGER.debugf( "Batch loading entity `%s#%s`", getLoadable().getNavigableRole().getFullPath(), key );
+			MULTI_KEY_LOAD_LOGGER.debugf(
+					"Collection keys to batch-fetch initialize (`%s#%s`) %s",
+					getLoadable().getNavigableRole().getFullPath(),
+					key,
+					keysToInitialize
+			);
 		}
 
-		final Object[] keys = resolveKeysToInitialize( key, session );
-
-		if ( hasSingleId( keys ) ) {
-			return singleKeyLoader.load( key, session );
-		}
-
-		initializeKeys( keys, session );
-
-		final CollectionKey collectionKey = new CollectionKey( getLoadable().getCollectionDescriptor(), key );
-		return session.getPersistenceContext().getCollection( collectionKey );
-	}
-
-	private Object[] resolveKeysToInitialize(Object keyBeingLoaded, SharedSessionContractImplementor session) {
-		final int length = getDomainBatchSize();
-		final Class<?> keyType = getKeyDescriptor().getJavaType().getJavaTypeClass();
-		final Object[] keysToInitialize = (Object[]) Array.newInstance( keyType, length );
-		session.getPersistenceContextInternal().getBatchFetchQueue()
-				.collectBatchLoadableCollectionKeys(
-						length,
-						(index, value) -> keysToInitialize[index] = value,
-						keyBeingLoaded,
-						getLoadable()
-				);
-		int newLength = length;
-		while ( newLength>1 && keysToInitialize[newLength-1] == null ) {
-			newLength--;
-		}
-		return newLength < length ? Arrays.copyOf( keysToInitialize, newLength ) : keysToInitialize;
-	}
-
-	private void initializeKeys(Object[] keysToInitialize, SharedSessionContractImplementor session) {
 		assert jdbcSelectOperation != null;
 		assert jdbcParameter != null;
 
@@ -169,8 +135,8 @@ public class CollectionBatchLoaderArrayParam
 				ListResultsConsumer.UniqueSemantic.FILTER
 		);
 
-		for ( int i = 0; i < keysToInitialize.length; i++ ) {
+		for ( Object initializedKey : keysToInitialize ) {
-			finishInitializingKey( keysToInitialize[i], session );
+			finishInitializingKey( initializedKey, session );
 		}
 	}
 }

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

@@ -10,15 +10,11 @@ import java.util.ArrayList;
 import java.util.List;
 
 import org.hibernate.LockOptions;
-import org.hibernate.collection.spi.PersistentCollection;
 import org.hibernate.engine.spi.BatchFetchQueue;
-import org.hibernate.engine.spi.CollectionKey;
 import org.hibernate.engine.spi.LoadQueryInfluencers;
 import org.hibernate.engine.spi.SessionFactoryImplementor;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.engine.spi.SubselectFetch;
-import org.hibernate.internal.util.MutableInteger;
-import org.hibernate.internal.util.collections.CollectionHelper;
 import org.hibernate.loader.ast.spi.CollectionBatchLoader;
 import org.hibernate.loader.ast.spi.SqlArrayMultiKeyLoader;
 import org.hibernate.metamodel.mapping.PluralAttributeMapping;
@@ -28,25 +24,24 @@ import org.hibernate.sql.ast.tree.select.SelectStatement;
 import org.hibernate.sql.exec.spi.JdbcOperationQuerySelect;
 import org.hibernate.sql.exec.spi.JdbcParameterBindings;
 
+import static org.hibernate.loader.ast.internal.MultiKeyLoadHelper.countIds;
 import static org.hibernate.loader.ast.internal.MultiKeyLoadLogging.MULTI_KEY_LOAD_DEBUG_ENABLED;
 import static org.hibernate.loader.ast.internal.MultiKeyLoadLogging.MULTI_KEY_LOAD_LOGGER;
 
 /**
- * CollectionLoader for batch fetching using a SQL IN predicate
+ * {@link CollectionBatchLoader} for batch fetching using a SQL {@code IN} predicate.
  *
  * @author Steve Ebersole
  */
 public class CollectionBatchLoaderInPredicate
 		extends AbstractCollectionBatchLoader
-		implements CollectionBatchLoader, SqlArrayMultiKeyLoader {
+		implements SqlArrayMultiKeyLoader {
 	private final int keyColumnCount;
 	private final int sqlBatchSize;
 	private final List<JdbcParameter> jdbcParameters;
 	private final SelectStatement sqlAst;
 	private final JdbcOperationQuerySelect jdbcSelect;
 
-	private CollectionLoaderSingleKey singleKeyLoader;
-
 	public CollectionBatchLoaderInPredicate(
 			int domainBatchSize,
 			LoadQueryInfluencers influencers,
@@ -54,8 +49,8 @@ public class CollectionBatchLoaderInPredicate
 			SessionFactoryImplementor sessionFactory) {
 		super( domainBatchSize, influencers, attributeMapping, sessionFactory );
 
-		this.keyColumnCount = attributeMapping.getKeyDescriptor().getJdbcTypeCount();
+		keyColumnCount = attributeMapping.getKeyDescriptor().getJdbcTypeCount();
-		this.sqlBatchSize = sessionFactory.getJdbcServices()
+		sqlBatchSize = sessionFactory.getJdbcServices()
 				.getDialect()
 				.getBatchLoadSizingStrategy()
 				.determineOptimalBatchLoadSize( keyColumnCount, domainBatchSize, false );
@@ -68,8 +63,8 @@ public class CollectionBatchLoaderInPredicate
 			);
 		}
 
-		this.jdbcParameters = new ArrayList<>();
+		jdbcParameters = new ArrayList<>();
-		this.sqlAst = LoaderSelectBuilder.createSelect(
+		sqlAst = LoaderSelectBuilder.createSelect(
 				attributeMapping,
 				null,
 				attributeMapping.getKeyDescriptor(),
@@ -80,9 +75,9 @@ public class CollectionBatchLoaderInPredicate
 				jdbcParameters::add,
 				sessionFactory
 		);
-		assert this.jdbcParameters.size() == this.sqlBatchSize * this.keyColumnCount;
+		assert jdbcParameters.size() == sqlBatchSize * keyColumnCount;
 
-		this.jdbcSelect = sessionFactory.getJdbcServices()
+		jdbcSelect = sessionFactory.getJdbcServices()
 				.getJdbcEnvironment()
 				.getSqlAstTranslatorFactory()
 				.buildSelectTranslator( sessionFactory, sqlAst )
@@ -90,53 +85,7 @@ public class CollectionBatchLoaderInPredicate
 	}
 
 	@Override
-	public PersistentCollection<?> load(
+	void initializeKeys(Object key, Object[] keysToInitialize, SharedSessionContractImplementor session) {
-			Object key,
-			SharedSessionContractImplementor session) {
-		if ( MULTI_KEY_LOAD_DEBUG_ENABLED ) {
-			MULTI_KEY_LOAD_LOGGER.debugf( "Loading collection `%s#%s` by batch-fetch", getLoadable().getNavigableRole().getFullPath(), key );
-		}
-
-		final MutableInteger nonNullCounter = new MutableInteger();
-		final ArrayList<Object> keysToInitialize = CollectionHelper.arrayList( getDomainBatchSize() );
-		session.getPersistenceContextInternal().getBatchFetchQueue().collectBatchLoadableCollectionKeys(
-				getDomainBatchSize(),
-				(index, batchableKey) -> {
-					keysToInitialize.add( batchableKey );
-					if ( batchableKey != null ) {
-						nonNullCounter.increment();
-					}
-				},
-				key,
-				getLoadable().asPluralAttributeMapping()
-		);
-
-		if ( nonNullCounter.get() <= 0 ) {
-			throw new IllegalStateException( "Number of non-null collection keys to batch fetch should never be 0" );
-		}
-
-		if ( nonNullCounter.get() == 1 ) {
-			prepareSingleKeyLoaderIfNeeded();
-			return singleKeyLoader.load( key, session );
-		}
-
-		initializeKeys( key, keysToInitialize.toArray( keysToInitialize.toArray( new Object[0] ) ), nonNullCounter.get(), session );
-
-		final CollectionKey collectionKey = new CollectionKey( getLoadable().getCollectionDescriptor(), key );
-		return session.getPersistenceContext().getCollection( collectionKey );
-	}
-
-	private void prepareSingleKeyLoaderIfNeeded() {
-		if ( singleKeyLoader == null ) {
-			singleKeyLoader = new CollectionLoaderSingleKey( getLoadable(), getInfluencers(), getSessionFactory() );
-		}
-	}
-
-	private <T> void initializeKeys(
-			T key,
-			T[] keysToInitialize,
-			int nonNullKeysToInitializeCount,
-			SharedSessionContractImplementor session) {
 		if ( MULTI_KEY_LOAD_DEBUG_ENABLED ) {
 			MULTI_KEY_LOAD_LOGGER.debugf(
 					"Collection keys to batch-fetch initialize (`%s#%s`) %s",
@@ -146,7 +95,7 @@ public class CollectionBatchLoaderInPredicate
 			);
 		}
 
-		final MultiKeyLoadChunker<T> chunker = new MultiKeyLoadChunker<>(
+		final MultiKeyLoadChunker<Object> chunker = new MultiKeyLoadChunker<>(
 				sqlBatchSize,
 				keyColumnCount,
 				getLoadable().getKeyDescriptor(),
@@ -159,7 +108,7 @@ public class CollectionBatchLoaderInPredicate
 
 		chunker.processChunks(
 				keysToInitialize,
-				nonNullKeysToInitializeCount,
+				countIds( keysToInitialize ),
 				(jdbcParameterBindings, session1) -> {
 					// Create a RegistrationHandler for handling any subselect fetches we encounter handling this chunk
 					final SubselectFetch.RegistrationHandler registrationHandler = SubselectFetch.createRegistrationHandler(
@@ -197,8 +146,7 @@ public class CollectionBatchLoaderInPredicate
 					for ( int i = 0; i < nonNullElementCount; i++ ) {
 						final int keyPosition = i + startIndex;
 						if ( keyPosition < keysToInitialize.length ) {
-							final T keyToInitialize = keysToInitialize[keyPosition];
+							finishInitializingKey( keysToInitialize[keyPosition], session );
-							finishInitializingKey( keyToInitialize, session );
 						}
 					}
 				},

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

@@ -11,11 +11,9 @@ import java.util.Arrays;
 import java.util.Locale;
 
 import org.hibernate.LockOptions;
-import org.hibernate.engine.spi.EntityKey;
 import org.hibernate.engine.spi.LoadQueryInfluencers;
 import org.hibernate.engine.spi.SessionFactoryImplementor;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
-import org.hibernate.loader.ast.spi.EntityBatchLoader;
 import org.hibernate.loader.ast.spi.SqlArrayMultiKeyLoader;
 import org.hibernate.metamodel.mapping.BasicEntityIdentifierMapping;
 import org.hibernate.metamodel.mapping.EntityIdentifierMapping;
@@ -29,7 +27,7 @@ import org.hibernate.sql.exec.spi.JdbcOperationQuerySelect;
 import org.hibernate.sql.exec.spi.JdbcParameterBindings;
 
 import static org.hibernate.engine.internal.BatchFetchQueueHelper.removeBatchLoadableEntityKey;
-import static org.hibernate.loader.ast.internal.MultiKeyLoadHelper.hasSingleId;
+import static org.hibernate.loader.ast.internal.MultiKeyLoadHelper.trimIdBatch;
 import static org.hibernate.loader.ast.internal.MultiKeyLoadLogging.MULTI_KEY_LOAD_DEBUG_ENABLED;
 import static org.hibernate.loader.ast.internal.MultiKeyLoadLogging.MULTI_KEY_LOAD_LOGGER;
 
@@ -41,8 +39,8 @@ import static org.hibernate.loader.ast.internal.MultiKeyLoadLogging.MULTI_KEY_LO
  * @author Steve Ebersole
  */
 public class EntityBatchLoaderArrayParam<T>
-		extends SingleIdEntityLoaderSupport<T>
+		extends AbstractEntityBatchLoader<T>
-		implements EntityBatchLoader<T>, SqlArrayMultiKeyLoader {
+		implements SqlArrayMultiKeyLoader {
 	private final int domainBatchSize;
 
 	private final BasicEntityIdentifierMapping identifierMapping;
@@ -50,7 +48,6 @@ public class EntityBatchLoaderArrayParam<T>
 	private final JdbcParameter jdbcParameter;
 	private final SelectStatement sqlAst;
 	private final JdbcOperationQuerySelect jdbcSelectOperation;
-	private final SingleIdEntityLoaderStandardImpl<T> singleIdLoader;
 
 
 	/**
@@ -104,8 +101,6 @@ public class EntityBatchLoaderArrayParam<T>
 				.getSqlAstTranslatorFactory()
 				.buildSelectTranslator( sessionFactory, sqlAst )
 				.translate( JdbcParameterBindings.NO_BINDINGS, QueryOptions.NONE );
-
-		singleIdLoader = new SingleIdEntityLoaderStandardImpl<>( entityDescriptor, sessionFactory );
 	}
 
 	@Override
@@ -113,31 +108,8 @@ public class EntityBatchLoaderArrayParam<T>
 		return domainBatchSize;
 	}
 
-	@Override
-	public final T load(
-			Object pkValue,
-			Object entityInstance,
-			LockOptions lockOptions,
-			Boolean readOnly,
-			SharedSessionContractImplementor session) {
-		if ( MULTI_KEY_LOAD_DEBUG_ENABLED ) {
-			MULTI_KEY_LOAD_LOGGER.debugf( "Batch fetching entity `%s#%s`", getLoadable().getEntityName(), pkValue );
-		}
-
-		final Object[] ids = resolveIdsToInitialize( pkValue, session );
-
-		if ( hasSingleId( ids ) ) {
-			return singleIdLoader.load( pkValue, entityInstance, lockOptions, readOnly, session );
-		}
-
-		initializeEntities( ids, pkValue, entityInstance, lockOptions, readOnly, session );
-
-		final EntityKey entityKey = session.generateEntityKey( pkValue, getLoadable().getEntityPersister() );
-		//noinspection unchecked
-		return (T) session.getPersistenceContext().getEntity( entityKey );
-	}
-
 	protected Object[] resolveIdsToInitialize(Object pkValue, SharedSessionContractImplementor session) {
+		//TODO: should this really be different to EntityBatchLoaderInPredicate impl?
 		final Class<?> idType = identifierMapping.getJavaType().getJavaTypeClass();
 		final Object[] idsToLoad = (Object[]) Array.newInstance( idType, domainBatchSize );
 		session.getPersistenceContextInternal().getBatchFetchQueue()
@@ -147,27 +119,29 @@ public class EntityBatchLoaderArrayParam<T>
 						pkValue,
 						getLoadable()
 				);
-		int newLength = domainBatchSize;
+		return trimIdBatch( domainBatchSize, idsToLoad );
-		while ( newLength>1 && idsToLoad[newLength-1] == null ) {
-			newLength--;
-		}
-		return newLength < domainBatchSize ? Arrays.copyOf( idsToLoad, newLength ) : idsToLoad;
 	}
 
-	private void initializeEntities(
+	@Override
+	protected void initializeEntities(
 			Object[] idsToInitialize,
-			Object pkValue,
+			Object id,
 			Object entityInstance,
 			LockOptions lockOptions,
 			Boolean readOnly,
 			SharedSessionContractImplementor session) {
+		if ( MULTI_KEY_LOAD_DEBUG_ENABLED ) {
+			MULTI_KEY_LOAD_LOGGER.debugf( "Ids to batch-fetch initialize (`%s#%s`) %s",
+					getLoadable().getEntityName(), id, Arrays.toString(idsToInitialize) );
+		}
+
 		LoaderHelper.loadByArrayParameter(
 				idsToInitialize,
 				sqlAst,
 				jdbcSelectOperation,
 				jdbcParameter,
 				arrayJdbcMapping,
-				pkValue,
+				id,
 				entityInstance,
 				getLoadable().getRootEntityDescriptor(),
 				lockOptions,
@@ -175,16 +149,11 @@ public class EntityBatchLoaderArrayParam<T>
 				session
 		);
 
-		//noinspection ForLoopReplaceableByForEach
+		for ( Object initializedId : idsToInitialize ) {
-		for ( int i = 0; i < idsToInitialize.length; i++ ) {
+			if ( initializedId != null ) {
-			final Object id = idsToInitialize[i];
+				// found or not, remove the key from the batch-fetch queue
-			if ( id == null ) {
+				removeBatchLoadableEntityKey( initializedId, getLoadable(), session );
-				// skip any of the null padded ids
-				//		- actually we could probably even break here
-				continue;
 			}
-			// found or not, remove the key from the batch-fetch queue
-			removeBatchLoadableEntityKey( id, getLoadable(), session );
 		}
 	}
 

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

@@ -24,11 +24,8 @@ import org.hibernate.metamodel.mapping.EntityMappingType;
 import org.hibernate.query.spi.QueryOptions;
 import org.hibernate.sql.ast.tree.expression.JdbcParameter;
 import org.hibernate.sql.ast.tree.select.SelectStatement;
-import org.hibernate.sql.exec.internal.JdbcParameterBindingsImpl;
 import org.hibernate.sql.exec.spi.JdbcOperationQuerySelect;
 import org.hibernate.sql.exec.spi.JdbcParameterBindings;
-import org.hibernate.sql.results.internal.RowTransformerStandardImpl;
-import org.hibernate.sql.results.spi.ListResultsConsumer;
 
 import static org.hibernate.internal.util.collections.CollectionHelper.arrayList;
 import static org.hibernate.loader.ast.internal.MultiKeyLoadLogging.MULTI_KEY_LOAD_DEBUG_ENABLED;
@@ -45,8 +42,8 @@ import static org.hibernate.loader.ast.internal.MultiKeyLoadLogging.MULTI_KEY_LO
  * @author Steve Ebersole
  */
 public class EntityBatchLoaderInPredicate<T>
-		extends SingleIdEntityLoaderSupport<T>
+		extends AbstractEntityBatchLoader<T>
-		implements EntityBatchLoader<T>, SqlInPredicateMultiKeyLoader {
+		implements SqlInPredicateMultiKeyLoader {
 	private final int domainBatchSize;
 	private final int sqlBatchSize;
 
@@ -78,7 +75,7 @@ public class EntityBatchLoaderInPredicate<T>
 			);
 		}
 
-		EntityIdentifierMapping identifierMapping = getLoadable().getIdentifierMapping();
+		final EntityIdentifierMapping identifierMapping = getLoadable().getIdentifierMapping();
 
 		final int expectedNumberOfParameters = identifierMapping.getJdbcTypeCount() * sqlBatchSize;
 
@@ -118,37 +115,12 @@ public class EntityBatchLoaderInPredicate<T>
 		return load( pkValue, null, lockOptions, readOnly, session );
 	}
 
+	protected Object[] resolveIdsToInitialize(Object id, SharedSessionContractImplementor session) {
+		return session.getPersistenceContextInternal().getBatchFetchQueue()
+				.getBatchLoadableEntityIds( getLoadable(), id, domainBatchSize );
+	}
+
 	@Override
-	public final T load(
-			Object pkValue,
-			Object entityInstance,
-			LockOptions lockOptions,
-			Boolean readOnly,
-			SharedSessionContractImplementor session) {
-		if ( MULTI_KEY_LOAD_DEBUG_ENABLED ) {
-			MULTI_KEY_LOAD_LOGGER.debugf( "Batch loading entity `%s#%s`", getLoadable().getEntityName(), pkValue );
-		}
-
-		final Object[] idsToInitialize = resolveIdsToLoad( pkValue, session );
-		if ( MULTI_KEY_LOAD_DEBUG_ENABLED ) {
-			MULTI_KEY_LOAD_LOGGER.debugf( "Ids to batch-fetch initialize (`%s#%s`) %s", getLoadable().getEntityName(), pkValue, Arrays.toString(idsToInitialize) );
-		}
-
-		initializeEntities( idsToInitialize, pkValue, entityInstance, lockOptions, readOnly, session );
-
-		final EntityKey entityKey = session.generateEntityKey( pkValue, getLoadable().getEntityPersister() );
-		//noinspection unchecked
-		return (T) session.getPersistenceContext().getEntity( entityKey );
-	}
-
-	protected Object[] resolveIdsToLoad(Object pkValue, SharedSessionContractImplementor session) {
-		return session.getPersistenceContextInternal().getBatchFetchQueue().getBatchLoadableEntityIds(
-				getLoadable(),
-				pkValue,
-				domainBatchSize
-		);
-	}
-
 	protected void initializeEntities(
 			Object[] idsToInitialize,
 			Object pkValue,
@@ -156,6 +128,10 @@ public class EntityBatchLoaderInPredicate<T>
 			LockOptions lockOptions,
 			Boolean readOnly,
 			SharedSessionContractImplementor session) {
+		if ( MULTI_KEY_LOAD_DEBUG_ENABLED ) {
+			MULTI_KEY_LOAD_LOGGER.debugf( "Ids to batch-fetch initialize (`%s#%s`) %s",
+					getLoadable().getEntityName(), pkValue, Arrays.toString(idsToInitialize) );
+		}
 		final MultiKeyLoadChunker<Object> chunker = new MultiKeyLoadChunker<>(
 				sqlBatchSize,
 				getLoadable().getIdentifierMapping().getJdbcTypeCount(),
@@ -227,103 +203,103 @@ public class EntityBatchLoaderInPredicate<T>
 //		}
 	}
 
-	private void initializeChunk(
+//	private void initializeChunk(
-			Object[] idsToInitialize,
+//			Object[] idsToInitialize,
-			int start,
+//			int start,
-			Object pkValue,
+//			Object pkValue,
-			Object entityInstance,
+//			Object entityInstance,
-			LockOptions lockOptions,
+//			LockOptions lockOptions,
-			Boolean readOnly,
+//			Boolean readOnly,
-			SharedSessionContractImplementor session) {
+//			SharedSessionContractImplementor session) {
-		initializeChunk(
+//		initializeChunk(
-				idsToInitialize,
+//				idsToInitialize,
-				getLoadable(),
+//				getLoadable(),
-				start,
+//				start,
-				sqlBatchSize,
+//				sqlBatchSize,
-				jdbcParameters,
+//				jdbcParameters,
-				sqlAst,
+//				sqlAst,
-				jdbcSelectOperation,
+//				jdbcSelectOperation,
-				pkValue,
+//				pkValue,
-				entityInstance,
+//				entityInstance,
-				lockOptions,
+//				lockOptions,
-				readOnly,
+//				readOnly,
-				session
+//				session
-		);
+//		);
-	}
+//	}
 
-	private static void initializeChunk(
+//	private static void initializeChunk(
-			Object[] idsToInitialize,
+//			Object[] idsToInitialize,
-			EntityMappingType entityMapping,
+//			EntityMappingType entityMapping,
-			int startIndex,
+//			int startIndex,
-			int numberOfKeys,
+//			int numberOfKeys,
-			List<JdbcParameter> jdbcParameters,
+//			List<JdbcParameter> jdbcParameters,
-			SelectStatement sqlAst,
+//			SelectStatement sqlAst,
-			JdbcOperationQuerySelect jdbcSelectOperation,
+//			JdbcOperationQuerySelect jdbcSelectOperation,
-			Object pkValue,
+//			Object pkValue,
-			Object entityInstance,
+//			Object entityInstance,
-			LockOptions lockOptions,
+//			LockOptions lockOptions,
-			Boolean readOnly,
+//			Boolean readOnly,
-			SharedSessionContractImplementor session) {
+//			SharedSessionContractImplementor session) {
-		final BatchFetchQueue batchFetchQueue = session.getPersistenceContext().getBatchFetchQueue();
+//		final BatchFetchQueue batchFetchQueue = session.getPersistenceContext().getBatchFetchQueue();
-
+//
-		final int numberOfJdbcParameters = entityMapping.getIdentifierMapping().getJdbcTypeCount() * numberOfKeys;
+//		final int numberOfJdbcParameters = entityMapping.getIdentifierMapping().getJdbcTypeCount() * numberOfKeys;
-		final JdbcParameterBindings jdbcParameterBindings = new JdbcParameterBindingsImpl( numberOfJdbcParameters );
+//		final JdbcParameterBindings jdbcParameterBindings = new JdbcParameterBindingsImpl( numberOfJdbcParameters );
-
+//
-		final List<EntityKey> entityKeys = arrayList( numberOfKeys );
+//		final List<EntityKey> entityKeys = arrayList( numberOfKeys );
-		int bindCount = 0;
+//		int bindCount = 0;
-		for ( int i = 0; i < numberOfKeys; i++ ) {
+//		for ( int i = 0; i < numberOfKeys; i++ ) {
-			final int idPosition = i + startIndex;
+//			final int idPosition = i + startIndex;
-			final Object value;
+//			final Object value;
-			if ( idPosition >= idsToInitialize.length ) {
+//			if ( idPosition >= idsToInitialize.length ) {
-				value = null;
+//				value = null;
-			}
+//			}
-			else {
+//			else {
-				value = idsToInitialize[idPosition];
+//				value = idsToInitialize[idPosition];
-			}
+//			}
-			if ( value != null ) {
+//			if ( value != null ) {
-				entityKeys.add( session.generateEntityKey( value, entityMapping.getEntityPersister() ) );
+//				entityKeys.add( session.generateEntityKey( value, entityMapping.getEntityPersister() ) );
-			}
+//			}
-			bindCount += jdbcParameterBindings.registerParametersForEachJdbcValue(
+//			bindCount += jdbcParameterBindings.registerParametersForEachJdbcValue(
-					value,
+//					value,
-					bindCount,
+//					bindCount,
-					entityMapping.getIdentifierMapping(),
+//					entityMapping.getIdentifierMapping(),
-					jdbcParameters,
+//					jdbcParameters,
-					session
+//					session
-			);
+//			);
-		}
+//		}
-		assert bindCount == jdbcParameters.size();
+//		assert bindCount == jdbcParameters.size();
-
+//
-		if ( entityKeys.isEmpty() ) {
+//		if ( entityKeys.isEmpty() ) {
-			// there are no non-null keys in the chunk
+//			// there are no non-null keys in the chunk
-			return;
+//			return;
-		}
+//		}
-
+//
-		// Create a SubselectFetch.RegistrationHandler for handling any subselect fetches we encounter here
+//		// Create a SubselectFetch.RegistrationHandler for handling any subselect fetches we encounter here
-		final SubselectFetch.RegistrationHandler subSelectFetchableKeysHandler = SubselectFetch.createRegistrationHandler(
+//		final SubselectFetch.RegistrationHandler subSelectFetchableKeysHandler = SubselectFetch.createRegistrationHandler(
-				batchFetchQueue,
+//				batchFetchQueue,
-				sqlAst,
+//				sqlAst,
-				jdbcParameters,
+//				jdbcParameters,
-				jdbcParameterBindings
+//				jdbcParameterBindings
-		);
+//		);
-
+//
-		session.getJdbcServices().getJdbcSelectExecutor().list(
+//		session.getJdbcServices().getJdbcSelectExecutor().list(
-				jdbcSelectOperation,
+//				jdbcSelectOperation,
-				jdbcParameterBindings,
+//				jdbcParameterBindings,
-				new SingleIdExecutionContext(
+//				new SingleIdExecutionContext(
-						pkValue,
+//						pkValue,
-						entityInstance,
+//						entityInstance,
-						entityMapping.getRootEntityDescriptor(),
+//						entityMapping.getRootEntityDescriptor(),
-						readOnly,
+//						readOnly,
-						lockOptions,
+//						lockOptions,
-						subSelectFetchableKeysHandler,
+//						subSelectFetchableKeysHandler,
-						session
+//						session
-				),
+//				),
-				RowTransformerStandardImpl.instance(),
+//				RowTransformerStandardImpl.instance(),
-				ListResultsConsumer.UniqueSemantic.FILTER
+//				ListResultsConsumer.UniqueSemantic.FILTER
-		);
+//		);
-
+//
-		entityKeys.forEach( batchFetchQueue::removeBatchLoadableEntityKey );
+//		entityKeys.forEach( batchFetchQueue::removeBatchLoadableEntityKey );
-	}
+//	}
 
 	@Override
 	public String toString() {

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

@@ -16,6 +16,8 @@ import org.hibernate.type.descriptor.java.JavaType;
 import org.hibernate.type.descriptor.java.spi.JavaTypeRegistry;
 import org.hibernate.type.spi.TypeConfiguration;
 
+import java.util.Arrays;
+
 /**
  * @author Steve Ebersole
  */
@@ -57,6 +59,16 @@ public class MultiKeyLoadHelper {
 		);
 	}
 
+	static int countIds(Object[] ids) {
+		int count = 0;
+		for ( int i=1; i<ids.length; i++ ) {
+			if ( ids[i] != null ) {
+				count++;
+			}
+		}
+		return count;
+	}
+
 	static boolean hasSingleId(Object[] ids) {
 		for ( int i=1; i<ids.length; i++ ) {
 			if ( ids[i] != null ) {
@@ -65,4 +77,12 @@ public class MultiKeyLoadHelper {
 		}
 		return true;
 	}
+
+	static Object[] trimIdBatch(int length, Object[] keysToInitialize) {
+		int newLength = length;
+		while ( newLength>1 && keysToInitialize[newLength-1] == null ) {
+			newLength--;
+		}
+		return newLength < length ? Arrays.copyOf(keysToInitialize, newLength) : keysToInitialize;
+	}
 }

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

@@ -8,7 +8,6 @@ package org.hibernate.loader.ast.internal;
 
 import java.util.Map;
 
-import org.hibernate.dialect.Dialect;
 import org.hibernate.engine.spi.LoadQueryInfluencers;
 import org.hibernate.engine.spi.SessionFactoryImplementor;
 import org.hibernate.loader.ast.spi.BatchLoaderFactory;

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

@@ -10,7 +10,6 @@ import java.util.function.IntFunction;
 
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.metamodel.mapping.internal.MappingModelCreationProcess;
-import org.hibernate.metamodel.mapping.internal.VirtualIdEmbeddable;
 import org.hibernate.spi.NavigablePath;
 import org.hibernate.sql.ast.spi.SqlAstCreationState;
 import org.hibernate.sql.ast.tree.from.TableGroup;