DRILL-8316: Convert Druid Storage Plugin to EVF & V2 JSON Reader

contrib/storage-druid/.gitignore

@@ -0,0 +1,2 @@
+# Directory to store oauth tokens for testing Googlesheets Storage plugin
+/src/test/resources/logback-test.xml

contrib/storage-druid/README.md

@@ -4,7 +4,7 @@ Drill druid storage plugin allows you to perform SQL queries against Druid datas
 This storage plugin is part of [Apache Drill](https://github.com/apache/drill)
 
 ### Tested with Druid version
-[0.22.0](https://github.com/apache/druid/releases/tag/druid-0.22.0)
+[30.0.0](https://github.com/apache/druid/releases/tag/druid-0.22.0)
 
 ### Druid API
 
@@ -33,27 +33,27 @@ Following is the default registration configuration.
 
 ### Druid storage plugin developer notes.
 
-* Building the plugin 
+* Building the plugin
 
     `mvn install -pl contrib/storage-druid`
 
 * Building DRILL
 
     `mvn clean install -DskipTests`
-    
+
 * Start Drill In Embedded Mode (mac)
 
     ```shell script
     distribution/target/apache-drill-1.20.0-SNAPSHOT/apache-drill-1.20.0-SNAPSHOT/bin/drill-embedded
     ```
-  
+
 * Starting Druid (Docker and Docker Compose required)
     ```
     cd contrib/storage-druid/src/test/resources/druid
     docker-compose up -d
     ```
-  
+
   * There is an `Indexing Task Json` in the same folder as the docker compose file. It can be used to ingest the wikipedia datasource.
-  
+
   * Make sure the druid storage plugin is enabled in Drill.
 

contrib/storage-druid/pom.xml

@@ -29,7 +29,7 @@
     <artifactId>drill-druid-storage</artifactId>
     <name>Drill : Contrib : Storage : Druid</name>
     <properties>
-        <druid.TestSuite>**/DruidTestSuit.class</druid.TestSuite>
+        <druid.TestSuite>**/DruidTestSuite.class</druid.TestSuite>
     </properties>
     <dependencies>
         <dependency>
@@ -53,13 +53,6 @@
             <version>${project.version}</version>
             <scope>test</scope>
         </dependency>
-        <dependency>
-            <groupId>org.assertj</groupId>
-            <artifactId>assertj-core</artifactId>
-            <!-- use 2.9.1 for Java 7 projects -->
-            <version>3.11.1</version>
-            <scope>test</scope>
-        </dependency>
     </dependencies>
 
     <build>

contrib/storage-druid/src/main/java/org/apache/drill/exec/store/druid/DruidBatchRecordReader.java

@@ -0,0 +1,162 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.drill.exec.store.druid;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import org.apache.drill.common.exceptions.CustomErrorContext;
+import org.apache.drill.common.exceptions.UserException;
+import org.apache.drill.common.expression.SchemaPath;
+import org.apache.drill.exec.physical.impl.scan.framework.ManagedReader;
+import org.apache.drill.exec.physical.impl.scan.framework.SchemaNegotiator;
+import org.apache.drill.exec.physical.resultSet.ResultSetLoader;
+import org.apache.drill.exec.store.druid.DruidSubScan.DruidSubScanSpec;
+import org.apache.drill.exec.store.druid.common.DruidFilter;
+import org.apache.drill.exec.store.druid.druid.DruidScanResponse;
+import org.apache.drill.exec.store.druid.druid.ScanQuery;
+import org.apache.drill.exec.store.druid.druid.ScanQueryBuilder;
+import org.apache.drill.exec.store.druid.rest.DruidQueryClient;
+import org.apache.drill.exec.store.easy.json.loader.JsonLoaderImpl;
+import org.apache.drill.exec.store.easy.json.loader.JsonLoaderImpl.JsonLoaderBuilder;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.math.BigInteger;
+import java.util.ArrayList;
+import java.util.List;
+
+public class DruidBatchRecordReader implements ManagedReader<SchemaNegotiator> {
+  private static final Logger logger = LoggerFactory.getLogger(DruidBatchRecordReader.class);
+  private static final int BATCH_SIZE = 4096;
+  private static final ObjectMapper objectMapper = new ObjectMapper();
+  private final DruidStoragePlugin plugin;
+  private final DruidSubScan.DruidSubScanSpec scanSpec;
+  private final List<String> columns;
+  private final DruidFilter filter;
+  private final DruidQueryClient druidQueryClient;
+  private final DruidOffsetTracker offsetTracker;
+  private int maxRecordsToRead = -1;
+  private JsonLoaderBuilder jsonBuilder;
+  private JsonLoaderImpl jsonLoader;
+  private SchemaNegotiator negotiator;
+  private ResultSetLoader resultSetLoader;
+  private CustomErrorContext errorContext;
+
+
+  public DruidBatchRecordReader(DruidSubScan subScan,
+                                DruidSubScanSpec subScanSpec,
+                                List<SchemaPath> projectedColumns,
+                                int maxRecordsToRead,
+                                DruidStoragePlugin plugin, DruidOffsetTracker offsetTracker) {
+    this.columns = new ArrayList<>();
+    this.maxRecordsToRead = maxRecordsToRead;
+    this.plugin = plugin;
+    this.scanSpec = subScanSpec;
+    this.filter = subScanSpec.getFilter();
+    this.druidQueryClient = plugin.getDruidQueryClient();
+    this.offsetTracker = offsetTracker;
+  }
+
+  @Override
+  public boolean open(SchemaNegotiator negotiator) {
+    this.negotiator = negotiator;
+    this.errorContext = this.negotiator.parentErrorContext();
+    this.negotiator.batchSize(BATCH_SIZE);
+    this.negotiator.setErrorContext(errorContext);
+
+    resultSetLoader = this.negotiator.build();
+
+
+    return true;
+  }
+
+  @Override
+  public boolean next() {
+    jsonBuilder = new JsonLoaderBuilder()
+        .resultSetLoader(resultSetLoader)
+        .standardOptions(negotiator.queryOptions())
+        .errorContext(errorContext);
+    int eventCounter = 0;
+    boolean result = false;
+    try {
+      String query = getQuery();
+      logger.debug("Executing query: {}", query);
+      DruidScanResponse druidScanResponse = druidQueryClient.executeQuery(query);
+      setNextOffset(druidScanResponse);
+
+      StringBuilder events = new StringBuilder();
+      for (ObjectNode eventNode : druidScanResponse.getEvents()) {
+        events.append(eventNode);
+        events.append("\n");
+        eventCounter++;
+      }
+
+
+      jsonLoader = (JsonLoaderImpl) jsonBuilder
+          .fromString(events.toString())
+          .build();
+
+      result = jsonLoader.readBatch();
+
+      if (eventCounter < BATCH_SIZE) {
+        return false;
+      } else {
+        return result;
+      }
+    } catch (Exception e) {
+      throw UserException
+        .dataReadError(e)
+        .message("Failure while executing druid query: " + e.getMessage())
+        .addContext(errorContext)
+        .build(logger);
+    }
+  }
+
+  @Override
+  public void close() {
+    if (jsonLoader != null) {
+      jsonLoader.close();
+      jsonLoader = null;
+    }
+  }
+
+  private String getQuery() throws JsonProcessingException {
+    int queryThreshold =
+      maxRecordsToRead >= 0
+        ? Math.min(BATCH_SIZE, maxRecordsToRead)
+        : BATCH_SIZE;
+    ScanQueryBuilder scanQueryBuilder = plugin.getScanQueryBuilder();
+    ScanQuery scanQuery =
+      scanQueryBuilder.build(
+        scanSpec.dataSourceName,
+        columns,
+        filter,
+        offsetTracker.getOffset(),
+        queryThreshold,
+        scanSpec.getMinTime(),
+        scanSpec.getMaxTime()
+      );
+    return objectMapper.writeValueAsString(scanQuery);
+  }
+
+  private void setNextOffset(DruidScanResponse druidScanResponse) {
+    offsetTracker.setNextOffset(BigInteger.valueOf(druidScanResponse.getEvents().size()));
+  }
+}

contrib/storage-druid/src/main/java/org/apache/drill/exec/store/druid/DruidGroupScan.java

@@ -25,12 +25,14 @@
 
 import org.apache.drill.common.PlanStringBuilder;
 import org.apache.drill.common.expression.SchemaPath;
+import org.apache.drill.exec.metastore.MetadataProviderManager;
 import org.apache.drill.exec.physical.EndpointAffinity;
 import org.apache.drill.exec.physical.base.AbstractGroupScan;
 import org.apache.drill.exec.physical.base.GroupScan;
 import org.apache.drill.exec.physical.base.PhysicalOperator;
 import org.apache.drill.exec.physical.base.ScanStats;
 import org.apache.drill.exec.proto.CoordinationProtos;
+import org.apache.drill.exec.record.metadata.TupleMetadata;
 import org.apache.drill.exec.store.StoragePluginRegistry;
 
 import org.apache.drill.exec.store.schedule.AffinityCreator;
@@ -44,6 +46,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
 
@@ -54,7 +57,7 @@ public class DruidGroupScan extends AbstractGroupScan {
   private static final long DEFAULT_TABLET_SIZE = 1000;
   private final DruidScanSpec scanSpec;
   private final DruidStoragePlugin storagePlugin;
-
+  private final MetadataProviderManager metadataProviderManager;
   private List<SchemaPath> columns;
   private boolean filterPushedDown = false;
   private int maxRecordsToRead;
@@ -73,19 +76,20 @@ public DruidGroupScan(@JsonProperty("userName") String userName,
         pluginRegistry.resolve(storagePluginConfig, DruidStoragePlugin.class),
         scanSpec,
         columns,
-        maxRecordsToRead);
+        maxRecordsToRead, null);
   }
 
   public DruidGroupScan(String userName,
                         DruidStoragePlugin storagePlugin,
                         DruidScanSpec scanSpec,
                         List<SchemaPath> columns,
-                        int maxRecordsToRead) {
+                        int maxRecordsToRead, MetadataProviderManager metadataProviderManager) {
     super(userName);
     this.storagePlugin = storagePlugin;
     this.scanSpec = scanSpec;
     this.columns = columns == null || columns.size() == 0? ALL_COLUMNS : columns;
     this.maxRecordsToRead = maxRecordsToRead;
+    this.metadataProviderManager = metadataProviderManager;
     init();
   }
 
@@ -102,6 +106,7 @@ private DruidGroupScan(DruidGroupScan that) {
     this.filterPushedDown = that.filterPushedDown;
     this.druidWorkList = that.druidWorkList;
     this.assignments = that.assignments;
+    this.metadataProviderManager = that.metadataProviderManager;
   }
 
   @Override
@@ -163,7 +168,8 @@ private void init() {
           getScanSpec().getFilter(),
           getDatasourceSize(),
           getDataSourceMinTime(),
-          getDataSourceMaxTime()
+          getDataSourceMaxTime(),
+          getSchema()
         )
       );
     druidWorkList.add(druidWork);
@@ -225,12 +231,13 @@ public DruidSubScan getSpecificScan(int minorFragmentId) {
             druidWork.getDruidSubScanSpec().getFilter(),
             druidWork.getDruidSubScanSpec().getDataSourceSize(),
             druidWork.getDruidSubScanSpec().getMinTime(),
-            druidWork.getDruidSubScanSpec().getMaxTime()
+            druidWork.getDruidSubScanSpec().getMaxTime(),
+            druidWork.getDruidSubScanSpec().getSchema()
           )
         );
     }
 
-    return new DruidSubScan(getUserName(), storagePlugin, scanSpecList, this.columns, this.maxRecordsToRead);
+    return new DruidSubScan(getUserName(), storagePlugin, scanSpecList, this.columns, this.maxRecordsToRead, getSchema());
   }
 
   @JsonIgnore
@@ -283,13 +290,30 @@ public int getMaxRecordsToRead() {
     return maxRecordsToRead;
   }
 
+  @JsonIgnore
+  public MetadataProviderManager getMetadataProviderManager() {
+    return metadataProviderManager;
+  }
+
+  public TupleMetadata getSchema() {
+    if (metadataProviderManager == null) {
+      return null;
+    }
+    try {
+      return metadataProviderManager.getSchemaProvider().read().getSchema();
+    } catch (IOException | NullPointerException e) {
+      return null;
+    }
+  }
+
   @Override
   public String toString() {
     return new PlanStringBuilder(this)
-        .field("druidScanSpec", scanSpec)
-        .field("columns", columns)
-        .field("druidStoragePlugin", storagePlugin)
-        .toString();
+      .field("druidScanSpec", scanSpec)
+      .field("columns", columns)
+      .field("druidStoragePlugin", storagePlugin)
+      .field("schema", getSchema())
+      .toString();
   }
 
   @Override

contrib/storage-druid/src/main/java/org/apache/drill/exec/store/druid/DruidOffsetTracker.java

@@ -0,0 +1,41 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.drill.exec.store.druid;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.math.BigInteger;
+
+public class DruidOffsetTracker {
+  private static final Logger logger = LoggerFactory.getLogger(DruidOffsetTracker.class);
+  private BigInteger nextOffset;
+
+  public DruidOffsetTracker() {
+    this.nextOffset = BigInteger.ZERO;
+  }
+
+  public BigInteger getOffset() {
+    return nextOffset;
+  }
+
+  public void setNextOffset(BigInteger offset) {
+    nextOffset = nextOffset.add(offset);
+    logger.debug("Incrementing offset by {}", offset);
+  }
+}