@Nonnull public static PinotTaskConfig fromHelixTaskConfig(@Nonnull TaskConfig helixTaskConfig) { Map<String, String> configs = new HashMap<>(helixTaskConfig.getConfigMap()); // Inside Helix task config map, there are 3 extra Helix properties: TASK_COMMAND, TASK_ID, TASK_TARGET_PARTITION configs.remove(TASK_ID_KEY); configs.remove(TASK_COMMAND_KEY); configs.remove(TASK_TARGET_PARTITION_KEY); return new PinotTaskConfig(helixTaskConfig.getCommand(), configs); }
@Override public List<SegmentConversionResult> executeTask(@Nonnull PinotTaskConfig pinotTaskConfig) throws Exception { String taskType = pinotTaskConfig.getTaskType(); Map<String, String> configs = pinotTaskConfig.getConfigs(); String tableNameWithType = configs.get(MinionConstants.TABLE_NAME_KEY); String inputSegmentNames = configs.get(MinionConstants.SEGMENT_NAME_KEY);
/** * Get the child task configs for the given task name. * * @param taskName Task name * @return List of child task configs */ @Nonnull public synchronized List<PinotTaskConfig> getTaskConfigs(@Nonnull String taskName) { Collection<TaskConfig> helixTaskConfigs = _taskDriver.getJobConfig(getHelixJobName(taskName)).getTaskConfigMap().values(); List<PinotTaskConfig> taskConfigs = new ArrayList<>(helixTaskConfigs.size()); for (TaskConfig helixTaskConfig : helixTaskConfigs) { taskConfigs.add(PinotTaskConfig.fromHelixTaskConfig(helixTaskConfig)); } return taskConfigs; }
Preconditions.checkState(numConcurrentTasksPerInstance > 0); String taskType = pinotTaskConfigs.get(0).getTaskType(); String parentTaskName = TASK_PREFIX + taskType + TASK_NAME_SEPARATOR + System.currentTimeMillis(); LOGGER for (int i = 0; i < numChildTasks; i++) { PinotTaskConfig pinotTaskConfig = pinotTaskConfigs.get(i); Preconditions.checkState(pinotTaskConfig.getTaskType().equals(taskType)); helixTaskConfigs.add(pinotTaskConfig.toHelixTaskConfig(parentTaskName + TASK_NAME_SEPARATOR + i));
/** * Returns all the segments that have been scheduled but not finished * * @param taskType Task type * @param clusterInfoProvider Cluster info provider * @return Set of running segments */ public static Set<Segment> getRunningSegments(@Nonnull String taskType, @Nonnull ClusterInfoProvider clusterInfoProvider) { Set<Segment> runningSegments = new HashSet<>(); Map<String, TaskState> taskStates = clusterInfoProvider.getTaskStates(taskType); for (Map.Entry<String, TaskState> entry : taskStates.entrySet()) { TaskState taskState = entry.getValue(); if (taskState == TaskState.NOT_STARTED || taskState == TaskState.IN_PROGRESS || taskState == TaskState.STOPPED) { for (PinotTaskConfig pinotTaskConfig : clusterInfoProvider.getTaskConfigs(entry.getKey())) { Map<String, String> configs = pinotTaskConfig.getConfigs(); runningSegments.add( new Segment(configs.get(MinionConstants.TABLE_NAME_KEY), configs.get(MinionConstants.SEGMENT_NAME_KEY))); } } } return runningSegments; } }
@Override protected SegmentConversionResult convert(@Nonnull PinotTaskConfig pinotTaskConfig, @Nonnull File originalIndexDir, @Nonnull File workingDir) throws Exception { Map<String, String> configs = pinotTaskConfig.getConfigs(); new RawIndexConverter(originalIndexDir, workingDir, configs.get(MinionConstants.ConvertToRawIndexTask.COLUMNS_TO_CONVERT_KEY)).convert(); return new SegmentConversionResult.Builder().setFile(workingDir) .setTableNameWithType(configs.get(MinionConstants.TABLE_NAME_KEY)) .setSegmentName(configs.get(MinionConstants.SEGMENT_NAME_KEY)).build(); }
@Override public SegmentConversionResult executeTask(@Nonnull PinotTaskConfig pinotTaskConfig) throws Exception { String taskType = pinotTaskConfig.getTaskType(); Map<String, String> configs = pinotTaskConfig.getConfigs(); String tableNameWithType = configs.get(MinionConstants.TABLE_NAME_KEY); String segmentName = configs.get(MinionConstants.SEGMENT_NAME_KEY);
configs.put(MinionConstants.ConvertToRawIndexTask.COLUMNS_TO_CONVERT_KEY, columnsToConvertConfig); pinotTaskConfigs.add(new PinotTaskConfig(MinionConstants.ConvertToRawIndexTask.TASK_TYPE, configs)); tableNumTasks++;
@Override protected SegmentConversionResult convert(@Nonnull PinotTaskConfig pinotTaskConfig, @Nonnull File originalIndexDir, @Nonnull File workingDir) throws Exception { Map<String, String> configs = pinotTaskConfig.getConfigs(); String tableNameWithType = configs.get(MinionConstants.TABLE_NAME_KEY); String rawTableName = TableNameBuilder.extractRawTableName(tableNameWithType); SegmentPurger.RecordPurgerFactory recordPurgerFactory = MINION_CONTEXT.getRecordPurgerFactory(); SegmentPurger.RecordPurger recordPurger = recordPurgerFactory != null ? recordPurgerFactory.getRecordPurger(rawTableName) : null; SegmentPurger.RecordModifierFactory recordModifierFactory = MINION_CONTEXT.getRecordModifierFactory(); SegmentPurger.RecordModifier recordModifier = recordModifierFactory != null ? recordModifierFactory.getRecordModifier(rawTableName) : null; SegmentPurger segmentPurger = new SegmentPurger(originalIndexDir, workingDir, recordPurger, recordModifier); File purgedSegmentFile = segmentPurger.purgeSegment(); if (purgedSegmentFile == null) { purgedSegmentFile = originalIndexDir; } return new SegmentConversionResult.Builder().setFile(purgedSegmentFile).setTableNameWithType(tableNameWithType) .setSegmentName(configs.get(MinionConstants.SEGMENT_NAME_KEY)) .setCustomProperty(RECORD_PURGER_KEY, segmentPurger.getRecordPurger()) .setCustomProperty(RECORD_MODIFIER_KEY, segmentPurger.getRecordModifier()) .setCustomProperty(NUM_RECORDS_PURGED_KEY, segmentPurger.getNumRecordsPurged()) .setCustomProperty(NUM_RECORDS_MODIFIED_KEY, segmentPurger.getNumRecordsModified()).build(); }
@Override public Boolean executeTask(@Nonnull PinotTaskConfig pinotTaskConfig) { assertTrue(MINION_CONTEXT.getDataDir().exists()); assertNotNull(MINION_CONTEXT.getMinionMetrics()); assertNotNull(MINION_CONTEXT.getMinionVersion()); assertEquals(pinotTaskConfig.getTaskType(), TestTaskGenerator.TASK_TYPE); Map<String, String> configs = pinotTaskConfig.getConfigs(); assertEquals(configs.size(), 2); String offlineTableName = configs.get("tableName"); assertEquals(TableNameBuilder.getTableTypeFromTableName(offlineTableName), TableType.OFFLINE); String rawTableName = TableNameBuilder.extractRawTableName(offlineTableName); assertTrue(rawTableName.equals(TABLE_NAME_1) || rawTableName.equals(TABLE_NAME_2)); assertEquals(configs.get("tableType"), TableType.OFFLINE.toString()); do { if (_cancelled) { throw new TaskCancelledException("Task has been cancelled"); } } while (HOLD.get()); return true; } };
@Nonnull @Override public List<PinotTaskConfig> generateTasks(@Nonnull List<TableConfig> tableConfigs) { assertEquals(tableConfigs.size(), 2); // Generate at most 2 tasks if (_clusterInfoProvider.getTaskStates(TASK_TYPE).size() >= 2) { return Collections.emptyList(); } List<PinotTaskConfig> taskConfigs = new ArrayList<>(); for (TableConfig tableConfig : tableConfigs) { Map<String, String> configs = new HashMap<>(); configs.put("tableName", tableConfig.getTableName()); configs.put("tableType", tableConfig.getTableType().toString()); taskConfigs.add(new PinotTaskConfig(TASK_TYPE, configs)); } return taskConfigs; }
@Test public void testConvert() throws Exception { PurgeTaskExecutor purgeTaskExecutor = new PurgeTaskExecutor(); PinotTaskConfig pinotTaskConfig = new PinotTaskConfig(MinionConstants.PurgeTask.TASK_TYPE, Collections .singletonMap(MinionConstants.TABLE_NAME_KEY, TableNameBuilder.OFFLINE.tableNameWithType(TABLE_NAME))); File purgedIndexDir = purgeTaskExecutor.convert(pinotTaskConfig, _originalIndexDir, PURGED_SEGMENT_DIR).getFile(); try (PinotSegmentRecordReader pinotSegmentRecordReader = new PinotSegmentRecordReader(purgedIndexDir)) { int numRecordsRemaining = 0; int numRecordsModified = 0; GenericRow row = new GenericRow(); while (pinotSegmentRecordReader.hasNext()) { row = pinotSegmentRecordReader.next(row); numRecordsRemaining++; if (row.getValue(D1).equals(Integer.MAX_VALUE)) { numRecordsModified++; } } Assert.assertEquals(numRecordsRemaining, NUM_ROWS - 1); Assert.assertEquals(numRecordsModified, NUM_ROWS - 1); } }