HHH-17779 misc improvements to key-based pagination

hibernate-core/src/main/java/org/hibernate/query/KeyedPage.java

@@ -105,4 +105,9 @@ public List<Comparable<?>> getKey() {
 	public KeyedPage<R> nextPage(List<Comparable<?>> keyOfLastResult) {
 		return new KeyedPage<>( keyDefinition, page.next(), keyOfLastResult );
 	}
+
+	@Internal
+	public KeyedPage<R> withKey(List<Comparable<?>> keyOfLastResult) {
+		return new KeyedPage<>( keyDefinition, page, keyOfLastResult );
+	}
 }

hibernate-core/src/main/java/org/hibernate/query/KeyedResultList.java

@@ -27,11 +27,13 @@
 @Incubating
 public class KeyedResultList<R> {
 	private final List<R> resultList;
+	private final List<List<?>> keyList;
 	private final KeyedPage<R> page;
 	private final KeyedPage<R> nextPage;
 
-	public KeyedResultList(List<R> resultList, KeyedPage<R> page, KeyedPage<R> nextPage) {
+	public KeyedResultList(List<R> resultList, List<List<?>> keyList, KeyedPage<R> page, KeyedPage<R> nextPage) {
 		this.resultList = resultList;
+		this.keyList = keyList;
 		this.page = page;
 		this.nextPage = nextPage;
 	}
@@ -43,6 +45,13 @@ public List<R> getResultList() {
 		return resultList;
 	}
 
+	/**
+	 * The keys of the results, in order.
+	 */
+	public List<List<?>> getKeyList() {
+		return keyList;
+	}
+
 	/**
 	 * The {@linkplain Page#getSize() size} and
 	 * approximate {@linkplain Page#getNumber()

hibernate-core/src/main/java/org/hibernate/query/sqm/internal/AbstractSqmSelectionQuery.java

@@ -33,6 +33,8 @@
 import static java.util.stream.Collectors.toList;
 import static org.hibernate.cfg.QuerySettings.FAIL_ON_PAGINATION_OVER_COLLECTION_FETCH;
 import static org.hibernate.query.sqm.internal.KeyBasedPagination.paginate;
+import static org.hibernate.query.sqm.internal.KeyedResult.collectKeys;
+import static org.hibernate.query.sqm.internal.KeyedResult.collectResults;
 import static org.hibernate.query.sqm.internal.SqmUtil.sortSpecification;
 import static org.hibernate.query.sqm.tree.SqmCopyContext.noParamCopyContext;
 
@@ -150,7 +152,7 @@ private SqmSelectStatement<KeyedResult<R>> paginateQuery(
 	@Override
 	public KeyedResultList<R> getKeyedResultList(KeyedPage<R> keyedPage) {
 		if ( keyedPage == null ) {
-			throw new IllegalStateException( "KeyedPage was not set" );
+			throw new IllegalStateException( "KeyedPage was not null" );
 		}
 		final Page page = keyedPage.getPage();
 		final List<Order<? super R>> keyDefinition = keyedPage.getKeyDefinition();
@@ -163,21 +165,16 @@ public KeyedResultList<R> getKeyedResultList(KeyedPage<R> keyedPage) {
 		}
 
 //		getQueryOptions().setQueryPlanCachingEnabled( false );
-		final List<KeyedResult<R>> executed =
+		final List<KeyedResult<R>> results =
 				buildConcreteQueryPlan( paginateQuery( keyDefinition, key ), getQueryOptions() )
 						.performList(this);
 
-		return new KeyedResultList<>( collectResults( executed, page.getSize() ),
+		return new KeyedResultList<>(
-				keyedPage, getNextPage( keyedPage, executed ) );
+				collectResults( results, page.getSize() ),
-	}
+				collectKeys( results, page.getSize() ),
-
+				keyedPage,
-	private static <R> List<R> collectResults(List<KeyedResult<R>> executed, int pageSize) {
+				getNextPage( keyedPage, results )
-		final int size = executed.size();
+		);
-		final List<R> resultList = new ArrayList<>( size );
-		for (int i = 0; i < size && i < pageSize; i++) {
-			resultList.add( executed.get(i).getResult() );
-		}
-		return resultList;
 	}
 
 	private static <R> KeyedPage<R> getNextPage(KeyedPage<R> keyedPage, List<KeyedResult<R>> executed) {

hibernate-core/src/main/java/org/hibernate/query/sqm/internal/KeyBasedPagination.java

@@ -88,11 +88,16 @@ private static <R> JpaCompoundSelection<KeyedResult<R>> keySelection(
 			NodeBuilder builder) {
 		final List<SqmPath<?>> items = new ArrayList<>();
 		for ( Order<? super R> key : keyDefinition ) {
-			if ( !key.getEntityClass().isAssignableFrom( selected.getJavaType() ) ) {
+			if ( key.getEntityClass() == null ) {
-				throw new IllegalQueryOperationException("Select item was of wrong entity type");
+				throw new IllegalQueryOperationException("Key-based pagination based on select list items is not yet supported");
+			}
+			else {
+				if ( !key.getEntityClass().isAssignableFrom( selected.getJavaType() ) ) {
+					throw new IllegalQueryOperationException("Select item was of wrong entity type");
+				}
+				// ordering by an attribute of the returned entity
+				items.add( root.get( key.getAttributeName() ) );
 			}
-			// ordering by an attribute of the returned entity
-			items.add( root.get( key.getAttributeName() ) );
 		}
 		return keyedResultConstructor( selected, builder, items );
 	}

hibernate-core/src/main/java/org/hibernate/query/sqm/internal/KeyedResult.java

@@ -6,6 +6,7 @@
  */
 package org.hibernate.query.sqm.internal;
 
+import java.util.ArrayList;
 import java.util.List;
 
 /**
@@ -30,4 +31,26 @@ public R getResult() {
 	public List<Comparable<?>> getKey() {
 		return key;
 	}
+
+	static <R> List<R> collectResults(List<KeyedResult<R>> executed, int pageSize) {
+		final int size = executed.size();
+		final List<R> resultList = new ArrayList<>( size );
+		for (int i = 0; i < size && i < pageSize; i++) {
+			resultList.add( executed.get(i).getResult() );
+		}
+		return resultList;
+	}
+
+	static List<List<?>> collectKeys(List<? extends KeyedResult<?>> executed, int pageSize) {
+		final int size = executed.size();
+		final List<List<?>> resultList = new ArrayList<>( size );
+		for ( int i = 0; i < size && i < pageSize; i++ ) {
+			final List<Comparable<?>> key = executed.get(i).getKey();
+			if ( key == null ) {
+				throw new IllegalArgumentException("Null keys in key-based pagination are not yet supported");
+			}
+			resultList.add( key );
+		}
+		return resultList;
+	}
 }