[multistage] support database in v2

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

@@ -81,7 +81,7 @@
 public class MultiStageBrokerRequestHandler extends BaseBrokerRequestHandler {
   private static final Logger LOGGER = LoggerFactory.getLogger(MultiStageBrokerRequestHandler.class);
 
-  private final QueryEnvironment _queryEnvironment;
+  private final WorkerManager _workerManager;
   private final MailboxService _mailboxService;
   private final QueryDispatcher _queryDispatcher;
 
@@ -93,9 +93,7 @@ public MultiStageBrokerRequestHandler(PinotConfiguration config, String brokerId
     LOGGER.info("Using Multi-stage BrokerRequestHandler.");
     String hostname = config.getProperty(CommonConstants.MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_HOSTNAME);
     int port = Integer.parseInt(config.getProperty(CommonConstants.MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_PORT));
-    _queryEnvironment = new QueryEnvironment(new TypeFactory(new TypeSystem()),
-        CalciteSchemaBuilder.asRootSchema(new PinotCatalog(tableCache)),
-        new WorkerManager(hostname, port, routingManager), _tableCache);
+    _workerManager = new WorkerManager(hostname, port, _routingManager);
     _mailboxService = new MailboxService(hostname, port, config);
     _queryDispatcher = new QueryDispatcher(_mailboxService);
 
@@ -128,9 +126,11 @@ protected BrokerResponse handleRequest(long requestId, String query, @Nullable S
       queryTimeoutMs = timeoutMsFromQueryOption == null ? _brokerTimeoutMs : timeoutMsFromQueryOption;
       // Compile the request
       compilationStartTimeNs = System.nanoTime();
+      QueryEnvironment queryEnvironment = new QueryEnvironment(new TypeFactory(new TypeSystem()),
+          CalciteSchemaBuilder.asRootSchema(new PinotCatalog(_tableCache)), _workerManager, _tableCache);
       switch (sqlNodeAndOptions.getSqlNode().getKind()) {
         case EXPLAIN:
-          queryPlanResult = _queryEnvironment.explainQuery(query, sqlNodeAndOptions, requestId);
+          queryPlanResult = queryEnvironment.explainQuery(query, sqlNodeAndOptions, requestId);
           String plan = queryPlanResult.getExplainPlan();
           Set<String> tableNames = queryPlanResult.getTableNames();
           if (!hasTableAccess(requesterIdentity, tableNames, requestContext, httpHeaders)) {
@@ -140,7 +140,7 @@ protected BrokerResponse handleRequest(long requestId, String query, @Nullable S
           return constructMultistageExplainPlan(query, plan);
         case SELECT:
         default:
-          queryPlanResult = _queryEnvironment.planQuery(query, sqlNodeAndOptions, requestId);
+          queryPlanResult = queryEnvironment.planQuery(query, sqlNodeAndOptions, requestId);
           break;
       }
     } catch (WebApplicationException e) {

pinot-common/src/main/java/org/apache/pinot/common/utils/DatabaseUtils.java

@@ -84,4 +84,13 @@ public static String translateTableName(String tableName, HttpHeaders headers, b
   public static String translateTableName(String tableName, HttpHeaders headers) {
     return translateTableName(tableName, headers, false);
   }
+
+  public static boolean isPartOfDatabase(String tableName, @Nullable String databaseName) {
+    // assumes tableName will not have default database prefix ('default.')
+    if (StringUtils.isEmpty(databaseName) || databaseName.equalsIgnoreCase(CommonConstants.DEFAULT_DATABASE)) {
+      return StringUtils.split(tableName, '.').length == 1;
+    } else {
+      return tableName.startsWith(databaseName + ".");
+    }
+  }
 }

pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java

@@ -732,22 +732,14 @@ public List<String> getAllTables() {
   public List<String> getAllTables(@Nullable String databaseName) {
     List<String> tableNames = new ArrayList<>();
     for (String resourceName : getAllResources()) {
-      if (TableNameBuilder.isTableResource(resourceName) && isPartOfDatabase(resourceName, databaseName)) {
+      if (TableNameBuilder.isTableResource(resourceName)
+          && DatabaseUtils.isPartOfDatabase(resourceName, databaseName)) {
         tableNames.add(resourceName);
       }
     }
     return tableNames;
   }
 
-  private boolean isPartOfDatabase(String tableName, @Nullable String databaseName) {
-    // assumes tableName will not have default database prefix ('default.')
-    if (StringUtils.isEmpty(databaseName) || databaseName.equalsIgnoreCase(CommonConstants.DEFAULT_DATABASE)) {
-      return StringUtils.split(tableName, '.').length == 1;
-    } else {
-      return tableName.startsWith(databaseName + ".");
-    }
-  }
-
   /**
    * Get all offline table names from default database.
    *
@@ -766,7 +758,8 @@ public List<String> getAllOfflineTables() {
   public List<String> getAllOfflineTables(@Nullable String databaseName) {
     List<String> offlineTableNames = new ArrayList<>();
     for (String resourceName : getAllResources()) {
-      if (isPartOfDatabase(resourceName, databaseName) && TableNameBuilder.isOfflineTableResource(resourceName)) {
+      if (DatabaseUtils.isPartOfDatabase(resourceName, databaseName)
+          && TableNameBuilder.isOfflineTableResource(resourceName)) {
         offlineTableNames.add(resourceName);
       }
     }
@@ -790,7 +783,7 @@ public List<String> getAllDimensionTables() {
    */
   public List<String> getAllDimensionTables(@Nullable String databaseName) {
     return _tableCache.getAllDimensionTables().stream()
-        .filter(table -> isPartOfDatabase(table, databaseName))
+        .filter(table -> DatabaseUtils.isPartOfDatabase(table, databaseName))
         .collect(Collectors.toList());
   }
 
@@ -812,7 +805,8 @@ public List<String> getAllRealtimeTables() {
   public List<String> getAllRealtimeTables(@Nullable String databaseName) {
     List<String> realtimeTableNames = new ArrayList<>();
     for (String resourceName : getAllResources()) {
-      if (isPartOfDatabase(resourceName, databaseName) && TableNameBuilder.isRealtimeTableResource(resourceName)) {
+      if (DatabaseUtils.isPartOfDatabase(resourceName, databaseName)
+          && TableNameBuilder.isRealtimeTableResource(resourceName)) {
         realtimeTableNames.add(resourceName);
       }
     }
@@ -837,7 +831,8 @@ public List<String> getAllRawTables() {
   public List<String> getAllRawTables(@Nullable String databaseName) {
     Set<String> rawTableNames = new HashSet<>();
     for (String resourceName : getAllResources()) {
-      if (TableNameBuilder.isTableResource(resourceName) && isPartOfDatabase(resourceName, databaseName)) {
+      if (TableNameBuilder.isTableResource(resourceName)
+          && DatabaseUtils.isPartOfDatabase(resourceName, databaseName)) {
         rawTableNames.add(TableNameBuilder.extractRawTableName(resourceName));
       }
     }
@@ -1616,7 +1611,7 @@ public List<String> getSchemaNames(@Nullable String databaseName) {
     List<String> schemas = _propertyStore.getChildNames(
         PinotHelixPropertyStoreZnRecordProvider.forSchema(_propertyStore).getRelativePath(), AccessOption.PERSISTENT);
     if (schemas != null) {
-      return schemas.stream().filter(schemaName -> isPartOfDatabase(schemaName, databaseName))
+      return schemas.stream().filter(schemaName -> DatabaseUtils.isPartOfDatabase(schemaName, databaseName))
           .collect(Collectors.toList());
     }
     return Collections.emptyList();
@@ -4011,7 +4006,7 @@ public Map<String, List<InstanceInfo>> getTableToLiveBrokersMapping(@Nullable St
     ZNRecord znRecord = ev.getRecord();
     for (Map.Entry<String, Map<String, String>> tableToBrokersEntry : znRecord.getMapFields().entrySet()) {
       String tableName = tableToBrokersEntry.getKey();
-      if (!isPartOfDatabase(tableName, databaseName)) {
+      if (!DatabaseUtils.isPartOfDatabase(tableName, databaseName)) {
         continue;
       }
       Map<String, String> brokersToState = tableToBrokersEntry.getValue();

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

@@ -25,12 +25,19 @@
 import java.time.Instant;
 import java.time.ZoneId;
 import java.time.format.DateTimeFormatter;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.TimeUnit;
 import java.util.regex.Pattern;
+import javax.annotation.Nullable;
 import org.apache.commons.io.FileUtils;
+import org.apache.commons.lang3.StringUtils;
 import org.apache.pinot.common.exception.QueryException;
 import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.ingestion.IngestionConfig;
+import org.apache.pinot.spi.config.table.ingestion.TransformConfig;
+import org.apache.pinot.spi.data.FieldSpec;
+import org.apache.pinot.spi.data.MetricFieldSpec;
 import org.apache.pinot.spi.data.Schema;
 import org.apache.pinot.util.TestUtils;
 import org.testng.Assert;
@@ -47,6 +54,8 @@
 
 public class MultiStageEngineIntegrationTest extends BaseClusterIntegrationTestSet {
   private static final String SCHEMA_FILE_NAME = "On_Time_On_Time_Performance_2014_100k_subset_nonulls.schema";
+  private String _tableName = DEFAULT_TABLE_NAME;
+  private List<File> _avroFiles = new ArrayList<>();
 
   @Override
   protected String getSchemaFileName() {
@@ -72,17 +81,17 @@ public void setUp()
     addTableConfig(tableConfig);
 
     // Unpack the Avro files
-    List<File> avroFiles = unpackAvroData(_tempDir);
+    _avroFiles = unpackAvroData(_tempDir);
 
     // Create and upload segments
-    ClusterIntegrationTestUtils.buildSegmentsFromAvro(avroFiles, tableConfig, schema, 0, _segmentDir, _tarDir);
+    ClusterIntegrationTestUtils.buildSegmentsFromAvro(_avroFiles, tableConfig, schema, 0, _segmentDir, _tarDir);
     uploadSegments(getTableName(), _tarDir);
 
     // Set up the H2 connection
-    setUpH2Connection(avroFiles);
+    setUpH2Connection(_avroFiles);
 
     // Initialize the query generator
-    setUpQueryGenerator(avroFiles);
+    setUpQueryGenerator(_avroFiles);
 
     // Wait for all documents loaded
     waitForAllDocsLoaded(600_000L);
@@ -749,6 +758,76 @@ public void testSearch()
     assertNoError(jsonNode);
   }
 
+  @Override
+  protected String getTableName() {
+    return _tableName;
+  }
+
+  @Test
+  public void testWithDatabaseContext()
+      throws Exception {
+    try {
+      _tableName = "db1." + DEFAULT_TABLE_NAME;
+      String defaultCol = "ActualElapsedTime";
+      String customCol = "ActualElapsedTime_2";
+      Schema schema = createSchema();
+      schema.addField(new MetricFieldSpec(customCol, FieldSpec.DataType.INT));
+      addSchema(schema);
+      TableConfig tableConfig = createOfflineTableConfig();
+      assert tableConfig.getIndexingConfig().getNoDictionaryColumns() != null;
+      List<String> noDicCols = new ArrayList<>(DEFAULT_NO_DICTIONARY_COLUMNS);
+      noDicCols.add(customCol);
+      tableConfig.getIndexingConfig().setNoDictionaryColumns(noDicCols);
+      IngestionConfig ingestionConfig = new IngestionConfig();
+      ingestionConfig.setTransformConfigs(List.of(new TransformConfig(customCol, defaultCol)));
+      tableConfig.setIngestionConfig(ingestionConfig);
+      addTableConfig(tableConfig);
+
+      // Create and upload segments to 'db1.mytable'
+      TestUtils.ensureDirectoriesExistAndEmpty(_segmentDir, _tarDir);
+      ClusterIntegrationTestUtils.buildSegmentsFromAvro(_avroFiles, tableConfig, schema, 0, _segmentDir, _tarDir);
+      uploadSegments(getTableName(), _tarDir);
+
+      // Wait for all documents loaded
+      waitForAllDocsLoaded(600_000L);
+
+      // default database check. No database context passed
+      checkQueryResultForDBTest(defaultCol, DEFAULT_TABLE_NAME);
+
+      // default database check. Default database context passed as table prefix
+      checkQueryResultForDBTest(defaultCol, "default." + DEFAULT_TABLE_NAME);
+
+      // default database check. Default database context passed as SET database='dbName'
+      checkQueryResultForDBTest(defaultCol, DEFAULT_TABLE_NAME, "default");
+
+      // Using renamed column "ActualElapsedTime_2" to ensure that the same table is not being queried.
+      // custom database check. Database context passed as table prefix
+      checkQueryResultForDBTest(customCol, _tableName);
+
+      // custom database check. Database context passed as SET database='dbName'
+      checkQueryResultForDBTest(customCol, DEFAULT_TABLE_NAME, "db1");
+    } finally {
+      dropOfflineTable(_tableName);
+      _tableName = DEFAULT_TABLE_NAME;
+    }
+  }
+
+  private void checkQueryResultForDBTest(String column, String tableName)
+      throws Exception {
+    checkQueryResultForDBTest(column, tableName, null);
+  }
+
+  private void checkQueryResultForDBTest(String column, String tableName, @Nullable String database)
+      throws Exception {
+    String query = (StringUtils.isNotBlank(database) ? "SET database='" + database + "'; " : "")
+        + "select max(" + column + ") from " + tableName + ";";
+    // max value of 'ActualElapsedTime'
+    long expectedValue = 678;
+    JsonNode jsonNode = postQuery(query);
+    long result = jsonNode.get("resultTable").get("rows").get(0).get(0).asLong();
+    assertEquals(result, expectedValue);
+  }
+
   @AfterClass
   public void tearDown()
       throws Exception {

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

@@ -2772,11 +2772,7 @@ public void testQuerySourceWithDatabaseNameV2()
     testQuery(pinotQuery, h2Query);
 
     pinotQuery = "SELECT DISTINCT Carrier FROM default.mytable LIMIT 1000000";
-    JsonNode response = postQuery(pinotQuery);
-    JsonNode exceptions = response.get("exceptions");
-    assertFalse(exceptions.isEmpty(), "At least one exception was expected");
-    JsonNode firstException = exceptions.get(0);
-    assertEquals(firstException.get("errorCode").asInt(), QueryException.QUERY_PLANNING_ERROR_CODE);
+    testQuery(pinotQuery, h2Query);
   }
 
   @Test(dataProvider = "useBothQueryEngines")

pinot-query-planner/src/main/java/org/apache/calcite/jdbc/CalciteSchemaBuilder.java

@@ -61,4 +61,15 @@ public static CalciteSchema asRootSchema(Schema root) {
     }
     return rootSchema;
   }
+
+  public static CalciteSchema asSubSchema(Schema root, String name) {
+    CalciteSchema subSchema = CalciteSchema.createRootSchema(false, false, name, root);
+    SchemaPlus schemaPlus = subSchema.plus();
+    for (Map.Entry<String, List<Function>> e : FunctionRegistry.getRegisteredCalciteFunctionMap().entrySet()) {
+      for (Function f : e.getValue()) {
+        schemaPlus.add(e.getKey(), f);
+      }
+    }
+    return subSchema;
+  }
 }