Skip to content

Commit

Permalink
Python: adds GEOSEARCHSTORE command (valkey-io#1581)
Browse files Browse the repository at this point in the history
  • Loading branch information
shohamazon authored and acarbonetto committed Jun 18, 2024
1 parent 416c7fe commit 9d1eb55
Show file tree
Hide file tree
Showing 8 changed files with 460 additions and 10 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
* Python: Added OBJECT FREQ command ([#1472](https://github.com/aws/glide-for-redis/pull/1472))
* Python: Added OBJECT IDLETIME command ([#1474](https://github.com/aws/glide-for-redis/pull/1474))
* Python: Added GEOSEARCH command ([#1482](https://github.com/aws/glide-for-redis/pull/1482))
* Python: Added GEOSEARCHSTORE command ([#1581](https://github.com/aws/glide-for-redis/pull/1581))
* Node: Added RENAMENX command ([#1483](https://github.com/aws/glide-for-redis/pull/1483))
* Python: Added OBJECT REFCOUNT command ([#1485](https://github.com/aws/glide-for-redis/pull/1485))
* Python: Added RENAMENX command ([#1492](https://github.com/aws/glide-for-redis/pull/1492))
Expand Down
5 changes: 3 additions & 2 deletions glide-core/src/protobuf/redis_request.proto
Original file line number Diff line number Diff line change
Expand Up @@ -222,8 +222,9 @@ enum RequestType {
GeoSearch = 182;
Watch = 183;
UnWatch = 184;
XGroupCreateConsumer = 185;
XGroupDelConsumer = 186;
GeoSearchStore = 185;
XGroupCreateConsumer = 188;
XGroupDelConsumer = 189;
}

message Command {
Expand Down
7 changes: 5 additions & 2 deletions glide-core/src/request_type.rs
Original file line number Diff line number Diff line change
Expand Up @@ -192,8 +192,9 @@ pub enum RequestType {
GeoSearch = 182,
Watch = 183,
UnWatch = 184,
XGroupCreateConsumer = 185,
XGroupDelConsumer = 186,
GeoSearchStore = 185,
XGroupCreateConsumer = 188,
XGroupDelConsumer = 189,
}

fn get_two_word_command(first: &str, second: &str) -> Cmd {
Expand Down Expand Up @@ -389,6 +390,7 @@ impl From<::protobuf::EnumOrUnknown<ProtobufRequestType>> for RequestType {
ProtobufRequestType::GeoSearch => RequestType::GeoSearch,
ProtobufRequestType::Watch => RequestType::Watch,
ProtobufRequestType::UnWatch => RequestType::UnWatch,
ProtobufRequestType::GeoSearchStore => RequestType::GeoSearchStore,
ProtobufRequestType::XGroupCreateConsumer => RequestType::XGroupCreateConsumer,
ProtobufRequestType::XGroupDelConsumer => RequestType::XGroupDelConsumer,
}
Expand Down Expand Up @@ -582,6 +584,7 @@ impl RequestType {
RequestType::GeoSearch => Some(cmd("GEOSEARCH")),
RequestType::Watch => Some(cmd("WATCH")),
RequestType::UnWatch => Some(cmd("UNWATCH")),
RequestType::GeoSearchStore => Some(cmd("GEOSEARCHSTORE")),
RequestType::XGroupCreateConsumer => {
Some(get_two_word_command("XGROUP", "CREATECONSUMER"))
}
Expand Down
68 changes: 67 additions & 1 deletion python/python/glide/async_commands/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -2708,7 +2708,7 @@ async def geosearch(
Since: Redis version 6.2.0.
"""
args = _create_geosearch_args(
key,
[key],
search_from,
seach_by,
order_by,
Expand All @@ -2723,6 +2723,72 @@ async def geosearch(
await self._execute_command(RequestType.GeoSearch, args),
)

async def geosearchstore(
self,
destination: str,
source: str,
search_from: Union[str, GeospatialData],
search_by: Union[GeoSearchByRadius, GeoSearchByBox],
count: Optional[GeoSearchCount] = None,
store_dist: bool = False,
) -> int:
"""
Searches for members in a sorted set stored at `key` representing geospatial data within a circular or rectangular area and stores the result in `destination`.
If `destination` already exists, it is overwritten. Otherwise, a new sorted set will be created.
To get the result directly, see `geosearch`.
Note:
When in cluster mode, both `source` and `destination` must map to the same hash slot.
Args:
destination (str): The key to store the search results.
source (str): The key of the sorted set representing geospatial data to search from.
search_from (Union[str, GeospatialData]): The location to search from. Can be specified either as a member
from the sorted set or as a geospatial data (see `GeospatialData`).
search_by (Union[GeoSearchByRadius, GeoSearchByBox]): The search criteria.
For circular area search, see `GeoSearchByRadius`.
For rectangular area search, see `GeoSearchByBox`.
count (Optional[GeoSearchCount]): Specifies the maximum number of results to store. See `GeoSearchCount`.
If not specified, stores all results.
store_dist (bool): Determines what is stored as the sorted set score. Defaults to False.
- If set to False, the geohash of the location will be stored as the sorted set score.
- If set to True, the distance from the center of the shape (circle or box) will be stored as the sorted set score.
The distance is represented as a floating-point number in the same unit specified for that shape.
Returns:
int: The number of elements in the resulting sorted set stored at `destination`.
Examples:
>>> await client.geoadd("my_geo_sorted_set", {"Palermo": GeospatialData(13.361389, 38.115556), "Catania": GeospatialData(15.087269, 37.502669)})
>>> await client.geosearchstore("my_dest_sorted_set", "my_geo_sorted_set", "Catania", GeoSearchByRadius(175, GeoUnit.MILES))
2 # Number of elements stored in "my_dest_sorted_set".
>>> await client.zrange_withscores("my_dest_sorted_set", RangeByIndex(0, -1))
{"Palermo": 3479099956230698.0, "Catania": 3479447370796909.0} # The elements within te search area, with their geohash as score.
>>> await client.geosearchstore("my_dest_sorted_set", "my_geo_sorted_set", GeospatialData(15, 37), GeoSearchByBox(400, 400, GeoUnit.KILOMETERS), store_dist=True)
2 # Number of elements stored in "my_dest_sorted_set", with distance as score.
>>> await client.zrange_withscores("my_dest_sorted_set", RangeByIndex(0, -1))
{"Catania": 56.4412578701582, "Palermo": 190.44242984775784} # The elements within te search area, with the distance as score.
Since: Redis version 6.2.0.
"""
args = _create_geosearch_args(
[destination, source],
search_from,
search_by,
None,
count,
False,
False,
False,
store_dist,
)

return cast(
int,
await self._execute_command(RequestType.GeoSearchStore, args),
)

async def zadd(
self,
key: str,
Expand Down
8 changes: 6 additions & 2 deletions python/python/glide/async_commands/sorted_set.py
Original file line number Diff line number Diff line change
Expand Up @@ -356,16 +356,17 @@ def _create_zinter_zunion_cmd_args(


def _create_geosearch_args(
key: str,
keys: List[str],
search_from: Union[str, GeospatialData],
seach_by: Union[GeoSearchByRadius, GeoSearchByBox],
order_by: Optional[OrderBy] = None,
count: Optional[GeoSearchCount] = None,
with_coord: bool = False,
with_dist: bool = False,
with_hash: bool = False,
store_dist: bool = False,
) -> List[str]:
args = [key]
args = keys
if isinstance(search_from, str):
args += ["FROMMEMBER", search_from]
else:
Expand All @@ -389,4 +390,7 @@ def _create_geosearch_args(
if with_hash:
args.append("WITHHASH")

if store_dist:
args.append("STOREDIST")

return args
53 changes: 52 additions & 1 deletion python/python/glide/async_commands/transaction.py
Original file line number Diff line number Diff line change
Expand Up @@ -1872,7 +1872,7 @@ def geosearch(
Since: Redis version 6.2.0.
"""
args = _create_geosearch_args(
key,
[key],
search_from,
seach_by,
order_by,
Expand All @@ -1884,6 +1884,57 @@ def geosearch(

return self.append_command(RequestType.GeoSearch, args)

def geosearchstore(
self: TTransaction,
destination: str,
source: str,
search_from: Union[str, GeospatialData],
search_by: Union[GeoSearchByRadius, GeoSearchByBox],
count: Optional[GeoSearchCount] = None,
store_dist: bool = False,
) -> TTransaction:
"""
Searches for members in a sorted set stored at `key` representing geospatial data within a circular or rectangular area and stores the result in `destination`.
If `destination` already exists, it is overwritten. Otherwise, a new sorted set will be created.
To get the result directly, see `geosearch`.
See https://valkey.io/commands/geosearch/ for more details.
Args:
destination (str): The key to store the search results.
source (str): The key of the sorted set representing geospatial data to search from.
search_from (Union[str, GeospatialData]): The location to search from. Can be specified either as a member
from the sorted set or as a geospatial data (see `GeospatialData`).
search_by (Union[GeoSearchByRadius, GeoSearchByBox]): The search criteria.
For circular area search, see `GeoSearchByRadius`.
For rectangular area search, see `GeoSearchByBox`.
count (Optional[GeoSearchCount]): Specifies the maximum number of results to store. See `GeoSearchCount`.
If not specified, stores all results.
store_dist (bool): Determines what is stored as the sorted set score. Defaults to False.
- If set to False, the geohash of the location will be stored as the sorted set score.
- If set to True, the distance from the center of the shape (circle or box) will be stored as the sorted set score.
The distance is represented as a floating-point number in the same unit specified for that shape.
Commands response:
int: The number of elements in the resulting sorted set stored at `destination`.s
Since: Redis version 6.2.0.
"""
args = _create_geosearch_args(
[destination, source],
search_from,
search_by,
None,
count,
False,
False,
False,
store_dist,
)

return self.append_command(RequestType.GeoSearchStore, args)

def zadd(
self: TTransaction,
key: str,
Expand Down
Loading

0 comments on commit 9d1eb55

Please sign in to comment.