protect usage MSQE compiler for empty schema polyfill with config param disabled by default (apache#15078)

Author: albertobastos

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

@@ -146,6 +146,7 @@ public abstract class BaseSingleStageBrokerRequestHandler extends BaseBrokerRequ
   // if >= 0, then overrides default limit of 10, otherwise setting is ignored
   protected final int _defaultQueryLimit;
   protected final boolean _enableMultistageMigrationMetric;
+  protected final boolean _useMSEToFillEmptyResponseSchema;
   protected ExecutorService _multistageCompileExecutor;
   protected BlockingQueue<Pair<String, String>> _multistageCompileQueryQueue;
 
@@ -179,6 +180,9 @@ public BaseSingleStageBrokerRequestHandler(PinotConfiguration config, String bro
       _multistageCompileQueryQueue = new LinkedBlockingQueue<>(1000);
     }
 
+    _useMSEToFillEmptyResponseSchema = _config.getProperty(Broker.USE_MSE_TO_FILL_EMPTY_RESPONSE_SCHEMA,
+        Broker.DEFAULT_USE_MSE_TO_FILL_EMPTY_RESPONSE_SCHEMA);
+
     LOGGER.info("Initialized {} with broker id: {}, timeout: {}ms, query response limit: {}, "
             + "default query limit {}, query log max length: {}, query log max rate: {}, query cancellation "
             + "enabled: {}", getClass().getSimpleName(), _brokerId, _brokerTimeoutMs, _queryResponseLimit,
@@ -869,7 +873,9 @@ protected BrokerResponse doHandleRequest(long requestId, String query, SqlNodeAn
     // server returns STRING as default dataType for all columns in (some) scenarios where no rows are returned
     // this is an attempt to return more faithful information based on other sources
     if (brokerResponse.getNumRowsResultSet() == 0) {
-      ParserUtils.fillEmptyResponseSchema(brokerResponse, _tableCache, schema, database, query);
+      boolean useMSE = QueryOptionsUtils.isUseMSEToFillEmptySchema(
+          pinotQuery.getQueryOptions(), _useMSEToFillEmptyResponseSchema);
+      ParserUtils.fillEmptyResponseSchema(useMSE, brokerResponse, _tableCache, schema, database, query);
     }
 
     // Set total query processing time
@@ -964,7 +970,9 @@ private BrokerResponseNative getEmptyBrokerOnlyResponse(PinotQuery pinotQuery, R
 
     // Send empty response since we don't need to evaluate either offline or realtime request.
     BrokerResponseNative brokerResponse = BrokerResponseNative.empty();
-    ParserUtils.fillEmptyResponseSchema(brokerResponse, _tableCache, schema, database, query);
+    boolean useMSE = QueryOptionsUtils.isUseMSEToFillEmptySchema(
+        pinotQuery.getQueryOptions(), _useMSEToFillEmptyResponseSchema);
+    ParserUtils.fillEmptyResponseSchema(useMSE, brokerResponse, _tableCache, schema, database, query);
     brokerResponse.setTimeUsedMs(System.currentTimeMillis() - requestContext.getRequestArrivalTimeMillis());
     _queryLogger.log(
         new QueryLogger.QueryLogParams(requestContext, tableName, brokerResponse,

pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java

@@ -362,6 +362,11 @@ public static boolean isSecondaryWorkload(Map<String, String> queryOptions) {
     return Boolean.parseBoolean(queryOptions.get(QueryOptionKey.IS_SECONDARY_WORKLOAD));
   }
 
+  public static Boolean isUseMSEToFillEmptySchema(Map<String, String> queryOptions, boolean defaultValue) {
+    String useMSEToFillEmptySchema = queryOptions.get(QueryOptionKey.USE_MSE_TO_FILL_EMPTY_RESPONSE_SCHEMA);
+    return useMSEToFillEmptySchema != null ? Boolean.parseBoolean(useMSEToFillEmptySchema) : defaultValue;
+  }
+
   @Nullable
   private static Integer uncheckedParseInt(String optionName, @Nullable String optionValue) {
     if (optionValue == null) {

pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/BasePauselessRealtimeIngestionTest.java

@@ -85,6 +85,12 @@ protected void overrideControllerConf(Map<String, Object> properties) {
         500);
   }
 
+  @Override
+  protected void overrideBrokerConf(PinotConfiguration brokerConf) {
+    super.overrideBrokerConf(brokerConf);
+    brokerConf.setProperty(CommonConstants.Broker.USE_MSE_TO_FILL_EMPTY_RESPONSE_SCHEMA, true);
+  }
+
   @Override
   protected void overrideServerConf(PinotConfiguration serverConf) {
     try {

pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/EmptyResponseIntegrationTest.java

@@ -36,6 +36,7 @@
 import org.apache.pinot.spi.config.table.TableConfig;
 import org.apache.pinot.spi.config.table.TableType;
 import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.env.PinotConfiguration;
 import org.apache.pinot.spi.utils.CommonConstants;
 import org.apache.pinot.spi.utils.InstanceTypeUtils;
 import org.apache.pinot.spi.utils.builder.TableNameBuilder;
@@ -82,6 +83,12 @@ protected List<FieldConfig> getFieldConfigs() {
             CompressionCodec.MV_ENTRY_DICT, null));
   }
 
+  @Override
+  protected void overrideBrokerConf(PinotConfiguration brokerConf) {
+    super.overrideBrokerConf(brokerConf);
+    brokerConf.setProperty(CommonConstants.Broker.USE_MSE_TO_FILL_EMPTY_RESPONSE_SCHEMA, true);
+  }
+
   @BeforeClass
   public void setUp()
       throws Exception {

pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PauselessRealtimeIngestionNewSegmentMetadataCreationFailureTest.java

@@ -70,7 +70,6 @@ public void testSegmentAssignment()
               segmentZKMetadata -> segmentZKMetadata.getStatus() == CommonConstants.Segment.Realtime.Status.COMMITTING)
           .count() == NUM_REALTIME_SEGMENTS_ZK_METADATA_WITH_FAILURE;
     }, 1000, 100000, "Some segments are still IN_PROGRESS");
-
     runValidationAndVerify();
   }
 }

pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/SegmentGenerationMinionClusterIntegrationTest.java

@@ -29,7 +29,9 @@
 import org.apache.pinot.spi.config.task.AdhocTaskConfig;
 import org.apache.pinot.spi.data.FieldSpec;
 import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.env.PinotConfiguration;
 import org.apache.pinot.spi.ingestion.batch.BatchConfigProperties;
+import org.apache.pinot.spi.utils.CommonConstants;
 import org.apache.pinot.spi.utils.JsonUtils;
 import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
 import org.apache.pinot.util.TestUtils;
@@ -46,6 +48,12 @@
 public class SegmentGenerationMinionClusterIntegrationTest extends BaseClusterIntegrationTest {
   private static final Logger LOGGER = LoggerFactory.getLogger(SegmentGenerationMinionClusterIntegrationTest.class);
 
+  @Override
+  protected void overrideBrokerConf(PinotConfiguration brokerConf) {
+    super.overrideBrokerConf(brokerConf);
+    brokerConf.setProperty(CommonConstants.Broker.USE_MSE_TO_FILL_EMPTY_RESPONSE_SCHEMA, true);
+  }
+
   @BeforeClass
   public void setUp()
       throws Exception {

pinot-query-planner/src/main/java/org/apache/pinot/query/parser/utils/ParserUtils.java

@@ -62,27 +62,32 @@ public static boolean canCompileWithMultiStageEngine(String query, String databa
    * 2. Data schema has all columns set to default type (STRING) (when all segments pruned on server).
    *
    * Priority is:
-   * - Types from multi-stage engine validation for the given query.
+   * - Types from multi-stage engine validation for the given query (if allowed).
    * - Types from schema for the given table (only applicable to selection fields).
    * - Types from single-stage engine response (no action).
    *
    * Multi-stage engine schema will be available only if query compiles.
    */
-  public static void fillEmptyResponseSchema(BrokerResponse response, TableCache tableCache, Schema schema,
-      String database, String query) {
+  public static void fillEmptyResponseSchema(boolean useMSE, BrokerResponse response, TableCache tableCache,
+      Schema schema, String database, String query) {
     Preconditions.checkState(response.getNumRowsResultSet() == 0, "Cannot fill schema for non-empty response");
 
     DataSchema dataSchema = response.getResultTable() != null ? response.getResultTable().getDataSchema() : null;
 
     List<RelDataTypeField> dataTypeFields = null;
-    try {
-      QueryEnvironment queryEnvironment = new QueryEnvironment(database, tableCache, null);
-      RelRoot node = queryEnvironment.getRelRootIfCanCompile(query);
-      if (node != null && node.validatedRowType != null) {
-        dataTypeFields = node.validatedRowType.getFieldList();
+    // Turn on (with pinot.broker.use.mse.to.fill.empty.response.schema=true or query option
+    // useMSEToFillEmptyResponseSchema=true) only for clusters where no queries with huge IN clauses are expected
+    // (see https://github.com/apache/pinot/issues/15064)
+    if (useMSE) {
+      try {
+        QueryEnvironment queryEnvironment = new QueryEnvironment(database, tableCache, null);
+        RelRoot node = queryEnvironment.getRelRootIfCanCompile(query);
+        if (node != null && node.validatedRowType != null) {
+          dataTypeFields = node.validatedRowType.getFieldList();
+        }
+      } catch (Exception ignored) {
+        // Ignored
       }
-    } catch (Exception ignored) {
-      // Ignored
     }
 
     if (dataSchema == null && dataTypeFields == null) {

pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java

@@ -584,6 +584,10 @@ public static class QueryOptionKey {
 
         // Custom Query ID provided by the client
         public static final String CLIENT_QUERY_ID = "clientQueryId";
+
+        // Use MSE compiler when trying to fill a response with no schema metadata
+        // (overrides the "pinot.broker.use.mse.to.fill.empty.response.schema" broker conf)
+        public static final String USE_MSE_TO_FILL_EMPTY_RESPONSE_SCHEMA = "useMSEToFillEmptyResponseSchema";
       }
 
       public static class QueryOptionValue {
@@ -704,6 +708,10 @@ public enum Type {
     }
 
     public static final String PREFIX_OF_CONFIG_OF_PINOT_FS_FACTORY = "pinot.broker.storage.factory";
+
+    public static final String USE_MSE_TO_FILL_EMPTY_RESPONSE_SCHEMA =
+        "pinot.broker.use.mse.to.fill.empty.response.schema";
+    public static final boolean DEFAULT_USE_MSE_TO_FILL_EMPTY_RESPONSE_SCHEMA = false;
   }
 
   public static class Server {