Refactor streaming agg query phase planning by bowenlan-amzn · Pull Request #20471 · opensearch-project/OpenSearch

bowenlan-amzn · 2026-01-23T22:24:22Z

Description

Streaming aggregation feature supports an agg defined with

Top level should be string terms agg and numeric terms agg
Second level/sub aggregation could be numeric terms, cardinality, max, min, sum (we should add avg and value_count soon)
If multiple aggs fit above constraints, it should also be good.

Segment streaming topN heuristic set to shard_size
Also user has a index level setting to provide the segment topN size.
The final topN is the larger one of these 2

The topN is compared with search.aggregations.streaming.max_estimated_bucket_count to decide whether streaming should be used or not. If exceed, streaming is considered to stream back too many buckets per segment so we will fall back to default aggregation path which send response per shard.

Related Issues

Resolves #[Issue number to be closed when this PR is merged]

Check List

Functionality includes testing.
API changes companion pull request created, if applicable.
Public documentation issue/PR created, if applicable.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

coderabbitai · 2026-01-23T22:24:35Z

📝 Walkthrough

Walkthrough

This PR refactors the streaming aggregation cost estimation architecture from a runtime collector-tree evaluation model to an early-stage aggregation-factory estimation model, introducing a new StreamingCostEstimable interface while removing the Streamable interface and eliminating the AggregatorTreeEvaluator class that previously managed streaming decisions.

Changes

Cohort / File(s)	Summary
Streaming Architecture Refactoring `CHANGELOG.md`, `server/src/main/java/org/opensearch/search/streaming/Streamable.java`, `server/src/main/java/org/opensearch/search/streaming/StreamingCostEstimable.java`, `server/src/main/java/org/opensearch/search/aggregations/AggregatorTreeEvaluator.java`, `server/src/test/java/org/opensearch/search/aggregations/AggregatorTreeEvaluatorTests.java`	Removes `Streamable` interface and replaces with new `StreamingCostEstimable` interface; eliminates `AggregatorTreeEvaluator` that previously evaluated streaming decisions at runtime; shifts cost estimation to factory-level before aggregator creation.
FlushModeResolver Refactoring `server/src/main/java/org/opensearch/search/streaming/FlushModeResolver.java`	Major refactoring from collector-tree analysis to aggregation-builder focused approach; replaces `resolve(Collector...)` with `isEligibleForStreaming(AggregatorFactories.Builder)` and adds new settings for streaming thresholds; introduces per-aggregation type streaming validation.
StreamingCostMetrics Simplification `server/src/main/java/org/opensearch/search/streaming/StreamingCostMetrics.java`, `server/src/test/java/org/opensearch/search/streaming/StreamingCostMetricsTests.java`, `server/src/test/java/org/opensearch/search/streaming/FlushModeResolverTests.java`	Reduces record fields from five to two (`streamable`, `topNSize`); adds `neutral()` factory method; updates combination logic to work only with topNSize; refactors tests from integration-style to isolated decision logic.
Terms Aggregator Factory Streaming `server/src/main/java/org/opensearch/search/aggregations/bucket/terms/TermsAggregatorFactory.java`	Implements `StreamingCostEstimable`; adds `computeSegTopN()` helper; requires `FlushMode.PER_SEGMENT` explicitly for streaming (previously allowed null); propagates segmentTopN parameter through streaming aggregator creation paths.
Terms Streaming Aggregators `server/src/main/java/org/opensearch/search/aggregations/bucket/terms/StreamNumericTermsAggregator.java`, `server/src/main/java/org/opensearch/search/aggregations/bucket/terms/StreamStringTermsAggregator.java`	Removes `Streamable` implementation; introduces new `segmentTopN` parameter; replaces internal segmentSize logic with segmentTopN for top-bucket selection; removes streaming cost metrics methods.
Cardinality Aggregator Streaming `server/src/main/java/org/opensearch/search/aggregations/metrics/CardinalityAggregatorFactory.java`, `server/src/main/java/org/opensearch/search/aggregations/metrics/StreamCardinalityAggregator.java`	Factory now implements `StreamingCostEstimable` with ordinals-aware estimation; aggregator removes `Streamable` and adds explicit lifecycle methods (getLeafCollector, doReset, doPostCollection, doClose).
Metric Aggregator Factories (Streamable Removal & StreamingCostEstimable Addition) `server/src/main/java/org/opensearch/search/aggregations/metrics/{AvgAggregatorFactory,MaxAggregatorFactory,MinAggregatorFactory,SumAggregatorFactory,ValueCountAggregatorFactory}.java`	Each now implements `StreamingCostEstimable` returning neutral cost metrics; removes responsibility for runtime streaming cost estimation.
Metric Aggregators (Streamable Removal) `server/src/main/java/org/opensearch/search/aggregations/metrics/{MaxAggregator,MinAggregator}.java`	Removes `Streamable` interface implementation and `getStreamingCostMetrics()` method from both classes.
Aggregator Factories & Collection Management `server/src/main/java/org/opensearch/search/aggregations/AggregatorFactories.java`, `server/src/main/java/org/opensearch/search/aggregations/AggregationCollectorManager.java`	`AggregatorFactories` adds streaming cost estimation workflow for top-level aggregators using new factory-level cost analysis; `AggregationCollectorManager` removes dynamic evaluation/recreation step in `createCollector`.
Profiling & Search Integration `server/src/main/java/org/opensearch/search/profile/aggregation/ProfilingAggregator.java`, `server/src/main/java/org/opensearch/rest/action/search/RestSearchAction.java`	Removes `Streamable` implementation from `ProfilingAggregator`; simplifies `canUseStreamSearch` to delegate to `FlushModeResolver.isEligibleForStreaming()` instead of checking specific aggregation types.
Test Coverage Updates `server/src/test/java/org/opensearch/rest/action/search/RestSearchActionTests.java`, `server/src/test/java/org/opensearch/search/aggregations/FactoryStreamingCostEstimationTests.java`, `server/src/test/java/org/opensearch/search/aggregations/bucket/terms/StreamNumericTermsAggregatorTests.java`, `server/src/test/java/org/opensearch/search/aggregations/bucket/terms/StreamStringTermsAggregatorTests.java`, `server/src/test/java/org/opensearch/search/aggregations/metrics/{MaxAggregatorTests,MinAggregatorTests,StreamCardinalityAggregatorTests}.java`	Removes old streaming cost metrics tests; adds comprehensive new factory-level streaming cost estimation tests; updates integration tests to use profile-based validation instead of assertions on Streamable interface.
Integration Tests & Test Framework `plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/streaming/aggregation/SubAggregationIT.java`, `test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java`	Integration tests refactored to use profile inspection for streaming validation with dynamic settings adjustments; test framework explicitly sets `FlushMode.PER_SEGMENT` in streaming aggregator creation paths.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related PRs

TopN selection for streaming terms aggregations #20481: Modifies StreamStringTermsAggregator and StreamNumericTermsAggregator to add segment-level TopN logic with segmentTopN parameter, directly implementing the factory changes introduced in this PR.
HLL field Mapper for Cardinality Rollups #20129: Modifies CardinalityAggregatorFactory streaming and flush-mode behavior including aggregator lifecycle changes, aligning with cardinality aggregation refactoring in this PR.
Add Hybrid Cardinality collector to prioritize Ordinals Collector #19524: Refactors cardinality aggregation internals and adds hybrid collector paths to CardinalityAggregator, touching the same CardinalityAggregatorFactory and streaming infrastructure modified here.

Suggested labels

enhancement, Search, Search:Aggregations, Search:Performance

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 22.12% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'Refactor streaming agg query phase planning' directly and concisely describes the main change—a refactoring of streaming aggregation query phase planning logic.
Description check	✅ Passed	The pull request description includes all required sections: a clear description of the feature (streaming aggregation constraints), related issues placeholder, and a completed checklist with testing and documentation confirmations.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

📝 Generate docstrings

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

github-actions · 2026-01-23T22:27:14Z

❌ Gradle check result for bb004fb: ABORTED