Star Tree Merge and Aggregation Fixes (#15274)

sarthakaggarwal97 · web-flow · commit a27d26bce258 · 2024-08-19T16:52:32.000+05:30
---------

Signed-off-by: Sarthak Aggarwal &lt;sarthagg@amazon.com&gt;
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/CountValueAggregator.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/CountValueAggregator.java
@@ -60,7 +60,7 @@ public Long mergeAggregatedValues(Long value, Long aggregatedValue) {
     }
 
     @Override
-    public Long toStarTreeNumericTypeValue(Long value) {
+    public Long toAggregatedValueType(Long value) {
         return value;
     }
 
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/StatelessDoubleValueAggregator.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/StatelessDoubleValueAggregator.java
@@ -52,12 +52,12 @@ public Double mergeAggregatedValues(Double value, Double aggregatedValue) {
     }
 
     @Override
-    public Double toStarTreeNumericTypeValue(Long value) {
+    public Double toAggregatedValueType(Long value) {
         try {
             if (value == null) {
                 return getIdentityMetricValue();
             }
-            return starTreeNumericType.getDoubleValue(value);
+            return VALUE_AGGREGATOR_TYPE.getDoubleValue(value);
         } catch (Exception e) {
             throw new IllegalStateException("Cannot convert " + value + " to sortable aggregation type", e);
         }
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/SumValueAggregator.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/SumValueAggregator.java
@@ -87,12 +87,12 @@ public Double getInitialAggregatedValue(Double value) {
     }
 
     @Override
-    public Double toStarTreeNumericTypeValue(Long value) {
+    public Double toAggregatedValueType(Long value) {
         try {
             if (value == null) {
                 return getIdentityMetricValue();
             }
-            return starTreeNumericType.getDoubleValue(value);
+            return VALUE_AGGREGATOR_TYPE.getDoubleValue(value);
         } catch (Exception e) {
             throw new IllegalStateException("Cannot convert " + value + " to sortable aggregation type", e);
         }
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/ValueAggregator.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/ValueAggregator.java
@@ -50,9 +50,9 @@ default A getInitialAggregatedValue(A value) {
     }
 
     /**
-     * Converts an aggregated value from a Long type.
+     * Converts a segment long value to an aggregated value.
      */
-    A toStarTreeNumericTypeValue(Long rawValue);
+    A toAggregatedValueType(Long rawValue);
 
     /**
      * Fetches a value that does not alter the result of aggregations
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/BaseStarTreeBuilder.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/BaseStarTreeBuilder.java
@@ -164,7 +164,7 @@ protected StarTreeDocument getStarTreeDocument(
             // actual indexing field they're based on
             metrics[i] = metricAggregatorInfos.get(i)
                 .getValueAggregators()
-                .toStarTreeNumericTypeValue(metricDocValuesIterator.value(currentDocId));
+                .toAggregatedValueType(metricDocValuesIterator.value(currentDocId));
             i++;
         }
         return new StarTreeDocument(dims, metrics);
diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/AbstractValueAggregatorTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/AbstractValueAggregatorTests.java
@@ -21,7 +21,7 @@
 public abstract class AbstractValueAggregatorTests extends OpenSearchTestCase {
 
     private ValueAggregator aggregator;
-    private StarTreeNumericType starTreeNumericType;
+    protected StarTreeNumericType starTreeNumericType;
 
     public AbstractValueAggregatorTests(StarTreeNumericType starTreeNumericType) {
         this.starTreeNumericType = starTreeNumericType;
@@ -69,7 +69,7 @@ public void testGetInitialAggregatedValueForSegmentDocValue() {
             assertEquals(CountValueAggregator.DEFAULT_INITIAL_VALUE, aggregator.getInitialAggregatedValueForSegmentDocValue(randomLong()));
         } else {
             assertEquals(
-                aggregator.toStarTreeNumericTypeValue(randomLong),
+                starTreeNumericType.getDoubleValue(randomLong),
                 aggregator.getInitialAggregatedValueForSegmentDocValue(randomLong)
             );
         }
diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/CountValueAggregatorTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/CountValueAggregatorTests.java
@@ -36,10 +36,10 @@ public void testGetInitialAggregatedValue() {
         assertEquals(randomLong, aggregator.getInitialAggregatedValue(randomLong), 0.0);
     }
 
-    public void testToStarTreeNumericTypeValue() {
+    public void testToAggregatedValueType() {
         long randomLong = randomLong();
-        assertEquals(randomLong, aggregator.toStarTreeNumericTypeValue(randomLong), 0.0);
-        assertNull(aggregator.toStarTreeNumericTypeValue(null));
+        assertEquals(randomLong, aggregator.toAggregatedValueType(randomLong), 0.0);
+        assertNull(aggregator.toAggregatedValueType(null));
     }
 
     public void testIdentityMetricValue() {
diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MaxValueAggregatorTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MaxValueAggregatorTests.java
@@ -23,21 +23,13 @@ public void testMergeAggregatedValueAndSegmentValue() {
         Long randomLong = randomLong();
         double randomDouble = randomDouble();
         assertEquals(
-            Math.max(aggregator.toStarTreeNumericTypeValue(randomLong), randomDouble),
+            Math.max(starTreeNumericType.getDoubleValue(randomLong), randomDouble),
             aggregator.mergeAggregatedValueAndSegmentValue(randomDouble, randomLong),
             0.0
         );
-        assertEquals(
-            aggregator.toStarTreeNumericTypeValue(randomLong),
-            aggregator.mergeAggregatedValueAndSegmentValue(null, randomLong),
-            0.0
-        );
+        assertEquals(starTreeNumericType.getDoubleValue(randomLong), aggregator.mergeAggregatedValueAndSegmentValue(null, randomLong), 0.0);
         assertEquals(randomDouble, aggregator.mergeAggregatedValueAndSegmentValue(randomDouble, null), 0.0);
-        assertEquals(
-            Math.max(2.0, aggregator.toStarTreeNumericTypeValue(3L)),
-            aggregator.mergeAggregatedValueAndSegmentValue(2.0, 3L),
-            0.0
-        );
+        assertEquals(Math.max(2.0, starTreeNumericType.getDoubleValue(3L)), aggregator.mergeAggregatedValueAndSegmentValue(2.0, 3L), 0.0);
     }
 
     public void testMergeAggregatedValues() {
@@ -53,10 +45,10 @@ public void testGetInitialAggregatedValue() {
         assertEquals(randomDouble, aggregator.getInitialAggregatedValue(randomDouble), 0.0);
     }
 
-    public void testToStarTreeNumericTypeValue() {
+    public void testToAggregatedValueType() {
         MaxValueAggregator aggregator = new MaxValueAggregator(StarTreeNumericType.DOUBLE);
         long randomLong = randomLong();
-        assertEquals(NumericUtils.sortableLongToDouble(randomLong), aggregator.toStarTreeNumericTypeValue(randomLong), 0.0);
+        assertEquals(NumericUtils.sortableLongToDouble(randomLong), aggregator.toAggregatedValueType(randomLong), 0.0);
     }
 
     public void testIdentityMetricValue() {
diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MinValueAggregatorTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/MinValueAggregatorTests.java
@@ -22,21 +22,13 @@ public void testMergeAggregatedValueAndSegmentValue() {
         Long randomLong = randomLong();
         double randomDouble = randomDouble();
         assertEquals(
-            Math.min(aggregator.toStarTreeNumericTypeValue(randomLong), randomDouble),
+            Math.min(starTreeNumericType.getDoubleValue(randomLong), randomDouble),
             aggregator.mergeAggregatedValueAndSegmentValue(randomDouble, randomLong),
             0.0
         );
-        assertEquals(
-            aggregator.toStarTreeNumericTypeValue(randomLong),
-            aggregator.mergeAggregatedValueAndSegmentValue(null, randomLong),
-            0.0
-        );
+        assertEquals(starTreeNumericType.getDoubleValue(randomLong), aggregator.mergeAggregatedValueAndSegmentValue(null, randomLong), 0.0);
         assertEquals(randomDouble, aggregator.mergeAggregatedValueAndSegmentValue(randomDouble, null), 0.0);
-        assertEquals(
-            Math.min(2.0, aggregator.toStarTreeNumericTypeValue(3L)),
-            aggregator.mergeAggregatedValueAndSegmentValue(2.0, 3L),
-            0.0
-        );
+        assertEquals(Math.min(2.0, starTreeNumericType.getDoubleValue(3L)), aggregator.mergeAggregatedValueAndSegmentValue(2.0, 3L), 0.0);
     }
 
     public void testMergeAggregatedValues() {
@@ -52,10 +44,10 @@ public void testGetInitialAggregatedValue() {
         assertEquals(randomDouble, aggregator.getInitialAggregatedValue(randomDouble), 0.0);
     }
 
-    public void testToStarTreeNumericTypeValue() {
+    public void testToAggregatedValueType() {
         MinValueAggregator aggregator = new MinValueAggregator(StarTreeNumericType.DOUBLE);
         long randomLong = randomLong();
-        assertEquals(NumericUtils.sortableLongToDouble(randomLong), aggregator.toStarTreeNumericTypeValue(randomLong), 0.0);
+        assertEquals(NumericUtils.sortableLongToDouble(randomLong), aggregator.toAggregatedValueType(randomLong), 0.0);
     }
 
     public void testIdentityMetricValue() {
diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/SumValueAggregatorTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/aggregators/SumValueAggregatorTests.java
@@ -29,7 +29,7 @@ public void testMergeAggregatedValueAndSegmentValue() {
         Long randomLong = randomLong();
         aggregator.getInitialAggregatedValue(randomDouble);
         assertEquals(
-            randomDouble + aggregator.toStarTreeNumericTypeValue(randomLong),
+            randomDouble + starTreeNumericType.getDoubleValue(randomLong),
             aggregator.mergeAggregatedValueAndSegmentValue(randomDouble, randomLong),
             0.0
         );
@@ -41,7 +41,7 @@ public void testMergeAggregatedValueAndSegmentValue_nullSegmentDocValue() {
         aggregator.getInitialAggregatedValue(randomDouble1);
         assertEquals(randomDouble1, aggregator.mergeAggregatedValueAndSegmentValue(randomDouble1, null), 0.0);
         assertEquals(
-            randomDouble1 + aggregator.toStarTreeNumericTypeValue(randomLong),
+            randomDouble1 + starTreeNumericType.getDoubleValue(randomLong),
             aggregator.mergeAggregatedValueAndSegmentValue(randomDouble1, randomLong),
             0.0
         );
@@ -50,11 +50,7 @@ public void testMergeAggregatedValueAndSegmentValue_nullSegmentDocValue() {
     public void testMergeAggregatedValueAndSegmentValue_nullInitialDocValue() {
         Long randomLong = randomLong();
         aggregator.getInitialAggregatedValue(null);
-        assertEquals(
-            aggregator.toStarTreeNumericTypeValue(randomLong),
-            aggregator.mergeAggregatedValueAndSegmentValue(null, randomLong),
-            0.0
-        );
+        assertEquals(starTreeNumericType.getDoubleValue(randomLong), aggregator.mergeAggregatedValueAndSegmentValue(null, randomLong), 0.0);
     }
 
     public void testMergeAggregatedValues() {
@@ -70,9 +66,9 @@ public void testGetInitialAggregatedValue() {
         assertEquals(randomDouble, aggregator.getInitialAggregatedValue(randomDouble), 0.0);
     }
 
-    public void testToStarTreeNumericTypeValue() {
+    public void testToAggregatedValueType() {
         long randomLong = randomLong();
-        assertEquals(aggregator.toStarTreeNumericTypeValue(randomLong), aggregator.toStarTreeNumericTypeValue(randomLong), 0.0);
+        assertEquals(aggregator.toAggregatedValueType(randomLong), aggregator.toAggregatedValueType(randomLong), 0.0);
     }
 
     public void testIdentityMetricValue() {
diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/AbstractStarTreeBuilderTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/AbstractStarTreeBuilderTests.java
@@ -1508,6 +1508,42 @@ private static StarTreeField getStarTreeFieldWithMultipleMetrics() {
         return sf;
     }
 
+    public void testMergeFlow_randomNumberTypes() throws Exception {
+
+        DocumentMapper documentMapper = mock(DocumentMapper.class);
+        when(mapperService.documentMapper()).thenReturn(documentMapper);
+        Settings settings = Settings.builder().put(settings(org.opensearch.Version.CURRENT).build()).build();
+        NumberFieldMapper numberFieldMapper1 = new NumberFieldMapper.Builder(
+            "field1",
+            randomFrom(NumberFieldMapper.NumberType.values()),
+            false,
+            true
+        ).build(new Mapper.BuilderContext(settings, new ContentPath()));
+        NumberFieldMapper numberFieldMapper2 = new NumberFieldMapper.Builder(
+            "field2",
+            randomFrom(NumberFieldMapper.NumberType.values()),
+            false,
+            true
+        ).build(new Mapper.BuilderContext(settings, new ContentPath()));
+        NumberFieldMapper numberFieldMapper3 = new NumberFieldMapper.Builder(
+            "field3",
+            randomFrom(NumberFieldMapper.NumberType.values()),
+            false,
+            true
+        ).build(new Mapper.BuilderContext(settings, new ContentPath()));
+        MappingLookup fieldMappers = new MappingLookup(
+            Set.of(numberFieldMapper1, numberFieldMapper2, numberFieldMapper3),
+            Collections.emptyList(),
+            Collections.emptyList(),
+            0,
+            null
+        );
+        when(documentMapper.mappers()).thenReturn(fieldMappers);
+        testMergeFlowWithSum();
+        builder.close();
+        testMergeFlowWithCount();
+    }
+
     public void testMergeFlowWithSum() throws IOException {
         List<Long> dimList = List.of(0L, 1L, 3L, 4L, 5L, 6L);
         List<Integer> docsWithField = List.of(0, 1, 3, 4, 5, 6);

Original file line number	Diff line number	Diff line change
`@@ -60,7 +60,7 @@ public Long mergeAggregatedValues(Long value, Long aggregatedValue) {`
`60`	`60`	`}`
`61`	`61`
`62`	`62`	`@Override`
`63`		`- public Long toStarTreeNumericTypeValue(Long value) {`
	`63`	`+ public Long toAggregatedValueType(Long value) {`
`64`	`64`	`return value;`
`65`	`65`	`}`
`66`	`66`
Original file line number	Diff line number	Diff line change
`@@ -52,12 +52,12 @@ public Double mergeAggregatedValues(Double value, Double aggregatedValue) {`
`52`	`52`	`}`
`53`	`53`
`54`	`54`	`@Override`
`55`		`- public Double toStarTreeNumericTypeValue(Long value) {`
	`55`	`+ public Double toAggregatedValueType(Long value) {`
`56`	`56`	`try {`
`57`	`57`	`if (value == null) {`
`58`	`58`	`return getIdentityMetricValue();`
`59`	`59`	`}`
`60`		`- return starTreeNumericType.getDoubleValue(value);`
	`60`	`+ return VALUE_AGGREGATOR_TYPE.getDoubleValue(value);`
`61`	`61`	`} catch (Exception e) {`
`62`	`62`	`throw new IllegalStateException("Cannot convert " + value + " to sortable aggregation type", e);`
`63`	`63`	`}`
Original file line number	Diff line number	Diff line change
`@@ -87,12 +87,12 @@ public Double getInitialAggregatedValue(Double value) {`
`87`	`87`	`}`
`88`	`88`
`89`	`89`	`@Override`
`90`		`- public Double toStarTreeNumericTypeValue(Long value) {`
	`90`	`+ public Double toAggregatedValueType(Long value) {`
`91`	`91`	`try {`
`92`	`92`	`if (value == null) {`
`93`	`93`	`return getIdentityMetricValue();`
`94`	`94`	`}`
`95`		`- return starTreeNumericType.getDoubleValue(value);`
	`95`	`+ return VALUE_AGGREGATOR_TYPE.getDoubleValue(value);`
`96`	`96`	`} catch (Exception e) {`
`97`	`97`	`throw new IllegalStateException("Cannot convert " + value + " to sortable aggregation type", e);`
`98`	`98`	`}`
Original file line number	Diff line number	Diff line change
`@@ -164,7 +164,7 @@ protected StarTreeDocument getStarTreeDocument(`
`164`	`164`	`// actual indexing field they're based on`
`165`	`165`	`metrics[i] = metricAggregatorInfos.get(i)`
`166`	`166`	`.getValueAggregators()`
`167`		`- .toStarTreeNumericTypeValue(metricDocValuesIterator.value(currentDocId));`
	`167`	`+ .toAggregatedValueType(metricDocValuesIterator.value(currentDocId));`
`168`	`168`	`i++;`
`169`	`169`	`}`
`170`	`170`	`return new StarTreeDocument(dims, metrics);`
Original file line number	Diff line number	Diff line change
`@@ -36,10 +36,10 @@ public void testGetInitialAggregatedValue() {`
`36`	`36`	`assertEquals(randomLong, aggregator.getInitialAggregatedValue(randomLong), 0.0);`
`37`	`37`	`}`
`38`	`38`
`39`		`- public void testToStarTreeNumericTypeValue() {`
	`39`	`+ public void testToAggregatedValueType() {`
`40`	`40`	`long randomLong = randomLong();`
`41`		`- assertEquals(randomLong, aggregator.toStarTreeNumericTypeValue(randomLong), 0.0);`
`42`		`- assertNull(aggregator.toStarTreeNumericTypeValue(null));`
	`41`	`+ assertEquals(randomLong, aggregator.toAggregatedValueType(randomLong), 0.0);`
	`42`	`+ assertNull(aggregator.toAggregatedValueType(null));`
`43`	`43`	`}`
`44`	`44`
`45`	`45`	`public void testIdentityMetricValue() {`