Adding support for multi-tier rollups in ISM

Signed-off-by: Kshitij Tandon <[email protected]>

Author: tandonks

src/main/kotlin/org/opensearch/indexmanagement/IndexManagementPlugin.kt

@@ -427,7 +427,7 @@ class IndexManagementPlugin :
                 .registerMapperService(RollupMapperService(client, clusterService, indexNameExpressionResolver))
                 .registerIndexer(RollupIndexer(settings, clusterService, client))
                 .registerSearcher(RollupSearchService(settings, clusterService, client))
-                .registerMetadataServices(RollupMetadataService(client, xContentRegistry))
+                .registerMetadataServices(RollupMetadataService(client, xContentRegistry, clusterService))
                 .registerConsumers()
                 .registerClusterConfigurationProvider(skipFlag)
         indexManagementIndices = IndexManagementIndices(settings, client.admin().indices(), clusterService)

src/main/kotlin/org/opensearch/indexmanagement/rollup/RollupMapperService.kt

@@ -293,6 +293,13 @@ class RollupMapperService(
 
             val indexMappingSource = indexTypeMappings.sourceAsMap
 
+            // TODO: Check if this is required or we want to fail in case of rollup on rolled up empty indices
+//            // If source is a rollup index with no properties (no data rolled up yet), skip field validation
+//            if (isSourceRollupIndex && !indexMappingSource.containsKey("properties")) {
+//                logger.info("Source rollup index [$index] has no properties yet, skipping field validation")
+//                return RollupJobValidationResult.Valid
+//            }
+
             val issues = mutableSetOf<String>()
             // Validate source fields in dimensions
             rollup.dimensions.forEach { dimension ->

src/main/kotlin/org/opensearch/indexmanagement/rollup/RollupMetadataService.kt

@@ -36,6 +36,7 @@ import org.opensearch.indexmanagement.rollup.model.Rollup
 import org.opensearch.indexmanagement.rollup.model.RollupMetadata
 import org.opensearch.indexmanagement.rollup.model.RollupStats
 import org.opensearch.indexmanagement.rollup.util.DATE_FIELD_STRICT_DATE_OPTIONAL_TIME_FORMAT
+import org.opensearch.indexmanagement.rollup.util.isRollupIndex
 import org.opensearch.indexmanagement.util.NO_ID
 import org.opensearch.search.aggregations.bucket.composite.InternalComposite
 import org.opensearch.search.aggregations.bucket.histogram.DateHistogramAggregationBuilder
@@ -48,7 +49,11 @@ import java.time.Instant
 // TODO: Wrap client calls in retry for transient failures
 // Service that handles CRUD operations for rollup metadata
 @Suppress("TooManyFunctions")
-class RollupMetadataService(val client: Client, val xContentRegistry: NamedXContentRegistry) {
+class RollupMetadataService(
+    val client: Client,
+    val xContentRegistry: NamedXContentRegistry,
+    val clusterService: org.opensearch.cluster.service.ClusterService,
+) {
     private val logger = LogManager.getLogger(javaClass)
 
     // If the job does not have a metadataID then we need to initialize the first metadata
@@ -178,6 +183,12 @@ class RollupMetadataService(val client: Client, val xContentRegistry: NamedXCont
     @Throws(Exception::class)
     private suspend fun getInitialStartTime(rollup: Rollup): StartingTimeResult {
         try {
+            // Check if source is a rollup index and use appropriate method
+            val isSourceRollupIndex = isRollupIndex(rollup.sourceIndex, clusterService.state())
+            if (isSourceRollupIndex) {
+                // Use min aggregation for rollup indices (RollupInterceptor blocks size > 0)
+                return getEarliestTimestampFromRollupIndex(rollup)
+            }
             // Rollup requires the first dimension to be the date histogram
             val dateHistogram = rollup.dimensions.first() as DateHistogram
             val searchSourceBuilder =
@@ -219,6 +230,69 @@ class RollupMetadataService(val client: Client, val xContentRegistry: NamedXCont
         }
     }
 
+    /**
+     * Get the earliest timestamp from a rollup index by finding the minimum value of the date histogram field.
+     * This is used to determine the starting point for continuous rollups on rollup indices.
+     * Uses sort instead of aggregation to avoid rollup interceptor validation.
+     */
+    @Suppress("ReturnCount")
+    @Throws(Exception::class)
+    private suspend fun getEarliestTimestampFromRollupIndex(rollup: Rollup): StartingTimeResult {
+        try {
+            val dateHistogram = rollup.dimensions.first() as DateHistogram
+            val dateField = dateHistogram.sourceField
+
+            logger.info("Idhr se ja rha hu mai tumhe kya 2")
+
+            val searchRequest = SearchRequest(rollup.sourceIndex)
+                .source(
+                    SearchSourceBuilder()
+                        .size(1)
+                        .query(MatchAllQueryBuilder())
+                        .sort("$dateField.date_histogram", SortOrder.ASC)
+                        .trackTotalHits(false)
+                        .fetchSource(false)
+                        .docValueField("$dateField.date_histogram", DATE_FIELD_STRICT_DATE_OPTIONAL_TIME_FORMAT),
+                )
+                .allowPartialSearchResults(false)
+
+            // Set BYPASS_SIZE_CHECK to allow size=1 when querying rollup index to get earliest timestamp
+            // This is needed for continuous rollup initialization on rollup indices (multi-tier rollup)
+            org.opensearch.indexmanagement.rollup.interceptor.RollupInterceptor.setBypass(
+                org.opensearch.indexmanagement.rollup.interceptor.RollupInterceptor.BYPASS_SIZE_CHECK,
+            )
+            try {
+                val response: SearchResponse = client.suspendUntil { search(searchRequest, it) }
+
+                if (response.hits.hits.isEmpty()) {
+                    return StartingTimeResult.NoDocumentsFound
+                }
+
+                logger.info("Idhr se ja rha hu mai tumhe kya")
+
+                // In rollup indices, date histogram fields are named as "field.date_histogram"
+                val rollupDateField = "$dateField.date_histogram"
+                val firstHitTimestampAsString: String =
+                    response.hits.hits.first().field(rollupDateField).getValue<String>()
+                        ?: return StartingTimeResult.NoDocumentsFound
+
+                val formatter = DateFormatter.forPattern(DATE_FIELD_STRICT_DATE_OPTIONAL_TIME_FORMAT)
+                val epochMillis = DateFormatters.from(formatter.parse(firstHitTimestampAsString), formatter.locale()).toInstant().toEpochMilli()
+                return StartingTimeResult.Success(getRoundedTime(epochMillis, dateHistogram))
+            } finally {
+                org.opensearch.indexmanagement.rollup.interceptor.RollupInterceptor.clearBypass()
+            }
+        } catch (e: RemoteTransportException) {
+            val unwrappedException = ExceptionsHelper.unwrapCause(e) as Exception
+            logger.error("Error when getting earliest timestamp from rollup index for rollup [{}]: {}", rollup.id, unwrappedException)
+            return StartingTimeResult.Failure(unwrappedException)
+        } catch (e: Exception) {
+            // TODO: Catching general exceptions for now, can make more granular
+            logger.error("Error when getting earliest timestamp from rollup index for rollup [{}]: {}", rollup.id, e)
+            return StartingTimeResult.Failure(e)
+        }
+    }
+
     /**
      * Return time rounded down to the nearest unit of time the interval is based on.
      * This should map to the equivalent bucket a document with the given timestamp would fall into for the date histogram.

src/main/kotlin/org/opensearch/indexmanagement/rollup/RollupRunner.kt

@@ -35,6 +35,9 @@ import org.opensearch.indexmanagement.rollup.model.RollupStats
 import org.opensearch.indexmanagement.rollup.model.incrementStats
 import org.opensearch.indexmanagement.rollup.model.mergeStats
 import org.opensearch.indexmanagement.rollup.settings.RollupSettings
+import org.opensearch.indexmanagement.rollup.util.getDateHistogram
+import org.opensearch.indexmanagement.rollup.util.getRollupJobs
+import org.opensearch.indexmanagement.rollup.util.isRollupIndex
 import org.opensearch.indexmanagement.util.acquireLockForScheduledJob
 import org.opensearch.indexmanagement.util.releaseLockForScheduledJob
 import org.opensearch.indexmanagement.util.renewLockForScheduledJob
@@ -44,6 +47,8 @@ import org.opensearch.jobscheduler.spi.ScheduledJobParameter
 import org.opensearch.jobscheduler.spi.ScheduledJobRunner
 import org.opensearch.script.ScriptService
 import org.opensearch.search.aggregations.bucket.composite.InternalComposite
+import org.opensearch.search.aggregations.bucket.histogram.DateHistogramAggregationBuilder
+import org.opensearch.search.aggregations.bucket.histogram.DateHistogramInterval
 import org.opensearch.threadpool.ThreadPool
 import org.opensearch.transport.client.Client
 
@@ -266,7 +271,16 @@ object RollupRunner :
                             withClosableContext(
                                 IndexManagementSecurityContext(job.id, settings, threadPool.threadContext, job.user),
                             ) {
-                                rollupSearchService.executeCompositeSearch(updatableJob, metadata)
+                                // Need to set this bypass as we are already doing the required aggregation re-writing in the composite aggregation
+                                // hence do not need to go through the interceptor for aggregation rewriting
+                                org.opensearch.indexmanagement.rollup.interceptor.RollupInterceptor.setBypass(
+                                    org.opensearch.indexmanagement.rollup.interceptor.RollupInterceptor.BYPASS_ROLLUP_SEARCH,
+                                )
+                                try {
+                                    rollupSearchService.executeCompositeSearch(updatableJob, metadata, clusterService)
+                                } finally {
+                                    org.opensearch.indexmanagement.rollup.interceptor.RollupInterceptor.clearBypass()
+                                }
                             }
                         val rollupResult =
                             when (rollupSearchResult) {
@@ -412,6 +426,15 @@ object RollupRunner :
                 else -> return@withClosableContext sourceIndexValidationResult
             }
 
+            // Additional validation for rollup-on-rollup scenarios
+            if (isRollupIndex(job.sourceIndex, clusterService.state())) {
+                when (val rollupValidationResult = validateRollupOnRollup(job)) {
+                    is RollupJobValidationResult.Valid -> {
+                    } // No action taken when valid
+                    else -> return@withClosableContext rollupValidationResult
+                }
+            }
+
             // we validate target index only if there is metadata document in the rollup
             if (metadata != null) {
                 logger.debug("Attempting to create/validate target index [${job.targetIndex}] for rollup job [${job.id}]")
@@ -470,6 +493,74 @@ object RollupRunner :
             }
         }
     }
+
+    @Suppress("ReturnCount")
+    private fun validateRollupOnRollup(job: Rollup): RollupJobValidationResult {
+        val sourceRollupJobs = clusterService.state().metadata.index(job.sourceIndex).getRollupJobs()
+        if (sourceRollupJobs == null) {
+            return RollupJobValidationResult.Invalid("Source rollup index has no rollup jobs")
+        }
+
+        val sourceJob = sourceRollupJobs.first()
+        val targetDateHistogram = job.getDateHistogram()
+        val sourceDateHistogram = sourceJob.getDateHistogram()
+
+        // Validate interval alignment
+        val sourceInterval = sourceDateHistogram.fixedInterval ?: sourceDateHistogram.calendarInterval!!
+        val targetInterval = targetDateHistogram.fixedInterval ?: targetDateHistogram.calendarInterval!!
+        val intervalValid = validateIntervalAlignment(sourceInterval, targetInterval)
+        if (!intervalValid) {
+            return RollupJobValidationResult.Invalid(
+                "Target interval [$targetInterval] must be an exact multiple of source interval [$sourceInterval]",
+            )
+        }
+
+        // Validate source field compatibility
+        val sourceDimensionFields = sourceJob.dimensions.map { it.sourceField }.toSet()
+        val sourceMetricFields = sourceJob.metrics.map { it.sourceField }.toSet()
+        val targetDimensionFields = job.dimensions.map { it.sourceField }.toSet()
+        val targetMetricFields = job.metrics.map { it.sourceField }.toSet()
+
+        val invalidDimensionFields = targetDimensionFields - sourceDimensionFields
+        val invalidMetricFields = targetMetricFields - sourceMetricFields
+
+        return when {
+            invalidDimensionFields.isNotEmpty() -> RollupJobValidationResult.Invalid(
+                "Cannot rollup on dimension fields $invalidDimensionFields that don't exist in source rollup",
+            )
+            invalidMetricFields.isNotEmpty() -> RollupJobValidationResult.Invalid(
+                "Cannot rollup on metric fields $invalidMetricFields that don't exist in source rollup",
+            )
+            else -> {
+                // Validate metric compatibility
+                val sourceMetrics = sourceJob.metrics.flatMap { it.metrics }.map { it.type.type }.toSet()
+                val targetMetrics = job.metrics.flatMap { it.metrics }.map { it.type.type }.toSet()
+                val unsupportedMetrics = targetMetrics - sourceMetrics
+
+                if (unsupportedMetrics.isNotEmpty()) {
+                    RollupJobValidationResult.Invalid(
+                        "Target rollup requests metrics $unsupportedMetrics that are not available in source rollup",
+                    )
+                } else {
+                    RollupJobValidationResult.Valid
+                }
+            }
+        }
+    }
+
+    private fun validateIntervalAlignment(sourceInterval: String, targetInterval: String): Boolean = try {
+        val sourceMillis = parseIntervalToMillis(sourceInterval)
+        val targetMillis = parseIntervalToMillis(targetInterval)
+        targetMillis % sourceMillis == 0L && targetMillis > sourceMillis
+    } catch (e: Exception) {
+        true // Let it through and fail later with better error
+    }
+
+    private fun parseIntervalToMillis(interval: String): Long = if (DateHistogramAggregationBuilder.DATE_FIELD_UNITS.containsKey(interval)) {
+        DateHistogramInterval(interval).estimateMillis()
+    } else {
+        TimeValue.parseTimeValue(interval, "parseIntervalToMillis").millis
+    }
 }
 
 sealed class RollupJobResult {

src/main/kotlin/org/opensearch/indexmanagement/rollup/RollupSearchService.kt

@@ -99,7 +99,7 @@ class RollupSearchService(
     }
 
     @Suppress("ComplexMethod")
-    suspend fun executeCompositeSearch(job: Rollup, metadata: RollupMetadata): RollupSearchResult = try {
+    suspend fun executeCompositeSearch(job: Rollup, metadata: RollupMetadata, clusterService: ClusterService): RollupSearchResult = try {
         var retryCount = 0
         RollupSearchResult.Success(
             retrySearchPolicy.retry(logger) {
@@ -113,7 +113,7 @@ class RollupSearchService(
                         )
                     }
 
-                    val searchRequest = job.copy(pageSize = pageSize).getRollupSearchRequest(metadata)
+                    val searchRequest = job.copy(pageSize = pageSize).getRollupSearchRequest(metadata, clusterService.state())
                     val cancelTimeoutTimeValue = TimeValue.timeValueMinutes(getCancelAfterTimeInterval(cancelAfterTimeInterval.minutes))
                     searchRequest.cancelAfterTimeInterval = cancelTimeoutTimeValue
 

src/main/kotlin/org/opensearch/indexmanagement/rollup/interceptor/RollupInterceptor.kt

@@ -65,6 +65,49 @@ class RollupInterceptor(
 
     @Volatile private var searchRawRollupIndices = RollupSettings.ROLLUP_SEARCH_SOURCE_INDICES.get(settings)
 
+    companion object {
+        /**
+         * Thread-local bypass mechanism for internal operations that need to query rollup indices
+         * without triggering interceptor validations.
+         */
+        private val bypassInterceptor = ThreadLocal<Int>()
+
+        /**
+         * Bypass level that skips all rollup search validations and rewriting.
+         * Used when the system needs to query rollup indices directly using composite aggregation
+         */
+        const val BYPASS_ROLLUP_SEARCH = 1
+
+        /**
+         * Bypass level that allows non-zero size in rollup searches.
+         * Used for internal operations like continuous rollup initialization that need to fetch
+         * actual documents from rollup indices (e.g., getEarliestTimestampFromRollupIndex).
+         */
+        const val BYPASS_SIZE_CHECK = 2
+
+        /**
+         * Sets the bypass level for the current thread.
+         * Must be followed by clearBypass() to avoid leaking the bypass state.
+         */
+        fun setBypass(bypassLevel: Int) {
+            bypassInterceptor.set(bypassLevel)
+        }
+
+        /**
+         * Clears the bypass level for the current thread.
+         * Should always be called in a finally block after setBypass().
+         */
+        fun clearBypass() {
+            bypassInterceptor.remove()
+        }
+
+        /**
+         * Gets the current bypass level for the thread.
+         * Returns 0 if no bypass is set.
+         */
+        fun getBypassLevel(): Int = bypassInterceptor.get() ?: 0
+    }
+
     init {
         clusterService.clusterSettings.addSettingsUpdateConsumer(RollupSettings.ROLLUP_SEARCH_ENABLED) {
             searchEnabled = it
@@ -89,7 +132,20 @@ class RollupInterceptor(
                 val index = request.shardId().indexName
                 val isRollupIndex = isRollupIndex(index, clusterService.state())
                 if (isRollupIndex) {
-                    if (request.source().size() != 0) {
+                    val bypassLevel = getBypassLevel()
+                    // BYPASS_ROLLUP_SEARCH: Skip all validations and query rewriting
+                    // Used for composite aggregation query which we do for retrieving the buckets
+                    // to write to target index during rollup of rolled up index
+                    if (bypassLevel == BYPASS_ROLLUP_SEARCH) {
+                        actualHandler.messageReceived(request, channel, task)
+                        return
+                    }
+
+                    // BYPASS_SIZE_CHECK: Allow non-zero size for internal operations that need to
+                    // fetch documents (e.g., fetching earliest timestamp document for rolled index
+                    // in case of continuous rollup job). Normal rollup searches must have size=0
+                    // since they should only return aggregations
+                    if (bypassLevel != BYPASS_SIZE_CHECK && request.source().size() != 0) {
                         throw IllegalArgumentException("Rollup search must have size explicitly set to 0, but found ${request.source().size()}")
                     }