Updated the Composite and Cardinality aggregators (#20173)

* Updated the Composite and Cardinality aggregators

Signed-off-by: Vinay Krishna Pudyodu <[email protected]>

* Removed the changes in extractSubAggsAndDocCount

Signed-off-by: Vinay Krishna Pudyodu <[email protected]>

---------

Signed-off-by: Vinay Krishna Pudyodu <[email protected]>

Author: vinaykpud

plugins/engine-datafusion/jni/src/partial_agg_optimizer.rs

@@ -55,8 +55,8 @@ impl PartialAggregationOptimizer {
             // println!("[DEBUG] Aggregate expressions: {:?}", agg.aggr_expr().iter().map(|e| e.name()).collect::<Vec<_>>());
 
             let needs_partial = agg.aggr_expr().iter().any(|e| {
-                let name = e.name().to_lowercase();
-                name.starts_with("approx_distinct(")
+                let name = e.fun().name().to_lowercase();
+                name.eq("approx_distinct")
             });
 
             if needs_partial && !matches!(agg.mode(), &AggregateMode::Partial) {

server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeAggregator.java

@@ -79,6 +79,7 @@
 import org.opensearch.search.aggregations.bucket.filterrewrite.FilterRewriteOptimizationContext;
 import org.opensearch.search.aggregations.bucket.missing.MissingOrder;
 import org.opensearch.search.aggregations.bucket.terms.LongKeyedBucketOrds;
+import org.opensearch.search.aggregations.bucket.terms.StringTerms;
 import org.opensearch.search.aggregations.metrics.InternalValueCount;
 import org.opensearch.search.aggregations.metrics.ValueCountAggregator;
 import org.opensearch.search.internal.SearchContext;
@@ -739,10 +740,13 @@ public List<InternalAggregation> convert(Map<String, Object[]> shardResult, Sear
         List<CompositeKey> compositeKeys = new ArrayList<>(shardResult.size());
         for (int i = 0; i < shardResult.get(shardResult.keySet().stream().findFirst().get()).length; i++) {
             for (CompositeValuesSourceConfig sourceConfig : sourceConfigs) {
-                if (sourceConfig.fieldType() == null) {
-                    throw new UnsupportedOperationException("Composite aggregation does not support script field types");
-                }
-                Object[] values = shardResult.get(sourceConfig.fieldType().name());
+//                if (sourceConfig.fieldType() == null) {
+//                    throw new UnsupportedOperationException("Composite aggregation does not support script field types");
+//                }
+                // source=hits | eval m = extract(minute from EventTime) | stats count() by UserID, m, SearchPhrase | sort - \`count()\` | head 10
+                // for above query without this change it will fail above
+                // We can get the name directly from sourceConfig
+                Object[] values = shardResult.get(sourceConfig.name());
                 // TODO : Would require conversion for certain types,
                 currentCompositeKey.add(searchContext.convertToComparable(values[i]));
             }

server/src/main/java/org/opensearch/search/aggregations/metrics/CardinalityAggregator.java

@@ -765,7 +765,10 @@ public long nextValue() throws IOException {
 
     @Override
     public InternalAggregation convertRow(Map<String, Object[]> shardResult, int row, SearchContext searchContext) {
-        Object[] hlls = shardResult.get(name);
+        Object[] hlls = shardResult.get(name + "[hll_registers]");
+        if (hlls == null) {
+            hlls = shardResult.get(name);
+        }
         HyperLogLogPlusPlus sketch = DataFusionHLLWrapper.getHyperLogLogPlusPlus((byte[]) hlls[row]);
         return new InternalCardinality(name, sketch, null);
     }

server/src/main/java/org/opensearch/search/query/SearchEngineResultConversionUtils.java

@@ -15,6 +15,8 @@
 import org.opensearch.search.aggregations.InternalAggregation;
 import org.opensearch.search.aggregations.InternalAggregations;
 import org.opensearch.search.aggregations.ShardResultConvertor;
+import org.opensearch.search.aggregations.metrics.CardinalityAggregator;
+import org.opensearch.search.aggregations.metrics.InternalCardinality;
 import org.opensearch.search.aggregations.metrics.InternalValueCount;
 import org.opensearch.search.aggregations.metrics.ValueCountAggregator;
 import org.opensearch.search.internal.SearchContext;
@@ -35,6 +37,11 @@ public static void convertDFResultGeneric(SearchContext searchContext) {
         if (searchContext.aggregations() != null) {
             Map<String, Object[]> dfResult = searchContext.getDFResults();
 
+//            LOGGER.info("DF Results at convertDFResultGeneric:");
+//            for (Map.Entry<String, Object[]> entry : dfResult.entrySet()) {
+//                LOGGER.info("{}: {}", entry.getKey(), java.util.Arrays.toString(entry.getValue()));
+//            }
+
             // Create aggregators which will process the result from DataFusion
             try {
 
@@ -76,6 +83,9 @@ public static Tuple<List<InternalAggregation>, Long> extractSubAggsAndDocCount(A
                 if (aggregator instanceof ValueCountAggregator) {
                     docCount = ((InternalValueCount) subAgg).getValue();
                 }
+//                if (aggregator instanceof CardinalityAggregator) {
+//                    docCount = ((InternalCardinality) subAgg).getValue();
+//                }
                 subAggs.add(subAgg);
             }
         }