HHH-16043 Correct single ID and init empty for batch collection loading

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

@@ -14,8 +14,10 @@ import org.hibernate.LockOptions;
 import org.hibernate.collection.spi.PersistentCollection;
 import org.hibernate.engine.jdbc.env.spi.JdbcEnvironment;
 import org.hibernate.engine.jdbc.spi.JdbcServices;
+import org.hibernate.engine.spi.CollectionEntry;
 import org.hibernate.engine.spi.CollectionKey;
 import org.hibernate.engine.spi.LoadQueryInfluencers;
+import org.hibernate.engine.spi.PersistenceContext;
 import org.hibernate.engine.spi.SessionFactoryImplementor;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.engine.spi.SubselectFetch;
@@ -30,6 +32,7 @@ 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.ResultsHelper;
 import org.hibernate.sql.results.internal.RowTransformerStandardImpl;
 import org.hibernate.sql.results.spi.ListResultsConsumer;
 
@@ -97,7 +100,7 @@ public class CollectionLoaderBatchKey implements CollectionLoader {
 					null,
 					attributeMapping.getKeyDescriptor(),
 					null,
-					batchSize,
+					1,
 					session.getLoadQueryInfluencers(),
 					LockOptions.NONE,
 					jdbcParameters::add,
@@ -201,6 +204,11 @@ public class CollectionLoaderBatchKey implements CollectionLoader {
 					ListResultsConsumer.UniqueSemantic.FILTER
 			);
 
+			for ( int i = smallBatchStart; i < smallBatchStart + smallBatchLength; i++ ) {
+				// collections that were not initialized here should be empty
+				finishLoadingCollection( batchIds[i], session );
+			}
+
 			// prepare for the next round...
 			smallBatchStart += smallBatchLength;
 			if ( smallBatchStart >= numberOfIds ) {
@@ -211,4 +219,20 @@ public class CollectionLoaderBatchKey implements CollectionLoader {
 		}
 	}
 
+	private void finishLoadingCollection(Object key, SharedSessionContractImplementor session) {
+		final PersistenceContext persistenceContext = session.getPersistenceContext();
+		final CollectionKey collectionKey = new CollectionKey( attributeMapping.getCollectionDescriptor(), key );
+		final PersistentCollection<?> collection = persistenceContext.getCollection( collectionKey );
+		if ( !collection.wasInitialized() ) {
+			final CollectionEntry entry = persistenceContext.getCollectionEntry( collection );
+			collection.initializeEmptyCollection( entry.getLoadedPersister() );
+			ResultsHelper.finalizeCollectionLoading(
+					persistenceContext,
+					entry.getLoadedPersister(),
+					collection,
+					collectionKey,
+					true
+			);
+		}
+	}
 }

hibernate-core/src/test/java/org/hibernate/orm/test/batchfetch/NestedLazyManyToOneTest.java

@@ -44,11 +44,11 @@ import static org.junit.jupiter.api.Assertions.assertEquals;
 		}
 )
 @SessionFactory(useCollectingStatementInspector = true)
-@ServiceRegistry(settings = {
+@ServiceRegistry(settings = @Setting(name = AvailableSettings.DEFAULT_BATCH_FETCH_SIZE, value = "5"))
-		@Setting(name = AvailableSettings.DEFAULT_BATCH_FETCH_SIZE, value = "5")
-})
 @JiraKey("HHH-16043")
 public class NestedLazyManyToOneTest {
+	private static final String QUESTION_MARK = "\\?";
+
 	@BeforeAll
 	public void prepareData(SessionFactoryScope scope) {
 		final Entity1 entity1 = new Entity1();
@@ -87,10 +87,9 @@ public class NestedLazyManyToOneTest {
 		scope.inTransaction( session -> {
 			Entity1 fromDb = session.find( Entity1.class, "0" );
 			Set<Entity2> children = fromDb.getChildren();
-
 			assertEquals( 8, children.size() );
-			statementInspector.assertExecutedCount( 2 );
+			statementInspector.assertExecutedCount( 2 ); // 1 for Entity1, 1 for Entity2
-			statementInspector.assertNumberOfOccurrenceInQueryNoSpace( 1, "\\?", 1 );
+			statementInspector.assertNumberOfOccurrenceInQueryNoSpace( 1, QUESTION_MARK, 1 );
 		} );
 	}
 
@@ -122,7 +121,8 @@ public class NestedLazyManyToOneTest {
 
 			assertEquals( 8, entity1.getChildren().size() );
 			statementInspector.assertExecutedCount( 3 ); // 1 for Entity1, 1 for Entity2, 1 for Entity3
-			statementInspector.assertNumberOfOccurrenceInQueryNoSpace( 2, "\\?", 5 );
+			statementInspector.assertNumberOfOccurrenceInQueryNoSpace( 1, QUESTION_MARK, 1 );
+			statementInspector.assertNumberOfOccurrenceInQueryNoSpace( 2, QUESTION_MARK, 5 );
 		} );
 	}
 
@@ -145,8 +145,9 @@ public class NestedLazyManyToOneTest {
 
 			assertEquals( 8, entity1.getChildren().size() );
 			statementInspector.assertExecutedCount( 4 ); // 1 for Entity1, 1 for Entity2, 2 for Entity3
-			statementInspector.assertNumberOfOccurrenceInQueryNoSpace( 2, "\\?", 5 );
+			statementInspector.assertNumberOfOccurrenceInQueryNoSpace( 1, QUESTION_MARK, 1 );
-			statementInspector.assertNumberOfOccurrenceInQueryNoSpace( 3, "\\?", 3 );
+			statementInspector.assertNumberOfOccurrenceInQueryNoSpace( 2, QUESTION_MARK, 5 );
+			statementInspector.assertNumberOfOccurrenceInQueryNoSpace( 3, QUESTION_MARK, 3 );
 		} );
 	}