log segment name at best effort upon query exception to help locate t…

…he error segment (#14440)

* wrap segment name on exception, rather than logging out separately

pinot-core/src/main/java/org/apache/pinot/core/operator/combine/BaseCombineOperator.java

@@ -35,6 +35,7 @@
 import org.apache.pinot.core.query.request.context.QueryContext;
 import org.apache.pinot.core.util.QueryMultiThreadingUtils;
 import org.apache.pinot.core.util.trace.TraceRunnable;
+import org.apache.pinot.segment.spi.IndexSegment;
 import org.apache.pinot.spi.accounting.ThreadExecutionContext;
 import org.apache.pinot.spi.accounting.ThreadResourceUsageProvider;
 import org.apache.pinot.spi.exception.EarlyTerminationException;
@@ -180,4 +181,20 @@ public List<Operator> getChildOperators() {
    * Invoked when {@link #processSegments()} is finished (called in the finally block).
    */
   protected abstract void onProcessSegmentsFinish();
+
+  protected static RuntimeException wrapOperatorException(Operator operator, RuntimeException e) {
+    // Skip early termination as that's not quite related to the segments.
+    if (e instanceof EarlyTerminationException) {
+      return e;
+    }
+    // Otherwise, try to get the segment name to help locate the segment when debugging query errors.
+    // Not all operators have associated segment, so do this at best effort.
+    IndexSegment segment = operator.getIndexSegment();
+    if (segment == null) {
+      return e;
+    }
+    throw new RuntimeException(
+        "Caught exception while doing operator: " + operator.getClass() + " on segment: " + segment.getSegmentName(),
+        e);
+  }
 }

pinot-core/src/main/java/org/apache/pinot/core/operator/combine/BaseSingleBlockCombineOperator.java

@@ -92,6 +92,8 @@ protected void processSegments() {
           ((AcquireReleaseColumnsSegmentOperator) operator).acquire();
         }
         resultsBlock = (T) operator.nextBlock();
+      } catch (RuntimeException e) {
+        throw wrapOperatorException(operator, e);
       } finally {
         if (operator instanceof AcquireReleaseColumnsSegmentOperator) {
           ((AcquireReleaseColumnsSegmentOperator) operator).release();

pinot-core/src/main/java/org/apache/pinot/core/operator/combine/GroupByCombineOperator.java

@@ -188,6 +188,8 @@ protected void processSegments() {
             mergedKeys++;
           }
         }
+      } catch (RuntimeException e) {
+        throw wrapOperatorException(operator, e);
       } finally {
         if (operator instanceof AcquireReleaseColumnsSegmentOperator) {
           ((AcquireReleaseColumnsSegmentOperator) operator).release();

pinot-core/src/main/java/org/apache/pinot/core/operator/combine/MinMaxValueBasedSelectionOrderByCombineOperator.java

@@ -205,6 +205,8 @@ protected void processSegments() {
           ((AcquireReleaseColumnsSegmentOperator) operator).acquire();
         }
         resultsBlock = (SelectionResultsBlock) operator.nextBlock();
+      } catch (RuntimeException e) {
+        throw wrapOperatorException(operator, e);
       } finally {
         if (operator instanceof AcquireReleaseColumnsSegmentOperator) {
           ((AcquireReleaseColumnsSegmentOperator) operator).release();

pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/BaseStreamingCombineOperator.java

@@ -142,6 +142,8 @@ protected void processSegments() {
             }
           }
         }
+      } catch (RuntimeException e) {
+        throw wrapOperatorException(operator, e);
       } finally {
         if (operator instanceof AcquireReleaseColumnsSegmentOperator) {
           ((AcquireReleaseColumnsSegmentOperator) operator).release();

pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/StreamingGroupByCombineOperator.java

@@ -216,6 +216,8 @@ public void processSegments() {
             mergedKeys++;
           }
         }
+      } catch (RuntimeException e) {
+        throw wrapOperatorException(operator, e);
       } finally {
         if (operator instanceof AcquireReleaseColumnsSegmentOperator) {
           ((AcquireReleaseColumnsSegmentOperator) operator).release();