Include information on whether a query is using SSQE or MSQE in Query…

…Logger (#14938)
apache · Feb 6, 2025 · db3f606 · db3f606
1 parent e9ee474
commit db3f606
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 6 deletions.
diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/querylog/QueryLogger.java b/pinot-broker/src/main/java/org/apache/pinot/broker/querylog/QueryLogger.java
@@ -120,20 +120,37 @@ public static class QueryLogParams {
     private final RequestContext _requestContext;
     private final String _table;
     private final BrokerResponse _response;
+    private final QueryEngine _queryEngine;
     @Nullable
     private final RequesterIdentity _identity;
     @Nullable
     private final ServerStats _serverStats;
 
     public QueryLogParams(RequestContext requestContext, String table, BrokerResponse response,
-        @Nullable RequesterIdentity identity, @Nullable ServerStats serverStats) {
+        QueryEngine queryEngine, @Nullable RequesterIdentity identity, @Nullable ServerStats serverStats) {
       _requestContext = requestContext;
       // NOTE: Passing table name separately because table name within request context is always raw table name.
       _table = table;
       _response = response;
+      _queryEngine = queryEngine;
       _identity = identity;
       _serverStats = serverStats;
     }
+
+    public enum QueryEngine {
+      SINGLE_STAGE("singleStage"),
+      MULTI_STAGE("multiStage");
+
+      private final String _name;
+
+      QueryEngine(String name) {
+        _name = name;
+      }
+
+      private String getName() {
+        return _name;
+      }
+    }
   }
 
   /**
@@ -255,6 +272,12 @@ void doFormat(StringBuilder builder, QueryLogger logger, QueryLogParams params)
           builder.append(CommonConstants.UNKNOWN);
         }
       }
+    },
+    QUERY_ENGINE("queryEngine") {
+      @Override
+      void doFormat(StringBuilder builder, QueryLogger logger, QueryLogParams params) {
+        builder.append(params._queryEngine.getName());
+      }
     };
 
     public final String _entryName;

diff --git a/...main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java b/...main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java
@@ -882,7 +882,8 @@ protected BrokerResponse handleRequest(long requestId, String query, SqlNodeAndO
 
       // Log query and stats
       _queryLogger.log(
-          new QueryLogger.QueryLogParams(requestContext, tableName, brokerResponse, requesterIdentity, serverStats));
+          new QueryLogger.QueryLogParams(requestContext, tableName, brokerResponse,
+              QueryLogger.QueryLogParams.QueryEngine.SINGLE_STAGE, requesterIdentity, serverStats));
 
       return brokerResponse;
     } finally {
@@ -931,7 +932,8 @@ private BrokerResponseNative getEmptyBrokerOnlyResponse(PinotQuery pinotQuery, R
     ParserUtils.fillEmptyResponseSchema(brokerResponse, _tableCache, schema, database, query);
     brokerResponse.setTimeUsedMs(System.currentTimeMillis() - requestContext.getRequestArrivalTimeMillis());
     _queryLogger.log(
-        new QueryLogger.QueryLogParams(requestContext, tableName, brokerResponse, requesterIdentity, null));
+        new QueryLogger.QueryLogParams(requestContext, tableName, brokerResponse,
+            QueryLogger.QueryLogParams.QueryEngine.SINGLE_STAGE, requesterIdentity, null));
     return brokerResponse;
   }
 

diff --git a/.../src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java b/.../src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java
@@ -328,8 +328,8 @@ protected BrokerResponse handleRequest(long requestId, String query, SqlNodeAndO
 
       // Log query and stats
       _queryLogger.log(
-          new QueryLogger.QueryLogParams(requestContext, tableNames.toString(), brokerResponse, requesterIdentity,
-              null));
+          new QueryLogger.QueryLogParams(requestContext, tableNames.toString(), brokerResponse,
+              QueryLogger.QueryLogParams.QueryEngine.MULTI_STAGE, requesterIdentity, null));
 
       return brokerResponse;
     } finally {

diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/querylog/QueryLoggerTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/querylog/QueryLoggerTest.java
@@ -110,6 +110,7 @@ public void shouldFormatLogLineProperly() {
         + "offlineThreadCpuTimeNs(total/thread/sysActivity/resSer):45/14/15/16,"
         + "realtimeThreadCpuTimeNs(total/thread/sysActivity/resSer):54/17/18/19,"
         + "clientIp=ip,"
+        + "queryEngine=singleStage,"
         + "query=SELECT * FROM foo");
     //@formatter:on
   }
@@ -281,6 +282,7 @@ public String getClientIp() {
     ServerStats serverStats = new ServerStats();
     serverStats.setServerStats("serverStats");
 
-    return new QueryLogger.QueryLogParams(requestContext, "table", response, identity, serverStats);
+    return new QueryLogger.QueryLogParams(requestContext, "table", response,
+        QueryLogger.QueryLogParams.QueryEngine.SINGLE_STAGE, identity, serverStats);
   }
 }