SERVER-58892 Pass through documents with a null densification value

3 files changed, 87 insertions, 9 deletions

diff --git a/jstests/aggregation/sources/densify/explicit_range.js b/jstests/aggregation/sources/densify/explicit_range.js
index 94031f9ed24..6f7e2d3d3cd 100644
--- a/jstests/aggregation/sources/densify/explicit_range.js
+++ b/jstests/aggregation/sources/densify/explicit_range.js
@@ -83,6 +83,32 @@ for (let i = 0; i < densifyUnits.length; i++) {
             {base, min: 0, max: 50, pred: i => i % 3 == 0 || i % 7 == 0, addFunc: add, coll: coll});
 
         runDensifyRangeTest({step, bounds: [10, 45]});
+
+        // Lots of off-step documents with nulls sprinkled in to confirm that a null value is
+        // treated the same as a missing value.
+        coll.drop();
+        insertDocumentsOnPredicate(
+            {base, min: 0, max: 10, pred: i => i % 3 == 0 || i % 7 == 0, addFunc: add, coll: coll});
+        coll.insert({val: null});
+        insertDocumentsOnPredicate({
+            base,
+            min: 10,
+            max: 20,
+            pred: i => i % 3 == 0 || i % 7 == 0,
+            addFunc: add,
+            coll: coll
+        });
+        coll.insert({val: null});
+        coll.insert({blah: base});  // Missing "val" key.
+        insertDocumentsOnPredicate({
+            base,
+            min: 20,
+            max: 50,
+            pred: i => i % 3 == 0 || i % 7 == 0,
+            addFunc: add,
+            coll: coll
+        });
+        runDensifyRangeTest({step, bounds: [10, 45]});
     }
 }
 })();
diff --git a/jstests/aggregation/sources/densify/full_range.js b/jstests/aggregation/sources/densify/full_range.js
index 1443a4f43ec..0fc8d247288 100644
--- a/jstests/aggregation/sources/densify/full_range.js
+++ b/jstests/aggregation/sources/densify/full_range.js
@@ -50,6 +50,33 @@ for (let i = 0; i < densifyUnits.length; i++) {
         insertDocumentsOnPredicate(
             {base, min: 0, max: 50, pred: i => i % 3 == 0 || i % 7 == 0, addFunc: add, coll: coll});
         runDensifyFullTest();
+
+        // Lots of off-step documents with nulls sprinkled in to confirm that a null value is
+        // treated the same as a missing value.
+        coll.drop();
+        insertDocumentsOnPredicate(
+            {base, min: 0, max: 10, pred: i => i % 3 == 0 || i % 7 == 0, addFunc: add, coll: coll});
+        coll.insert({val: null});
+        insertDocumentsOnPredicate({
+            base,
+            min: 10,
+            max: 20,
+            pred: i => i % 3 == 0 || i % 7 == 0,
+            addFunc: add,
+            coll: coll
+        });
+        coll.insert({val: null});
+        coll.insert({blah: base});  // Missing "val" key.
+        insertDocumentsOnPredicate({
+            base,
+            min: 20,
+            max: 25,
+            pred: i => i % 3 == 0 || i % 7 == 0,
+            addFunc: add,
+            coll: coll
+        });
+
+        runDensifyFullTest();
     }
 }
 })();
diff --git a/jstests/aggregation/sources/densify/libs/densify_in_js.js b/jstests/aggregation/sources/densify/libs/densify_in_js.js
index a89f03bfd23..48b63298499 100644
--- a/jstests/aggregation/sources/densify/libs/densify_in_js.js
+++ b/jstests/aggregation/sources/densify/libs/densify_in_js.js
@@ -91,7 +91,20 @@ function densifyInJS(stage, docs) {
     const {step, bounds, unit} = stage.range;
     const stream = [];
 
-    docs.sort((a, b) => a[field] - b[field]);
+    // $densify is translated into a $sort on `field` and then $internalDensify, so replicate that
+    // behavior here by sorting the array of documents by the field.
+    docs.sort((a, b) => {
+        if (a[field] == null && b[field] == null) {
+            return 0;
+        } else if (a[field] == null) {  // null << any value.
+            return -1;
+        } else if (b[field] == null) {
+            return 1;
+        } else {
+            return a[field] - b[field];
+        }
+    });
+    const docsWithoutNulls = docs.filter(doc => doc[field] != null);
 
     const {add, sub, getNextStepFromBase} = getArithmeticFunctionsForUnit(unit);
 
@@ -110,27 +123,35 @@ function densifyInJS(stage, docs) {
         if (docs.length == 0) {
             return stream;
         }
-        return densifyInJS({
-            field: stage.field,
-            range: {step, unit, bounds: [docs[0][field], docs[docs.length - 1][field]]}
-        },
+        const minValue = docsWithoutNulls[0][field];
+        const maxValue = docsWithoutNulls[docsWithoutNulls.length - 1][field];
+        return densifyInJS({field: stage.field, range: {step, unit, bounds: [minValue, maxValue]}},
                            docs);
 
     } else if (bounds === "partition") {
         throw new Error("Partitioning not supported by JS densify.");
     } else if (bounds.length == 2) {
         const [lower, upper] = bounds;
-        let currentVal =
-            docs.length > 0 ? Math.min(docs[0][field], sub(lower, step)) : sub(lower, step);
+        let currentVal = docsWithoutNulls.length > 0
+            ? Math.min(docsWithoutNulls[0], sub(lower, step))
+            : sub(lower, step);
         for (let i = 0; i < docs.length; i++) {
             const nextVal = docs[i][field];
+            if (nextVal === null || nextVal === undefined) {
+                // If the next value in the stream is missing or null, let the doc pass through
+                // without modifying anything else.
+                stream.push(docs[i]);
+                continue;
+            }
             stream.push(...generateDocuments(getNextStepFromBase(currentVal, lower, step),
                                              nextVal,
                                              (val) => val >= lower && val < upper));
             stream.push(docs[i]);
             currentVal = nextVal;
         }
-        const lastVal = docs.length > 0 ? docs[docs.length - 1][field] : sub(lower, step);
+        const lastVal = docsWithoutNulls.length > 0
+            ? docsWithoutNulls[docsWithoutNulls.length - 1][field]
+            : sub(lower, step);
         if (lastVal < upper) {
             stream.push(...generateDocuments(getNextStepFromBase(currentVal, lower, step), upper));
         }
@@ -161,6 +182,10 @@ const densifyUnits = [null, "millisecond", "second", "day", "month", "quarter",
 
 const interestingSteps = [1, 2, 3, 4, 5, 7, 11, 13];
 
+function buildErrorString(found, expected) {
+    return "Expected:\n" + tojson(expected) + "\nGot:\n" + tojson(found);
+}
+
 function testDensifyStage(stage, coll, msg) {
     if (stage.range.unit === null) {
         delete stage.range.unit;
@@ -168,5 +193,5 @@ function testDensifyStage(stage, coll, msg) {
     const result = coll.aggregate([{"$densify": stage}]).toArray();
     const expected = densifyInJS(stage, coll.find({}).toArray());
     const newMsg = (msg || "") + " | stage: " + tojson(stage);
-    assert.eq(expected, result, newMsg);
+    assert(arrayEq(expected, result), newMsg + buildErrorString(result, expected));
 }