[hive] branch master updated: HIVE-23429 : LLAP: Optimize retrieving queryId details in LlapTaskCommunicator (Rajesh Balamohan via Ashutosh Chauhan)

[ha...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

hashutosh pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/hive.git


The following commit(s) were added to refs/heads/master by this push:
     new 7e77f25  HIVE-23429 : LLAP: Optimize retrieving queryId details in LlapTaskCommunicator (Rajesh Balamohan via Ashutosh Chauhan)
7e77f25 is described below

commit 7e77f25f1361d0fff8b97e6fc45052bc280e5f7b
Author: Rajesh Balamohan <rb...@apache.org>
AuthorDate: Mon May 11 07:24:04 2020 -0700

    HIVE-23429 : LLAP: Optimize retrieving queryId details in LlapTaskCommunicator (Rajesh Balamohan via Ashutosh Chauhan)
    
    Signed-off-by: Ashutosh Chauhan <ha...@apache.org>
---
 .../hive/llap/tezplugins/LlapTaskCommunicator.java     | 18 ++++++++++++++----
 .../org/apache/hadoop/hive/ql/exec/tez/TezTask.java    |  3 +++
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/llap-tez/src/java/org/apache/hadoop/hive/llap/tezplugins/LlapTaskCommunicator.java b/llap-tez/src/java/org/apache/hadoop/hive/llap/tezplugins/LlapTaskCommunicator.java
index b168f76..36a2d6b 100644
--- a/llap-tez/src/java/org/apache/hadoop/hive/llap/tezplugins/LlapTaskCommunicator.java
+++ b/llap-tez/src/java/org/apache/hadoop/hive/llap/tezplugins/LlapTaskCommunicator.java
@@ -94,12 +94,14 @@ import org.apache.tez.dag.api.TezUncheckedException;
 import org.apache.tez.dag.api.UserPayload;
 import org.apache.tez.dag.api.event.VertexStateUpdate;
 import org.apache.tez.dag.app.TezTaskCommunicatorImpl;
+import org.apache.tez.dag.app.dag.DAG;
 import org.apache.tez.dag.records.TezTaskAttemptID;
 import org.apache.tez.runtime.api.TaskFailureType;
 import org.apache.tez.runtime.api.impl.TaskSpec;
 import org.apache.tez.runtime.api.impl.TezHeartbeatRequest;
 import org.apache.tez.runtime.api.impl.TezHeartbeatResponse;
 import org.apache.tez.serviceplugins.api.ContainerEndReason;
+import org.apache.tez.serviceplugins.api.DagInfo;
 import org.apache.tez.serviceplugins.api.ServicePluginErrorDefaults;
 import org.apache.tez.serviceplugins.api.TaskAttemptEndReason;
 import org.apache.tez.serviceplugins.api.TaskCommunicatorContext;
@@ -398,11 +400,9 @@ public class LlapTaskCommunicator extends TezTaskCommunicatorImpl {
         credentialsChanged, priority);
     int dagId = taskSpec.getTaskAttemptID().getTaskID().getVertexID().getDAGId().getId();
     if (currentQueryIdentifierProto == null || (dagId != currentQueryIdentifierProto.getDagIndex())) {
-      // TODO HiveQueryId extraction by parsing the Processor payload is ugly. This can be improved
-      // once TEZ-2672 is fixed.
-      String hiveQueryId;
+      String hiveQueryId = extractQueryIdFromContext();
       try {
-        hiveQueryId = extractQueryId(taskSpec);
+        hiveQueryId = (hiveQueryId == null) ? extractQueryId(taskSpec) : hiveQueryId;
       } catch (IOException e) {
         throw new RuntimeException("Failed to extract query id from task spec: " + taskSpec, e);
       }
@@ -820,12 +820,22 @@ public class LlapTaskCommunicator extends TezTaskCommunicatorImpl {
     // is likely already happening.
   }
 
+  // Needed for GenericUDTFGetSplits, where TaskSpecs are generated
   private String extractQueryId(TaskSpec taskSpec) throws IOException {
     UserPayload processorPayload = taskSpec.getProcessorDescriptor().getUserPayload();
     Configuration conf = TezUtils.createConfFromUserPayload(processorPayload);
     return HiveConf.getVar(conf, HiveConf.ConfVars.HIVEQUERYID);
   }
 
+  private String extractQueryIdFromContext() {
+    //TODO: Remove following instance of check, When TEZ-2672 exposes getConf from DagInfo
+    DagInfo dagInfo = getContext().getCurrentDagInfo();
+    if (dagInfo instanceof DAG) {
+      return ((DAG)dagInfo).getConf().get(ConfVars.HIVEQUERYID.varname);
+    }
+    return null;
+  }
+
   private SubmitWorkRequestProto constructSubmitWorkRequest(ContainerId containerId,
                                                             TaskSpec taskSpec,
                                                             FragmentRuntimeInfo fragmentRuntimeInfo,
diff --git a/ql/src/java/org/apache/hadoop/hive/ql/exec/tez/TezTask.java b/ql/src/java/org/apache/hadoop/hive/ql/exec/tez/TezTask.java
index 854bc89..b1bf2f8 100644
--- a/ql/src/java/org/apache/hadoop/hive/ql/exec/tez/TezTask.java
+++ b/ql/src/java/org/apache/hadoop/hive/ql/exec/tez/TezTask.java
@@ -424,6 +424,9 @@ public class TezTask extends Task<TezWork> {
         .put("description", ctx.getCmd());
     String dagInfo = json.toString();
 
+    String queryId = HiveConf.getVar(conf, HiveConf.ConfVars.HIVEQUERYID);
+    dag.setConf(HiveConf.ConfVars.HIVEQUERYID.varname, queryId);
+
     if (LOG.isDebugEnabled()) {
       LOG.debug("DagInfo: " + dagInfo);
     }