[incubator-uniffle] branch master updated: [#872] feat(tez): Get parameter from Inputcontext then provide util function (#915)

[ro...@apache.org]

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

roryqi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-uniffle.git


The following commit(s) were added to refs/heads/master by this push:
     new 1033085f [#872] feat(tez): Get parameter from Inputcontext then provide util function (#915)
1033085f is described below

commit 1033085f018b86d1b2cc9a8ec8991ff9a3b77f13
Author: Qing <11...@qq.com>
AuthorDate: Wed May 31 21:59:11 2023 +0800

    [#872] feat(tez): Get parameter from Inputcontext then provide util function (#915)
    
    ### What changes were proposed in this pull request?
    Get parameter from Inputcontext and then provide util function
    
    ### Why are the changes needed?
    Fix: https://github.com/apache/incubator-uniffle/issues/872
    
    ### Does this PR introduce _any_ user-facing change?
    No.
    
    ### How was this patch tested?
    unit test
---
 .../org/apache/tez/common/InputContextUtils.java   | 44 +++++++++++++++++++
 .../apache/tez/common/InputContextUtilsTest.java   | 49 ++++++++++++++++++++++
 2 files changed, 93 insertions(+)

diff --git a/client-tez/src/main/java/org/apache/tez/common/InputContextUtils.java b/client-tez/src/main/java/org/apache/tez/common/InputContextUtils.java
new file mode 100644
index 00000000..88cdacf7
--- /dev/null
+++ b/client-tez/src/main/java/org/apache/tez/common/InputContextUtils.java
@@ -0,0 +1,44 @@
+/*
+ * 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.tez.common;
+
+import org.apache.tez.dag.records.TezTaskAttemptID;
+import org.apache.tez.runtime.api.InputContext;
+
+public class InputContextUtils {
+
+  private InputContextUtils() {
+  }
+
+  public static TezTaskAttemptID getTezTaskAttemptID(InputContext inputContext) {
+    String uniqueIdentifier = inputContext.getUniqueIdentifier();
+    return TezTaskAttemptID.fromString(uniqueIdentifier.substring(0, uniqueIdentifier.length() - 6));
+  }
+
+  /**
+   *
+   * @param inputContext
+   * @return Compute shuffle id using InputContext
+   */
+  public static int computeShuffleId(InputContext inputContext) {
+    int dagIdentifier = inputContext.getDagIdentifier();
+    String sourceVertexName = inputContext.getSourceVertexName();
+    String taskVertexName = inputContext.getTaskVertexName();
+    return RssTezUtils.computeShuffleId(dagIdentifier, sourceVertexName, taskVertexName);
+  }
+}
diff --git a/client-tez/src/test/java/org/apache/tez/common/InputContextUtilsTest.java b/client-tez/src/test/java/org/apache/tez/common/InputContextUtilsTest.java
new file mode 100644
index 00000000..350291d5
--- /dev/null
+++ b/client-tez/src/test/java/org/apache/tez/common/InputContextUtilsTest.java
@@ -0,0 +1,49 @@
+/*
+ * 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.tez.common;
+
+import org.apache.tez.dag.records.TezTaskAttemptID;
+import org.apache.tez.runtime.api.InputContext;
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+public class InputContextUtilsTest {
+
+  @Test
+  public void testGetTezTaskAttemptID() {
+    InputContext inputContext = mock(InputContext.class);
+    when(inputContext.getUniqueIdentifier()).thenReturn("attempt_1685094627632_0157_1_01_000000_0_10006");
+
+    TezTaskAttemptID rightTaskAttemptID = TezTaskAttemptID.fromString("attempt_1685094627632_0157_1_01_000000_0");
+    assertEquals(rightTaskAttemptID, InputContextUtils.getTezTaskAttemptID(inputContext));
+  }
+
+
+  @Test
+  public void testComputeShuffleId() {
+    InputContext inputContext = mock(InputContext.class);
+    when(inputContext.getDagIdentifier()).thenReturn(1);
+    when(inputContext.getSourceVertexName()).thenReturn("Map 1");
+    when(inputContext.getTaskVertexName()).thenReturn("Reducer 1");
+
+    assertEquals(1001601, InputContextUtils.computeShuffleId(inputContext));
+  }
+}