Skip to content

Commit

Permalink
Added support for nested sorting for script sorting and geo sorting.
Browse files Browse the repository at this point in the history
  • Loading branch information
martijnvg committed May 16, 2013
1 parent b147235 commit 8365af4
Show file tree
Hide file tree
Showing 9 changed files with 520 additions and 22 deletions.
Expand Up @@ -32,7 +32,7 @@
*
*/
// LUCENE MONITOR: Monitor against FieldComparator.Double
public class DoubleScriptDataComparator extends FieldComparator<Double> {
public class DoubleScriptDataComparator extends NumberComparatorBase<Double> {

public static IndexFieldData.XFieldComparatorSource comparatorSource(SearchScript script) {
return new InnerSource(script);
Expand Down Expand Up @@ -126,4 +126,25 @@ public void setBottom(final int bottom) {
public Double value(int slot) {
return values[slot];
}

@Override
public void add(int slot, int doc) {
script.setNextDocId(doc);
values[slot] += script.runAsDouble();
}

@Override
public void divide(int slot, int divisor) {
values[slot] /= divisor;
}

@Override
public void missing(int slot) {
values[slot] = Double.MAX_VALUE;
}

@Override
public int compareBottomMissing() {
return Double.compare(bottom, Double.MAX_VALUE);
}
}
Expand Up @@ -31,7 +31,7 @@

/**
*/
public class GeoDistanceComparator extends FieldComparator<Double> {
public class GeoDistanceComparator extends NumberComparatorBase<Double> {

protected final IndexGeoPointFieldData<?> indexFieldData;

Expand Down Expand Up @@ -121,6 +121,26 @@ public Double value(int slot) {
return values[slot];
}

@Override
public void add(int slot, int doc) {
values[slot] += geoDistanceValues.computeDistance(doc);
}

@Override
public void divide(int slot, int divisor) {
values[slot] /= divisor;
}

@Override
public void missing(int slot) {
values[slot] = Double.MAX_VALUE;
}

@Override
public int compareBottomMissing() {
return Double.compare(bottom, Double.MAX_VALUE);
}

// Computes the distance based on geo points.
// Due to this abstractions the geo distance comparator doesn't need to deal with whether fields have one
// or multiple geo points per document.
Expand Down
Expand Up @@ -64,7 +64,7 @@ public FieldComparator<?> newComparator(String fieldname, int numHits, int sortP
return new NestedFieldComparator.Avg((NumberComparatorBase) wrappedComparator, rootDocumentsFilter, innerDocumentsFilter, numHits);
default:
throw new ElasticSearchIllegalArgumentException(
String.format("Unsupported sort_mode[%s] for nested type", sortMode)
String.format("Unsupported sort_mode[%s] for nested type", sortMode)
);
}
}
Expand Down
Expand Up @@ -22,6 +22,7 @@
import org.elasticsearch.common.geo.GeoDistance;
import org.elasticsearch.common.unit.DistanceUnit;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.index.query.FilterBuilder;

import java.io.IOException;

Expand All @@ -40,6 +41,8 @@ public class GeoDistanceSortBuilder extends SortBuilder {
private DistanceUnit unit;
private SortOrder order;
private String sortMode;
private FilterBuilder nestedFilter;
private String nestedPath;

/**
* Constructs a new distance based sort on a geo point like field.
Expand Down Expand Up @@ -107,11 +110,29 @@ public SortBuilder missing(Object missing) {
* Defines which distance to use for sorting in the case a document contains multiple geo points.
* Possible values: min and max
*/
public SortBuilder sortMode(String sortMode) {
public GeoDistanceSortBuilder sortMode(String sortMode) {
this.sortMode = sortMode;
return this;
}

/**
* Sets the nested filter that the nested objects should match with in order to be taken into account
* for sorting.
*/
public GeoDistanceSortBuilder setNestedFilter(FilterBuilder nestedFilter) {
this.nestedFilter = nestedFilter;
return this;
}

/**
* Sets the nested path if sorting occurs on a field that is inside a nested object. By default when sorting on a
* field inside a nested object, the nearest upper nested object is selected as nested path.
*/
public GeoDistanceSortBuilder setNestedPath(String nestedPath) {
this.nestedPath = nestedPath;
return this;
}

@Override
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
builder.startObject("_geo_distance");
Expand All @@ -135,6 +156,13 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
builder.field("mode", sortMode);
}

if (nestedPath != null) {
builder.field("nested_path", nestedPath);
}
if (nestedFilter != null) {
builder.field("nested_filter", nestedFilter, params);
}

builder.endObject();
return builder;
}
Expand Down
Expand Up @@ -19,6 +19,7 @@

package org.elasticsearch.search.sort;

import org.apache.lucene.search.Filter;
import org.apache.lucene.search.SortField;
import org.elasticsearch.ElasticSearchIllegalArgumentException;
import org.elasticsearch.common.geo.GeoDistance;
Expand All @@ -27,11 +28,16 @@
import org.elasticsearch.common.geo.GeoUtils;
import org.elasticsearch.common.unit.DistanceUnit;
import org.elasticsearch.common.xcontent.XContentParser;
import org.elasticsearch.index.fielddata.IndexFieldData;
import org.elasticsearch.index.fielddata.IndexGeoPointFieldData;
import org.elasticsearch.index.fielddata.fieldcomparator.GeoDistanceComparatorSource;
import org.elasticsearch.index.fielddata.fieldcomparator.SortMode;
import org.elasticsearch.index.mapper.FieldMapper;
import org.elasticsearch.index.mapper.ObjectMappers;
import org.elasticsearch.index.mapper.geo.GeoPointFieldMapper;
import org.elasticsearch.index.mapper.object.ObjectMapper;
import org.elasticsearch.index.search.nested.NestedFieldComparatorSource;
import org.elasticsearch.index.search.nested.NonNestedDocsFilter;
import org.elasticsearch.search.internal.SearchContext;

/**
Expand All @@ -52,6 +58,8 @@ public SortField parse(XContentParser parser, SearchContext context) throws Exce
GeoDistance geoDistance = GeoDistance.ARC;
boolean reverse = false;
SortMode sortMode = null;
String nestedPath = null;
Filter nestedFilter = null;

boolean normalizeLon = true;
boolean normalizeLat = true;
Expand All @@ -72,17 +80,21 @@ public SortField parse(XContentParser parser, SearchContext context) throws Exce
fieldName = currentName;
} else if (token == XContentParser.Token.START_OBJECT) {
// the json in the format of -> field : { lat : 30, lon : 12 }
fieldName = currentName;
while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
if (token == XContentParser.Token.FIELD_NAME) {
currentName = parser.currentName();
} else if (token.isValue()) {
if (currentName.equals(GeoPointFieldMapper.Names.LAT)) {
point.resetLat(parser.doubleValue());
} else if (currentName.equals(GeoPointFieldMapper.Names.LON)) {
point.resetLon(parser.doubleValue());
} else if (currentName.equals(GeoPointFieldMapper.Names.GEOHASH)) {
GeoHashUtils.decode(parser.text(), point);
if ("nested_filter".equals(currentName) || "nestedFilter".equals(currentName)) {
nestedFilter = context.queryParserService().parseInnerFilter(parser);
} else {
fieldName = currentName;
while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
if (token == XContentParser.Token.FIELD_NAME) {
currentName = parser.currentName();
} else if (token.isValue()) {
if (currentName.equals(GeoPointFieldMapper.Names.LAT)) {
point.resetLat(parser.doubleValue());
} else if (currentName.equals(GeoPointFieldMapper.Names.LON)) {
point.resetLon(parser.doubleValue());
} else if (currentName.equals(GeoPointFieldMapper.Names.GEOHASH)) {
GeoHashUtils.decode(parser.text(), point);
}
}
}
}
Expand All @@ -100,6 +112,8 @@ public SortField parse(XContentParser parser, SearchContext context) throws Exce
normalizeLon = parser.booleanValue();
} else if ("mode".equals(currentName)) {
sortMode = SortMode.fromString(parser.text());
} else if ("nested_path".equals(currentName) || "nestedPath".equals(currentName)) {
nestedPath = parser.text();
} else {
point.resetFromString(parser.text());
fieldName = currentName;
Expand All @@ -125,6 +139,35 @@ public SortField parse(XContentParser parser, SearchContext context) throws Exce
}
IndexGeoPointFieldData indexFieldData = context.fieldData().getForField(mapper);

return new SortField(fieldName, new GeoDistanceComparatorSource(indexFieldData, point.lat(), point.lon(), unit, geoDistance, sortMode), reverse);
IndexFieldData.XFieldComparatorSource geoDistanceComparatorSource = new GeoDistanceComparatorSource(
indexFieldData, point.lat(), point.lon(), unit, geoDistance, sortMode
);
ObjectMapper objectMapper;
if (nestedPath != null) {
ObjectMappers objectMappers = context.mapperService().objectMapper(nestedPath);
if (objectMappers == null) {
throw new ElasticSearchIllegalArgumentException("failed to find nested object mapping for explicit nested path [" + nestedPath + "]");
}
objectMapper = objectMappers.mapper();
if (!objectMapper.nested().isNested()) {
throw new ElasticSearchIllegalArgumentException("mapping for explicit nested path is not mapped as nested: [" + nestedPath + "]");
}
} else {
objectMapper = context.mapperService().resolveClosestNestedObjectMapper(fieldName);
}
if (objectMapper != null && objectMapper.nested().isNested()) {
Filter rootDocumentsFilter = context.filterCache().cache(NonNestedDocsFilter.INSTANCE);
Filter innerDocumentsFilter;
if (nestedFilter != null) {
innerDocumentsFilter = context.filterCache().cache(nestedFilter);
} else {
innerDocumentsFilter = context.filterCache().cache(objectMapper.nestedTypeFilter());
}
geoDistanceComparatorSource = new NestedFieldComparatorSource(
sortMode, geoDistanceComparatorSource, rootDocumentsFilter, innerDocumentsFilter
);
}

return new SortField(fieldName, geoDistanceComparatorSource, reverse);
}
}
43 changes: 43 additions & 0 deletions src/main/java/org/elasticsearch/search/sort/ScriptSortBuilder.java
Expand Up @@ -21,6 +21,7 @@

import com.google.common.collect.Maps;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.index.query.FilterBuilder;

import java.io.IOException;
import java.util.Map;
Expand All @@ -40,6 +41,12 @@ public class ScriptSortBuilder extends SortBuilder {

private Map<String, Object> params;

private String sortMode;

private FilterBuilder nestedFilter;

private String nestedPath;

/**
* Constructs a script sort builder with the script and the type.
*
Expand Down Expand Up @@ -100,6 +107,33 @@ public ScriptSortBuilder lang(String lang) {
return this;
}

/**
* Defines which distance to use for sorting in the case a document contains multiple geo points.
* Possible values: min and max
*/
public ScriptSortBuilder sortMode(String sortMode) {
this.sortMode = sortMode;
return this;
}

/**
* Sets the nested filter that the nested objects should match with in order to be taken into account
* for sorting.
*/
public ScriptSortBuilder setNestedFilter(FilterBuilder nestedFilter) {
this.nestedFilter = nestedFilter;
return this;
}

/**
* Sets the nested path if sorting occurs on a field that is inside a nested object. For sorting by script this
* needs to be specified.
*/
public ScriptSortBuilder setNestedPath(String nestedPath) {
this.nestedPath = nestedPath;
return this;
}

@Override
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
builder.startObject("_script");
Expand All @@ -114,6 +148,15 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
if (this.params != null) {
builder.field("params", this.params);
}
if (sortMode != null) {
builder.field("mode", sortMode);
}
if (nestedPath != null) {
builder.field("nested_path", nestedPath);
}
if (nestedFilter != null) {
builder.field("nested_filter", nestedFilter, params);
}
builder.endObject();
return builder;
}
Expand Down

0 comments on commit 8365af4

Please sign in to comment.