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/CollectionBatchLoaderArrayParam.java

@@ -17,6 +17,7 @@ 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;
@@ -65,7 +66,7 @@ public class CollectionBatchLoaderArrayParam
 			);
 		}
 
-		final SimpleForeignKeyDescriptor keyDescriptor = (SimpleForeignKeyDescriptor) getLoadable().getKeyDescriptor();
+		final ForeignKeyDescriptor keyDescriptor = getLoadable().getKeyDescriptor();
 
 		arrayElementType = keyDescriptor.getJavaType().getJavaTypeClass();
 		Class<?> arrayClass = Array.newInstance( arrayElementType, 0 ).getClass();
@@ -75,7 +76,7 @@ public class CollectionBatchLoaderArrayParam
 				.getRegisteredType( arrayClass );
 		arrayJdbcMapping = MultiKeyLoadHelper.resolveArrayJdbcMapping(
 				arrayBasicType,
-				keyDescriptor.getJdbcMapping(),
+				keyDescriptor.getSingleJdbcMapping(),
 				arrayClass,
 				getSessionFactory()
 		);
@@ -102,15 +103,62 @@ public class CollectionBatchLoaderArrayParam
 		if ( MULTI_KEY_LOAD_DEBUG_ENABLED ) {
 			MULTI_KEY_LOAD_LOGGER.debugf( "Batch loading entity `%s#%s`", getLoadable().getNavigableRole().getFullPath(), key );
 		}
+		final ForeignKeyDescriptor keyDescriptor = getLoadable().getKeyDescriptor();
+		if ( keyDescriptor.isEmbedded() ) {
+			assert keyDescriptor.getJdbcTypeCount() == 1;
+			return loadEmbeddable( key, session, keyDescriptor );
+		}
+		else {
+
+			final Object[] keysToInitialize = resolveKeysToInitialize( key, session );
+			initializeKeys( keysToInitialize, session );
+
+			for ( int i = 0; i < keysToInitialize.length; i++ ) {
+				finishInitializingKey( keysToInitialize[i], session );
+			}
+
+			final CollectionKey collectionKey = new CollectionKey( getLoadable().getCollectionDescriptor(), key );
+			return session.getPersistenceContext().getCollection( collectionKey );
+		}
+	}
+
+	private PersistentCollection<?> loadEmbeddable(
+			Object keyBeingLoaded,
+			SharedSessionContractImplementor session,
+			ForeignKeyDescriptor keyDescriptor) {
+
+		final int length = getDomainBatchSize();
+		final Object[] keysToInitialize = (Object[]) Array.newInstance(
+				keyDescriptor.getSingleJdbcMapping().getJdbcJavaType().getJavaTypeClass(),
+				length
+		);
+		final Object[] embeddedKeys = (Object[]) Array.newInstance(
+				arrayElementType,
+				length
+		);
+		session.getPersistenceContextInternal().getBatchFetchQueue()
+				.collectBatchLoadableCollectionKeys(
+						length,
+						(index, key) ->
+								keyDescriptor.forEachJdbcValue( key, (i, value, jdbcMapping) -> {
+									keysToInitialize[index] = value;
+									embeddedKeys[index] = key;
+								}, session )
+						,
+						keyBeingLoaded,
+						getLoadable()
+				);
 
-		final Object[] keysToInitialize = resolveKeysToInitialize( key, session );
 		initializeKeys( keysToInitialize, session );
 
-		for ( int i = 0; i < keysToInitialize.length; i++ ) {
+		for ( Object initializedKey : embeddedKeys ) {
-			finishInitializingKey( keysToInitialize[i], session );
+				finishInitializingKey( initializedKey, session );
 		}
 
-		final CollectionKey collectionKey = new CollectionKey( getLoadable().getCollectionDescriptor(), key );
+		final CollectionKey collectionKey = new CollectionKey(
+				getLoadable().getCollectionDescriptor(),
+				keysToInitialize
+		);
 		return session.getPersistenceContext().getCollection( collectionKey );
 	}
 

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,7 +203,7 @@ public class LoaderSelectBuilder {
 			QuerySpec rootQuerySpec,
 			NavigablePath rootNavigablePath,
 			TableGroup rootTableGroup,
-			BasicValuedModelPart restrictedPart,
+			ValuedModelPart restrictedPart,
 			JdbcParameter jdbcArrayParameter,
 			LoaderSqlAstCreationState sqlAstCreationState) {
 		final SqlExpressionResolver sqlExpressionResolver = sqlAstCreationState.getSqlExpressionResolver();

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

@@ -199,4 +199,6 @@ 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

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