diff --git a/ql/src/java/org/apache/hadoop/hive/ql/exec/Utilities.java b/ql/src/java/org/apache/hadoop/hive/ql/exec/Utilities.java index 024e05cdff5e3171521eec2cbdb7566472bab815..14aaceb30eb66f370f0fa7b8baef80503333b166 100644 --- a/ql/src/java/org/apache/hadoop/hive/ql/exec/Utilities.java +++ b/ql/src/java/org/apache/hadoop/hive/ql/exec/Utilities.java @@ -1777,8 +1777,26 @@ public static String replaceTaskIdFromFilename(String filename, String fileId) { return (ret); } - private static String replaceTaskId(String taskId, int bucketNum) { - return replaceTaskId(taskId, String.valueOf(bucketNum)); + /** + * Replace taskId with input bucketNum. For example, if taskId is 000000 and bucketNum is 1, + * return should be 000001; if taskId is (ds%3D1)000000 and bucketNum is 1, return should be + * (ds%3D1)000001. This method is different from the replaceTaskId(String, String) method. + * In this method, the pattern is in taskId. + * @param taskId + * @param bucketNum + * @return + */ + public static String replaceTaskId(String taskId, int bucketNum) { + String bucketNumStr = String.valueOf(bucketNum); + Matcher m = PREFIXED_TASK_ID_REGEX.matcher(taskId); + if (!m.matches()) { + LOG.warn("Unable to determine bucket number from task id: " + taskId + ". Using " + + "task ID as bucket number."); + return adjustBucketNumLen(bucketNumStr, taskId); + } else { + String adjustedBucketNum = adjustBucketNumLen(bucketNumStr, m.group(2)); + return (m.group(1) == null ? "" : m.group(1)) + adjustedBucketNum; + } } /** diff --git a/ql/src/test/org/apache/hadoop/hive/ql/exec/TestUtilities.java b/ql/src/test/org/apache/hadoop/hive/ql/exec/TestUtilities.java index 69f88893f2eaca99d3d999b9b628acee9007773b..ccc21e9eb2113cefb073f82fa7f24ebb52402235 100644 --- a/ql/src/test/org/apache/hadoop/hive/ql/exec/TestUtilities.java +++ b/ql/src/test/org/apache/hadoop/hive/ql/exec/TestUtilities.java @@ -140,4 +140,15 @@ public void testGetJarFilesByPath() { FileUtils.deleteQuietly(f); } } + + public void testReplaceTaskId() { + String taskID = "000000"; + int bucketNum = 1; + String newTaskID = Utilities.replaceTaskId(taskID, bucketNum); + Assert.assertEquals("000001", newTaskID); + taskID = "(ds%3D1)000001"; + newTaskID = Utilities.replaceTaskId(taskID, 5); + Assert.assertEquals("(ds%3D1)000005", newTaskID); + } + }