From ef54e06495fc70bd3cb85ca7f4d6b916c7a01768 Mon Sep 17 00:00:00 2001
From: Christian Beikov <christian.beikov@gmail.com>
Date: Wed, 25 Mar 2026 14:21:43 +0100
Subject: [PATCH] HHH-20283 Fix key based pagination issue after HHH-20274

---
 .../sqm/internal/KeyBasedPagination.java      | 37 ++++++++++++-------
 .../paging/keybased/KeyBasedPagingTest.java   | 15 ++++++++
 2 files changed, 38 insertions(+), 14 deletions(-)

diff --git a/hibernate-core/src/main/java/org/hibernate/query/sqm/internal/KeyBasedPagination.java b/hibernate-core/src/main/java/org/hibernate/query/sqm/internal/KeyBasedPagination.java
index ee7a6ee389dd..d9a7465f0bc1 100644
--- a/hibernate-core/src/main/java/org/hibernate/query/sqm/internal/KeyBasedPagination.java
+++ b/hibernate-core/src/main/java/org/hibernate/query/sqm/internal/KeyBasedPagination.java
@@ -14,6 +14,7 @@
 import org.hibernate.query.sqm.NodeBuilder;
 import org.hibernate.query.sqm.tree.domain.SqmPath;
 import org.hibernate.query.sqm.tree.from.SqmFrom;
+import org.hibernate.query.sqm.tree.from.SqmRoot;
 import org.hibernate.query.sqm.tree.predicate.SqmPredicate;
 import org.hibernate.query.sqm.tree.select.SqmQuerySpec;
 import org.hibernate.query.sqm.tree.select.SqmSelectStatement;
@@ -37,25 +38,33 @@ static <R> SqmSelectStatement<KeyedResult<R>> paginate(
 			SqmSelectStatement<KeyedResult<R>> statement, NodeBuilder builder) {
 		final SqmQuerySpec<?> querySpec = statement.getQuerySpec();
 		final List<? extends JpaSelection<?>> items = querySpec.getSelectClause().getSelectionItems();
-		if ( items.size() == 1 ) {
-			final JpaSelection<?> selected = items.get(0);
-			if ( selected instanceof SqmFrom<?, ?> root ) {
-				statement.orderBy( keyDefinition.stream().map( order -> sortSpecification( statement, order ) )
-						.collect( toList() ) );
-				statement.select( keySelection( keyDefinition, root, selected, builder ) );
-				if ( keyValues != null ) {
-					final SqmPredicate restriction = keyRestriction( keyDefinition, keyValues, root, builder );
-					final SqmPredicate queryWhere = querySpec.getRestriction();
-					statement.where( queryWhere == null ? restriction : builder.and( queryWhere, restriction ) );
+		final JpaSelection<?> selected = switch ( items.size() ) {
+			case 0 -> {
+				final List<SqmRoot<?>> sqmRoots = querySpec.getRootList();
+				if ( sqmRoots == null || sqmRoots.isEmpty() ) {
+					throw new IllegalArgumentException( "Query did not define any query roots" );
 				}
-				return statement;
+				if ( sqmRoots.size() != 1 ) {
+					throw new IllegalArgumentException( "Query has multiple query roots" );
+				}
+				yield sqmRoots.get( 0 );
 			}
-			else {
-				throw new IllegalQueryOperationException("Select item was not an entity type");
+			case 1 -> items.get( 0 );
+			default -> throw new IllegalQueryOperationException("Query has multiple items in the select list");
+		};
+		if ( selected instanceof SqmFrom<?, ?> root ) {
+			statement.orderBy( keyDefinition.stream().map( order -> sortSpecification( statement, order ) )
+					.collect( toList() ) );
+			statement.select( keySelection( keyDefinition, root, selected, builder ) );
+			if ( keyValues != null ) {
+				final SqmPredicate restriction = keyRestriction( keyDefinition, keyValues, root, builder );
+				final SqmPredicate queryWhere = querySpec.getRestriction();
+				statement.where( queryWhere == null ? restriction : builder.and( queryWhere, restriction ) );
 			}
+			return statement;
 		}
 		else {
-			throw new IllegalQueryOperationException("Query has multiple items in the select list");
+			throw new IllegalQueryOperationException("Select item was not an entity type");
 		}
 	}
 
diff --git a/hibernate-core/src/test/java/org/hibernate/orm/test/paging/keybased/KeyBasedPagingTest.java b/hibernate-core/src/test/java/org/hibernate/orm/test/paging/keybased/KeyBasedPagingTest.java
index 7267b313d29c..29da0aa984e1 100644
--- a/hibernate-core/src/test/java/org/hibernate/orm/test/paging/keybased/KeyBasedPagingTest.java
+++ b/hibernate-core/src/test/java/org/hibernate/orm/test/paging/keybased/KeyBasedPagingTest.java
@@ -11,7 +11,10 @@
 import org.hibernate.query.Page;
 import org.hibernate.query.SelectionQuery;
 
+import org.hibernate.query.criteria.HibernateCriteriaBuilder;
+import org.hibernate.query.criteria.JpaCriteriaQuery;
 import org.hibernate.testing.orm.junit.DomainModel;
+import org.hibernate.testing.orm.junit.Jira;
 import org.hibernate.testing.orm.junit.SessionFactory;
 import org.hibernate.testing.orm.junit.SessionFactoryScope;
 import org.junit.jupiter.api.Test;
@@ -153,6 +156,18 @@ public class KeyBasedPagingTest {
 		});
 	}
 
+	@Test
+	@Jira("https://hibernate.atlassian.net/browse/HHH-20283")
+	void testCriteria(SessionFactoryScope scope) {
+		scope.inSession(session -> {
+			HibernateCriteriaBuilder criteriaBuilder = session.getCriteriaBuilder();
+			JpaCriteriaQuery<Person> query = criteriaBuilder.createQuery( Person.class );
+			query.from( Person.class );
+			session.createSelectionQuery( query )
+					.getKeyedResultList( Page.first( 5 ).keyedBy( Order.asc( Person.class, "ssn" ) ) );
+		});
+	}
+
 	@Entity(name = "Person")
 	static class Person {
 		@Id