Streaming Aggregation

[bowenlan-amzn]

Description

RFC: #16774

Change based on #18424 Streaming Transport.

High Level Idea

Instead of return one response per shard search, we can return one partial aggregation response per segment and use streaming transport to send it back immediately.
For the existing shard search response, we can still put in all the other results apart from aggregation and send that back as the last response of the stream and complete it.

This way moves the reduce logic and aggregation result memory consumption from data node to coordinator node, making coordinator node the single place to scale up to handle high cardinality aggregation scenarios.

Run the Code

Please see SubAggregationIT under org.opensearch.streaming.aggregation

Add these to JUnit run configuration in Intellij

--add-opens=java.base/java.nio=org.apache.arrow.memory.core,ALL-UNNAMED
-Dio.netty.allocator.numDirectArenas=1
-Dio.netty.noUnsafe=false
-Dio.netty.tryUnsafe=true
-Dio.netty.tryReflectionSetAccessible=true

Change List

Entry

• New experimental stream query parameter when calling search rest endpoint. User needs to opt in by set stream=true to use the streaming aggregation feature
• Route to StreamSearchAction for doing stream transport search action.

Transport Layer

• StreamTransportSearchAction extends TransportSearchAction and override async action provider to provide a stream async action.
• StreamSearchTransportService to register request handler and response handler.
• StreamSearchChannelListener to send results back from data node.
• At coordinator node, handle StreamTransportResponse which provides stream of results from data node.

Coordinator

• Extends current search async action with StreamSearchQueryThenFetchAsyncAction, provide a streaming action listener StreamSearchActionListener
• Stream responses are treated as 2 types, the last response is the shard result which still goes through the existing code path to be consumed. All the stream results before the last one go through a new consumption path.
• All stream responses are consumed by StreamQueryPhaseResultConsumer which extends the existing QueryPhaseResultConsumer. It adds a new consumeStreamResult to push stream results into the same buffer pendingMerges.
• Stream results and shard result have 2 different result consumption callbacks. Since the callbacks are triggered from same synchronized queue, pendingMerges. We can sync these 2 callbacks in StreamSearchQueryThenFetchAsyncAction and trigger the onPhaseDone without race condition.

Data Node

• StreamingStringTermsAggregator to do the per segment aggregation and able to reset after building aggregation. Comparing to current default GlobalOrdinalsStringTermsAggregator, the key changes are
  • we don't need global ordinal but just per segment doc values ordinal
  • we don't need to sort and reduce buckets when building aggregation, but just directly build aggregation from all buckets collected.
• BucketCollectorProcessor build aggregation batch
• ContextIndexSearcher trigger the aggregation batch build, reset, and send back the batch result using StreamSearchChannelListener

Coverage

• Code cov Streaming Aggregation #18874 (comment)

Things Left

• Coordinator merge reduce bounded by a new reduce_size parameter. And research what's the right reduce_size to keep similar accuracy as the current terms aggregation

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.

[github-actions]

❌ Gradle check result for cfce243: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[github-actions]

❌ Gradle check result for 064e610: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[github-actions]

❌ Gradle check result for d936286: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[github-actions]

❌ Gradle check result for bc86e80: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[github-actions]

❌ Gradle check result for 31a6479: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[github-actions]

❌ Gradle check result for 3215bc1: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[github-actions]

❌ Gradle check result for e79f0fc: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[github-actions]

❌ Gradle check result for 11c8299: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[github-actions]

❌ Gradle check result for ef194ae: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[github-actions]

❌ Gradle check result for 44f645f: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[github-actions]

❌ Gradle check result for 2496fb0: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[github-actions]

❌ Gradle check result for 4875792: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[rishabhmaurya]

good work isolating streaming aggs flow with regular search flow @bowenlan-amzn @harshavamsi.
It looks pretty safe change to me, I will wait for some time for other to validate before I merge them in.

[github-actions]

❌ Gradle check result for 49ebf96: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[github-actions]

❌ Gradle check result for 49ebf96: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[github-actions]

✅ Gradle check result for 49ebf96: SUCCESS