SERVER-71763 CSI planning heuristics enhancements

1 files changed, 83 insertions, 32 deletions

diff --git a/src/mongo/db/query/query_planner_columnar_test.cpp b/src/mongo/db/query/query_planner_columnar_test.cpp
index 226dc60464c..8c922259dfa 100644
--- a/src/mongo/db/query/query_planner_columnar_test.cpp
+++ b/src/mongo/db/query/query_planner_columnar_test.cpp
@@ -68,13 +68,24 @@ protected:
         params.collectionStats.approximateDataSizeBytes = 100000;
         internalQueryColumnScanMinCollectionSizeBytes.store(0);
         internalQueryColumnScanMinAvgDocSizeBytes.store(0);
+        internalQueryColumnScanMinNumColumnFilters.store(0);
     }
 
     void tearDown() final {
+        resetPlannerHeuristics();
+    }
+
+    void resetPlannerHeuristics() {
         internalQueryMaxNumberOfFieldsToChooseUnfilteredColumnScan.store(
             kInternalQueryMaxNumberOfFieldsToChooseUnfilteredColumnScanDefault);
         internalQueryMaxNumberOfFieldsToChooseFilteredColumnScan.store(
             kInternalQueryMaxNumberOfFieldsToChooseFilteredColumnScanDefault);
+        internalQueryColumnScanMinCollectionSizeBytes.store(
+            kInternalQueryColumnScanMinCollectionSizeBytesDefault);
+        internalQueryColumnScanMinAvgDocSizeBytes.store(
+            kInternalQueryColumnScanMinAvgDocSizeBytesDefault);
+        internalQueryColumnScanMinNumColumnFilters.store(
+            kInternalQueryColumnScanMinNumColumnFiltersDefault);
     }
 
     void addColumnStoreIndexAndEnableFilterSplitting(bool genPerColFilter = true,
@@ -112,11 +123,11 @@ protected:
                                                                            columnstoreProjection);
     }
 
-    long long collectionSizeBytes() {
+    double collectionSizeBytes() {
         return params.collectionStats.approximateDataSizeBytes;
     }
 
-    long long avgDocumentSizeBytes() {
+    double avgDocumentSizeBytes() {
         return static_cast<double>(params.collectionStats.approximateDataSizeBytes) /
             params.collectionStats.noOfRecords;
     }
@@ -1263,62 +1274,92 @@ TEST_F(QueryPlannerColumnarTest, ColumnIndexForCountWithPostAssemblyFilter) {
     })");
 }
 
-TEST_F(QueryPlannerColumnarTest, AvgDocSizeTooSmall) {
+TEST_F(QueryPlannerColumnarTest, PlanningHeuristics_NotMet) {
     addColumnStoreIndexAndEnableFilterSplitting();
-    internalQueryColumnScanMinAvgDocSizeBytes.store(avgDocumentSizeBytes() + 1);
+    resetPlannerHeuristics();
+    params.availableMemoryBytes = 10 * 1024;
 
-    runQuerySortProj(BSONObj(), BSONObj(), BSON("a" << 1));
+    // Update the collection's stats to just below the expected defaults.
+    params.collectionStats.approximateDataSizeBytes = params.availableMemoryBytes - 1;
+    params.collectionStats.noOfRecords =
+        collectionSizeBytes() / internalQueryColumnScanMinAvgDocSizeBytes.load() + 1;
+
+    runQuerySortProj(BSON("a" << 1), BSONObj(), BSON("a" << 1));
 
     assertNumSolutions(1U);
     assertSolutionExists(R"({proj: {spec: {a: 1}, node: {cscan: {dir: 1}}}})");
 }
 
-TEST_F(QueryPlannerColumnarTest, AvgDocSizeLargeEnough) {
+TEST_F(QueryPlannerColumnarTest, PlanningHeuristics_AvgDocSizeLargeEnough) {
     addColumnStoreIndexAndEnableFilterSplitting();
-    internalQueryColumnScanMinAvgDocSizeBytes.store(avgDocumentSizeBytes());
+    resetPlannerHeuristics();
+    params.availableMemoryBytes = 10 * 1024;
 
-    runQuerySortProj(BSONObj(), BSONObj(), BSON("a" << 1));
+    params.collectionStats.approximateDataSizeBytes = params.availableMemoryBytes - 1;
+    params.collectionStats.noOfRecords =
+        collectionSizeBytes() / internalQueryColumnScanMinAvgDocSizeBytes.load();
+
+    runQuerySortProj(BSON("a" << 1), BSONObj(), BSON("a" << 1));
 
     assertNumSolutions(1U);
     assertSolutionExists(R"({
         column_scan: {
-            filtersByPath: {},
+            filtersByPath: {a: {a: {$eq: 1}}},
             outputFields: ['a', '_id'],
-            matchFields: []
+            matchFields: ['a']
         }
     })");
 }
 
-TEST_F(QueryPlannerColumnarTest, CollectionTooSmall) {
+TEST_F(QueryPlannerColumnarTest, PlanningHeuristics_CollectionLargeEnough) {
     addColumnStoreIndexAndEnableFilterSplitting();
-    internalQueryColumnScanMinCollectionSizeBytes.store(collectionSizeBytes() + 1);
+    resetPlannerHeuristics();
+    params.availableMemoryBytes = 10 * 1024;
 
-    runQuerySortProj(BSONObj(), BSONObj(), BSON("a" << 1));
+    params.collectionStats.approximateDataSizeBytes = params.availableMemoryBytes;
+    params.collectionStats.noOfRecords =
+        collectionSizeBytes() / internalQueryColumnScanMinAvgDocSizeBytes.load() + 1;
+
+    runQuerySortProj(BSON("a" << 1), BSONObj(), BSON("a" << 1));
 
     assertNumSolutions(1U);
-    assertSolutionExists(R"({proj: {spec: {a: 1}, node: {cscan: {dir: 1}}}})");
+    assertSolutionExists(R"({
+        column_scan: {
+            filtersByPath: {a: {a: {$eq: 1}}},
+            outputFields: ['a', '_id'],
+            matchFields: ['a']
+        }
+    })");
 }
 
-TEST_F(QueryPlannerColumnarTest, CollectionLargeEnough) {
+TEST_F(QueryPlannerColumnarTest, PlanningHeuristics_EnoughColumnFilters) {
     addColumnStoreIndexAndEnableFilterSplitting();
-    internalQueryColumnScanMinCollectionSizeBytes.store(collectionSizeBytes());
+    resetPlannerHeuristics();
+    params.availableMemoryBytes = 10 * 1024;
 
-    runQuerySortProj(BSONObj(), BSONObj(), BSON("a" << 1));
+    params.collectionStats.approximateDataSizeBytes = params.availableMemoryBytes - 1;
+    params.collectionStats.noOfRecords =
+        collectionSizeBytes() / internalQueryColumnScanMinAvgDocSizeBytes.load() + 1;
+
+    runQuerySortProj(
+        BSON("a" << 3 << "b" << 4 << "c" << 5), BSONObj(), BSON("a" << 1 << "b" << 1 << "c" << 1));
 
     assertNumSolutions(1U);
     assertSolutionExists(R"({
         column_scan: {
-            filtersByPath: {},
-            outputFields: ['a', '_id'],
-            matchFields: []
+            filtersByPath: {a: {a: {$eq: 3}}, b: {b: {$eq: 4}}, c: {c: {$eq: 5}}},
+            outputFields: ['_id', 'a', 'b', 'c'],
+            matchFields: ['a', 'b', 'c']
         }
     })");
 }
 
-TEST_F(QueryPlannerColumnarTest, EmptyCollectionWithAvgDocSizeThreshold) {
+TEST_F(QueryPlannerColumnarTest, PlanningHeuristics_EmptyCollection) {
     addColumnStoreIndexAndEnableFilterSplitting();
-    // Set a non-zero threshold.
-    internalQueryColumnScanMinAvgDocSizeBytes.store(avgDocumentSizeBytes());
+    // Set non-zero thresholds.
+    internalQueryColumnScanMinCollectionSizeBytes.store(1);
+    internalQueryColumnScanMinAvgDocSizeBytes.store(1);
+    internalQueryColumnScanMinNumColumnFilters.store(1);
     // Update the collection's stats to be zero/empty.
     params.collectionStats.noOfRecords = 0;
     params.collectionStats.approximateDataSizeBytes = 0;
@@ -1329,18 +1370,28 @@ TEST_F(QueryPlannerColumnarTest, EmptyCollectionWithAvgDocSizeThreshold) {
     assertSolutionExists(R"({proj: {spec: {a: 1}, node: {cscan: {dir: 1}}}})");
 }
 
-TEST_F(QueryPlannerColumnarTest, EmptyCollectionWithCollectionSizeThreshold) {
-    addColumnStoreIndexAndEnableFilterSplitting();
-    // Set a non-zero threshold.
-    internalQueryColumnScanMinCollectionSizeBytes.store(collectionSizeBytes());
-    // Update the collection's stats to be zero/empty.
-    params.collectionStats.noOfRecords = 0;
-    params.collectionStats.approximateDataSizeBytes = 0;
+TEST_F(QueryPlannerColumnarTest, PlanningHeuristics_HintOverridesHeuristics) {
+    addColumnStoreIndexAndEnableFilterSplitting(true, "csi");
+    internalQueryColumnScanMinCollectionSizeBytes.store(collectionSizeBytes() + 1);
+    internalQueryColumnScanMinAvgDocSizeBytes.store(avgDocumentSizeBytes() + 1);
+    internalQueryColumnScanMinNumColumnFilters.store(1);
 
-    runQuerySortProj(BSONObj(), BSONObj(), BSON("a" << 1));
+    runQuerySortProjSkipLimitHint(BSONObj(),
+                                  BSONObj(),
+                                  BSON("a" << 1),
+                                  0,
+                                  0,
+                                  BSON("$hint"
+                                       << "csi"));
 
     assertNumSolutions(1U);
-    assertSolutionExists(R"({proj: {spec: {a: 1}, node: {cscan: {dir: 1}}}})");
+    assertSolutionExists(R"({
+        column_scan: {
+            filtersByPath: {},
+            outputFields: ['a', '_id'],
+            matchFields: []
+        }
+    })");
 }
 
 TEST_F(QueryPlannerColumnarTest, HintIndexWithNonStandardKeyPattern) {