Making terms agg work with string (#19874)

Signed-off-by: expani <[email protected]>

Author: expani

plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/DatafusionContext.java

@@ -8,6 +8,7 @@
 
 package org.opensearch.datafusion.search;
 
+import org.apache.arrow.vector.util.Text;
 import org.apache.lucene.search.Collector;
 import org.apache.lucene.search.CollectorManager;
 import org.apache.lucene.search.FieldDoc;
@@ -807,4 +808,14 @@ public Map<String, Object[]> getDFResults() {
         return dfResults;
     }
 
+    @Override
+    public Comparable<?> convertToComparable(Object rawValue) {
+        if (rawValue instanceof Number) {
+            return (Comparable) rawValue;
+        } else if (rawValue instanceof Text) {
+            return rawValue.toString();
+        } else {
+            throw new IllegalArgumentException("Conversion to Comparable not supported for type " + rawValue.getClass());
+        }
+    }
 }

server/src/main/java/org/opensearch/search/aggregations/ShardResultConvertor.java

@@ -8,22 +8,24 @@
 
 package org.opensearch.search.aggregations;
 
+import org.opensearch.search.internal.SearchContext;
+
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 
 public interface ShardResultConvertor {
 
-    default List<InternalAggregation> convert(Map<String, Object[]> shardResult) {
+    default List<InternalAggregation> convert(Map<String, Object[]> shardResult, SearchContext searchContext) {
         int rows = shardResult.entrySet().stream().findFirst().get().getValue().length;
         List<InternalAggregation> internalAggregations = new ArrayList<>();
         for (int i = 0; i < rows; i++) {
-            internalAggregations.add(convertRow(shardResult, i));
+            internalAggregations.add(convertRow(shardResult, i, searchContext));
         }
         return internalAggregations;
     }
 
-    default InternalAggregation convertRow(Map<String, Object[]> shardResult, int row) {
+    default InternalAggregation convertRow(Map<String, Object[]> shardResult, int row, SearchContext searchContext) {
         throw new UnsupportedOperationException("Row conversion not supported");
     }
 

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

@@ -729,7 +729,7 @@ public void collect(int doc, long zeroBucket) throws IOException {
     }
 
     @Override
-    public List<InternalAggregation> convert(Map<String, Object[]> shardResult) {
+    public List<InternalAggregation> convert(Map<String, Object[]> shardResult, SearchContext searchContext) {
         // Generate the composite keys
         List<Comparable> currentCompositeKey = new ArrayList<>(sourceConfigs.length);
         List<CompositeKey> compositeKeys = new ArrayList<>(shardResult.size());
@@ -740,7 +740,7 @@ public List<InternalAggregation> convert(Map<String, Object[]> shardResult) {
                 }
                 Object[] values = shardResult.get(sourceConfig.fieldType().name());
                 // TODO : Would require conversion for certain types,
-                currentCompositeKey.add((Comparable) values[i]);
+                currentCompositeKey.add(searchContext.convertToComparable(values[i]));
             }
             compositeKeys.add(new CompositeKey(currentCompositeKey.toArray(new Comparable[0])));
             currentCompositeKey.clear();
@@ -754,7 +754,7 @@ public List<InternalAggregation> convert(Map<String, Object[]> shardResult) {
                     throw new UnsupportedOperationException(String.format("Aggregation [%s] doesn't support shard result conversion Impl [%s]", subAgg.name(), subAgg.getClass().getName()));
                 }
                 ShardResultConvertor convertor = (ShardResultConvertor) subAgg;
-                subAggs.add(convertor.convertRow(shardResult, row));
+                subAggs.add(convertor.convertRow(shardResult, row, searchContext));
             }
             buckets.add(new InternalComposite.InternalBucket(
                 sourceNames,

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

@@ -282,7 +282,7 @@ public void collectStarTreeEntry(int starTreeEntryBit, long bucket) throws IOExc
     }
 
     @Override
-    public InternalAggregation convertRow(Map<String, Object[]> shardResult, int row) {
+    public InternalAggregation convertRow(Map<String, Object[]> shardResult, int row, SearchContext searchContext) {
         Object[] counts = shardResult.get(name + "_count");
         Object[] sums = shardResult.get(name + "_sum");
         return new InternalAvg(name, ((Number) sums[row]).doubleValue(), ((Number) sums[row]).longValue(), format, metadata());

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

@@ -285,7 +285,7 @@ public void doReset() {
     }
 
     @Override
-    public InternalAggregation convertRow(Map<String, Object[]> shardResult, int row) {
+    public InternalAggregation convertRow(Map<String, Object[]> shardResult, int row, SearchContext searchContext) {
         Object[] values = shardResult.get(name);
         return new InternalMax(name, ((Number) values[row]).doubleValue(), formatter, metadata());
     }

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

@@ -276,7 +276,7 @@ public StarTreeBucketCollector getStarTreeBucketCollector(
     }
 
     @Override
-    public InternalAggregation convertRow(Map<String, Object[]> shardResult, int row) {
+    public InternalAggregation convertRow(Map<String, Object[]> shardResult, int row, SearchContext searchContext) {
         Object[] values = shardResult.get(name);
         return new InternalMin(name, ((Number) values[row]).doubleValue(), format, metadata());
     }

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

@@ -220,7 +220,7 @@ public void doClose() {
     }
 
     @Override
-    public InternalAggregation convertRow(Map<String, Object[]> shardResult, int row) {
+    public InternalAggregation convertRow(Map<String, Object[]> shardResult, int row, SearchContext searchContext) {
         Object[] values = shardResult.get(name);
         return new InternalSum(name, ((Number) values[row]).doubleValue(), format, metadata());
     }

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

@@ -214,7 +214,7 @@ public StarTreeBucketCollector getStarTreeBucketCollector(
     }
 
     @Override
-    public InternalAggregation convertRow(Map<String, Object[]> shardResult, int row) {
+    public InternalAggregation convertRow(Map<String, Object[]> shardResult, int row, SearchContext searchContext) {
             Object[] values = shardResult.get(name);
             return new InternalValueCount(name, ((Number) values[row]).longValue(), metadata());
     }

server/src/main/java/org/opensearch/search/internal/SearchContext.java

@@ -574,4 +574,9 @@ public void setDFResults(Map<String, Object[]> dfResults) {
     public Map<String, Object[]> getDFResults() {
         return Collections.emptyMap();
     }
+
+    // TODO : This should be a part of mapper given by DataFormat or SearchEngine as related to Field type.
+    public Comparable convertToComparable(Object rawValue) {
+        throw new UnsupportedOperationException("Engine doesn't implement response value conversion");
+    }
 }

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

@@ -23,8 +23,6 @@
 
 public class SearchEngineResultConversionUtils {
 
-    private static final Logger LOGGER = LogManager.getLogger(SearchEngineResultConversionUtils.class);
-
     public static void convertDFResultGeneric(SearchContext searchContext) {
         if (searchContext.aggregations() != null) {
             Map<String, Object[]> dfResult = searchContext.getDFResults();
@@ -51,9 +49,8 @@ public static void convertDFResultGeneric(SearchContext searchContext) {
                 }).toList();
 
                 InternalAggregations internalAggregations = InternalAggregations.from(
-                    shardResultConvertors.stream().flatMap(x -> x.convert(dfResult).stream()).collect(Collectors.toList())
+                    shardResultConvertors.stream().flatMap(x -> x.convert(dfResult, searchContext).stream()).collect(Collectors.toList())
                 );
-//                LOGGER.info("Internal Aggregations converted {}", internalAggregations.asMap());
                 searchContext.queryResult().aggregations(internalAggregations);
             } catch (IOException e) {
                 throw new RuntimeException(e);