Support GEOSHAPE field type in RediSearch

redis · Sep 26, 2023 · 5ccf0ff · 5ccf0ff
1 parent e0e2319
commit 5ccf0ff
Show file tree

Hide file tree

Showing 5 changed files with 159 additions and 11 deletions.
diff --git a/pom.xml b/pom.xml
@@ -72,6 +72,12 @@
 			<artifactId>gson</artifactId>
 			<version>2.10.1</version>
 		</dependency>
+		<dependency>
+			<groupId>org.locationtech.jts</groupId>
+			<artifactId>jts-core</artifactId>
+			<version>1.19.0</version>
+                        <optional>true</optional>
+		</dependency>
 
 		<dependency>
 			<groupId>junit</groupId>

diff --git a/src/main/java/redis/clients/jedis/search/SearchProtocol.java b/src/main/java/redis/clients/jedis/search/SearchProtocol.java
@@ -49,13 +49,13 @@ public byte[] getRaw() {
 
   public enum SearchKeyword implements Rawable {
 
-    SCHEMA, TEXT, TAG, NUMERIC, GEO, VECTOR, VERBATIM, NOCONTENT, NOSTOPWORDS, WITHSCORES, LANGUAGE,
-    INFIELDS, SORTBY, ASC, DESC, LIMIT, HIGHLIGHT, FIELDS, TAGS, SUMMARIZE, FRAGS, LEN, SEPARATOR,
-    INKEYS, RETURN, FILTER, GEOFILTER, ADD, INCR, MAX, FUZZY, READ, DEL, DD, TEMPORARY, STOPWORDS,
-    NOFREQS, NOFIELDS, NOOFFSETS, NOHL, SET, GET, ON, SORTABLE, UNF, PREFIX, LANGUAGE_FIELD, SCORE,
-    SCORE_FIELD, SCORER, PARAMS, AS, DIALECT, SLOP, TIMEOUT, INORDER, EXPANDER, MAXTEXTFIELDS,
-    SKIPINITIALSCAN, WITHSUFFIXTRIE, NOSTEM, NOINDEX, PHONETIC, WEIGHT, CASESENSITIVE,
-    LOAD, APPLY, GROUPBY, MAXIDLE, WITHCURSOR, DISTANCE, TERMS, INCLUDE, EXCLUDE,
+    SCHEMA, TEXT, TAG, NUMERIC, GEO, GEOSHAPE, VECTOR, VERBATIM, NOCONTENT, NOSTOPWORDS, WITHSCORES,
+    LANGUAGE, INFIELDS, SORTBY, ASC, DESC, LIMIT, HIGHLIGHT, FIELDS, TAGS, SUMMARIZE, FRAGS, LEN,
+    SEPARATOR, INKEYS, RETURN, FILTER, GEOFILTER, ADD, INCR, MAX, FUZZY, READ, DEL, DD, TEMPORARY,
+    STOPWORDS, NOFREQS, NOFIELDS, NOOFFSETS, NOHL, SET, GET, ON, SORTABLE, UNF, PREFIX,
+    LANGUAGE_FIELD, SCORE, SCORE_FIELD, SCORER, PARAMS, AS, DIALECT, SLOP, TIMEOUT, INORDER,
+    EXPANDER, MAXTEXTFIELDS, SKIPINITIALSCAN, WITHSUFFIXTRIE, NOSTEM, NOINDEX, PHONETIC, WEIGHT,
+    CASESENSITIVE, LOAD, APPLY, GROUPBY, MAXIDLE, WITHCURSOR, DISTANCE, TERMS, INCLUDE, EXCLUDE,
     SEARCH, AGGREGATE, QUERY, LIMITED, COUNT, REDUCE;
 
     private final byte[] raw;

diff --git a/src/main/java/redis/clients/jedis/search/schemafields/GeoShapeField.java b/src/main/java/redis/clients/jedis/search/schemafields/GeoShapeField.java
@@ -0,0 +1,41 @@
+package redis.clients.jedis.search.schemafields;
+
+import static redis.clients.jedis.search.SearchProtocol.SearchKeyword.GEOSHAPE;
+
+import redis.clients.jedis.CommandArguments;
+import redis.clients.jedis.search.FieldName;
+
+public class GeoShapeField extends SchemaField {
+
+  public enum CoordinateSystem {
+    FLAT,
+    SPHERICAL
+  }
+
+  private final CoordinateSystem system;
+
+  public GeoShapeField(String fieldName, CoordinateSystem system) {
+    super(fieldName);
+    this.system = system;
+  }
+
+  public GeoShapeField(FieldName fieldName, CoordinateSystem system) {
+    super(fieldName);
+    this.system = system;
+  }
+
+  public static GeoShapeField of(String fieldName, CoordinateSystem system) {
+    return new GeoShapeField(fieldName, system);
+  }
+
+  @Override
+  public GeoShapeField as(String attribute) {
+    super.as(attribute);
+    return this;
+  }
+
+  @Override
+  public void addParams(CommandArguments args) {
+    args.addParams(fieldName).add(GEOSHAPE).add(system);
+  }
+}
diff --git a/src/test/java/redis/clients/jedis/modules/search/QueryBuilderTest.java b/src/test/java/redis/clients/jedis/modules/search/QueryBuilderTest.java
@@ -91,10 +91,10 @@ public void testIntersectionBasic() {
 
   @Test
   public void testIntersectionNested() {
-    Node n = intersect().
-        add(union("name", value("mark"), value("dvir"))).
-        add("time", between(100, 200)).
-        add(disjunct("created", lt(1000)));
+    Node n = intersect()
+        .add(union("name", value("mark"), value("dvir")))
+        .add("time", between(100, 200))
+        .add(disjunct("created", lt(1000)));
     assertEquals("(@name:(mark|dvir) @time:[100 200] -@created:[-inf (1000])", n.toString());
   }
 

diff --git a/src/test/java/redis/clients/jedis/modules/search/SearchWithParamsTest.java b/src/test/java/redis/clients/jedis/modules/search/SearchWithParamsTest.java
@@ -9,6 +9,13 @@
 import org.junit.BeforeClass;
 import org.junit.Test;
 
+import org.locationtech.jts.geom.Coordinate;
+import org.locationtech.jts.geom.GeometryFactory;
+import org.locationtech.jts.geom.Point;
+import org.locationtech.jts.geom.Polygon;
+import org.locationtech.jts.io.ParseException;
+import org.locationtech.jts.io.WKTReader;
+
 import redis.clients.jedis.GeoCoordinate;
 import redis.clients.jedis.RedisProtocol;
 import redis.clients.jedis.args.GeoUnit;
@@ -333,6 +340,100 @@ public void geoFilterAndGeoCoordinateObject() {
     assertEquals(2, res.getTotalResults());
   }
 
+  @Test
+  public void geoShapeFilterSpherical() throws ParseException {
+    final WKTReader reader = new WKTReader();
+    final GeometryFactory factory = new GeometryFactory();
+
+    assertOK(client.ftCreate(index, GeoShapeField.of("geom", GeoShapeField.CoordinateSystem.SPHERICAL)));
+
+    // polygon type
+    final Polygon small = factory.createPolygon(new Coordinate[]{new Coordinate(34.9001, 29.7001),
+        new Coordinate(34.9001, 29.7100), new Coordinate(34.9100, 29.7100),
+        new Coordinate(34.9100, 29.7001), new Coordinate(34.9001, 29.7001)});
+    client.hset("small", RediSearchUtil.toStringMap(Collections.singletonMap("geom", small)));
+
+    final Polygon large = factory.createPolygon(new Coordinate[]{new Coordinate(34.9001, 29.7001),
+        new Coordinate(34.9001, 29.7200), new Coordinate(34.9200, 29.7200),
+        new Coordinate(34.9200, 29.7001), new Coordinate(34.9001, 29.7001)});
+    client.hset("large", RediSearchUtil.toStringMap(Collections.singletonMap("geom", large)));
+
+    // within condition
+    final Polygon within = factory.createPolygon(new Coordinate[]{new Coordinate(34.9000, 29.7000),
+        new Coordinate(34.9000, 29.7150), new Coordinate(34.9150, 29.7150),
+        new Coordinate(34.9150, 29.7000), new Coordinate(34.9000, 29.7000)});
+
+    SearchResult res = client.ftSearch(index, "@geom:[within $poly]",
+        FTSearchParams.searchParams().addParam("poly", within).dialect(3));
+    assertEquals(1, res.getTotalResults());
+    assertEquals(1, res.getDocuments().size());
+    assertEquals(small, reader.read(res.getDocuments().get(0).getString("geom")));
+
+    // contains condition
+    final Polygon contains = factory.createPolygon(new Coordinate[]{new Coordinate(34.9002, 29.7002),
+        new Coordinate(34.9002, 29.7050), new Coordinate(34.9050, 29.7050),
+        new Coordinate(34.9050, 29.7002), new Coordinate(34.9002, 29.7002)});
+
+    res = client.ftSearch(index, "@geom:[contains $poly]",
+        FTSearchParams.searchParams().addParam("poly", contains).dialect(3));
+    assertEquals(2, res.getTotalResults());
+    assertEquals(2, res.getDocuments().size());
+
+    // point type
+    final Point point = factory.createPoint(new Coordinate(34.9010, 29.7010));
+    client.hset("point", RediSearchUtil.toStringMap(Collections.singletonMap("geom", point)));
+
+    res = client.ftSearch(index, "@geom:[within $poly]",
+        FTSearchParams.searchParams().addParam("poly", within).dialect(3));
+    assertEquals(2, res.getTotalResults());
+    assertEquals(2, res.getDocuments().size());
+  }
+
+  @Test
+  public void geoShapeFilterFlat() throws ParseException {
+    final WKTReader reader = new WKTReader();
+    final GeometryFactory factory = new GeometryFactory();
+
+    assertOK(client.ftCreate(index, GeoShapeField.of("geom", GeoShapeField.CoordinateSystem.FLAT)));
+
+    // polygon type
+    final Polygon small = factory.createPolygon(new Coordinate[]{new Coordinate(1, 1),
+        new Coordinate(1, 100), new Coordinate(100, 100), new Coordinate(100, 1), new Coordinate(1, 1)});
+    client.hset("small", RediSearchUtil.toStringMap(Collections.singletonMap("geom", small)));
+
+    final Polygon large = factory.createPolygon(new Coordinate[]{new Coordinate(1, 1),
+        new Coordinate(1, 200), new Coordinate(200, 200), new Coordinate(200, 1), new Coordinate(1, 1)});
+    client.hset("large", RediSearchUtil.toStringMap(Collections.singletonMap("geom", large)));
+
+    // within condition
+    final Polygon within = factory.createPolygon(new Coordinate[]{new Coordinate(0, 0),
+        new Coordinate(0, 150), new Coordinate(150, 150), new Coordinate(150, 0), new Coordinate(0, 0)});
+
+    SearchResult res = client.ftSearch(index, "@geom:[within $poly]",
+        FTSearchParams.searchParams().addParam("poly", within).dialect(3));
+    assertEquals(1, res.getTotalResults());
+    assertEquals(1, res.getDocuments().size());
+    assertEquals(small, reader.read(res.getDocuments().get(0).getString("geom")));
+
+    // contains condition
+    final Polygon contains = factory.createPolygon(new Coordinate[]{new Coordinate(2, 2),
+        new Coordinate(2, 50), new Coordinate(50, 50), new Coordinate(50, 2), new Coordinate(2, 2)});
+
+    res = client.ftSearch(index, "@geom:[contains $poly]",
+        FTSearchParams.searchParams().addParam("poly", contains).dialect(3));
+    assertEquals(2, res.getTotalResults());
+    assertEquals(2, res.getDocuments().size());
+
+    // point type
+    final Point point = factory.createPoint(new Coordinate(10, 10));
+    client.hset("point", RediSearchUtil.toStringMap(Collections.singletonMap("geom", point)));
+
+    res = client.ftSearch(index, "@geom:[within $poly]",
+        FTSearchParams.searchParams().addParam("poly", within).dialect(3));
+    assertEquals(2, res.getTotalResults());
+    assertEquals(2, res.getDocuments().size());
+  }
+
   @Test
   public void testQueryFlags() {
     assertOK(client.ftCreate(index, TextField.of("title")));