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 ZRANDMEMBER command #276

Merged
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
* Python: Added BLPOP and BRPOP commands ([#1369](https://github.com/aws/glide-for-redis/pull/1369))
* Python: Added ZRANGESTORE command ([#1377](https://github.com/aws/glide-for-redis/pull/1377))
* Python: Added ZDIFFSTORE command ([#1378](https://github.com/aws/glide-for-redis/pull/1378))
* Python: Added ZRANDMEMBER command (TODO: add PR link)


#### Fixes
Expand Down
170 changes: 119 additions & 51 deletions glide-core/src/client/value_conversion.rs

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh, lack of RESP2 support in java client strikes again!

Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,8 @@ pub(crate) enum ExpectedReturnType {
ArrayOfDoubleOrNull,
Lolwut,
ArrayOfArraysOfDoubleOrNull,
ArrayOfKeyValuePairs,
ArrayOfPairs,
ArrayOfMemberScorePairs,
}

pub(crate) fn convert_to_expected_type(
Expand Down Expand Up @@ -270,23 +271,25 @@ pub(crate) fn convert_to_expected_type(
.into()),
}
}
ExpectedReturnType::ArrayOfKeyValuePairs => match value {
Value::Nil => Ok(value),
Value::Array(ref array) if array.is_empty() || matches!(array[0], Value::Array(_)) => {
Ok(value)
}
Value::Array(array)
if matches!(array[0], Value::BulkString(_) | Value::SimpleString(_)) =>
{
convert_flat_array_to_key_value_pairs(array)
}
_ => Err((
ErrorKind::TypeError,
"Response couldn't be converted to an array of key-value pairs",
format!("(response was {:?})", value),
)
.into()),
},
// Used by HRANDFIELD when the WITHVALUES arg is passed.
// The server response can be a nil value, an empty array, a flat array of key-value pairs, or a two-dimensional array of key-value pairs.
// The conversions we do here are as follows:
//
// - if the server returned nil, return nil
// - if the server returned an empty array, return an empty array
// - otherwise, return a two-dimensional array of key-value pairs
ExpectedReturnType::ArrayOfPairs => convert_to_array_of_pairs(value, None),
// Used by ZRANDMEMBER when the WITHSCORES arg is passed.
// The server response can be a nil value, an empty array, a flat array of member-score pairs, or a two-dimensional array of member-score pairs.
// The server response scores can be strings or doubles. The conversions we do here are as follows:
//
// - if the server returned nil, return nil
// - if the server returned an empty array, return an empty array

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

does this ever happen? I think the commands will return nil, and we won't ever get an empty array

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

From the redis docs: "Array reply: when the additional count argument is passed, the command returns an array of members, or an empty array when key doesn't exist."

// - otherwise, return a two-dimensional array of member-score pairs, where scores are of type double
ExpectedReturnType::ArrayOfMemberScorePairs => {
aaron-congo marked this conversation as resolved.
Show resolved Hide resolved
// RESP2 returns scores as strings, but we want scores as type double.
convert_to_array_of_pairs(value, Some(ExpectedReturnType::Double))
}
}
}

Expand Down Expand Up @@ -367,10 +370,53 @@ fn convert_array_to_map(
Ok(Value::Map(map))
}

/// Converts a flat array of values to a two-dimensional array, where the inner arrays are two-length arrays representing key-value pairs. Normally a map would be more suitable for these responses, but some commands (eg HRANDFIELD) may return duplicate key-value pairs depending on the command arguments. These duplicated pairs cannot be represented by a map.
/// Used by commands like ZRANDMEMBER and HRANDFIELD. Normally a map would be more suitable for these key-value responses, but these commands may return duplicate key-value pairs depending on the command arguments. These duplicated pairs cannot be represented by a map.
///
/// Converts a server response as follows:
/// - if the server returned nil, return nil.
/// - if the server returned an empty array, return an empty array.
/// - if the server returned a flat array (RESP2), convert it to a two-dimensional array, where the inner arrays are length=2 arrays representing key-value pairs.
/// - if the server returned a two-dimensional array (RESP3), return the response as is, since it is already in the correct format.
/// - otherwise, return an error.
///
/// `response` is a server response that we should attempt to convert as described above.
/// `value_expected_return_type` indicates the desired return type of the values in the key-value pairs. The values will only be converted if the response was a flat array, since RESP3 already returns an array of pairs with values already of the correct type.
fn convert_to_array_of_pairs(
response: Value,
value_expected_return_type: Option<ExpectedReturnType>,
) -> RedisResult<Value> {
match response {
Value::Nil => Ok(response),
Value::Array(ref array) if array.is_empty() || matches!(array[0], Value::Array(_)) => {
// The server response is an empty array or a RESP3 array of pairs. In RESP3, the values in the pairs are
// already of the correct type, so we do not need to convert them and `response` is in the correct format.
Ok(response)
}
Value::Array(array)
if array.len() % 2 == 0
&& matches!(array[0], Value::BulkString(_) | Value::SimpleString(_)) =>
{
// The server response is a RESP2 flat array with keys at even indices and their associated values at
// odd indices.
convert_flat_array_to_array_of_pairs(array, value_expected_return_type)
}
_ => Err((
ErrorKind::TypeError,
"Response couldn't be converted to an array of member-score pairs",
format!("(response was {:?})", response),
)
.into()),
}
}

/// Converts a flat array of values to a two-dimensional array, where the inner arrays are length=2 arrays representing key-value pairs. Normally a map would be more suitable for these responses, but some commands (eg HRANDFIELD) may return duplicate key-value pairs depending on the command arguments. These duplicated pairs cannot be represented by a map.
///
/// `array` is a flat array containing keys at even-positioned elements and their associated values at odd-positioned elements.
fn convert_flat_array_to_key_value_pairs(array: Vec<Value>) -> RedisResult<Value> {
/// `value_expected_return_type` indicates the desired return type of the values in the key-value pairs.
fn convert_flat_array_to_array_of_pairs(
array: Vec<Value>,
value_expected_return_type: Option<ExpectedReturnType>,
) -> RedisResult<Value> {
if array.len() % 2 != 0 {
return Err((
ErrorKind::TypeError,
Expand All @@ -381,7 +427,9 @@ fn convert_flat_array_to_key_value_pairs(array: Vec<Value>) -> RedisResult<Value

let mut result = Vec::with_capacity(array.len() / 2);
for i in (0..array.len()).step_by(2) {
let pair = vec![array[i].clone(), array[i + 1].clone()];
let key = array[i].clone();
let value = convert_to_expected_type(array[i + 1].clone(), value_expected_return_type)?;
let pair = vec![key, value];
result.push(Value::Array(pair));
}
Ok(Value::Array(result))
Expand All @@ -407,10 +455,10 @@ pub(crate) fn expected_type_for_cmd(cmd: &Cmd) -> Option<ExpectedReturnType> {
b"GEOPOS" => Some(ExpectedReturnType::ArrayOfArraysOfDoubleOrNull),
b"HRANDFIELD" => cmd
.position(b"WITHVALUES")
.map(|_| ExpectedReturnType::ArrayOfKeyValuePairs),
.map(|_| ExpectedReturnType::ArrayOfPairs),
b"ZRANDMEMBER" => cmd
.position(b"WITHSCORES")
.map(|_| ExpectedReturnType::ArrayOfKeyValuePairs),
.map(|_| ExpectedReturnType::ArrayOfMemberScorePairs),
b"ZADD" => cmd
.position(b"INCR")
.map(|_| ExpectedReturnType::DoubleOrNull),
Expand Down Expand Up @@ -517,33 +565,20 @@ mod tests {
}

#[test]
fn convert_array_of_kv_pairs() {
fn convert_to_array_of_pairs_return_type() {
assert!(matches!(
expected_type_for_cmd(
redis::cmd("HRANDFIELD")
.arg("key")
.arg("1")
.arg("withvalues")
),
Some(ExpectedReturnType::ArrayOfKeyValuePairs)
Some(ExpectedReturnType::ArrayOfPairs)
));

assert!(expected_type_for_cmd(redis::cmd("HRANDFIELD").arg("key").arg("1")).is_none());
assert!(expected_type_for_cmd(redis::cmd("HRANDFIELD").arg("key")).is_none());

assert!(matches!(
expected_type_for_cmd(
redis::cmd("ZRANDMEMBER")
.arg("key")
.arg("1")
.arg("withscores")
),
Some(ExpectedReturnType::ArrayOfKeyValuePairs)
));

assert!(expected_type_for_cmd(redis::cmd("ZRANDMEMBER").arg("key").arg("1")).is_none());
assert!(expected_type_for_cmd(redis::cmd("ZRANDMEMBER").arg("key")).is_none());

let flat_array = Value::Array(vec![
Value::BulkString(b"key1".to_vec()),
Value::BulkString(b"value1".to_vec()),
Expand All @@ -561,38 +596,71 @@ mod tests {
]),
]);
let converted_flat_array =
convert_to_expected_type(flat_array, Some(ExpectedReturnType::ArrayOfKeyValuePairs))
.unwrap();
convert_to_expected_type(flat_array, Some(ExpectedReturnType::ArrayOfPairs)).unwrap();
assert_eq!(two_dimensional_array, converted_flat_array);

let converted_two_dimensional_array = convert_to_expected_type(
two_dimensional_array.clone(),
Some(ExpectedReturnType::ArrayOfKeyValuePairs),
Some(ExpectedReturnType::ArrayOfPairs),
)
.unwrap();
assert_eq!(two_dimensional_array, converted_two_dimensional_array);

let empty_array = Value::Array(vec![]);
let converted_empty_array = convert_to_expected_type(
empty_array.clone(),
Some(ExpectedReturnType::ArrayOfKeyValuePairs),
)
.unwrap();
let converted_empty_array =
convert_to_expected_type(empty_array.clone(), Some(ExpectedReturnType::ArrayOfPairs))
.unwrap();
assert_eq!(empty_array, converted_empty_array);

let converted_nil_value =
convert_to_expected_type(Value::Nil, Some(ExpectedReturnType::ArrayOfKeyValuePairs))
.unwrap();
convert_to_expected_type(Value::Nil, Some(ExpectedReturnType::ArrayOfPairs)).unwrap();
assert_eq!(Value::Nil, converted_nil_value);

let array_of_doubles = Value::Array(vec![Value::Double(5.5)]);
let flat_array_unexpected_length =
Value::Array(vec![Value::BulkString(b"somekey".to_vec())]);
assert!(convert_to_expected_type(
array_of_doubles,
Some(ExpectedReturnType::ArrayOfKeyValuePairs)
flat_array_unexpected_length,
Some(ExpectedReturnType::ArrayOfPairs)
)
.is_err());
}

#[test]
fn convert_to_member_score_pairs_return_type() {
assert!(matches!(
expected_type_for_cmd(
redis::cmd("ZRANDMEMBER")
.arg("key")
.arg("1")
.arg("withscores")
),
Some(ExpectedReturnType::ArrayOfMemberScorePairs)
));

assert!(expected_type_for_cmd(redis::cmd("ZRANDMEMBER").arg("key").arg("1")).is_none());
assert!(expected_type_for_cmd(redis::cmd("ZRANDMEMBER").arg("key")).is_none());

// convert_to_array_of_pairs_return_type already tests most functionality since the conversion for ArrayOfPairs
// and ArrayOfMemberScorePairs is mostly the same. Here we also test that the scores are converted to double
// when the server response was a RESP2 flat array.
let flat_array = Value::Array(vec![
Value::BulkString(b"one".to_vec()),
Value::BulkString(b"1.0".to_vec()),
Value::BulkString(b"two".to_vec()),
Value::BulkString(b"2.0".to_vec()),
]);
let expected_response = Value::Array(vec![
Value::Array(vec![Value::BulkString(b"one".to_vec()), Value::Double(1.0)]),
Value::Array(vec![Value::BulkString(b"two".to_vec()), Value::Double(2.0)]),
]);
let converted_flat_array = convert_to_expected_type(
flat_array,
Some(ExpectedReturnType::ArrayOfMemberScorePairs),
)
.unwrap();
assert_eq!(expected_response, converted_flat_array);
}

#[test]
fn convert_zadd_only_if_incr_is_included() {
assert!(matches!(
Expand Down
87 changes: 85 additions & 2 deletions python/python/glide/async_commands/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -983,7 +983,7 @@ async def hrandfield_count(self, key: str, count: int) -> List[str]:
key (str): The key of the hash.
count (int): The number of field names to return.
If `count` is positive, returns unique elements.
If negative, allows for duplicates.
If `count` is negative, allows for duplicates elements.

Returns:
List[str]: A list of random field names from the hash.
Expand All @@ -1010,7 +1010,7 @@ async def hrandfield_withvalues(self, key: str, count: int) -> List[List[str]]:
key (str): The key of the hash.
count (int): The number of field names to return.
If `count` is positive, returns unique elements.
If negative, allows for duplicates.
If `count` is negative, allows for duplicates elements.

Returns:
List[List[str]]: A list of `[field_name, value]` lists, where `field_name` is a random field name from the
Expand Down Expand Up @@ -2746,6 +2746,89 @@ async def zdiffstore(self, destination: str, keys: List[str]) -> int:
),
)

async def zrandmember(self, key: str) -> Optional[str]:
"""
Returns a random element from the sorted set stored at 'key'.

See https://valkey.io/commands/zrandmember for more details.

Args:
key (str): The key of the sorted set.

Returns:
Optional[str]: A string representing a random element from the sorted set.
If the sorted set does not exist or is empty, the response will be None.

Examples:
>>> await client.zrandmember("my_sorted_set")
"GLIDE" # "GLIDE" is a random member of "my_sorted_set".
>>> await client.zrandmember("non_existing_sorted_set")

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we shouldn't include error cases - examples should demonstrate how to use the command, not how to not use it...

Copy link
Author

@aaron-congo aaron-congo May 14, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I copied this example from the Java client, which also includes this example of executing this command against a non-existing key. There are also other commands in core.py that include examples of executing against a non-existing key.

None # "non_existing_sorted_set" is not an existing key, so None was returned.
"""
return cast(
Optional[str],
await self._execute_command(RequestType.ZRandMember, [key]),
)

async def zrandmember_count(self, key: str, count: int) -> List[str]:
Yury-Fridlyand marked this conversation as resolved.
Show resolved Hide resolved
"""
Retrieves random elements from the sorted set stored at 'key'.

See https://valkey.io/commands/zrandmember for more details.

Args:
key (str): The key of the sorted set.
count (int): The number of elements to return.
If `count` is positive, returns unique elements.
If `count` is negative, allows for duplicates elements.

Returns:
List[str]: A list of elements from the sorted set.
If the sorted set does not exist or is empty, the response will be an empty list.

Examples:
>>> await client.zrandmember("my_sorted_set", -3)
["GLIDE", "GLIDE", "PYTHON"] # "GLIDE" and "PYTHON" are random members of "my_sorted_set".
>>> await client.zrandmember("non_existing_sorted_set", 3)

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we should remove this example

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

See my comment above

[] # "non_existing_sorted_set" is not an existing key, so an empty list was returned.
"""
return cast(
List[str],
await self._execute_command(RequestType.ZRandMember, [key, str(count)]),
)

async def zrandmember_withscores(
self, key: str, count: int
) -> List[List[Union[str, float]]]:
"""
Retrieves random elements along with their scores from the sorted set stored at 'key'.

See https://valkey.io/commands/zrandmember for more details.

Args:
key (str): The key of the sorted set.
count (int): The number of elements to return.
If `count` is positive, returns unique elements.
If `count` is negative, allows for duplicates elements.

Returns:
List[List[Union[str, float]]]: A list of `[member, score]` lists, where `member` is a random member from
the sorted set and `score` is the associated score.
If the sorted set does not exist or is empty, the response will be an empty list.

Examples:
>>> await client.zrandmember_withscores("my_sorted_set", -3)
[["GLIDE", 1.0], ["GLIDE", 1.0], ["PYTHON", 2.0]] # "GLIDE" and "PYTHON" are random members of "my_sorted_set", and have scores of 1.0 and 2.0, respectively.
>>> await client.zrandmember_withscores("non_existing_sorted_set", 3)
[] # "non_existing_sorted_set" is not an existing key, so an empty list was returned.
"""
return cast(
List[List[Union[str, float]]],
await self._execute_command(
RequestType.ZRandMember, [key, str(count), "WITHSCORES"]
aaron-congo marked this conversation as resolved.
Show resolved Hide resolved
),
)

async def invoke_script(
self,
script: Script,
Expand Down
Loading
Loading