HHH-17779 support for key-based pagination

lots of fixes / improvements

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

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

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

@@ -36,15 +36,45 @@ public KeyedResultList(List<R> resultList, KeyedPage<R> page, KeyedPage<R> nextP
 		this.nextPage = nextPage;
 	}
 
+	/**
+	 * The results on the current page.
+	 */
 	public List<R> getResultList() {
 		return resultList;
 	}
 
+	/**
+	 * The {@linkplain Page#getSize() size} and
+	 * approximate {@linkplain Page#getNumber()
+	 * page number} of the current page.
+	 */
 	public KeyedPage<R> getPage() {
 		return page;
 	}
 
+	/**
+	 * The specification of the next page of results,
+	 * if there are more results, or {@code null} if
+	 * it is known that there are no more results
+	 * after this page.
+	 */
 	public KeyedPage<R> getNextPage() {
 		return nextPage;
 	}
+
+	/**
+	 * @return {@code true} if this is known to be the
+	 *         last page of results.
+	 */
+	public boolean isLastPage() {
+		return nextPage == null;
+	}
+
+	/**
+	 * @return {@code true} if this is the first page
+	 *         of results.
+	 */
+	public boolean isFirstPage() {
+		return page.getPage().isFirst();
+	}
 }

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

@@ -90,6 +90,9 @@ public <R> KeyedPage<R> keyedBy(Order<? super R> keyDefinition) {
 		return new KeyedPage<>( List.of(keyDefinition), this );
 	}
 	public <R> KeyedPage<R> keyedBy(List<Order<? super R>> keyDefinition) {
+		if ( keyDefinition.isEmpty() )  {
+			throw new IllegalArgumentException("Key definition must not be empty");
+		}
 		return new KeyedPage<>( keyDefinition, this );
 	}
 }

hibernate-core/src/main/java/org/hibernate/query/sqm/NodeBuilder.java

@@ -32,6 +32,7 @@
 import org.hibernate.query.criteria.JpaSelection;
 import org.hibernate.query.criteria.JpaSimpleCase;
 import org.hibernate.query.criteria.JpaWindow;
+import org.hibernate.query.criteria.ValueHandlingMode;
 import org.hibernate.query.spi.QueryEngine;
 import org.hibernate.query.sqm.tree.delete.SqmDeleteStatement;
 import org.hibernate.query.sqm.tree.domain.SqmBagJoin;
@@ -85,6 +86,8 @@ public interface NodeBuilder extends HibernateCriteriaBuilder {
 
 	QueryEngine getQueryEngine();
 
+	void setCriteriaValueHandlingMode(ValueHandlingMode criteriaValueHandlingMode);
+
 	<R> SqmTuple<R> tuple(
 			Class<R> tupleType,
 			SqmExpression<?>... expressions);

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

@@ -6,12 +6,9 @@
  */
 package org.hibernate.query.sqm.internal;
 
-import jakarta.persistence.criteria.Expression;
-import org.hibernate.AssertionFailure;
 import org.hibernate.HibernateException;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.graph.spi.AppliedGraph;
-import org.hibernate.query.IllegalQueryOperationException;
 import org.hibernate.query.IllegalSelectQueryException;
 import org.hibernate.query.KeyedPage;
 import org.hibernate.query.KeyedResultList;
@@ -19,7 +16,7 @@
 import org.hibernate.query.Page;
 import org.hibernate.query.QueryLogging;
 import org.hibernate.query.SelectionQuery;
-import org.hibernate.query.criteria.JpaSelection;
+import org.hibernate.query.criteria.ValueHandlingMode;
 import org.hibernate.query.hql.internal.QuerySplitter;
 import org.hibernate.query.spi.AbstractSelectionQuery;
 import org.hibernate.query.spi.MutableQueryOptions;
@@ -27,21 +24,15 @@
 import org.hibernate.query.spi.SelectQueryPlan;
 import org.hibernate.query.sqm.NodeBuilder;
 import org.hibernate.query.sqm.tree.SqmStatement;
-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.predicate.SqmWhereClause;
-import org.hibernate.query.sqm.tree.select.SqmQuerySpec;
 import org.hibernate.query.sqm.tree.select.SqmSelectStatement;
 import org.hibernate.sql.results.internal.TupleMetadata;
 
 import java.util.ArrayList;
 import java.util.List;
 
-import static java.util.Arrays.asList;
 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.SqmUtil.sortSpecification;
 import static org.hibernate.query.sqm.tree.SqmCopyContext.noParamCopyContext;
 
@@ -143,98 +134,18 @@ public SelectionQuery<R> setPage(Page page) {
 		return this;
 	}
 
-	static class KeyedResult<R> {
+	private SqmSelectStatement<KeyedResult<R>> paginateQuery(
-		final R result;
-		final List<Comparable<?>> key;
-		public KeyedResult(R result, List<Comparable<?>> key) {
-			this.result = result;
-			this.key = key;
-		}
-
-		public R getResult() {
-			return result;
-		}
-
-		public List<Comparable<?>> getKey() {
-			return key;
-		}
-	}
-
-	private SqmSelectStatement<KeyedResult<R>> keyedQuery(
 			List<Order<? super R>> keyDefinition, List<Comparable<?>> keyValues) {
 		@SuppressWarnings("unchecked")
 		final SqmSelectStatement<KeyedResult<R>> sqm =
 				(SqmSelectStatement<KeyedResult<R>>)
 						getSqmSelectStatement().copy( noParamCopyContext() );
 		final NodeBuilder builder = sqm.nodeBuilder();
-		final SqmQuerySpec<?> querySpec = sqm.getQuerySpec();
+		//TODO: find a better way handle parameters
-		final List<? extends JpaSelection<?>> items = querySpec.getSelectClause().getSelectionItems();
+		builder.setCriteriaValueHandlingMode(ValueHandlingMode.INLINE);
-		if ( items.size() == 1 ) {
+		return paginate( keyDefinition, keyValues, sqm, builder );
-			final JpaSelection<?> selected = items.get(0);
-			if ( selected instanceof SqmRoot ) {
-				final List<SqmPath<?>> newItems = new ArrayList<>();
-				final SqmFrom<?,?> root = (SqmFrom<?,?>) selected;
-				SqmPredicate restriction = querySpec.getRestriction();
-				if ( restriction==null && keyValues != null ) {
-					restriction = builder.disjunction();
-				}
-				for ( int i = 0; i < keyDefinition.size(); i++ ) {
-					// ordering by an attribute of the returned entity
-					final Order<? super R> key = keyDefinition.get(i);
-					if ( keyValues != null ) {
-						@SuppressWarnings("rawtypes")
-						final Comparable keyValue = keyValues.get(i);
-						restriction = builder.or( restriction,
-								keyPredicate( root, key, keyValue, newItems, keyValues, builder ) );
-					}
-					newItems.add( root.get( key.getAttributeName() ) );
-				}
-				sqm.select( builder.construct((Class) KeyedResult.class,
-						asList(selected, builder.construct(List.class, newItems)) ) );
-				if ( keyValues != null ) {
-					sqm.where( restriction );
-				}
-				return sqm;
-			}
-			else {
-				throw new IllegalQueryOperationException("Select item was not an entity type");
-			}
-		}
-		else {
-			throw new IllegalQueryOperationException("Query has multiple items in the select list");
-		}
 	}
 
-	private <C extends Comparable<? super C>> SqmPredicate keyPredicate(
-			SqmFrom<?, ?> selected, Order<? super R> key, C keyValue,
-			List<SqmPath<?>> previousKeys, List<Comparable<?>> keyValues,
-			NodeBuilder builder) {
-		if ( !key.getEntityClass().isAssignableFrom( selected.getJavaType() ) ) {
-			throw new IllegalQueryOperationException("Select item was of wrong entity type");
-		}
-		final Expression<? extends C> path = selected.get( key.getAttributeName() );
-		// TODO: use a parameter here and create a binding for it
-//		@SuppressWarnings("unchecked")
-//		final Class<C> valueClass = (Class<C>) keyValue.getClass();
-//		final JpaParameterExpression<C> parameter = builder.parameter(valueClass);
-//		setParameter( parameter, keyValue );
-		SqmPredicate predicate;
-		switch ( key.getDirection() ) {
-			case ASCENDING:
-				predicate = builder.greaterThan( path, keyValue );
-				break;
-			case DESCENDING:
-				predicate = builder.lessThan( path, keyValue );
-				break;
-			default:
-				throw new AssertionFailure("Unrecognized key direction");
-		}
-		for ( int i = 0; i < previousKeys.size(); i++ ) {
-			final SqmPath keyPath = previousKeys.get(i);
-			predicate = builder.and( predicate, keyPath.equalTo( keyValues.get(i) ) );
-		}
-		return predicate;
-	}
 
 	@Override
 	public KeyedResultList<R> getKeyedResultList(KeyedPage<R> keyedPage) {
@@ -246,27 +157,34 @@ public KeyedResultList<R> getKeyedResultList(KeyedPage<R> keyedPage) {
 		final List<Comparable<?>> key = keyedPage.getKey();
 
 		setOrder( keyDefinition );
-		setMaxResults( page.getMaxResults() );
+		setMaxResults( page.getMaxResults() + 1 );
 		if ( key == null ) {
 			setFirstResult( page.getFirstResult() );
 		}
 
+//		getQueryOptions().setQueryPlanCachingEnabled( false );
 		final List<KeyedResult<R>> executed =
-				buildConcreteQueryPlan( keyedQuery( keyDefinition, key ), getQueryOptions() )
+				buildConcreteQueryPlan( paginateQuery( keyDefinition, key ), getQueryOptions() )
 						.performList(this);
-		final KeyedPage<R> nextPage = keyedPage.nextPage( getKeyOfLastResult( executed, key ) );
+
-		return new KeyedResultList<>( collectResultList( executed ), keyedPage, nextPage );
+		return new KeyedResultList<>( collectResults( executed, page.getSize() ),
+				keyedPage, getNextPage( keyedPage, executed ) );
 	}
 
-	private static <R> List<R> collectResultList(List<KeyedResult<R>> executed) {
+	private static <R> List<R> collectResults(List<KeyedResult<R>> executed, int pageSize) {
-		return executed.stream()
+		final int size = executed.size();
-				.map(KeyedResult::getResult)
+		final List<R> resultList = new ArrayList<>( size );
-				.collect(toList());
+		for (int i = 0; i < size && i < pageSize; i++) {
+			resultList.add( executed.get(i).getResult() );
+		}
+		return resultList;
 	}
 
-	private static <R> List<Comparable<?>> getKeyOfLastResult(
+	private static <R> KeyedPage<R> getNextPage(KeyedPage<R> keyedPage, List<KeyedResult<R>> executed) {
-			List<KeyedResult<R>> executed, List<Comparable<?>> key) {
+		final int pageSize = keyedPage.getPage().getSize();
-		return executed.isEmpty() ? key : executed.get(executed.size() - 1).getKey();
+		return executed.size() == pageSize + 1
+				? keyedPage.nextPage( executed.get(pageSize - 1).getKey() )
+				: null;
 	}
 
 	public abstract Class<R> getExpectedResultType();

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

@@ -0,0 +1,135 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * License: GNU Lesser General Public License (LGPL), version 2.1 or later
+ * See the lgpl.txt file in the root directory or http://www.gnu.org/licenses/lgpl-2.1.html
+ */
+package org.hibernate.query.sqm.internal;
+
+import jakarta.persistence.criteria.Expression;
+import org.hibernate.AssertionFailure;
+import org.hibernate.query.IllegalQueryOperationException;
+import org.hibernate.query.Order;
+import org.hibernate.query.SortDirection;
+import org.hibernate.query.criteria.JpaCompoundSelection;
+import org.hibernate.query.criteria.JpaSelection;
+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;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import static java.util.Arrays.asList;
+
+/**
+ * Manipulation of SQM query tree for key-based pagination.
+ *
+ * @author Gavin King
+ */
+public class KeyBasedPagination {
+
+	static <R> SqmSelectStatement<KeyedResult<R>> paginate(
+			List<Order<? super R>> keyDefinition, List<Comparable<?>> keyValues,
+			SqmSelectStatement<KeyedResult<R>> sqm, NodeBuilder builder) {
+		final SqmQuerySpec<?> querySpec = sqm.getQuerySpec();
+		final List<? extends JpaSelection<?>> items = querySpec.getSelectClause().getSelectionItems();
+		if ( items.size() == 1 ) {
+			final JpaSelection<?> selected = items.get(0);
+			if ( selected instanceof SqmRoot) {
+				final SqmFrom<?,?> root = (SqmFrom<?,?>) selected;
+				sqm.select( keySelection( keyDefinition, root, selected, builder ) );
+				if ( keyValues != null ) {
+					final SqmPredicate restriction = keyRestriction( keyDefinition, keyValues, root, builder );
+					final SqmPredicate queryWhere = querySpec.getRestriction();
+					sqm.where( queryWhere == null ? restriction : builder.and( queryWhere, restriction ) );
+				}
+				return sqm;
+			}
+			else {
+				throw new IllegalQueryOperationException("Select item was not an entity type");
+			}
+		}
+		else {
+			throw new IllegalQueryOperationException("Query has multiple items in the select list");
+		}
+	}
+
+	@SuppressWarnings({"rawtypes", "unchecked"})
+	private static <R> SqmPredicate keyRestriction(
+			List<Order<? super R>> keyDefinition,
+			List<Comparable<?>> keyValues,
+			SqmFrom<?, ?> root,
+			NodeBuilder builder) {
+		final List<SqmPath<?>> keyPaths = new ArrayList<>();
+		for ( Order<? super R> key : keyDefinition ) {
+			keyPaths.add( root.get( key.getAttributeName() ) );
+		}
+		SqmPredicate restriction = null;
+		for (int i = 0; i < keyDefinition.size(); i++ ) {
+			// ordering by an attribute of the returned entity
+			final SortDirection direction = keyDefinition.get(i).getDirection();
+			final SqmPath key = keyPaths.get(i);
+			final Comparable keyValue = keyValues.get(i);
+			final List<SqmPath<?>> previousKeys = keyPaths.subList(0, i);
+			final SqmPredicate predicate = keyPredicate( key, keyValue, direction, previousKeys, keyValues, builder );
+			restriction = restriction == null ? predicate : builder.or( restriction, predicate );
+		}
+		return restriction;
+	}
+
+	private static <R> JpaCompoundSelection<KeyedResult<R>> keySelection(
+			List<Order<? super R>> keyDefinition,
+			SqmFrom<?, ?> root, JpaSelection<?> selected,
+			NodeBuilder builder) {
+		final List<SqmPath<?>> items = new ArrayList<>();
+		for ( Order<? super R> key : keyDefinition ) {
+			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() ) );
+		}
+		return keyedResultConstructor( selected, builder, items );
+	}
+
+	private static <R> JpaCompoundSelection<KeyedResult<R>> keyedResultConstructor(
+			JpaSelection<?> selected, NodeBuilder builder, List<SqmPath<?>> newItems) {
+		@SuppressWarnings({"rawtypes", "unchecked"})
+		final Class<KeyedResult<R>> resultClass = (Class) KeyedResult.class;
+		return builder.construct( resultClass, asList( selected, builder.construct(List.class, newItems ) ) );
+	}
+
+	@SuppressWarnings({"rawtypes", "unchecked"})
+	private static <C extends Comparable<? super C>> SqmPredicate keyPredicate(
+			Expression<? extends C> key, C keyValue, SortDirection direction,
+			List<SqmPath<?>> previousKeys, List<Comparable<?>> keyValues,
+			NodeBuilder builder) {
+		// TODO: use a parameter here and create a binding for it
+//		@SuppressWarnings("unchecked")
+//		final Class<C> valueClass = (Class<C>) keyValue.getClass();
+//		final JpaParameterExpression<C> parameter = builder.parameter(valueClass);
+//		setParameter( parameter, keyValue );
+		SqmPredicate predicate;
+		switch ( direction ) {
+			case ASCENDING:
+				predicate = builder.greaterThan( key, keyValue );
+				break;
+			case DESCENDING:
+				predicate = builder.lessThan( key, keyValue );
+				break;
+			default:
+				throw new AssertionFailure("Unrecognized key direction");
+		}
+		for ( int i = 0; i < previousKeys.size(); i++ ) {
+			final SqmPath keyPath = previousKeys.get(i);
+			predicate = builder.and( predicate, keyPath.equalTo( keyValues.get(i) ) );
+		}
+		return predicate;
+	}
+
+}

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

@@ -0,0 +1,33 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * License: GNU Lesser General Public License (LGPL), version 2.1 or later
+ * See the lgpl.txt file in the root directory or http://www.gnu.org/licenses/lgpl-2.1.html
+ */
+package org.hibernate.query.sqm.internal;
+
+import java.util.List;
+
+/**
+ * Intermediate holder class for results of queries
+ * executed using key-based pagination.
+ *
+ * @author Gavin King
+ */
+class KeyedResult<R> {
+	final R result;
+	final List<Comparable<?>> key;
+
+	public KeyedResult(R result, List<Comparable<?>> key) {
+		this.result = result;
+		this.key = key;
+	}
+
+	public R getResult() {
+		return result;
+	}
+
+	public List<Comparable<?>> getKey() {
+		return key;
+	}
+}

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

@@ -197,7 +197,7 @@ public class SqmCriteriaNodeBuilder implements NodeBuilder, SqmCreationContext,
 	private final transient boolean jpaComplianceEnabled;
 	private final transient QueryEngine queryEngine;
 	private final transient Supplier<SessionFactoryImplementor> sessionFactory;
-	private final transient ValueHandlingMode criteriaValueHandlingMode;
+	private transient ValueHandlingMode criteriaValueHandlingMode;
 	private transient BasicType<Boolean> booleanType;
 	private transient BasicType<Integer> integerType;
 	private transient BasicType<Long> longType;
@@ -224,6 +224,10 @@ public SqmCriteriaNodeBuilder(
 		}
 	}
 
+	public void setCriteriaValueHandlingMode(ValueHandlingMode criteriaValueHandlingMode) {
+		this.criteriaValueHandlingMode = criteriaValueHandlingMode;
+	}
+
 	@Override
 	public JpaMetamodel getDomainModel() {
 		return getSessionFactory().getJpaMetamodel();