2B POC code

.gitignore

@@ -58,3 +58,6 @@ kubernetes/helm/**/Chart.lock
 #Develocity
 .mvn/.gradle-enterprise/
 .mvn/.develocity/
+
+pinot-integration-tests/src/test/resources/udf-test-results/*.md
+!/pinot-integration-tests/src/test/resources/udf-test-results/README.md

pinot-broker/src/main/java/org/apache/pinot/broker/api/resources/PinotBrokerDebug.java

@@ -55,8 +55,8 @@
 import org.apache.pinot.core.auth.ManualAuthorization;
 import org.apache.pinot.core.auth.TargetType;
 import org.apache.pinot.core.routing.RoutingTable;
-import org.apache.pinot.core.routing.ServerRouteInfo;
-import org.apache.pinot.core.routing.TimeBoundaryInfo;
+import org.apache.pinot.core.routing.SegmentsToQuery;
+import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo;
 import org.apache.pinot.core.transport.ServerInstance;
 import org.apache.pinot.core.transport.server.routing.stats.ServerRoutingStatsEntry;
 import org.apache.pinot.core.transport.server.routing.stats.ServerRoutingStatsManager;
@@ -158,11 +158,11 @@ public Map<String, Map<ServerInstance, List<String>>> getRoutingTable(
       @ApiResponse(code = 404, message = "Routing not found"),
       @ApiResponse(code = 500, message = "Internal server error")
   })
-  public Map<String, Map<ServerInstance, ServerRouteInfo>> getRoutingTableWithOptionalSegments(
+  public Map<String, Map<ServerInstance, SegmentsToQuery>> getRoutingTableWithOptionalSegments(
       @ApiParam(value = "Name of the table") @PathParam("tableName") String tableName,
       @Context HttpHeaders headers) {
     tableName = DatabaseUtils.translateTableName(tableName, headers);
-    Map<String, Map<ServerInstance, ServerRouteInfo>> result = new TreeMap<>();
+    Map<String, Map<ServerInstance, SegmentsToQuery>> result = new TreeMap<>();
     getRoutingTable(tableName, (tableNameWithType, routingTable) -> result.put(tableNameWithType,
         routingTable.getServerInstanceToSegmentsMap()));
     if (!result.isEmpty()) {
@@ -193,7 +193,7 @@ private void getRoutingTable(String tableName, BiConsumer<String, RoutingTable>
   }
 
   private static Map<ServerInstance, List<String>> removeOptionalSegments(
-      Map<ServerInstance, ServerRouteInfo> serverInstanceToSegmentsMap) {
+      Map<ServerInstance, SegmentsToQuery> serverInstanceToSegmentsMap) {
     Map<ServerInstance, List<String>> ret = new HashMap<>();
     serverInstanceToSegmentsMap.forEach((k, v) -> ret.put(k, v.getSegments()));
     return ret;
@@ -232,7 +232,7 @@ public Map<ServerInstance, List<String>> getRoutingTableForQuery(
       @ApiResponse(code = 404, message = "Routing not found"),
       @ApiResponse(code = 500, message = "Internal server error")
   })
-  public Map<ServerInstance, ServerRouteInfo> getRoutingTableForQueryWithOptionalSegments(
+  public Map<ServerInstance, SegmentsToQuery> getRoutingTableForQueryWithOptionalSegments(
       @ApiParam(value = "SQL query (table name should have type suffix)") @QueryParam("query") String query,
       @Context HttpHeaders httpHeaders) {
     BrokerRequest brokerRequest = CalciteSqlCompiler.compileToBrokerRequest(query);

pinot-broker/src/main/java/org/apache/pinot/broker/api/resources/PinotClientRequest.java

@@ -280,6 +280,33 @@ public void processSqlWithMultiStageQueryEnginePost(String query, @Suspended Asy
     }
   }
 
+  @POST
+  @ManagedAsync
+  @Produces(MediaType.APPLICATION_JSON)
+  @Path("query/timeseries")
+  @ApiOperation(value = "Query Pinot using the Time Series Engine")
+  @ManualAuthorization
+  public void processTimeSeriesQueryEngine(Map<String, String> queryParams, @Suspended AsyncResponse asyncResponse,
+      @Context org.glassfish.grizzly.http.server.Request requestCtx, @Context HttpHeaders httpHeaders) {
+    try {
+      if (!queryParams.containsKey(Request.QUERY)) {
+        throw new IllegalStateException("Payload is missing the query string field 'query'");
+      }
+      String language = queryParams.get(Request.LANGUAGE);
+      String queryString = queryParams.get(Request.QUERY);
+      try (RequestScope requestContext = Tracing.getTracer().createRequestScope()) {
+        PinotBrokerTimeSeriesResponse response = executeTimeSeriesQuery(language, queryString, queryParams,
+            requestContext, makeHttpIdentity(requestCtx), httpHeaders);
+        asyncResponse.resume(response.toBrokerResponse());
+      }
+    } catch (Exception e) {
+      LOGGER.error("Caught exception while processing POST timeseries request", e);
+      _brokerMetrics.addMeteredGlobalValue(BrokerMeter.UNCAUGHT_POST_EXCEPTIONS, 1L);
+      asyncResponse.resume(new WebApplicationException(e,
+          Response.status(Response.Status.INTERNAL_SERVER_ERROR).entity(e.getMessage()).build()));
+    }
+  }
+
   @GET
   @ManagedAsync
   @Produces(MediaType.APPLICATION_JSON)
@@ -293,7 +320,7 @@ public void processTimeSeriesQueryEngine(@Suspended AsyncResponse asyncResponse,
     try {
       try (RequestScope requestContext = Tracing.getTracer().createRequestScope()) {
         String queryString = requestCtx.getQueryString();
-        PinotBrokerTimeSeriesResponse response = executeTimeSeriesQuery(language, queryString, requestContext,
+        PinotBrokerTimeSeriesResponse response = executeTimeSeriesQuery(language, queryString, Map.of(), requestContext,
           makeHttpIdentity(requestCtx), httpHeaders);
         if (response.getErrorType() != null && !response.getErrorType().isEmpty()) {
           asyncResponse.resume(Response.serverError().entity(response).build());
@@ -538,9 +565,10 @@ private BrokerResponse executeSqlQuery(ObjectNode sqlRequestJson, HttpRequesterI
   }
 
   private PinotBrokerTimeSeriesResponse executeTimeSeriesQuery(String language, String queryString,
-    RequestContext requestContext, RequesterIdentity requesterIdentity, HttpHeaders httpHeaders) {
-    return _requestHandler.handleTimeSeriesRequest(language, queryString, requestContext, requesterIdentity,
-      httpHeaders);
+      Map<String, String> queryParams, RequestContext requestContext, RequesterIdentity requesterIdentity,
+      HttpHeaders httpHeaders) {
+    return _requestHandler.handleTimeSeriesRequest(language, queryString, queryParams, requestContext,
+        requesterIdentity, httpHeaders);
   }
 
   public static HttpRequesterIdentity makeHttpIdentity(org.glassfish.grizzly.http.server.Request context) {

pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java

@@ -19,7 +19,6 @@
 package org.apache.pinot.broker.broker.helix;
 
 import com.google.common.base.Preconditions;
-import com.google.common.collect.ImmutableList;
 import java.io.IOException;
 import java.net.InetAddress;
 import java.util.ArrayList;
@@ -85,6 +84,7 @@
 import org.apache.pinot.query.mailbox.MailboxService;
 import org.apache.pinot.query.service.dispatch.QueryDispatcher;
 import org.apache.pinot.segment.local.function.GroovyFunctionEvaluator;
+import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant;
 import org.apache.pinot.spi.accounting.ThreadResourceUsageProvider;
 import org.apache.pinot.spi.cursors.ResponseStoreService;
 import org.apache.pinot.spi.env.PinotConfiguration;
@@ -153,13 +153,14 @@ public abstract class BaseBrokerStarter implements ServiceStartable {
   protected AbstractResponseStore _responseStore;
   protected BrokerGrpcServer _brokerGrpcServer;
   protected FailureDetector _failureDetector;
+  protected ThreadResourceUsageAccountant _resourceUsageAccountant;
 
   @Override
   public void init(PinotConfiguration brokerConf)
       throws Exception {
     _brokerConf = brokerConf;
     // Remove all white-spaces from the list of zkServers (if any).
-    _zkServers = brokerConf.getProperty(Helix.CONFIG_OF_ZOOKEEPR_SERVER).replaceAll("\\s+", "");
+    _zkServers = brokerConf.getProperty(Helix.CONFIG_OF_ZOOKEEPER_SERVER).replaceAll("\\s+", "");
     _clusterName = brokerConf.getProperty(Helix.CONFIG_OF_CLUSTER_NAME);
     ServiceStartableUtils.applyClusterConfig(_brokerConf, _zkServers, _clusterName, ServiceRole.BROKER);
     applyCustomConfigs(brokerConf);
@@ -345,14 +346,27 @@ public void start()
         instanceId -> _routingManager.excludeServerFromRouting(instanceId));
     _failureDetector.start();
 
+    // Enable/disable thread CPU time measurement through instance config.
+    ThreadResourceUsageProvider.setThreadCpuTimeMeasurementEnabled(
+        _brokerConf.getProperty(CommonConstants.Broker.CONFIG_OF_ENABLE_THREAD_CPU_TIME_MEASUREMENT,
+            CommonConstants.Broker.DEFAULT_ENABLE_THREAD_CPU_TIME_MEASUREMENT));
+    // Enable/disable thread memory allocation tracking through instance config
+    ThreadResourceUsageProvider.setThreadMemoryMeasurementEnabled(
+        _brokerConf.getProperty(CommonConstants.Broker.CONFIG_OF_ENABLE_THREAD_ALLOCATED_BYTES_MEASUREMENT,
+            CommonConstants.Broker.DEFAULT_THREAD_ALLOCATED_BYTES_MEASUREMENT));
+    _resourceUsageAccountant = Tracing.ThreadAccountantOps.createThreadAccountant(
+        _brokerConf.subset(CommonConstants.PINOT_QUERY_SCHEDULER_PREFIX), _instanceId,
+        org.apache.pinot.spi.config.instance.InstanceType.BROKER);
+    Preconditions.checkNotNull(_resourceUsageAccountant);
+
     // Create Broker request handler.
     String brokerId = _brokerConf.getProperty(Broker.CONFIG_OF_BROKER_ID, getDefaultBrokerId());
     String brokerRequestHandlerType =
         _brokerConf.getProperty(Broker.BROKER_REQUEST_HANDLER_TYPE, Broker.DEFAULT_BROKER_REQUEST_HANDLER_TYPE);
     BaseSingleStageBrokerRequestHandler singleStageBrokerRequestHandler;
     if (brokerRequestHandlerType.equalsIgnoreCase(Broker.GRPC_BROKER_REQUEST_HANDLER_TYPE)) {
       singleStageBrokerRequestHandler = new GrpcBrokerRequestHandler(_brokerConf, brokerId, _routingManager,
-          _accessControlFactory, _queryQuotaManager, tableCache, _failureDetector);
+          _accessControlFactory, _queryQuotaManager, tableCache, _failureDetector, _resourceUsageAccountant);
     } else {
       // Default request handler type, i.e. netty
       NettyConfig nettyDefaults = NettyConfig.extractNettyConfig(_brokerConf, Broker.BROKER_NETTY_PREFIX);
@@ -364,7 +378,7 @@ public void start()
       singleStageBrokerRequestHandler =
           new SingleConnectionBrokerRequestHandler(_brokerConf, brokerId, _routingManager, _accessControlFactory,
               _queryQuotaManager, tableCache, nettyDefaults, tlsDefaults, _serverRoutingStatsManager,
-              _failureDetector);
+              _failureDetector, _resourceUsageAccountant);
     }
     MultiStageBrokerRequestHandler multiStageBrokerRequestHandler = null;
     QueryDispatcher queryDispatcher = null;
@@ -377,13 +391,13 @@ public void start()
       queryDispatcher = createQueryDispatcher(_brokerConf);
       multiStageBrokerRequestHandler =
           new MultiStageBrokerRequestHandler(_brokerConf, brokerId, _routingManager, _accessControlFactory,
-              _queryQuotaManager, tableCache, _multiStageQueryThrottler, _failureDetector);
+              _queryQuotaManager, tableCache, _multiStageQueryThrottler, _failureDetector, _resourceUsageAccountant);
     }
     TimeSeriesRequestHandler timeSeriesRequestHandler = null;
     if (StringUtils.isNotBlank(_brokerConf.getProperty(PinotTimeSeriesConfiguration.getEnabledLanguagesConfigKey()))) {
       Preconditions.checkNotNull(queryDispatcher, "Multistage Engine should be enabled to use time-series engine");
       timeSeriesRequestHandler = new TimeSeriesRequestHandler(_brokerConf, brokerId, _routingManager,
-          _accessControlFactory, _queryQuotaManager, tableCache, queryDispatcher);
+          _accessControlFactory, _queryQuotaManager, tableCache, queryDispatcher, _resourceUsageAccountant);
     }
 
     LOGGER.info("Initializing PinotFSFactory");
@@ -407,17 +421,6 @@ public void start()
             timeSeriesRequestHandler, _responseStore);
     _brokerRequestHandler.start();
 
-    // Enable/disable thread CPU time measurement through instance config.
-    ThreadResourceUsageProvider.setThreadCpuTimeMeasurementEnabled(
-        _brokerConf.getProperty(CommonConstants.Broker.CONFIG_OF_ENABLE_THREAD_CPU_TIME_MEASUREMENT,
-            CommonConstants.Broker.DEFAULT_ENABLE_THREAD_CPU_TIME_MEASUREMENT));
-    // Enable/disable thread memory allocation tracking through instance config
-    ThreadResourceUsageProvider.setThreadMemoryMeasurementEnabled(
-        _brokerConf.getProperty(CommonConstants.Broker.CONFIG_OF_ENABLE_THREAD_ALLOCATED_BYTES_MEASUREMENT,
-            CommonConstants.Broker.DEFAULT_THREAD_ALLOCATED_BYTES_MEASUREMENT));
-    Tracing.ThreadAccountantOps.initializeThreadAccountant(
-        _brokerConf.subset(CommonConstants.PINOT_QUERY_SCHEDULER_PREFIX), _instanceId,
-        org.apache.pinot.spi.config.instance.InstanceType.BROKER);
     Tracing.ThreadAccountantOps.startThreadAccountant();
 
     String controllerUrl = _brokerConf.getProperty(Broker.CONTROLLER_URL);
@@ -568,7 +571,7 @@ private void updateInstanceConfigAndBrokerResourceIfNeeded() {
       if (StringUtils.isNotEmpty(instanceTagsConfig)) {
         for (String instanceTag : StringUtils.split(instanceTagsConfig, ',')) {
           Preconditions.checkArgument(TagNameUtils.isBrokerTag(instanceTag), "Illegal broker instance tag: %s",
-                  instanceTag);
+              instanceTag);
           instanceConfig.addTag(instanceTag);
         }
         shouldUpdateBrokerResource = true;
@@ -618,11 +621,12 @@ private void registerServiceStatusHandler() {
 
     LOGGER.info("Registering service status handler");
     ServiceStatus.setServiceStatusCallback(_instanceId, new ServiceStatus.MultipleCallbackServiceStatusCallback(
-        ImmutableList.of(new ServiceStatus.IdealStateAndCurrentStateMatchServiceStatusCallback(_participantHelixManager,
+        List.of(
+            new ServiceStatus.LifecycleServiceStatusCallback(this::isStarting, this::isShuttingDown),
+            new ServiceStatus.IdealStateAndCurrentStateMatchServiceStatusCallback(_participantHelixManager,
                 _clusterName, _instanceId, resourcesToMonitor, minResourcePercentForStartup),
             new ServiceStatus.IdealStateAndExternalViewMatchServiceStatusCallback(_participantHelixManager,
-                _clusterName, _instanceId, resourcesToMonitor, minResourcePercentForStartup),
-            new ServiceStatus.LifecycleServiceStatusCallback(this::isStarting, this::isShuttingDown))));
+                _clusterName, _instanceId, resourcesToMonitor, minResourcePercentForStartup))));
   }
 
   private String getDefaultBrokerId() {

pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/HelixBrokerStarter.java

@@ -53,7 +53,7 @@ public HelixBrokerStarter(PinotConfiguration brokerConf, String clusterName, Str
   private static PinotConfiguration applyBrokerConfigs(PinotConfiguration brokerConf, String clusterName,
       String zkServers, @Nullable String brokerHost) {
     brokerConf.setProperty(Helix.CONFIG_OF_CLUSTER_NAME, clusterName);
-    brokerConf.setProperty(Helix.CONFIG_OF_ZOOKEEPR_SERVER, zkServers);
+    brokerConf.setProperty(Helix.CONFIG_OF_ZOOKEEPER_SERVER, zkServers);
     if (brokerHost == null) {
       brokerConf.clearProperty(Broker.CONFIG_OF_BROKER_HOSTNAME);
     } else {
@@ -75,7 +75,7 @@ public static HelixBrokerStarter getDefault()
     properties.put(Helix.KEY_OF_BROKER_QUERY_PORT, 5001);
     properties.put(Broker.CONFIG_OF_BROKER_TIMEOUT_MS, 60 * 1000L);
     properties.put(Helix.CONFIG_OF_CLUSTER_NAME, "quickstart");
-    properties.put(Helix.CONFIG_OF_ZOOKEEPR_SERVER, "localhost:2122");
+    properties.put(Helix.CONFIG_OF_ZOOKEEPER_SERVER, "localhost:2122");
     HelixBrokerStarter helixBrokerStarter = new HelixBrokerStarter();
     helixBrokerStarter.init(new PinotConfiguration(properties));
     return helixBrokerStarter;

pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java

@@ -42,7 +42,6 @@
 import org.apache.pinot.broker.broker.AccessControlFactory;
 import org.apache.pinot.broker.querylog.QueryLogger;
 import org.apache.pinot.broker.queryquota.QueryQuotaManager;
-import org.apache.pinot.broker.routing.BrokerRoutingManager;
 import org.apache.pinot.common.config.provider.TableCache;
 import org.apache.pinot.common.metrics.BrokerMeter;
 import org.apache.pinot.common.metrics.BrokerMetrics;
@@ -53,6 +52,8 @@
 import org.apache.pinot.common.utils.request.RequestUtils;
 import org.apache.pinot.core.auth.Actions;
 import org.apache.pinot.core.auth.TargetType;
+import org.apache.pinot.core.routing.RoutingManager;
+import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant;
 import org.apache.pinot.spi.auth.AuthorizationResult;
 import org.apache.pinot.spi.auth.TableAuthorizationResult;
 import org.apache.pinot.spi.auth.broker.RequesterIdentity;
@@ -76,7 +77,7 @@ public abstract class BaseBrokerRequestHandler implements BrokerRequestHandler {
   private static final Logger LOGGER = LoggerFactory.getLogger(BaseBrokerRequestHandler.class);
   protected final PinotConfiguration _config;
   protected final String _brokerId;
-  protected final BrokerRoutingManager _routingManager;
+  protected final RoutingManager _routingManager;
   protected final AccessControlFactory _accessControlFactory;
   protected final QueryQuotaManager _queryQuotaManager;
   protected final TableCache _tableCache;
@@ -89,6 +90,8 @@ public abstract class BaseBrokerRequestHandler implements BrokerRequestHandler {
   protected final QueryLogger _queryLogger;
   @Nullable
   protected final String _enableNullHandling;
+  protected final ThreadResourceUsageAccountant _resourceUsageAccountant;
+
   /**
    * Maps broker-generated query id to the query string.
    */
@@ -98,8 +101,9 @@ public abstract class BaseBrokerRequestHandler implements BrokerRequestHandler {
    */
   protected final Map<Long, String> _clientQueryIds;
 
-  public BaseBrokerRequestHandler(PinotConfiguration config, String brokerId, BrokerRoutingManager routingManager,
-      AccessControlFactory accessControlFactory, QueryQuotaManager queryQuotaManager, TableCache tableCache) {
+  public BaseBrokerRequestHandler(PinotConfiguration config, String brokerId, RoutingManager routingManager,
+      AccessControlFactory accessControlFactory, QueryQuotaManager queryQuotaManager, TableCache tableCache,
+      ThreadResourceUsageAccountant resourceUsageAccountant) {
     _config = config;
     _brokerId = brokerId;
     _routingManager = routingManager;
@@ -125,6 +129,7 @@ public BaseBrokerRequestHandler(PinotConfiguration config, String brokerId, Brok
       _queriesById = null;
       _clientQueryIds = null;
     }
+    _resourceUsageAccountant = resourceUsageAccountant;
   }
 
   @Override
@@ -389,6 +394,12 @@ protected String extractClientRequestId(SqlNodeAndOptions sqlNodeAndOptions) {
         ? sqlNodeAndOptions.getOptions().get(Broker.Request.QueryOptionKey.CLIENT_QUERY_ID) : null;
   }
 
+  /**
+   * Called when a query starts
+   * TODO: This method was created to keep track of running queries for cancellation, but it is useful for other uses.
+   *   But right now the semantics are not clear. For example, while MSE calls this method once, SSE calls it once per
+   *   query AND subquery, which means this method is called multiple times for the same query.
+   */
   protected void onQueryStart(long requestId, String clientRequestId, String query, Object... extras) {
     if (isQueryCancellationEnabled()) {
       _queriesById.put(requestId, query);