From f17a66a737e958c5506e9a4740a9d220c9b907da Mon Sep 17 00:00:00 2001
From: Michael Froh <froh@amazon.com>
Date: Mon, 18 Dec 2023 16:14:39 +0000
Subject: [PATCH] Add support for reverse nested agg too

Signed-off-by: Michael Froh <froh@amazon.com>
---
 .../test/search.aggregation/230_composite.yml | 56 +++++++++++++++++++
 .../CompositeAggregationBuilder.java          |  5 +-
 2 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/rest-api-spec/src/main/resources/rest-api-spec/test/search.aggregation/230_composite.yml b/rest-api-spec/src/main/resources/rest-api-spec/test/search.aggregation/230_composite.yml
index a3c1a961fe699..599a93117e669 100644
--- a/rest-api-spec/src/main/resources/rest-api-spec/test/search.aggregation/230_composite.yml
+++ b/rest-api-spec/src/main/resources/rest-api-spec/test/search.aggregation/230_composite.yml
@@ -152,6 +152,8 @@ setup:
   - match: { aggregations.test.buckets.4.key.kw: "bar" }
   - match: { aggregations.test.buckets.4.doc_count: 1 }
 
+
+
 ---
 "Aggregate After":
   - do:
@@ -576,6 +578,60 @@ setup:
   - match: { aggregations.1.2.3.buckets.0.key.nested: 20 }
   - match: { aggregations.1.2.3.buckets.0.doc_count:  2 }
 
+---
+"Composite aggregation with filtered reverse nested parent":
+  - do:
+      search:
+        rest_total_hits_as_int: true
+        index: test
+        body:
+          aggregations:
+            1:
+              nested:
+                path: nested
+              aggs:
+                2:
+                  filter:
+                    range:
+                      nested.nested_long:
+                        gt: 0
+                        lt: 20
+                  aggs:
+                    3:
+                      reverse_nested: {}
+                      aggs:
+                        4:
+                          composite:
+                            sources: [
+                              {
+                                "long": {
+                                  "terms": {
+                                    "field": "long"
+                                  }
+                                }
+                              },
+                              {
+                                "kw": {
+                                  "terms": {
+                                    "field": "keyword"
+                                  }
+                                }
+                              }
+                            ]
+  - match: {hits.total: 6}
+  - length: { aggregations.1.2.3.4.buckets: 4 }
+  - match: { aggregations.1.2.3.4.buckets.0.key.long: 0 }
+  - match: { aggregations.1.2.3.4.buckets.0.key.kw: "bar" }
+  - match: { aggregations.1.2.3.4.buckets.0.doc_count: 1 }
+  - match: { aggregations.1.2.3.4.buckets.1.key.long: 10 }
+  - match: { aggregations.1.2.3.4.buckets.1.key.kw: "foo" }
+  - match: { aggregations.1.2.3.4.buckets.1.doc_count: 1 }
+  - match: { aggregations.1.2.3.4.buckets.2.key.long: 20 }
+  - match: { aggregations.1.2.3.4.buckets.2.key.kw: "foo" }
+  - match: { aggregations.1.2.3.4.buckets.2.doc_count: 1 }
+  - match: { aggregations.1.2.3.4.buckets.3.key.long: 100 }
+  - match: { aggregations.1.2.3.4.buckets.3.key.kw: "bar" }
+  - match: { aggregations.1.2.3.4.buckets.3.doc_count: 1 }
 
 ---
 "Composite aggregation with unmapped field":
diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeAggregationBuilder.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeAggregationBuilder.java
index a94a41a736ba0..b1ff0050d551a 100644
--- a/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeAggregationBuilder.java
+++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeAggregationBuilder.java
@@ -46,6 +46,7 @@
 import org.opensearch.search.aggregations.AggregatorFactory;
 import org.opensearch.search.aggregations.bucket.filter.FilterAggregatorFactory;
 import org.opensearch.search.aggregations.bucket.nested.NestedAggregatorFactory;
+import org.opensearch.search.aggregations.bucket.nested.ReverseNestedAggregatorFactory;
 import org.opensearch.search.aggregations.support.ValuesSourceRegistry;
 
 import java.io.IOException;
@@ -244,7 +245,9 @@ public BucketCardinality bucketCardinality() {
     private static AggregatorFactory checkParentIsSafe(AggregatorFactory factory) {
         if (factory == null) {
             return null;
-        } else if (factory instanceof NestedAggregatorFactory || factory instanceof FilterAggregatorFactory) {
+        } else if (factory instanceof NestedAggregatorFactory ||
+            factory instanceof FilterAggregatorFactory ||
+            factory instanceof ReverseNestedAggregatorFactory) {
             return checkParentIsSafe(factory.getParent());
         } else {
             return factory;