Skip to content

Commit

Permalink
Add geo bounding box support for query serializer
Browse files Browse the repository at this point in the history 
Signed-off-by: Peter Alfonsi <[email protected]>
  • Loading branch information
Peter Alfonsi committed Dec 11, 2024
1 parent cfa5076 commit 06b91fa
Show file tree
Hide file tree
Showing 2 changed files with 93 additions and 9 deletions.
86 changes: 77 additions & 9 deletions modules/cache-common/src/main/java/org/opensearch/cache/common/query/QuerySerializer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,19 +8,29 @@

package org.opensearch.cache.common.query;

import org.apache.lucene.document.LatLonPoint;
import org.apache.lucene.document.LongPoint;
import org.apache.lucene.search.PointRangeQuery;
import org.apache.lucene.search.Query;
import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.NumericUtils;
import org.opensearch.OpenSearchException;
import org.opensearch.common.cache.serializer.Serializer;
import org.opensearch.common.io.stream.BytesStreamOutput;
import org.opensearch.core.common.bytes.BytesReference;
import org.opensearch.core.common.io.stream.BytesStreamInput;

import java.io.IOError;
import java.io.IOException;
import java.util.Arrays;

import static org.apache.lucene.geo.GeoEncodingUtils.decodeLatitude;
import static org.apache.lucene.geo.GeoEncodingUtils.decodeLongitude;
import static org.apache.lucene.geo.GeoEncodingUtils.encodeLatitude;
import static org.apache.lucene.geo.GeoEncodingUtils.encodeLatitudeCeil;
import static org.apache.lucene.geo.GeoEncodingUtils.encodeLongitude;
import static org.apache.lucene.geo.GeoEncodingUtils.encodeLongitudeCeil;

/**
* A class to serialize Query objects. Not all Query objects are supported. You can check with isAllowed().
*/
Expand All @@ -33,6 +43,11 @@ public class QuerySerializer implements Serializer<Query, byte[]> {

// TODO: Used as part of gross hack to determine which impl of PointRangeQuery comes in.
static final Query longPointRangeQuery = LongPoint.newRangeQuery("", 0, 1);
static final Query geoBoundingBoxQuery = LatLonPoint.newBoxQuery("", 0.0, 0.0, 0.0, 0.0);

// Bytes for different type of PointRangeQuery
static final byte LONG_POINT_BYTE = 0x01;
static final byte GEO_BOUNDING_BOX_BYTE = 0x02;

/**
* Required for javadocs.
Expand Down Expand Up @@ -89,10 +104,19 @@ public boolean equals(Query object, byte[] bytes) {

private void serializePointRangeQuery(BytesStreamOutput os, Query query) throws IOException {
PointRangeQuery prQuery = (PointRangeQuery) query;
// Determine if it's LongPoint's implementation
if (!isLongPointRangeQuery(prQuery)) {
throw new UnsupportedOperationException("can only serialize LongPoint's implementation of PointRangeQuery");
final byte pointRangeTypeByte = getPointRangeTypeByte(prQuery);
os.writeByte(pointRangeTypeByte);
switch (pointRangeTypeByte) {
case LONG_POINT_BYTE:
serializeLongPoint(os, prQuery);
break;
case GEO_BOUNDING_BOX_BYTE:
serializeGeoBoundingBox(os, prQuery);
break;
}
}

private void serializeLongPoint(BytesStreamOutput os, PointRangeQuery prQuery) throws IOException {
os.writeString(prQuery.getField());

byte[] lowerPoint = prQuery.getLowerPoint();
Expand All @@ -102,13 +126,9 @@ private void serializePointRangeQuery(BytesStreamOutput os, Query query) throws
byte[] upperPoint = prQuery.getUpperPoint();
os.writeVInt(upperPoint.length);
os.writeBytes(upperPoint);
// TODO: for PointRangeQuery these are byte[] but for LongPoint's impl they're long[]
// It looks like LongPoint calls pack() (public) before shoving it into bytes. So I need to do the reverse of pack() on the way
// out...
}

private Query deserializePointRangeQuery(BytesStreamInput is) throws IOException {
// TODO: Currently ONLY reads LongPoint's implementation of PointRangeQuery
private Query deserializeLongPoint(BytesStreamInput is) throws IOException {
String field = is.readString();

int lowerPointLength = is.readVInt();
Expand All @@ -126,6 +146,45 @@ private Query deserializePointRangeQuery(BytesStreamInput is) throws IOException
return LongPoint.newRangeQuery(field, lowerPointLong, upperPointLong);
}

private void serializeGeoBoundingBox(BytesStreamOutput os, PointRangeQuery prQuery) throws IOException {
// Serializes queries produced via LatLonPoint.newBoxQuery()
// TODO: For now, assume the query does NOT cross the dateline. If it does, newBoxQuery() does something complex, so avoid this case for now.
// I think in this case the query type ends up being ConstantScoreQuery, so it's probably going to fail on its own anyway.

os.writeString(prQuery.getField());

byte[] lowerPoint = prQuery.getLowerPoint();
assert lowerPoint.length == 2 * Integer.BYTES;
os.writeBytes(lowerPoint);

byte[] upperPoint = prQuery.getUpperPoint();
assert upperPoint.length == 2 * Integer.BYTES;
os.writeBytes(upperPoint);
}

private Query deserializeGeoBoundingBox(BytesStreamInput is) throws IOException {
String field = is.readString();
byte[] values = new byte[4 * Integer.BYTES];
is.readBytes(values, 0, 4 * Integer.BYTES);
double minLatitude = decodeLatitude(values, 0);
double minLongitude = decodeLongitude(values, Integer.BYTES);
double maxLatitude = decodeLatitude(values, 2 * Integer.BYTES);
double maxLongitude = decodeLongitude(values, 3 * Integer.BYTES);

return LatLonPoint.newBoxQuery(field, minLatitude, maxLatitude, minLongitude, maxLongitude);
}

private Query deserializePointRangeQuery(BytesStreamInput is) throws IOException {
final byte pointRangeTypeByte = is.readByte();
switch (pointRangeTypeByte) {
case LONG_POINT_BYTE:
return deserializeLongPoint(is);
case GEO_BOUNDING_BOX_BYTE:
return deserializeGeoBoundingBox(is);
}
throw new UnsupportedOperationException("Unrecognized PointRangeQuery type byte" + pointRangeTypeByte);
}

private void serializeDummyQuery(BytesStreamOutput os, Query query) throws IOException {
os.writeInt(((DummyQuery) query).getId());
}
Expand All @@ -134,12 +193,21 @@ private Query deserializeDummyQuery(BytesStreamInput is) throws IOException {
return new DummyQuery(is.readInt());
}

private byte getPointRangeTypeByte(PointRangeQuery query) {
if (isLongPointRangeQuery(query)) return LONG_POINT_BYTE;
if (isGeoBoundingBoxQuery(query)) return GEO_BOUNDING_BOX_BYTE;
throw new UnsupportedOperationException("Cannot serialize this type of PointRangeQuery");
}

private boolean isLongPointRangeQuery(PointRangeQuery query) {
// TODO: A gross hack - but ok for the PoC
// TODO: also - does it work??
return query.getClass() == longPointRangeQuery.getClass();
}

private boolean isGeoBoundingBoxQuery(PointRangeQuery query) {
return query.getClass() == geoBoundingBoxQuery.getClass();
}

/**
* Checks whether the query is serializable.
* @param query The query to serialize.
Expand Down
16 changes: 16 additions & 0 deletions ...es/cache-common/src/test/java/org/opensearch/cache/common/query/QuerySerializerTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
package org.opensearch.cache.common.query;

import org.apache.lucene.document.DoublePoint;
import org.apache.lucene.document.LatLonPoint;
import org.apache.lucene.document.LongPoint;
import org.apache.lucene.search.Query;
import org.opensearch.test.OpenSearchTestCase;
Expand Down Expand Up @@ -48,4 +49,19 @@ public void testDummyQuery() throws Exception {
assertNotEquals(new DummyQuery(id - 1), original);
assertNotEquals(new DummyQuery(id - 1), deser);
}

public void testGeoBoundingBoxQuery() throws Exception {
String field = "location";
double minLatitude = -4.9;
double minLongitude = -122.0;
double maxLatitude = 44.0;
double maxLongitude = -121.9;

Query original = LatLonPoint.newBoxQuery(field, minLatitude, maxLatitude, minLongitude, maxLongitude);
QuerySerializer serializer = new QuerySerializer();
byte[] ser = serializer.serialize(original);
Query deser = serializer.deserialize(ser);
assertTrue(serializer.equals(original, ser));
assertEquals(original, deser);
}
}

0 comments on commit 06b91fa

Please sign in to comment.