Skip to content

Commit 8ae8fb4

Browse files
Fix wrong type and non existing key issues in S/Z UNION, DIFF, INTER commands (#1162)
* Fix wrong type and non existing key issues in S/Z UNION, DIFF, INTER commants * Fixed review comment --------- Co-authored-by: Badrish Chandramouli <[email protected]>
1 parent 14fb5f9 commit 8ae8fb4

File tree

4 files changed

+453
-107
lines changed

4 files changed

+453
-107
lines changed

Diff for: libs/server/Storage/Session/ObjectStore/SetOps.cs

+26
Original file line numberDiff line numberDiff line change
@@ -464,6 +464,17 @@ private GarnetStatus SetIntersect<TObjectContext>(ReadOnlySpan<ArgSlice> keys, r
464464
}
465465

466466
var status = GET(keys[0].ToArray(), out var first, ref objectContext);
467+
468+
if (status == GarnetStatus.NOTFOUND)
469+
{
470+
return GarnetStatus.OK;
471+
}
472+
473+
if (status == GarnetStatus.WRONGTYPE)
474+
{
475+
return GarnetStatus.WRONGTYPE;
476+
}
477+
467478
if (status == GarnetStatus.OK)
468479
{
469480
if (first.GarnetObject is not SetObject firstObject)
@@ -490,6 +501,8 @@ private GarnetStatus SetIntersect<TObjectContext>(ReadOnlySpan<ArgSlice> keys, r
490501
}
491502

492503
status = GET(keys[i].ToArray(), out var next, ref objectContext);
504+
if (status == GarnetStatus.WRONGTYPE)
505+
return GarnetStatus.WRONGTYPE;
493506
if (status == GarnetStatus.OK)
494507
{
495508
if (next.GarnetObject is not SetObject nextObject)
@@ -890,6 +903,17 @@ private GarnetStatus SetDiff<TObjectContext>(ArgSlice[] keys, ref TObjectContext
890903

891904
// first SetObject
892905
var status = GET(keys[0].ToArray(), out var first, ref objectContext);
906+
907+
if (status == GarnetStatus.NOTFOUND)
908+
{
909+
return GarnetStatus.OK;
910+
}
911+
912+
if (status == GarnetStatus.WRONGTYPE)
913+
{
914+
return GarnetStatus.WRONGTYPE;
915+
}
916+
893917
if (status == GarnetStatus.OK)
894918
{
895919
if (first.GarnetObject is not SetObject firstObject)
@@ -909,6 +933,8 @@ private GarnetStatus SetDiff<TObjectContext>(ArgSlice[] keys, ref TObjectContext
909933
for (var i = 1; i < keys.Length; i++)
910934
{
911935
status = GET(keys[i].ToArray(), out var next, ref objectContext);
936+
if (status == GarnetStatus.WRONGTYPE)
937+
return GarnetStatus.WRONGTYPE;
912938
if (status == GarnetStatus.OK)
913939
{
914940
if (next.GarnetObject is not SetObject nextObject)

Diff for: libs/server/Storage/Session/ObjectStore/SortedSetOps.cs

+142-107
Original file line numberDiff line numberDiff line change
@@ -1126,57 +1126,71 @@ private GarnetStatus SortedSetUnion<TObjectContext>(ReadOnlySpan<ArgSlice> keys,
11261126

11271127
// Get the first sorted set
11281128
var status = GET(keys[0].ToArray(), out var firstObj, ref objectContext);
1129+
1130+
if (status == GarnetStatus.WRONGTYPE)
1131+
{
1132+
return GarnetStatus.WRONGTYPE;
1133+
}
1134+
1135+
Dictionary<byte[], double> sortedSetDictionary = null;
1136+
11291137
if (status == GarnetStatus.OK)
11301138
{
11311139
if (firstObj.GarnetObject is not SortedSetObject firstSortedSet)
11321140
{
11331141
return GarnetStatus.WRONGTYPE;
11341142
}
1143+
sortedSetDictionary = firstSortedSet.Dictionary;
1144+
}
11351145

1136-
// Initialize pairs with the first set
1137-
if (weights is null)
1138-
{
1139-
pairs = new Dictionary<byte[], double>(firstSortedSet.Dictionary, ByteArrayComparer.Instance);
1140-
}
1141-
else
1146+
// Initialize pairs with the first set
1147+
if (weights is null)
1148+
{
1149+
pairs = sortedSetDictionary is null ? new Dictionary<byte[], double>(ByteArrayComparer.Instance) : new Dictionary<byte[], double>(sortedSetDictionary, ByteArrayComparer.Instance);
1150+
}
1151+
else
1152+
{
1153+
pairs = new Dictionary<byte[], double>(ByteArrayComparer.Instance);
1154+
if (sortedSetDictionary is not null)
11421155
{
1143-
pairs = new Dictionary<byte[], double>(ByteArrayComparer.Instance);
1144-
foreach (var (key, score) in firstSortedSet.Dictionary)
1156+
foreach (var (key, score) in sortedSetDictionary)
11451157
{
11461158
pairs[key] = weights[0] * score;
11471159
}
11481160
}
1161+
}
11491162

1150-
// Process remaining sets
1151-
for (var i = 1; i < keys.Length; i++)
1163+
// Process remaining sets
1164+
for (var i = 1; i < keys.Length; i++)
1165+
{
1166+
status = GET(keys[i].ToArray(), out var nextObj, ref objectContext);
1167+
if (status == GarnetStatus.WRONGTYPE)
1168+
return GarnetStatus.WRONGTYPE;
1169+
if (status != GarnetStatus.OK)
1170+
continue;
1171+
1172+
if (nextObj.GarnetObject is not SortedSetObject nextSortedSet)
11521173
{
1153-
status = GET(keys[i].ToArray(), out var nextObj, ref objectContext);
1154-
if (status != GarnetStatus.OK)
1155-
continue;
1174+
pairs = default;
1175+
return GarnetStatus.WRONGTYPE;
1176+
}
11561177

1157-
if (nextObj.GarnetObject is not SortedSetObject nextSortedSet)
1178+
foreach (var (key, score) in nextSortedSet.Dictionary)
1179+
{
1180+
var weightedScore = weights is null ? score : score * weights[i];
1181+
if (pairs.TryGetValue(key, out var existingScore))
11581182
{
1159-
pairs = default;
1160-
return GarnetStatus.WRONGTYPE;
1183+
pairs[key] = aggregateType switch
1184+
{
1185+
SortedSetAggregateType.Sum => existingScore + weightedScore,
1186+
SortedSetAggregateType.Min => Math.Min(existingScore, weightedScore),
1187+
SortedSetAggregateType.Max => Math.Max(existingScore, weightedScore),
1188+
_ => existingScore + weightedScore // Default to SUM
1189+
};
11611190
}
1162-
1163-
foreach (var (key, score) in nextSortedSet.Dictionary)
1191+
else
11641192
{
1165-
var weightedScore = weights is null ? score : score * weights[i];
1166-
if (pairs.TryGetValue(key, out var existingScore))
1167-
{
1168-
pairs[key] = aggregateType switch
1169-
{
1170-
SortedSetAggregateType.Sum => existingScore + weightedScore,
1171-
SortedSetAggregateType.Min => Math.Min(existingScore, weightedScore),
1172-
SortedSetAggregateType.Max => Math.Max(existingScore, weightedScore),
1173-
_ => existingScore + weightedScore // Default to SUM
1174-
};
1175-
}
1176-
else
1177-
{
1178-
pairs[key] = weightedScore;
1179-
}
1193+
pairs[key] = weightedScore;
11801194
}
11811195
}
11821196
}
@@ -1190,37 +1204,47 @@ private GarnetStatus SortedSetDifference<TObjectContext>(ReadOnlySpan<ArgSlice>
11901204
pairs = default;
11911205

11921206
var statusOp = GET(keys[0].ToArray(), out var firstObj, ref objectContext);
1193-
if (statusOp == GarnetStatus.OK)
1207+
if (statusOp == GarnetStatus.WRONGTYPE)
11941208
{
1195-
if (firstObj.GarnetObject is not SortedSetObject firstSortedSet)
1196-
{
1197-
return GarnetStatus.WRONGTYPE;
1198-
}
1209+
return GarnetStatus.WRONGTYPE;
1210+
}
11991211

1200-
if (keys.Length == 1)
1201-
{
1202-
pairs = firstSortedSet.Dictionary;
1203-
return GarnetStatus.OK;
1204-
}
1212+
if (statusOp == GarnetStatus.NOTFOUND)
1213+
{
1214+
pairs = new Dictionary<byte[], double>(ByteArrayComparer.Instance);
1215+
return GarnetStatus.OK;
1216+
}
12051217

1206-
// read the rest of the keys
1207-
for (var item = 1; item < keys.Length; item++)
1208-
{
1209-
statusOp = GET(keys[item].ToArray(), out var nextObj, ref objectContext);
1210-
if (statusOp != GarnetStatus.OK)
1211-
continue;
1218+
if (firstObj.GarnetObject is not SortedSetObject firstSortedSet)
1219+
{
1220+
return GarnetStatus.WRONGTYPE;
1221+
}
12121222

1213-
if (nextObj.GarnetObject is not SortedSetObject nextSortedSet)
1214-
{
1215-
pairs = default;
1216-
return GarnetStatus.WRONGTYPE;
1217-
}
1223+
if (keys.Length == 1)
1224+
{
1225+
pairs = firstSortedSet.Dictionary;
1226+
return GarnetStatus.OK;
1227+
}
12181228

1219-
if (pairs == default)
1220-
pairs = SortedSetObject.CopyDiff(firstSortedSet, nextSortedSet);
1221-
else
1222-
SortedSetObject.InPlaceDiff(pairs, nextSortedSet);
1229+
// read the rest of the keys
1230+
for (var item = 1; item < keys.Length; item++)
1231+
{
1232+
statusOp = GET(keys[item].ToArray(), out var nextObj, ref objectContext);
1233+
if (statusOp == GarnetStatus.WRONGTYPE)
1234+
return GarnetStatus.WRONGTYPE;
1235+
if (statusOp != GarnetStatus.OK)
1236+
continue;
1237+
1238+
if (nextObj.GarnetObject is not SortedSetObject nextSortedSet)
1239+
{
1240+
pairs = default;
1241+
return GarnetStatus.WRONGTYPE;
12231242
}
1243+
1244+
if (pairs == default)
1245+
pairs = SortedSetObject.CopyDiff(firstSortedSet, nextSortedSet);
1246+
else
1247+
SortedSetObject.InPlaceDiff(pairs, nextSortedSet);
12241248
}
12251249

12261250
return GarnetStatus.OK;
@@ -1410,71 +1434,82 @@ private GarnetStatus SortedSetIntersection<TObjectContext>(ReadOnlySpan<ArgSlice
14101434
pairs = default;
14111435

14121436
var statusOp = GET(keys[0].ToArray(), out var firstObj, ref objectContext);
1413-
if (statusOp == GarnetStatus.OK)
1437+
1438+
if (statusOp == GarnetStatus.WRONGTYPE)
14141439
{
1415-
if (firstObj.GarnetObject is not SortedSetObject firstSortedSet)
1416-
{
1417-
return GarnetStatus.WRONGTYPE;
1418-
}
1440+
return GarnetStatus.WRONGTYPE;
1441+
}
14191442

1420-
// Initialize result with first set
1421-
if (weights is null)
1443+
if (statusOp == GarnetStatus.NOTFOUND)
1444+
{
1445+
pairs = new Dictionary<byte[], double>(ByteArrayComparer.Instance);
1446+
return GarnetStatus.OK;
1447+
}
1448+
1449+
if (firstObj.GarnetObject is not SortedSetObject firstSortedSet)
1450+
{
1451+
return GarnetStatus.WRONGTYPE;
1452+
}
1453+
1454+
// Initialize result with first set
1455+
if (weights is null)
1456+
{
1457+
pairs = keys.Length == 1 ? firstSortedSet.Dictionary : new Dictionary<byte[], double>(firstSortedSet.Dictionary, ByteArrayComparer.Instance);
1458+
}
1459+
else
1460+
{
1461+
pairs = new Dictionary<byte[], double>(ByteArrayComparer.Instance);
1462+
foreach (var kvp in firstSortedSet.Dictionary)
14221463
{
1423-
pairs = keys.Length == 1 ? firstSortedSet.Dictionary : new Dictionary<byte[], double>(firstSortedSet.Dictionary, ByteArrayComparer.Instance);
1464+
pairs[kvp.Key] = kvp.Value * weights[0];
14241465
}
1425-
else
1466+
}
1467+
1468+
if (keys.Length == 1)
1469+
{
1470+
return GarnetStatus.OK;
1471+
}
1472+
1473+
// Intersect with remaining sets
1474+
for (var i = 1; i < keys.Length; i++)
1475+
{
1476+
statusOp = GET(keys[i].ToArray(), out var nextObj, ref objectContext);
1477+
if (statusOp == GarnetStatus.WRONGTYPE)
1478+
return GarnetStatus.WRONGTYPE;
1479+
if (statusOp != GarnetStatus.OK)
14261480
{
1427-
pairs = new Dictionary<byte[], double>(ByteArrayComparer.Instance);
1428-
foreach (var kvp in firstSortedSet.Dictionary)
1429-
{
1430-
pairs[kvp.Key] = kvp.Value * weights[0];
1431-
}
1481+
pairs = default;
1482+
return statusOp;
14321483
}
14331484

1434-
if (keys.Length == 1)
1485+
if (nextObj.GarnetObject is not SortedSetObject nextSortedSet)
14351486
{
1436-
return GarnetStatus.OK;
1487+
pairs = default;
1488+
return GarnetStatus.WRONGTYPE;
14371489
}
14381490

1439-
// Intersect with remaining sets
1440-
for (var i = 1; i < keys.Length; i++)
1491+
foreach (var kvp in pairs)
14411492
{
1442-
statusOp = GET(keys[i].ToArray(), out var nextObj, ref objectContext);
1443-
if (statusOp != GarnetStatus.OK)
1493+
if (!nextSortedSet.TryGetScore(kvp.Key, out var score))
14441494
{
1445-
pairs = default;
1446-
return statusOp;
1447-
}
1448-
1449-
if (nextObj.GarnetObject is not SortedSetObject nextSortedSet)
1450-
{
1451-
pairs = default;
1452-
return GarnetStatus.WRONGTYPE;
1495+
pairs.Remove(kvp.Key);
1496+
continue;
14531497
}
14541498

1455-
foreach (var kvp in pairs)
1499+
var weightedScore = weights is null ? score : score * weights[i];
1500+
pairs[kvp.Key] = aggregateType switch
14561501
{
1457-
if (!nextSortedSet.TryGetScore(kvp.Key, out var score))
1458-
{
1459-
pairs.Remove(kvp.Key);
1460-
continue;
1461-
}
1462-
1463-
var weightedScore = weights is null ? score : score * weights[i];
1464-
pairs[kvp.Key] = aggregateType switch
1465-
{
1466-
SortedSetAggregateType.Sum => kvp.Value + weightedScore,
1467-
SortedSetAggregateType.Min => Math.Min(kvp.Value, weightedScore),
1468-
SortedSetAggregateType.Max => Math.Max(kvp.Value, weightedScore),
1469-
_ => kvp.Value + weightedScore // Default to SUM
1470-
};
1471-
}
1502+
SortedSetAggregateType.Sum => kvp.Value + weightedScore,
1503+
SortedSetAggregateType.Min => Math.Min(kvp.Value, weightedScore),
1504+
SortedSetAggregateType.Max => Math.Max(kvp.Value, weightedScore),
1505+
_ => kvp.Value + weightedScore // Default to SUM
1506+
};
1507+
}
14721508

1473-
// If intersection becomes empty, we can stop early
1474-
if (pairs.Count == 0)
1475-
{
1476-
break;
1477-
}
1509+
// If intersection becomes empty, we can stop early
1510+
if (pairs.Count == 0)
1511+
{
1512+
break;
14781513
}
14791514
}
14801515

0 commit comments

Comments
 (0)