HHH-16750 ClassCastException when an Entity with an ElementCollection has an EmbeddableId with just one field and Batch is enabled

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

@@ -35,7 +35,7 @@ public abstract class AbstractCollectionBatchLoader implements CollectionBatchLo
 
 	private final int keyJdbcCount;
 
-	private final CollectionLoaderSingleKey singleKeyLoader;
+	final CollectionLoaderSingleKey singleKeyLoader;
 
 	public AbstractCollectionBatchLoader(
 			int domainBatchSize,
@@ -91,10 +91,14 @@ public abstract class AbstractCollectionBatchLoader implements CollectionBatchLo
 
 		initializeKeys( key, keys, session );
 
+		finishInitializingKeys( keys, session );
+
 		final CollectionKey collectionKey = new CollectionKey( getLoadable().getCollectionDescriptor(), key );
 		return session.getPersistenceContext().getCollection( collectionKey );
 	}
 
+	abstract void finishInitializingKeys(Object[] key, SharedSessionContractImplementor session);
+
 	protected void finishInitializingKey(Object key, SharedSessionContractImplementor session) {
 		if ( key == null ) {
 			return;

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

@@ -9,6 +9,8 @@ package org.hibernate.loader.ast.internal;
 import java.lang.reflect.Array;
 
 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;
@@ -18,7 +20,6 @@ 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;
 import org.hibernate.query.spi.QueryOptions;
 import org.hibernate.sql.ast.tree.expression.JdbcParameter;
 import org.hibernate.sql.ast.tree.select.SelectStatement;
@@ -32,6 +33,8 @@ import org.hibernate.sql.results.internal.RowTransformerStandardImpl;
 import org.hibernate.sql.results.spi.ListResultsConsumer;
 import org.hibernate.type.BasicType;
 
+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;
 
@@ -67,15 +70,12 @@ public class CollectionBatchLoaderArrayParam
 		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,
-				simpleKeyDescriptor.getJdbcMapping(),
+				keyDescriptor.getSingleJdbcMapping(),
 				arrayClass,
 				getSessionFactory()
 		);
@@ -83,7 +83,7 @@ public class CollectionBatchLoaderArrayParam
 		jdbcParameter = new JdbcParameterImpl( arrayJdbcMapping );
 		sqlSelect = LoaderSelectBuilder.createSelectBySingleArrayParameter(
 				getLoadable(),
-				simpleKeyDescriptor.getKeyPart(),
+				keyDescriptor.getKeyPart(),
 				getInfluencers(),
 				LockOptions.NONE,
 				jdbcParameter,
@@ -97,6 +97,72 @@ public class CollectionBatchLoaderArrayParam
 				.translate( JdbcParameterBindings.NO_BINDINGS, QueryOptions.NONE );
 	}
 
+	@Override
+	public PersistentCollection<?> load(Object keyBeingLoaded, SharedSessionContractImplementor session) {
+		final ForeignKeyDescriptor keyDescriptor = getLoadable().getKeyDescriptor();
+		if ( keyDescriptor.isEmbedded() ) {
+			assert keyDescriptor.getJdbcTypeCount() == 1;
+			return loadEmbeddable( keyBeingLoaded, session, keyDescriptor );
+		}
+		else {
+			return super.load( keyBeingLoaded, session );
+		}
+
+	}
+
+	private PersistentCollection<?> loadEmbeddable(
+			Object keyBeingLoaded,
+			SharedSessionContractImplementor session,
+			ForeignKeyDescriptor keyDescriptor) {
+		if ( MULTI_KEY_LOAD_DEBUG_ENABLED ) {
+			MULTI_KEY_LOAD_LOGGER.debugf(
+					"Batch fetching collection: %s.%s",
+					getLoadable().getNavigableRole().getFullPath(), keyBeingLoaded
+			);
+		}
+
+		final int length = getDomainBatchSize();
+		final Object[] keysToInitialize = (Object[]) Array.newInstance(
+				keyDescriptor.getSingleJdbcMapping().getJdbcJavaType().getJavaTypeClass(),
+				length
+		);
+		final Object[] embeddedKeys = (Object[]) Array.newInstance(
+				keyDescriptor.getJavaType().getJavaTypeClass(),
+				length
+		);
+		session.getPersistenceContextInternal().getBatchFetchQueue()
+				.collectBatchLoadableCollectionKeys(
+						length,
+						(index, key) ->
+								keyDescriptor.forEachJdbcValue( key, (i, value, jdbcMapping) -> {
+									keysToInitialize[index] = value;
+									embeddedKeys[index] = key;
+								}, session )
+						,
+						keyBeingLoaded,
+						getLoadable()
+				);
+		// now trim down the array to the number of keys we found
+		final Object[] keys = trimIdBatch( length, keysToInitialize );
+
+		if ( hasSingleId( keys ) ) {
+			return singleKeyLoader.load( keyBeingLoaded, session );
+		}
+
+		initializeKeys( keyBeingLoaded, keys, session );
+
+		for ( Object initializedKey : embeddedKeys ) {
+			if ( initializedKey != null ) {
+				finishInitializingKey( initializedKey, session );
+			}
+		}
+		final CollectionKey collectionKey = new CollectionKey(
+				getLoadable().getCollectionDescriptor(),
+				keysToInitialize
+		);
+		return session.getPersistenceContext().getCollection( collectionKey );
+	}
+
 	@Override
 	void initializeKeys(Object key, Object[] keysToInitialize, SharedSessionContractImplementor session) {
 		if ( MULTI_KEY_LOAD_DEBUG_ENABLED ) {
@@ -134,9 +200,36 @@ public class CollectionBatchLoaderArrayParam
 				RowTransformerStandardImpl.instance(),
 				ListResultsConsumer.UniqueSemantic.FILTER
 		);
+	}
 
-		for ( Object initializedKey : keysToInitialize ) {
+	@Override
+	void finishInitializingKeys(Object[] keys, SharedSessionContractImplementor session) {
+		for ( Object initializedKey : keys ) {
 			finishInitializingKey( initializedKey, session );
 		}
 	}
+
+	@Override
+	Object[] resolveKeysToInitialize(Object keyBeingLoaded, SharedSessionContractImplementor session) {
+		final ForeignKeyDescriptor keyDescriptor = getLoadable().getKeyDescriptor();
+		if( keyDescriptor.isEmbedded()){
+			assert keyDescriptor.getJdbcTypeCount() == 1;
+			final int length = getDomainBatchSize();
+			final Object[] keysToInitialize = (Object[]) Array.newInstance( keyDescriptor.getSingleJdbcMapping().getJdbcJavaType().getJavaTypeClass(), length );
+			session.getPersistenceContextInternal().getBatchFetchQueue()
+					.collectBatchLoadableCollectionKeys(
+							length,
+							(index, key) ->
+								keyDescriptor.forEachJdbcValue( key, (i, value, jdbcMapping) -> {
+									keysToInitialize[index] = value;
+								}, session )
+							,
+							keyBeingLoaded,
+							getLoadable()
+					);
+			// now trim down the array to the number of keys we found
+			return trimIdBatch( length, keysToInitialize );
+		}
+		return super.resolveKeysToInitialize( keyBeingLoaded, session );
+	}
 }

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

@@ -152,4 +152,8 @@ public class CollectionBatchLoaderInPredicate
 		);
 	}
 
+	@Override
+	void finishInitializingKeys(Object[] key, SharedSessionContractImplementor session) {
+		// do nothing
+	}
 }

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

@@ -40,6 +40,7 @@ import org.hibernate.metamodel.mapping.NonAggregatedIdentifierMapping;
 import org.hibernate.metamodel.mapping.PluralAttributeMapping;
 import org.hibernate.metamodel.mapping.Restrictable;
 import org.hibernate.metamodel.mapping.SelectableMapping;
+import org.hibernate.metamodel.mapping.ValuedModelPart;
 import org.hibernate.metamodel.mapping.internal.EmbeddedAttributeMapping;
 import org.hibernate.metamodel.mapping.internal.SimpleForeignKeyDescriptor;
 import org.hibernate.metamodel.mapping.internal.ToOneAttributeMapping;
@@ -142,7 +143,7 @@ public class LoaderSelectBuilder {
 	 */
 	public static SelectStatement createSelectBySingleArrayParameter(
 			Loadable loadable,
-			BasicValuedModelPart restrictedPart,
+			ValuedModelPart restrictedPart,
 			LoadQueryInfluencers influencers,
 			LockOptions lockOptions,
 			JdbcParameter jdbcArrayParameter,
@@ -202,9 +203,10 @@ public class LoaderSelectBuilder {
 			QuerySpec rootQuerySpec,
 			NavigablePath rootNavigablePath,
 			TableGroup rootTableGroup,
-			BasicValuedModelPart restrictedPart,
+			ValuedModelPart restrictedPart,
 			JdbcParameter jdbcArrayParameter,
 			LoaderSqlAstCreationState sqlAstCreationState) {
+		assert restrictedPart.getJdbcTypeCount() == 1;
 		final SqlExpressionResolver sqlExpressionResolver = sqlAstCreationState.getSqlExpressionResolver();
 		final SelectableMapping restrictedPartMapping = restrictedPart.getSelectable( 0 );
 		final NavigablePath restrictionPath = rootNavigablePath.append( restrictedPart.getNavigableRole().getNavigableName() );

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

@@ -198,4 +198,7 @@ public interface ForeignKeyDescriptor extends VirtualModelPart, ValuedModelPart
 		ValuedModelPart getModelPart();
 
 	}
+
+	boolean isEmbedded();
+
 }

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

@@ -670,4 +670,9 @@ public class EmbeddedForeignKeyDescriptor implements ForeignKeyDescriptor {
 	public boolean hasPartitionedSelectionMapping() {
 		return keySide.getModelPart().hasPartitionedSelectionMapping();
 	}
+
+	@Override
+	public boolean isEmbedded() {
+		return true;
+	}
 }

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

@@ -719,4 +719,9 @@ public class SimpleForeignKeyDescriptor implements ForeignKeyDescriptor, BasicVa
 				targetSide.getModelPart().getSelectionExpression()
 		);
 	}
+
+	@Override
+	public boolean isEmbedded() {
+		return false;
+	}
 }