Skip to content

Commit

Permalink
Python: add ZRANGESTORE command (valkey-io#1377)
Browse files Browse the repository at this point in the history
* Python: add ZRANGESTORE command (#258)

* Update with PR link
  • Loading branch information
aaron-congo authored and cyip10 committed Jun 24, 2024
1 parent 2bc8b2d commit 3326b4f
Show file tree
Hide file tree
Showing 6 changed files with 372 additions and 0 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
* Python: Added ZMSCORE command ([#1357](https://github.com/aws/glide-for-redis/pull/1357))
* Python: Added HRANDFIELD command ([#1334](https://github.com/aws/glide-for-redis/pull/1334))
* Python: Added XADD, XTRIM commands ([#1320](https://github.com/aws/glide-for-redis/pull/1320))
* Python: Added ZRANGESTORE command ([#1377](https://github.com/aws/glide-for-redis/pull/1377))

#### Fixes
* Python: Fix typing error "‘type’ object is not subscriptable" ([#1203](https://github.com/aws/glide-for-redis/pull/1203))
Expand Down
41 changes: 41 additions & 0 deletions python/python/glide/async_commands/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
RangeByScore,
ScoreBoundary,
_create_zrange_args,
_create_zrangestore_args,
)
from glide.constants import TOK, TResult
from glide.protobuf.redis_request_pb2 import RequestType
Expand Down Expand Up @@ -2338,6 +2339,46 @@ async def zrange_withscores(
Mapping[str, float], await self._execute_command(RequestType.Zrange, args)
)

async def zrangestore(
self,
destination: str,
source: str,
range_query: Union[RangeByIndex, RangeByLex, RangeByScore],
reverse: bool = False,
) -> int:
"""
Stores a specified range of elements from the sorted set at `source`, into a new sorted set at `destination`. If
`destination` doesn't exist, a new sorted set is created; if it exists, it's overwritten.
When in Cluster mode, all keys must map to the same hash slot.
ZRANGESTORE can perform different types of range queries: by index (rank), by the score, or by lexicographical
order.
See https://valkey.io/commands/zrangestore for more details.
Args:
destination (str): The key for the destination sorted set.
source (str): The key of the source sorted set.
range_query (Union[RangeByIndex, RangeByLex, RangeByScore]): The range query object representing the type of range query to perform.
- For range queries by index (rank), use RangeByIndex.
- For range queries by lexicographical order, use RangeByLex.
- For range queries by score, use RangeByScore.
reverse (bool): If True, reverses the sorted set, with index 0 as the element with the highest score.
Returns:
int: The number of elements in the resulting sorted set.
Examples:
>>> await client.zrangestore("destination_key", "my_sorted_set", RangeByIndex(0, 2), True)
3 # The 3 members with the highest scores from "my_sorted_set" were stored in the sorted set at "destination_key".
>>> await client.zrangestore("destination_key", "my_sorted_set", RangeByScore(InfBound.NEG_INF, ScoreBoundary(3)))
2 # The 2 members with scores between negative infinity and 3 (inclusive) from "my_sorted_set" were stored in the sorted set at "destination_key".
"""
args = _create_zrangestore_args(destination, source, range_query, reverse)

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

async def zrank(
self,
key: str,
Expand Down
26 changes: 26 additions & 0 deletions python/python/glide/async_commands/sorted_set.py
Original file line number Diff line number Diff line change
Expand Up @@ -158,3 +158,29 @@ def _create_zrange_args(
args.append("WITHSCORES")

return args


def _create_zrangestore_args(
destination: str,
source: str,
range_query: Union[RangeByLex, RangeByScore, RangeByIndex],
reverse: bool,
) -> List[str]:
args = [destination, source, str(range_query.start), str(range_query.stop)]

if isinstance(range_query, RangeByScore):
args.append("BYSCORE")
elif isinstance(range_query, RangeByLex):
args.append("BYLEX")
if reverse:
args.append("REV")
if hasattr(range_query, "limit") and range_query.limit is not None:
args.extend(
[
"LIMIT",
str(range_query.limit.offset),
str(range_query.limit.count),
]
)

return args
33 changes: 33 additions & 0 deletions python/python/glide/async_commands/transaction.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
RangeByScore,
ScoreBoundary,
_create_zrange_args,
_create_zrangestore_args,
)
from glide.protobuf.redis_request_pb2 import RequestType

Expand Down Expand Up @@ -1664,6 +1665,38 @@ def zrange_withscores(

return self.append_command(RequestType.Zrange, args)

def zrangestore(
self: TTransaction,
destination: str,
source: str,
range_query: Union[RangeByIndex, RangeByLex, RangeByScore],
reverse: bool = False,
) -> TTransaction:
"""
Stores a specified range of elements from the sorted set at `source`, into a new sorted set at `destination`. If
`destination` doesn't exist, a new sorted set is created; if it exists, it's overwritten.
ZRANGESTORE can perform different types of range queries: by index (rank), by the score, or by lexicographical
order.
See https://valkey.io/commands/zrangestore for more details.
Args:
destination (str): The key for the destination sorted set.
source (str): The key of the source sorted set.
range_query (Union[RangeByIndex, RangeByLex, RangeByScore]): The range query object representing the type of range query to perform.
- For range queries by index (rank), use RangeByIndex.
- For range queries by lexicographical order, use RangeByLex.
- For range queries by score, use RangeByScore.
reverse (bool): If True, reverses the sorted set, with index 0 as the element with the highest score.
Command response:
int: The number of elements in the resulting sorted set.
"""
args = _create_zrangestore_args(destination, source, range_query, reverse)

return self.append_command(RequestType.ZRangeStore, args)

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

0 comments on commit 3326b4f

Please sign in to comment.