fix: more optimal pg query

Signed-off-by: Artemiy Degtyarev <[email protected]>
Signed-off-by: Artemiy Chereshnevvv <[email protected]>

Author: Artemiy Degtyarev

spring-data-jdbc/src/main/java/org/springframework/data/jdbc/repository/query/StatementFactory.java

@@ -238,47 +238,46 @@ Sort applyCriteriaOrderBy(Sort sort, @Nullable ScrollPosition scrollPosition) {
 		}
 
 		Criteria applyScrollCriteria(Criteria criteria, @Nullable ScrollPosition position, Sort sort) {
-			if (!(position instanceof KeysetScrollPosition) || position.isInitial() || ((KeysetScrollPosition) position).getKeys().isEmpty())
+			if (!(position instanceof KeysetScrollPosition keyset)
+				|| position.isInitial()
+				|| keyset.getKeys().isEmpty()) {
 				return criteria;
+			}
 
-			Map<String, Object> keys = ((KeysetScrollPosition) position).getKeys();
+			Map<String, Object> keys = keyset.getKeys();
 			List<String> columns = new ArrayList<>(keys.keySet());
 			List<Object> values = new ArrayList<>(keys.values());
 
-			Criteria result = null;
+			Map<String, Sort.Direction> directions =
+				sort.stream().collect(Collectors.toMap(Sort.Order::getProperty, Sort.Order::getDirection));
 
-			for (int i = 0; i < keys.size(); i++) {
-				Map<String, Sort.Direction> orders =
-					sort.get().collect(Collectors.toMap(Sort.Order::getProperty, Sort.Order::getDirection));
+			String primary = columns.get(0);
+			Object primaryValue = values.get(0);
 
-				Criteria.CriteriaStep orCriteriaStep = Criteria.where(columns.get(i));
-				Criteria orCriteria;
+			Criteria primaryCompare;
+			Sort.Direction dir = directions.getOrDefault(primary, Sort.DEFAULT_DIRECTION);
 
-				Sort.Direction defaultDirection = orders.values().stream().findFirst().orElse(Sort.DEFAULT_DIRECTION);
-				Sort.Direction curDirection = orders.get(columns.get(i));
+			if (keyset.scrollsForward() ^ dir.isDescending()) {
+				primaryCompare = Criteria.where(primary).greaterThanOrEquals(primaryValue);
+			} else {
+				primaryCompare = Criteria.where(primary).lessThanOrEquals(primaryValue);
+			}
 
-				if (curDirection == null)
-					curDirection = defaultDirection;
+			Criteria tupleCompare = null;
+			for (int i = 0; i < columns.size(); i++) {
+				Criteria orCriteria;
+				String col = columns.get(i);
+				Object val = values.get(i);
 
-				if (((KeysetScrollPosition) position).scrollsForward()) {
-					if (curDirection.isAscending())
-						orCriteria = orCriteriaStep.greaterThan(values.get(i));
-					else
-						orCriteria = orCriteriaStep.lessThan(values.get(i));
-				} else {
-					if (curDirection.isAscending())
-						orCriteria = orCriteriaStep.lessThan(values.get(i));
-					else
-						orCriteria = orCriteriaStep.greaterThan(values.get(i));
-				}
+				Sort.Direction colDir = directions.getOrDefault(col, dir);
+				boolean asc = keyset.scrollsForward() ? colDir.isAscending() : colDir.isDescending();
 
-				for (int j = 0; j < i; j++)
-					orCriteria = Criteria.where(columns.get(j)).is(values.get(j)).and(orCriteria);
+				orCriteria = asc ? Criteria.where(col).greaterThan(val) : Criteria.where(col).lessThan(val);
 
-				result = (result == null) ? orCriteria : result.or(orCriteria);
+				tupleCompare = (tupleCompare == null) ? orCriteria : tupleCompare.or(orCriteria);
 			}
 
-			return criteria.and(result);
+			return criteria.and(primaryCompare).and(tupleCompare);
 		}
 
 		SelectBuilder.SelectOrdered applyOrderBy(Sort sort, RelationalPersistentEntity<?> entity, Table table,

spring-data-jdbc/src/test/java/org/springframework/data/jdbc/repository/JdbcRepositoryIntegrationTests.java

@@ -1559,8 +1559,8 @@ void queryByWindowKeySetTwoKeys() {
 		DummyEntity four = repository.save(createEntity("four", it -> it.setPointInTime(Instant.ofEpochSecond(4000))));
 		DummyEntity five = repository.save(createEntity("five", it -> it.setPointInTime(Instant.ofEpochSecond(5000))));
 
-		WindowIterator<DummyEntity> iter = WindowIterator.of(position -> repository.findFirst2ByOrderByIdPropDesc(position))
-			.startingAt(ScrollPosition.forward(Map.of("idProp", 5, "pointInTime", Instant.now())));
+		WindowIterator<DummyEntity> iter = WindowIterator.of(position -> repository.findFirst2ByOrderByIdPropAsc(position))
+			.startingAt(ScrollPosition.forward(Map.of("idProp", 1, "pointInTime", Instant.ofEpochSecond(1000))));
 
 		List<DummyEntity> entities = new ArrayList<>();
 		while (iter.hasNext())