HHH-16651 use single id loader when there is one id in the batch

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

@@ -7,7 +7,6 @@
 package org.hibernate.loader.ast.internal;
 
 import java.lang.reflect.Array;
-import java.util.Collections;
 
 import org.hibernate.LockOptions;
 import org.hibernate.collection.spi.PersistentCollection;
@@ -33,6 +32,8 @@ import org.hibernate.sql.results.internal.RowTransformerStandardImpl;
 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;
 
@@ -49,6 +50,7 @@ public class CollectionBatchLoaderArrayParam
 	private final JdbcParameter jdbcParameter;
 	private final SelectStatement sqlSelect;
 	private final JdbcOperationQuerySelect jdbcSelectOperation;
+	private final CollectionLoaderSingleKey singleKeyLoader;
 
 	public CollectionBatchLoaderArrayParam(
 			int domainBatchSize,
@@ -68,7 +70,7 @@ public class CollectionBatchLoaderArrayParam
 		final SimpleForeignKeyDescriptor keyDescriptor = (SimpleForeignKeyDescriptor) getLoadable().getKeyDescriptor();
 
 		arrayElementType = keyDescriptor.getJavaType().getJavaTypeClass();
-		Class<?> arrayClass = Array.newInstance( arrayElementType, 0 ).getClass();
+		final Class<?> arrayClass = Array.newInstance( arrayElementType, 0 ).getClass();
 
 		final BasicType<?> arrayBasicType = getSessionFactory().getTypeConfiguration()
 				.getBasicTypeRegistry()
@@ -95,6 +97,8 @@ public class CollectionBatchLoaderArrayParam
 				.getSqlAstTranslatorFactory()
 				.buildSelectTranslator( getSessionFactory(), sqlSelect )
 				.translate( JdbcParameterBindings.NO_BINDINGS, QueryOptions.NONE );
+
+		singleKeyLoader = new CollectionLoaderSingleKey( attributeMapping, loadQueryInfluencers, sessionFactory );
 	}
 
 	@Override
@@ -103,13 +107,14 @@ public class CollectionBatchLoaderArrayParam
 			MULTI_KEY_LOAD_LOGGER.debugf( "Batch loading entity `%s#%s`", getLoadable().getNavigableRole().getFullPath(), key );
 		}
 
-		final Object[] keysToInitialize = resolveKeysToInitialize( key, session );
+		final Object[] keys = resolveKeysToInitialize( key, session );
-		initializeKeys( keysToInitialize, session );
 
-		for ( int i = 0; i < keysToInitialize.length; i++ ) {
+		if ( hasSingleId( keys ) ) {
-			finishInitializingKey( keysToInitialize[i], session );
+			return singleKeyLoader.load( key, session );
 		}
 
+		initializeKeys( keys, session );
+
 		final CollectionKey collectionKey = new CollectionKey( getLoadable().getCollectionDescriptor(), key );
 		return session.getPersistenceContext().getCollection( collectionKey );
 	}
@@ -138,7 +143,7 @@ public class CollectionBatchLoaderArrayParam
 		final SubselectFetch.RegistrationHandler subSelectFetchableKeysHandler = SubselectFetch.createRegistrationHandler(
 				session.getPersistenceContext().getBatchFetchQueue(),
 				sqlSelect,
-				Collections.singletonList( jdbcParameter ),
+				singletonList( jdbcParameter ),
 				jdbcParameterBindings
 		);
 
@@ -152,6 +157,10 @@ public class CollectionBatchLoaderArrayParam
 				RowTransformerStandardImpl.instance(),
 				ListResultsConsumer.UniqueSemantic.FILTER
 		);
+
+		for ( int i = 0; i < keysToInitialize.length; i++ ) {
+			finishInitializingKey( keysToInitialize[i], session );
+		}
 	}
 
 	public void prepare() {

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

@@ -10,7 +10,6 @@ import java.lang.reflect.Array;
 import java.util.Locale;
 
 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;
@@ -28,6 +27,8 @@ import org.hibernate.sql.exec.internal.JdbcParameterImpl;
 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.MultiKeyLoadLogging.MULTI_KEY_LOAD_DEBUG_ENABLED;
 import static org.hibernate.loader.ast.internal.MultiKeyLoadLogging.MULTI_KEY_LOAD_LOGGER;
 
@@ -48,6 +49,7 @@ public class EntityBatchLoaderArrayParam<T>
 	private final JdbcParameter jdbcParameter;
 	private final SelectStatement sqlAst;
 	private final JdbcOperationQuerySelect jdbcSelectOperation;
+	private final SingleIdEntityLoaderStandardImpl<T> singleIdLoader;
 
 
 	/**
@@ -100,7 +102,10 @@ public class EntityBatchLoaderArrayParam<T>
 				.getJdbcEnvironment()
 				.getSqlAstTranslatorFactory()
 				.buildSelectTranslator( sessionFactory, sqlAst )
-				.translate( JdbcParameterBindings.NO_BINDINGS, QueryOptions.NONE );	}
+				.translate( JdbcParameterBindings.NO_BINDINGS, QueryOptions.NONE );
+
+		singleIdLoader = new SingleIdEntityLoaderStandardImpl<>( entityDescriptor, sessionFactory );
+	}
 
 	@Override
 	public int getDomainBatchSize() {
@@ -119,6 +124,11 @@ public class EntityBatchLoaderArrayParam<T>
 		}
 
 		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() );
@@ -167,7 +177,7 @@ public class EntityBatchLoaderArrayParam<T>
 				continue;
 			}
 			// found or not, remove the key from the batch-fetch queue
-			BatchFetchQueueHelper.removeBatchLoadableEntityKey( id, getLoadable(), session );
+			removeBatchLoadableEntityKey( id, getLoadable(), session );
 		}
 	}
 

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

@@ -56,4 +56,13 @@ public class MultiKeyLoadHelper {
 				typeConfiguration.getCurrentBaseSqlTypeIndicators()
 		);
 	}
+
+	static boolean hasSingleId(Object[] ids) {
+		for ( int i=1; i<ids.length; i++ ) {
+			if ( ids[i] != null ) {
+				return false;
+			}
+		}
+		return true;
+	}
 }