Skip to content

Add support for generating geospatial/vector search queries & expression parameter bindings & regular expressions during AOT run. #5005

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 7 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@

<groupId>org.springframework.data</groupId>
<artifactId>spring-data-mongodb-parent</artifactId>
<version>5.0.0-SNAPSHOT</version>
<version>5.0.x-GH-5004-SNAPSHOT</version>
<packaging>pom</packaging>

<name>Spring Data MongoDB</name>
Expand Down
2 changes: 1 addition & 1 deletion spring-data-mongodb-distribution/pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
<parent>
<groupId>org.springframework.data</groupId>
<artifactId>spring-data-mongodb-parent</artifactId>
<version>5.0.0-SNAPSHOT</version>
<version>5.0.x-GH-5004-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>

Expand Down
4 changes: 2 additions & 2 deletions spring-data-mongodb/pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
<parent>
<groupId>org.springframework.data</groupId>
<artifactId>spring-data-mongodb-parent</artifactId>
<version>5.0.0-SNAPSHOT</version>
<version>5.0.x-GH-5004-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>

Expand Down Expand Up @@ -264,7 +264,7 @@
<dependency>
<groupId>org.junit-pioneer</groupId>
<artifactId>junit-pioneer</artifactId>
<version>0.5.3</version>
<version>2.3.0</version>
<scope>test</scope>
</dependency>

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -225,6 +225,14 @@ interface TerminatingFindNear<T> {
* @return never {@literal null}.
*/
GeoResults<T> all();

/**
* Count matching elements.
*
* @return number of elements matching the query.
* @since 5.0
*/
long count();
}

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -243,6 +243,11 @@ public <R> TerminatingFindNear<R> map(QueryResultConverter<? super G, ? extends
public GeoResults<G> all() {
return template.doGeoNear(nearQuery, domainType, getCollectionName(), returnType, resultConverter);
}

@Override
public long count() {
return template.doGeoNearCount(nearQuery, domainType, getCollectionName());
}
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@
import org.springframework.dao.support.PersistenceExceptionTranslator;
import org.springframework.data.convert.EntityReader;
import org.springframework.data.domain.OffsetScrollPosition;
import org.springframework.data.domain.Pageable;
import org.springframework.data.domain.Window;
import org.springframework.data.geo.Distance;
import org.springframework.data.geo.GeoResult;
Expand Down Expand Up @@ -1044,6 +1045,31 @@ public <T> GeoResults<T> geoNear(NearQuery near, Class<?> domainType, String col
return doGeoNear(near, domainType, collectionName, returnType, QueryResultConverter.entity());
}

long doGeoNearCount(NearQuery near, Class<?> domainType, String collectionName) {

Builder optionsBuilder = AggregationOptions.builder().collation(near.getCollation());

if (near.hasReadPreference()) {
optionsBuilder.readPreference(near.getReadPreference());
}

if (near.hasReadConcern()) {
optionsBuilder.readConcern(near.getReadConcern());
}

String distanceField = operations.nearQueryDistanceFieldName(domainType);
Aggregation $geoNear = TypedAggregation.newAggregation(domainType,
Aggregation.geoNear(near, distanceField).skip(-1).limit(-1), Aggregation.count().as("_totalCount"))
.withOptions(optionsBuilder.build());

AggregationResults<Document> results = doAggregate($geoNear, collectionName, Document.class,
queryOperations.createAggregation($geoNear, (AggregationOperationContext) null));
Iterator<Document> iterator = results.iterator();
return iterator.hasNext()
? NumberUtils.convertNumberToTargetClass(iterator.next().get("_totalCount", Integer.class), Long.class)
: 0L;
}

<T, R> GeoResults<R> doGeoNear(NearQuery near, Class<?> domainType, String collectionName, Class<T> returnType,
QueryResultConverter<? super T, ? extends R> resultConverter) {

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,8 @@ public class GeoNearOperation implements AggregationOperation {
private final NearQuery nearQuery;
private final String distanceField;
private final @Nullable String indexKey;
private final @Nullable Long skip;
private final @Nullable Integer limit;

/**
* Creates a new {@link GeoNearOperation} from the given {@link NearQuery} and the given distance field. The
Expand All @@ -51,7 +53,7 @@ public class GeoNearOperation implements AggregationOperation {
* @param distanceField must not be {@literal null}.
*/
public GeoNearOperation(NearQuery nearQuery, String distanceField) {
this(nearQuery, distanceField, null);
this(nearQuery, distanceField, null, nearQuery.getSkip(), null);
}

/**
Expand All @@ -63,14 +65,17 @@ public GeoNearOperation(NearQuery nearQuery, String distanceField) {
* @param indexKey can be {@literal null};
* @since 2.1
*/
private GeoNearOperation(NearQuery nearQuery, String distanceField, @Nullable String indexKey) {
private GeoNearOperation(NearQuery nearQuery, String distanceField, @Nullable String indexKey, @Nullable Long skip,
@Nullable Integer limit) {

Assert.notNull(nearQuery, "NearQuery must not be null");
Assert.hasLength(distanceField, "Distance field must not be null or empty");

this.nearQuery = nearQuery;
this.distanceField = distanceField;
this.indexKey = indexKey;
this.skip = skip;
this.limit = limit;
}

/**
Expand All @@ -83,7 +88,30 @@ private GeoNearOperation(NearQuery nearQuery, String distanceField, @Nullable St
*/
@Contract("_ -> new")
public GeoNearOperation useIndex(String key) {
return new GeoNearOperation(nearQuery, distanceField, key);
return new GeoNearOperation(nearQuery, distanceField, key, skip, limit);
}

/**
* Override potential skip applied via {@link NearQuery#getSkip()}. Adds an additional {@link SkipOperation} if value
* is non negative.
*
* @param skip
* @return new instance of {@link GeoNearOperation}.
* @since 5.0
*/
public GeoNearOperation skip(long skip) {
return new GeoNearOperation(nearQuery, distanceField, indexKey, skip, limit);
}

/**
* Override potential limit value. Adds an additional {@link LimitOperation} if value is non negative.
*
* @param limit
* @return new instance of {@link GeoNearOperation}.
* @since 5.0
*/
public GeoNearOperation limit(Integer limit) {
return new GeoNearOperation(nearQuery, distanceField, indexKey, skip, limit);
}

@Override
Expand All @@ -92,7 +120,13 @@ public Document toDocument(AggregationOperationContext context) {
Document command = context.getMappedObject(nearQuery.toDocument());

if (command.containsKey("query")) {
command.replace("query", context.getMappedObject(command.get("query", Document.class)));
Document query = command.get("query", Document.class);
if (query == null || query.isEmpty()) {
command.remove("query");
} else {
command.replace("query", context.getMappedObject(query));
}

}

command.remove("collation");
Expand All @@ -115,15 +149,18 @@ public List<Document> toPipelineStages(AggregationOperationContext context) {

Document command = toDocument(context);
Number limit = (Number) command.get("$geoNear", Document.class).remove("num");
if (limit != null && this.limit != null) {
limit = this.limit;
}

List<Document> stages = new ArrayList<>(3);
stages.add(command);

if (nearQuery.getSkip() != null && nearQuery.getSkip() > 0) {
stages.add(new Document("$skip", nearQuery.getSkip()));
if (this.skip != null && this.skip > 0) {
stages.add(new Document("$skip", this.skip));
}

if (limit != null) {
if (limit != null && limit.longValue() > 0) {
stages.add(new Document("$limit", limit.longValue()));
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -46,33 +46,37 @@ public interface CriteriaDefinition {
* @since 5.0
* @author Christoph Strobl
*/
class Placeholder {

private final Object expression;
interface Placeholder {

/**
* Create a new placeholder for index bindable parameter.
*
* @param position the index of the parameter to bind.
* @return new instance of {@link Placeholder}.
*/
public static Placeholder indexed(int position) {
return new Placeholder("?%s".formatted(position));
static Placeholder indexed(int position) {
return new PlaceholderImpl("?%s".formatted(position));
}

public static Placeholder placeholder(String expression) {
return new Placeholder(expression);
static Placeholder placeholder(String expression) {
return new PlaceholderImpl(expression);
}

Placeholder(Object value) {
this.expression = value;
Object getValue();
}

static class PlaceholderImpl implements Placeholder {
private final Object expression;

public PlaceholderImpl(Object expression) {
this.expression = expression;
}

@Override
public Object getValue() {
return expression;
}

@Override
public String toString() {
return getValue().toString();
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
import org.springframework.data.geo.Circle;
import org.springframework.data.geo.Polygon;
import org.springframework.data.geo.Shape;
import org.springframework.data.mongodb.core.geo.GeoJson;
import org.springframework.data.mongodb.core.geo.Sphere;
import org.springframework.util.Assert;

Expand Down Expand Up @@ -75,6 +76,9 @@ private String getCommand(Shape shape) {

Assert.notNull(shape, "Shape must not be null");

if(shape instanceof GeoJson<?>) {
return "$geometry";
}
if (shape instanceof Box) {
return "$box";
} else if (shape instanceof Circle) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -671,7 +671,7 @@ public Document toDocument() {
document.put("distanceMultiplier", getDistanceMultiplier());
}

if (limit != null) {
if (limit != null && limit > 0) {
document.put("num", limit);
}

Expand Down
Loading