handle unsigned long in flush operations for star tree

opensearch-project · Nov 26, 2024 · 7011249 · 7011249
1 parent 3da97f2
commit 7011249
Show file tree

Hide file tree

Showing 8 changed files with 150 additions and 61 deletions.
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/DimensionFactory.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/DimensionFactory.java
@@ -36,14 +36,15 @@ public class DimensionFactory {
     public static Dimension parseAndCreateDimension(
         String name,
         String type,
+        Boolean isUnsignedLong,
         Map<String, Object> dimensionMap,
         Mapper.TypeParser.ParserContext c
     ) {
         switch (type) {
             case DateDimension.DATE:
                 return parseAndCreateDateDimension(name, dimensionMap, c);
             case NumericDimension.NUMERIC:
-                return new NumericDimension(name);
+                return new NumericDimension(name, isUnsignedLong);
             case KEYWORD:
                 return new KeywordDimension(name);
             default:
@@ -53,6 +54,15 @@ public static Dimension parseAndCreateDimension(
         }
     }
 
+    public static Dimension parseAndCreateDimension(
+        String name,
+        String type,
+        Map<String, Object> dimensionMap,
+        Mapper.TypeParser.ParserContext c
+    ) {
+        return parseAndCreateDimension(name, type, false, dimensionMap, c);
+    }
+
     public static Dimension parseAndCreateDimension(
         String name,
         Mapper.Builder builder,
@@ -68,7 +78,7 @@ public static Dimension parseAndCreateDimension(
             case DATE:
                 return parseAndCreateDateDimension(name, dimensionMap, c);
             case NUMERIC:
-                return new NumericDimension(name);
+                return new NumericDimension(name, builder.isUnsignedLong());
             case KEYWORD:
                 return new KeywordDimension(name);
             default:

diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/NumericDimension.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/NumericDimension.java
@@ -26,10 +26,19 @@
 @ExperimentalApi
 public class NumericDimension implements Dimension {
     public static final String NUMERIC = "numeric";
+    public static final String IS_UNSIGNED_LONG_FIELD = "isUnsignedLong";
+
     private final String field;
+    private final Boolean isUnsignedLong;
 
     public NumericDimension(String field) {
         this.field = field;
+        isUnsignedLong = false;
+    }
+
+    public NumericDimension(String field, Boolean isUnsignedLong) {
+        this.field = field;
+        this.isUnsignedLong = isUnsignedLong;
     }
 
     public String getField() {
@@ -61,6 +70,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         builder.startObject();
         builder.field(CompositeDataCubeFieldType.NAME, field);
         builder.field(CompositeDataCubeFieldType.TYPE, NUMERIC);
+        builder.field(IS_UNSIGNED_LONG_FIELD, isUnsignedLong);
         builder.endObject();
         return builder;
     }
@@ -77,4 +87,8 @@ public boolean equals(Object o) {
     public int hashCode() {
         return Objects.hash(field);
     }
+
+    public Boolean isUnsignedLong() {
+        return isUnsignedLong;
+    }
 }
diff --git a/...org/opensearch/index/compositeindex/datacube/startree/builder/OffHeapStarTreeBuilder.java b/...org/opensearch/index/compositeindex/datacube/startree/builder/OffHeapStarTreeBuilder.java
@@ -17,6 +17,7 @@
 import org.apache.lucene.util.LongValues;
 import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.common.util.io.IOUtils;
+import org.opensearch.index.compositeindex.datacube.Dimension;
 import org.opensearch.index.compositeindex.datacube.startree.StarTreeDocument;
 import org.opensearch.index.compositeindex.datacube.startree.StarTreeField;
 import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
@@ -228,14 +229,15 @@ private Iterator<StarTreeDocument> sortAndReduceDocuments(int[] sortedDocIds, in
                 logger.debug("Sorted doc ids array is null");
                 return Collections.emptyIterator();
             }
+            List<Dimension> dimensionsOrder = starTreeDocumentFileManager.starTreeField.getDimensionsOrder();
             try {
                 StarTreeDocumentsSorter.sort(sortedDocIds, -1, numDocs, index -> {
                     try {
                         return segmentDocumentFileManager.readDimensions(sortedDocIds[index]);
                     } catch (IOException e) {
                         throw new UncheckedIOException(e);
                     }
-                });
+                }, dimensionsOrder);
             } catch (UncheckedIOException ex) {
                 // Unwrap UncheckedIOException and throw as IOException
                 if (ex.getCause() != null) {
@@ -308,6 +310,7 @@ public List<StarTreeDocument> getStarTreeDocuments() throws IOException {
     @Override
     public Long getDimensionValue(int docId, int dimensionId) throws IOException {
         return starTreeDocumentFileManager.getDimensionValue(docId, dimensionId);
+
     }
 
     /**
@@ -328,13 +331,15 @@ public Iterator<StarTreeDocument> generateStarTreeDocumentsForStarNode(int start
         for (int i = 0; i < numDocs; i++) {
             sortedDocIds[i] = startDocId + i;
         }
+        List<Dimension> dimensionsOrder = starTreeDocumentFileManager.starTreeField.getDimensionsOrder();
         StarTreeDocumentsSorter.sort(sortedDocIds, dimensionId, numDocs, index -> {
             try {
                 return starTreeDocumentFileManager.readDimensions(sortedDocIds[index]);
             } catch (IOException e) {
                 throw new RuntimeException(e);
             }
-        });
+        }, dimensionsOrder);
+
         // Create an iterator for aggregated documents
         return new Iterator<StarTreeDocument>() {
             boolean hasNext = true;

diff --git a/.../org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeDocumentsSorter.java b/.../org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeDocumentsSorter.java
@@ -9,7 +9,10 @@
 package org.opensearch.index.compositeindex.datacube.startree.utils;
 
 import org.apache.lucene.util.IntroSorter;
+import org.opensearch.index.compositeindex.datacube.Dimension;
+import org.opensearch.index.compositeindex.datacube.NumericDimension;
 
+import java.util.List;
 import java.util.Objects;
 import java.util.function.IntFunction;
 
@@ -24,7 +27,8 @@ public static void sort(
         final int[] sortedDocIds,
         final int dimensionId,
         final int numDocs,
-        final IntFunction<Long[]> dimensionsReader
+        final IntFunction<Long[]> dimensionsReader,
+        final List<Dimension> dimensionsOrder
     ) {
         new IntroSorter() {
             private Long[] dimensions;
@@ -45,19 +49,27 @@ protected void setPivot(int i) {
             protected int comparePivot(int j) {
                 Long[] currentDimensions = dimensionsReader.apply(j);
                 for (int i = dimensionId + 1; i < dimensions.length; i++) {
-                    Long dimension = currentDimensions[i];
-                    if (!Objects.equals(dimensions[i], dimension)) {
-                        if (dimensions[i] == null && dimension == null) {
+                    Dimension dimension = dimensionsOrder.get(i);
+                    Long dimensionValue = currentDimensions[i];
+                    if (!Objects.equals(dimensions[i], dimensionValue)) {
+                        if (dimensions[i] == null && dimensionValue == null) {
                             return 0;
                         }
-                        if (dimension == null) {
+                        if (dimensionValue == null) {
                             return -1;
                         }
                         if (dimensions[i] == null) {
                             return 1;
                         }
-                        return Long.compare(dimensions[i], dimension);
+                        if (dimension instanceof NumericDimension) {
+                            NumericDimension numericDimension = (NumericDimension) dimension;
+                            if (numericDimension.isUnsignedLong()) {
+                                return Long.compareUnsigned(dimensions[i], dimensionValue);
+                            }
+                        }
+                        return Long.compare(dimensions[i], dimensionValue);
                     }
+
                 }
                 return 0;
             }

diff --git a/server/src/main/java/org/opensearch/index/mapper/MapperBuilderProperties.java b/server/src/main/java/org/opensearch/index/mapper/MapperBuilderProperties.java
@@ -29,6 +29,15 @@ default Optional<DimensionType> getSupportedDataCubeDimensionType() {
         return Optional.empty();
     }
 
+    /**
+     * Indicates whether the current dimension is unsigned long.
+     *
+     * @return true if the dimension is unsigned long, false otherwise.
+     */
+    default Boolean isUnsignedLong() {
+        return false;
+    }
+
     /**
      * Indicates whether the implementation supports data cube metrics.
      *

diff --git a/server/src/main/java/org/opensearch/index/mapper/NumberFieldMapper.java b/server/src/main/java/org/opensearch/index/mapper/NumberFieldMapper.java
@@ -177,15 +177,14 @@ public NumberFieldMapper build(BuilderContext context) {
 
         @Override
         public Optional<DimensionType> getSupportedDataCubeDimensionType() {
-
-            // unsigned long is not supported as dimension for star tree
-            if (type.numericType.equals(NumericType.UNSIGNED_LONG)) {
-                return Optional.empty();
-            }
-
             return Optional.of(DimensionType.NUMERIC);
         }
 
+        @Override
+        public Boolean isUnsignedLong() {
+            return type.numericType.equals(NumericType.UNSIGNED_LONG);
+        }
+
         @Override
         public boolean isDataCubeMetricSupported() {
             return true;

diff --git a/server/src/main/java/org/opensearch/index/mapper/StarTreeMapper.java b/server/src/main/java/org/opensearch/index/mapper/StarTreeMapper.java
@@ -34,6 +34,8 @@
 import java.util.Set;
 import java.util.stream.Collectors;
 
+import static org.opensearch.index.compositeindex.datacube.NumericDimension.IS_UNSIGNED_LONG_FIELD;
+
 /**
  * A field mapper for star tree fields
  *
@@ -239,12 +241,14 @@ private Dimension getDimension(String fieldName, Object dimensionMapping, Mapper
             if (this.objbuilder == null || this.objbuilder.mappersBuilders == null) {
                 String type = (String) XContentMapValues.extractValue(CompositeDataCubeFieldType.TYPE, dimensionMap);
                 dimensionMap.remove(CompositeDataCubeFieldType.TYPE);
+                Boolean isUnsignedLong = (Boolean) XContentMapValues.extractValue(IS_UNSIGNED_LONG_FIELD, dimensionMap);
+                dimensionMap.remove(IS_UNSIGNED_LONG_FIELD);
                 if (type == null) {
                     throw new MapperParsingException(
                         String.format(Locale.ROOT, "unable to parse ordered_dimensions for star tree field [%s]", fieldName)
                     );
                 }
-                return DimensionFactory.parseAndCreateDimension(name, type, dimensionMap, context);
+                return DimensionFactory.parseAndCreateDimension(name, type, isUnsignedLong, dimensionMap, context);
             } else {
                 Optional<Mapper.Builder> dimBuilder = findMapperBuilderByName(name, this.objbuilder.mappersBuilders);
                 if (dimBuilder.isEmpty()) {