diff --git a/pinot-core/src/main/java/org/apache/pinot/core/common/MinionConstants.java b/pinot-core/src/main/java/org/apache/pinot/core/common/MinionConstants.java index f5d1e8ce94..b91a4a9c2b 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/common/MinionConstants.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/common/MinionConstants.java @@ -312,4 +312,31 @@ public static class UpsertCompactMergeTask { public static final String MERGED_SEGMENTS_ZK_SUFFIX = ".mergedSegments"; } + + // Purges rows inside segment that match filter function. + public static class RowFilterPurgeTask { + public static final String TASK_TYPE = "RowFilterPurgeTask"; + + /** + * Task-config key: Groovy/expression filter function applied row-by-row. + * Rows for which the function evaluates to {@code true} are removed from the segment. + * This is intentionally separate from {@code ingestionConfig.filterConfig.filterFunction} so + * that the filter only affects the purge operation and does not alter regular ingestion behaviour + * of the table. + *

Example: {@code "Groovy({status == 'DELETED'}, status)"} + */ + public static final String FILTER_FUNCTION_KEY = "filterFunction"; + + /** + * Task-config key (boolean, default {@code false}): when {@code true} a row that throws during + * filter evaluation is silently marked incomplete instead of failing the task. + */ + public static final String FILTER_CONTINUE_ON_ERROR_KEY = "filterContinueOnError"; + + /** Custom-property key reported in SegmentConversionResult: number of rows retained in the segment. */ + public static final String NUM_ROWS_RETAINED_KEY = "numRowsRetained"; + + /** Custom-property key reported in SegmentConversionResult: number of rows dropped by the filter. */ + public static final String NUM_ROWS_FILTERED_KEY = "numRowsFiltered"; + } } diff --git a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/rowfilterpurge/RowFilterPurgeTaskExecutor.java b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/rowfilterpurge/RowFilterPurgeTaskExecutor.java new file mode 100644 index 0000000000..07a7b3cc01 --- /dev/null +++ b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/rowfilterpurge/RowFilterPurgeTaskExecutor.java @@ -0,0 +1,118 @@ +/** + * 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.pinot.plugin.minion.tasks.rowfilterpurge; + +import com.google.common.base.Preconditions; +import java.io.File; +import java.util.Collections; +import java.util.Map; +import org.apache.pinot.common.metadata.segment.SegmentZKMetadataCustomMapModifier; +import org.apache.pinot.core.common.MinionConstants; +import org.apache.pinot.core.minion.PinotTaskConfig; +import org.apache.pinot.core.minion.SegmentPurger; +import org.apache.pinot.plugin.minion.tasks.BaseSingleSegmentConversionExecutor; +import org.apache.pinot.plugin.minion.tasks.SegmentConversionResult; +import org.apache.pinot.plugin.minion.tasks.purge.PurgeTaskExecutor; +import org.apache.pinot.segment.local.function.FunctionEvaluator; +import org.apache.pinot.segment.local.function.FunctionEvaluatorFactory; +import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.data.Schema; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +public class RowFilterPurgeTaskExecutor extends BaseSingleSegmentConversionExecutor { + private static final Logger LOGGER = LoggerFactory.getLogger(RowFilterPurgeTaskExecutor.class); + + @Override + protected SegmentConversionResult convert(PinotTaskConfig pinotTaskConfig, File indexDir, File workingDir) + throws Exception { + Map configs = pinotTaskConfig.getConfigs(); + String tableNameWithType = configs.get(MinionConstants.TABLE_NAME_KEY); + String segmentName = configs.get(MinionConstants.SEGMENT_NAME_KEY); + + TableConfig tableConfig = getTableConfig(tableNameWithType); + Schema schema = getSchema(tableNameWithType); + + String filterFunction = configs.get(MinionConstants.RowFilterPurgeTask.FILTER_FUNCTION_KEY); + Preconditions.checkArgument(filterFunction != null, "'%s' is required", + MinionConstants.RowFilterPurgeTask.FILTER_FUNCTION_KEY); + boolean continueOnError = Boolean.parseBoolean( + configs.getOrDefault(MinionConstants.RowFilterPurgeTask.FILTER_CONTINUE_ON_ERROR_KEY, "false")); + + FunctionEvaluator filterEvaluator = FunctionEvaluatorFactory.getExpressionEvaluator(filterFunction); + + LOGGER.info("RowFilterPurgeTask: purging segment {} of table {}", segmentName, tableNameWithType); + _eventObserver.notifyProgress(pinotTaskConfig, + "Purging segment: " + segmentName + " of table: " + tableNameWithType); + + SegmentPurger.RecordPurger recordPurger = row -> { + try { + return Boolean.TRUE.equals(filterEvaluator.evaluate(row)); + } catch (Exception e) { + if (continueOnError) { + LOGGER.debug("Ignoring filter evaluation error for row, treating as retained. Error: {}", e.getMessage()); + return false; + } + throw new RuntimeException( + "Filter evaluation failed. Set filterContinueOnError=true to skip bad rows. Error: " + e.getMessage(), e); + } + }; + + SegmentPurger segmentPurger = + new SegmentPurger(indexDir, workingDir, tableConfig, schema, recordPurger, null, null); + File purgedSegmentFile = segmentPurger.purgeSegment(); + + int numFiltered = segmentPurger.getNumRecordsPurged(); + int numRetained = new SegmentMetadataImpl(indexDir).getTotalDocs() - numFiltered; + + LOGGER.info("RowFilterPurgeTask: finished segment {}, rows retained: {}, rows filtered: {}", + segmentName, numRetained, numFiltered); + + if (purgedSegmentFile == null) { + // Nothing was purged; return original so the executor still stamps ZK. + return new SegmentConversionResult.Builder() + .setFile(indexDir) + .setTableNameWithType(tableNameWithType) + .setSegmentName(segmentName) + .setCustomProperty(MinionConstants.RowFilterPurgeTask.NUM_ROWS_RETAINED_KEY, numRetained) + .setCustomProperty(PurgeTaskExecutor.NUM_RECORDS_PURGED_KEY, numFiltered) + .build(); + } + + return new SegmentConversionResult.Builder() + .setFile(purgedSegmentFile) + .setTableNameWithType(tableNameWithType) + .setSegmentName(segmentName) + .setCustomProperty(MinionConstants.RowFilterPurgeTask.NUM_ROWS_RETAINED_KEY, numRetained) + .setCustomProperty(MinionConstants.RowFilterPurgeTask.NUM_ROWS_FILTERED_KEY, numFiltered) + .build(); + } + + @Override + protected SegmentZKMetadataCustomMapModifier getSegmentZKMetadataCustomMapModifier( + PinotTaskConfig pinotTaskConfig, SegmentConversionResult segmentConversionResult) { + return new SegmentZKMetadataCustomMapModifier( + SegmentZKMetadataCustomMapModifier.ModifyMode.UPDATE, + Collections.singletonMap( + MinionConstants.RowFilterPurgeTask.TASK_TYPE + MinionConstants.TASK_TIME_SUFFIX, + String.valueOf(System.currentTimeMillis()))); + } +} diff --git a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/rowfilterpurge/RowFilterPurgeTaskExecutorFactory.java b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/rowfilterpurge/RowFilterPurgeTaskExecutorFactory.java new file mode 100644 index 0000000000..6eed94f9de --- /dev/null +++ b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/rowfilterpurge/RowFilterPurgeTaskExecutorFactory.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.pinot.plugin.minion.tasks.rowfilterpurge; + +import org.apache.pinot.core.common.MinionConstants; +import org.apache.pinot.minion.MinionConf; +import org.apache.pinot.minion.executor.MinionTaskZkMetadataManager; +import org.apache.pinot.minion.executor.PinotTaskExecutor; +import org.apache.pinot.minion.executor.PinotTaskExecutorFactory; +import org.apache.pinot.spi.annotations.minion.TaskExecutorFactory; + + +@TaskExecutorFactory +public class RowFilterPurgeTaskExecutorFactory implements PinotTaskExecutorFactory { + + @Override + public void init(MinionTaskZkMetadataManager zkMetadataManager) { + } + + @Override + public void init(MinionTaskZkMetadataManager zkMetadataManager, MinionConf minionConf) { + } + + @Override + public String getTaskType() { + return MinionConstants.RowFilterPurgeTask.TASK_TYPE; + } + + @Override + public PinotTaskExecutor create() { + return new RowFilterPurgeTaskExecutor(); + } +} diff --git a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/rowfilterpurge/RowFilterPurgeTaskExecutorV1.java b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/rowfilterpurge/RowFilterPurgeTaskExecutorV1.java new file mode 100644 index 0000000000..2778a06521 --- /dev/null +++ b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/rowfilterpurge/RowFilterPurgeTaskExecutorV1.java @@ -0,0 +1,273 @@ +/** + * 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.pinot.plugin.minion.tasks.rowfilterpurge; + +import java.io.File; +import java.io.IOException; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Set; +import javax.annotation.Nullable; +import org.apache.pinot.common.metadata.segment.SegmentZKMetadataCustomMapModifier; +import org.apache.pinot.core.common.MinionConstants; +import org.apache.pinot.core.minion.PinotTaskConfig; +import org.apache.pinot.core.startree.StarTreeUtils; +import org.apache.pinot.plugin.minion.tasks.BaseSingleSegmentConversionExecutor; +import org.apache.pinot.plugin.minion.tasks.SegmentConversionResult; +import org.apache.pinot.segment.local.function.FunctionEvaluator; +import org.apache.pinot.segment.local.function.FunctionEvaluatorFactory; +import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; +import org.apache.pinot.segment.local.segment.readers.PinotSegmentRecordReader; +import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; +import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl; +import org.apache.pinot.segment.spi.index.startree.StarTreeV2Metadata; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.data.readers.GenericRow; +import org.apache.pinot.spi.data.readers.RecordReader; +import org.apache.pinot.spi.data.readers.RecordReaderConfig; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +public class RowFilterPurgeTaskExecutorV1 extends BaseSingleSegmentConversionExecutor { + private static final Logger LOGGER = LoggerFactory.getLogger(RowFilterPurgeTaskExecutorV1.class); + + @Override + protected SegmentConversionResult convert(PinotTaskConfig pinotTaskConfig, File indexDir, File workingDir) + throws Exception { + Map configs = pinotTaskConfig.getConfigs(); + String tableNameWithType = configs.get(MinionConstants.TABLE_NAME_KEY); + String segmentName = configs.get(MinionConstants.SEGMENT_NAME_KEY); + + TableConfig tableConfig = getTableConfig(tableNameWithType); + Schema schema = getSchema(tableNameWithType); + + String filterFunction = configs.get(MinionConstants.RowFilterPurgeTask.FILTER_FUNCTION_KEY); + boolean continueOnError = Boolean.parseBoolean( + configs.getOrDefault(MinionConstants.RowFilterPurgeTask.FILTER_CONTINUE_ON_ERROR_KEY, "false")); + FunctionEvaluator filterEvaluator = + filterFunction != null ? FunctionEvaluatorFactory.getExpressionEvaluator(filterFunction) : null; + + LOGGER.info("RowFilterPurgeTask: purging segment {} of table {}", segmentName, tableNameWithType); + _eventObserver.notifyProgress(pinotTaskConfig, + "Purging segment: " + segmentName + " of table: " + tableNameWithType); + + SegmentMetadataImpl segmentMetadata = new SegmentMetadataImpl(indexDir); + + int[] counters = new int[2]; // [0] = numRetained, [1] = numFiltered + + try (FilteringRecordReader filteringRecordReader = + new FilteringRecordReader(indexDir, filterEvaluator, continueOnError, counters)) { + + // Pre-pass: scan once to decide whether any rows need to be removed. + // Mirrors SegmentPurger's first-pass check so we avoid rebuilding segments + // when the filter matches nothing. + GenericRow reuse = new GenericRow(); + while (filteringRecordReader.hasNext()) { + filteringRecordReader.next(reuse); + } + + if (counters[1] == 0) { + // Nothing matched the filter; return the original segment so the executor + // still uploads it and stamps the ZK purge-time marker. + LOGGER.info("RowFilterPurgeTask: no rows matched filter for segment {}, skipping rebuild", segmentName); + return new SegmentConversionResult.Builder() + .setFile(indexDir) + .setTableNameWithType(tableNameWithType) + .setSegmentName(segmentName) + .setCustomProperty(MinionConstants.RowFilterPurgeTask.NUM_ROWS_RETAINED_KEY, counters[0]) + .setCustomProperty(MinionConstants.RowFilterPurgeTask.NUM_ROWS_FILTERED_KEY, 0) + .build(); + } + + // Some rows matched; rebuild the segment without them. + filteringRecordReader.rewind(); + + SegmentGeneratorConfig segmentGeneratorConfig = new SegmentGeneratorConfig(tableConfig, schema); + segmentGeneratorConfig.setOutDir(workingDir.getPath()); + segmentGeneratorConfig.setSegmentName(segmentName); + // Preserve original index-creation time for idempotent refreshes. + segmentGeneratorConfig.setCreationTime(String.valueOf(segmentMetadata.getIndexCreationTime())); + if (segmentMetadata.getTimeInterval() != null) { + segmentGeneratorConfig.setTimeColumnName(tableConfig.getValidationConfig().getTimeColumnName()); + segmentGeneratorConfig.setStartTime(Long.toString(segmentMetadata.getStartTime())); + segmentGeneratorConfig.setEndTime(Long.toString(segmentMetadata.getEndTime())); + segmentGeneratorConfig.setSegmentTimeUnit(segmentMetadata.getTimeUnit()); + } + // Preserve the segment's existing star-tree state when dynamic creation is disabled, + // matching SegmentPurger behaviour to avoid dropping star-tree indexes after purge. + if (!tableConfig.getIndexingConfig().isEnableDynamicStarTreeCreation()) { + List starTreeMetadata = segmentMetadata.getStarTreeV2MetadataList(); + segmentGeneratorConfig.setStarTreeIndexConfigs( + starTreeMetadata != null ? StarTreeUtils.toStarTreeIndexConfigs(starTreeMetadata) : null); + segmentGeneratorConfig.setEnableDefaultStarTree(false); + } + + SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl(); + driver.init(segmentGeneratorConfig, filteringRecordReader); + driver.build(); + } + + LOGGER.info("RowFilterPurgeTask: finished segment {}, rows retained: {}, rows filtered: {}", + segmentName, counters[0], counters[1]); + + File outputSegmentDir = new File(workingDir, segmentName); + return new SegmentConversionResult.Builder() + .setFile(outputSegmentDir) + .setTableNameWithType(tableNameWithType) + .setSegmentName(segmentName) + .setCustomProperty(MinionConstants.RowFilterPurgeTask.NUM_ROWS_RETAINED_KEY, counters[0]) + .setCustomProperty(MinionConstants.RowFilterPurgeTask.NUM_ROWS_FILTERED_KEY, counters[1]) + .build(); + } + + @Override + protected SegmentZKMetadataCustomMapModifier getSegmentZKMetadataCustomMapModifier( + PinotTaskConfig pinotTaskConfig, SegmentConversionResult segmentConversionResult) { + return new SegmentZKMetadataCustomMapModifier( + SegmentZKMetadataCustomMapModifier.ModifyMode.UPDATE, + Collections.singletonMap( + MinionConstants.RowFilterPurgeTask.TASK_TYPE + MinionConstants.TASK_TIME_SUFFIX, + String.valueOf(System.currentTimeMillis()))); + } + + /** + * A {@link RecordReader} wrapping {@link PinotSegmentRecordReader} that evaluates a + * {@link FunctionEvaluator} on each row and skips rows where the evaluator returns {@code true}. + * When no evaluator is provided all rows pass through. + * + *

The segment is opened in the constructor so the reader can be iterated for the pre-pass + * before being handed to {@link SegmentIndexCreationDriverImpl}. The driver calls only + * {@link #rewind()} (via {@code RecordReaderSegmentCreationDataSource.getRecordReader()}), + * never {@link #init}, so {@link #init} is a no-op. + * + *

Rows are eagerly pre-fetched in {@link #hasNext()} (matching the pattern used in + * {@link org.apache.pinot.core.minion.SegmentPurger}) so the underlying reader advances + * exactly once per logical record. + */ + private static class FilteringRecordReader implements RecordReader { + + private final PinotSegmentRecordReader _pinotSegmentRecordReader; + @Nullable + private final FunctionEvaluator _filterEvaluator; + private final boolean _continueOnError; + private final int[] _counters; // [0] numRetained, [1] numFiltered + + private GenericRow _nextRow = new GenericRow(); + private boolean _nextRowReturned = true; + private boolean _finished = false; + + FilteringRecordReader(File indexDir, @Nullable FunctionEvaluator filterEvaluator, + boolean continueOnError, int[] counters) + throws Exception { + _pinotSegmentRecordReader = new PinotSegmentRecordReader(); + _pinotSegmentRecordReader.init(indexDir, null, null); + _filterEvaluator = filterEvaluator; + _continueOnError = continueOnError; + _counters = counters; + } + + @Override + public void init(File dataFile, @Nullable Set fieldsToRead, + @Nullable RecordReaderConfig recordReaderConfig) { + // Already initialised in constructor; the driver calls rewind(), not init(). + } + + @Override + public boolean hasNext() { + if (_filterEvaluator == null) { + return _pinotSegmentRecordReader.hasNext(); + } + + if (_finished) { + return false; + } + if (!_nextRowReturned) { + return true; + } + + // Advance until we find a row that passes the filter. + while (_pinotSegmentRecordReader.hasNext()) { + _nextRow.clear(); + _nextRow = _pinotSegmentRecordReader.next(_nextRow); + if (shouldFilter(_nextRow)) { + _counters[1]++; + } else { + _nextRowReturned = false; + return true; + } + } + + _finished = true; + return false; + } + + /** + * Returns {@code true} if the row should be dropped. + * When {@code continueOnError} is set, evaluation errors treat the row as passing (retained) + * rather than failing the entire task. + */ + private boolean shouldFilter(GenericRow row) { + try { + return Boolean.TRUE.equals(_filterEvaluator.evaluate(row)); + } catch (Exception e) { + if (_continueOnError) { + LOGGER.debug("Ignoring filter evaluation error for row, treating as retained. Error: {}", e.getMessage()); + return false; + } + throw new RuntimeException( + "Filter evaluation failed. Set filterContinueOnError=true to skip bad rows. Error: " + e.getMessage(), e); + } + } + + @Override + public GenericRow next() { + return next(new GenericRow()); + } + + @Override + public GenericRow next(GenericRow reuse) { + if (_filterEvaluator == null) { + reuse = _pinotSegmentRecordReader.next(reuse); + } else { + reuse.init(_nextRow); + _nextRowReturned = true; + } + _counters[0]++; + return reuse; + } + + @Override + public void rewind() { + _pinotSegmentRecordReader.rewind(); + _nextRowReturned = true; + _finished = false; + _counters[0] = 0; + _counters[1] = 0; + } + + @Override + public void close() + throws IOException { + _pinotSegmentRecordReader.close(); + } + } +} diff --git a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/rowfilterpurge/RowFilterPurgeTaskGenerator.java b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/rowfilterpurge/RowFilterPurgeTaskGenerator.java new file mode 100644 index 0000000000..d52060be38 --- /dev/null +++ b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/rowfilterpurge/RowFilterPurgeTaskGenerator.java @@ -0,0 +1,210 @@ +/** + * 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.pinot.plugin.minion.tasks.rowfilterpurge; + +import com.google.common.base.Preconditions; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.pinot.common.data.Segment; +import org.apache.pinot.common.metadata.segment.SegmentZKMetadata; +import org.apache.pinot.controller.helix.core.minion.generator.BaseTaskGenerator; +import org.apache.pinot.controller.helix.core.minion.generator.TaskGeneratorUtils; +import org.apache.pinot.core.common.MinionConstants; +import org.apache.pinot.core.minion.PinotTaskConfig; +import org.apache.pinot.spi.annotations.minion.TaskGenerator; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableTaskConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.Schema; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/** + * Task generator for {@link MinionConstants.RowFilterPurgeTask}. + * + *

This task is designed as a one-off purge: segments that have already been purged + * (detected via the {@code RowFilterPurgeTask.time} marker in ZK metadata) are skipped by the + * periodic scheduler path. To re-purge segments (e.g. with a different filter), use the adhoc + * path via {@code POST /tasks/execute}. + * + *

Two trigger paths are supported: + *

    + *
  1. Periodic / manual schedule — {@code POST /tasks/schedule?taskType=RowFilterPurgeTask&tableName=...} + * reads {@code filterFunction} from the table task config and skips already-purged segments.
    2. + *
  2. Adhoc — {@code POST /tasks/execute} with an {@code AdhocTaskConfig} body; the + * {@code filterFunction} is taken from the request payload and ALL segments (including + * already-purged ones) are scheduled, allowing re-purge with a new filter.
    3. + *
+ */ +@TaskGenerator +public class RowFilterPurgeTaskGenerator extends BaseTaskGenerator { + private static final Logger LOGGER = LoggerFactory.getLogger(RowFilterPurgeTaskGenerator.class); + + @Override + public String getTaskType() { + return MinionConstants.RowFilterPurgeTask.TASK_TYPE; + } + + // ── Periodic / manual-schedule path ──────────────────────────────────────── + + @Override + public List generateTasks(List tableConfigs) { + String taskType = MinionConstants.RowFilterPurgeTask.TASK_TYPE; + List pinotTaskConfigs = new ArrayList<>(); + + for (TableConfig tableConfig : tableConfigs) { + String tableName = tableConfig.getTableName(); + + TableTaskConfig tableTaskConfig = tableConfig.getTaskConfig(); + if (tableTaskConfig == null) { + LOGGER.warn("No task config found for table: {}", tableName); + continue; + } + Map taskConfigs = tableTaskConfig.getConfigsForTaskType(taskType); + if (taskConfigs == null) { + LOGGER.warn("No {} task config found for table: {}", taskType, tableName); + continue; + } + if (taskConfigs.get(MinionConstants.RowFilterPurgeTask.FILTER_FUNCTION_KEY) == null) { + LOGGER.warn("Skipping table {} – '{}' is not configured", tableName, + MinionConstants.RowFilterPurgeTask.FILTER_FUNCTION_KEY); + continue; + } + + // Periodic path: skip segments that have already been purged (one-off enforcement). + pinotTaskConfigs.addAll(generateTasksForTable(tableConfig, taskConfigs, true)); + } + + return pinotTaskConfigs; + } + + // ── Adhoc path (POST /tasks/execute) ─────────────────────────────────────── + + /** + * Generates tasks for an adhoc trigger, allowing the caller to supply {@code filterFunction} + * and other configs at request time without modifying the table config. All segments are + * candidates (including already-purged ones) so a new filter can be applied to the full table. + */ + @Override + public List generateTasks(TableConfig tableConfig, Map taskConfigs) + throws Exception { + Preconditions.checkArgument( + taskConfigs != null && taskConfigs.containsKey(MinionConstants.RowFilterPurgeTask.FILTER_FUNCTION_KEY), + "'%s' is required in taskConfigs for adhoc RowFilterPurgeTask", + MinionConstants.RowFilterPurgeTask.FILTER_FUNCTION_KEY); + + // Adhoc path: process all segments regardless of prior-purge marker. + return generateTasksForTable(tableConfig, taskConfigs, false); + } + + // ── Validation ───────────────────────────────────────────────────────────── + + @Override + public void validateTaskConfigs(TableConfig tableConfig, Schema schema, Map taskConfigs) { + Preconditions.checkArgument( + taskConfigs.containsKey(MinionConstants.RowFilterPurgeTask.FILTER_FUNCTION_KEY), + "'%s' must be configured for RowFilterPurgeTask on table: %s", + MinionConstants.RowFilterPurgeTask.FILTER_FUNCTION_KEY, tableConfig.getTableName()); + } + + // ── Shared helper ────────────────────────────────────────────────────────── + + /** + * @param skipPurged if {@code true}, segments that already carry the purge-time marker are + * skipped (periodic one-off enforcement); if {@code false}, all segments are + * scheduled (adhoc re-purge path). + */ + private List generateTasksForTable(TableConfig tableConfig, Map taskConfigs, + boolean skipPurged) { + String taskType = MinionConstants.RowFilterPurgeTask.TASK_TYPE; + String tableName = tableConfig.getTableName(); + + int tableMaxNumTasks; + String tableMaxNumTasksConfig = taskConfigs.get(MinionConstants.TABLE_MAX_NUM_TASKS_KEY); + if (tableMaxNumTasksConfig != null) { + try { + tableMaxNumTasks = Integer.parseInt(tableMaxNumTasksConfig); + } catch (NumberFormatException e) { + tableMaxNumTasks = Integer.MAX_VALUE; + LOGGER.warn("Invalid {} for table: {}, task: {}", MinionConstants.TABLE_MAX_NUM_TASKS_KEY, + tableName, taskType); + } + } else { + tableMaxNumTasks = Integer.MAX_VALUE; + } + + List segmentsZKMetadata = + tableConfig.getTableType() == TableType.OFFLINE + ? getSegmentsZKMetadataForTable(tableName) + : getNonConsumingSegmentsZKMetadataForRealtimeTable(tableName); + + Set runningSegments = + TaskGeneratorUtils.getRunningSegments(taskType, _clusterInfoAccessor); + + List pinotTaskConfigs = new ArrayList<>(); + int numTasksScheduled = 0; + + for (SegmentZKMetadata segmentMetadata : segmentsZKMetadata) { + if (numTasksScheduled >= tableMaxNumTasks) { + break; + } + + String segmentName = segmentMetadata.getSegmentName(); + + if (skipPurged && segmentMetadata.getCustomMap() != null && segmentMetadata.getCustomMap() + .containsKey(taskType + MinionConstants.TASK_TIME_SUFFIX)) { + continue; + } + + if (runningSegments.contains(new Segment(tableName, segmentName))) { + LOGGER.debug("Skipping segment {} – already running", segmentName); + continue; + } + + String downloadUrl = segmentMetadata.getDownloadUrl(); + if (downloadUrl == null || downloadUrl.isEmpty()) { + LOGGER.warn("Skipping segment {} – no download URL available", segmentName); + continue; + } + + Map configs = new HashMap<>(getBaseTaskConfigs(tableConfig, List.of(segmentName))); + configs.put(MinionConstants.DOWNLOAD_URL_KEY, downloadUrl); + configs.put(MinionConstants.UPLOAD_URL_KEY, _clusterInfoAccessor.getVipUrl() + "/segments"); + configs.put(MinionConstants.ORIGINAL_SEGMENT_CRC_KEY, String.valueOf(segmentMetadata.getCrc())); + + String filterFunction = taskConfigs.get(MinionConstants.RowFilterPurgeTask.FILTER_FUNCTION_KEY); + configs.put(MinionConstants.RowFilterPurgeTask.FILTER_FUNCTION_KEY, filterFunction); + + String continueOnError = taskConfigs.get(MinionConstants.RowFilterPurgeTask.FILTER_CONTINUE_ON_ERROR_KEY); + if (continueOnError != null) { + configs.put(MinionConstants.RowFilterPurgeTask.FILTER_CONTINUE_ON_ERROR_KEY, continueOnError); + } + + pinotTaskConfigs.add(new PinotTaskConfig(taskType, configs)); + numTasksScheduled++; + } + + LOGGER.info("Scheduled {} {} tasks for table: {}", numTasksScheduled, taskType, tableName); + return pinotTaskConfigs; + } +}