star tree keyword changes

Signed-off-by: Bharathwaj G <[email protected]>
opensearch-project · Sep 27, 2024 · 249fa53 · 249fa53
1 parent eb43e37
commit 249fa53
Show file tree

Hide file tree

Showing 8 changed files with 184 additions and 24 deletions.
diff --git a/server/src/internalClusterTest/java/org/opensearch/index/mapper/StarTreeMapperIT.java b/server/src/internalClusterTest/java/org/opensearch/index/mapper/StarTreeMapperIT.java
@@ -56,7 +56,7 @@ public class StarTreeMapperIT extends OpenSearchIntegTestCase {
         .put(IndexSettings.INDEX_TRANSLOG_FLUSH_THRESHOLD_SIZE_SETTING.getKey(), new ByteSizeValue(512, ByteSizeUnit.MB))
         .build();
 
-    private static XContentBuilder createMinimalTestMapping(boolean invalidDim, boolean invalidMetric, boolean keywordDim) {
+    private static XContentBuilder createMinimalTestMapping(boolean invalidDim, boolean invalidMetric, boolean ipdim) {
         try {
             return jsonBuilder().startObject()
                 .startObject("composite")
@@ -68,12 +68,15 @@ private static XContentBuilder createMinimalTestMapping(boolean invalidDim, bool
                 .endObject()
                 .startArray("ordered_dimensions")
                 .startObject()
-                .field("name", getDim(invalidDim, keywordDim))
+                .field("name", getDim(invalidDim, ipdim))
+                .endObject()
+                .startObject()
+                .field("name", "keyword_dv")
                 .endObject()
                 .endArray()
                 .startArray("metrics")
                 .startObject()
-                .field("name", getDim(invalidMetric, false))
+                .field("name", getMetric(invalidMetric, false))
                 .endObject()
                 .endArray()
                 .endObject()
@@ -99,6 +102,10 @@ private static XContentBuilder createMinimalTestMapping(boolean invalidDim, bool
                 .field("type", "keyword")
                 .field("doc_values", false)
                 .endObject()
+                .startObject("ip")
+                .field("type", "ip")
+                .field("doc_values", false)
+                .endObject()
                 .endObject()
                 .endObject();
         } catch (IOException e) {
@@ -356,10 +363,19 @@ private XContentBuilder getMappingWithDuplicateFields(boolean isDuplicateDim, bo
     }
 
     private static String getDim(boolean hasDocValues, boolean isKeyword) {
+        if (hasDocValues) {
+            return random().nextBoolean() ? "numeric" : "keyword";
+        } else if (isKeyword) {
+            return "ip";
+        }
+        return "numeric_dv";
+    }
+
+    private static String getMetric(boolean hasDocValues, boolean isKeyword) {
         if (hasDocValues) {
             return "numeric";
         } else if (isKeyword) {
-            return "keyword";
+            return "ip";
         }
         return "numeric_dv";
     }
@@ -398,6 +414,7 @@ public void testValidCompositeIndex() {
                         assertEquals(expectedTimeUnits.get(i).shortName(), dateDim.getSortedCalendarIntervals().get(i).shortName());
                     }
                     assertEquals("numeric_dv", starTreeFieldType.getDimensions().get(1).getField());
+                    assertEquals("keyword_dv", starTreeFieldType.getDimensions().get(2).getField());
                     assertEquals("numeric_dv", starTreeFieldType.getMetrics().get(0).getField());
                     List<MetricStat> expectedMetrics = Arrays.asList(MetricStat.VALUE_COUNT, MetricStat.SUM, MetricStat.AVG);
                     assertEquals(expectedMetrics, starTreeFieldType.getMetrics().get(0).getMetrics());
@@ -665,10 +682,7 @@ public void testInvalidDimCompositeIndex() {
             IllegalArgumentException.class,
             () -> prepareCreate(TEST_INDEX).setSettings(settings).setMapping(createMinimalTestMapping(true, false, false)).get()
         );
-        assertEquals(
-            "Aggregations not supported for the dimension field [numeric] with field type [integer] as part of star tree field",
-            ex.getMessage()
-        );
+        assertTrue(ex.getMessage().startsWith("Aggregations not supported for the dimension field "));
     }
 
     public void testMaxDimsCompositeIndex() {
@@ -734,7 +748,7 @@ public void testUnsupportedDim() {
             () -> prepareCreate(TEST_INDEX).setSettings(settings).setMapping(createMinimalTestMapping(false, false, true)).get()
         );
         assertEquals(
-            "Failed to parse mapping [_doc]: unsupported field type associated with dimension [keyword] as part of star tree field [startree-1]",
+            "Failed to parse mapping [_doc]: unsupported field type associated with dimension [ip] as part of star tree field [startree-1]",
             ex.getMessage()
         );
     }

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
@@ -24,6 +24,7 @@
 import java.util.stream.Collectors;
 
 import static org.opensearch.index.compositeindex.datacube.DateDimension.CALENDAR_INTERVALS;
+import static org.opensearch.index.compositeindex.datacube.KeywordDimension.KEYWORD;
 
 /**
  * Dimension factory class mainly used to parse and create dimension from the mappings
@@ -43,6 +44,8 @@ public static Dimension parseAndCreateDimension(
                 return parseAndCreateDateDimension(name, dimensionMap, c);
             case NumericDimension.NUMERIC:
                 return new NumericDimension(name);
+            case KEYWORD:
+                return new KeywordDimension(name);
             default:
                 throw new IllegalArgumentException(
                     String.format(Locale.ROOT, "unsupported field type associated with dimension [%s] as part of star tree field", name)
@@ -56,16 +59,23 @@ public static Dimension parseAndCreateDimension(
         Map<String, Object> dimensionMap,
         Mapper.TypeParser.ParserContext c
     ) {
-        if (builder.getSupportedDataCubeDimensionType().isPresent()
-            && builder.getSupportedDataCubeDimensionType().get().equals(DimensionType.DATE)) {
-            return parseAndCreateDateDimension(name, dimensionMap, c);
-        } else if (builder.getSupportedDataCubeDimensionType().isPresent()
-            && builder.getSupportedDataCubeDimensionType().get().equals(DimensionType.NUMERIC)) {
+        if (builder.getSupportedDataCubeDimensionType().isEmpty()) {
+            throw new IllegalArgumentException(
+                String.format(Locale.ROOT, "unsupported field type associated with star tree dimension [%s]", name)
+            );
+        }
+        switch (builder.getSupportedDataCubeDimensionType().get()) {
+            case DATE:
+                return parseAndCreateDateDimension(name, dimensionMap, c);
+            case NUMERIC:
                 return new NumericDimension(name);
-            }
-        throw new IllegalArgumentException(
-            String.format(Locale.ROOT, "unsupported field type associated with star tree dimension [%s]", name)
-        );
+            case KEYWORD:
+                return new KeywordDimension(name);
+            default:
+                throw new IllegalArgumentException(
+                    String.format(Locale.ROOT, "unsupported field type associated with star tree dimension [%s]", name)
+                );
+        }
     }
 
     private static DateDimension parseAndCreateDateDimension(

diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/DimensionType.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/DimensionType.java
@@ -27,5 +27,11 @@ public enum DimensionType {
      * Represents a date dimension type.
      * This is used for dimensions that contain date or timestamp values.
      */
-    DATE
+    DATE,
+
+    /**
+     * Represents a keyword dimension type.
+     * This is used for dimensions that contain keyword ordinals.
+     */
+    KEYWORD
 }
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/KeywordDimension.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/KeywordDimension.java
@@ -0,0 +1,75 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.index.compositeindex.datacube;
+
+import org.opensearch.common.annotation.ExperimentalApi;
+import org.opensearch.core.xcontent.XContentBuilder;
+import org.opensearch.index.mapper.CompositeDataCubeFieldType;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Objects;
+
+/**
+ * Composite index keyword dimension class
+ *
+ * @opensearch.experimental
+ */
+@ExperimentalApi
+public class KeywordDimension implements Dimension {
+    public static final String KEYWORD = "keyword";
+    private final String field;
+
+    public KeywordDimension(String field) {
+        this.field = field;
+    }
+
+    @Override
+    public String getField() {
+        return field;
+    }
+
+    @Override
+    public int getNumSubDimensions() {
+        return 1;
+    }
+
+    @Override
+    public int setDimensionValues(Long value, Long[] dims, int index) {
+        dims[index++] = value;
+        return index;
+    }
+
+    @Override
+    public List<String> getDimensionFieldsNames() {
+        return List.of(field);
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        builder.startObject();
+        builder.field(CompositeDataCubeFieldType.NAME, field);
+        builder.field(CompositeDataCubeFieldType.TYPE, KEYWORD);
+        builder.endObject();
+        return builder;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        KeywordDimension dimension = (KeywordDimension) o;
+        return Objects.equals(field, dimension.getField());
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(field);
+    }
+}
diff --git a/...rg/opensearch/index/compositeindex/datacube/startree/builder/StarTreeDocsFileManager.java b/...rg/opensearch/index/compositeindex/datacube/startree/builder/StarTreeDocsFileManager.java
@@ -14,6 +14,7 @@
 import org.apache.lucene.store.IndexInput;
 import org.apache.lucene.store.IndexOutput;
 import org.apache.lucene.store.RandomAccessInput;
+import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.common.util.io.IOUtils;
 import org.opensearch.index.compositeindex.datacube.startree.StarTreeDocument;
 import org.opensearch.index.compositeindex.datacube.startree.StarTreeField;
@@ -45,7 +46,9 @@
  * <p>The set of 'star-tree.documents' files is maintained, and a tracker array is used to keep track of the start document ID for each file.
  * Once the number of files reaches a set threshold, the files are merged.
  *
+ * @opensearch.experimental
  */
+@ExperimentalApi
 public class StarTreeDocsFileManager extends AbstractDocumentsFileManager implements Closeable {
     private static final Logger logger = LogManager.getLogger(StarTreeDocsFileManager.class);
     private static final String STAR_TREE_DOC_FILE_NAME = "star-tree.documents";

diff --git a/server/src/main/java/org/opensearch/index/mapper/KeywordFieldMapper.java b/server/src/main/java/org/opensearch/index/mapper/KeywordFieldMapper.java
@@ -59,6 +59,7 @@
 import org.opensearch.core.xcontent.XContentParser;
 import org.opensearch.index.analysis.IndexAnalyzers;
 import org.opensearch.index.analysis.NamedAnalyzer;
+import org.opensearch.index.compositeindex.datacube.DimensionType;
 import org.opensearch.index.fielddata.IndexFieldData;
 import org.opensearch.index.fielddata.plain.SortedSetOrdinalsIndexFieldData;
 import org.opensearch.index.query.QueryShardContext;
@@ -73,6 +74,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.function.Supplier;
 
 import static org.opensearch.search.SearchService.ALLOW_EXPENSIVE_QUERIES;
@@ -254,6 +256,11 @@ public KeywordFieldMapper build(BuilderContext context) {
                 this
             );
         }
+
+        @Override
+        public Optional<DimensionType> getSupportedDataCubeDimensionType() {
+            return Optional.of(DimensionType.KEYWORD);
+        }
     }
 
     public static final TypeParser PARSER = new TypeParser((n, c) -> new Builder(n, c.getIndexAnalyzers()));

diff --git a/...ensearch/index/compositeindex/datacube/startree/builder/AbstractStarTreeBuilderTests.java b/...ensearch/index/compositeindex/datacube/startree/builder/AbstractStarTreeBuilderTests.java
@@ -4264,9 +4264,15 @@ public void testFlushFlowWithTimestamps() throws IOException {
 
         compositeField = getStarTreeFieldWithDateDimension();
         SortedNumericStarTreeValuesIterator d1sndv = new SortedNumericStarTreeValuesIterator(getSortedNumericMock(dimList, docsWithField));
-        SortedNumericStarTreeValuesIterator d2sndv = new SortedNumericStarTreeValuesIterator(getSortedNumericMock(dimList2, docsWithField2));
-        SortedNumericStarTreeValuesIterator m1sndv = new SortedNumericStarTreeValuesIterator(getSortedNumericMock(metricsList, metricsWithField));
-        SortedNumericStarTreeValuesIterator m2sndv = new SortedNumericStarTreeValuesIterator(getSortedNumericMock(metricsList, metricsWithField));
+        SortedNumericStarTreeValuesIterator d2sndv = new SortedNumericStarTreeValuesIterator(
+            getSortedNumericMock(dimList2, docsWithField2)
+        );
+        SortedNumericStarTreeValuesIterator m1sndv = new SortedNumericStarTreeValuesIterator(
+            getSortedNumericMock(metricsList, metricsWithField)
+        );
+        SortedNumericStarTreeValuesIterator m2sndv = new SortedNumericStarTreeValuesIterator(
+            getSortedNumericMock(metricsList, metricsWithField)
+        );
         this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec(
             writeState,
             Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC,