Use by-id lookup for queries referring to identifier values.

Closes: #2851
Original Pull Request: #2854

Author: mp911de

src/main/java/org/springframework/data/redis/core/RedisQueryEngine.java

@@ -33,18 +33,20 @@
 import org.springframework.data.keyvalue.core.SortAccessor;
 import org.springframework.data.keyvalue.core.SpelSortAccessor;
 import org.springframework.data.keyvalue.core.query.KeyValueQuery;
+import org.springframework.data.mapping.PersistentPropertyPath;
 import org.springframework.data.redis.connection.RedisConnection;
 import org.springframework.data.redis.connection.RedisGeoCommands.GeoLocation;
 import org.springframework.data.redis.connection.RedisGeoCommands.GeoRadiusCommandArgs;
 import org.springframework.data.redis.connection.util.ByteArrayWrapper;
 import org.springframework.data.redis.core.convert.GeoIndexedPropertyValue;
+import org.springframework.data.redis.core.convert.RedisConverter;
 import org.springframework.data.redis.core.convert.RedisData;
+import org.springframework.data.redis.core.mapping.RedisPersistentProperty;
 import org.springframework.data.redis.repository.query.RedisOperationChain;
 import org.springframework.data.redis.repository.query.RedisOperationChain.NearPath;
 import org.springframework.data.redis.repository.query.RedisOperationChain.PathAndValue;
 import org.springframework.data.redis.util.ByteUtils;
 import org.springframework.expression.spel.standard.SpelExpressionParser;
-import org.springframework.lang.NonNullApi;
 import org.springframework.lang.Nullable;
 import org.springframework.util.CollectionUtils;
 
@@ -100,7 +102,7 @@ private <T> List<T> doFind(RedisOperationChain criteria, long offset, int rows,
 
 		RedisCallback<Map<byte[], Map<byte[], byte[]>>> callback = connection -> {
 
-			List<byte[]> keys = findKeys(criteria, rows, keyspace, connection);
+			List<byte[]> keys = findKeys(criteria, rows, keyspace, type, connection);
 			byte[] keyspaceBin = getRequiredAdapter().getConverter().getConversionService().convert(keyspace + ":",
 					byte[].class);
 
@@ -143,18 +145,35 @@ private <T> List<T> doFind(RedisOperationChain criteria, long offset, int rows,
 		return result;
 	}
 
-	private List<byte[]> findKeys(RedisOperationChain criteria, int rows, String keyspace, RedisConnection connection) {
+	private List<byte[]> findKeys(RedisOperationChain criteria, int rows, String keyspace, Class<?> domainType,
+			RedisConnection connection) {
 
 		List<byte[]> allKeys = new ArrayList<>();
 
 		if (!criteria.getSismember().isEmpty()) {
-			allKeys.addAll(connection.sInter(keys(keyspace + ":", criteria.getSismember())));
+
+			Set<PathAndValue> sismember = criteria.getSismember();
+			if (sismember.size() == 1) {
+				KeySelector keySelector = KeySelector.of(getRequiredAdapter().getConverter(), sismember, domainType);
+				if (!keySelector.setValueLookup().isEmpty()) {
+					allKeys.addAll(connection.sInter(keys(keyspace + ":", keySelector.setValueLookup())));
+				}
+
+				allKeys.addAll(keySelector.keys());
+			} else {
+				allKeys.addAll(connection.sInter(keys(keyspace + ":", sismember)));
+			}
 		}
 
-		if (!criteria.getOrSismember().isEmpty()) {
+		KeySelector keySelector = KeySelector.of(getRequiredAdapter().getConverter(), criteria.getOrSismember(),
+				domainType);
+
+		if (!keySelector.setValueLookup().isEmpty()) {
 			allKeys.addAll(connection.sUnion(keys(keyspace + ":", criteria.getOrSismember())));
 		}
 
+		allKeys.addAll(keySelector.keys());
+
 		if (criteria.getNear() != null) {
 
 			GeoRadiusCommandArgs limit = GeoRadiusCommandArgs.newGeoRadiusArgs();
@@ -170,7 +189,6 @@ private List<byte[]> findKeys(RedisOperationChain criteria, int rows, String key
 			}
 		}
 
-
 		Set<ByteArrayWrapper> unique = new LinkedHashSet<>(allKeys.size());
 		allKeys.forEach(key -> unique.add(new ByteArrayWrapper(key)));
 
@@ -244,4 +262,35 @@ public RedisOperationChain resolve(KeyValueQuery<?> query) {
 			return (RedisOperationChain) query.getCriteria();
 		}
 	}
+
+	/**
+	 * Value object capturing the direct object keys and set of values that need to be looked up from the secondary
+	 * indexes.
+	 *
+	 * @param keys
+	 * @param setValueLookup
+	 * @since 3.2.4
+	 */
+	record KeySelector(Collection<byte[]> keys, Set<PathAndValue> setValueLookup) {
+
+		static KeySelector of(RedisConverter converter, Set<PathAndValue> pathAndValues, Class<?> domainType) {
+
+			Set<byte[]> keys = new LinkedHashSet<>();
+			Set<PathAndValue> remainder = new LinkedHashSet<>();
+
+			for (PathAndValue pathAndValue : pathAndValues) {
+
+				PersistentPropertyPath<RedisPersistentProperty> path = converter.getMappingContext()
+						.getPersistentPropertyPath(pathAndValue.getPath(), domainType);
+				if (path.getLeafProperty().isIdProperty()) {
+					byte[] key = converter.getConversionService().convert(pathAndValue.getFirstValue(), byte[].class);
+					keys.add(key);
+				} else {
+					remainder.add(pathAndValue);
+				}
+			}
+
+			return new KeySelector(keys, remainder);
+		}
+	}
 }

src/test/java/org/springframework/data/redis/repository/RedisRepositoryIntegrationTestBase.java

@@ -27,6 +27,7 @@
 
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.data.annotation.Id;
 import org.springframework.data.annotation.Reference;
@@ -101,6 +102,57 @@ void simpleFindShouldReturnEntitiesCorrectly() {
 		assertThat(repo.findByLastname("al'thor")).contains(rand);
 	}
 
+	@Test // GH-2851
+	void shouldReturnSingleEntityByIdViaQueryMethod() {
+
+		Person rand = new Person();
+		rand.firstname = "rand";
+		rand.lastname = "al'thor";
+
+		Person egwene = new Person();
+		egwene.firstname = "egwene";
+
+		repo.saveAll(Arrays.asList(rand, egwene));
+
+		assertThat(repo.findEntityById(rand.getId())).isEqualTo(rand);
+		assertThat(repo.findEntityById(egwene.getId())).isEqualTo(egwene);
+	}
+
+	@Test // GH-2851
+	void shouldProjectSingleResult() {
+
+		Person rand = new Person();
+		rand.firstname = "rand";
+		rand.lastname = "al'thor";
+
+		Person egwene = new Person();
+		egwene.firstname = "egwene";
+
+		repo.saveAll(Arrays.asList(rand, egwene));
+
+		PersonProjection projectionById = repo.findProjectionById(rand.getId());
+		assertThat(projectionById).isNotNull();
+		assertThat(projectionById.getFirstname()).isEqualTo(rand.firstname);
+	}
+
+	@Test // GH-2851
+	void shouldProjectCollection() {
+
+		Person rand = new Person();
+		rand.firstname = "rand";
+		rand.lastname = "al'thor";
+
+		Person egwene = new Person();
+		egwene.firstname = "egwene";
+
+		repo.saveAll(Arrays.asList(rand, egwene));
+
+		List<PersonProjection> projectionById = repo.findProjectionBy();
+		assertThat(projectionById).hasSize(2) //
+				.extracting(PersonProjection::getFirstname) //
+				.contains(rand.getFirstname(), egwene.getFirstname());
+	}
+
 	@Test // DATAREDIS-425
 	void simpleFindByMultipleProperties() {
 
@@ -570,10 +622,20 @@ public interface PersonRepository extends PagingAndSortingRepository<Person, Str
 
 		Slice<Person> findByHometownLocationNear(Point point, Distance distance, Pageable pageable);
 
+		Person findEntityById(String id);
+
+		PersonProjection findProjectionById(String id);
+
+		List<PersonProjection> findProjectionBy();
+
 		@Override
 		<S extends Person> List<S> findAll(Example<S> example);
 	}
 
+	public interface PersonProjection {
+		String getFirstname();
+	}
+
 	public interface CityRepository extends CrudRepository<City, String> {
 
 		List<City> findByLocationNear(Point point, Distance distance);