Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Python: add ZRANGESTORE command #1377

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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
Yury-Fridlyand marked this conversation as resolved.
Show resolved Hide resolved
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
Loading