Skip to content

Commit 05124de

Browse files
authored
Add GEOSHAPE field type for index creation of RediSearch (#2957)
* first pass of geoshape index type * first attempt at test, but demonstrates the initial commit is broken * fix new field + fix tests * work on linter * more linter * try to mark test with correct fixture * fix linter
1 parent 7867c62 commit 05124de

File tree

2 files changed

+34
-0
lines changed

2 files changed

+34
-0
lines changed

redis/commands/search/field.py

+16
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ class Field:
1313
SORTABLE = "SORTABLE"
1414
NOINDEX = "NOINDEX"
1515
AS = "AS"
16+
GEOSHAPE = "GEOSHAPE"
1617

1718
def __init__(
1819
self,
@@ -91,6 +92,21 @@ def __init__(self, name: str, **kwargs):
9192
Field.__init__(self, name, args=[Field.NUMERIC], **kwargs)
9293

9394

95+
class GeoShapeField(Field):
96+
"""
97+
GeoShapeField is used to enable within/contain indexing/searching
98+
"""
99+
100+
SPHERICAL = "SPHERICAL"
101+
FLAT = "FLAT"
102+
103+
def __init__(self, name: str, coord_system=None, **kwargs):
104+
args = [Field.GEOSHAPE]
105+
if coord_system:
106+
args.append(coord_system)
107+
Field.__init__(self, name, args=args, **kwargs)
108+
109+
94110
class GeoField(Field):
95111
"""
96112
GeoField is used to define a geo-indexing field in a schema definition

tests/test_search.py

+18
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
from redis.commands.search import Search
1414
from redis.commands.search.field import (
1515
GeoField,
16+
GeoShapeField,
1617
NumericField,
1718
TagField,
1819
TextField,
@@ -2266,3 +2267,20 @@ def test_query_timeout(r: redis.Redis):
22662267
q2 = Query("foo").timeout("not_a_number")
22672268
with pytest.raises(redis.ResponseError):
22682269
r.ft().search(q2)
2270+
2271+
2272+
@pytest.mark.redismod
2273+
def test_geoshape(client: redis.Redis):
2274+
client.ft().create_index((GeoShapeField("geom", GeoShapeField.FLAT)))
2275+
waitForIndex(client, getattr(client.ft(), "index_name", "idx"))
2276+
client.hset("small", "geom", "POLYGON((1 1, 1 100, 100 100, 100 1, 1 1))")
2277+
client.hset("large", "geom", "POLYGON((1 1, 1 200, 200 200, 200 1, 1 1))")
2278+
q1 = Query("@geom:[WITHIN $poly]").dialect(3)
2279+
qp1 = {"poly": "POLYGON((0 0, 0 150, 150 150, 150 0, 0 0))"}
2280+
q2 = Query("@geom:[CONTAINS $poly]").dialect(3)
2281+
qp2 = {"poly": "POLYGON((2 2, 2 50, 50 50, 50 2, 2 2))"}
2282+
result = client.ft().search(q1, query_params=qp1)
2283+
assert len(result.docs) == 1
2284+
assert result.docs[0]["id"] == "small"
2285+
result = client.ft().search(q2, query_params=qp2)
2286+
assert len(result.docs) == 2

0 commit comments

Comments
 (0)