LUCENE-9965: Add tooling to introspect query execution time (#144)

This change adds new IndexSearcher and Collector implementations to profile
search execution and break down the timings. The breakdown includes the total
time spent in each of the following categories along with the number of times
visited: create weight, build scorer, next doc, advance, score, match.

Co-authored-by: Julie Tibshirani <[email protected]>

Author: jdconrad

lucene/CHANGES.txt

@@ -354,6 +354,33 @@ Other
 
 * LUCENE-9976: Fix WANDScorer assertion error. (Zach Chen, Adrien Grand, Dawid Weiss)
 
+======================= Lucene 8.10.0 =======================
+
+API Changes
+---------------------
+(No changes)
+
+New Features
+---------------------
+(No changes)
+
+Improvements
+---------------------
+* LUCENE-9965: Added QueryProfilerIndexSearcher and ProfilerCollector to support debugging
+  query execution strategy and timing. (Jack Conradson, Julie Tibshirani)
+
+Optimizations
+---------------------
+(No changes)
+
+Bug Fixes
+---------------------
+(No changes)
+
+Other
+---------------------
+(No changes)
+
 ======================= Lucene 8.9.0 =======================
 
 API Changes

lucene/sandbox/src/java/org/apache/lucene/sandbox/search/ProfilerCollector.java

@@ -0,0 +1,103 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.lucene.sandbox.search;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.search.Collector;
+import org.apache.lucene.search.LeafCollector;
+import org.apache.lucene.search.ScoreMode;
+
+/**
+ * This class wraps a Collector and times the execution of: - setScorer() - collect() -
+ * doSetNextReader() - needsScores()
+ *
+ * <p>QueryProfiler facilitates the linking of the Collector graph
+ */
+public class ProfilerCollector implements Collector {
+
+  /** A more friendly representation of the Collector's class name */
+  private final String collectorName;
+
+  /** A "hint" to help provide some context about this Collector */
+  private final String reason;
+
+  /** The wrapped collector */
+  private final ProfilerCollectorWrapper collector;
+
+  /** A list of "embedded" children collectors */
+  private final List<ProfilerCollector> children;
+
+  public ProfilerCollector(Collector collector, String reason, List<ProfilerCollector> children) {
+    this.collector = new ProfilerCollectorWrapper(collector);
+    this.reason = reason;
+    this.collectorName = deriveCollectorName(collector);
+    this.children = children;
+  }
+
+  /** @return the profiled time for this collector (inclusive of children) */
+  public long getTime() {
+    return collector.getTime();
+  }
+
+  /** @return a human readable "hint" about what this collector was used for */
+  public String getReason() {
+    return this.reason;
+  }
+
+  /** @return the lucene class name of the collector */
+  public String getName() {
+    return this.collectorName;
+  }
+
+  /**
+   * Creates a human-friendly representation of the Collector name.
+   *
+   * @param c The Collector to derive a name from
+   * @return A (hopefully) prettier name
+   */
+  private String deriveCollectorName(Collector c) {
+    return c.getClass().getSimpleName();
+  }
+
+  @Override
+  public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
+    return collector.getLeafCollector(context);
+  }
+
+  @Override
+  public ScoreMode scoreMode() {
+    return collector.scoreMode();
+  }
+
+  public ProfilerCollectorResult getProfileResult() {
+    return ProfilerCollector.doGetCollectorTree(this);
+  }
+
+  private static ProfilerCollectorResult doGetCollectorTree(ProfilerCollector collector) {
+    List<ProfilerCollectorResult> childResults = new ArrayList<>(collector.children.size());
+    for (ProfilerCollector child : collector.children) {
+      ProfilerCollectorResult result = doGetCollectorTree(child);
+      childResults.add(result);
+    }
+    return new ProfilerCollectorResult(
+        collector.getName(), collector.getReason(), collector.getTime(), childResults);
+  }
+}

lucene/sandbox/src/java/org/apache/lucene/sandbox/search/ProfilerCollectorResult.java

@@ -0,0 +1,67 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.lucene.sandbox.search;
+
+import java.util.List;
+
+/**
+ * Public class for profiled timings of the Collectors used in the search. Children
+ * CollectorResult's may be embedded inside of a parent CollectorResult
+ */
+public class ProfilerCollectorResult {
+
+  /** A more friendly representation of the Collector's class name */
+  private final String collectorName;
+
+  /** A "hint" to help provide some context about this Collector */
+  private final String reason;
+
+  /** The total elapsed time for this Collector */
+  private final long time;
+
+  /** A list of children collectors "embedded" inside this collector */
+  private final List<ProfilerCollectorResult> children;
+
+  public ProfilerCollectorResult(
+      String collectorName, String reason, Long time, List<ProfilerCollectorResult> children) {
+    this.collectorName = collectorName;
+    this.reason = reason;
+    this.time = time;
+    this.children = children;
+  }
+
+  /** @return the profiled time for this collector (inclusive of children) */
+  public long getTime() {
+    return this.time;
+  }
+
+  /** @return a human readable "hint" about what this collector was used for */
+  public String getReason() {
+    return this.reason;
+  }
+
+  /** @return the lucene class name of the collector */
+  public String getName() {
+    return this.collectorName;
+  }
+
+  /** @return a list of children collectors */
+  public List<ProfilerCollectorResult> getProfiledChildren() {
+    return children;
+  }
+}

lucene/sandbox/src/java/org/apache/lucene/sandbox/search/ProfilerCollectorWrapper.java

@@ -0,0 +1,86 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.lucene.sandbox.search;
+
+import java.io.IOException;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.search.Collector;
+import org.apache.lucene.search.FilterCollector;
+import org.apache.lucene.search.FilterLeafCollector;
+import org.apache.lucene.search.LeafCollector;
+import org.apache.lucene.search.Scorable;
+import org.apache.lucene.search.ScoreMode;
+
+/** A collector that profiles how much time is spent calling it. */
+class ProfilerCollectorWrapper extends FilterCollector {
+
+  private long time;
+
+  /** Sole constructor. */
+  ProfilerCollectorWrapper(Collector in) {
+    super(in);
+  }
+
+  @Override
+  public ScoreMode scoreMode() {
+    final long start = System.nanoTime();
+    try {
+      return super.scoreMode();
+    } finally {
+      time += Math.max(1, System.nanoTime() - start);
+    }
+  }
+
+  @Override
+  public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
+    final long start = System.nanoTime();
+    final LeafCollector inLeafCollector;
+    try {
+      inLeafCollector = super.getLeafCollector(context);
+    } finally {
+      time += Math.max(1, System.nanoTime() - start);
+    }
+    return new FilterLeafCollector(inLeafCollector) {
+
+      @Override
+      public void collect(int doc) throws IOException {
+        final long start = System.nanoTime();
+        try {
+          super.collect(doc);
+        } finally {
+          time += Math.max(1, System.nanoTime() - start);
+        }
+      }
+
+      @Override
+      public void setScorer(Scorable scorer) throws IOException {
+        final long start = System.nanoTime();
+        try {
+          super.setScorer(scorer);
+        } finally {
+          time += Math.max(1, System.nanoTime() - start);
+        }
+      }
+    };
+  }
+
+  /** Return the total time spent on this collector. */
+  public long getTime() {
+    return time;
+  }
+}

lucene/sandbox/src/java/org/apache/lucene/sandbox/search/QueryProfilerBreakdown.java

@@ -0,0 +1,62 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.lucene.sandbox.search;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * A record of timings for the various operations that may happen during query execution. A node's
+ * time may be composed of several internal attributes (rewriting, weighting, scoring, etc).
+ */
+class QueryProfilerBreakdown {
+
+  /** The accumulated timings for this query node */
+  private final QueryProfilerTimer[] timers;
+
+  /** Sole constructor. */
+  public QueryProfilerBreakdown() {
+    timers = new QueryProfilerTimer[QueryProfilerTimingType.values().length];
+    for (int i = 0; i < timers.length; ++i) {
+      timers[i] = new QueryProfilerTimer();
+    }
+  }
+
+  public QueryProfilerTimer getTimer(QueryProfilerTimingType type) {
+    return timers[type.ordinal()];
+  }
+
+  /** Build a timing count breakdown. */
+  public final Map<String, Long> toBreakdownMap() {
+    Map<String, Long> map = new HashMap<>(timers.length * 2);
+    for (QueryProfilerTimingType type : QueryProfilerTimingType.values()) {
+      map.put(type.toString(), timers[type.ordinal()].getApproximateTiming());
+      map.put(type.toString() + "_count", timers[type.ordinal()].getCount());
+    }
+    return Collections.unmodifiableMap(map);
+  }
+
+  public final long toTotalTime() {
+    long total = 0;
+    for (QueryProfilerTimer timer : timers) {
+      total += timer.getApproximateTiming();
+    }
+    return total;
+  }
+}