From 7832d3c9fe50a47ea9f3cf0f29b72d021b804141 Mon Sep 17 00:00:00 2001
From: Shintaro Murakami <mrkm4ntr@gmail.com>
Date: Mon, 29 Jan 2024 17:11:42 +0900
Subject: [PATCH] Propagate topLevelScoringClause from QueryProfiler (#13031)

---
 lucene/CHANGES.txt                            |  2 +
 .../sandbox/search/QueryProfilerWeight.java   |  5 ++
 .../search/TestQueryProfilerWeight.java       | 60 +++++++++++++++----
 3 files changed, 54 insertions(+), 13 deletions(-)

diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
index fa460204306..582ecd966ab 100644
--- a/lucene/CHANGES.txt
+++ b/lucene/CHANGES.txt
@@ -240,6 +240,8 @@ Bug Fixes
 
 * GITHUB#12287: Fix a bug in ShapeTestUtil. (Heemin Kim)
 
+* GITHUB#13031: ScorerSupplier created by QueryProfilerWeight will propagate topLevelScoringClause to the sub ScorerSupplier. (Shintaro Murakami)
+
 Build
 ---------------------
 
diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/search/QueryProfilerWeight.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/search/QueryProfilerWeight.java
index 8caaac02139..47c16c4dc22 100644
--- a/lucene/sandbox/src/java/org/apache/lucene/sandbox/search/QueryProfilerWeight.java
+++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/search/QueryProfilerWeight.java
@@ -94,6 +94,11 @@ class QueryProfilerWeight extends FilterWeight {
           timer.stop();
         }
       }
+
+      @Override
+      public void setTopLevelScoringClause() throws IOException {
+        subQueryScorerSupplier.setTopLevelScoringClause();
+      }
     };
   }
 
diff --git a/lucene/sandbox/src/test/org/apache/lucene/sandbox/search/TestQueryProfilerWeight.java b/lucene/sandbox/src/test/org/apache/lucene/sandbox/search/TestQueryProfilerWeight.java
index 41b4007155b..7c9e6736cc5 100644
--- a/lucene/sandbox/src/test/org/apache/lucene/sandbox/search/TestQueryProfilerWeight.java
+++ b/lucene/sandbox/src/test/org/apache/lucene/sandbox/search/TestQueryProfilerWeight.java
@@ -29,6 +29,7 @@ import org.apache.lucene.search.Matches;
 import org.apache.lucene.search.MatchesIterator;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.Scorer;
+import org.apache.lucene.search.ScorerSupplier;
 import org.apache.lucene.search.Weight;
 import org.apache.lucene.tests.util.LuceneTestCase;
 
@@ -45,26 +46,49 @@ public class TestQueryProfilerWeight extends LuceneTestCase {
     }
 
     @Override
-    public Scorer scorer(LeafReaderContext context) {
-      return new Scorer(this) {
+    public Scorer scorer(LeafReaderContext context) throws IOException {
+      return scorerSupplier(context).get(Long.MAX_VALUE);
+    }
+
+    @Override
+    public ScorerSupplier scorerSupplier(LeafReaderContext context) {
+      Weight weight = this;
+      return new ScorerSupplier() {
+        private long cost = 0;
+
         @Override
-        public DocIdSetIterator iterator() {
-          return null;
+        public Scorer get(long leadCost) {
+          return new Scorer(weight) {
+            @Override
+            public DocIdSetIterator iterator() {
+              return null;
+            }
+
+            @Override
+            public float getMaxScore(int upTo) {
+              return 42f;
+            }
+
+            @Override
+            public float score() {
+              return 0;
+            }
+
+            @Override
+            public int docID() {
+              return 0;
+            }
+          };
         }
 
         @Override
-        public float getMaxScore(int upTo) {
-          return 42f;
+        public long cost() {
+          return cost;
         }
 
         @Override
-        public float score() {
-          return 0;
-        }
-
-        @Override
-        public int docID() {
-          return 0;
+        public void setTopLevelScoringClause() {
+          cost = 42;
         }
       };
     }
@@ -153,4 +177,14 @@ public class TestQueryProfilerWeight extends LuceneTestCase {
     QueryProfilerWeight profileWeight = new QueryProfilerWeight(fakeWeight, profile);
     assertEquals(42f, profileWeight.scorer(null).getMaxScore(DocIdSetIterator.NO_MORE_DOCS), 0f);
   }
+
+  public void testPropagateTopLevelScoringClause() throws IOException {
+    Query query = new MatchAllDocsQuery();
+    Weight fakeWeight = new FakeWeight(query);
+    QueryProfilerBreakdown profile = new QueryProfilerBreakdown();
+    QueryProfilerWeight profileWeight = new QueryProfilerWeight(fakeWeight, profile);
+    ScorerSupplier scorerSupplier = profileWeight.scorerSupplier(null);
+    scorerSupplier.setTopLevelScoringClause();
+    assertEquals(42, scorerSupplier.cost());
+  }
 }