SERVER-61868 Switch aggColl(Min|Max) to use comparison that is aware of MQL semantics for NaN

3 files changed, 101 insertions, 135 deletions

diff --git a/src/mongo/db/exec/sbe/vm/vm.cpp b/src/mongo/db/exec/sbe/vm/vm.cpp
index cd2f781a2ee..d4096ee646b 100644
--- a/src/mongo/db/exec/sbe/vm/vm.cpp
+++ b/src/mongo/db/exec/sbe/vm/vm.cpp
@@ -1104,7 +1104,8 @@ std::tuple<bool, value::TypeTags, value::Value> ByteCode::builtinStdDevSampFinal
 std::tuple<bool, value::TypeTags, value::Value> ByteCode::aggMin(value::TypeTags accTag,
                                                                  value::Value accValue,
                                                                  value::TypeTags fieldTag,
-                                                                 value::Value fieldValue) {
+                                                                 value::Value fieldValue,
+                                                                 CollatorInterface* collator) {
     // Skip aggregation step if we don't have the input.
     if (fieldTag == value::TypeTags::Nothing) {
         auto [tag, val] = value::copyValue(accTag, accValue);
@@ -1117,7 +1118,7 @@ std::tuple<bool, value::TypeTags, value::Value> ByteCode::aggMin(value::TypeTags
         return {true, tag, val};
     }
 
-    auto [tag, val] = compare3way(accTag, accValue, fieldTag, fieldValue);
+    auto [tag, val] = compare3way(accTag, accValue, fieldTag, fieldValue, collator);
 
     if (tag == value::TypeTags::NumberInt32 && value::bitcastTo<int>(val) < 0) {
         auto [tag, val] = value::copyValue(accTag, accValue);
@@ -1128,42 +1129,12 @@ std::tuple<bool, value::TypeTags, value::Value> ByteCode::aggMin(value::TypeTags
     }
 }
 
-std::tuple<bool, value::TypeTags, value::Value> ByteCode::aggCollMin(value::TypeTags accTag,
-                                                                     value::Value accValue,
-                                                                     value::TypeTags collTag,
-                                                                     value::Value collValue,
-                                                                     value::TypeTags fieldTag,
-                                                                     value::Value fieldValue) {
-    // Skip aggregation step if we don't have the input or if the collation is Nothing or an
-    // unexpected type.
-    if (fieldTag == value::TypeTags::Nothing || collTag != value::TypeTags::collator) {
-        auto [tag, val] = value::copyValue(accTag, accValue);
-        return {true, tag, val};
-    }
-
-    // Initialize the accumulator.
-    if (accTag == value::TypeTags::Nothing) {
-        auto [tag, val] = value::copyValue(fieldTag, fieldValue);
-        return {true, tag, val};
-    }
-
-    auto collator = value::getCollatorView(collValue);
-
-    auto [tag, val] = genericCompare<std::less<>>(accTag, accValue, fieldTag, fieldValue, collator);
-
-    if (tag == value::TypeTags::Boolean && value::bitcastTo<bool>(val)) {
-        auto [tag, val] = value::copyValue(accTag, accValue);
-        return {true, tag, val};
-    } else {
-        auto [tag, val] = value::copyValue(fieldTag, fieldValue);
-        return {true, tag, val};
-    }
-}
 
 std::tuple<bool, value::TypeTags, value::Value> ByteCode::aggMax(value::TypeTags accTag,
                                                                  value::Value accValue,
                                                                  value::TypeTags fieldTag,
-                                                                 value::Value fieldValue) {
+                                                                 value::Value fieldValue,
+                                                                 CollatorInterface* collator) {
     // Skip aggregation step if we don't have the input.
     if (fieldTag == value::TypeTags::Nothing) {
         auto [tag, val] = value::copyValue(accTag, accValue);
@@ -1176,7 +1147,7 @@ std::tuple<bool, value::TypeTags, value::Value> ByteCode::aggMax(value::TypeTags
         return {true, tag, val};
     }
 
-    auto [tag, val] = compare3way(accTag, accValue, fieldTag, fieldValue);
+    auto [tag, val] = compare3way(accTag, accValue, fieldTag, fieldValue, collator);
 
     if (tag == value::TypeTags::NumberInt32 && value::bitcastTo<int>(val) > 0) {
         auto [tag, val] = value::copyValue(accTag, accValue);
@@ -1187,39 +1158,6 @@ std::tuple<bool, value::TypeTags, value::Value> ByteCode::aggMax(value::TypeTags
     }
 }
 
-std::tuple<bool, value::TypeTags, value::Value> ByteCode::aggCollMax(value::TypeTags accTag,
-                                                                     value::Value accValue,
-                                                                     value::TypeTags collTag,
-                                                                     value::Value collValue,
-                                                                     value::TypeTags fieldTag,
-                                                                     value::Value fieldValue) {
-    // Skip aggregation step if we don't have the input or if the collation is Nothing or an
-    // unexpected type.
-    if (fieldTag == value::TypeTags::Nothing || collTag != value::TypeTags::collator) {
-        auto [tag, val] = value::copyValue(accTag, accValue);
-        return {true, tag, val};
-    }
-
-    // Initialize the accumulator.
-    if (accTag == value::TypeTags::Nothing) {
-        auto [tag, val] = value::copyValue(fieldTag, fieldValue);
-        return {true, tag, val};
-    }
-
-    auto collator = value::getCollatorView(collValue);
-
-    auto [tag, val] =
-        genericCompare<std::greater<>>(accTag, accValue, fieldTag, fieldValue, collator);
-
-    if (tag == value::TypeTags::Boolean && value::bitcastTo<bool>(val)) {
-        auto [tag, val] = value::copyValue(accTag, accValue);
-        return {true, tag, val};
-    } else {
-        auto [tag, val] = value::copyValue(fieldTag, fieldValue);
-        return {true, tag, val};
-    }
-}
-
 std::tuple<bool, value::TypeTags, value::Value> ByteCode::aggFirst(value::TypeTags accTag,
                                                                    value::Value accValue,
                                                                    value::TypeTags fieldTag,
@@ -4735,133 +4673,147 @@ void ByteCode::runInternal(const CodeFragment* code, int64_t position) {
                     break;
                 }
                 case Instruction::aggSum: {
-                    auto [rhsOwned, rhsTag, rhsVal] = getFromStack(0);
+                    auto [fieldOwned, fieldTag, fieldVal] = getFromStack(0);
                     popStack();
-                    auto [lhsOwned, lhsTag, lhsVal] = getFromStack(0);
+                    auto [accOwned, accTag, accVal] = getFromStack(0);
 
-                    auto [owned, tag, val] = aggSum(lhsTag, lhsVal, rhsTag, rhsVal);
+                    auto [owned, tag, val] = aggSum(accTag, accVal, fieldTag, fieldVal);
 
                     topStack(owned, tag, val);
 
-                    if (rhsOwned) {
-                        value::releaseValue(rhsTag, rhsVal);
+                    if (fieldOwned) {
+                        value::releaseValue(fieldTag, fieldVal);
                     }
-                    if (lhsOwned) {
-                        value::releaseValue(lhsTag, lhsVal);
+                    if (accOwned) {
+                        value::releaseValue(accTag, accVal);
                     }
                     break;
                 }
                 case Instruction::aggMin: {
-                    auto [rhsOwned, rhsTag, rhsVal] = getFromStack(0);
+                    auto [fieldOwned, fieldTag, fieldVal] = getFromStack(0);
                     popStack();
-                    auto [lhsOwned, lhsTag, lhsVal] = getFromStack(0);
+                    auto [accOwned, accTag, accVal] = getFromStack(0);
 
-                    auto [owned, tag, val] = aggMin(lhsTag, lhsVal, rhsTag, rhsVal);
+                    auto [owned, tag, val] = aggMin(accTag, accVal, fieldTag, fieldVal);
 
                     topStack(owned, tag, val);
 
-                    if (rhsOwned) {
-                        value::releaseValue(rhsTag, rhsVal);
+                    if (fieldOwned) {
+                        value::releaseValue(fieldTag, fieldVal);
                     }
-                    if (lhsOwned) {
-                        value::releaseValue(lhsTag, lhsVal);
+                    if (accOwned) {
+                        value::releaseValue(accTag, accVal);
                     }
                     break;
                 }
                 case Instruction::aggCollMin: {
-                    auto [rhsOwned, rhsTag, rhsVal] = getFromStack(0);
+                    auto [fieldOwned, fieldTag, fieldVal] = getFromStack(0);
                     popStack();
                     auto [collOwned, collTag, collVal] = getFromStack(0);
                     popStack();
-                    auto [lhsOwned, lhsTag, lhsVal] = getFromStack(0);
+                    auto [accOwned, accTag, accVal] = getFromStack(0);
 
-                    auto [owned, tag, val] =
-                        aggCollMin(lhsTag, lhsVal, collTag, collVal, rhsTag, rhsVal);
+                    // Skip aggregation step if the collation is Nothing or an unexpected type.
+                    if (collTag != value::TypeTags::collator) {
+                        auto [tag, val] = value::copyValue(accTag, accVal);
+                        topStack(true, tag, val);
+                        break;
+                    }
+                    auto collator = value::getCollatorView(collVal);
+
+                    auto [owned, tag, val] = aggMin(accTag, accVal, fieldTag, fieldVal, collator);
 
                     topStack(owned, tag, val);
 
-                    if (rhsOwned) {
-                        value::releaseValue(rhsTag, rhsVal);
+                    if (fieldOwned) {
+                        value::releaseValue(fieldTag, fieldVal);
                     }
                     if (collOwned) {
                         value::releaseValue(collTag, collVal);
                     }
-                    if (lhsOwned) {
-                        value::releaseValue(lhsTag, lhsVal);
+                    if (accOwned) {
+                        value::releaseValue(accTag, accVal);
                     }
                     break;
                 }
                 case Instruction::aggMax: {
-                    auto [rhsOwned, rhsTag, rhsVal] = getFromStack(0);
+                    auto [fieldOwned, fieldTag, fieldVal] = getFromStack(0);
                     popStack();
-                    auto [lhsOwned, lhsTag, lhsVal] = getFromStack(0);
+                    auto [accOwned, accTag, accVal] = getFromStack(0);
 
-                    auto [owned, tag, val] = aggMax(lhsTag, lhsVal, rhsTag, rhsVal);
+                    auto [owned, tag, val] = aggMax(accTag, accVal, fieldTag, fieldVal);
 
                     topStack(owned, tag, val);
 
-                    if (rhsOwned) {
-                        value::releaseValue(rhsTag, rhsVal);
+                    if (fieldOwned) {
+                        value::releaseValue(fieldTag, fieldVal);
                     }
-                    if (lhsOwned) {
-                        value::releaseValue(lhsTag, lhsVal);
+                    if (accOwned) {
+                        value::releaseValue(accTag, accVal);
                     }
                     break;
                 }
                 case Instruction::aggCollMax: {
-                    auto [rhsOwned, rhsTag, rhsVal] = getFromStack(0);
+                    auto [fieldOwned, fieldTag, fieldVal] = getFromStack(0);
                     popStack();
                     auto [collOwned, collTag, collVal] = getFromStack(0);
                     popStack();
-                    auto [lhsOwned, lhsTag, lhsVal] = getFromStack(0);
+                    auto [accOwned, accTag, accVal] = getFromStack(0);
 
-                    auto [owned, tag, val] =
-                        aggCollMax(lhsTag, lhsVal, collTag, collVal, rhsTag, rhsVal);
+                    // Skip aggregation step if the collation is Nothing or an unexpected type.
+                    if (collTag != value::TypeTags::collator) {
+                        auto [tag, val] = value::copyValue(accTag, accVal);
+                        topStack(true, tag, val);
+                        break;
+                    }
+                    auto collator = value::getCollatorView(collVal);
+
+                    auto [owned, tag, val] = aggMax(accTag, accVal, fieldTag, fieldVal, collator);
 
                     topStack(owned, tag, val);
 
-                    if (rhsOwned) {
-                        value::releaseValue(rhsTag, rhsVal);
+                    if (fieldOwned) {
+                        value::releaseValue(fieldTag, fieldVal);
                     }
                     if (collOwned) {
                         value::releaseValue(collTag, collVal);
                     }
-                    if (lhsOwned) {
-                        value::releaseValue(lhsTag, lhsVal);
+                    if (accOwned) {
+                        value::releaseValue(accTag, accVal);
                     }
                     break;
                 }
                 case Instruction::aggFirst: {
-                    auto [rhsOwned, rhsTag, rhsVal] = getFromStack(0);
+                    auto [fieldOwned, fieldTag, fieldVal] = getFromStack(0);
                     popStack();
-                    auto [lhsOwned, lhsTag, lhsVal] = getFromStack(0);
+                    auto [accOwned, accTag, accVal] = getFromStack(0);
 
-                    auto [owned, tag, val] = aggFirst(lhsTag, lhsVal, rhsTag, rhsVal);
+                    auto [owned, tag, val] = aggFirst(accTag, accVal, fieldTag, fieldVal);
 
                     topStack(owned, tag, val);
 
-                    if (rhsOwned) {
-                        value::releaseValue(rhsTag, rhsVal);
+                    if (fieldOwned) {
+                        value::releaseValue(fieldTag, fieldVal);
                     }
-                    if (lhsOwned) {
-                        value::releaseValue(lhsTag, lhsVal);
+                    if (accOwned) {
+                        value::releaseValue(accTag, accVal);
                     }
                     break;
                 }
                 case Instruction::aggLast: {
-                    auto [rhsOwned, rhsTag, rhsVal] = getFromStack(0);
+                    auto [fieldOwned, fieldTag, fieldVal] = getFromStack(0);
                     popStack();
-                    auto [lhsOwned, lhsTag, lhsVal] = getFromStack(0);
+                    auto [accOwned, accTag, accVal] = getFromStack(0);
 
-                    auto [owned, tag, val] = aggLast(lhsTag, lhsVal, rhsTag, rhsVal);
+                    auto [owned, tag, val] = aggLast(accTag, accVal, fieldTag, fieldVal);
 
                     topStack(owned, tag, val);
 
-                    if (rhsOwned) {
-                        value::releaseValue(rhsTag, rhsVal);
+                    if (fieldOwned) {
+                        value::releaseValue(fieldTag, fieldVal);
                     }
-                    if (lhsOwned) {
-                        value::releaseValue(lhsTag, lhsVal);
+                    if (accOwned) {
+                        value::releaseValue(accTag, accVal);
                     }
                     break;
                 }
diff --git a/src/mongo/db/exec/sbe/vm/vm.h b/src/mongo/db/exec/sbe/vm/vm.h
index b457288bd1d..fb1042174b4 100644
--- a/src/mongo/db/exec/sbe/vm/vm.h
+++ b/src/mongo/db/exec/sbe/vm/vm.h
@@ -851,12 +851,14 @@ private:
     std::tuple<bool, value::TypeTags, value::Value> aggMin(value::TypeTags accTag,
                                                            value::Value accValue,
                                                            value::TypeTags fieldTag,
-                                                           value::Value fieldValue);
+                                                           value::Value fieldValue,
+                                                           CollatorInterface* collator = nullptr);
 
     std::tuple<bool, value::TypeTags, value::Value> aggMax(value::TypeTags accTag,
                                                            value::Value accValue,
                                                            value::TypeTags fieldTag,
-                                                           value::Value fieldValue);
+                                                           value::Value fieldValue,
+                                                           CollatorInterface* collator = nullptr);
 
     std::tuple<bool, value::TypeTags, value::Value> aggFirst(value::TypeTags accTag,
                                                              value::Value accValue,
@@ -868,20 +870,6 @@ private:
                                                             value::TypeTags fieldTag,
                                                             value::Value fieldValue);
 
-    std::tuple<bool, value::TypeTags, value::Value> aggCollMin(value::TypeTags accTag,
-                                                               value::Value accValue,
-                                                               value::TypeTags collTag,
-                                                               value::Value collValue,
-                                                               value::TypeTags fieldTag,
-                                                               value::Value fieldValue);
-
-    std::tuple<bool, value::TypeTags, value::Value> aggCollMax(value::TypeTags accTag,
-                                                               value::Value accValue,
-                                                               value::TypeTags collTag,
-                                                               value::Value collValue,
-                                                               value::TypeTags fieldTag,
-                                                               value::Value fieldValue);
-
     std::tuple<bool, value::TypeTags, value::Value> genericAcos(value::TypeTags operandTag,
                                                                 value::Value operandValue);
     std::tuple<bool, value::TypeTags, value::Value> genericAcosh(value::TypeTags operandTag,
diff --git a/src/mongo/db/query/sbe_stage_builder_accumulator_test.cpp b/src/mongo/db/query/sbe_stage_builder_accumulator_test.cpp
index 587b6d1d3c9..a42ced91e2b 100644
--- a/src/mongo/db/query/sbe_stage_builder_accumulator_test.cpp
+++ b/src/mongo/db/query/sbe_stage_builder_accumulator_test.cpp
@@ -419,6 +419,19 @@ TEST_F(SbeStageBuilderGroupTest, MinAccumulatorTranslationStringsReverseStringCo
         std::make_unique<CollatorInterfaceMock>(CollatorInterfaceMock::MockType::kReverseString));
 }
 
+TEST_F(SbeStageBuilderGroupTest, MinAccumulatorTranslationNaN) {
+    auto docs = std::vector<BSONArray>{
+        BSON_ARRAY(BSON("a" << 1 << "b" << 42ll)),
+        BSON_ARRAY(BSON("a" << 1 << "b" << std::numeric_limits<double>::quiet_NaN()))};
+    // The collator doesn't affect comparison of numbers but until SERVER-61868 just providing a
+    // collator used to trigger a different codepath, so let's throw it in for a good measure.
+    runGroupAggregationTest(
+        "{_id: null, x: {$min: '$b'}}",
+        docs,
+        BSON_ARRAY(BSON("_id" << BSONNULL << "x" << std::numeric_limits<double>::quiet_NaN())),
+        std::make_unique<CollatorInterfaceMock>(CollatorInterfaceMock::MockType::kReverseString));
+}
+
 TEST_F(SbeStageBuilderGroupTest, MaxAccumulatorTranslationBasic) {
     auto docs = std::vector<BSONArray>{BSON_ARRAY(BSON("a" << 1 << "b" << 100ll)),
                                        BSON_ARRAY(BSON("a" << 1 << "b" << Decimal128(10.0))),
@@ -503,6 +516,19 @@ TEST_F(SbeStageBuilderGroupTest, MaxAccumulatorTranslationStringsReverseStringCo
         std::make_unique<CollatorInterfaceMock>(CollatorInterfaceMock::MockType::kReverseString));
 }
 
+TEST_F(SbeStageBuilderGroupTest, MaxAccumulatorTranslationNaN) {
+    auto docs = std::vector<BSONArray>{
+        BSON_ARRAY(BSON("a" << 1 << "b" << 42ll)),
+        BSON_ARRAY(BSON("a" << 1 << "b" << std::numeric_limits<double>::quiet_NaN()))};
+    // The collator doesn't affect comparison of numbers but until SERVER-61868 just providing a
+    // collator used to trigger a different codepath, so let's throw it in for a good measure.
+    runGroupAggregationTest(
+        "{_id: null, x: {$max: '$b'}}",
+        docs,
+        BSON_ARRAY(BSON("_id" << BSONNULL << "x" << 42ll)),
+        std::make_unique<CollatorInterfaceMock>(CollatorInterfaceMock::MockType::kReverseString));
+}
+
 TEST_F(SbeStageBuilderGroupTest, FirstAccumulatorTranslationOneDoc) {
     auto docs = std::vector<BSONArray>{BSON_ARRAY(BSON("a" << 1 << "b" << BSON_ARRAY(1 << 2)))};
     runGroupAggregationTest("{_id: null, x: {$first: '$b'}}",