Star Tree Merge and Aggregation Fixes (opensearch-project#15274)

--------- Signed-off-by: Sarthak Aggarwal <sarthagg@amazon.com>
akolarkunnu · Sep 10, 2024 · 9e1a444 · 9e1a444
1 parent 777dfd5
commit 9e1a444
Show file tree

Hide file tree

Showing 11 changed files with 64 additions and 48 deletions.
diff --git a/...g/opensearch/index/compositeindex/datacube/startree/aggregators/CountValueAggregator.java b/...g/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/...ch/index/compositeindex/datacube/startree/aggregators/StatelessDoubleValueAggregator.java b/...ch/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/...org/opensearch/index/compositeindex/datacube/startree/aggregators/SumValueAggregator.java b/...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/...va/org/opensearch/index/compositeindex/datacube/startree/aggregators/ValueAggregator.java b/...va/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/...va/org/opensearch/index/compositeindex/datacube/startree/builder/BaseStarTreeBuilder.java b/...va/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/...arch/index/compositeindex/datacube/startree/aggregators/AbstractValueAggregatorTests.java b/...arch/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/...nsearch/index/compositeindex/datacube/startree/aggregators/CountValueAggregatorTests.java b/...nsearch/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/...pensearch/index/compositeindex/datacube/startree/aggregators/MaxValueAggregatorTests.java b/...pensearch/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/...pensearch/index/compositeindex/datacube/startree/aggregators/MinValueAggregatorTests.java b/...pensearch/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/...pensearch/index/compositeindex/datacube/startree/aggregators/SumValueAggregatorTests.java b/...pensearch/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/...ensearch/index/compositeindex/datacube/startree/builder/AbstractStarTreeBuilderTests.java b/...ensearch/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);