From 146b0f7264b1426d475e757729d213097b776212 Mon Sep 17 00:00:00 2001 From: Andriy Redko Date: Mon, 7 Oct 2024 08:26:43 -0400 Subject: [PATCH 1/7] Update Apache Lucene to 9.12.0 (#15333) * Update Apache Lucene to 9.12.0 Signed-off-by: Andriy Redko * change to IOContext READONCE in locations where the file is not expected to be read multiple times. Signed-off-by: Marc Handalian * Use READ IOContext for all non Segment* files when copying node-node Signed-off-by: Marc Handalian * Fixing more test failures Signed-off-by: Andriy Redko * Move Composite912Codec under org.opensearch.index.codec.composite.composite912 package Signed-off-by: Andriy Redko --------- Signed-off-by: Andriy Redko Signed-off-by: Marc Handalian Co-authored-by: Marc Handalian --- CHANGELOG.md | 1 + buildSrc/version.properties | 2 +- ...cene-core-9.12.0-snapshot-847316d.jar.sha1 | 1 - .../core/licenses/lucene-core-9.12.0.jar.sha1 | 1 + ...pressions-9.12.0-snapshot-847316d.jar.sha1 | 1 - .../lucene-expressions-9.12.0.jar.sha1 | 1 + ...lysis-icu-9.12.0-snapshot-847316d.jar.sha1 | 1 - .../lucene-analysis-icu-9.12.0.jar.sha1 | 1 + ...-kuromoji-9.12.0-snapshot-847316d.jar.sha1 | 1 - .../lucene-analysis-kuromoji-9.12.0.jar.sha1 | 1 + ...ysis-nori-9.12.0-snapshot-847316d.jar.sha1 | 1 - .../lucene-analysis-nori-9.12.0.jar.sha1 | 1 + ...-phonetic-9.12.0-snapshot-847316d.jar.sha1 | 1 - .../lucene-analysis-phonetic-9.12.0.jar.sha1 | 1 + ...s-smartcn-9.12.0-snapshot-847316d.jar.sha1 | 1 - .../lucene-analysis-smartcn-9.12.0.jar.sha1 | 1 + ...s-stempel-9.12.0-snapshot-847316d.jar.sha1 | 1 - .../lucene-analysis-stempel-9.12.0.jar.sha1 | 1 + ...orfologik-9.12.0-snapshot-847316d.jar.sha1 | 1 - ...lucene-analysis-morfologik-9.12.0.jar.sha1 | 1 + .../index/codec/CorrelationCodecVersion.java | 2 +- ...is-common-9.12.0-snapshot-847316d.jar.sha1 | 1 - .../lucene-analysis-common-9.12.0.jar.sha1 | 1 + ...rd-codecs-9.12.0-snapshot-847316d.jar.sha1 | 1 - .../lucene-backward-codecs-9.12.0.jar.sha1 | 1 + ...cene-core-9.12.0-snapshot-847316d.jar.sha1 | 1 - server/licenses/lucene-core-9.12.0.jar.sha1 | 1 + ...-grouping-9.12.0-snapshot-847316d.jar.sha1 | 1 - .../licenses/lucene-grouping-9.12.0.jar.sha1 | 1 + ...ghlighter-9.12.0-snapshot-847316d.jar.sha1 | 1 - .../lucene-highlighter-9.12.0.jar.sha1 | 1 + ...cene-join-9.12.0-snapshot-847316d.jar.sha1 | 1 - server/licenses/lucene-join-9.12.0.jar.sha1 | 1 + ...ne-memory-9.12.0-snapshot-847316d.jar.sha1 | 1 - server/licenses/lucene-memory-9.12.0.jar.sha1 | 1 + ...cene-misc-9.12.0-snapshot-847316d.jar.sha1 | 1 - server/licenses/lucene-misc-9.12.0.jar.sha1 | 1 + ...e-queries-9.12.0-snapshot-847316d.jar.sha1 | 1 - .../licenses/lucene-queries-9.12.0.jar.sha1 | 1 + ...eryparser-9.12.0-snapshot-847316d.jar.sha1 | 1 - .../lucene-queryparser-9.12.0.jar.sha1 | 1 + ...e-sandbox-9.12.0-snapshot-847316d.jar.sha1 | 1 - .../licenses/lucene-sandbox-9.12.0.jar.sha1 | 1 + ...al-extras-9.12.0-snapshot-847316d.jar.sha1 | 1 - .../lucene-spatial-extras-9.12.0.jar.sha1 | 1 + ...spatial3d-9.12.0-snapshot-847316d.jar.sha1 | 1 - .../licenses/lucene-spatial3d-9.12.0.jar.sha1 | 1 + ...e-suggest-9.12.0-snapshot-847316d.jar.sha1 | 1 - .../licenses/lucene-suggest-9.12.0.jar.sha1 | 1 + .../node/tasks/ConcurrentSearchTasksIT.java | 2 +- .../recovery/TruncatedRecoveryIT.java | 3 +- .../similarity/LegacyBM25Similarity.java | 10 +--- .../apache/lucene/util/CombinedBitSet.java | 26 ++++++--- .../org/opensearch/common/lucene/Lucene.java | 4 +- .../opensearch/index/codec/CodecService.java | 12 ++-- .../PerFieldMappingPostingFormatCodec.java | 4 +- .../composite/CompositeCodecFactory.java | 11 ++-- .../composite912/Composite912Codec.java | 58 +++++++++++++++++++ .../composite/composite912/package-info.java | 12 ++++ .../composite99/Composite99Codec.java | 8 +-- .../index/mapper/CompletionFieldMapper.java | 4 +- .../opensearch/index/shard/IndexShard.java | 4 +- .../shard/RemoteStoreRefreshListener.java | 4 +- .../index/similarity/SimilarityProviders.java | 3 +- .../index/store/RemoteDirectory.java | 2 +- .../store/RemoteSegmentStoreDirectory.java | 2 +- .../org/opensearch/index/store/Store.java | 8 +-- .../SegmentFileTransferHandler.java | 39 ++++++++++--- .../replication/SegmentReplicationTarget.java | 2 +- .../blobstore/BlobStoreRepository.java | 2 +- .../search/internal/ContextIndexSearcher.java | 3 +- .../services/org.apache.lucene.codecs.Codec | 1 + .../opensearch/index/codec/CodecTests.java | 40 ++++++------- .../LuceneDocValuesConsumerFactoryTests.java | 6 +- .../LuceneDocValuesProducerFactoryTests.java | 6 +- .../StarTreeDocValuesFormatTests.java | 8 +-- .../builder/AbstractStarTreeBuilderTests.java | 12 ++-- .../builder/BaseStarTreeBuilderTests.java | 6 +- .../builder/StarTreesBuilderTests.java | 6 +- .../meta/StarTreeMetadataTests.java | 6 +- .../engine/CompletionStatsCacheTests.java | 8 +-- .../RemoteIndexShardCorruptionTests.java | 2 +- .../opensearch/index/store/StoreTests.java | 2 +- .../org/opensearch/test/CorruptionUtils.java | 2 +- 84 files changed, 231 insertions(+), 148 deletions(-) delete mode 100644 libs/core/licenses/lucene-core-9.12.0-snapshot-847316d.jar.sha1 create mode 100644 libs/core/licenses/lucene-core-9.12.0.jar.sha1 delete mode 100644 modules/lang-expression/licenses/lucene-expressions-9.12.0-snapshot-847316d.jar.sha1 create mode 100644 modules/lang-expression/licenses/lucene-expressions-9.12.0.jar.sha1 delete mode 100644 plugins/analysis-icu/licenses/lucene-analysis-icu-9.12.0-snapshot-847316d.jar.sha1 create mode 100644 plugins/analysis-icu/licenses/lucene-analysis-icu-9.12.0.jar.sha1 delete mode 100644 plugins/analysis-kuromoji/licenses/lucene-analysis-kuromoji-9.12.0-snapshot-847316d.jar.sha1 create mode 100644 plugins/analysis-kuromoji/licenses/lucene-analysis-kuromoji-9.12.0.jar.sha1 delete mode 100644 plugins/analysis-nori/licenses/lucene-analysis-nori-9.12.0-snapshot-847316d.jar.sha1 create mode 100644 plugins/analysis-nori/licenses/lucene-analysis-nori-9.12.0.jar.sha1 delete mode 100644 plugins/analysis-phonetic/licenses/lucene-analysis-phonetic-9.12.0-snapshot-847316d.jar.sha1 create mode 100644 plugins/analysis-phonetic/licenses/lucene-analysis-phonetic-9.12.0.jar.sha1 delete mode 100644 plugins/analysis-smartcn/licenses/lucene-analysis-smartcn-9.12.0-snapshot-847316d.jar.sha1 create mode 100644 plugins/analysis-smartcn/licenses/lucene-analysis-smartcn-9.12.0.jar.sha1 delete mode 100644 plugins/analysis-stempel/licenses/lucene-analysis-stempel-9.12.0-snapshot-847316d.jar.sha1 create mode 100644 plugins/analysis-stempel/licenses/lucene-analysis-stempel-9.12.0.jar.sha1 delete mode 100644 plugins/analysis-ukrainian/licenses/lucene-analysis-morfologik-9.12.0-snapshot-847316d.jar.sha1 create mode 100644 plugins/analysis-ukrainian/licenses/lucene-analysis-morfologik-9.12.0.jar.sha1 delete mode 100644 server/licenses/lucene-analysis-common-9.12.0-snapshot-847316d.jar.sha1 create mode 100644 server/licenses/lucene-analysis-common-9.12.0.jar.sha1 delete mode 100644 server/licenses/lucene-backward-codecs-9.12.0-snapshot-847316d.jar.sha1 create mode 100644 server/licenses/lucene-backward-codecs-9.12.0.jar.sha1 delete mode 100644 server/licenses/lucene-core-9.12.0-snapshot-847316d.jar.sha1 create mode 100644 server/licenses/lucene-core-9.12.0.jar.sha1 delete mode 100644 server/licenses/lucene-grouping-9.12.0-snapshot-847316d.jar.sha1 create mode 100644 server/licenses/lucene-grouping-9.12.0.jar.sha1 delete mode 100644 server/licenses/lucene-highlighter-9.12.0-snapshot-847316d.jar.sha1 create mode 100644 server/licenses/lucene-highlighter-9.12.0.jar.sha1 delete mode 100644 server/licenses/lucene-join-9.12.0-snapshot-847316d.jar.sha1 create mode 100644 server/licenses/lucene-join-9.12.0.jar.sha1 delete mode 100644 server/licenses/lucene-memory-9.12.0-snapshot-847316d.jar.sha1 create mode 100644 server/licenses/lucene-memory-9.12.0.jar.sha1 delete mode 100644 server/licenses/lucene-misc-9.12.0-snapshot-847316d.jar.sha1 create mode 100644 server/licenses/lucene-misc-9.12.0.jar.sha1 delete mode 100644 server/licenses/lucene-queries-9.12.0-snapshot-847316d.jar.sha1 create mode 100644 server/licenses/lucene-queries-9.12.0.jar.sha1 delete mode 100644 server/licenses/lucene-queryparser-9.12.0-snapshot-847316d.jar.sha1 create mode 100644 server/licenses/lucene-queryparser-9.12.0.jar.sha1 delete mode 100644 server/licenses/lucene-sandbox-9.12.0-snapshot-847316d.jar.sha1 create mode 100644 server/licenses/lucene-sandbox-9.12.0.jar.sha1 delete mode 100644 server/licenses/lucene-spatial-extras-9.12.0-snapshot-847316d.jar.sha1 create mode 100644 server/licenses/lucene-spatial-extras-9.12.0.jar.sha1 delete mode 100644 server/licenses/lucene-spatial3d-9.12.0-snapshot-847316d.jar.sha1 create mode 100644 server/licenses/lucene-spatial3d-9.12.0.jar.sha1 delete mode 100644 server/licenses/lucene-suggest-9.12.0-snapshot-847316d.jar.sha1 create mode 100644 server/licenses/lucene-suggest-9.12.0.jar.sha1 create mode 100644 server/src/main/java/org/opensearch/index/codec/composite/composite912/Composite912Codec.java create mode 100644 server/src/main/java/org/opensearch/index/codec/composite/composite912/package-info.java rename server/src/test/java/org/opensearch/index/codec/{composite99 => composite912}/datacube/startree/StarTreeDocValuesFormatTests.java (97%) diff --git a/CHANGELOG.md b/CHANGELOG.md index cbf19fb0d35be..2e3a515fb8008 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -35,6 +35,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Bump `dnsjava:dnsjava` from 3.6.1 to 3.6.2 ([#16041](https://github.com/opensearch-project/OpenSearch/pull/16041)) - Bump `com.maxmind.geoip2:geoip2` from 4.2.0 to 4.2.1 ([#16042](https://github.com/opensearch-project/OpenSearch/pull/16042)) - Bump `com.maxmind.db:maxmind-db` from 3.1.0 to 3.1.1 ([#16137](https://github.com/opensearch-project/OpenSearch/pull/16137)) +- Bump Apache lucene from 9.11.1 to 9.12.0 ([#15333](https://github.com/opensearch-project/OpenSearch/pull/15333)) - Bump `com.azure:azure-core-http-netty` from 1.15.3 to 1.15.4 ([#16133](https://github.com/opensearch-project/OpenSearch/pull/16133)) - Bump `org.jline:jline` from 3.26.3 to 3.27.0 ([#16135](https://github.com/opensearch-project/OpenSearch/pull/16135)) - Bump `netty` from 4.1.112.Final to 4.1.114.Final ([#16182](https://github.com/opensearch-project/OpenSearch/pull/16182)) diff --git a/buildSrc/version.properties b/buildSrc/version.properties index 8cede6b717883..19e2ba2e2c5bb 100644 --- a/buildSrc/version.properties +++ b/buildSrc/version.properties @@ -1,5 +1,5 @@ opensearch = 3.0.0 -lucene = 9.12.0-snapshot-847316d +lucene = 9.12.0 bundled_jdk_vendor = adoptium bundled_jdk = 21.0.4+7 diff --git a/libs/core/licenses/lucene-core-9.12.0-snapshot-847316d.jar.sha1 b/libs/core/licenses/lucene-core-9.12.0-snapshot-847316d.jar.sha1 deleted file mode 100644 index e3fd1708ea428..0000000000000 --- a/libs/core/licenses/lucene-core-9.12.0-snapshot-847316d.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -51ff4940eb1024184bbaa5dae39695d2392c5bab \ No newline at end of file diff --git a/libs/core/licenses/lucene-core-9.12.0.jar.sha1 b/libs/core/licenses/lucene-core-9.12.0.jar.sha1 new file mode 100644 index 0000000000000..e55f896dedb63 --- /dev/null +++ b/libs/core/licenses/lucene-core-9.12.0.jar.sha1 @@ -0,0 +1 @@ +fdb055d569bb20bfce9618fe2b01c29bab7f290c \ No newline at end of file diff --git a/modules/lang-expression/licenses/lucene-expressions-9.12.0-snapshot-847316d.jar.sha1 b/modules/lang-expression/licenses/lucene-expressions-9.12.0-snapshot-847316d.jar.sha1 deleted file mode 100644 index 83dd8e657bdd5..0000000000000 --- a/modules/lang-expression/licenses/lucene-expressions-9.12.0-snapshot-847316d.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -b866103bbaca4141c152deca9252bd137026dafc \ No newline at end of file diff --git a/modules/lang-expression/licenses/lucene-expressions-9.12.0.jar.sha1 b/modules/lang-expression/licenses/lucene-expressions-9.12.0.jar.sha1 new file mode 100644 index 0000000000000..476049a66cc08 --- /dev/null +++ b/modules/lang-expression/licenses/lucene-expressions-9.12.0.jar.sha1 @@ -0,0 +1 @@ +5ba843374a0aab3dfe0b11cb28b251844d85bf5b \ No newline at end of file diff --git a/plugins/analysis-icu/licenses/lucene-analysis-icu-9.12.0-snapshot-847316d.jar.sha1 b/plugins/analysis-icu/licenses/lucene-analysis-icu-9.12.0-snapshot-847316d.jar.sha1 deleted file mode 100644 index 80e254ed3d098..0000000000000 --- a/plugins/analysis-icu/licenses/lucene-analysis-icu-9.12.0-snapshot-847316d.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -04436942995a4952ce5654126dfb767d6335674e \ No newline at end of file diff --git a/plugins/analysis-icu/licenses/lucene-analysis-icu-9.12.0.jar.sha1 b/plugins/analysis-icu/licenses/lucene-analysis-icu-9.12.0.jar.sha1 new file mode 100644 index 0000000000000..31398b27708a3 --- /dev/null +++ b/plugins/analysis-icu/licenses/lucene-analysis-icu-9.12.0.jar.sha1 @@ -0,0 +1 @@ +a9232b6a4882979118d3281b98dfdb6e0e1cb5ca \ No newline at end of file diff --git a/plugins/analysis-kuromoji/licenses/lucene-analysis-kuromoji-9.12.0-snapshot-847316d.jar.sha1 b/plugins/analysis-kuromoji/licenses/lucene-analysis-kuromoji-9.12.0-snapshot-847316d.jar.sha1 deleted file mode 100644 index 3baed2a6e660b..0000000000000 --- a/plugins/analysis-kuromoji/licenses/lucene-analysis-kuromoji-9.12.0-snapshot-847316d.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -85918e24fc3bf63fcd953807ab2eb3fa55c987c2 \ No newline at end of file diff --git a/plugins/analysis-kuromoji/licenses/lucene-analysis-kuromoji-9.12.0.jar.sha1 b/plugins/analysis-kuromoji/licenses/lucene-analysis-kuromoji-9.12.0.jar.sha1 new file mode 100644 index 0000000000000..fa4c9d2d09d6e --- /dev/null +++ b/plugins/analysis-kuromoji/licenses/lucene-analysis-kuromoji-9.12.0.jar.sha1 @@ -0,0 +1 @@ +a3a6950ffc22e76a082e1b3cefb022b9f7870d29 \ No newline at end of file diff --git a/plugins/analysis-nori/licenses/lucene-analysis-nori-9.12.0-snapshot-847316d.jar.sha1 b/plugins/analysis-nori/licenses/lucene-analysis-nori-9.12.0-snapshot-847316d.jar.sha1 deleted file mode 100644 index 4e9327112d412..0000000000000 --- a/plugins/analysis-nori/licenses/lucene-analysis-nori-9.12.0-snapshot-847316d.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -15e425e9cc0ab9d65fac3c919199a24dfa3631eb \ No newline at end of file diff --git a/plugins/analysis-nori/licenses/lucene-analysis-nori-9.12.0.jar.sha1 b/plugins/analysis-nori/licenses/lucene-analysis-nori-9.12.0.jar.sha1 new file mode 100644 index 0000000000000..576b924286d2d --- /dev/null +++ b/plugins/analysis-nori/licenses/lucene-analysis-nori-9.12.0.jar.sha1 @@ -0,0 +1 @@ +e71f85b72ed3939039ba8897b28b065dd11918b9 \ No newline at end of file diff --git a/plugins/analysis-phonetic/licenses/lucene-analysis-phonetic-9.12.0-snapshot-847316d.jar.sha1 b/plugins/analysis-phonetic/licenses/lucene-analysis-phonetic-9.12.0-snapshot-847316d.jar.sha1 deleted file mode 100644 index 7e7e9fe5b22b4..0000000000000 --- a/plugins/analysis-phonetic/licenses/lucene-analysis-phonetic-9.12.0-snapshot-847316d.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -3d16c18348e7d4a00cb83100c43f3e21239d224e \ No newline at end of file diff --git a/plugins/analysis-phonetic/licenses/lucene-analysis-phonetic-9.12.0.jar.sha1 b/plugins/analysis-phonetic/licenses/lucene-analysis-phonetic-9.12.0.jar.sha1 new file mode 100644 index 0000000000000..c8c146bbd0d25 --- /dev/null +++ b/plugins/analysis-phonetic/licenses/lucene-analysis-phonetic-9.12.0.jar.sha1 @@ -0,0 +1 @@ +6baa3ae7ab20d6e644cf0bedb271c50a44c0e259 \ No newline at end of file diff --git a/plugins/analysis-smartcn/licenses/lucene-analysis-smartcn-9.12.0-snapshot-847316d.jar.sha1 b/plugins/analysis-smartcn/licenses/lucene-analysis-smartcn-9.12.0-snapshot-847316d.jar.sha1 deleted file mode 100644 index 98e0ecc9cbb89..0000000000000 --- a/plugins/analysis-smartcn/licenses/lucene-analysis-smartcn-9.12.0-snapshot-847316d.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -2ef6d9dffc6816d3cd04a54fe1ee43e13f850a37 \ No newline at end of file diff --git a/plugins/analysis-smartcn/licenses/lucene-analysis-smartcn-9.12.0.jar.sha1 b/plugins/analysis-smartcn/licenses/lucene-analysis-smartcn-9.12.0.jar.sha1 new file mode 100644 index 0000000000000..54ea0b19f2a7b --- /dev/null +++ b/plugins/analysis-smartcn/licenses/lucene-analysis-smartcn-9.12.0.jar.sha1 @@ -0,0 +1 @@ +f183e1e8b1eaaa4dec444774a285bb8b66518522 \ No newline at end of file diff --git a/plugins/analysis-stempel/licenses/lucene-analysis-stempel-9.12.0-snapshot-847316d.jar.sha1 b/plugins/analysis-stempel/licenses/lucene-analysis-stempel-9.12.0-snapshot-847316d.jar.sha1 deleted file mode 100644 index ef675f2b9702e..0000000000000 --- a/plugins/analysis-stempel/licenses/lucene-analysis-stempel-9.12.0-snapshot-847316d.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -e72b2262f5393d9ff255fb901297d4e6790e9102 \ No newline at end of file diff --git a/plugins/analysis-stempel/licenses/lucene-analysis-stempel-9.12.0.jar.sha1 b/plugins/analysis-stempel/licenses/lucene-analysis-stempel-9.12.0.jar.sha1 new file mode 100644 index 0000000000000..5442a40f5bba2 --- /dev/null +++ b/plugins/analysis-stempel/licenses/lucene-analysis-stempel-9.12.0.jar.sha1 @@ -0,0 +1 @@ +b81a609934e65d12ab9d2d84bc2ea6f56a360e57 \ No newline at end of file diff --git a/plugins/analysis-ukrainian/licenses/lucene-analysis-morfologik-9.12.0-snapshot-847316d.jar.sha1 b/plugins/analysis-ukrainian/licenses/lucene-analysis-morfologik-9.12.0-snapshot-847316d.jar.sha1 deleted file mode 100644 index d8bbac27fd360..0000000000000 --- a/plugins/analysis-ukrainian/licenses/lucene-analysis-morfologik-9.12.0-snapshot-847316d.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -416ac44b2e76592c9e85338798cae93c3cf5475e \ No newline at end of file diff --git a/plugins/analysis-ukrainian/licenses/lucene-analysis-morfologik-9.12.0.jar.sha1 b/plugins/analysis-ukrainian/licenses/lucene-analysis-morfologik-9.12.0.jar.sha1 new file mode 100644 index 0000000000000..60fd4015cfde0 --- /dev/null +++ b/plugins/analysis-ukrainian/licenses/lucene-analysis-morfologik-9.12.0.jar.sha1 @@ -0,0 +1 @@ +bec069f286b45f20b743c81e84202369cd0467e7 \ No newline at end of file diff --git a/plugins/events-correlation-engine/src/main/java/org/opensearch/plugin/correlation/core/index/codec/CorrelationCodecVersion.java b/plugins/events-correlation-engine/src/main/java/org/opensearch/plugin/correlation/core/index/codec/CorrelationCodecVersion.java index 3fcc995fb4199..9dbb695f14b78 100644 --- a/plugins/events-correlation-engine/src/main/java/org/opensearch/plugin/correlation/core/index/codec/CorrelationCodecVersion.java +++ b/plugins/events-correlation-engine/src/main/java/org/opensearch/plugin/correlation/core/index/codec/CorrelationCodecVersion.java @@ -8,8 +8,8 @@ package org.opensearch.plugin.correlation.core.index.codec; +import org.apache.lucene.backward_codecs.lucene99.Lucene99Codec; import org.apache.lucene.codecs.Codec; -import org.apache.lucene.codecs.lucene99.Lucene99Codec; import org.opensearch.index.mapper.MapperService; import org.opensearch.plugin.correlation.core.index.codec.correlation990.CorrelationCodec; import org.opensearch.plugin.correlation.core.index.codec.correlation990.PerFieldCorrelationVectorsFormat; diff --git a/server/licenses/lucene-analysis-common-9.12.0-snapshot-847316d.jar.sha1 b/server/licenses/lucene-analysis-common-9.12.0-snapshot-847316d.jar.sha1 deleted file mode 100644 index f1249066d10f2..0000000000000 --- a/server/licenses/lucene-analysis-common-9.12.0-snapshot-847316d.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -7e282aab7388efc911348f1eacd90e661580dda7 \ No newline at end of file diff --git a/server/licenses/lucene-analysis-common-9.12.0.jar.sha1 b/server/licenses/lucene-analysis-common-9.12.0.jar.sha1 new file mode 100644 index 0000000000000..fd952034f3742 --- /dev/null +++ b/server/licenses/lucene-analysis-common-9.12.0.jar.sha1 @@ -0,0 +1 @@ +4c2503cfaba37249e20ea877555cb52ee89d1ae1 \ No newline at end of file diff --git a/server/licenses/lucene-backward-codecs-9.12.0-snapshot-847316d.jar.sha1 b/server/licenses/lucene-backward-codecs-9.12.0-snapshot-847316d.jar.sha1 deleted file mode 100644 index ac50c5e110a72..0000000000000 --- a/server/licenses/lucene-backward-codecs-9.12.0-snapshot-847316d.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -69e59ba4bed4c58836d2727d72b7f0095d2dcb92 \ No newline at end of file diff --git a/server/licenses/lucene-backward-codecs-9.12.0.jar.sha1 b/server/licenses/lucene-backward-codecs-9.12.0.jar.sha1 new file mode 100644 index 0000000000000..2993134edd610 --- /dev/null +++ b/server/licenses/lucene-backward-codecs-9.12.0.jar.sha1 @@ -0,0 +1 @@ +68fe98c94e9644a584ea1bf525e68d9406fc61ec \ No newline at end of file diff --git a/server/licenses/lucene-core-9.12.0-snapshot-847316d.jar.sha1 b/server/licenses/lucene-core-9.12.0-snapshot-847316d.jar.sha1 deleted file mode 100644 index e3fd1708ea428..0000000000000 --- a/server/licenses/lucene-core-9.12.0-snapshot-847316d.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -51ff4940eb1024184bbaa5dae39695d2392c5bab \ No newline at end of file diff --git a/server/licenses/lucene-core-9.12.0.jar.sha1 b/server/licenses/lucene-core-9.12.0.jar.sha1 new file mode 100644 index 0000000000000..e55f896dedb63 --- /dev/null +++ b/server/licenses/lucene-core-9.12.0.jar.sha1 @@ -0,0 +1 @@ +fdb055d569bb20bfce9618fe2b01c29bab7f290c \ No newline at end of file diff --git a/server/licenses/lucene-grouping-9.12.0-snapshot-847316d.jar.sha1 b/server/licenses/lucene-grouping-9.12.0-snapshot-847316d.jar.sha1 deleted file mode 100644 index cc5bf5bfd8ec0..0000000000000 --- a/server/licenses/lucene-grouping-9.12.0-snapshot-847316d.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -5847a7d47f13ecb7f039fb9adf6f3b8e4bddde77 \ No newline at end of file diff --git a/server/licenses/lucene-grouping-9.12.0.jar.sha1 b/server/licenses/lucene-grouping-9.12.0.jar.sha1 new file mode 100644 index 0000000000000..48388974bb38f --- /dev/null +++ b/server/licenses/lucene-grouping-9.12.0.jar.sha1 @@ -0,0 +1 @@ +ccf99f8db57aa97b2c1f95c5cc2a11156a043921 \ No newline at end of file diff --git a/server/licenses/lucene-highlighter-9.12.0-snapshot-847316d.jar.sha1 b/server/licenses/lucene-highlighter-9.12.0-snapshot-847316d.jar.sha1 deleted file mode 100644 index eb14059d2cd8c..0000000000000 --- a/server/licenses/lucene-highlighter-9.12.0-snapshot-847316d.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -7cc0a26777a479f06fbcfae7abc23e784e1a00dc \ No newline at end of file diff --git a/server/licenses/lucene-highlighter-9.12.0.jar.sha1 b/server/licenses/lucene-highlighter-9.12.0.jar.sha1 new file mode 100644 index 0000000000000..3d457579da892 --- /dev/null +++ b/server/licenses/lucene-highlighter-9.12.0.jar.sha1 @@ -0,0 +1 @@ +e93429f66fbcd3b58d81f01223d6ce5688047296 \ No newline at end of file diff --git a/server/licenses/lucene-join-9.12.0-snapshot-847316d.jar.sha1 b/server/licenses/lucene-join-9.12.0-snapshot-847316d.jar.sha1 deleted file mode 100644 index b87170c39c78c..0000000000000 --- a/server/licenses/lucene-join-9.12.0-snapshot-847316d.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -9cd99401c826d910da3c2beab8e42f1af8be6ea4 \ No newline at end of file diff --git a/server/licenses/lucene-join-9.12.0.jar.sha1 b/server/licenses/lucene-join-9.12.0.jar.sha1 new file mode 100644 index 0000000000000..c5f6d16598a60 --- /dev/null +++ b/server/licenses/lucene-join-9.12.0.jar.sha1 @@ -0,0 +1 @@ +14c802d6955eaf11107375a2ada8fe8ec53b3e01 \ No newline at end of file diff --git a/server/licenses/lucene-memory-9.12.0-snapshot-847316d.jar.sha1 b/server/licenses/lucene-memory-9.12.0-snapshot-847316d.jar.sha1 deleted file mode 100644 index de591dd659cb5..0000000000000 --- a/server/licenses/lucene-memory-9.12.0-snapshot-847316d.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -cfee136ecbc3df7adc38b38e020dca5e61c22773 \ No newline at end of file diff --git a/server/licenses/lucene-memory-9.12.0.jar.sha1 b/server/licenses/lucene-memory-9.12.0.jar.sha1 new file mode 100644 index 0000000000000..e7ac44089c006 --- /dev/null +++ b/server/licenses/lucene-memory-9.12.0.jar.sha1 @@ -0,0 +1 @@ +ffe090339540876b40df792aee51a42af6b3f37f \ No newline at end of file diff --git a/server/licenses/lucene-misc-9.12.0-snapshot-847316d.jar.sha1 b/server/licenses/lucene-misc-9.12.0-snapshot-847316d.jar.sha1 deleted file mode 100644 index 1a999bb9c6686..0000000000000 --- a/server/licenses/lucene-misc-9.12.0-snapshot-847316d.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -afbc5adf93d4eb1a1b109ad828d1968bf16ef292 \ No newline at end of file diff --git a/server/licenses/lucene-misc-9.12.0.jar.sha1 b/server/licenses/lucene-misc-9.12.0.jar.sha1 new file mode 100644 index 0000000000000..afb546be4e032 --- /dev/null +++ b/server/licenses/lucene-misc-9.12.0.jar.sha1 @@ -0,0 +1 @@ +ad17704ee90eb926b6d3105f7027485cdadbecd9 \ No newline at end of file diff --git a/server/licenses/lucene-queries-9.12.0-snapshot-847316d.jar.sha1 b/server/licenses/lucene-queries-9.12.0-snapshot-847316d.jar.sha1 deleted file mode 100644 index 783a26551ae8c..0000000000000 --- a/server/licenses/lucene-queries-9.12.0-snapshot-847316d.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -16907c36f6adb8ba8f260e05738c66afb37c72d3 \ No newline at end of file diff --git a/server/licenses/lucene-queries-9.12.0.jar.sha1 b/server/licenses/lucene-queries-9.12.0.jar.sha1 new file mode 100644 index 0000000000000..e24756e38dad2 --- /dev/null +++ b/server/licenses/lucene-queries-9.12.0.jar.sha1 @@ -0,0 +1 @@ +3ac2a62b0b55c5725bb65f0c5454f9f8a401cf43 \ No newline at end of file diff --git a/server/licenses/lucene-queryparser-9.12.0-snapshot-847316d.jar.sha1 b/server/licenses/lucene-queryparser-9.12.0-snapshot-847316d.jar.sha1 deleted file mode 100644 index b3e9e4de96174..0000000000000 --- a/server/licenses/lucene-queryparser-9.12.0-snapshot-847316d.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -72baa9bddcf2efb71ffb695f1e9f548699ec13a0 \ No newline at end of file diff --git a/server/licenses/lucene-queryparser-9.12.0.jar.sha1 b/server/licenses/lucene-queryparser-9.12.0.jar.sha1 new file mode 100644 index 0000000000000..e93e00a063dd0 --- /dev/null +++ b/server/licenses/lucene-queryparser-9.12.0.jar.sha1 @@ -0,0 +1 @@ +55959399373876f4c184944315458dc6b88fbd81 \ No newline at end of file diff --git a/server/licenses/lucene-sandbox-9.12.0-snapshot-847316d.jar.sha1 b/server/licenses/lucene-sandbox-9.12.0-snapshot-847316d.jar.sha1 deleted file mode 100644 index 2aefa435b1e9a..0000000000000 --- a/server/licenses/lucene-sandbox-9.12.0-snapshot-847316d.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -dd3c63066f583d90b563ebaa6fbe61c603403acb \ No newline at end of file diff --git a/server/licenses/lucene-sandbox-9.12.0.jar.sha1 b/server/licenses/lucene-sandbox-9.12.0.jar.sha1 new file mode 100644 index 0000000000000..a3fd8446e0dbc --- /dev/null +++ b/server/licenses/lucene-sandbox-9.12.0.jar.sha1 @@ -0,0 +1 @@ +f65882536d681c11a1cbc920e5679201101e3603 \ No newline at end of file diff --git a/server/licenses/lucene-spatial-extras-9.12.0-snapshot-847316d.jar.sha1 b/server/licenses/lucene-spatial-extras-9.12.0-snapshot-847316d.jar.sha1 deleted file mode 100644 index d27112c6db6ab..0000000000000 --- a/server/licenses/lucene-spatial-extras-9.12.0-snapshot-847316d.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -69b99530e0b05251c12863bee6a9325cafd5fdaa \ No newline at end of file diff --git a/server/licenses/lucene-spatial-extras-9.12.0.jar.sha1 b/server/licenses/lucene-spatial-extras-9.12.0.jar.sha1 new file mode 100644 index 0000000000000..b0f11fb667faf --- /dev/null +++ b/server/licenses/lucene-spatial-extras-9.12.0.jar.sha1 @@ -0,0 +1 @@ +9d00cc7cc2279822ef6740f0677cafacfb439fa8 \ No newline at end of file diff --git a/server/licenses/lucene-spatial3d-9.12.0-snapshot-847316d.jar.sha1 b/server/licenses/lucene-spatial3d-9.12.0-snapshot-847316d.jar.sha1 deleted file mode 100644 index 29423ac0ababd..0000000000000 --- a/server/licenses/lucene-spatial3d-9.12.0-snapshot-847316d.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -a67d193b4b08790169db7cf005a2429991260287 \ No newline at end of file diff --git a/server/licenses/lucene-spatial3d-9.12.0.jar.sha1 b/server/licenses/lucene-spatial3d-9.12.0.jar.sha1 new file mode 100644 index 0000000000000..858eee25ac191 --- /dev/null +++ b/server/licenses/lucene-spatial3d-9.12.0.jar.sha1 @@ -0,0 +1 @@ +e3092632ca1d4427d3ebb2c866ac89d90f5b61ec \ No newline at end of file diff --git a/server/licenses/lucene-suggest-9.12.0-snapshot-847316d.jar.sha1 b/server/licenses/lucene-suggest-9.12.0-snapshot-847316d.jar.sha1 deleted file mode 100644 index 6ce1f639ccbb7..0000000000000 --- a/server/licenses/lucene-suggest-9.12.0-snapshot-847316d.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -7a1625ae39071ccbfb3af11df5a74291758f4b47 \ No newline at end of file diff --git a/server/licenses/lucene-suggest-9.12.0.jar.sha1 b/server/licenses/lucene-suggest-9.12.0.jar.sha1 new file mode 100644 index 0000000000000..973a7726d845d --- /dev/null +++ b/server/licenses/lucene-suggest-9.12.0.jar.sha1 @@ -0,0 +1 @@ +e1c6636499317ebe498f3490a1ec8b86b8a363dd \ No newline at end of file diff --git a/server/src/internalClusterTest/java/org/opensearch/action/admin/cluster/node/tasks/ConcurrentSearchTasksIT.java b/server/src/internalClusterTest/java/org/opensearch/action/admin/cluster/node/tasks/ConcurrentSearchTasksIT.java index 8b3c40c43e2d2..e2a86b9358829 100644 --- a/server/src/internalClusterTest/java/org/opensearch/action/admin/cluster/node/tasks/ConcurrentSearchTasksIT.java +++ b/server/src/internalClusterTest/java/org/opensearch/action/admin/cluster/node/tasks/ConcurrentSearchTasksIT.java @@ -110,7 +110,7 @@ public void testConcurrentSearchTaskTracking() { Map> threadStats = getThreadStats(SearchAction.NAME + "[*]", taskInfo.getTaskId()); // Concurrent search forks each slice of 5 segments to different thread (see please // https://github.com/apache/lucene/issues/12498) - assertEquals((int) Math.ceil(getSegmentCount(INDEX_NAME) / 5.0) + 1, threadStats.size()); + assertEquals((int) Math.ceil(getSegmentCount(INDEX_NAME) / 5.0), threadStats.size()); // assert that all task descriptions have non-zero length MatcherAssert.assertThat(taskInfo.getDescription().length(), greaterThan(0)); diff --git a/server/src/internalClusterTest/java/org/opensearch/recovery/TruncatedRecoveryIT.java b/server/src/internalClusterTest/java/org/opensearch/recovery/TruncatedRecoveryIT.java index 692beb86279b9..e9515700199ba 100644 --- a/server/src/internalClusterTest/java/org/opensearch/recovery/TruncatedRecoveryIT.java +++ b/server/src/internalClusterTest/java/org/opensearch/recovery/TruncatedRecoveryIT.java @@ -58,6 +58,7 @@ import java.util.Collections; import java.util.List; import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import static org.opensearch.indices.recovery.RecoverySettings.INDICES_RECOVERY_CHUNK_SIZE_SETTING; @@ -184,7 +185,7 @@ public void testCancelRecoveryAndResume() throws Exception { ) .get(); - latch.await(); + latch.await(5, TimeUnit.MINUTES); // at this point we got some truncated left overs on the replica on the unlucky node // now we are allowing the recovery to allocate again and finish to see if we wipe the truncated files diff --git a/server/src/main/java/org/apache/lucene/misc/search/similarity/LegacyBM25Similarity.java b/server/src/main/java/org/apache/lucene/misc/search/similarity/LegacyBM25Similarity.java index 3b812e1c70368..e48a6704ec217 100644 --- a/server/src/main/java/org/apache/lucene/misc/search/similarity/LegacyBM25Similarity.java +++ b/server/src/main/java/org/apache/lucene/misc/search/similarity/LegacyBM25Similarity.java @@ -70,6 +70,7 @@ public LegacyBM25Similarity(float k1, float b) { * within the range {@code [0..1]} */ public LegacyBM25Similarity(float k1, float b, boolean discountOverlaps) { + super(discountOverlaps); this.bm25Similarity = new BM25Similarity(k1, b, discountOverlaps); } @@ -101,15 +102,6 @@ public final float getB() { return bm25Similarity.getB(); } - /** - * Returns true if overlap tokens are discounted from the document's length. - * - * @see #LegacyBM25Similarity(float, float, boolean) - */ - public boolean getDiscountOverlaps() { - return bm25Similarity.getDiscountOverlaps(); - } - @Override public String toString() { return bm25Similarity.toString(); diff --git a/server/src/main/java/org/apache/lucene/util/CombinedBitSet.java b/server/src/main/java/org/apache/lucene/util/CombinedBitSet.java index 1b4f31892c7f8..027a7ed292863 100644 --- a/server/src/main/java/org/apache/lucene/util/CombinedBitSet.java +++ b/server/src/main/java/org/apache/lucene/util/CombinedBitSet.java @@ -87,15 +87,7 @@ public int prevSetBit(int index) { @Override public int nextSetBit(int index) { - assert index >= 0 && index < length : "index=" + index + " numBits=" + length(); - int next = first.nextSetBit(index); - while (next != DocIdSetIterator.NO_MORE_DOCS && second.get(next) == false) { - if (next == length() - 1) { - return DocIdSetIterator.NO_MORE_DOCS; - } - next = first.nextSetBit(next + 1); - } - return next; + return nextSetBit(index, length() - 1); } @Override @@ -132,4 +124,20 @@ public void clear(int startIndex, int endIndex) { public boolean getAndSet(int i) { throw new UnsupportedOperationException("not implemented"); } + + @Override + public int nextSetBit(int start, int end) { + assert start >= 0 && start < length() : "start=" + start + " numBits=" + length(); + assert end >= start && end < length() : "end=" + end + " numBits=" + length(); + + int next = first.nextSetBit(start); + while (next != DocIdSetIterator.NO_MORE_DOCS && second.get(next) == false) { + if (next >= end) { + return DocIdSetIterator.NO_MORE_DOCS; + } + next = first.nextSetBit(next + 1); + } + return next; + + } } diff --git a/server/src/main/java/org/opensearch/common/lucene/Lucene.java b/server/src/main/java/org/opensearch/common/lucene/Lucene.java index 2c7b6b552b43f..361ee859e3e6a 100644 --- a/server/src/main/java/org/opensearch/common/lucene/Lucene.java +++ b/server/src/main/java/org/opensearch/common/lucene/Lucene.java @@ -110,7 +110,7 @@ * @opensearch.internal */ public class Lucene { - public static final String LATEST_CODEC = "Lucene99"; + public static final String LATEST_CODEC = "Lucene912"; public static final String SOFT_DELETES_FIELD = "__soft_deletes"; @@ -272,7 +272,7 @@ public static void checkSegmentInfoIntegrity(final Directory directory) throws I @Override protected Object doBody(String segmentFileName) throws IOException { - try (IndexInput input = directory.openInput(segmentFileName, IOContext.READ)) { + try (IndexInput input = directory.openInput(segmentFileName, IOContext.READONCE)) { CodecUtil.checksumEntireFile(input); } return null; diff --git a/server/src/main/java/org/opensearch/index/codec/CodecService.java b/server/src/main/java/org/opensearch/index/codec/CodecService.java index 59fafdf1ba74e..3a93795ef61ec 100644 --- a/server/src/main/java/org/opensearch/index/codec/CodecService.java +++ b/server/src/main/java/org/opensearch/index/codec/CodecService.java @@ -34,8 +34,8 @@ import org.apache.logging.log4j.Logger; import org.apache.lucene.codecs.Codec; -import org.apache.lucene.codecs.lucene99.Lucene99Codec; -import org.apache.lucene.codecs.lucene99.Lucene99Codec.Mode; +import org.apache.lucene.codecs.lucene912.Lucene912Codec; +import org.apache.lucene.codecs.lucene912.Lucene912Codec.Mode; import org.opensearch.common.Nullable; import org.opensearch.common.collect.MapBuilder; import org.opensearch.index.IndexSettings; @@ -70,10 +70,10 @@ public CodecService(@Nullable MapperService mapperService, IndexSettings indexSe final MapBuilder codecs = MapBuilder.newMapBuilder(); assert null != indexSettings; if (mapperService == null) { - codecs.put(DEFAULT_CODEC, new Lucene99Codec()); - codecs.put(LZ4, new Lucene99Codec()); - codecs.put(BEST_COMPRESSION_CODEC, new Lucene99Codec(Mode.BEST_COMPRESSION)); - codecs.put(ZLIB, new Lucene99Codec(Mode.BEST_COMPRESSION)); + codecs.put(DEFAULT_CODEC, new Lucene912Codec()); + codecs.put(LZ4, new Lucene912Codec()); + codecs.put(BEST_COMPRESSION_CODEC, new Lucene912Codec(Mode.BEST_COMPRESSION)); + codecs.put(ZLIB, new Lucene912Codec(Mode.BEST_COMPRESSION)); } else { // CompositeCodec still delegates to PerFieldMappingPostingFormatCodec // We can still support all the compression codecs when composite index is present diff --git a/server/src/main/java/org/opensearch/index/codec/PerFieldMappingPostingFormatCodec.java b/server/src/main/java/org/opensearch/index/codec/PerFieldMappingPostingFormatCodec.java index 1ad17f121560c..0ed8c3880f2ea 100644 --- a/server/src/main/java/org/opensearch/index/codec/PerFieldMappingPostingFormatCodec.java +++ b/server/src/main/java/org/opensearch/index/codec/PerFieldMappingPostingFormatCodec.java @@ -37,7 +37,7 @@ import org.apache.lucene.codecs.DocValuesFormat; import org.apache.lucene.codecs.PostingsFormat; import org.apache.lucene.codecs.lucene90.Lucene90DocValuesFormat; -import org.apache.lucene.codecs.lucene99.Lucene99Codec; +import org.apache.lucene.codecs.lucene912.Lucene912Codec; import org.opensearch.common.lucene.Lucene; import org.opensearch.index.codec.fuzzy.FuzzyFilterPostingsFormat; import org.opensearch.index.codec.fuzzy.FuzzySetFactory; @@ -59,7 +59,7 @@ * * @opensearch.internal */ -public class PerFieldMappingPostingFormatCodec extends Lucene99Codec { +public class PerFieldMappingPostingFormatCodec extends Lucene912Codec { private final Logger logger; private final MapperService mapperService; private final DocValuesFormat dvFormat = new Lucene90DocValuesFormat(); diff --git a/server/src/main/java/org/opensearch/index/codec/composite/CompositeCodecFactory.java b/server/src/main/java/org/opensearch/index/codec/composite/CompositeCodecFactory.java index 99691d7061ac9..f5003417c1298 100644 --- a/server/src/main/java/org/opensearch/index/codec/composite/CompositeCodecFactory.java +++ b/server/src/main/java/org/opensearch/index/codec/composite/CompositeCodecFactory.java @@ -10,8 +10,9 @@ import org.apache.logging.log4j.Logger; import org.apache.lucene.codecs.Codec; -import org.apache.lucene.codecs.lucene99.Lucene99Codec; +import org.apache.lucene.codecs.lucene912.Lucene912Codec; import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.index.codec.composite.composite912.Composite912Codec; import org.opensearch.index.codec.composite.composite99.Composite99Codec; import org.opensearch.index.mapper.MapperService; @@ -38,10 +39,10 @@ public CompositeCodecFactory() {} public Map getCompositeIndexCodecs(MapperService mapperService, Logger logger) { Map codecs = new HashMap<>(); - codecs.put(DEFAULT_CODEC, new Composite99Codec(Lucene99Codec.Mode.BEST_SPEED, mapperService, logger)); - codecs.put(LZ4, new Composite99Codec(Lucene99Codec.Mode.BEST_SPEED, mapperService, logger)); - codecs.put(BEST_COMPRESSION_CODEC, new Composite99Codec(Lucene99Codec.Mode.BEST_COMPRESSION, mapperService, logger)); - codecs.put(ZLIB, new Composite99Codec(Lucene99Codec.Mode.BEST_COMPRESSION, mapperService, logger)); + codecs.put(DEFAULT_CODEC, new Composite912Codec(Lucene912Codec.Mode.BEST_SPEED, mapperService, logger)); + codecs.put(LZ4, new Composite912Codec(Lucene912Codec.Mode.BEST_SPEED, mapperService, logger)); + codecs.put(BEST_COMPRESSION_CODEC, new Composite912Codec(Lucene912Codec.Mode.BEST_COMPRESSION, mapperService, logger)); + codecs.put(ZLIB, new Composite912Codec(Lucene912Codec.Mode.BEST_COMPRESSION, mapperService, logger)); return codecs; } } diff --git a/server/src/main/java/org/opensearch/index/codec/composite/composite912/Composite912Codec.java b/server/src/main/java/org/opensearch/index/codec/composite/composite912/Composite912Codec.java new file mode 100644 index 0000000000000..9bbff7d344c2f --- /dev/null +++ b/server/src/main/java/org/opensearch/index/codec/composite/composite912/Composite912Codec.java @@ -0,0 +1,58 @@ +/* + * 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.codec.composite.composite912; + +import org.apache.logging.log4j.Logger; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.DocValuesFormat; +import org.apache.lucene.codecs.FilterCodec; +import org.apache.lucene.codecs.lucene912.Lucene912Codec; +import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.index.codec.PerFieldMappingPostingFormatCodec; +import org.opensearch.index.codec.composite.composite99.Composite99DocValuesFormat; +import org.opensearch.index.mapper.MapperService; + +/** + * Extends the Codec to support new file formats for composite indices eg: star tree index + * based on the mappings. + * + * @opensearch.experimental + */ +@ExperimentalApi +public class Composite912Codec extends FilterCodec { + public static final String COMPOSITE_INDEX_CODEC_NAME = "Composite912Codec"; + private final MapperService mapperService; + + // needed for SPI - this is used in reader path + public Composite912Codec() { + this(COMPOSITE_INDEX_CODEC_NAME, new Lucene912Codec(), null); + } + + public Composite912Codec(Lucene912Codec.Mode compressionMode, MapperService mapperService, Logger logger) { + this(COMPOSITE_INDEX_CODEC_NAME, new PerFieldMappingPostingFormatCodec(compressionMode, mapperService, logger), mapperService); + } + + /** + * Sole constructor. When subclassing this codec, create a no-arg ctor and pass the delegate codec and a unique name to + * this ctor. + * + * @param name name of the codec + * @param delegate codec delegate + * @param mapperService mapper service instance + */ + protected Composite912Codec(String name, Codec delegate, MapperService mapperService) { + super(name, delegate); + this.mapperService = mapperService; + } + + @Override + public DocValuesFormat docValuesFormat() { + return new Composite99DocValuesFormat(mapperService); + } +} diff --git a/server/src/main/java/org/opensearch/index/codec/composite/composite912/package-info.java b/server/src/main/java/org/opensearch/index/codec/composite/composite912/package-info.java new file mode 100644 index 0000000000000..84520ed546e60 --- /dev/null +++ b/server/src/main/java/org/opensearch/index/codec/composite/composite912/package-info.java @@ -0,0 +1,12 @@ +/* + * 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. + */ + +/** + * Responsible for handling all composite index codecs and operations associated with Composite912 codec + */ +package org.opensearch.index.codec.composite.composite912; diff --git a/server/src/main/java/org/opensearch/index/codec/composite/composite99/Composite99Codec.java b/server/src/main/java/org/opensearch/index/codec/composite/composite99/Composite99Codec.java index 8422932e937c2..20f1c7e9b7614 100644 --- a/server/src/main/java/org/opensearch/index/codec/composite/composite99/Composite99Codec.java +++ b/server/src/main/java/org/opensearch/index/codec/composite/composite99/Composite99Codec.java @@ -8,13 +8,11 @@ package org.opensearch.index.codec.composite.composite99; -import org.apache.logging.log4j.Logger; +import org.apache.lucene.backward_codecs.lucene99.Lucene99Codec; import org.apache.lucene.codecs.Codec; import org.apache.lucene.codecs.DocValuesFormat; import org.apache.lucene.codecs.FilterCodec; -import org.apache.lucene.codecs.lucene99.Lucene99Codec; import org.opensearch.common.annotation.ExperimentalApi; -import org.opensearch.index.codec.PerFieldMappingPostingFormatCodec; import org.opensearch.index.mapper.MapperService; /** @@ -33,10 +31,6 @@ public Composite99Codec() { this(COMPOSITE_INDEX_CODEC_NAME, new Lucene99Codec(), null); } - public Composite99Codec(Lucene99Codec.Mode compressionMode, MapperService mapperService, Logger logger) { - this(COMPOSITE_INDEX_CODEC_NAME, new PerFieldMappingPostingFormatCodec(compressionMode, mapperService, logger), mapperService); - } - /** * Sole constructor. When subclassing this codec, create a no-arg ctor and pass the delegate codec and a unique name to * this ctor. diff --git a/server/src/main/java/org/opensearch/index/mapper/CompletionFieldMapper.java b/server/src/main/java/org/opensearch/index/mapper/CompletionFieldMapper.java index a9d9f6cb35fcb..fe48adf3249a3 100644 --- a/server/src/main/java/org/opensearch/index/mapper/CompletionFieldMapper.java +++ b/server/src/main/java/org/opensearch/index/mapper/CompletionFieldMapper.java @@ -35,7 +35,7 @@ import org.apache.lucene.document.FieldType; import org.apache.lucene.index.IndexOptions; import org.apache.lucene.index.Term; -import org.apache.lucene.search.suggest.document.Completion99PostingsFormat; +import org.apache.lucene.search.suggest.document.Completion912PostingsFormat; import org.apache.lucene.search.suggest.document.CompletionAnalyzer; import org.apache.lucene.search.suggest.document.CompletionQuery; import org.apache.lucene.search.suggest.document.FuzzyCompletionQuery; @@ -330,7 +330,7 @@ public ContextMappings getContextMappings() { */ public static synchronized PostingsFormat postingsFormat() { if (postingsFormat == null) { - postingsFormat = new Completion99PostingsFormat(); + postingsFormat = new Completion912PostingsFormat(); } return postingsFormat; } diff --git a/server/src/main/java/org/opensearch/index/shard/IndexShard.java b/server/src/main/java/org/opensearch/index/shard/IndexShard.java index 6e12e4ed3da1a..eb3999718ca5b 100644 --- a/server/src/main/java/org/opensearch/index/shard/IndexShard.java +++ b/server/src/main/java/org/opensearch/index/shard/IndexShard.java @@ -5214,7 +5214,7 @@ public void syncSegmentsFromGivenRemoteSegmentStore( } else if (segmentsNFile != null) { try ( ChecksumIndexInput indexInput = new BufferedChecksumIndexInput( - storeDirectory.openInput(segmentsNFile, IOContext.DEFAULT) + storeDirectory.openInput(segmentsNFile, IOContext.READONCE) ) ) { long commitGeneration = SegmentInfos.generationFromSegmentsFileName(segmentsNFile); @@ -5288,7 +5288,7 @@ private String copySegmentFiles( // Visible for testing boolean localDirectoryContains(Directory localDirectory, String file, long checksum) throws IOException { - try (IndexInput indexInput = localDirectory.openInput(file, IOContext.DEFAULT)) { + try (IndexInput indexInput = localDirectory.openInput(file, IOContext.READONCE)) { if (checksum == CodecUtil.retrieveChecksum(indexInput)) { return true; } else { diff --git a/server/src/main/java/org/opensearch/index/shard/RemoteStoreRefreshListener.java b/server/src/main/java/org/opensearch/index/shard/RemoteStoreRefreshListener.java index 3db5e0b7f4b0c..702928771f131 100644 --- a/server/src/main/java/org/opensearch/index/shard/RemoteStoreRefreshListener.java +++ b/server/src/main/java/org/opensearch/index/shard/RemoteStoreRefreshListener.java @@ -459,7 +459,7 @@ private void uploadNewSegments( batchUploadListener.onFailure(ex); }); statsListener.beforeUpload(src); - remoteDirectory.copyFrom(storeDirectory, src, IOContext.DEFAULT, aggregatedListener, isLowPriorityUpload()); + remoteDirectory.copyFrom(storeDirectory, src, IOContext.READONCE, aggregatedListener, isLowPriorityUpload()); } } @@ -488,7 +488,7 @@ private boolean skipUpload(String file) { private String getChecksumOfLocalFile(String file) throws IOException { if (!localSegmentChecksumMap.containsKey(file)) { - try (IndexInput indexInput = storeDirectory.openInput(file, IOContext.DEFAULT)) { + try (IndexInput indexInput = storeDirectory.openInput(file, IOContext.READONCE)) { String checksum = Long.toString(CodecUtil.retrieveChecksum(indexInput)); localSegmentChecksumMap.put(file, checksum); } diff --git a/server/src/main/java/org/opensearch/index/similarity/SimilarityProviders.java b/server/src/main/java/org/opensearch/index/similarity/SimilarityProviders.java index 139b8fffbac3a..b485b827ca3af 100644 --- a/server/src/main/java/org/opensearch/index/similarity/SimilarityProviders.java +++ b/server/src/main/java/org/opensearch/index/similarity/SimilarityProviders.java @@ -291,8 +291,7 @@ public static ClassicSimilarity createClassicSimilarity(Settings settings, Versi boolean discountOverlaps = settings.getAsBoolean(DISCOUNT_OVERLAPS, true); - ClassicSimilarity similarity = new ClassicSimilarity(); - similarity.setDiscountOverlaps(discountOverlaps); + ClassicSimilarity similarity = new ClassicSimilarity(discountOverlaps); return similarity; } diff --git a/server/src/main/java/org/opensearch/index/store/RemoteDirectory.java b/server/src/main/java/org/opensearch/index/store/RemoteDirectory.java index c61fae74c0bc0..aa856aa22313d 100644 --- a/server/src/main/java/org/opensearch/index/store/RemoteDirectory.java +++ b/server/src/main/java/org/opensearch/index/store/RemoteDirectory.java @@ -424,7 +424,7 @@ private void uploadBlob( } private long calculateChecksumOfChecksum(Directory directory, String file) throws IOException { - try (IndexInput indexInput = directory.openInput(file, IOContext.DEFAULT)) { + try (IndexInput indexInput = directory.openInput(file, IOContext.READONCE)) { try { return checksumOfChecksum(indexInput, SEGMENT_CHECKSUM_BYTES); } catch (Exception e) { diff --git a/server/src/main/java/org/opensearch/index/store/RemoteSegmentStoreDirectory.java b/server/src/main/java/org/opensearch/index/store/RemoteSegmentStoreDirectory.java index 5be516166803e..25d6e24cade14 100644 --- a/server/src/main/java/org/opensearch/index/store/RemoteSegmentStoreDirectory.java +++ b/server/src/main/java/org/opensearch/index/store/RemoteSegmentStoreDirectory.java @@ -759,7 +759,7 @@ private void tryAndDeleteLocalFile(String filename, Directory directory) { } private String getChecksumOfLocalFile(Directory directory, String file) throws IOException { - try (IndexInput indexInput = directory.openInput(file, IOContext.DEFAULT)) { + try (IndexInput indexInput = directory.openInput(file, IOContext.READONCE)) { return Long.toString(CodecUtil.retrieveChecksum(indexInput)); } } diff --git a/server/src/main/java/org/opensearch/index/store/Store.java b/server/src/main/java/org/opensearch/index/store/Store.java index 56fc5b1ffa90d..e7de61ff764f6 100644 --- a/server/src/main/java/org/opensearch/index/store/Store.java +++ b/server/src/main/java/org/opensearch/index/store/Store.java @@ -171,12 +171,6 @@ public class Store extends AbstractIndexShardComponent implements Closeable, Ref Property.IndexScope ); - /** - * Specific {@link IOContext} used to verify Lucene files footer checksums. - * See {@link MetadataSnapshot#checksumFromLuceneFile(Directory, String, Map, Logger, Version, boolean)} - */ - public static final IOContext READONCE_CHECKSUM = new IOContext(IOContext.READONCE.context); - private final AtomicBoolean isClosed = new AtomicBoolean(false); private final StoreDirectory directory; private final ReentrantReadWriteLock metadataLock = new ReentrantReadWriteLock(); @@ -1209,7 +1203,7 @@ private static void checksumFromLuceneFile( ) throws IOException { final String checksum; final BytesRefBuilder fileHash = new BytesRefBuilder(); - try (IndexInput in = directory.openInput(file, READONCE_CHECKSUM)) { + try (IndexInput in = directory.openInput(file, IOContext.READONCE)) { final long length; try { length = in.length(); diff --git a/server/src/main/java/org/opensearch/indices/replication/SegmentFileTransferHandler.java b/server/src/main/java/org/opensearch/indices/replication/SegmentFileTransferHandler.java index a70700b64eb6e..80ba5146535aa 100644 --- a/server/src/main/java/org/opensearch/indices/replication/SegmentFileTransferHandler.java +++ b/server/src/main/java/org/opensearch/indices/replication/SegmentFileTransferHandler.java @@ -11,6 +11,7 @@ import org.apache.logging.log4j.Logger; import org.apache.logging.log4j.message.ParameterizedMessage; import org.apache.lucene.index.CorruptIndexException; +import org.apache.lucene.index.IndexFileNames; import org.apache.lucene.store.IOContext; import org.apache.lucene.store.IndexInput; import org.apache.lucene.util.ArrayUtil; @@ -97,20 +98,26 @@ public MultiChunkTransfer createTransfer( return new MultiChunkTransfer<>(logger, threadPool.getThreadContext(), listener, maxConcurrentFileChunks, Arrays.asList(files)) { final Deque buffers = new ConcurrentLinkedDeque<>(); - InputStreamIndexInput currentInput = null; + volatile InputStreamIndexInput currentInput = null; long offset = 0; @Override protected void onNewResource(StoreFileMetadata md) throws IOException { offset = 0; IOUtils.close(currentInput, () -> currentInput = null); - final IndexInput indexInput = store.directory().openInput(md.name(), IOContext.READONCE); - currentInput = new InputStreamIndexInput(indexInput, md.length()) { - @Override - public void close() throws IOException { - IOUtils.close(indexInput, super::close); // InputStreamIndexInput's close is a noop - } - }; + // Open all files other than Segments* using IOContext.READ. + // With Lucene9_12 a READONCE context will confine the underlying IndexInput (MemorySegmentIndexInput) to a single thread. + // Segments* files require IOContext.READONCE + // https://github.com/apache/lucene/blob/b2d3a2b37e00f19a74949097736be8fd64745f61/lucene/test-framework/src/java/org/apache/lucene/tests/store/MockDirectoryWrapper.java#L817 + if (md.name().startsWith(IndexFileNames.SEGMENTS) == false) { + final IndexInput indexInput = store.directory().openInput(md.name(), IOContext.READ); + currentInput = new InputStreamIndexInput(indexInput, md.length()) { + @Override + public void close() throws IOException { + IOUtils.close(indexInput, super::close); // InputStreamIndexInput's close is a noop + } + }; + } } private byte[] acquireBuffer() { @@ -126,7 +133,7 @@ protected FileChunk nextChunkRequest(StoreFileMetadata md) throws IOException { assert Transports.assertNotTransportThread("read file chunk"); cancellableThreads.checkForCancel(); final byte[] buffer = acquireBuffer(); - final int bytesRead = currentInput.read(buffer); + final int bytesRead = readBytes(md, buffer); if (bytesRead == -1) { throw new CorruptIndexException("file truncated; length=" + md.length() + " offset=" + offset, md.name()); } @@ -142,6 +149,20 @@ protected FileChunk nextChunkRequest(StoreFileMetadata md) throws IOException { return chunk; } + private int readBytes(StoreFileMetadata md, byte[] buffer) throws IOException { + // if we don't have a currentInput by now open once to create the chunk. + if (currentInput == null) { + try (IndexInput indexInput = store.directory().openInput(md.name(), IOContext.READONCE)) { + try (InputStreamIndexInput in = new InputStreamIndexInput(indexInput, md.length())) { + in.skip(offset); + return in.read(buffer); + } + } + } else { + return currentInput.read(buffer); + } + } + @Override protected void executeChunkRequest(FileChunk request, ActionListener listener1) { cancellableThreads.checkForCancel(); diff --git a/server/src/main/java/org/opensearch/indices/replication/SegmentReplicationTarget.java b/server/src/main/java/org/opensearch/indices/replication/SegmentReplicationTarget.java index af764556b7549..7131b49a41834 100644 --- a/server/src/main/java/org/opensearch/indices/replication/SegmentReplicationTarget.java +++ b/server/src/main/java/org/opensearch/indices/replication/SegmentReplicationTarget.java @@ -244,7 +244,7 @@ private List getFiles(CheckpointInfoResponse checkpointInfo) // pkg private for tests private boolean validateLocalChecksum(StoreFileMetadata file) { - try (IndexInput indexInput = indexShard.store().directory().openInput(file.name(), IOContext.DEFAULT)) { + try (IndexInput indexInput = indexShard.store().directory().openInput(file.name(), IOContext.READONCE)) { String checksum = Store.digestToString(CodecUtil.retrieveChecksum(indexInput)); if (file.checksum().equals(checksum)) { return true; diff --git a/server/src/main/java/org/opensearch/repositories/blobstore/BlobStoreRepository.java b/server/src/main/java/org/opensearch/repositories/blobstore/BlobStoreRepository.java index 14c201e819994..8482146530860 100644 --- a/server/src/main/java/org/opensearch/repositories/blobstore/BlobStoreRepository.java +++ b/server/src/main/java/org/opensearch/repositories/blobstore/BlobStoreRepository.java @@ -4530,7 +4530,7 @@ private void snapshotFile( ) throws IOException { final BlobContainer shardContainer = shardContainer(indexId, shardId); final String file = fileInfo.physicalName(); - try (IndexInput indexInput = store.openVerifyingInput(file, IOContext.READONCE, fileInfo.metadata())) { + try (IndexInput indexInput = store.openVerifyingInput(file, IOContext.READ, fileInfo.metadata())) { for (int i = 0; i < fileInfo.numberOfParts(); i++) { final long partBytes = fileInfo.partBytes(i); diff --git a/server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java b/server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java index f118e4106db83..aa8212e8dad69 100644 --- a/server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java +++ b/server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java @@ -303,7 +303,8 @@ protected void search(List leaves, Weight weight, Collector c * {@link LeafCollector#collect(int)} is called for every matching document in * the provided ctx. */ - private void searchLeaf(LeafReaderContext ctx, Weight weight, Collector collector) throws IOException { + @Override + protected void searchLeaf(LeafReaderContext ctx, Weight weight, Collector collector) throws IOException { // Check if at all we need to call this leaf for collecting results. if (canMatch(ctx) == false) { diff --git a/server/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec b/server/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec index f51452c57f975..432d7e4906c08 100644 --- a/server/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec +++ b/server/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec @@ -1 +1,2 @@ org.opensearch.index.codec.composite.composite99.Composite99Codec +org.opensearch.index.codec.composite.composite912.Composite912Codec diff --git a/server/src/test/java/org/opensearch/index/codec/CodecTests.java b/server/src/test/java/org/opensearch/index/codec/CodecTests.java index bbf98b5c32918..cdd2c6d92f788 100644 --- a/server/src/test/java/org/opensearch/index/codec/CodecTests.java +++ b/server/src/test/java/org/opensearch/index/codec/CodecTests.java @@ -35,7 +35,7 @@ import org.apache.logging.log4j.LogManager; import org.apache.lucene.codecs.Codec; import org.apache.lucene.codecs.lucene90.Lucene90StoredFieldsFormat; -import org.apache.lucene.codecs.lucene99.Lucene99Codec; +import org.apache.lucene.codecs.lucene912.Lucene912Codec; import org.apache.lucene.document.Document; import org.apache.lucene.index.DirectoryReader; import org.apache.lucene.index.IndexWriter; @@ -48,7 +48,7 @@ import org.opensearch.env.Environment; import org.opensearch.index.IndexSettings; import org.opensearch.index.analysis.IndexAnalyzers; -import org.opensearch.index.codec.composite.composite99.Composite99Codec; +import org.opensearch.index.codec.composite.composite912.Composite912Codec; import org.opensearch.index.engine.EngineConfig; import org.opensearch.index.mapper.MapperService; import org.opensearch.index.similarity.SimilarityService; @@ -71,58 +71,58 @@ public class CodecTests extends OpenSearchTestCase { public void testResolveDefaultCodecs() throws Exception { CodecService codecService = createCodecService(false); assertThat(codecService.codec("default"), instanceOf(PerFieldMappingPostingFormatCodec.class)); - assertThat(codecService.codec("default"), instanceOf(Lucene99Codec.class)); + assertThat(codecService.codec("default"), instanceOf(Lucene912Codec.class)); } public void testDefault() throws Exception { Codec codec = createCodecService(false).codec("default"); - assertStoredFieldsCompressionEquals(Lucene99Codec.Mode.BEST_SPEED, codec); + assertStoredFieldsCompressionEquals(Lucene912Codec.Mode.BEST_SPEED, codec); } public void testDefaultWithCompositeIndex() throws Exception { Codec codec = createCodecService(false, true).codec("default"); - assertStoredFieldsCompressionEquals(Lucene99Codec.Mode.BEST_SPEED, codec); - assert codec instanceof Composite99Codec; + assertStoredFieldsCompressionEquals(Lucene912Codec.Mode.BEST_SPEED, codec); + assert codec instanceof Composite912Codec; } public void testBestCompression() throws Exception { Codec codec = createCodecService(false).codec("best_compression"); - assertStoredFieldsCompressionEquals(Lucene99Codec.Mode.BEST_COMPRESSION, codec); + assertStoredFieldsCompressionEquals(Lucene912Codec.Mode.BEST_COMPRESSION, codec); } public void testBestCompressionWithCompositeIndex() throws Exception { Codec codec = createCodecService(false, true).codec("best_compression"); - assertStoredFieldsCompressionEquals(Lucene99Codec.Mode.BEST_COMPRESSION, codec); - assert codec instanceof Composite99Codec; + assertStoredFieldsCompressionEquals(Lucene912Codec.Mode.BEST_COMPRESSION, codec); + assert codec instanceof Composite912Codec; } public void testLZ4() throws Exception { Codec codec = createCodecService(false).codec("lz4"); - assertStoredFieldsCompressionEquals(Lucene99Codec.Mode.BEST_SPEED, codec); + assertStoredFieldsCompressionEquals(Lucene912Codec.Mode.BEST_SPEED, codec); assert codec instanceof PerFieldMappingPostingFormatCodec; } public void testLZ4WithCompositeIndex() throws Exception { Codec codec = createCodecService(false, true).codec("lz4"); - assertStoredFieldsCompressionEquals(Lucene99Codec.Mode.BEST_SPEED, codec); - assert codec instanceof Composite99Codec; + assertStoredFieldsCompressionEquals(Lucene912Codec.Mode.BEST_SPEED, codec); + assert codec instanceof Composite912Codec; } public void testZlib() throws Exception { Codec codec = createCodecService(false).codec("zlib"); - assertStoredFieldsCompressionEquals(Lucene99Codec.Mode.BEST_COMPRESSION, codec); + assertStoredFieldsCompressionEquals(Lucene912Codec.Mode.BEST_COMPRESSION, codec); assert codec instanceof PerFieldMappingPostingFormatCodec; } public void testZlibWithCompositeIndex() throws Exception { Codec codec = createCodecService(false, true).codec("zlib"); - assertStoredFieldsCompressionEquals(Lucene99Codec.Mode.BEST_COMPRESSION, codec); - assert codec instanceof Composite99Codec; + assertStoredFieldsCompressionEquals(Lucene912Codec.Mode.BEST_COMPRESSION, codec); + assert codec instanceof Composite912Codec; } public void testResolveDefaultCodecsWithCompositeIndex() throws Exception { CodecService codecService = createCodecService(false, true); - assertThat(codecService.codec("default"), instanceOf(Composite99Codec.class)); + assertThat(codecService.codec("default"), instanceOf(Composite912Codec.class)); } public void testBestCompressionWithCompressionLevel() { @@ -157,12 +157,12 @@ public void testLuceneCodecsWithCompressionLevel() { public void testDefaultMapperServiceNull() throws Exception { Codec codec = createCodecService(true).codec("default"); - assertStoredFieldsCompressionEquals(Lucene99Codec.Mode.BEST_SPEED, codec); + assertStoredFieldsCompressionEquals(Lucene912Codec.Mode.BEST_SPEED, codec); } public void testBestCompressionMapperServiceNull() throws Exception { Codec codec = createCodecService(true).codec("best_compression"); - assertStoredFieldsCompressionEquals(Lucene99Codec.Mode.BEST_COMPRESSION, codec); + assertStoredFieldsCompressionEquals(Lucene912Codec.Mode.BEST_COMPRESSION, codec); } public void testExceptionCodecNull() { @@ -174,11 +174,11 @@ public void testExceptionIndexSettingsNull() { } // write some docs with it, inspect .si to see this was the used compression - private void assertStoredFieldsCompressionEquals(Lucene99Codec.Mode expected, Codec actual) throws Exception { + private void assertStoredFieldsCompressionEquals(Lucene912Codec.Mode expected, Codec actual) throws Exception { SegmentReader sr = getSegmentReader(actual); String v = sr.getSegmentInfo().info.getAttribute(Lucene90StoredFieldsFormat.MODE_KEY); assertNotNull(v); - assertEquals(expected, Lucene99Codec.Mode.valueOf(v)); + assertEquals(expected, Lucene912Codec.Mode.valueOf(v)); } private CodecService createCodecService(boolean isMapperServiceNull) throws IOException { diff --git a/server/src/test/java/org/opensearch/index/codec/composite/LuceneDocValuesConsumerFactoryTests.java b/server/src/test/java/org/opensearch/index/codec/composite/LuceneDocValuesConsumerFactoryTests.java index 7fb8fe7f68f45..5abb01dd02e1b 100644 --- a/server/src/test/java/org/opensearch/index/codec/composite/LuceneDocValuesConsumerFactoryTests.java +++ b/server/src/test/java/org/opensearch/index/codec/composite/LuceneDocValuesConsumerFactoryTests.java @@ -9,7 +9,7 @@ package org.opensearch.index.codec.composite; import org.apache.lucene.codecs.DocValuesConsumer; -import org.apache.lucene.codecs.lucene99.Lucene99Codec; +import org.apache.lucene.codecs.lucene912.Lucene912Codec; import org.apache.lucene.index.FieldInfo; import org.apache.lucene.index.FieldInfos; import org.apache.lucene.index.SegmentInfo; @@ -44,12 +44,12 @@ public void testGetDocValuesConsumerForCompositeCodec() throws IOException { SegmentInfo segmentInfo = new SegmentInfo( directory, Version.LATEST, - Version.LUCENE_9_11_0, + Version.LUCENE_9_12_0, "test_segment", randomInt(), false, false, - new Lucene99Codec(), + new Lucene912Codec(), new HashMap<>(), UUID.randomUUID().toString().substring(0, 16).getBytes(StandardCharsets.UTF_8), new HashMap<>(), diff --git a/server/src/test/java/org/opensearch/index/codec/composite/LuceneDocValuesProducerFactoryTests.java b/server/src/test/java/org/opensearch/index/codec/composite/LuceneDocValuesProducerFactoryTests.java index 55d637dfb9cae..67a953d923b22 100644 --- a/server/src/test/java/org/opensearch/index/codec/composite/LuceneDocValuesProducerFactoryTests.java +++ b/server/src/test/java/org/opensearch/index/codec/composite/LuceneDocValuesProducerFactoryTests.java @@ -10,7 +10,7 @@ import org.apache.lucene.codecs.DocValuesConsumer; import org.apache.lucene.codecs.DocValuesProducer; -import org.apache.lucene.codecs.lucene99.Lucene99Codec; +import org.apache.lucene.codecs.lucene912.Lucene912Codec; import org.apache.lucene.index.FieldInfo; import org.apache.lucene.index.FieldInfos; import org.apache.lucene.index.SegmentInfo; @@ -48,12 +48,12 @@ public void testGetDocValuesProducerForCompositeCodec99() throws IOException { SegmentInfo segmentInfo = new SegmentInfo( directory, Version.LATEST, - Version.LUCENE_9_11_0, + Version.LUCENE_9_12_0, "test_segment", randomInt(), false, false, - new Lucene99Codec(), + new Lucene912Codec(), new HashMap<>(), UUID.randomUUID().toString().substring(0, 16).getBytes(StandardCharsets.UTF_8), new HashMap<>(), diff --git a/server/src/test/java/org/opensearch/index/codec/composite99/datacube/startree/StarTreeDocValuesFormatTests.java b/server/src/test/java/org/opensearch/index/codec/composite912/datacube/startree/StarTreeDocValuesFormatTests.java similarity index 97% rename from server/src/test/java/org/opensearch/index/codec/composite99/datacube/startree/StarTreeDocValuesFormatTests.java rename to server/src/test/java/org/opensearch/index/codec/composite912/datacube/startree/StarTreeDocValuesFormatTests.java index 1c267c67e60ed..4fe0199f89f41 100644 --- a/server/src/test/java/org/opensearch/index/codec/composite99/datacube/startree/StarTreeDocValuesFormatTests.java +++ b/server/src/test/java/org/opensearch/index/codec/composite912/datacube/startree/StarTreeDocValuesFormatTests.java @@ -6,14 +6,14 @@ * compatible open source license. */ -package org.opensearch.index.codec.composite99.datacube.startree; +package org.opensearch.index.codec.composite912.datacube.startree; import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.apache.lucene.codecs.Codec; -import org.apache.lucene.codecs.lucene99.Lucene99Codec; +import org.apache.lucene.codecs.lucene912.Lucene912Codec; import org.apache.lucene.document.Document; import org.apache.lucene.document.SortedNumericDocValuesField; import org.apache.lucene.index.DirectoryReader; @@ -41,7 +41,7 @@ import org.opensearch.index.MapperTestUtils; import org.opensearch.index.codec.composite.CompositeIndexFieldInfo; import org.opensearch.index.codec.composite.CompositeIndexReader; -import org.opensearch.index.codec.composite.composite99.Composite99Codec; +import org.opensearch.index.codec.composite.composite912.Composite912Codec; import org.opensearch.index.compositeindex.datacube.startree.StarTreeDocument; import org.opensearch.index.compositeindex.datacube.startree.StarTreeFieldConfiguration; import org.opensearch.index.compositeindex.datacube.startree.StarTreeIndexSettings; @@ -107,7 +107,7 @@ protected Codec getCodec() { } catch (IOException e) { throw new RuntimeException(e); } - Codec codec = new Composite99Codec(Lucene99Codec.Mode.BEST_SPEED, mapperService, testLogger); + Codec codec = new Composite912Codec(Lucene912Codec.Mode.BEST_SPEED, mapperService, testLogger); return codec; } 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 index b77200f173e71..95fd1579c2f63 100644 --- 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 @@ -10,7 +10,7 @@ import org.apache.lucene.codecs.DocValuesConsumer; import org.apache.lucene.codecs.DocValuesProducer; -import org.apache.lucene.codecs.lucene99.Lucene99Codec; +import org.apache.lucene.codecs.lucene912.Lucene912Codec; import org.apache.lucene.index.DocValues; import org.apache.lucene.index.DocValuesType; import org.apache.lucene.index.EmptyDocValuesProducer; @@ -273,7 +273,7 @@ private SegmentReadState getReadState(int numDocs, List dimensionFields, numDocs, false, false, - new Lucene99Codec(), + new Lucene912Codec(), new HashMap<>(), writeState.segmentInfo.getId(), new HashMap<>(), @@ -287,12 +287,12 @@ private SegmentWriteState getWriteState(int numDocs, byte[] id) { SegmentInfo segmentInfo = new SegmentInfo( directory, Version.LATEST, - Version.LUCENE_9_11_0, + Version.LUCENE_9_12_0, "test_segment", numDocs, false, false, - new Lucene99Codec(), + new Lucene912Codec(), new HashMap<>(), id, new HashMap<>(), @@ -1453,7 +1453,7 @@ public void test_build_multipleStarTrees() throws IOException { 7, false, false, - new Lucene99Codec(), + new Lucene912Codec(), new HashMap<>(), UUID.randomUUID().toString().substring(0, 16).getBytes(StandardCharsets.UTF_8), new HashMap<>(), @@ -1724,7 +1724,7 @@ public void test_build_starTreeDataset() throws IOException { 7, false, false, - new Lucene99Codec(), + new Lucene912Codec(), new HashMap<>(), UUID.randomUUID().toString().substring(0, 16).getBytes(StandardCharsets.UTF_8), new HashMap<>(), diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/BaseStarTreeBuilderTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/BaseStarTreeBuilderTests.java index 8f750b09c75ee..705700aada2ff 100644 --- a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/BaseStarTreeBuilderTests.java +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/BaseStarTreeBuilderTests.java @@ -10,7 +10,7 @@ import org.apache.lucene.codecs.DocValuesConsumer; import org.apache.lucene.codecs.DocValuesProducer; -import org.apache.lucene.codecs.lucene99.Lucene99Codec; +import org.apache.lucene.codecs.lucene912.Lucene912Codec; import org.apache.lucene.index.DocValuesType; import org.apache.lucene.index.FieldInfo; import org.apache.lucene.index.FieldInfos; @@ -108,12 +108,12 @@ public static void setup() throws IOException { SegmentInfo segmentInfo = new SegmentInfo( directory, Version.LATEST, - Version.LUCENE_9_11_0, + Version.LUCENE_9_12_0, "test_segment", 5, false, false, - new Lucene99Codec(), + new Lucene912Codec(), new HashMap<>(), UUID.randomUUID().toString().substring(0, 16).getBytes(StandardCharsets.UTF_8), new HashMap<>(), diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreesBuilderTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreesBuilderTests.java index 94fe148eab6b3..4ab21dbce059f 100644 --- a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreesBuilderTests.java +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreesBuilderTests.java @@ -10,7 +10,7 @@ import org.apache.lucene.codecs.DocValuesConsumer; import org.apache.lucene.codecs.DocValuesProducer; -import org.apache.lucene.codecs.lucene99.Lucene99Codec; +import org.apache.lucene.codecs.lucene912.Lucene912Codec; import org.apache.lucene.index.FieldInfo; import org.apache.lucene.index.FieldInfos; import org.apache.lucene.index.SegmentInfo; @@ -60,12 +60,12 @@ public void setUp() throws Exception { SegmentInfo segmentInfo = new SegmentInfo( directory, Version.LATEST, - Version.LUCENE_9_11_0, + Version.LUCENE_9_12_0, "test_segment", 5, false, false, - new Lucene99Codec(), + new Lucene912Codec(), new HashMap<>(), UUID.randomUUID().toString().substring(0, 16).getBytes(StandardCharsets.UTF_8), new HashMap<>(), diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/meta/StarTreeMetadataTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/meta/StarTreeMetadataTests.java index 2127eeb9636d3..c8636426449ad 100644 --- a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/meta/StarTreeMetadataTests.java +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/meta/StarTreeMetadataTests.java @@ -8,7 +8,7 @@ package org.opensearch.index.compositeindex.datacube.startree.fileformats.meta; -import org.apache.lucene.codecs.lucene99.Lucene99Codec; +import org.apache.lucene.codecs.lucene912.Lucene912Codec; import org.apache.lucene.index.DocValuesType; import org.apache.lucene.index.FieldInfo; import org.apache.lucene.index.FieldInfos; @@ -73,12 +73,12 @@ public void setup() throws IOException { SegmentInfo segmentInfo = new SegmentInfo( directory, Version.LATEST, - Version.LUCENE_9_11_0, + Version.LUCENE_9_12_0, "test_segment", 6, false, false, - new Lucene99Codec(), + new Lucene912Codec(), new HashMap<>(), UUID.randomUUID().toString().substring(0, 16).getBytes(StandardCharsets.UTF_8), new HashMap<>(), diff --git a/server/src/test/java/org/opensearch/index/engine/CompletionStatsCacheTests.java b/server/src/test/java/org/opensearch/index/engine/CompletionStatsCacheTests.java index 0c87c384e0749..6a2db930aff9b 100644 --- a/server/src/test/java/org/opensearch/index/engine/CompletionStatsCacheTests.java +++ b/server/src/test/java/org/opensearch/index/engine/CompletionStatsCacheTests.java @@ -32,14 +32,14 @@ package org.opensearch.index.engine; import org.apache.lucene.codecs.PostingsFormat; -import org.apache.lucene.codecs.lucene99.Lucene99Codec; +import org.apache.lucene.codecs.lucene912.Lucene912Codec; import org.apache.lucene.document.Document; import org.apache.lucene.index.DirectoryReader; import org.apache.lucene.index.IndexWriter; import org.apache.lucene.index.IndexWriterConfig; import org.apache.lucene.search.Query; import org.apache.lucene.search.QueryCachingPolicy; -import org.apache.lucene.search.suggest.document.Completion99PostingsFormat; +import org.apache.lucene.search.suggest.document.Completion912PostingsFormat; import org.apache.lucene.search.suggest.document.SuggestField; import org.apache.lucene.store.Directory; import org.opensearch.OpenSearchException; @@ -69,8 +69,8 @@ public void testExceptionsAreNotCached() { public void testCompletionStatsCache() throws IOException, InterruptedException { final IndexWriterConfig indexWriterConfig = newIndexWriterConfig(); - final PostingsFormat postingsFormat = new Completion99PostingsFormat(); - indexWriterConfig.setCodec(new Lucene99Codec() { + final PostingsFormat postingsFormat = new Completion912PostingsFormat(); + indexWriterConfig.setCodec(new Lucene912Codec() { @Override public PostingsFormat getPostingsFormatForField(String field) { return postingsFormat; // all fields are suggest fields diff --git a/server/src/test/java/org/opensearch/index/shard/RemoteIndexShardCorruptionTests.java b/server/src/test/java/org/opensearch/index/shard/RemoteIndexShardCorruptionTests.java index 21bf580712761..9f31fb84ce6af 100644 --- a/server/src/test/java/org/opensearch/index/shard/RemoteIndexShardCorruptionTests.java +++ b/server/src/test/java/org/opensearch/index/shard/RemoteIndexShardCorruptionTests.java @@ -43,7 +43,7 @@ public void testLocalDirectoryContains() throws IOException { } boolean corrupted = randomBoolean(); long checksum = 0; - try (IndexInput indexInput = localDirectory.openInput(file, IOContext.DEFAULT)) { + try (IndexInput indexInput = localDirectory.openInput(file, IOContext.READONCE)) { checksum = CodecUtil.retrieveChecksum(indexInput); } if (corrupted) { diff --git a/server/src/test/java/org/opensearch/index/store/StoreTests.java b/server/src/test/java/org/opensearch/index/store/StoreTests.java index ab30a4c1c435f..3777753cb2de8 100644 --- a/server/src/test/java/org/opensearch/index/store/StoreTests.java +++ b/server/src/test/java/org/opensearch/index/store/StoreTests.java @@ -392,7 +392,7 @@ public void testNewChecksums() throws IOException { metadata = store.getMetadata(); assertThat(metadata.asMap().isEmpty(), is(false)); for (StoreFileMetadata meta : metadata) { - try (IndexInput input = store.directory().openInput(meta.name(), IOContext.DEFAULT)) { + try (IndexInput input = store.directory().openInput(meta.name(), IOContext.READONCE)) { String checksum = Store.digestToString(CodecUtil.retrieveChecksum(input)); assertThat("File: " + meta.name() + " has a different checksum", meta.checksum(), equalTo(checksum)); assertThat(meta.writtenBy(), equalTo(Version.LATEST)); diff --git a/test/framework/src/main/java/org/opensearch/test/CorruptionUtils.java b/test/framework/src/main/java/org/opensearch/test/CorruptionUtils.java index 67522bb618cf1..8e551a51a2094 100644 --- a/test/framework/src/main/java/org/opensearch/test/CorruptionUtils.java +++ b/test/framework/src/main/java/org/opensearch/test/CorruptionUtils.java @@ -96,7 +96,7 @@ public static void corruptFile(Random random, Path... files) throws IOException long checksumAfterCorruption; long actualChecksumAfterCorruption; - try (ChecksumIndexInput input = dir.openChecksumInput(fileToCorrupt.getFileName().toString(), IOContext.DEFAULT)) { + try (ChecksumIndexInput input = dir.openChecksumInput(fileToCorrupt.getFileName().toString(), IOContext.READONCE)) { assertThat(input.getFilePointer(), is(0L)); input.seek(input.length() - CodecUtil.footerLength()); checksumAfterCorruption = input.getChecksum(); From e885aa93279342c9ab219ecf612d887ff8de8af6 Mon Sep 17 00:00:00 2001 From: expani1729 <110471048+expani@users.noreply.github.com> Date: Tue, 8 Oct 2024 00:31:55 +0530 Subject: [PATCH 2/7] Latency improvements to Multi Term Aggregations (#14993) * Avoid deep copy and other allocation improvements * Refactoring based on PR Comments and added JavaDocs * Added more comments * Added character for Triggering Jenkins build * Changes to cover collectZeroDocEntries method * Updated comment based on change in method's functionality * Added test to cover branches in collectZeroDocEntriesIfRequired * Rebased and resolved changelog conflict --------- Signed-off-by: expani --- CHANGELOG.md | 1 + .../bucket/terms/MultiTermsAggregator.java | 96 ++++++++++++------- .../terms/MultiTermsAggregatorTests.java | 70 +++++++++++--- 3 files changed, 116 insertions(+), 51 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2e3a515fb8008..c2d021273f271 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,6 +19,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Add changes to block calls in cat shards, indices and segments based on dynamic limit settings ([#15986](https://github.com/opensearch-project/OpenSearch/pull/15986)) - New `phone` & `phone-search` analyzer + tokenizer ([#15915](https://github.com/opensearch-project/OpenSearch/pull/15915)) - Add _list/shards API as paginated alternate to _cat/shards ([#14641](https://github.com/opensearch-project/OpenSearch/pull/14641)) +- Latency and Memory allocation improvements to Multi Term Aggregation queries ([#14993](https://github.com/opensearch-project/OpenSearch/pull/14993)) ### Dependencies - Bump `com.azure:azure-identity` from 1.13.0 to 1.13.2 ([#15578](https://github.com/opensearch-project/OpenSearch/pull/15578)) diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/terms/MultiTermsAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/terms/MultiTermsAggregator.java index 59f48bd7fbaba..07edf487af670 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/terms/MultiTermsAggregator.java +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/terms/MultiTermsAggregator.java @@ -33,6 +33,7 @@ import org.opensearch.search.aggregations.InternalAggregation; import org.opensearch.search.aggregations.InternalOrder; import org.opensearch.search.aggregations.LeafBucketCollector; +import org.opensearch.search.aggregations.bucket.BucketsAggregator; import org.opensearch.search.aggregations.bucket.DeferableBucketAggregator; import org.opensearch.search.aggregations.bucket.LocalBucketCountThresholds; import org.opensearch.search.aggregations.support.AggregationPath; @@ -215,19 +216,11 @@ public InternalAggregation buildEmptyAggregation() { @Override protected LeafBucketCollector getLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub) throws IOException { - MultiTermsValuesSourceCollector collector = multiTermsValue.getValues(ctx); + MultiTermsValuesSourceCollector collector = multiTermsValue.getValues(ctx, bucketOrds, this, sub); return new LeafBucketCollector() { @Override public void collect(int doc, long owningBucketOrd) throws IOException { - for (BytesRef compositeKey : collector.apply(doc)) { - long bucketOrd = bucketOrds.add(owningBucketOrd, compositeKey); - if (bucketOrd < 0) { - bucketOrd = -1 - bucketOrd; - collectExistingBucket(sub, doc, bucketOrd); - } else { - collectBucket(sub, doc, bucketOrd); - } - } + collector.apply(doc, owningBucketOrd); } }; } @@ -268,12 +261,10 @@ private void collectZeroDocEntriesIfNeeded(long owningBucketOrd) throws IOExcept } // we need to fill-in the blanks for (LeafReaderContext ctx : context.searcher().getTopReaderContext().leaves()) { - MultiTermsValuesSourceCollector collector = multiTermsValue.getValues(ctx); // brute force + MultiTermsValuesSourceCollector collector = multiTermsValue.getValues(ctx, bucketOrds, null, null); for (int docId = 0; docId < ctx.reader().maxDoc(); ++docId) { - for (BytesRef compositeKey : collector.apply(docId)) { - bucketOrds.add(owningBucketOrd, compositeKey); - } + collector.apply(docId, owningBucketOrd); } } } @@ -284,10 +275,11 @@ private void collectZeroDocEntriesIfNeeded(long owningBucketOrd) throws IOExcept @FunctionalInterface interface MultiTermsValuesSourceCollector { /** - * Collect a list values of multi_terms on each doc. - * Each terms could have multi_values, so the result is the cartesian product of each term's values. + * Generates the cartesian product of all fields used in aggregation and + * collects them in buckets using the composite key of their field values. */ - List apply(int doc) throws IOException; + void apply(int doc, long owningBucketOrd) throws IOException; + } @FunctionalInterface @@ -361,47 +353,72 @@ public MultiTermsValuesSource(List valuesSources) { this.valuesSources = valuesSources; } - public MultiTermsValuesSourceCollector getValues(LeafReaderContext ctx) throws IOException { + public MultiTermsValuesSourceCollector getValues( + LeafReaderContext ctx, + BytesKeyedBucketOrds bucketOrds, + BucketsAggregator aggregator, + LeafBucketCollector sub + ) throws IOException { List collectors = new ArrayList<>(); for (InternalValuesSource valuesSource : valuesSources) { collectors.add(valuesSource.apply(ctx)); } + boolean collectBucketOrds = aggregator != null && sub != null; return new MultiTermsValuesSourceCollector() { + + /** + * This method does the following :
+ *
  • Fetches the values of every field present in the doc List>> via @{@link InternalValuesSourceCollector}
    • + *
  • Generates Composite keys from the fetched values for all fields present in the aggregation.
    • + *
  • Adds every composite key to the @{@link BytesKeyedBucketOrds} and Optionally collects them via @{@link BucketsAggregator#collectBucket(LeafBucketCollector, int, long)}
    • + */ @Override - public List apply(int doc) throws IOException { + public void apply(int doc, long owningBucketOrd) throws IOException { + // TODO A new list creation can be avoided for every doc. List>> collectedValues = new ArrayList<>(); for (InternalValuesSourceCollector collector : collectors) { collectedValues.add(collector.apply(doc)); } - List result = new ArrayList<>(); scratch.seek(0); scratch.writeVInt(collectors.size()); // number of fields per composite key - cartesianProduct(result, scratch, collectedValues, 0); - return result; + generateAndCollectCompositeKeys(collectedValues, 0, owningBucketOrd, doc); } /** - * Cartesian product using depth first search. - * - *

    - * Composite keys are encoded to a {@link BytesRef} in a format compatible with {@link StreamOutput::writeGenericValue}, - * but reuses the encoding of the shared prefixes from the previous levels to avoid wasteful work. + * This generates and collects all Composite keys in their buckets by performing a cartesian product
    + * of all the values in all the fields ( used in agg ) for the given doc recursively. + * @param collectedValues : Values of all fields present in the aggregation for the @doc + * @param index : Points to the field being added to generate the composite key */ - private void cartesianProduct( - List compositeKeys, - BytesStreamOutput scratch, + private void generateAndCollectCompositeKeys( List>> collectedValues, - int index + int index, + long owningBucketOrd, + int doc ) throws IOException { if (collectedValues.size() == index) { - compositeKeys.add(BytesRef.deepCopyOf(scratch.bytes().toBytesRef())); + // Avoid performing a deep copy of the composite key by inlining. + long bucketOrd = bucketOrds.add(owningBucketOrd, scratch.bytes().toBytesRef()); + if (collectBucketOrds) { + if (bucketOrd < 0) { + bucketOrd = -1 - bucketOrd; + aggregator.collectExistingBucket(sub, doc, bucketOrd); + } else { + aggregator.collectBucket(sub, doc, bucketOrd); + } + } return; } long position = scratch.position(); - for (TermValue value : collectedValues.get(index)) { + List> values = collectedValues.get(index); + int numIterations = values.size(); + // For each loop is not done to reduce the allocations done for Iterator objects + // once for every field in every doc. + for (int i = 0; i < numIterations; i++) { + TermValue value = values.get(i); value.writeTo(scratch); // encode the value - cartesianProduct(compositeKeys, scratch, collectedValues, index + 1); // dfs + generateAndCollectCompositeKeys(collectedValues, index + 1, owningBucketOrd, doc); // dfs scratch.seek(position); // backtrack } } @@ -441,9 +458,14 @@ static InternalValuesSource bytesValuesSource(ValuesSource valuesSource, Include if (i > 0 && bytes.equals(previous)) { continue; } - BytesRef copy = BytesRef.deepCopyOf(bytes); - termValues.add(TermValue.of(copy)); - previous = copy; + // Performing a deep copy is not required for field containing only one value. + if (valuesCount > 1) { + BytesRef copy = BytesRef.deepCopyOf(bytes); + termValues.add(TermValue.of(copy)); + previous = copy; + } else { + termValues.add(TermValue.of(bytes)); + } } return termValues; }; diff --git a/server/src/test/java/org/opensearch/search/aggregations/bucket/terms/MultiTermsAggregatorTests.java b/server/src/test/java/org/opensearch/search/aggregations/bucket/terms/MultiTermsAggregatorTests.java index d550c4c354c0f..bb46c5607a4a7 100644 --- a/server/src/test/java/org/opensearch/search/aggregations/bucket/terms/MultiTermsAggregatorTests.java +++ b/server/src/test/java/org/opensearch/search/aggregations/bucket/terms/MultiTermsAggregatorTests.java @@ -126,6 +126,19 @@ public class MultiTermsAggregatorTests extends AggregatorTestCase { private static final Consumer NONE_DECORATOR = null; + private static final Consumer IP_AND_KEYWORD_DESC_ORDER_VERIFY = h -> { + MatcherAssert.assertThat(h.getBuckets(), hasSize(3)); + MatcherAssert.assertThat(h.getBuckets().get(0).getKey(), contains(equalTo("a"), equalTo("192.168.0.0"))); + MatcherAssert.assertThat(h.getBuckets().get(0).getKeyAsString(), equalTo("a|192.168.0.0")); + MatcherAssert.assertThat(h.getBuckets().get(0).getDocCount(), equalTo(2L)); + MatcherAssert.assertThat(h.getBuckets().get(1).getKey(), contains(equalTo("b"), equalTo("192.168.0.1"))); + MatcherAssert.assertThat(h.getBuckets().get(1).getKeyAsString(), equalTo("b|192.168.0.1")); + MatcherAssert.assertThat(h.getBuckets().get(1).getDocCount(), equalTo(1L)); + MatcherAssert.assertThat(h.getBuckets().get(2).getKey(), contains(equalTo("c"), equalTo("192.168.0.2"))); + MatcherAssert.assertThat(h.getBuckets().get(2).getKeyAsString(), equalTo("c|192.168.0.2")); + MatcherAssert.assertThat(h.getBuckets().get(2).getDocCount(), equalTo(1L)); + }; + @Override protected List getSupportedValuesSourceTypes() { return Collections.unmodifiableList( @@ -672,8 +685,48 @@ public void testDatesFieldFormat() throws IOException { ); } - public void testIpAndKeyword() throws IOException { - testAggregation(new MatchAllDocsQuery(), fieldConfigs(asList(KEYWORD_FIELD, IP_FIELD)), NONE_DECORATOR, iw -> { + public void testIpAndKeywordDefaultDescOrder() throws IOException { + ipAndKeywordTest(NONE_DECORATOR, IP_AND_KEYWORD_DESC_ORDER_VERIFY); + } + + public void testIpAndKeywordWithBucketCountSameAsSize() throws IOException { + ipAndKeywordTest(multiTermsAggregationBuilder -> { + multiTermsAggregationBuilder.minDocCount(0); + multiTermsAggregationBuilder.size(3); + multiTermsAggregationBuilder.order(BucketOrder.compound(BucketOrder.count(false))); + }, IP_AND_KEYWORD_DESC_ORDER_VERIFY); + } + + public void testIpAndKeywordWithBucketCountGreaterThanSize() throws IOException { + ipAndKeywordTest(multiTermsAggregationBuilder -> { + multiTermsAggregationBuilder.minDocCount(0); + multiTermsAggregationBuilder.size(10); + multiTermsAggregationBuilder.order(BucketOrder.compound(BucketOrder.count(false))); + }, IP_AND_KEYWORD_DESC_ORDER_VERIFY); + } + + public void testIpAndKeywordAscOrder() throws IOException { + ipAndKeywordTest(multiTermsAggregationBuilder -> { + multiTermsAggregationBuilder.minDocCount(0); + multiTermsAggregationBuilder.size(3); + multiTermsAggregationBuilder.order(BucketOrder.compound(BucketOrder.count(true))); + }, h -> { + MatcherAssert.assertThat(h.getBuckets(), hasSize(3)); + MatcherAssert.assertThat(h.getBuckets().get(0).getKey(), contains(equalTo("b"), equalTo("192.168.0.1"))); + MatcherAssert.assertThat(h.getBuckets().get(0).getKeyAsString(), equalTo("b|192.168.0.1")); + MatcherAssert.assertThat(h.getBuckets().get(0).getDocCount(), equalTo(1L)); + MatcherAssert.assertThat(h.getBuckets().get(1).getKey(), contains(equalTo("c"), equalTo("192.168.0.2"))); + MatcherAssert.assertThat(h.getBuckets().get(1).getKeyAsString(), equalTo("c|192.168.0.2")); + MatcherAssert.assertThat(h.getBuckets().get(1).getDocCount(), equalTo(1L)); + MatcherAssert.assertThat(h.getBuckets().get(2).getKey(), contains(equalTo("a"), equalTo("192.168.0.0"))); + MatcherAssert.assertThat(h.getBuckets().get(2).getKeyAsString(), equalTo("a|192.168.0.0")); + MatcherAssert.assertThat(h.getBuckets().get(2).getDocCount(), equalTo(2L)); + }); + } + + private void ipAndKeywordTest(Consumer builderDecorator, Consumer verify) + throws IOException { + testAggregation(new MatchAllDocsQuery(), fieldConfigs(asList(KEYWORD_FIELD, IP_FIELD)), builderDecorator, iw -> { iw.addDocument( asList( new SortedDocValuesField(KEYWORD_FIELD, new BytesRef("a")), @@ -698,18 +751,7 @@ public void testIpAndKeyword() throws IOException { new SortedDocValuesField(IP_FIELD, new BytesRef(InetAddressPoint.encode(InetAddresses.forString("192.168.0.0")))) ) ); - }, h -> { - MatcherAssert.assertThat(h.getBuckets(), hasSize(3)); - MatcherAssert.assertThat(h.getBuckets().get(0).getKey(), contains(equalTo("a"), equalTo("192.168.0.0"))); - MatcherAssert.assertThat(h.getBuckets().get(0).getKeyAsString(), equalTo("a|192.168.0.0")); - MatcherAssert.assertThat(h.getBuckets().get(0).getDocCount(), equalTo(2L)); - MatcherAssert.assertThat(h.getBuckets().get(1).getKey(), contains(equalTo("b"), equalTo("192.168.0.1"))); - MatcherAssert.assertThat(h.getBuckets().get(1).getKeyAsString(), equalTo("b|192.168.0.1")); - MatcherAssert.assertThat(h.getBuckets().get(1).getDocCount(), equalTo(1L)); - MatcherAssert.assertThat(h.getBuckets().get(2).getKey(), contains(equalTo("c"), equalTo("192.168.0.2"))); - MatcherAssert.assertThat(h.getBuckets().get(2).getKeyAsString(), equalTo("c|192.168.0.2")); - MatcherAssert.assertThat(h.getBuckets().get(2).getDocCount(), equalTo(1L)); - }); + }, verify); } public void testEmpty() throws IOException { From e24b4c98405d95b781a43fe73ca370c1a1ca452a Mon Sep 17 00:00:00 2001 From: Kaushal Kumar Date: Mon, 7 Oct 2024 17:21:20 -0700 Subject: [PATCH 3/7] Add wlm resiliency orchestrator (query group service) (#15925) * cancellation related Signed-off-by: Kiran Prakash * Update CHANGELOG.md Signed-off-by: Kiran Prakash * add better cancellation reason Signed-off-by: Kiran Prakash * Update DefaultTaskCancellationTests.java Signed-off-by: Kiran Prakash * refactor Signed-off-by: Kiran Prakash * refactor Signed-off-by: Kiran Prakash * Update DefaultTaskCancellation.java Signed-off-by: Kiran Prakash * Update DefaultTaskCancellation.java Signed-off-by: Kiran Prakash * Update DefaultTaskCancellation.java Signed-off-by: Kiran Prakash * Update DefaultTaskSelectionStrategy.java Signed-off-by: Kiran Prakash * refactor Signed-off-by: Kiran Prakash * refactor node level threshold Signed-off-by: Kiran Prakash * use query group task Signed-off-by: Kaushal Kumar * code clean up and refactorings Signed-off-by: Kaushal Kumar * add unit tests and fix existing ones Signed-off-by: Kaushal Kumar * uncomment the test case Signed-off-by: Kaushal Kumar * update CHANGELOG Signed-off-by: Kaushal Kumar * fix imports Signed-off-by: Kaushal Kumar * add queryGroupService Signed-off-by: Kaushal Kumar * refactor and add UTs for new constructs Signed-off-by: Kaushal Kumar * fix javadocs Signed-off-by: Kaushal Kumar * remove code clutter Signed-off-by: Kaushal Kumar * change annotation version and task selection strategy Signed-off-by: Kaushal Kumar * rename a util class Signed-off-by: Kaushal Kumar * remove wrappers from resource type Signed-off-by: Kaushal Kumar * apply spotless Signed-off-by: Kaushal Kumar * address comments Signed-off-by: Kaushal Kumar * add rename changes Signed-off-by: Kaushal Kumar * address comments Signed-off-by: Kaushal Kumar * initial changes Signed-off-by: Kaushal Kumar * refactor changes and logical bug fix Signed-off-by: Kaushal Kumar * add chanegs Signed-off-by: Kaushal Kumar * address comments Signed-off-by: Kaushal Kumar * temp changes Signed-off-by: Kaushal Kumar * add UTs Signed-off-by: Kaushal Kumar * add changelog Signed-off-by: Kaushal Kumar * add task completion listener hook Signed-off-by: Kaushal Kumar * add remaining pieces to make the feature functional Signed-off-by: Kaushal Kumar * extend stats and fix bugs Signed-off-by: Kaushal Kumar * fix bugs and add logic to make SBP work with wlm Signed-off-by: Kaushal Kumar * address comments Signed-off-by: Kaushal Kumar * fix bugs and SBP ITs Signed-off-by: Kaushal Kumar * add missed applyCluster state change Signed-off-by: Kaushal Kumar * address comments Signed-off-by: Kaushal Kumar * decouple queryGroupService and cancellationService Signed-off-by: Kaushal Kumar * replace StateApplier with StateListener interface Signed-off-by: Kaushal Kumar * fix precommit errors Signed-off-by: Kaushal Kumar --------- Signed-off-by: Kiran Prakash Signed-off-by: Kaushal Kumar Co-authored-by: Kiran Prakash --- CHANGELOG.md | 1 + .../backpressure/SearchBackpressureIT.java | 2 + .../org/opensearch/action/ActionModule.java | 6 +- .../common/settings/ClusterSettings.java | 3 + .../main/java/org/opensearch/node/Node.java | 39 +- .../SearchBackpressureService.java | 13 +- .../org/opensearch/wlm/QueryGroupService.java | 290 ++++++++++- .../org/opensearch/wlm/QueryGroupTask.java | 9 +- .../wlm/QueryGroupsStateAccessor.java | 64 +++ .../main/java/org/opensearch/wlm/WlmMode.java | 40 ++ .../wlm/WorkloadManagementSettings.java | 105 ++++ .../QueryGroupTaskCancellationService.java | 123 ++++- .../opensearch/wlm/stats/QueryGroupState.java | 29 +- .../opensearch/wlm/stats/QueryGroupStats.java | 11 +- ...QueryGroupResourceUsageTrackerService.java | 2 + .../SearchBackpressureServiceTests.java | 52 +- .../wlm/QueryGroupServiceTests.java | 489 ++++++++++++++++++ ...adManagementTransportInterceptorTests.java | 36 +- ...ueryGroupTaskCancellationServiceTests.java | 100 +++- ...eryGroupRequestOperationListenerTests.java | 97 +++- .../wlm/stats/QueryGroupStateTests.java | 10 +- .../wlm/stats/QueryGroupStatsTests.java | 6 +- 22 files changed, 1410 insertions(+), 117 deletions(-) create mode 100644 server/src/main/java/org/opensearch/wlm/QueryGroupsStateAccessor.java create mode 100644 server/src/main/java/org/opensearch/wlm/WlmMode.java create mode 100644 server/src/test/java/org/opensearch/wlm/QueryGroupServiceTests.java diff --git a/CHANGELOG.md b/CHANGELOG.md index c2d021273f271..ab8cdd882bbb9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Add support for async deletion in S3BlobContainer ([#15621](https://github.com/opensearch-project/OpenSearch/pull/15621)) - MultiTermQueries in keyword fields now default to `indexed` approach and gated behind cluster setting ([#15637](https://github.com/opensearch-project/OpenSearch/pull/15637)) - [Workload Management] QueryGroup resource cancellation framework changes ([#15651](https://github.com/opensearch-project/OpenSearch/pull/15651)) +- [Workload Management] Add orchestrator for wlm resiliency (QueryGroupService) ([#15925](https://github.com/opensearch-project/OpenSearch/pull/15925)) - Fallback to Remote cluster-state on Term-Version check mismatch - ([#15424](https://github.com/opensearch-project/OpenSearch/pull/15424)) - Implement WithFieldName interface in ValuesSourceAggregationBuilder & FieldSortBuilder ([#15916](https://github.com/opensearch-project/OpenSearch/pull/15916)) - Add successfulSearchShardIndices in searchRequestContext ([#15967](https://github.com/opensearch-project/OpenSearch/pull/15967)) diff --git a/server/src/internalClusterTest/java/org/opensearch/search/backpressure/SearchBackpressureIT.java b/server/src/internalClusterTest/java/org/opensearch/search/backpressure/SearchBackpressureIT.java index fb84134120e00..40c9301ef4bce 100644 --- a/server/src/internalClusterTest/java/org/opensearch/search/backpressure/SearchBackpressureIT.java +++ b/server/src/internalClusterTest/java/org/opensearch/search/backpressure/SearchBackpressureIT.java @@ -39,6 +39,7 @@ import org.opensearch.test.ParameterizedStaticSettingsOpenSearchIntegTestCase; import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.TransportService; +import org.opensearch.wlm.QueryGroupTask; import org.hamcrest.MatcherAssert; import org.junit.After; import org.junit.Before; @@ -411,6 +412,7 @@ protected void doExecute(Task task, TestRequest request, ActionListener { try { CancellableTask cancellableTask = (CancellableTask) task; + ((QueryGroupTask) task).setQueryGroupId(threadPool.getThreadContext()); long startTime = System.nanoTime(); // Doing a busy-wait until task cancellation or timeout. diff --git a/server/src/main/java/org/opensearch/action/ActionModule.java b/server/src/main/java/org/opensearch/action/ActionModule.java index d6b4a93b34217..12c7626e385a4 100644 --- a/server/src/main/java/org/opensearch/action/ActionModule.java +++ b/server/src/main/java/org/opensearch/action/ActionModule.java @@ -481,6 +481,7 @@ import org.opensearch.tasks.Task; import org.opensearch.threadpool.ThreadPool; import org.opensearch.usage.UsageService; +import org.opensearch.wlm.QueryGroupTask; import java.util.ArrayList; import java.util.Collections; @@ -565,7 +566,10 @@ public ActionModule( destructiveOperations = new DestructiveOperations(settings, clusterSettings); Set headers = Stream.concat( actionPlugins.stream().flatMap(p -> p.getRestHeaders().stream()), - Stream.of(new RestHeaderDefinition(Task.X_OPAQUE_ID, false)) + Stream.of( + new RestHeaderDefinition(Task.X_OPAQUE_ID, false), + new RestHeaderDefinition(QueryGroupTask.QUERY_GROUP_ID_HEADER, false) + ) ).collect(Collectors.toSet()); UnaryOperator restWrapper = null; for (ActionPlugin plugin : actionPlugins) { diff --git a/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java b/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java index e18c4dabb29eb..a84a29256ee19 100644 --- a/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java +++ b/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java @@ -798,6 +798,9 @@ public void apply(Settings value, Settings current, Settings previous) { WorkloadManagementSettings.NODE_LEVEL_CPU_CANCELLATION_THRESHOLD, WorkloadManagementSettings.NODE_LEVEL_MEMORY_REJECTION_THRESHOLD, WorkloadManagementSettings.NODE_LEVEL_MEMORY_CANCELLATION_THRESHOLD, + WorkloadManagementSettings.WLM_MODE_SETTING, + WorkloadManagementSettings.QUERYGROUP_SERVICE_RUN_INTERVAL_SETTING, + WorkloadManagementSettings.QUERYGROUP_SERVICE_DURESS_STREAK_SETTING, // Settings to be used for limiting rest requests ResponseLimitSettings.CAT_INDICES_RESPONSE_LIMIT_SETTING, diff --git a/server/src/main/java/org/opensearch/node/Node.java b/server/src/main/java/org/opensearch/node/Node.java index fe4bfa1e64ba1..ddba10369ff05 100644 --- a/server/src/main/java/org/opensearch/node/Node.java +++ b/server/src/main/java/org/opensearch/node/Node.java @@ -270,8 +270,13 @@ import org.opensearch.usage.UsageService; import org.opensearch.watcher.ResourceWatcherService; import org.opensearch.wlm.QueryGroupService; +import org.opensearch.wlm.QueryGroupsStateAccessor; +import org.opensearch.wlm.WorkloadManagementSettings; import org.opensearch.wlm.WorkloadManagementTransportInterceptor; +import org.opensearch.wlm.cancellation.MaximumResourceTaskSelectionStrategy; +import org.opensearch.wlm.cancellation.QueryGroupTaskCancellationService; import org.opensearch.wlm.listeners.QueryGroupRequestOperationListener; +import org.opensearch.wlm.tracker.QueryGroupResourceUsageTrackerService; import javax.net.ssl.SNIHostName; @@ -1022,8 +1027,30 @@ protected Node( List identityAwarePlugins = pluginsService.filterPlugins(IdentityAwarePlugin.class); identityService.initializeIdentityAwarePlugins(identityAwarePlugins); - final QueryGroupService queryGroupService = new QueryGroupService(); // We will need to replace this with actual instance of the - // queryGroupService + final QueryGroupResourceUsageTrackerService queryGroupResourceUsageTrackerService = new QueryGroupResourceUsageTrackerService( + taskResourceTrackingService + ); + final WorkloadManagementSettings workloadManagementSettings = new WorkloadManagementSettings( + settings, + settingsModule.getClusterSettings() + ); + + final QueryGroupsStateAccessor queryGroupsStateAccessor = new QueryGroupsStateAccessor(); + + final QueryGroupService queryGroupService = new QueryGroupService( + new QueryGroupTaskCancellationService( + workloadManagementSettings, + new MaximumResourceTaskSelectionStrategy(), + queryGroupResourceUsageTrackerService, + queryGroupsStateAccessor + ), + clusterService, + threadPool, + workloadManagementSettings, + queryGroupsStateAccessor + ); + taskResourceTrackingService.addTaskCompletionListener(queryGroupService); + final QueryGroupRequestOperationListener queryGroupRequestOperationListener = new QueryGroupRequestOperationListener( queryGroupService, threadPool @@ -1089,7 +1116,7 @@ protected Node( WorkloadManagementTransportInterceptor workloadManagementTransportInterceptor = new WorkloadManagementTransportInterceptor( threadPool, - new QueryGroupService() // We will need to replace this with actual implementation + queryGroupService ); final Collection secureSettingsFactories = pluginsService.filterPlugins(Plugin.class) @@ -1184,7 +1211,8 @@ protected Node( searchBackpressureSettings, taskResourceTrackingService, threadPool, - transportService.getTaskManager() + transportService.getTaskManager(), + queryGroupService ); final SegmentReplicationStatsTracker segmentReplicationStatsTracker = new SegmentReplicationStatsTracker(indicesService); @@ -1396,6 +1424,7 @@ protected Node( b.bind(IndexingPressureService.class).toInstance(indexingPressureService); b.bind(TaskResourceTrackingService.class).toInstance(taskResourceTrackingService); b.bind(SearchBackpressureService.class).toInstance(searchBackpressureService); + b.bind(QueryGroupService.class).toInstance(queryGroupService); b.bind(AdmissionControlService.class).toInstance(admissionControlService); b.bind(UsageService.class).toInstance(usageService); b.bind(AggregationUsageService.class).toInstance(searchModule.getValuesSourceRegistry().getUsageService()); @@ -1589,6 +1618,7 @@ public Node start() throws NodeValidationException { nodeService.getMonitorService().start(); nodeService.getSearchBackpressureService().start(); nodeService.getTaskCancellationMonitoringService().start(); + injector.getInstance(QueryGroupService.class).start(); final ClusterService clusterService = injector.getInstance(ClusterService.class); @@ -1762,6 +1792,7 @@ private Node stop() { injector.getInstance(FsHealthService.class).stop(); injector.getInstance(NodeResourceUsageTracker.class).stop(); injector.getInstance(ResourceUsageCollectorService.class).stop(); + injector.getInstance(QueryGroupService.class).stop(); nodeService.getMonitorService().stop(); nodeService.getSearchBackpressureService().stop(); injector.getInstance(GatewayService.class).stop(); diff --git a/server/src/main/java/org/opensearch/search/backpressure/SearchBackpressureService.java b/server/src/main/java/org/opensearch/search/backpressure/SearchBackpressureService.java index e98046ba1dede..d6eac45ca0288 100644 --- a/server/src/main/java/org/opensearch/search/backpressure/SearchBackpressureService.java +++ b/server/src/main/java/org/opensearch/search/backpressure/SearchBackpressureService.java @@ -42,6 +42,7 @@ import org.opensearch.tasks.TaskResourceTrackingService.TaskCompletionListener; import org.opensearch.threadpool.Scheduler; import org.opensearch.threadpool.ThreadPool; +import org.opensearch.wlm.QueryGroupService; import org.opensearch.wlm.ResourceType; import java.io.IOException; @@ -86,12 +87,14 @@ public class SearchBackpressureService extends AbstractLifecycleComponent implem private final Map, SearchBackpressureState> searchBackpressureStates; private final TaskManager taskManager; + private final QueryGroupService queryGroupService; public SearchBackpressureService( SearchBackpressureSettings settings, TaskResourceTrackingService taskResourceTrackingService, ThreadPool threadPool, - TaskManager taskManager + TaskManager taskManager, + QueryGroupService queryGroupService ) { this(settings, taskResourceTrackingService, threadPool, System::nanoTime, new NodeDuressTrackers(new EnumMap<>(ResourceType.class) { { @@ -131,7 +134,8 @@ public SearchBackpressureService( settings.getClusterSettings(), SearchShardTaskSettings.SETTING_HEAP_MOVING_AVERAGE_WINDOW_SIZE ), - taskManager + taskManager, + queryGroupService ); } @@ -143,7 +147,8 @@ public SearchBackpressureService( NodeDuressTrackers nodeDuressTrackers, TaskResourceUsageTrackers searchTaskTrackers, TaskResourceUsageTrackers searchShardTaskTrackers, - TaskManager taskManager + TaskManager taskManager, + QueryGroupService queryGroupService ) { this.settings = settings; this.taskResourceTrackingService = taskResourceTrackingService; @@ -151,6 +156,7 @@ public SearchBackpressureService( this.threadPool = threadPool; this.nodeDuressTrackers = nodeDuressTrackers; this.taskManager = taskManager; + this.queryGroupService = queryGroupService; this.searchBackpressureStates = Map.of( SearchTask.class, @@ -346,6 +352,7 @@ List getTa .stream() .filter(type::isInstance) .map(type::cast) + .filter(queryGroupService::shouldSBPHandle) .collect(Collectors.toUnmodifiableList()); } diff --git a/server/src/main/java/org/opensearch/wlm/QueryGroupService.java b/server/src/main/java/org/opensearch/wlm/QueryGroupService.java index 6545598dd9951..cda5916db26f3 100644 --- a/server/src/main/java/org/opensearch/wlm/QueryGroupService.java +++ b/server/src/main/java/org/opensearch/wlm/QueryGroupService.java @@ -8,36 +8,192 @@ package org.opensearch.wlm; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.action.search.SearchShardTask; +import org.opensearch.cluster.ClusterChangedEvent; +import org.opensearch.cluster.ClusterStateListener; +import org.opensearch.cluster.metadata.Metadata; +import org.opensearch.cluster.metadata.QueryGroup; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.lifecycle.AbstractLifecycleComponent; import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException; +import org.opensearch.monitor.jvm.JvmStats; +import org.opensearch.monitor.process.ProcessProbe; +import org.opensearch.search.backpressure.trackers.NodeDuressTrackers; +import org.opensearch.search.backpressure.trackers.NodeDuressTrackers.NodeDuressTracker; +import org.opensearch.tasks.Task; +import org.opensearch.tasks.TaskResourceTrackingService; +import org.opensearch.threadpool.Scheduler; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.wlm.cancellation.QueryGroupTaskCancellationService; import org.opensearch.wlm.stats.QueryGroupState; import org.opensearch.wlm.stats.QueryGroupStats; import org.opensearch.wlm.stats.QueryGroupStats.QueryGroupStatsHolder; +import java.io.IOException; import java.util.HashMap; +import java.util.HashSet; import java.util.Map; +import java.util.Optional; +import java.util.Set; + +import static org.opensearch.wlm.tracker.QueryGroupResourceUsageTrackerService.TRACKED_RESOURCES; /** * As of now this is a stub and main implementation PR will be raised soon.Coming PR will collate these changes with core QueryGroupService changes */ -public class QueryGroupService { - // This map does not need to be concurrent since we will process the cluster state change serially and update - // this map with new additions and deletions of entries. QueryGroupState is thread safe - private final Map queryGroupStateMap; +public class QueryGroupService extends AbstractLifecycleComponent + implements + ClusterStateListener, + TaskResourceTrackingService.TaskCompletionListener { + + private static final Logger logger = LogManager.getLogger(QueryGroupService.class); + + private final QueryGroupTaskCancellationService taskCancellationService; + private volatile Scheduler.Cancellable scheduledFuture; + private final ThreadPool threadPool; + private final ClusterService clusterService; + private final WorkloadManagementSettings workloadManagementSettings; + private Set activeQueryGroups; + private final Set deletedQueryGroups; + private final NodeDuressTrackers nodeDuressTrackers; + private final QueryGroupsStateAccessor queryGroupsStateAccessor; + + public QueryGroupService( + QueryGroupTaskCancellationService taskCancellationService, + ClusterService clusterService, + ThreadPool threadPool, + WorkloadManagementSettings workloadManagementSettings, + QueryGroupsStateAccessor queryGroupsStateAccessor + ) { + + this( + taskCancellationService, + clusterService, + threadPool, + workloadManagementSettings, + new NodeDuressTrackers( + Map.of( + ResourceType.CPU, + new NodeDuressTracker( + () -> workloadManagementSettings.getNodeLevelCpuCancellationThreshold() < ProcessProbe.getInstance() + .getProcessCpuPercent() / 100.0, + workloadManagementSettings::getDuressStreak + ), + ResourceType.MEMORY, + new NodeDuressTracker( + () -> workloadManagementSettings.getNodeLevelMemoryCancellationThreshold() <= JvmStats.jvmStats() + .getMem() + .getHeapUsedPercent() / 100.0, + workloadManagementSettings::getDuressStreak + ) + ) + ), + queryGroupsStateAccessor, + new HashSet<>(), + new HashSet<>() + ); + } + + public QueryGroupService( + QueryGroupTaskCancellationService taskCancellationService, + ClusterService clusterService, + ThreadPool threadPool, + WorkloadManagementSettings workloadManagementSettings, + NodeDuressTrackers nodeDuressTrackers, + QueryGroupsStateAccessor queryGroupsStateAccessor, + Set activeQueryGroups, + Set deletedQueryGroups + ) { + this.taskCancellationService = taskCancellationService; + this.clusterService = clusterService; + this.threadPool = threadPool; + this.workloadManagementSettings = workloadManagementSettings; + this.nodeDuressTrackers = nodeDuressTrackers; + this.activeQueryGroups = activeQueryGroups; + this.deletedQueryGroups = deletedQueryGroups; + this.queryGroupsStateAccessor = queryGroupsStateAccessor; + activeQueryGroups.forEach(queryGroup -> this.queryGroupsStateAccessor.addNewQueryGroup(queryGroup.get_id())); + this.queryGroupsStateAccessor.addNewQueryGroup(QueryGroupTask.DEFAULT_QUERY_GROUP_ID_SUPPLIER.get()); + this.clusterService.addListener(this); + } + + /** + * run at regular interval + */ + void doRun() { + if (workloadManagementSettings.getWlmMode() == WlmMode.DISABLED) { + return; + } + taskCancellationService.cancelTasks(nodeDuressTrackers::isNodeInDuress, activeQueryGroups, deletedQueryGroups); + taskCancellationService.pruneDeletedQueryGroups(deletedQueryGroups); + } + + /** + * {@link AbstractLifecycleComponent} lifecycle method + */ + @Override + protected void doStart() { + scheduledFuture = threadPool.scheduleWithFixedDelay(() -> { + try { + doRun(); + } catch (Exception e) { + logger.debug("Exception occurred in Query Sandbox service", e); + } + }, this.workloadManagementSettings.getQueryGroupServiceRunInterval(), ThreadPool.Names.GENERIC); + } - public QueryGroupService() { - this(new HashMap<>()); + @Override + protected void doStop() { + if (scheduledFuture != null) { + scheduledFuture.cancel(); + } } - public QueryGroupService(Map queryGroupStateMap) { - this.queryGroupStateMap = queryGroupStateMap; + @Override + protected void doClose() throws IOException {} + + @Override + public void clusterChanged(ClusterChangedEvent event) { + // Retrieve the current and previous cluster states + Metadata previousMetadata = event.previousState().metadata(); + Metadata currentMetadata = event.state().metadata(); + + // Extract the query groups from both the current and previous cluster states + Map previousQueryGroups = previousMetadata.queryGroups(); + Map currentQueryGroups = currentMetadata.queryGroups(); + + // Detect new query groups added in the current cluster state + for (String queryGroupName : currentQueryGroups.keySet()) { + if (!previousQueryGroups.containsKey(queryGroupName)) { + // New query group detected + QueryGroup newQueryGroup = currentQueryGroups.get(queryGroupName); + // Perform any necessary actions with the new query group + queryGroupsStateAccessor.addNewQueryGroup(newQueryGroup.get_id()); + } + } + + // Detect query groups deleted in the current cluster state + for (String queryGroupName : previousQueryGroups.keySet()) { + if (!currentQueryGroups.containsKey(queryGroupName)) { + // Query group deleted + QueryGroup deletedQueryGroup = previousQueryGroups.get(queryGroupName); + // Perform any necessary actions with the deleted query group + this.deletedQueryGroups.add(deletedQueryGroup); + queryGroupsStateAccessor.removeQueryGroup(deletedQueryGroup.get_id()); + } + } + this.activeQueryGroups = new HashSet<>(currentMetadata.queryGroups().values()); } /** * updates the failure stats for the query group + * * @param queryGroupId query group identifier */ public void incrementFailuresFor(final String queryGroupId) { - QueryGroupState queryGroupState = queryGroupStateMap.get(queryGroupId); + QueryGroupState queryGroupState = queryGroupsStateAccessor.getQueryGroupState(queryGroupId); // This can happen if the request failed for a deleted query group // or new queryGroup is being created and has not been acknowledged yet if (queryGroupState == null) { @@ -47,12 +203,11 @@ public void incrementFailuresFor(final String queryGroupId) { } /** - * * @return node level query group stats */ public QueryGroupStats nodeStats() { final Map statsHolderMap = new HashMap<>(); - for (Map.Entry queryGroupsState : queryGroupStateMap.entrySet()) { + for (Map.Entry queryGroupsState : queryGroupsStateAccessor.getQueryGroupStateMap().entrySet()) { final String queryGroupId = queryGroupsState.getKey(); final QueryGroupState currentState = queryGroupsState.getValue(); @@ -63,18 +218,113 @@ public QueryGroupStats nodeStats() { } /** - * * @param queryGroupId query group identifier */ public void rejectIfNeeded(String queryGroupId) { - if (queryGroupId == null) return; - boolean reject = false; - final StringBuilder reason = new StringBuilder(); - // TODO: At this point this is dummy and we need to decide whether to cancel the request based on last - // reported resource usage for the queryGroup. We also need to increment the rejection count here for the - // query group - if (reject) { - throw new OpenSearchRejectedExecutionException("QueryGroup " + queryGroupId + " is already contended." + reason.toString()); + if (workloadManagementSettings.getWlmMode() != WlmMode.ENABLED) { + return; + } + + if (queryGroupId == null || queryGroupId.equals(QueryGroupTask.DEFAULT_QUERY_GROUP_ID_SUPPLIER.get())) return; + QueryGroupState queryGroupState = queryGroupsStateAccessor.getQueryGroupState(queryGroupId); + + // This can happen if the request failed for a deleted query group + // or new queryGroup is being created and has not been acknowledged yet or invalid query group id + if (queryGroupState == null) { + return; + } + + // rejections will not happen for SOFT mode QueryGroups + Optional optionalQueryGroup = activeQueryGroups.stream().filter(x -> x.get_id().equals(queryGroupId)).findFirst(); + + if (optionalQueryGroup.isPresent() && optionalQueryGroup.get().getResiliencyMode() == MutableQueryGroupFragment.ResiliencyMode.SOFT) + return; + + optionalQueryGroup.ifPresent(queryGroup -> { + boolean reject = false; + final StringBuilder reason = new StringBuilder(); + for (ResourceType resourceType : TRACKED_RESOURCES) { + if (queryGroup.getResourceLimits().containsKey(resourceType)) { + final double threshold = getNormalisedRejectionThreshold( + queryGroup.getResourceLimits().get(resourceType), + resourceType + ); + final double lastRecordedUsage = queryGroupState.getResourceState().get(resourceType).getLastRecordedUsage(); + if (threshold < lastRecordedUsage) { + reject = true; + reason.append(resourceType) + .append(" limit is breaching for ENFORCED type QueryGroup: (") + .append(threshold) + .append(" < ") + .append(lastRecordedUsage) + .append("). "); + queryGroupState.getResourceState().get(resourceType).rejections.inc(); + // should not double count even if both the resource limits are breaching + break; + } + } + } + if (reject) { + queryGroupState.totalRejections.inc(); + throw new OpenSearchRejectedExecutionException( + "QueryGroup " + queryGroupId + " is already contended. " + reason.toString() + ); + } + }); + } + + private double getNormalisedRejectionThreshold(double limit, ResourceType resourceType) { + if (resourceType == ResourceType.CPU) { + return limit * workloadManagementSettings.getNodeLevelCpuRejectionThreshold(); + } else if (resourceType == ResourceType.MEMORY) { + return limit * workloadManagementSettings.getNodeLevelMemoryRejectionThreshold(); + } + throw new IllegalArgumentException(resourceType + " is not supported in WLM yet"); + } + + public Set getActiveQueryGroups() { + return activeQueryGroups; + } + + public Set getDeletedQueryGroups() { + return deletedQueryGroups; + } + + /** + * This method determines whether the task should be accounted by SBP if both features co-exist + * @param t QueryGroupTask + * @return whether or not SBP handle it + */ + public boolean shouldSBPHandle(Task t) { + QueryGroupTask task = (QueryGroupTask) t; + boolean isInvalidQueryGroupTask = true; + if (!task.getQueryGroupId().equals(QueryGroupTask.DEFAULT_QUERY_GROUP_ID_SUPPLIER.get())) { + isInvalidQueryGroupTask = activeQueryGroups.stream() + .noneMatch(queryGroup -> queryGroup.get_id().equals(task.getQueryGroupId())); + } + return workloadManagementSettings.getWlmMode() != WlmMode.ENABLED || isInvalidQueryGroupTask; + } + + @Override + public void onTaskCompleted(Task task) { + if (!(task instanceof QueryGroupTask)) { + return; + } + final QueryGroupTask queryGroupTask = (QueryGroupTask) task; + String queryGroupId = queryGroupTask.getQueryGroupId(); + + // set the default queryGroupId if not existing in the active query groups + String finalQueryGroupId = queryGroupId; + boolean exists = activeQueryGroups.stream().anyMatch(queryGroup -> queryGroup.get_id().equals(finalQueryGroupId)); + + if (!exists) { + queryGroupId = QueryGroupTask.DEFAULT_QUERY_GROUP_ID_SUPPLIER.get(); + } + + if (task instanceof SearchShardTask) { + queryGroupsStateAccessor.getQueryGroupState(queryGroupId).shardCompletions.inc(); + } else { + queryGroupsStateAccessor.getQueryGroupState(queryGroupId).completions.inc(); } } } diff --git a/server/src/main/java/org/opensearch/wlm/QueryGroupTask.java b/server/src/main/java/org/opensearch/wlm/QueryGroupTask.java index a1cb766579d43..97c48bd828978 100644 --- a/server/src/main/java/org/opensearch/wlm/QueryGroupTask.java +++ b/server/src/main/java/org/opensearch/wlm/QueryGroupTask.java @@ -17,7 +17,6 @@ import org.opensearch.tasks.CancellableTask; import java.util.Map; -import java.util.Optional; import java.util.function.LongSupplier; import java.util.function.Supplier; @@ -82,9 +81,11 @@ public final String getQueryGroupId() { * @param threadContext current threadContext */ public final void setQueryGroupId(final ThreadContext threadContext) { - this.queryGroupId = Optional.ofNullable(threadContext) - .map(threadContext1 -> threadContext1.getHeader(QUERY_GROUP_ID_HEADER)) - .orElse(DEFAULT_QUERY_GROUP_ID_SUPPLIER.get()); + if (threadContext != null && threadContext.getHeader(QUERY_GROUP_ID_HEADER) != null) { + this.queryGroupId = threadContext.getHeader(QUERY_GROUP_ID_HEADER); + } else { + this.queryGroupId = DEFAULT_QUERY_GROUP_ID_SUPPLIER.get(); + } } public long getElapsedTime() { diff --git a/server/src/main/java/org/opensearch/wlm/QueryGroupsStateAccessor.java b/server/src/main/java/org/opensearch/wlm/QueryGroupsStateAccessor.java new file mode 100644 index 0000000000000..7f93e41f12092 --- /dev/null +++ b/server/src/main/java/org/opensearch/wlm/QueryGroupsStateAccessor.java @@ -0,0 +1,64 @@ +/* + * 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.wlm; + +import org.opensearch.wlm.stats.QueryGroupState; + +import java.util.HashMap; +import java.util.Map; + +/** + * This class is used to decouple {@link QueryGroupService} and {@link org.opensearch.wlm.cancellation.QueryGroupTaskCancellationService} to share the + * {@link QueryGroupState}s + */ +public class QueryGroupsStateAccessor { + // This map does not need to be concurrent since we will process the cluster state change serially and update + // this map with new additions and deletions of entries. QueryGroupState is thread safe + private final Map queryGroupStateMap; + + public QueryGroupsStateAccessor() { + this(new HashMap<>()); + } + + public QueryGroupsStateAccessor(Map queryGroupStateMap) { + this.queryGroupStateMap = queryGroupStateMap; + } + + /** + * returns the query groups state + */ + public Map getQueryGroupStateMap() { + return queryGroupStateMap; + } + + /** + * return QueryGroupState for the given queryGroupId + * @param queryGroupId + * @return QueryGroupState for the given queryGroupId, if id is invalid return default query group state + */ + public QueryGroupState getQueryGroupState(String queryGroupId) { + return queryGroupStateMap.getOrDefault(queryGroupId, queryGroupStateMap.get(QueryGroupTask.DEFAULT_QUERY_GROUP_ID_SUPPLIER.get())); + } + + /** + * adds new QueryGroupState against given queryGroupId + * @param queryGroupId + */ + public void addNewQueryGroup(String queryGroupId) { + this.queryGroupStateMap.putIfAbsent(queryGroupId, new QueryGroupState()); + } + + /** + * removes QueryGroupState against given queryGroupId + * @param queryGroupId + */ + public void removeQueryGroup(String queryGroupId) { + this.queryGroupStateMap.remove(queryGroupId); + } +} diff --git a/server/src/main/java/org/opensearch/wlm/WlmMode.java b/server/src/main/java/org/opensearch/wlm/WlmMode.java new file mode 100644 index 0000000000000..40407525cc24d --- /dev/null +++ b/server/src/main/java/org/opensearch/wlm/WlmMode.java @@ -0,0 +1,40 @@ +/* + * 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.wlm; + +import org.opensearch.common.annotation.PublicApi; + +/** + * Enum to hold the values whether wlm is enabled or not + */ +@PublicApi(since = "2.18.0") +public enum WlmMode { + ENABLED("enabled"), + MONITOR_ONLY("monitor_only"), + DISABLED("disabled"); + + private final String name; + + WlmMode(String name) { + this.name = name; + } + + public String getName() { + return name; + } + + public static WlmMode fromName(String name) { + for (WlmMode wlmMode : values()) { + if (wlmMode.getName().equals(name)) { + return wlmMode; + } + } + throw new IllegalArgumentException(name + " is an invalid WlmMode"); + } +} diff --git a/server/src/main/java/org/opensearch/wlm/WorkloadManagementSettings.java b/server/src/main/java/org/opensearch/wlm/WorkloadManagementSettings.java index b3577c1b3219d..af25eedd7eed5 100644 --- a/server/src/main/java/org/opensearch/wlm/WorkloadManagementSettings.java +++ b/server/src/main/java/org/opensearch/wlm/WorkloadManagementSettings.java @@ -12,6 +12,7 @@ import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.Setting; import org.opensearch.common.settings.Settings; +import org.opensearch.common.unit.TimeValue; /** * Main class to declare Workload Management related settings @@ -22,16 +23,66 @@ public class WorkloadManagementSettings { private static final Double DEFAULT_NODE_LEVEL_MEMORY_CANCELLATION_THRESHOLD = 0.9; private static final Double DEFAULT_NODE_LEVEL_CPU_REJECTION_THRESHOLD = 0.8; private static final Double DEFAULT_NODE_LEVEL_CPU_CANCELLATION_THRESHOLD = 0.9; + private static final Long DEFAULT_QUERYGROUP_SERVICE_RUN_INTERVAL_MILLIS = 1000L; public static final double NODE_LEVEL_MEMORY_CANCELLATION_THRESHOLD_MAX_VALUE = 0.95; public static final double NODE_LEVEL_MEMORY_REJECTION_THRESHOLD_MAX_VALUE = 0.9; public static final double NODE_LEVEL_CPU_CANCELLATION_THRESHOLD_MAX_VALUE = 0.95; public static final double NODE_LEVEL_CPU_REJECTION_THRESHOLD_MAX_VALUE = 0.9; + public static final String DEFAULT_WLM_MODE = "monitor_only"; private Double nodeLevelMemoryCancellationThreshold; private Double nodeLevelMemoryRejectionThreshold; private Double nodeLevelCpuCancellationThreshold; private Double nodeLevelCpuRejectionThreshold; + /** + * Setting name for QueryGroupService node duress streak + */ + public static final String QUERYGROUP_DURESS_STREAK_SETTING_NAME = "wlm.query_group.duress_streak"; + private int duressStreak; + public static final Setting QUERYGROUP_SERVICE_DURESS_STREAK_SETTING = Setting.intSetting( + QUERYGROUP_DURESS_STREAK_SETTING_NAME, + 3, + 3, + Setting.Property.Dynamic, + Setting.Property.NodeScope + ); + + /** + * Setting name for Query Group Service run interval + */ + public static final String QUERYGROUP_ENFORCEMENT_INTERVAL_SETTING_NAME = "wlm.query_group.enforcement_interval"; + + private TimeValue queryGroupServiceRunInterval; + /** + * Setting to control the run interval of Query Group Service + */ + public static final Setting QUERYGROUP_SERVICE_RUN_INTERVAL_SETTING = Setting.longSetting( + QUERYGROUP_ENFORCEMENT_INTERVAL_SETTING_NAME, + DEFAULT_QUERYGROUP_SERVICE_RUN_INTERVAL_MILLIS, + 1000, + Setting.Property.Dynamic, + Setting.Property.NodeScope + ); + + /** + * WLM mode setting name + */ + public static final String WLM_MODE_SETTING_NAME = "wlm.query_group.mode"; + + private volatile WlmMode wlmMode; + + /** + * WLM mode setting, which determines which mode WLM is operating in + */ + public static final Setting WLM_MODE_SETTING = new Setting( + WLM_MODE_SETTING_NAME, + DEFAULT_WLM_MODE, + WlmMode::fromName, + Setting.Property.Dynamic, + Setting.Property.NodeScope + ); + /** * Setting name for node level memory based rejection threshold for QueryGroup service */ @@ -91,10 +142,13 @@ public class WorkloadManagementSettings { * @param clusterSettings - QueryGroup cluster settings */ public WorkloadManagementSettings(Settings settings, ClusterSettings clusterSettings) { + this.wlmMode = WLM_MODE_SETTING.get(settings); nodeLevelMemoryCancellationThreshold = NODE_LEVEL_MEMORY_CANCELLATION_THRESHOLD.get(settings); nodeLevelMemoryRejectionThreshold = NODE_LEVEL_MEMORY_REJECTION_THRESHOLD.get(settings); nodeLevelCpuCancellationThreshold = NODE_LEVEL_CPU_CANCELLATION_THRESHOLD.get(settings); nodeLevelCpuRejectionThreshold = NODE_LEVEL_CPU_REJECTION_THRESHOLD.get(settings); + this.queryGroupServiceRunInterval = TimeValue.timeValueMillis(QUERYGROUP_SERVICE_RUN_INTERVAL_SETTING.get(settings)); + duressStreak = QUERYGROUP_SERVICE_DURESS_STREAK_SETTING.get(settings); ensureRejectionThresholdIsLessThanCancellation( nodeLevelMemoryRejectionThreshold, @@ -113,6 +167,57 @@ public WorkloadManagementSettings(Settings settings, ClusterSettings clusterSett clusterSettings.addSettingsUpdateConsumer(NODE_LEVEL_MEMORY_REJECTION_THRESHOLD, this::setNodeLevelMemoryRejectionThreshold); clusterSettings.addSettingsUpdateConsumer(NODE_LEVEL_CPU_CANCELLATION_THRESHOLD, this::setNodeLevelCpuCancellationThreshold); clusterSettings.addSettingsUpdateConsumer(NODE_LEVEL_CPU_REJECTION_THRESHOLD, this::setNodeLevelCpuRejectionThreshold); + clusterSettings.addSettingsUpdateConsumer(WLM_MODE_SETTING, this::setWlmMode); + clusterSettings.addSettingsUpdateConsumer(QUERYGROUP_SERVICE_RUN_INTERVAL_SETTING, this::setQueryGroupServiceRunInterval); + clusterSettings.addSettingsUpdateConsumer(QUERYGROUP_SERVICE_DURESS_STREAK_SETTING, this::setDuressStreak); + } + + /** + * node duress streak getter + * @return current duressStreak value + */ + public int getDuressStreak() { + return duressStreak; + } + + /** + * node duress streak setter + * @param duressStreak new value + */ + private void setDuressStreak(int duressStreak) { + this.duressStreak = duressStreak; + } + + /** + * queryGroupServiceRunInterval setter + * @param newIntervalInMillis new value + */ + private void setQueryGroupServiceRunInterval(long newIntervalInMillis) { + this.queryGroupServiceRunInterval = TimeValue.timeValueMillis(newIntervalInMillis); + } + + /** + * queryGroupServiceRunInterval getter + * @return current queryGroupServiceRunInterval value + */ + public TimeValue getQueryGroupServiceRunInterval() { + return this.queryGroupServiceRunInterval; + } + + /** + * WlmMode setter + * @param mode new mode value + */ + private void setWlmMode(final WlmMode mode) { + this.wlmMode = mode; + } + + /** + * WlmMode getter + * @return the current wlmMode + */ + public WlmMode getWlmMode() { + return this.wlmMode; } /** diff --git a/server/src/main/java/org/opensearch/wlm/cancellation/QueryGroupTaskCancellationService.java b/server/src/main/java/org/opensearch/wlm/cancellation/QueryGroupTaskCancellationService.java index a2c97c8d8635b..e82a19c5f7af2 100644 --- a/server/src/main/java/org/opensearch/wlm/cancellation/QueryGroupTaskCancellationService.java +++ b/server/src/main/java/org/opensearch/wlm/cancellation/QueryGroupTaskCancellationService.java @@ -8,20 +8,26 @@ package org.opensearch.wlm.cancellation; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; import org.opensearch.cluster.metadata.QueryGroup; -import org.opensearch.tasks.CancellableTask; import org.opensearch.tasks.TaskCancellation; import org.opensearch.wlm.MutableQueryGroupFragment.ResiliencyMode; import org.opensearch.wlm.QueryGroupLevelResourceUsageView; import org.opensearch.wlm.QueryGroupTask; +import org.opensearch.wlm.QueryGroupsStateAccessor; import org.opensearch.wlm.ResourceType; +import org.opensearch.wlm.WlmMode; import org.opensearch.wlm.WorkloadManagementSettings; +import org.opensearch.wlm.stats.QueryGroupState; import org.opensearch.wlm.tracker.QueryGroupResourceUsageTrackerService; import java.util.ArrayList; import java.util.Collection; +import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.function.BooleanSupplier; import java.util.function.Consumer; import java.util.stream.Collectors; @@ -47,46 +53,78 @@ */ public class QueryGroupTaskCancellationService { public static final double MIN_VALUE = 1e-9; + private static final Logger log = LogManager.getLogger(QueryGroupTaskCancellationService.class); private final WorkloadManagementSettings workloadManagementSettings; private final TaskSelectionStrategy taskSelectionStrategy; private final QueryGroupResourceUsageTrackerService resourceUsageTrackerService; // a map of QueryGroupId to its corresponding QueryGroupLevelResourceUsageView object Map queryGroupLevelResourceUsageViews; - private final Collection activeQueryGroups; - private final Collection deletedQueryGroups; + private final QueryGroupsStateAccessor queryGroupStateAccessor; public QueryGroupTaskCancellationService( WorkloadManagementSettings workloadManagementSettings, TaskSelectionStrategy taskSelectionStrategy, QueryGroupResourceUsageTrackerService resourceUsageTrackerService, - Collection activeQueryGroups, - Collection deletedQueryGroups + QueryGroupsStateAccessor queryGroupStateAccessor ) { this.workloadManagementSettings = workloadManagementSettings; this.taskSelectionStrategy = taskSelectionStrategy; this.resourceUsageTrackerService = resourceUsageTrackerService; - this.activeQueryGroups = activeQueryGroups; - this.deletedQueryGroups = deletedQueryGroups; + this.queryGroupStateAccessor = queryGroupStateAccessor; } /** * Cancel tasks based on the implemented strategy. */ - public final void cancelTasks(BooleanSupplier isNodeInDuress) { + public void cancelTasks( + BooleanSupplier isNodeInDuress, + Collection activeQueryGroups, + Collection deletedQueryGroups + ) { queryGroupLevelResourceUsageViews = resourceUsageTrackerService.constructQueryGroupLevelUsageViews(); // cancel tasks from QueryGroups that are in Enforced mode that are breaching their resource limits - cancelTasks(ResiliencyMode.ENFORCED); + cancelTasks(ResiliencyMode.ENFORCED, activeQueryGroups); // if the node is in duress, cancel tasks accordingly. - handleNodeDuress(isNodeInDuress); + handleNodeDuress(isNodeInDuress, activeQueryGroups, deletedQueryGroups); + + updateResourceUsageInQueryGroupState(activeQueryGroups); + } + + private void updateResourceUsageInQueryGroupState(Collection activeQueryGroups) { + Set isSearchWorkloadRunning = new HashSet<>(); + for (Map.Entry queryGroupLevelResourceUsageViewEntry : queryGroupLevelResourceUsageViews + .entrySet()) { + isSearchWorkloadRunning.add(queryGroupLevelResourceUsageViewEntry.getKey()); + QueryGroupState queryGroupState = getQueryGroupState(queryGroupLevelResourceUsageViewEntry.getKey()); + TRACKED_RESOURCES.forEach(resourceType -> { + final double currentUsage = queryGroupLevelResourceUsageViewEntry.getValue().getResourceUsageData().get(resourceType); + queryGroupState.getResourceState().get(resourceType).setLastRecordedUsage(currentUsage); + }); + } + + activeQueryGroups.forEach(queryGroup -> { + if (!isSearchWorkloadRunning.contains(queryGroup.get_id())) { + TRACKED_RESOURCES.forEach( + resourceType -> getQueryGroupState(queryGroup.get_id()).getResourceState().get(resourceType).setLastRecordedUsage(0.0) + ); + } + }); } - private void handleNodeDuress(BooleanSupplier isNodeInDuress) { + private void handleNodeDuress( + BooleanSupplier isNodeInDuress, + Collection activeQueryGroups, + Collection deletedQueryGroups + ) { if (!isNodeInDuress.getAsBoolean()) { return; } // List of tasks to be executed in order if the node is in duress - List> duressActions = List.of(v -> cancelTasksFromDeletedQueryGroups(), v -> cancelTasks(ResiliencyMode.SOFT)); + List> duressActions = List.of( + v -> cancelTasksFromDeletedQueryGroups(deletedQueryGroups), + v -> cancelTasks(ResiliencyMode.SOFT, activeQueryGroups) + ); for (Consumer duressAction : duressActions) { if (!isNodeInDuress.getAsBoolean()) { @@ -96,8 +134,8 @@ private void handleNodeDuress(BooleanSupplier isNodeInDuress) { } } - private void cancelTasksFromDeletedQueryGroups() { - cancelTasks(getAllCancellableTasks(this.deletedQueryGroups)); + private void cancelTasksFromDeletedQueryGroups(Collection deletedQueryGroups) { + cancelTasks(getAllCancellableTasks(deletedQueryGroups)); } /** @@ -105,9 +143,9 @@ private void cancelTasksFromDeletedQueryGroups() { * * @return List of tasks that can be cancelled */ - List getAllCancellableTasks(ResiliencyMode resiliencyMode) { + List getAllCancellableTasks(ResiliencyMode resiliencyMode, Collection queryGroups) { return getAllCancellableTasks( - activeQueryGroups.stream().filter(queryGroup -> queryGroup.getResiliencyMode() == resiliencyMode).collect(Collectors.toList()) + queryGroups.stream().filter(queryGroup -> queryGroup.getResiliencyMode() == resiliencyMode).collect(Collectors.toList()) ); } @@ -118,6 +156,7 @@ List getAllCancellableTasks(ResiliencyMode resiliencyMode) { */ List getAllCancellableTasks(Collection queryGroups) { List taskCancellations = new ArrayList<>(); + final List onCancelCallbacks = new ArrayList<>(); for (QueryGroup queryGroup : queryGroups) { final List reasons = new ArrayList<>(); List selectedTasks = new ArrayList<>(); @@ -127,8 +166,7 @@ List getAllCancellableTasks(Collection queryGroups .calculateResourceUsage(selectedTasks); if (excessUsage > MIN_VALUE) { reasons.add(new TaskCancellation.Reason(generateReasonString(queryGroup, resourceType), 1)); - // TODO: We will need to add the cancellation callback for these resources for the queryGroup to reflect stats - + onCancelCallbacks.add(this.getResourceTypeOnCancelCallback(queryGroup.get_id(), resourceType)); // Only add tasks not already added to avoid double cancellations selectedTasks.addAll( taskSelectionStrategy.selectTasksForCancellation(getTasksFor(queryGroup), excessUsage, resourceType) @@ -140,8 +178,9 @@ List getAllCancellableTasks(Collection queryGroups } if (!reasons.isEmpty()) { + onCancelCallbacks.add(getQueryGroupState(queryGroup.get_id()).totalCancellations::inc); taskCancellations.addAll( - selectedTasks.stream().map(task -> createTaskCancellation(task, reasons)).collect(Collectors.toList()) + selectedTasks.stream().map(task -> new TaskCancellation(task, reasons, onCancelCallbacks)).collect(Collectors.toList()) ); } } @@ -164,16 +203,27 @@ private List getTasksFor(QueryGroup queryGroup) { return queryGroupLevelResourceUsageViews.get(queryGroup.get_id()).getActiveTasks(); } - private void cancelTasks(ResiliencyMode resiliencyMode) { - cancelTasks(getAllCancellableTasks(resiliencyMode)); + private void cancelTasks(ResiliencyMode resiliencyMode, Collection queryGroups) { + cancelTasks(getAllCancellableTasks(resiliencyMode, queryGroups)); } private void cancelTasks(List cancellableTasks) { - cancellableTasks.forEach(TaskCancellation::cancel); - } - private TaskCancellation createTaskCancellation(CancellableTask task, List reasons) { - return new TaskCancellation(task, reasons, List.of(this::callbackOnCancel)); + Consumer cancellationLoggingConsumer = (taskCancellation -> { + log.warn( + "Task {} is eligible for cancellation for reason {}", + taskCancellation.getTask().getId(), + taskCancellation.getReasonString() + ); + }); + Consumer cancellationConsumer = cancellationLoggingConsumer; + if (workloadManagementSettings.getWlmMode() == WlmMode.ENABLED) { + cancellationConsumer = (taskCancellation -> { + cancellationLoggingConsumer.accept(taskCancellation); + taskCancellation.cancel(); + }); + } + cancellableTasks.forEach(cancellationConsumer); } private double getExcessUsage(QueryGroup queryGroup, ResourceType resourceType) { @@ -199,7 +249,26 @@ private double getNormalisedThreshold(QueryGroup queryGroup, ResourceType resour return queryGroup.getResourceLimits().get(resourceType) * nodeLevelCancellationThreshold; } - private void callbackOnCancel() { - // TODO Implement callback logic here mostly used for Stats + private Runnable getResourceTypeOnCancelCallback(String queryGroupId, ResourceType resourceType) { + QueryGroupState queryGroupState = getQueryGroupState(queryGroupId); + return queryGroupState.getResourceState().get(resourceType).cancellations::inc; + } + + private QueryGroupState getQueryGroupState(String queryGroupId) { + assert queryGroupId != null : "queryGroupId should never be null at this point."; + + return queryGroupStateAccessor.getQueryGroupState(queryGroupId); + } + + /** + * Removes the queryGroups from deleted list if it doesn't have any tasks running + */ + public void pruneDeletedQueryGroups(Collection deletedQueryGroups) { + List currentDeletedQueryGroups = new ArrayList<>(deletedQueryGroups); + for (QueryGroup queryGroup : currentDeletedQueryGroups) { + if (queryGroupLevelResourceUsageViews.get(queryGroup.get_id()).getActiveTasks().isEmpty()) { + deletedQueryGroups.remove(queryGroup); + } + } } } diff --git a/server/src/main/java/org/opensearch/wlm/stats/QueryGroupState.java b/server/src/main/java/org/opensearch/wlm/stats/QueryGroupState.java index 376d34dd7c8ca..cbc7046a79464 100644 --- a/server/src/main/java/org/opensearch/wlm/stats/QueryGroupState.java +++ b/server/src/main/java/org/opensearch/wlm/stats/QueryGroupState.java @@ -19,14 +19,19 @@ */ public class QueryGroupState { /** - * completions at the query group level, this is a cumulative counter since the Opensearch start time + * co-ordinator level completions at the query group level, this is a cumulative counter since the Opensearch start time */ - final CounterMetric completions = new CounterMetric(); + public final CounterMetric completions = new CounterMetric(); + + /** + * shard level completions at the query group level, this is a cumulative counter since the Opensearch start time + */ + public final CounterMetric shardCompletions = new CounterMetric(); /** * rejections at the query group level, this is a cumulative counter since the OpenSearch start time */ - final CounterMetric totalRejections = new CounterMetric(); + public final CounterMetric totalRejections = new CounterMetric(); /** * this will track the cumulative failures in a query group @@ -36,7 +41,7 @@ public class QueryGroupState { /** * This will track total number of cancellations in the query group due to all resource type breaches */ - final CounterMetric totalCancellations = new CounterMetric(); + public final CounterMetric totalCancellations = new CounterMetric(); /** * This is used to store the resource type state both for CPU and MEMORY @@ -54,12 +59,20 @@ public QueryGroupState() { /** * - * @return completions in the query group + * @return co-ordinator completions in the query group */ public long getCompletions() { return completions.count(); } + /** + * + * @return shard completions in the query group + */ + public long getShardCompletions() { + return shardCompletions.count(); + } + /** * * @return rejections in the query group @@ -92,9 +105,9 @@ public Map getResourceState() { * This class holds the resource level stats for the query group */ public static class ResourceTypeState { - final ResourceType resourceType; - final CounterMetric cancellations = new CounterMetric(); - final CounterMetric rejections = new CounterMetric(); + public final ResourceType resourceType; + public final CounterMetric cancellations = new CounterMetric(); + public final CounterMetric rejections = new CounterMetric(); private double lastRecordedUsage = 0; public ResourceTypeState(ResourceType resourceType) { diff --git a/server/src/main/java/org/opensearch/wlm/stats/QueryGroupStats.java b/server/src/main/java/org/opensearch/wlm/stats/QueryGroupStats.java index 2b389c2167778..9d74201de252b 100644 --- a/server/src/main/java/org/opensearch/wlm/stats/QueryGroupStats.java +++ b/server/src/main/java/org/opensearch/wlm/stats/QueryGroupStats.java @@ -91,7 +91,9 @@ public static class QueryGroupStatsHolder implements ToXContentObject, Writeable public static final String REJECTIONS = "rejections"; public static final String TOTAL_CANCELLATIONS = "total_cancellations"; public static final String FAILURES = "failures"; + public static final String SHARD_COMPLETIONS = "shard_completions"; private long completions; + private long shardCompletions; private long rejections; private long failures; private long totalCancellations; @@ -105,11 +107,13 @@ public QueryGroupStatsHolder( long rejections, long failures, long totalCancellations, + long shardCompletions, Map resourceStats ) { this.completions = completions; this.rejections = rejections; this.failures = failures; + this.shardCompletions = shardCompletions; this.totalCancellations = totalCancellations; this.resourceStats = resourceStats; } @@ -119,6 +123,7 @@ public QueryGroupStatsHolder(StreamInput in) throws IOException { this.rejections = in.readVLong(); this.failures = in.readVLong(); this.totalCancellations = in.readVLong(); + this.shardCompletions = in.readVLong(); this.resourceStats = in.readMap((i) -> ResourceType.fromName(i.readString()), ResourceStats::new); } @@ -140,6 +145,7 @@ public static QueryGroupStatsHolder from(QueryGroupState queryGroupState) { statsHolder.rejections = queryGroupState.getTotalRejections(); statsHolder.failures = queryGroupState.getFailures(); statsHolder.totalCancellations = queryGroupState.getTotalCancellations(); + statsHolder.shardCompletions = queryGroupState.getShardCompletions(); statsHolder.resourceStats = resourceStatsMap; return statsHolder; } @@ -155,6 +161,7 @@ public static void writeTo(StreamOutput out, QueryGroupStatsHolder statsHolder) out.writeVLong(statsHolder.rejections); out.writeVLong(statsHolder.failures); out.writeVLong(statsHolder.totalCancellations); + out.writeVLong(statsHolder.shardCompletions); out.writeMap(statsHolder.resourceStats, (o, val) -> o.writeString(val.getName()), ResourceStats::writeTo); } @@ -166,6 +173,7 @@ public void writeTo(StreamOutput out) throws IOException { @Override public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { builder.field(COMPLETIONS, completions); + builder.field(SHARD_COMPLETIONS, shardCompletions); builder.field(REJECTIONS, rejections); builder.field(FAILURES, failures); builder.field(TOTAL_CANCELLATIONS, totalCancellations); @@ -187,6 +195,7 @@ public boolean equals(Object o) { QueryGroupStatsHolder that = (QueryGroupStatsHolder) o; return completions == that.completions && rejections == that.rejections + && shardCompletions == that.shardCompletions && Objects.equals(resourceStats, that.resourceStats) && failures == that.failures && totalCancellations == that.totalCancellations; @@ -194,7 +203,7 @@ public boolean equals(Object o) { @Override public int hashCode() { - return Objects.hash(completions, rejections, totalCancellations, failures, resourceStats); + return Objects.hash(completions, shardCompletions, rejections, totalCancellations, failures, resourceStats); } } diff --git a/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java b/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java index b23d9ff342139..19f7bf48d8421 100644 --- a/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java +++ b/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java @@ -47,6 +47,8 @@ public Map constructQueryGroupLevelUsa // Iterate over each QueryGroup entry for (Map.Entry> queryGroupEntry : tasksByQueryGroup.entrySet()) { + // refresh the resource stats + taskResourceTrackingService.refreshResourceStats(queryGroupEntry.getValue().toArray(new QueryGroupTask[0])); // Compute the QueryGroup resource usage final Map queryGroupUsage = new EnumMap<>(ResourceType.class); for (ResourceType resourceType : TRACKED_RESOURCES) { diff --git a/server/src/test/java/org/opensearch/search/backpressure/SearchBackpressureServiceTests.java b/server/src/test/java/org/opensearch/search/backpressure/SearchBackpressureServiceTests.java index a444eb42eac2e..8c4685bee6a0b 100644 --- a/server/src/test/java/org/opensearch/search/backpressure/SearchBackpressureServiceTests.java +++ b/server/src/test/java/org/opensearch/search/backpressure/SearchBackpressureServiceTests.java @@ -39,6 +39,8 @@ import org.opensearch.test.transport.MockTransportService; import org.opensearch.threadpool.TestThreadPool; import org.opensearch.threadpool.ThreadPool; +import org.opensearch.wlm.QueryGroupService; +import org.opensearch.wlm.QueryGroupTask; import org.opensearch.wlm.ResourceType; import org.junit.After; import org.junit.Before; @@ -75,10 +77,12 @@ public class SearchBackpressureServiceTests extends OpenSearchTestCase { MockTransportService transportService; TaskManager taskManager; ThreadPool threadPool; + QueryGroupService queryGroupService; @Before public void setup() { threadPool = new TestThreadPool(getClass().getName()); + queryGroupService = mock(QueryGroupService.class); transportService = MockTransportService.createNewService(Settings.EMPTY, Version.CURRENT, threadPool, NoopTracer.INSTANCE); transportService.start(); transportService.acceptIncomingRequests(); @@ -120,9 +124,12 @@ public void testIsNodeInDuress() { new NodeDuressTrackers(duressTrackers), new TaskResourceUsageTrackers(), new TaskResourceUsageTrackers(), - taskManager + taskManager, + queryGroupService ); + when(queryGroupService.shouldSBPHandle(any())).thenReturn(true); + // Node not in duress. cpuUsage.set(0.0); heapUsage.set(0.0); @@ -163,9 +170,12 @@ public void testTrackerStateUpdateOnSearchTaskCompletion() { new NodeDuressTrackers(new EnumMap<>(ResourceType.class)), taskResourceUsageTrackers, new TaskResourceUsageTrackers(), - taskManager + taskManager, + queryGroupService ); + when(queryGroupService.shouldSBPHandle(any())).thenReturn(true); + for (int i = 0; i < 100; i++) { // service.onTaskCompleted(new SearchTask(1, "test", "test", () -> "Test", TaskId.EMPTY_TASK_ID, new HashMap<>())); service.onTaskCompleted(createMockTaskWithResourceStats(SearchTask.class, 100, 200, i)); @@ -194,9 +204,12 @@ public void testTrackerStateUpdateOnSearchShardTaskCompletion() { new NodeDuressTrackers(new EnumMap<>(ResourceType.class)), new TaskResourceUsageTrackers(), taskResourceUsageTrackers, - taskManager + taskManager, + queryGroupService ); + when(queryGroupService.shouldSBPHandle(any())).thenReturn(true); + // Record task completions to update the tracker state. Tasks other than SearchTask & SearchShardTask are ignored. service.onTaskCompleted(createMockTaskWithResourceStats(CancellableTask.class, 100, 200, 101)); for (int i = 0; i < 100; i++) { @@ -246,9 +259,12 @@ public void testSearchTaskInFlightCancellation() { new NodeDuressTrackers(duressTrackers), taskResourceUsageTrackers, new TaskResourceUsageTrackers(), - mockTaskManager + mockTaskManager, + queryGroupService ); + when(queryGroupService.shouldSBPHandle(any())).thenReturn(true); + // Run two iterations so that node is marked 'in duress' from the third iteration onwards. service.doRun(); service.doRun(); @@ -261,7 +277,7 @@ public void testSearchTaskInFlightCancellation() { when(settings.getSearchTaskSettings()).thenReturn(searchTaskSettings); // Create a mix of low and high resource usage SearchTasks (50 low + 25 high resource usage tasks). - Map activeSearchTasks = new HashMap<>(); + Map activeSearchTasks = new HashMap<>(); for (long i = 0; i < 75; i++) { if (i % 3 == 0) { activeSearchTasks.put(i, createMockTaskWithResourceStats(SearchTask.class, 500, taskHeapUsageBytes, i)); @@ -269,6 +285,7 @@ public void testSearchTaskInFlightCancellation() { activeSearchTasks.put(i, createMockTaskWithResourceStats(SearchTask.class, 100, taskHeapUsageBytes, i)); } } + activeSearchTasks.values().forEach(task -> task.setQueryGroupId(threadPool.getThreadContext())); doReturn(activeSearchTasks).when(mockTaskResourceTrackingService).getResourceAwareTasks(); // There are 25 SearchTasks eligible for cancellation but only 5 will be cancelled (burst limit). @@ -338,9 +355,12 @@ public void testSearchShardTaskInFlightCancellation() { nodeDuressTrackers, new TaskResourceUsageTrackers(), taskResourceUsageTrackers, - mockTaskManager + mockTaskManager, + queryGroupService ); + when(queryGroupService.shouldSBPHandle(any())).thenReturn(true); + // Run two iterations so that node is marked 'in duress' from the third iteration onwards. service.doRun(); service.doRun(); @@ -353,7 +373,7 @@ public void testSearchShardTaskInFlightCancellation() { when(settings.getSearchShardTaskSettings()).thenReturn(searchShardTaskSettings); // Create a mix of low and high resource usage tasks (60 low + 15 high resource usage tasks). - Map activeSearchShardTasks = new HashMap<>(); + Map activeSearchShardTasks = new HashMap<>(); for (long i = 0; i < 75; i++) { if (i % 5 == 0) { activeSearchShardTasks.put(i, createMockTaskWithResourceStats(SearchShardTask.class, 500, taskHeapUsageBytes, i)); @@ -361,6 +381,7 @@ public void testSearchShardTaskInFlightCancellation() { activeSearchShardTasks.put(i, createMockTaskWithResourceStats(SearchShardTask.class, 100, taskHeapUsageBytes, i)); } } + activeSearchShardTasks.values().forEach(task -> task.setQueryGroupId(threadPool.getThreadContext())); doReturn(activeSearchShardTasks).when(mockTaskResourceTrackingService).getResourceAwareTasks(); // There are 15 SearchShardTasks eligible for cancellation but only 10 will be cancelled (burst limit). @@ -439,9 +460,12 @@ public void testNonCancellationOfHeapBasedTasksWhenHeapNotInDuress() { nodeDuressTrackers, taskResourceUsageTrackers, new TaskResourceUsageTrackers(), - mockTaskManager + mockTaskManager, + queryGroupService ); + when(queryGroupService.shouldSBPHandle(any())).thenReturn(true); + service.doRun(); service.doRun(); @@ -451,7 +475,7 @@ public void testNonCancellationOfHeapBasedTasksWhenHeapNotInDuress() { when(settings.getSearchTaskSettings()).thenReturn(searchTaskSettings); // Create a mix of low and high resource usage tasks (60 low + 15 high resource usage tasks). - Map activeSearchTasks = new HashMap<>(); + Map activeSearchTasks = new HashMap<>(); for (long i = 0; i < 75; i++) { if (i % 5 == 0) { activeSearchTasks.put(i, createMockTaskWithResourceStats(SearchTask.class, 500, 800, i)); @@ -459,6 +483,7 @@ public void testNonCancellationOfHeapBasedTasksWhenHeapNotInDuress() { activeSearchTasks.put(i, createMockTaskWithResourceStats(SearchTask.class, 100, 800, i)); } } + activeSearchTasks.values().forEach(task -> task.setQueryGroupId(threadPool.getThreadContext())); doReturn(activeSearchTasks).when(mockTaskResourceTrackingService).getResourceAwareTasks(); // this will trigger cancellation but these cancellation should only be cpu based @@ -534,10 +559,12 @@ public void testNonCancellationWhenSearchTrafficIsNotQualifyingForCancellation() nodeDuressTrackers, taskResourceUsageTrackers, new TaskResourceUsageTrackers(), - mockTaskManager + mockTaskManager, + queryGroupService ) ); + when(queryGroupService.shouldSBPHandle(any())).thenReturn(true); when(service.isHeapUsageDominatedBySearch(anyList(), anyDouble())).thenReturn(false); service.doRun(); @@ -549,15 +576,16 @@ public void testNonCancellationWhenSearchTrafficIsNotQualifyingForCancellation() when(settings.getSearchTaskSettings()).thenReturn(searchTaskSettings); // Create a mix of low and high resource usage tasks (60 low + 15 high resource usage tasks). - Map activeSearchTasks = new HashMap<>(); + Map activeSearchTasks = new HashMap<>(); for (long i = 0; i < 75; i++) { - Class taskType = randomBoolean() ? SearchTask.class : SearchShardTask.class; + Class taskType = randomBoolean() ? SearchTask.class : SearchShardTask.class; if (i % 5 == 0) { activeSearchTasks.put(i, createMockTaskWithResourceStats(taskType, 500, 800, i)); } else { activeSearchTasks.put(i, createMockTaskWithResourceStats(taskType, 100, 800, i)); } } + activeSearchTasks.values().forEach(task -> task.setQueryGroupId(threadPool.getThreadContext())); doReturn(activeSearchTasks).when(mockTaskResourceTrackingService).getResourceAwareTasks(); // this will trigger cancellation but the cancellation should not happen as the node is not is duress because of search traffic diff --git a/server/src/test/java/org/opensearch/wlm/QueryGroupServiceTests.java b/server/src/test/java/org/opensearch/wlm/QueryGroupServiceTests.java new file mode 100644 index 0000000000000..c5cf0dac4f807 --- /dev/null +++ b/server/src/test/java/org/opensearch/wlm/QueryGroupServiceTests.java @@ -0,0 +1,489 @@ +/* + * 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.wlm; + +import org.opensearch.action.search.SearchTask; +import org.opensearch.cluster.ClusterChangedEvent; +import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.metadata.Metadata; +import org.opensearch.cluster.metadata.QueryGroup; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.unit.TimeValue; +import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException; +import org.opensearch.search.backpressure.trackers.NodeDuressTrackers; +import org.opensearch.tasks.Task; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.threadpool.Scheduler; +import org.opensearch.threadpool.TestThreadPool; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.wlm.cancellation.QueryGroupTaskCancellationService; +import org.opensearch.wlm.cancellation.TaskSelectionStrategy; +import org.opensearch.wlm.stats.QueryGroupState; +import org.opensearch.wlm.tracker.QueryGroupResourceUsageTrackerService; + +import java.util.Collection; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; +import java.util.function.BooleanSupplier; + +import org.mockito.ArgumentCaptor; +import org.mockito.Mockito; + +import static org.opensearch.wlm.tracker.ResourceUsageCalculatorTests.createMockTaskWithResourceStats; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class QueryGroupServiceTests extends OpenSearchTestCase { + private QueryGroupService queryGroupService; + private QueryGroupTaskCancellationService mockCancellationService; + private ClusterService mockClusterService; + private ThreadPool mockThreadPool; + private WorkloadManagementSettings mockWorkloadManagementSettings; + private Scheduler.Cancellable mockScheduledFuture; + private Map mockQueryGroupStateMap; + NodeDuressTrackers mockNodeDuressTrackers; + QueryGroupsStateAccessor mockQueryGroupsStateAccessor; + + public void setUp() throws Exception { + super.setUp(); + mockClusterService = Mockito.mock(ClusterService.class); + mockThreadPool = Mockito.mock(ThreadPool.class); + mockScheduledFuture = Mockito.mock(Scheduler.Cancellable.class); + mockWorkloadManagementSettings = Mockito.mock(WorkloadManagementSettings.class); + mockQueryGroupStateMap = new HashMap<>(); + mockNodeDuressTrackers = Mockito.mock(NodeDuressTrackers.class); + mockCancellationService = Mockito.mock(TestQueryGroupCancellationService.class); + mockQueryGroupsStateAccessor = new QueryGroupsStateAccessor(); + + queryGroupService = new QueryGroupService( + mockCancellationService, + mockClusterService, + mockThreadPool, + mockWorkloadManagementSettings, + mockNodeDuressTrackers, + mockQueryGroupsStateAccessor, + new HashSet<>(), + new HashSet<>() + ); + } + + public void tearDown() throws Exception { + super.tearDown(); + mockThreadPool.shutdown(); + } + + public void testClusterChanged() { + ClusterChangedEvent mockClusterChangedEvent = Mockito.mock(ClusterChangedEvent.class); + ClusterState mockPreviousClusterState = Mockito.mock(ClusterState.class); + ClusterState mockClusterState = Mockito.mock(ClusterState.class); + Metadata mockPreviousMetadata = Mockito.mock(Metadata.class); + Metadata mockMetadata = Mockito.mock(Metadata.class); + QueryGroup addedQueryGroup = new QueryGroup( + "addedQueryGroup", + "4242", + new MutableQueryGroupFragment(MutableQueryGroupFragment.ResiliencyMode.ENFORCED, Map.of(ResourceType.MEMORY, 0.5)), + 1L + ); + QueryGroup deletedQueryGroup = new QueryGroup( + "deletedQueryGroup", + "4241", + new MutableQueryGroupFragment(MutableQueryGroupFragment.ResiliencyMode.ENFORCED, Map.of(ResourceType.MEMORY, 0.5)), + 1L + ); + Map previousQueryGroups = new HashMap<>(); + previousQueryGroups.put("4242", addedQueryGroup); + Map currentQueryGroups = new HashMap<>(); + currentQueryGroups.put("4241", deletedQueryGroup); + + when(mockClusterChangedEvent.previousState()).thenReturn(mockPreviousClusterState); + when(mockClusterChangedEvent.state()).thenReturn(mockClusterState); + when(mockPreviousClusterState.metadata()).thenReturn(mockPreviousMetadata); + when(mockClusterState.metadata()).thenReturn(mockMetadata); + when(mockPreviousMetadata.queryGroups()).thenReturn(previousQueryGroups); + when(mockMetadata.queryGroups()).thenReturn(currentQueryGroups); + queryGroupService.clusterChanged(mockClusterChangedEvent); + + Set currentQueryGroupsExpected = Set.of(currentQueryGroups.get("4241")); + Set previousQueryGroupsExpected = Set.of(previousQueryGroups.get("4242")); + + assertEquals(currentQueryGroupsExpected, queryGroupService.getActiveQueryGroups()); + assertEquals(previousQueryGroupsExpected, queryGroupService.getDeletedQueryGroups()); + } + + public void testDoStart_SchedulesTask() { + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + when(mockWorkloadManagementSettings.getQueryGroupServiceRunInterval()).thenReturn(TimeValue.timeValueSeconds(1)); + queryGroupService.doStart(); + Mockito.verify(mockThreadPool).scheduleWithFixedDelay(any(Runnable.class), any(TimeValue.class), eq(ThreadPool.Names.GENERIC)); + } + + public void testDoStop_CancelsScheduledTask() { + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + when(mockThreadPool.scheduleWithFixedDelay(any(), any(), any())).thenReturn(mockScheduledFuture); + queryGroupService.doStart(); + queryGroupService.doStop(); + Mockito.verify(mockScheduledFuture).cancel(); + } + + public void testDoRun_WhenModeEnabled() { + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + when(mockNodeDuressTrackers.isNodeInDuress()).thenReturn(true); + // Call the method + queryGroupService.doRun(); + + // Verify that refreshQueryGroups was called + + // Verify that cancelTasks was called with a BooleanSupplier + ArgumentCaptor booleanSupplierCaptor = ArgumentCaptor.forClass(BooleanSupplier.class); + Mockito.verify(mockCancellationService).cancelTasks(booleanSupplierCaptor.capture(), any(), any()); + + // Assert the behavior of the BooleanSupplier + BooleanSupplier capturedSupplier = booleanSupplierCaptor.getValue(); + assertTrue(capturedSupplier.getAsBoolean()); + + } + + public void testDoRun_WhenModeDisabled() { + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.DISABLED); + when(mockNodeDuressTrackers.isNodeInDuress()).thenReturn(false); + queryGroupService.doRun(); + // Verify that refreshQueryGroups was called + + Mockito.verify(mockCancellationService, never()).cancelTasks(any(), any(), any()); + + } + + public void testRejectIfNeeded_whenQueryGroupIdIsNullOrDefaultOne() { + QueryGroup testQueryGroup = new QueryGroup( + "testQueryGroup", + "queryGroupId1", + new MutableQueryGroupFragment(MutableQueryGroupFragment.ResiliencyMode.ENFORCED, Map.of(ResourceType.CPU, 0.10)), + 1L + ); + Set activeQueryGroups = new HashSet<>() { + { + add(testQueryGroup); + } + }; + mockQueryGroupStateMap = new HashMap<>(); + mockQueryGroupsStateAccessor = new QueryGroupsStateAccessor(mockQueryGroupStateMap); + mockQueryGroupStateMap.put("queryGroupId1", new QueryGroupState()); + + Map spyMap = spy(mockQueryGroupStateMap); + + queryGroupService = new QueryGroupService( + mockCancellationService, + mockClusterService, + mockThreadPool, + mockWorkloadManagementSettings, + mockNodeDuressTrackers, + mockQueryGroupsStateAccessor, + activeQueryGroups, + new HashSet<>() + ); + queryGroupService.rejectIfNeeded(null); + + verify(spyMap, never()).get(any()); + + queryGroupService.rejectIfNeeded(QueryGroupTask.DEFAULT_QUERY_GROUP_ID_SUPPLIER.get()); + verify(spyMap, never()).get(any()); + } + + public void testRejectIfNeeded_whenQueryGroupIsSoftMode() { + QueryGroup testQueryGroup = new QueryGroup( + "testQueryGroup", + "queryGroupId1", + new MutableQueryGroupFragment(MutableQueryGroupFragment.ResiliencyMode.SOFT, Map.of(ResourceType.CPU, 0.10)), + 1L + ); + Set activeQueryGroups = new HashSet<>() { + { + add(testQueryGroup); + } + }; + mockQueryGroupStateMap = new HashMap<>(); + QueryGroupState spyState = spy(new QueryGroupState()); + mockQueryGroupStateMap.put("queryGroupId1", spyState); + + mockQueryGroupsStateAccessor = new QueryGroupsStateAccessor(mockQueryGroupStateMap); + + Map spyMap = spy(mockQueryGroupStateMap); + + queryGroupService = new QueryGroupService( + mockCancellationService, + mockClusterService, + mockThreadPool, + mockWorkloadManagementSettings, + mockNodeDuressTrackers, + mockQueryGroupsStateAccessor, + activeQueryGroups, + new HashSet<>() + ); + queryGroupService.rejectIfNeeded("queryGroupId1"); + + verify(spyState, never()).getResourceState(); + } + + public void testRejectIfNeeded_whenQueryGroupIsEnforcedMode_andNotBreaching() { + QueryGroup testQueryGroup = new QueryGroup( + "testQueryGroup", + "queryGroupId1", + new MutableQueryGroupFragment(MutableQueryGroupFragment.ResiliencyMode.ENFORCED, Map.of(ResourceType.CPU, 0.10)), + 1L + ); + QueryGroup spuQueryGroup = spy(testQueryGroup); + Set activeQueryGroups = new HashSet<>() { + { + add(spuQueryGroup); + } + }; + mockQueryGroupStateMap = new HashMap<>(); + QueryGroupState queryGroupState = new QueryGroupState(); + queryGroupState.getResourceState().get(ResourceType.CPU).setLastRecordedUsage(0.05); + + mockQueryGroupStateMap.put("queryGroupId1", queryGroupState); + + mockQueryGroupsStateAccessor = new QueryGroupsStateAccessor(mockQueryGroupStateMap); + + queryGroupService = new QueryGroupService( + mockCancellationService, + mockClusterService, + mockThreadPool, + mockWorkloadManagementSettings, + mockNodeDuressTrackers, + mockQueryGroupsStateAccessor, + activeQueryGroups, + new HashSet<>() + ); + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + when(mockWorkloadManagementSettings.getNodeLevelCpuRejectionThreshold()).thenReturn(0.8); + queryGroupService.rejectIfNeeded("queryGroupId1"); + + // verify the check to compare the current usage and limit + // this should happen 3 times => 2 to check whether the resource limit has the TRACKED resource type and 1 to get the value + verify(spuQueryGroup, times(3)).getResourceLimits(); + assertEquals(0, queryGroupState.getResourceState().get(ResourceType.CPU).rejections.count()); + assertEquals(0, queryGroupState.totalRejections.count()); + } + + public void testRejectIfNeeded_whenQueryGroupIsEnforcedMode_andBreaching() { + QueryGroup testQueryGroup = new QueryGroup( + "testQueryGroup", + "queryGroupId1", + new MutableQueryGroupFragment( + MutableQueryGroupFragment.ResiliencyMode.ENFORCED, + Map.of(ResourceType.CPU, 0.10, ResourceType.MEMORY, 0.10) + ), + 1L + ); + QueryGroup spuQueryGroup = spy(testQueryGroup); + Set activeQueryGroups = new HashSet<>() { + { + add(spuQueryGroup); + } + }; + mockQueryGroupStateMap = new HashMap<>(); + QueryGroupState queryGroupState = new QueryGroupState(); + queryGroupState.getResourceState().get(ResourceType.CPU).setLastRecordedUsage(0.18); + queryGroupState.getResourceState().get(ResourceType.MEMORY).setLastRecordedUsage(0.18); + QueryGroupState spyState = spy(queryGroupState); + + mockQueryGroupsStateAccessor = new QueryGroupsStateAccessor(mockQueryGroupStateMap); + + mockQueryGroupStateMap.put("queryGroupId1", spyState); + + queryGroupService = new QueryGroupService( + mockCancellationService, + mockClusterService, + mockThreadPool, + mockWorkloadManagementSettings, + mockNodeDuressTrackers, + mockQueryGroupsStateAccessor, + activeQueryGroups, + new HashSet<>() + ); + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + assertThrows(OpenSearchRejectedExecutionException.class, () -> queryGroupService.rejectIfNeeded("queryGroupId1")); + + // verify the check to compare the current usage and limit + // this should happen 3 times => 1 to check whether the resource limit has the TRACKED resource type and 1 to get the value + // because it will break out of the loop since the limits are breached + verify(spuQueryGroup, times(2)).getResourceLimits(); + assertEquals( + 1, + queryGroupState.getResourceState().get(ResourceType.CPU).rejections.count() + queryGroupState.getResourceState() + .get(ResourceType.MEMORY).rejections.count() + ); + assertEquals(1, queryGroupState.totalRejections.count()); + } + + public void testRejectIfNeeded_whenFeatureIsNotEnabled() { + QueryGroup testQueryGroup = new QueryGroup( + "testQueryGroup", + "queryGroupId1", + new MutableQueryGroupFragment(MutableQueryGroupFragment.ResiliencyMode.ENFORCED, Map.of(ResourceType.CPU, 0.10)), + 1L + ); + Set activeQueryGroups = new HashSet<>() { + { + add(testQueryGroup); + } + }; + mockQueryGroupStateMap = new HashMap<>(); + mockQueryGroupStateMap.put("queryGroupId1", new QueryGroupState()); + + Map spyMap = spy(mockQueryGroupStateMap); + + mockQueryGroupsStateAccessor = new QueryGroupsStateAccessor(mockQueryGroupStateMap); + + queryGroupService = new QueryGroupService( + mockCancellationService, + mockClusterService, + mockThreadPool, + mockWorkloadManagementSettings, + mockNodeDuressTrackers, + mockQueryGroupsStateAccessor, + activeQueryGroups, + new HashSet<>() + ); + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.DISABLED); + + queryGroupService.rejectIfNeeded(testQueryGroup.get_id()); + verify(spyMap, never()).get(any()); + } + + public void testOnTaskCompleted() { + Task task = createMockTaskWithResourceStats(SearchTask.class, 100, 200, 0, 12); + mockThreadPool = new TestThreadPool("queryGroupServiceTests"); + mockThreadPool.getThreadContext().putHeader(QueryGroupTask.QUERY_GROUP_ID_HEADER, "testId"); + QueryGroupState queryGroupState = new QueryGroupState(); + mockQueryGroupStateMap.put("testId", queryGroupState); + mockQueryGroupsStateAccessor = new QueryGroupsStateAccessor(mockQueryGroupStateMap); + queryGroupService = new QueryGroupService( + mockCancellationService, + mockClusterService, + mockThreadPool, + mockWorkloadManagementSettings, + mockNodeDuressTrackers, + mockQueryGroupsStateAccessor, + new HashSet<>() { + { + add( + new QueryGroup( + "testQueryGroup", + "testId", + new MutableQueryGroupFragment( + MutableQueryGroupFragment.ResiliencyMode.ENFORCED, + Map.of(ResourceType.CPU, 0.10, ResourceType.MEMORY, 0.10) + ), + 1L + ) + ); + } + }, + new HashSet<>() + ); + + ((QueryGroupTask) task).setQueryGroupId(mockThreadPool.getThreadContext()); + queryGroupService.onTaskCompleted(task); + + assertEquals(1, queryGroupState.completions.count()); + + // test non QueryGroupTask + task = new Task(1, "simple", "test", "mock task", null, null); + queryGroupService.onTaskCompleted(task); + + // It should still be 1 + assertEquals(1, queryGroupState.completions.count()); + + mockThreadPool.shutdown(); + } + + public void testShouldSBPHandle() { + QueryGroupTask task = createMockTaskWithResourceStats(SearchTask.class, 100, 200, 0, 12); + QueryGroupState queryGroupState = new QueryGroupState(); + Set activeQueryGroups = new HashSet<>(); + mockQueryGroupStateMap.put("testId", queryGroupState); + mockQueryGroupsStateAccessor = new QueryGroupsStateAccessor(mockQueryGroupStateMap); + queryGroupService = new QueryGroupService( + mockCancellationService, + mockClusterService, + mockThreadPool, + mockWorkloadManagementSettings, + mockNodeDuressTrackers, + mockQueryGroupsStateAccessor, + activeQueryGroups, + Collections.emptySet() + ); + + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + + // Default queryGroupId + mockThreadPool = new TestThreadPool("queryGroupServiceTests"); + mockThreadPool.getThreadContext() + .putHeader(QueryGroupTask.QUERY_GROUP_ID_HEADER, QueryGroupTask.DEFAULT_QUERY_GROUP_ID_SUPPLIER.get()); + task.setQueryGroupId(mockThreadPool.getThreadContext()); + assertTrue(queryGroupService.shouldSBPHandle(task)); + + mockThreadPool.shutdownNow(); + + // invalid queryGroup task + mockThreadPool = new TestThreadPool("queryGroupServiceTests"); + mockThreadPool.getThreadContext().putHeader(QueryGroupTask.QUERY_GROUP_ID_HEADER, "testId"); + task.setQueryGroupId(mockThreadPool.getThreadContext()); + assertTrue(queryGroupService.shouldSBPHandle(task)); + + // Valid query group task but wlm not enabled + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.DISABLED); + activeQueryGroups.add( + new QueryGroup( + "testQueryGroup", + "testId", + new MutableQueryGroupFragment( + MutableQueryGroupFragment.ResiliencyMode.ENFORCED, + Map.of(ResourceType.CPU, 0.10, ResourceType.MEMORY, 0.10) + ), + 1L + ) + ); + assertTrue(queryGroupService.shouldSBPHandle(task)); + + } + + // This is needed to test the behavior of QueryGroupService#doRun method + static class TestQueryGroupCancellationService extends QueryGroupTaskCancellationService { + public TestQueryGroupCancellationService( + WorkloadManagementSettings workloadManagementSettings, + TaskSelectionStrategy taskSelectionStrategy, + QueryGroupResourceUsageTrackerService resourceUsageTrackerService, + QueryGroupsStateAccessor queryGroupsStateAccessor, + Collection activeQueryGroups, + Collection deletedQueryGroups + ) { + super(workloadManagementSettings, taskSelectionStrategy, resourceUsageTrackerService, queryGroupsStateAccessor); + } + + @Override + public void cancelTasks( + BooleanSupplier isNodeInDuress, + Collection activeQueryGroups, + Collection deletedQueryGroups + ) { + + } + } +} diff --git a/server/src/test/java/org/opensearch/wlm/WorkloadManagementTransportInterceptorTests.java b/server/src/test/java/org/opensearch/wlm/WorkloadManagementTransportInterceptorTests.java index 4668b845150a9..d4cd7b79455a3 100644 --- a/server/src/test/java/org/opensearch/wlm/WorkloadManagementTransportInterceptorTests.java +++ b/server/src/test/java/org/opensearch/wlm/WorkloadManagementTransportInterceptorTests.java @@ -8,24 +8,56 @@ package org.opensearch.wlm; +import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.metadata.Metadata; +import org.opensearch.cluster.service.ClusterService; import org.opensearch.test.OpenSearchTestCase; import org.opensearch.threadpool.TestThreadPool; import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.TransportRequest; import org.opensearch.transport.TransportRequestHandler; import org.opensearch.wlm.WorkloadManagementTransportInterceptor.RequestHandler; +import org.opensearch.wlm.cancellation.QueryGroupTaskCancellationService; + +import java.util.Collections; import static org.opensearch.threadpool.ThreadPool.Names.SAME; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; public class WorkloadManagementTransportInterceptorTests extends OpenSearchTestCase { - + private QueryGroupTaskCancellationService mockTaskCancellationService; + private ClusterService mockClusterService; + private ThreadPool mockThreadPool; + private WorkloadManagementSettings mockWorkloadManagementSettings; private ThreadPool threadPool; private WorkloadManagementTransportInterceptor sut; + private QueryGroupsStateAccessor stateAccessor; public void setUp() throws Exception { super.setUp(); + mockTaskCancellationService = mock(QueryGroupTaskCancellationService.class); + mockClusterService = mock(ClusterService.class); + mockThreadPool = mock(ThreadPool.class); + mockWorkloadManagementSettings = mock(WorkloadManagementSettings.class); threadPool = new TestThreadPool(getTestName()); - sut = new WorkloadManagementTransportInterceptor(threadPool, new QueryGroupService()); + stateAccessor = new QueryGroupsStateAccessor(); + + ClusterState state = mock(ClusterState.class); + Metadata metadata = mock(Metadata.class); + when(mockClusterService.state()).thenReturn(state); + when(state.metadata()).thenReturn(metadata); + when(metadata.queryGroups()).thenReturn(Collections.emptyMap()); + sut = new WorkloadManagementTransportInterceptor( + threadPool, + new QueryGroupService( + mockTaskCancellationService, + mockClusterService, + mockThreadPool, + mockWorkloadManagementSettings, + stateAccessor + ) + ); } public void tearDown() throws Exception { diff --git a/server/src/test/java/org/opensearch/wlm/cancellation/QueryGroupTaskCancellationServiceTests.java b/server/src/test/java/org/opensearch/wlm/cancellation/QueryGroupTaskCancellationServiceTests.java index f7a49235efc69..13e8e2c527073 100644 --- a/server/src/test/java/org/opensearch/wlm/cancellation/QueryGroupTaskCancellationServiceTests.java +++ b/server/src/test/java/org/opensearch/wlm/cancellation/QueryGroupTaskCancellationServiceTests.java @@ -17,12 +17,16 @@ import org.opensearch.wlm.MutableQueryGroupFragment.ResiliencyMode; import org.opensearch.wlm.QueryGroupLevelResourceUsageView; import org.opensearch.wlm.QueryGroupTask; +import org.opensearch.wlm.QueryGroupsStateAccessor; import org.opensearch.wlm.ResourceType; +import org.opensearch.wlm.WlmMode; import org.opensearch.wlm.WorkloadManagementSettings; +import org.opensearch.wlm.stats.QueryGroupState; import org.opensearch.wlm.tracker.QueryGroupResourceUsageTrackerService; import org.opensearch.wlm.tracker.ResourceUsageCalculatorTrackerServiceTests.TestClock; import org.junit.Before; +import java.util.ArrayList; import java.util.Collection; import java.util.Collections; import java.util.HashMap; @@ -31,7 +35,9 @@ import java.util.Map; import java.util.Set; import java.util.stream.Collectors; +import java.util.stream.IntStream; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -47,6 +53,7 @@ public class QueryGroupTaskCancellationServiceTests extends OpenSearchTestCase { private QueryGroupTaskCancellationService taskCancellation; private WorkloadManagementSettings workloadManagementSettings; private QueryGroupResourceUsageTrackerService resourceUsageTrackerService; + private QueryGroupsStateAccessor stateAccessor; @Before public void setup() { @@ -59,12 +66,13 @@ public void setup() { when(workloadManagementSettings.getNodeLevelCpuCancellationThreshold()).thenReturn(0.9); when(workloadManagementSettings.getNodeLevelMemoryCancellationThreshold()).thenReturn(0.9); resourceUsageTrackerService = mock(QueryGroupResourceUsageTrackerService.class); + stateAccessor = mock(QueryGroupsStateAccessor.class); + when(stateAccessor.getQueryGroupState(any())).thenReturn(new QueryGroupState()); taskCancellation = new QueryGroupTaskCancellationService( workloadManagementSettings, new MaximumResourceTaskSelectionStrategy(), resourceUsageTrackerService, - activeQueryGroups, - deletedQueryGroups + stateAccessor ); } @@ -138,7 +146,7 @@ public void testGetCancellableTasksFrom_returnsTasksWhenBreachingThresholdForMem activeQueryGroups.add(queryGroup1); taskCancellation.queryGroupLevelResourceUsageViews = queryGroupLevelViews; - List cancellableTasksFrom = taskCancellation.getAllCancellableTasks(ResiliencyMode.ENFORCED); + List cancellableTasksFrom = taskCancellation.getAllCancellableTasks(ResiliencyMode.ENFORCED, activeQueryGroups); assertEquals(2, cancellableTasksFrom.size()); assertEquals(1234, cancellableTasksFrom.get(0).getTask().getId()); assertEquals(4321, cancellableTasksFrom.get(1).getTask().getId()); @@ -187,11 +195,10 @@ public void testGetCancellableTasksFrom_filtersQueryGroupCorrectly() { workloadManagementSettings, new MaximumResourceTaskSelectionStrategy(), resourceUsageTrackerService, - activeQueryGroups, - deletedQueryGroups + stateAccessor ); - List cancellableTasksFrom = taskCancellation.getAllCancellableTasks(ResiliencyMode.SOFT); + List cancellableTasksFrom = taskCancellation.getAllCancellableTasks(ResiliencyMode.SOFT, activeQueryGroups); assertEquals(0, cancellableTasksFrom.size()); } @@ -219,19 +226,19 @@ public void testCancelTasks_cancelsGivenTasks() { workloadManagementSettings, new MaximumResourceTaskSelectionStrategy(), resourceUsageTrackerService, - activeQueryGroups, - deletedQueryGroups + stateAccessor ); taskCancellation.queryGroupLevelResourceUsageViews = queryGroupLevelViews; - List cancellableTasksFrom = taskCancellation.getAllCancellableTasks(ResiliencyMode.ENFORCED); + List cancellableTasksFrom = taskCancellation.getAllCancellableTasks(ResiliencyMode.ENFORCED, activeQueryGroups); assertEquals(2, cancellableTasksFrom.size()); assertEquals(1234, cancellableTasksFrom.get(0).getTask().getId()); assertEquals(4321, cancellableTasksFrom.get(1).getTask().getId()); when(resourceUsageTrackerService.constructQueryGroupLevelUsageViews()).thenReturn(queryGroupLevelViews); - taskCancellation.cancelTasks(() -> false); + when(workloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + taskCancellation.cancelTasks(() -> false, activeQueryGroups, deletedQueryGroups); assertTrue(cancellableTasksFrom.get(0).getTask().isCancelled()); assertTrue(cancellableTasksFrom.get(1).getTask().isCancelled()); } @@ -281,13 +288,11 @@ public void testCancelTasks_cancelsTasksFromDeletedQueryGroups() { workloadManagementSettings, new MaximumResourceTaskSelectionStrategy(), resourceUsageTrackerService, - activeQueryGroups, - deletedQueryGroups + stateAccessor ); - taskCancellation.queryGroupLevelResourceUsageViews = queryGroupLevelViews; - List cancellableTasksFrom = taskCancellation.getAllCancellableTasks(ResiliencyMode.ENFORCED); + List cancellableTasksFrom = taskCancellation.getAllCancellableTasks(ResiliencyMode.ENFORCED, activeQueryGroups); assertEquals(2, cancellableTasksFrom.size()); assertEquals(1234, cancellableTasksFrom.get(0).getTask().getId()); assertEquals(4321, cancellableTasksFrom.get(1).getTask().getId()); @@ -298,7 +303,8 @@ public void testCancelTasks_cancelsTasksFromDeletedQueryGroups() { assertEquals(1001, cancellableTasksFromDeletedQueryGroups.get(1).getTask().getId()); when(resourceUsageTrackerService.constructQueryGroupLevelUsageViews()).thenReturn(queryGroupLevelViews); - taskCancellation.cancelTasks(() -> true); + when(workloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + taskCancellation.cancelTasks(() -> true, activeQueryGroups, deletedQueryGroups); assertTrue(cancellableTasksFrom.get(0).getTask().isCancelled()); assertTrue(cancellableTasksFrom.get(1).getTask().isCancelled()); @@ -352,12 +358,11 @@ public void testCancelTasks_does_not_cancelTasksFromDeletedQueryGroups_whenNodeN workloadManagementSettings, new MaximumResourceTaskSelectionStrategy(), resourceUsageTrackerService, - activeQueryGroups, - deletedQueryGroups + stateAccessor ); taskCancellation.queryGroupLevelResourceUsageViews = queryGroupLevelViews; - List cancellableTasksFrom = taskCancellation.getAllCancellableTasks(ResiliencyMode.ENFORCED); + List cancellableTasksFrom = taskCancellation.getAllCancellableTasks(ResiliencyMode.ENFORCED, activeQueryGroups); assertEquals(2, cancellableTasksFrom.size()); assertEquals(1234, cancellableTasksFrom.get(0).getTask().getId()); assertEquals(4321, cancellableTasksFrom.get(1).getTask().getId()); @@ -368,7 +373,8 @@ public void testCancelTasks_does_not_cancelTasksFromDeletedQueryGroups_whenNodeN assertEquals(1001, cancellableTasksFromDeletedQueryGroups.get(1).getTask().getId()); when(resourceUsageTrackerService.constructQueryGroupLevelUsageViews()).thenReturn(queryGroupLevelViews); - taskCancellation.cancelTasks(() -> false); + when(workloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + taskCancellation.cancelTasks(() -> false, activeQueryGroups, deletedQueryGroups); assertTrue(cancellableTasksFrom.get(0).getTask().isCancelled()); assertTrue(cancellableTasksFrom.get(1).getTask().isCancelled()); @@ -411,24 +417,24 @@ public void testCancelTasks_cancelsGivenTasks_WhenNodeInDuress() { workloadManagementSettings, new MaximumResourceTaskSelectionStrategy(), resourceUsageTrackerService, - activeQueryGroups, - deletedQueryGroups + stateAccessor ); taskCancellation.queryGroupLevelResourceUsageViews = queryGroupLevelViews; - List cancellableTasksFrom = taskCancellation.getAllCancellableTasks(ResiliencyMode.ENFORCED); + List cancellableTasksFrom = taskCancellation.getAllCancellableTasks(ResiliencyMode.ENFORCED, activeQueryGroups); assertEquals(2, cancellableTasksFrom.size()); assertEquals(1234, cancellableTasksFrom.get(0).getTask().getId()); assertEquals(4321, cancellableTasksFrom.get(1).getTask().getId()); - List cancellableTasksFrom1 = taskCancellation.getAllCancellableTasks(ResiliencyMode.SOFT); + List cancellableTasksFrom1 = taskCancellation.getAllCancellableTasks(ResiliencyMode.SOFT, activeQueryGroups); assertEquals(2, cancellableTasksFrom1.size()); assertEquals(5678, cancellableTasksFrom1.get(0).getTask().getId()); assertEquals(8765, cancellableTasksFrom1.get(1).getTask().getId()); when(resourceUsageTrackerService.constructQueryGroupLevelUsageViews()).thenReturn(queryGroupLevelViews); - taskCancellation.cancelTasks(() -> true); + when(workloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + taskCancellation.cancelTasks(() -> true, activeQueryGroups, deletedQueryGroups); assertTrue(cancellableTasksFrom.get(0).getTask().isCancelled()); assertTrue(cancellableTasksFrom.get(1).getTask().isCancelled()); assertTrue(cancellableTasksFrom1.get(0).getTask().isCancelled()); @@ -456,7 +462,7 @@ public void testGetAllCancellableTasks_ReturnsNoTasksWhenNotBreachingThresholds( activeQueryGroups.add(queryGroup1); taskCancellation.queryGroupLevelResourceUsageViews = queryGroupLevelViews; - List allCancellableTasks = taskCancellation.getAllCancellableTasks(ResiliencyMode.ENFORCED); + List allCancellableTasks = taskCancellation.getAllCancellableTasks(ResiliencyMode.ENFORCED, activeQueryGroups); assertTrue(allCancellableTasks.isEmpty()); } @@ -479,7 +485,7 @@ public void testGetAllCancellableTasks_ReturnsTasksWhenBreachingThresholds() { activeQueryGroups.add(queryGroup1); taskCancellation.queryGroupLevelResourceUsageViews = queryGroupLevelViews; - List allCancellableTasks = taskCancellation.getAllCancellableTasks(ResiliencyMode.ENFORCED); + List allCancellableTasks = taskCancellation.getAllCancellableTasks(ResiliencyMode.ENFORCED, activeQueryGroups); assertEquals(2, allCancellableTasks.size()); assertEquals(1234, allCancellableTasks.get(0).getTask().getId()); assertEquals(4321, allCancellableTasks.get(1).getTask().getId()); @@ -513,6 +519,48 @@ public void testGetCancellableTasksFrom_doesNotReturnTasksWhenQueryGroupIdNotFou assertEquals(0, cancellableTasksFrom.size()); } + public void testPruneDeletedQueryGroups() { + QueryGroup queryGroup1 = new QueryGroup( + "testQueryGroup1", + queryGroupId1, + new MutableQueryGroupFragment(ResiliencyMode.ENFORCED, Map.of(ResourceType.CPU, 0.2)), + 1L + ); + QueryGroup queryGroup2 = new QueryGroup( + "testQueryGroup2", + queryGroupId2, + new MutableQueryGroupFragment(ResiliencyMode.ENFORCED, Map.of(ResourceType.CPU, 0.1)), + 1L + ); + List deletedQueryGroups = new ArrayList<>(); + deletedQueryGroups.add(queryGroup1); + deletedQueryGroups.add(queryGroup2); + QueryGroupLevelResourceUsageView resourceUsageView1 = createResourceUsageViewMock(); + + List activeTasks = IntStream.range(0, 5).mapToObj(this::getRandomSearchTask).collect(Collectors.toList()); + when(resourceUsageView1.getActiveTasks()).thenReturn(activeTasks); + + QueryGroupLevelResourceUsageView resourceUsageView2 = createResourceUsageViewMock(); + when(resourceUsageView2.getActiveTasks()).thenReturn(new ArrayList<>()); + + queryGroupLevelViews.put(queryGroupId1, resourceUsageView1); + queryGroupLevelViews.put(queryGroupId2, resourceUsageView2); + + QueryGroupTaskCancellationService taskCancellation = new QueryGroupTaskCancellationService( + workloadManagementSettings, + new MaximumResourceTaskSelectionStrategy(), + resourceUsageTrackerService, + stateAccessor + ); + taskCancellation.queryGroupLevelResourceUsageViews = queryGroupLevelViews; + + taskCancellation.pruneDeletedQueryGroups(deletedQueryGroups); + + assertEquals(1, deletedQueryGroups.size()); + assertEquals(queryGroupId1, deletedQueryGroups.get(0).get_id()); + + } + private QueryGroupLevelResourceUsageView createResourceUsageViewMock() { QueryGroupLevelResourceUsageView mockView = mock(QueryGroupLevelResourceUsageView.class); when(mockView.getActiveTasks()).thenReturn(List.of(getRandomSearchTask(1234), getRandomSearchTask(4321))); diff --git a/server/src/test/java/org/opensearch/wlm/listeners/QueryGroupRequestOperationListenerTests.java b/server/src/test/java/org/opensearch/wlm/listeners/QueryGroupRequestOperationListenerTests.java index 0307ff623c408..1127b50399d24 100644 --- a/server/src/test/java/org/opensearch/wlm/listeners/QueryGroupRequestOperationListenerTests.java +++ b/server/src/test/java/org/opensearch/wlm/listeners/QueryGroupRequestOperationListenerTests.java @@ -8,6 +8,9 @@ package org.opensearch.wlm.listeners; +import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.metadata.Metadata; +import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException; import org.opensearch.test.OpenSearchTestCase; @@ -15,12 +18,16 @@ import org.opensearch.threadpool.ThreadPool; import org.opensearch.wlm.QueryGroupService; import org.opensearch.wlm.QueryGroupTask; +import org.opensearch.wlm.QueryGroupsStateAccessor; import org.opensearch.wlm.ResourceType; +import org.opensearch.wlm.WorkloadManagementSettings; +import org.opensearch.wlm.cancellation.QueryGroupTaskCancellationService; import org.opensearch.wlm.stats.QueryGroupState; import org.opensearch.wlm.stats.QueryGroupStats; import java.io.IOException; import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -28,18 +35,24 @@ import static org.mockito.Mockito.doNothing; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; public class QueryGroupRequestOperationListenerTests extends OpenSearchTestCase { public static final int ITERATIONS = 20; ThreadPool testThreadPool; QueryGroupService queryGroupService; - + private QueryGroupTaskCancellationService taskCancellationService; + private ClusterService mockClusterService; + private WorkloadManagementSettings mockWorkloadManagementSettings; Map queryGroupStateMap; String testQueryGroupId; QueryGroupRequestOperationListener sut; public void setUp() throws Exception { super.setUp(); + taskCancellationService = mock(QueryGroupTaskCancellationService.class); + mockClusterService = mock(ClusterService.class); + mockWorkloadManagementSettings = mock(WorkloadManagementSettings.class); queryGroupStateMap = new HashMap<>(); testQueryGroupId = "safjgagnakg-3r3fads"; testThreadPool = new TestThreadPool("RejectionTestThreadPool"); @@ -77,6 +90,21 @@ public void testValidQueryGroupRequestFailure() throws IOException { 0, 1, 0, + 0, + Map.of( + ResourceType.CPU, + new QueryGroupStats.ResourceStats(0, 0, 0), + ResourceType.MEMORY, + new QueryGroupStats.ResourceStats(0, 0, 0) + ) + ), + QueryGroupTask.DEFAULT_QUERY_GROUP_ID_SUPPLIER.get(), + new QueryGroupStats.QueryGroupStatsHolder( + 0, + 0, + 0, + 0, + 0, Map.of( ResourceType.CPU, new QueryGroupStats.ResourceStats(0, 0, 0), @@ -93,8 +121,18 @@ public void testValidQueryGroupRequestFailure() throws IOException { public void testMultiThreadedValidQueryGroupRequestFailures() { queryGroupStateMap.put(testQueryGroupId, new QueryGroupState()); - - queryGroupService = new QueryGroupService(queryGroupStateMap); + QueryGroupsStateAccessor accessor = new QueryGroupsStateAccessor(queryGroupStateMap); + setupMockedQueryGroupsFromClusterState(); + queryGroupService = new QueryGroupService( + taskCancellationService, + mockClusterService, + testThreadPool, + mockWorkloadManagementSettings, + null, + accessor, + Collections.emptySet(), + Collections.emptySet() + ); sut = new QueryGroupRequestOperationListener(queryGroupService, testThreadPool); @@ -127,6 +165,21 @@ public void testMultiThreadedValidQueryGroupRequestFailures() { 0, ITERATIONS, 0, + 0, + Map.of( + ResourceType.CPU, + new QueryGroupStats.ResourceStats(0, 0, 0), + ResourceType.MEMORY, + new QueryGroupStats.ResourceStats(0, 0, 0) + ) + ), + QueryGroupTask.DEFAULT_QUERY_GROUP_ID_SUPPLIER.get(), + new QueryGroupStats.QueryGroupStatsHolder( + 0, + 0, + 0, + 0, + 0, Map.of( ResourceType.CPU, new QueryGroupStats.ResourceStats(0, 0, 0), @@ -149,6 +202,21 @@ public void testInvalidQueryGroupFailure() throws IOException { 0, 0, 0, + 0, + Map.of( + ResourceType.CPU, + new QueryGroupStats.ResourceStats(0, 0, 0), + ResourceType.MEMORY, + new QueryGroupStats.ResourceStats(0, 0, 0) + ) + ), + QueryGroupTask.DEFAULT_QUERY_GROUP_ID_SUPPLIER.get(), + new QueryGroupStats.QueryGroupStatsHolder( + 0, + 0, + 1, + 0, + 0, Map.of( ResourceType.CPU, new QueryGroupStats.ResourceStats(0, 0, 0), @@ -169,12 +237,23 @@ private void assertSuccess( QueryGroupStats expectedStats, String threadContextQG_Id ) { - + QueryGroupsStateAccessor stateAccessor = new QueryGroupsStateAccessor(queryGroupStateMap); try (ThreadContext.StoredContext currentContext = testThreadPool.getThreadContext().stashContext()) { testThreadPool.getThreadContext().putHeader(QueryGroupTask.QUERY_GROUP_ID_HEADER, threadContextQG_Id); queryGroupStateMap.put(testQueryGroupId, new QueryGroupState()); - queryGroupService = new QueryGroupService(queryGroupStateMap); + setupMockedQueryGroupsFromClusterState(); + + queryGroupService = new QueryGroupService( + taskCancellationService, + mockClusterService, + testThreadPool, + mockWorkloadManagementSettings, + null, + stateAccessor, + Collections.emptySet(), + Collections.emptySet() + ); sut = new QueryGroupRequestOperationListener(queryGroupService, testThreadPool); sut.onRequestFailure(null, null); @@ -184,4 +263,12 @@ private void assertSuccess( } } + + private void setupMockedQueryGroupsFromClusterState() { + ClusterState state = mock(ClusterState.class); + Metadata metadata = mock(Metadata.class); + when(mockClusterService.state()).thenReturn(state); + when(state.metadata()).thenReturn(metadata); + when(metadata.queryGroups()).thenReturn(Collections.emptyMap()); + } } diff --git a/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStateTests.java b/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStateTests.java index 576eec7be1888..566c4261d6878 100644 --- a/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStateTests.java +++ b/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStateTests.java @@ -23,7 +23,13 @@ public void testRandomQueryGroupsStateUpdates() { for (int i = 0; i < 25; i++) { if (i % 5 == 0) { - updaterThreads.add(new Thread(() -> queryGroupState.completions.inc())); + updaterThreads.add(new Thread(() -> { + if (randomBoolean()) { + queryGroupState.completions.inc(); + } else { + queryGroupState.shardCompletions.inc(); + } + })); } else if (i % 5 == 1) { updaterThreads.add(new Thread(() -> { queryGroupState.totalRejections.inc(); @@ -57,7 +63,7 @@ public void testRandomQueryGroupsStateUpdates() { } }); - assertEquals(5, queryGroupState.getCompletions()); + assertEquals(5, queryGroupState.getCompletions() + queryGroupState.getShardCompletions()); assertEquals(5, queryGroupState.getTotalRejections()); final long sumOfRejectionsDueToResourceTypes = queryGroupState.getResourceState().get(ResourceType.CPU).rejections.count() diff --git a/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStatsTests.java b/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStatsTests.java index 661c3a7beae40..ac6d19580dacb 100644 --- a/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStatsTests.java +++ b/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStatsTests.java @@ -28,9 +28,10 @@ public void testToXContent() throws IOException { queryGroupId, new QueryGroupStats.QueryGroupStatsHolder( 123456789, + 13, 2, 0, - 13, + 1213718, Map.of(ResourceType.CPU, new QueryGroupStats.ResourceStats(0.3, 13, 2)) ) ); @@ -40,7 +41,7 @@ public void testToXContent() throws IOException { queryGroupStats.toXContent(builder, ToXContent.EMPTY_PARAMS); builder.endObject(); assertEquals( - "{\"query_groups\":{\"afakjklaj304041-afaka\":{\"completions\":123456789,\"rejections\":2,\"failures\":0,\"total_cancellations\":13,\"cpu\":{\"current_usage\":0.3,\"cancellations\":13,\"rejections\":2}}}}", + "{\"query_groups\":{\"afakjklaj304041-afaka\":{\"completions\":123456789,\"shard_completions\":1213718,\"rejections\":13,\"failures\":2,\"total_cancellations\":0,\"cpu\":{\"current_usage\":0.3,\"cancellations\":13,\"rejections\":2}}}}", builder.toString() ); } @@ -60,6 +61,7 @@ protected QueryGroupStats createTestInstance() { randomNonNegativeLong(), randomNonNegativeLong(), randomNonNegativeLong(), + randomNonNegativeLong(), Map.of( ResourceType.CPU, new QueryGroupStats.ResourceStats( From 266bdd30054f5d9110bf136c7b4f8f4564f9d767 Mon Sep 17 00:00:00 2001 From: Sagar <99425694+sgup432@users.noreply.github.com> Date: Mon, 7 Oct 2024 19:25:09 -0700 Subject: [PATCH 4/7] [Tiered Caching] Segmented cache changes (#16047) * Segmented cache changes for TieredCache Signed-off-by: Sagar Upadhyaya * Adding change log Signed-off-by: Sagar Upadhyaya * Allow segment number to be power of two Signed-off-by: Sagar Upadhyaya * Moving common tiered cache IT methods to a common base class Signed-off-by: Sagar Upadhyaya * Adding disk took time IT test with multiple segment Signed-off-by: Sagar Upadhyaya * Correcting changelog Signed-off-by: Sagar Upadhyaya * Addressing comments Signed-off-by: Sagar Upadhyaya * Fixing invalid segment count variable name Signed-off-by: Sagar Upadhyaya * Introducing new settings for size for respective cache tier Signed-off-by: Sagar Upadhyaya * Changing the default segmentCount logic Signed-off-by: Sagar Upadhyaya * Fixing missing java doc issue Signed-off-by: Sagar Upadhyaya --------- Signed-off-by: Sagar Upadhyaya Signed-off-by: Sagar <99425694+sgup432@users.noreply.github.com> --- CHANGELOG.md | 1 + .../tier/TieredSpilloverCacheBaseIT.java | 57 + .../common/tier/TieredSpilloverCacheIT.java | 169 ++- .../tier/TieredSpilloverCacheStatsIT.java | 100 +- .../common/tier/TieredSpilloverCache.java | 884 +++++++++----- .../tier/TieredSpilloverCachePlugin.java | 9 + .../tier/TieredSpilloverCacheSettings.java | 69 ++ .../cache/common/tier/MockDiskCache.java | 27 +- .../tier/TieredSpilloverCacheTests.java | 1087 +++++++++++++---- .../cache/EhcacheDiskCacheSettings.java | 2 +- .../cache/store/disk/EhcacheDiskCache.java | 55 +- .../store/disk/EhCacheDiskCacheTests.java | 68 ++ .../org/opensearch/common/cache/Cache.java | 37 +- .../opensearch/common/cache/CacheBuilder.java | 15 +- .../common/cache/settings/CacheSettings.java | 13 + .../cache/store/OpenSearchOnHeapCache.java | 27 +- .../cache/store/builders/ICacheBuilder.java | 11 + .../cache/store/config/CacheConfig.java | 39 + .../indices/IndicesRequestCache.java | 5 +- .../opensearch/indices/IndicesService.java | 2 +- .../opensearch/common/cache/CacheTests.java | 64 +- .../store/OpenSearchOnHeapCacheTests.java | 31 + .../indices/IndicesRequestCacheTests.java | 63 +- 23 files changed, 2170 insertions(+), 665 deletions(-) create mode 100644 modules/cache-common/src/internalClusterTest/java/org/opensearch/cache/common/tier/TieredSpilloverCacheBaseIT.java diff --git a/CHANGELOG.md b/CHANGELOG.md index ab8cdd882bbb9..06ac7740eea6e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Fallback to Remote cluster-state on Term-Version check mismatch - ([#15424](https://github.com/opensearch-project/OpenSearch/pull/15424)) - Implement WithFieldName interface in ValuesSourceAggregationBuilder & FieldSortBuilder ([#15916](https://github.com/opensearch-project/OpenSearch/pull/15916)) - Add successfulSearchShardIndices in searchRequestContext ([#15967](https://github.com/opensearch-project/OpenSearch/pull/15967)) +- [Tiered Caching] Segmented cache changes ([#16047](https://github.com/opensearch-project/OpenSearch/pull/16047)) - Add support for msearch API to pass search pipeline name - ([#15923](https://github.com/opensearch-project/OpenSearch/pull/15923)) - Add _list/indices API as paginated alternate to _cat/indices ([#14718](https://github.com/opensearch-project/OpenSearch/pull/14718)) - Add success and failure metrics for async shard fetch ([#15976](https://github.com/opensearch-project/OpenSearch/pull/15976)) diff --git a/modules/cache-common/src/internalClusterTest/java/org/opensearch/cache/common/tier/TieredSpilloverCacheBaseIT.java b/modules/cache-common/src/internalClusterTest/java/org/opensearch/cache/common/tier/TieredSpilloverCacheBaseIT.java new file mode 100644 index 0000000000000..01371ca8eeefb --- /dev/null +++ b/modules/cache-common/src/internalClusterTest/java/org/opensearch/cache/common/tier/TieredSpilloverCacheBaseIT.java @@ -0,0 +1,57 @@ +/* + * 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.cache.common.tier; + +import org.opensearch.common.cache.CacheType; +import org.opensearch.common.cache.settings.CacheSettings; +import org.opensearch.common.cache.store.OpenSearchOnHeapCache; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.util.FeatureFlags; +import org.opensearch.test.OpenSearchIntegTestCase; + +public class TieredSpilloverCacheBaseIT extends OpenSearchIntegTestCase { + + public Settings defaultSettings(String onHeapCacheSizeInBytesOrPercentage, int numberOfSegments) { + return Settings.builder() + .put(FeatureFlags.PLUGGABLE_CACHE, "true") + .put( + CacheSettings.getConcreteStoreNameSettingForCacheType(CacheType.INDICES_REQUEST_CACHE).getKey(), + TieredSpilloverCache.TieredSpilloverCacheFactory.TIERED_SPILLOVER_CACHE_NAME + ) + .put( + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_NAME.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), + OpenSearchOnHeapCache.OpenSearchOnHeapCacheFactory.NAME + ) + .put( + TieredSpilloverCacheSettings.TIERED_SPILLOVER_DISK_STORE_NAME.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), + MockDiskCache.MockDiskCacheFactory.NAME + ) + .put( + TieredSpilloverCacheSettings.TIERED_SPILLOVER_SEGMENTS.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), + numberOfSegments + ) + .put( + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), + onHeapCacheSizeInBytesOrPercentage + ) + .build(); + } + + public int getNumberOfSegments() { + return randomFrom(1, 2, 4, 8, 16, 64, 128, 256); + } +} diff --git a/modules/cache-common/src/internalClusterTest/java/org/opensearch/cache/common/tier/TieredSpilloverCacheIT.java b/modules/cache-common/src/internalClusterTest/java/org/opensearch/cache/common/tier/TieredSpilloverCacheIT.java index 02be0990eb136..d58e36c036510 100644 --- a/modules/cache-common/src/internalClusterTest/java/org/opensearch/cache/common/tier/TieredSpilloverCacheIT.java +++ b/modules/cache-common/src/internalClusterTest/java/org/opensearch/cache/common/tier/TieredSpilloverCacheIT.java @@ -22,12 +22,8 @@ import org.opensearch.cluster.metadata.IndexMetadata; import org.opensearch.common.cache.CacheType; import org.opensearch.common.cache.ICache; -import org.opensearch.common.cache.settings.CacheSettings; -import org.opensearch.common.cache.store.OpenSearchOnHeapCache; -import org.opensearch.common.cache.store.settings.OpenSearchOnHeapCacheSettings; import org.opensearch.common.settings.Settings; import org.opensearch.common.unit.TimeValue; -import org.opensearch.common.util.FeatureFlags; import org.opensearch.index.cache.request.RequestCacheStats; import org.opensearch.index.query.QueryBuilders; import org.opensearch.indices.IndicesRequestCache; @@ -43,13 +39,15 @@ import java.util.Arrays; import java.util.Collection; import java.util.List; +import java.util.Locale; import java.util.Map; +import java.util.UUID; import java.util.concurrent.TimeUnit; import java.util.function.Function; import java.util.stream.Collectors; import java.util.stream.Stream; -import static org.opensearch.common.cache.store.settings.OpenSearchOnHeapCacheSettings.MAXIMUM_SIZE_IN_BYTES_KEY; +import static org.opensearch.common.cache.settings.CacheSettings.INVALID_SEGMENT_COUNT_EXCEPTION_MESSAGE; import static org.opensearch.indices.IndicesService.INDICES_CACHE_CLEAN_INTERVAL_SETTING; import static org.opensearch.search.aggregations.AggregationBuilders.dateHistogram; import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertAcked; @@ -58,43 +56,15 @@ import static org.hamcrest.Matchers.greaterThan; @OpenSearchIntegTestCase.ClusterScope(numDataNodes = 0, scope = OpenSearchIntegTestCase.Scope.TEST) -public class TieredSpilloverCacheIT extends OpenSearchIntegTestCase { +public class TieredSpilloverCacheIT extends TieredSpilloverCacheBaseIT { @Override protected Collection> nodePlugins() { return Arrays.asList(TieredSpilloverCachePlugin.class, MockDiskCachePlugin.class); } - static Settings defaultSettings(String onHeapCacheSizeInBytesOrPercentage) { - return Settings.builder() - .put(FeatureFlags.PLUGGABLE_CACHE, "true") - .put( - CacheSettings.getConcreteStoreNameSettingForCacheType(CacheType.INDICES_REQUEST_CACHE).getKey(), - TieredSpilloverCache.TieredSpilloverCacheFactory.TIERED_SPILLOVER_CACHE_NAME - ) - .put( - TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_NAME.getConcreteSettingForNamespace( - CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() - ).getKey(), - OpenSearchOnHeapCache.OpenSearchOnHeapCacheFactory.NAME - ) - .put( - TieredSpilloverCacheSettings.TIERED_SPILLOVER_DISK_STORE_NAME.getConcreteSettingForNamespace( - CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() - ).getKey(), - MockDiskCache.MockDiskCacheFactory.NAME - ) - .put( - OpenSearchOnHeapCacheSettings.getSettingListForCacheType(CacheType.INDICES_REQUEST_CACHE) - .get(MAXIMUM_SIZE_IN_BYTES_KEY) - .getKey(), - onHeapCacheSizeInBytesOrPercentage - ) - .build(); - } - public void testPluginsAreInstalled() { - internalCluster().startNode(Settings.builder().put(defaultSettings("1%")).build()); + internalCluster().startNode(Settings.builder().put(defaultSettings("1%", getNumberOfSegments())).build()); NodesInfoRequest nodesInfoRequest = new NodesInfoRequest(); nodesInfoRequest.addMetric(NodesInfoRequest.Metric.PLUGINS.metricName()); NodesInfoResponse nodesInfoResponse = OpenSearchIntegTestCase.client().admin().cluster().nodesInfo(nodesInfoRequest).actionGet(); @@ -111,7 +81,8 @@ public void testPluginsAreInstalled() { } public void testSanityChecksWithIndicesRequestCache() throws InterruptedException { - internalCluster().startNodes(3, Settings.builder().put(defaultSettings("1%")).build()); + int numberOfSegments = getNumberOfSegments(); + internalCluster().startNodes(3, Settings.builder().put(defaultSettings("1%", numberOfSegments)).build()); Client client = client(); assertAcked( client.admin() @@ -147,9 +118,97 @@ public void testSanityChecksWithIndicesRequestCache() throws InterruptedExceptio ); } + public void testWithDynamicTookTimePolicyWithMultiSegments() throws Exception { + int numberOfSegments = getNumberOfSegments(); + int onHeapCacheSizePerSegmentInBytes = 800; // Per cache entry below is around ~700 bytes, so keeping this + // just a bit higher so that each segment can atleast hold 1 entry. + int onHeapCacheSizeInBytes = onHeapCacheSizePerSegmentInBytes * numberOfSegments; + internalCluster().startNode(Settings.builder().put(defaultSettings(onHeapCacheSizeInBytes + "b", numberOfSegments)).build()); + Client client = client(); + assertAcked( + client.admin() + .indices() + .prepareCreate("index") + .setMapping("k", "type=keyword") + .setSettings( + Settings.builder() + .put(IndicesRequestCache.INDEX_CACHE_REQUEST_ENABLED_SETTING.getKey(), true) + .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1) + .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0) + .put("index.refresh_interval", -1) + ) + .get() + ); + // Set a very high value for took time policy so that no items evicted from onHeap cache are spilled + // to disk. And then hit requests so that few items are cached into cache. + ClusterUpdateSettingsRequest updateSettingsRequest = new ClusterUpdateSettingsRequest().transientSettings( + Settings.builder() + .put( + TieredSpilloverCacheSettings.TOOK_TIME_POLICY_CONCRETE_SETTINGS_MAP.get(CacheType.INDICES_REQUEST_CACHE).getKey(), + new TimeValue(100, TimeUnit.SECONDS) + ) + .build() + ); + assertAcked(internalCluster().client().admin().cluster().updateSettings(updateSettingsRequest).get()); + int numberOfIndexedItems = numberOfSegments + 1; // Best case if all keys are distributed among different + // segment, atleast one of the segment will have 2 entries and we will see evictions. + for (int iterator = 0; iterator < numberOfIndexedItems; iterator++) { + indexRandom(true, client.prepareIndex("index").setSource("k" + iterator, "hello" + iterator)); + } + ensureSearchable("index"); + refreshAndWaitForReplication(); + // Force merge the index to ensure there can be no background merges during the subsequent searches that would invalidate the cache + ForceMergeResponse forceMergeResponse = client.admin().indices().prepareForceMerge("index").setFlush(true).get(); + OpenSearchAssertions.assertAllSuccessful(forceMergeResponse); + long perQuerySizeInCacheInBytes = -1; + for (int iterator = 0; iterator < numberOfIndexedItems; iterator++) { + SearchResponse resp = client.prepareSearch("index") + .setRequestCache(true) + .setQuery(QueryBuilders.termQuery("k" + iterator, "hello" + iterator)) + .get(); + if (perQuerySizeInCacheInBytes == -1) { + RequestCacheStats requestCacheStats = getRequestCacheStats(client, "index"); + perQuerySizeInCacheInBytes = requestCacheStats.getMemorySizeInBytes(); + } + assertSearchResponse(resp); + } + RequestCacheStats requestCacheStats = getRequestCacheStats(client, "index"); + // Considering disk cache won't be used due to took time policy having a high value, we expect overall cache + // size to be less than or equal to onHeapCache size. + assertTrue(requestCacheStats.getMemorySizeInBytes() <= onHeapCacheSizeInBytes); + assertEquals(numberOfIndexedItems, requestCacheStats.getMissCount()); + // We should atleast one eviction considering disk cache isn't able to hold anything due to policy. + assertTrue(requestCacheStats.getEvictions() > 0); + assertEquals(0, requestCacheStats.getHitCount()); + long lastEvictionSeen = requestCacheStats.getEvictions(); + + // Decrease took time policy to zero so that disk cache also comes into play. Now we should be able + // to cache all entries. + updateSettingsRequest = new ClusterUpdateSettingsRequest().transientSettings( + Settings.builder() + .put( + TieredSpilloverCacheSettings.TOOK_TIME_POLICY_CONCRETE_SETTINGS_MAP.get(CacheType.INDICES_REQUEST_CACHE).getKey(), + new TimeValue(0, TimeUnit.MILLISECONDS) + ) + .build() + ); + assertAcked(internalCluster().client().admin().cluster().updateSettings(updateSettingsRequest).get()); + for (int iterator = 0; iterator < numberOfIndexedItems * 2; iterator++) { + SearchResponse resp = client.prepareSearch("index") + .setRequestCache(true) + .setQuery(QueryBuilders.termQuery(UUID.randomUUID().toString(), UUID.randomUUID().toString())) + .get(); + assertSearchResponse(resp); + } + + requestCacheStats = getRequestCacheStats(client, "index"); + // We shouldn't see any new evictions now. + assertEquals(lastEvictionSeen, requestCacheStats.getEvictions()); + } + public void testWithDynamicTookTimePolicy() throws Exception { int onHeapCacheSizeInBytes = 2000; - internalCluster().startNode(Settings.builder().put(defaultSettings(onHeapCacheSizeInBytes + "b")).build()); + internalCluster().startNode(Settings.builder().put(defaultSettings(onHeapCacheSizeInBytes + "b", 1)).build()); Client client = client(); assertAcked( client.admin() @@ -271,9 +330,10 @@ public void testWithDynamicTookTimePolicy() throws Exception { public void testInvalidationWithIndicesRequestCache() throws Exception { int onHeapCacheSizeInBytes = 2000; + int numberOfSegments = getNumberOfSegments(); internalCluster().startNode( Settings.builder() - .put(defaultSettings(onHeapCacheSizeInBytes + "b")) + .put(defaultSettings(onHeapCacheSizeInBytes + "b", numberOfSegments)) .put(INDICES_CACHE_CLEAN_INTERVAL_SETTING.getKey(), new TimeValue(1)) .build() ); @@ -354,10 +414,11 @@ public void testInvalidationWithIndicesRequestCache() throws Exception { } public void testWithExplicitCacheClear() throws Exception { + int numberOfSegments = getNumberOfSegments(); int onHeapCacheSizeInBytes = 2000; internalCluster().startNode( Settings.builder() - .put(defaultSettings(onHeapCacheSizeInBytes + "b")) + .put(defaultSettings(onHeapCacheSizeInBytes + "b", numberOfSegments)) .put(INDICES_CACHE_CLEAN_INTERVAL_SETTING.getKey(), new TimeValue(1)) .build() ); @@ -426,10 +487,13 @@ public void testWithExplicitCacheClear() throws Exception { } public void testWithDynamicDiskCacheSetting() throws Exception { - int onHeapCacheSizeInBytes = 10; // Keep it low so that all items are cached onto disk. + int numberOfSegments = getNumberOfSegments(); + int onHeapCacheSizeInBytes = randomIntBetween(numberOfSegments + 1, numberOfSegments * 2); // Keep it low so + // that all items are + // cached onto disk. internalCluster().startNode( Settings.builder() - .put(defaultSettings(onHeapCacheSizeInBytes + "b")) + .put(defaultSettings(onHeapCacheSizeInBytes + "b", numberOfSegments)) .put(INDICES_CACHE_CLEAN_INTERVAL_SETTING.getKey(), new TimeValue(1)) .build() ); @@ -540,6 +604,27 @@ public void testWithDynamicDiskCacheSetting() throws Exception { assertEquals(0, lastKnownHitCount - requestCacheStats.getHitCount()); } + public void testWithInvalidSegmentNumberSetting() throws Exception { + int numberOfSegments = getNumberOfSegments(); + int onHeapCacheSizeInBytes = randomIntBetween(numberOfSegments + 1, numberOfSegments * 2); // Keep it low so + // that all items are + // cached onto disk. + assertThrows( + String.format( + Locale.ROOT, + INVALID_SEGMENT_COUNT_EXCEPTION_MESSAGE, + TieredSpilloverCache.TieredSpilloverCacheFactory.TIERED_SPILLOVER_CACHE_NAME + ), + IllegalArgumentException.class, + () -> internalCluster().startNode( + Settings.builder() + .put(defaultSettings(onHeapCacheSizeInBytes + "b", 300)) + .put(INDICES_CACHE_CLEAN_INTERVAL_SETTING.getKey(), new TimeValue(1)) + .build() + ) + ); + } + private RequestCacheStats getRequestCacheStats(Client client, String indexName) { return client.admin().indices().prepareStats(indexName).setRequestCache(true).get().getTotal().getRequestCache(); } @@ -550,7 +635,7 @@ public MockDiskCachePlugin() {} @Override public Map getCacheFactoryMap() { - return Map.of(MockDiskCache.MockDiskCacheFactory.NAME, new MockDiskCache.MockDiskCacheFactory(0, 1000, false)); + return Map.of(MockDiskCache.MockDiskCacheFactory.NAME, new MockDiskCache.MockDiskCacheFactory(0, 10000, false, 1)); } @Override diff --git a/modules/cache-common/src/internalClusterTest/java/org/opensearch/cache/common/tier/TieredSpilloverCacheStatsIT.java b/modules/cache-common/src/internalClusterTest/java/org/opensearch/cache/common/tier/TieredSpilloverCacheStatsIT.java index 783b6083e9226..fe6bd7050a8f3 100644 --- a/modules/cache-common/src/internalClusterTest/java/org/opensearch/cache/common/tier/TieredSpilloverCacheStatsIT.java +++ b/modules/cache-common/src/internalClusterTest/java/org/opensearch/cache/common/tier/TieredSpilloverCacheStatsIT.java @@ -45,16 +45,16 @@ // Use a single data node to simplify accessing cache stats across different shards. @OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.TEST, numDataNodes = 0) -public class TieredSpilloverCacheStatsIT extends OpenSearchIntegTestCase { +public class TieredSpilloverCacheStatsIT extends TieredSpilloverCacheBaseIT { @Override protected Collection> nodePlugins() { return Arrays.asList(TieredSpilloverCachePlugin.class, TieredSpilloverCacheIT.MockDiskCachePlugin.class); } - private final String HEAP_CACHE_SIZE_STRING = "10000B"; - private final int HEAP_CACHE_SIZE = 10_000; - private final String index1Name = "index1"; - private final String index2Name = "index2"; + private static final String HEAP_CACHE_SIZE_STRING = "10000B"; + private static final int HEAP_CACHE_SIZE = 10_000; + private static final String index1Name = "index1"; + private static final String index2Name = "index2"; /** * Test aggregating by indices @@ -63,7 +63,7 @@ public void testIndicesLevelAggregation() throws Exception { internalCluster().startNodes( 1, Settings.builder() - .put(TieredSpilloverCacheIT.defaultSettings(HEAP_CACHE_SIZE_STRING)) + .put(defaultSettings(HEAP_CACHE_SIZE_STRING, 1)) .put( TieredSpilloverCacheSettings.TOOK_TIME_POLICY_CONCRETE_SETTINGS_MAP.get(CacheType.INDICES_REQUEST_CACHE).getKey(), new TimeValue(0, TimeUnit.SECONDS) @@ -116,7 +116,7 @@ public void testIndicesAndTierLevelAggregation() throws Exception { internalCluster().startNodes( 1, Settings.builder() - .put(TieredSpilloverCacheIT.defaultSettings(HEAP_CACHE_SIZE_STRING)) + .put(defaultSettings(HEAP_CACHE_SIZE_STRING, 1)) .put( TieredSpilloverCacheSettings.TOOK_TIME_POLICY_CONCRETE_SETTINGS_MAP.get(CacheType.INDICES_REQUEST_CACHE).getKey(), new TimeValue(0, TimeUnit.SECONDS) @@ -196,7 +196,7 @@ public void testTierLevelAggregation() throws Exception { internalCluster().startNodes( 1, Settings.builder() - .put(TieredSpilloverCacheIT.defaultSettings(HEAP_CACHE_SIZE_STRING)) + .put(defaultSettings(HEAP_CACHE_SIZE_STRING, 1)) .put( TieredSpilloverCacheSettings.TOOK_TIME_POLICY_CONCRETE_SETTINGS_MAP.get(CacheType.INDICES_REQUEST_CACHE).getKey(), new TimeValue(0, TimeUnit.SECONDS) @@ -205,7 +205,6 @@ public void testTierLevelAggregation() throws Exception { ); Client client = client(); Map values = setupCacheForAggregationTests(client); - // Get values for tiers alone and check they add correctly across indices ImmutableCacheStatsHolder tiersOnlyStatsHolder = getNodeCacheStatsResult(client, List.of(TIER_DIMENSION_NAME)); ImmutableCacheStats totalHeapExpectedStats = returnNullIfAllZero( @@ -238,7 +237,7 @@ public void testInvalidLevelsAreIgnored() throws Exception { internalCluster().startNodes( 1, Settings.builder() - .put(TieredSpilloverCacheIT.defaultSettings(HEAP_CACHE_SIZE_STRING)) + .put(defaultSettings(HEAP_CACHE_SIZE_STRING, getNumberOfSegments())) .put( TieredSpilloverCacheSettings.TOOK_TIME_POLICY_CONCRETE_SETTINGS_MAP.get(CacheType.INDICES_REQUEST_CACHE).getKey(), new TimeValue(0, TimeUnit.SECONDS) @@ -289,7 +288,7 @@ public void testStatsMatchOldApi() throws Exception { internalCluster().startNodes( 1, Settings.builder() - .put(TieredSpilloverCacheIT.defaultSettings(HEAP_CACHE_SIZE_STRING)) + .put(defaultSettings(HEAP_CACHE_SIZE_STRING, getNumberOfSegments())) .put( TieredSpilloverCacheSettings.TOOK_TIME_POLICY_CONCRETE_SETTINGS_MAP.get(CacheType.INDICES_REQUEST_CACHE).getKey(), new TimeValue(0, TimeUnit.SECONDS) @@ -342,6 +341,82 @@ public void testStatsMatchOldApi() throws Exception { assertEquals(oldAPIStats.getMemorySizeInBytes(), totalStats.getSizeInBytes()); } + public void testStatsWithMultipleSegments() throws Exception { + int numberOfSegments = randomFrom(2, 4, 8, 16, 64); + int singleSearchSizeApproxUpperBound = 700; // We know this from other tests and manually verifying + int heap_cache_size_per_segment = singleSearchSizeApproxUpperBound * numberOfSegments; // Worst case if all + // keys land up in same segment, it would still be able to accommodate. + internalCluster().startNodes( + 1, + Settings.builder() + .put(defaultSettings(heap_cache_size_per_segment * numberOfSegments + "B", numberOfSegments)) + .put( + TieredSpilloverCacheSettings.TOOK_TIME_POLICY_CONCRETE_SETTINGS_MAP.get(CacheType.INDICES_REQUEST_CACHE).getKey(), + new TimeValue(0, TimeUnit.SECONDS) + ) + .build() + ); + Client client = client(); + startIndex(client, index1Name); + // First search one time to calculate item size + searchIndex(client, index1Name, 0); + // get total stats + long singleSearchSize = getTotalStats(client).getSizeInBytes(); + // Now try to hit queries same as number of segments. All these should be able to reside inside onHeap cache. + for (int i = 1; i < numberOfSegments; i++) { + searchIndex(client, index1Name, i); + } + ImmutableCacheStatsHolder allLevelsStatsHolder = getNodeCacheStatsResult( + client, + List.of(IndicesRequestCache.INDEX_DIMENSION_NAME, TIER_DIMENSION_NAME) + ); + ImmutableCacheStats index1OnHeapExpectedStats = returnNullIfAllZero( + new ImmutableCacheStats(0, numberOfSegments, 0, singleSearchSize * numberOfSegments, numberOfSegments) + ); + assertEquals( + index1OnHeapExpectedStats, + allLevelsStatsHolder.getStatsForDimensionValues(List.of(index1Name, TIER_DIMENSION_VALUE_ON_HEAP)) + ); + ImmutableCacheStats index1DiskCacheExpectedStats = returnNullIfAllZero(new ImmutableCacheStats(0, numberOfSegments, 0, 0, 0)); + assertEquals( + index1DiskCacheExpectedStats, + allLevelsStatsHolder.getStatsForDimensionValues(List.of(index1Name, TIER_DIMENSION_VALUE_DISK)) + ); + + // Now fire same queries to get some hits + for (int i = 0; i < numberOfSegments; i++) { + searchIndex(client, index1Name, i); + } + allLevelsStatsHolder = getNodeCacheStatsResult(client, List.of(IndicesRequestCache.INDEX_DIMENSION_NAME, TIER_DIMENSION_NAME)); + index1OnHeapExpectedStats = returnNullIfAllZero( + new ImmutableCacheStats(numberOfSegments, numberOfSegments, 0, singleSearchSize * numberOfSegments, numberOfSegments) + ); + assertEquals( + index1OnHeapExpectedStats, + allLevelsStatsHolder.getStatsForDimensionValues(List.of(index1Name, TIER_DIMENSION_VALUE_ON_HEAP)) + ); + + // Now try to evict from onheap cache by adding numberOfSegments ^ 2 which will guarantee this. + for (int i = numberOfSegments; i < numberOfSegments + numberOfSegments * numberOfSegments; i++) { + searchIndex(client, index1Name, i); + } + allLevelsStatsHolder = getNodeCacheStatsResult(client, List.of(IndicesRequestCache.INDEX_DIMENSION_NAME, TIER_DIMENSION_NAME)); + ImmutableCacheStats onHeapCacheStat = allLevelsStatsHolder.getStatsForDimensionValues( + List.of(index1Name, TIER_DIMENSION_VALUE_ON_HEAP) + ); + // Jut verifying evictions happened as can't fetch the exact number considering we don't have a way to get + // segment number for queries. + assertTrue(onHeapCacheStat.getEvictions() > 0); + ImmutableCacheStats diskCacheStat = allLevelsStatsHolder.getStatsForDimensionValues(List.of(index1Name, TIER_DIMENSION_VALUE_DISK)); + + // Similarly verify items are present on disk cache now + assertEquals(onHeapCacheStat.getEvictions(), diskCacheStat.getItems()); + assertTrue(diskCacheStat.getSizeInBytes() > 0); + assertTrue(diskCacheStat.getMisses() > 0); + assertTrue(diskCacheStat.getHits() == 0); + assertTrue(diskCacheStat.getEvictions() == 0); + } + private void startIndex(Client client, String indexName) throws InterruptedException { assertAcked( client.admin() @@ -373,6 +448,7 @@ private Map setupCacheForAggregationTests(Client client) throws searchIndex(client, index1Name, 0); // get total stats long singleSearchSize = getTotalStats(client).getSizeInBytes(); + int itemsOnHeapAfterTest = HEAP_CACHE_SIZE / (int) singleSearchSize; // As the heap tier evicts, the items on it after the test will // be the same as its max capacity int itemsOnDiskAfterTest = 1 + randomInt(30); // The first one we search (to get the size) always goes to disk @@ -416,7 +492,6 @@ private Map setupCacheForAggregationTests(Client client) throws for (int i = itemsOnDiskAfterTest + itemsOnHeapIndex1AfterTest; i < itemsOnDiskAfterTest + itemsOnHeapAfterTest; i++) { searchIndex(client, index2Name, i); } - // Get some hits on all combinations of indices and tiers for (int i = itemsOnDiskAfterTest; i < itemsOnDiskAfterTest + hitsOnHeapIndex1; i++) { // heap hits for index 1 @@ -499,6 +574,7 @@ private static ImmutableCacheStatsHolder getNodeCacheStatsResult(Client client, .addMetric(NodesStatsRequest.Metric.CACHE_STATS.metricName()) .setIndices(statsFlags) .get(); + // Can always get the first data node as there's only one in this test suite assertEquals(1, nodeStatsResponse.getNodes().size()); NodeCacheStats ncs = nodeStatsResponse.getNodes().get(0).getNodeCacheStats(); diff --git a/modules/cache-common/src/main/java/org/opensearch/cache/common/tier/TieredSpilloverCache.java b/modules/cache-common/src/main/java/org/opensearch/cache/common/tier/TieredSpilloverCache.java index f69c56808b2a1..ab5335ca0ca66 100644 --- a/modules/cache-common/src/main/java/org/opensearch/cache/common/tier/TieredSpilloverCache.java +++ b/modules/cache-common/src/main/java/org/opensearch/cache/common/tier/TieredSpilloverCache.java @@ -34,6 +34,7 @@ import java.util.Iterator; import java.util.LinkedHashMap; import java.util.List; +import java.util.Locale; import java.util.Map; import java.util.NoSuchElementException; import java.util.Objects; @@ -49,8 +50,13 @@ import java.util.function.ToLongBiFunction; import static org.opensearch.cache.common.tier.TieredSpilloverCacheSettings.DISK_CACHE_ENABLED_SETTING_MAP; +import static org.opensearch.cache.common.tier.TieredSpilloverCacheSettings.TIERED_SPILLOVER_DISK_STORE_SIZE; +import static org.opensearch.cache.common.tier.TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE; +import static org.opensearch.cache.common.tier.TieredSpilloverCacheSettings.TIERED_SPILLOVER_SEGMENTS; import static org.opensearch.cache.common.tier.TieredSpilloverCacheStatsHolder.TIER_DIMENSION_VALUE_DISK; import static org.opensearch.cache.common.tier.TieredSpilloverCacheStatsHolder.TIER_DIMENSION_VALUE_ON_HEAP; +import static org.opensearch.common.cache.settings.CacheSettings.INVALID_SEGMENT_COUNT_EXCEPTION_MESSAGE; +import static org.opensearch.common.cache.settings.CacheSettings.VALID_SEGMENT_COUNT_VALUES; /** * This cache spillover the evicted items from heap tier to disk tier. All the new items are first cached on heap @@ -69,29 +75,16 @@ public class TieredSpilloverCache implements ICache { private static final List SPILLOVER_REMOVAL_REASONS = List.of(RemovalReason.EVICTED, RemovalReason.CAPACITY); private static final Logger logger = LogManager.getLogger(TieredSpilloverCache.class); - private final ICache diskCache; - private final ICache onHeapCache; - - // Removal listeners for the individual tiers - private final RemovalListener, V> onDiskRemovalListener; - private final RemovalListener, V> onHeapRemovalListener; - - // Removal listener from the spillover cache as a whole - private final RemovalListener, V> removalListener; + static final String ZERO_SEGMENT_COUNT_EXCEPTION_MESSAGE = "Segment count cannot be less than one for tiered cache"; // In future we want to just read the stats from the individual tiers' statsHolder objects, but this isn't // possible right now because of the way computeIfAbsent is implemented. private final TieredSpilloverCacheStatsHolder statsHolder; - private ToLongBiFunction, V> weigher; private final List dimensionNames; - ReadWriteLock readWriteLock = new ReentrantReadWriteLock(); - ReleasableLock readLock = new ReleasableLock(readWriteLock.readLock()); - ReleasableLock writeLock = new ReleasableLock(readWriteLock.writeLock()); - /** - * Maintains caching tiers in ascending order of cache latency. - */ - private final Map, TierInfo> caches; - private final List> policies; + + private final int numberOfSegments; + + final TieredSpilloverCacheSegment[] tieredSpilloverCacheSegments; /** * This map is used to handle concurrent requests for same key in computeIfAbsent() to ensure we load the value @@ -99,224 +92,525 @@ public class TieredSpilloverCache implements ICache { */ Map, CompletableFuture, V>>> completableFutureMap = new ConcurrentHashMap<>(); + @SuppressWarnings({ "unchecked", "rawtypes" }) TieredSpilloverCache(Builder builder) { - Objects.requireNonNull(builder.onHeapCacheFactory, "onHeap cache builder can't be null"); - Objects.requireNonNull(builder.diskCacheFactory, "disk cache builder can't be null"); Objects.requireNonNull(builder.cacheConfig, "cache config can't be null"); - Objects.requireNonNull(builder.cacheConfig.getClusterSettings(), "cluster settings can't be null"); - this.removalListener = Objects.requireNonNull(builder.removalListener, "Removal listener can't be null"); - - this.onHeapRemovalListener = new HeapTierRemovalListener(this); - this.onDiskRemovalListener = new DiskTierRemovalListener(this); - this.weigher = Objects.requireNonNull(builder.cacheConfig.getWeigher(), "Weigher can't be null"); - - this.onHeapCache = builder.onHeapCacheFactory.create( - new CacheConfig.Builder().setRemovalListener(onHeapRemovalListener) - .setKeyType(builder.cacheConfig.getKeyType()) - .setValueType(builder.cacheConfig.getValueType()) - .setSettings(builder.cacheConfig.getSettings()) - .setWeigher(builder.cacheConfig.getWeigher()) - .setDimensionNames(builder.cacheConfig.getDimensionNames()) - .setMaxSizeInBytes(builder.cacheConfig.getMaxSizeInBytes()) - .setExpireAfterAccess(builder.cacheConfig.getExpireAfterAccess()) - .setClusterSettings(builder.cacheConfig.getClusterSettings()) - .setStatsTrackingEnabled(false) - .build(), - builder.cacheType, - builder.cacheFactories - - ); - this.diskCache = builder.diskCacheFactory.create( - new CacheConfig.Builder().setRemovalListener(onDiskRemovalListener) - .setKeyType(builder.cacheConfig.getKeyType()) - .setValueType(builder.cacheConfig.getValueType()) - .setSettings(builder.cacheConfig.getSettings()) - .setWeigher(builder.cacheConfig.getWeigher()) - .setKeySerializer(builder.cacheConfig.getKeySerializer()) - .setValueSerializer(builder.cacheConfig.getValueSerializer()) - .setDimensionNames(builder.cacheConfig.getDimensionNames()) - .setStatsTrackingEnabled(false) - .build(), - builder.cacheType, - builder.cacheFactories - ); + Objects.requireNonNull(builder.cacheConfig.getSettings(), "settings can't be null"); + if (builder.numberOfSegments <= 0) { + throw new IllegalArgumentException(ZERO_SEGMENT_COUNT_EXCEPTION_MESSAGE); + } + this.numberOfSegments = builder.numberOfSegments; Boolean isDiskCacheEnabled = DISK_CACHE_ENABLED_SETTING_MAP.get(builder.cacheType).get(builder.cacheConfig.getSettings()); - LinkedHashMap, TierInfo> cacheListMap = new LinkedHashMap<>(); - cacheListMap.put(onHeapCache, new TierInfo(true, TIER_DIMENSION_VALUE_ON_HEAP)); - cacheListMap.put(diskCache, new TierInfo(isDiskCacheEnabled, TIER_DIMENSION_VALUE_DISK)); - this.caches = Collections.synchronizedMap(cacheListMap); - this.dimensionNames = builder.cacheConfig.getDimensionNames(); // Pass "tier" as the innermost dimension name, in addition to whatever dimensions are specified for the cache as a whole this.statsHolder = new TieredSpilloverCacheStatsHolder(dimensionNames, isDiskCacheEnabled); - this.policies = builder.policies; // Will never be null; builder initializes it to an empty list + long onHeapCachePerSegmentSizeInBytes = builder.onHeapCacheSizeInBytes / this.numberOfSegments; + long diskCachePerSegmentSizeInBytes = builder.diskCacheSizeInBytes / this.numberOfSegments; + if (onHeapCachePerSegmentSizeInBytes <= 0) { + throw new IllegalArgumentException("Per segment size for onHeap cache within Tiered cache should be " + "greater than 0"); + } + if (diskCachePerSegmentSizeInBytes <= 0) { + throw new IllegalArgumentException("Per segment size for disk cache within Tiered cache should be " + "greater than 0"); + } + this.tieredSpilloverCacheSegments = new TieredSpilloverCacheSegment[this.numberOfSegments]; + for (int i = 0; i < numberOfSegments; i++) { + tieredSpilloverCacheSegments[i] = new TieredSpilloverCacheSegment( + builder, + statsHolder, + i + 1, + this.numberOfSegments, + onHeapCachePerSegmentSizeInBytes, + diskCachePerSegmentSizeInBytes + ); + } builder.cacheConfig.getClusterSettings() .addSettingsUpdateConsumer(DISK_CACHE_ENABLED_SETTING_MAP.get(builder.cacheType), this::enableDisableDiskCache); } - // Package private for testing - ICache getOnHeapCache() { - return onHeapCache; - } + static class TieredSpilloverCacheSegment implements ICache { - // Package private for testing - ICache getDiskCache() { - return diskCache; - } + private final ICache diskCache; + private final ICache onHeapCache; - // Package private for testing. - void enableDisableDiskCache(Boolean isDiskCacheEnabled) { - // When disk cache is disabled, we are not clearing up the disk cache entries yet as that should be part of - // separate cache/clear API. - this.caches.put(diskCache, new TierInfo(isDiskCacheEnabled, TIER_DIMENSION_VALUE_DISK)); - this.statsHolder.setDiskCacheEnabled(isDiskCacheEnabled); - } + // Removal listeners for the individual tiers + private final RemovalListener, V> onDiskRemovalListener; + private final RemovalListener, V> onHeapRemovalListener; - @Override - public V get(ICacheKey key) { - Tuple cacheValueTuple = getValueFromTieredCache(true).apply(key); - if (cacheValueTuple == null) { - return null; + // Removal listener from the spillover cache as a whole + private final RemovalListener, V> removalListener; + + private ToLongBiFunction, V> weigher; + ReadWriteLock readWriteLock = new ReentrantReadWriteLock(); + ReleasableLock readLock = new ReleasableLock(readWriteLock.readLock()); + ReleasableLock writeLock = new ReleasableLock(readWriteLock.writeLock()); + + private final Map, TierInfo> caches; + + private final List> policies; + + private final TieredSpilloverCacheStatsHolder statsHolder; + + /** + * This map is used to handle concurrent requests for same key in computeIfAbsent() to ensure we load the value + * only once. + */ + Map, CompletableFuture, V>>> completableFutureMap = new ConcurrentHashMap<>(); + + TieredSpilloverCacheSegment( + Builder builder, + TieredSpilloverCacheStatsHolder statsHolder, + int segmentNumber, + int numberOfSegments, + long onHeapCacheSizeInBytes, + long diskCacheSizeInBytes + ) { + Objects.requireNonNull(builder.onHeapCacheFactory, "onHeap cache builder can't be null"); + Objects.requireNonNull(builder.diskCacheFactory, "disk cache builder can't be null"); + Objects.requireNonNull(builder.cacheConfig, "cache config can't be null"); + Objects.requireNonNull(builder.cacheConfig.getClusterSettings(), "cluster settings can't be null"); + Objects.requireNonNull(builder.cacheConfig.getStoragePath(), "Storage path shouldn't be null"); + this.removalListener = Objects.requireNonNull(builder.removalListener, "Removal listener can't be null"); + this.statsHolder = statsHolder; + + this.onHeapRemovalListener = new HeapTierRemovalListener(this); + this.onDiskRemovalListener = new DiskTierRemovalListener(this); + this.weigher = Objects.requireNonNull(builder.cacheConfig.getWeigher(), "Weigher can't be null"); + this.onHeapCache = builder.onHeapCacheFactory.create( + new CacheConfig.Builder().setRemovalListener(onHeapRemovalListener) + .setKeyType(builder.cacheConfig.getKeyType()) + .setValueType(builder.cacheConfig.getValueType()) + .setSettings(builder.cacheConfig.getSettings()) + .setWeigher(builder.cacheConfig.getWeigher()) + .setDimensionNames(builder.cacheConfig.getDimensionNames()) + .setMaxSizeInBytes(onHeapCacheSizeInBytes) + .setExpireAfterAccess(builder.cacheConfig.getExpireAfterAccess()) + .setClusterSettings(builder.cacheConfig.getClusterSettings()) + .setSegmentCount(1) // We don't need to make underlying caches multi-segmented + .setStatsTrackingEnabled(false) + .setCacheAlias("tiered_on_heap#" + segmentNumber) + .build(), + builder.cacheType, + builder.cacheFactories + + ); + this.diskCache = builder.diskCacheFactory.create( + new CacheConfig.Builder().setRemovalListener(onDiskRemovalListener) + .setKeyType(builder.cacheConfig.getKeyType()) + .setValueType(builder.cacheConfig.getValueType()) + .setSettings(builder.cacheConfig.getSettings()) + .setWeigher(builder.cacheConfig.getWeigher()) + .setKeySerializer(builder.cacheConfig.getKeySerializer()) + .setValueSerializer(builder.cacheConfig.getValueSerializer()) + .setDimensionNames(builder.cacheConfig.getDimensionNames()) + .setSegmentCount(1) // We don't need to make underlying caches multi-segmented + .setStatsTrackingEnabled(false) + .setMaxSizeInBytes(diskCacheSizeInBytes) + .setStoragePath(builder.cacheConfig.getStoragePath() + "/" + segmentNumber) + .setCacheAlias("tiered_disk_cache#" + segmentNumber) + .build(), + builder.cacheType, + builder.cacheFactories + ); + + Boolean isDiskCacheEnabled = DISK_CACHE_ENABLED_SETTING_MAP.get(builder.cacheType).get(builder.cacheConfig.getSettings()); + LinkedHashMap, TierInfo> cacheListMap = new LinkedHashMap<>(); + cacheListMap.put(onHeapCache, new TierInfo(true, TIER_DIMENSION_VALUE_ON_HEAP)); + cacheListMap.put(diskCache, new TierInfo(isDiskCacheEnabled, TIER_DIMENSION_VALUE_DISK)); + this.caches = Collections.synchronizedMap(cacheListMap); + this.policies = builder.policies; // Will never be null; builder initializes it to an empty list } - return cacheValueTuple.v1(); - } - @Override - public void put(ICacheKey key, V value) { - try (ReleasableLock ignore = writeLock.acquire()) { - onHeapCache.put(key, value); - updateStatsOnPut(TIER_DIMENSION_VALUE_ON_HEAP, key, value); + // Package private for testing + ICache getOnHeapCache() { + return onHeapCache; } - } - @Override - public V computeIfAbsent(ICacheKey key, LoadAwareCacheLoader, V> loader) throws Exception { - // Don't capture stats in the initial getValueFromTieredCache(). If we have concurrent requests for the same key, - // and it only has to be loaded one time, we should report one miss and the rest hits. But, if we do stats in - // getValueFromTieredCache(), - // we will see all misses. Instead, handle stats in computeIfAbsent(). - Tuple cacheValueTuple; - CompletableFuture, V>> future = null; - try (ReleasableLock ignore = readLock.acquire()) { - cacheValueTuple = getValueFromTieredCache(false).apply(key); + // Package private for testing + ICache getDiskCache() { + return diskCache; + } + + void enableDisableDiskCache(Boolean isDiskCacheEnabled) { + // When disk cache is disabled, we are not clearing up the disk cache entries yet as that should be part of + // separate cache/clear API. + this.caches.put(diskCache, new TierInfo(isDiskCacheEnabled, TIER_DIMENSION_VALUE_DISK)); + this.statsHolder.setDiskCacheEnabled(isDiskCacheEnabled); + } + + @Override + public V get(ICacheKey key) { + Tuple cacheValueTuple = getValueFromTieredCache(true).apply(key); if (cacheValueTuple == null) { - // Only one of the threads will succeed putting a future into map for the same key. - // Rest will fetch existing future and wait on that to complete. - future = completableFutureMap.putIfAbsent(key, new CompletableFuture<>()); + return null; } + return cacheValueTuple.v1(); } - List heapDimensionValues = statsHolder.getDimensionsWithTierValue(key.dimensions, TIER_DIMENSION_VALUE_ON_HEAP); - List diskDimensionValues = statsHolder.getDimensionsWithTierValue(key.dimensions, TIER_DIMENSION_VALUE_DISK); - - if (cacheValueTuple == null) { - // Add the value to the onHeap cache. We are calling computeIfAbsent which does another get inside. - // This is needed as there can be many requests for the same key at the same time and we only want to load - // the value once. - V value = compute(key, loader, future); - // Handle stats - if (loader.isLoaded()) { - // The value was just computed and added to the cache by this thread. Register a miss for the heap cache, and the disk cache - // if present + + @Override + public void put(ICacheKey key, V value) { + // First check in case the key is already present in either of tiers. + Tuple cacheValueTuple = getValueFromTieredCache(true).apply(key); + if (cacheValueTuple == null) { + // In case it is not present in any tier, put it inside onHeap cache by default. + try (ReleasableLock ignore = writeLock.acquire()) { + onHeapCache.put(key, value); + } updateStatsOnPut(TIER_DIMENSION_VALUE_ON_HEAP, key, value); - statsHolder.incrementMisses(heapDimensionValues); - if (caches.get(diskCache).isEnabled()) { - statsHolder.incrementMisses(diskDimensionValues); + } else { + // Put it inside desired tier. + try (ReleasableLock ignore = writeLock.acquire()) { + for (Map.Entry, TierInfo> entry : this.caches.entrySet()) { + if (cacheValueTuple.v2().equals(entry.getValue().tierName)) { + entry.getKey().put(key, value); + } + } + updateStatsOnPut(cacheValueTuple.v2(), key, value); + } + } + } + + @Override + public V computeIfAbsent(ICacheKey key, LoadAwareCacheLoader, V> loader) throws Exception { + // Don't capture stats in the initial getValueFromTieredCache(). If we have concurrent requests for the same key, + // and it only has to be loaded one time, we should report one miss and the rest hits. But, if we do stats in + // getValueFromTieredCache(), + // we will see all misses. Instead, handle stats in computeIfAbsent(). + Tuple cacheValueTuple; + CompletableFuture, V>> future = null; + try (ReleasableLock ignore = readLock.acquire()) { + cacheValueTuple = getValueFromTieredCache(false).apply(key); + if (cacheValueTuple == null) { + // Only one of the threads will succeed putting a future into map for the same key. + // Rest will fetch existing future and wait on that to complete. + future = completableFutureMap.putIfAbsent(key, new CompletableFuture<>()); + } + } + List heapDimensionValues = statsHolder.getDimensionsWithTierValue(key.dimensions, TIER_DIMENSION_VALUE_ON_HEAP); + List diskDimensionValues = statsHolder.getDimensionsWithTierValue(key.dimensions, TIER_DIMENSION_VALUE_DISK); + + if (cacheValueTuple == null) { + // Add the value to the onHeap cache. We are calling computeIfAbsent which does another get inside. + // This is needed as there can be many requests for the same key at the same time and we only want to load + // the value once. + V value = compute(key, loader, future); + // Handle stats + if (loader.isLoaded()) { + // The value was just computed and added to the cache by this thread. Register a miss for the heap cache, and the disk + // cache + // if present + updateStatsOnPut(TIER_DIMENSION_VALUE_ON_HEAP, key, value); + statsHolder.incrementMisses(heapDimensionValues); + if (caches.get(diskCache).isEnabled()) { + statsHolder.incrementMisses(diskDimensionValues); + } + } else { + // Another thread requesting this key already loaded the value. Register a hit for the heap cache + statsHolder.incrementHits(heapDimensionValues); } + return value; } else { - // Another thread requesting this key already loaded the value. Register a hit for the heap cache - statsHolder.incrementHits(heapDimensionValues); + // Handle stats for an initial hit from getValueFromTieredCache() + if (cacheValueTuple.v2().equals(TIER_DIMENSION_VALUE_ON_HEAP)) { + // A hit for the heap tier + statsHolder.incrementHits(heapDimensionValues); + } else if (cacheValueTuple.v2().equals(TIER_DIMENSION_VALUE_DISK)) { + // Miss for the heap tier, hit for the disk tier + statsHolder.incrementMisses(heapDimensionValues); + statsHolder.incrementHits(diskDimensionValues); + } + } + return cacheValueTuple.v1(); + } + + private V compute(ICacheKey key, LoadAwareCacheLoader, V> loader, CompletableFuture, V>> future) + throws Exception { + // Handler to handle results post-processing. Takes a tuple or exception as an input and returns + // the value. Also before returning value, puts the value in cache. + BiFunction, V>, Throwable, Void> handler = (pair, ex) -> { + if (pair != null) { + try (ReleasableLock ignore = writeLock.acquire()) { + onHeapCache.put(pair.v1(), pair.v2()); + } catch (Exception e) { + // TODO: Catch specific exceptions to know whether this resulted from cache or underlying removal + // listeners/stats. Needs better exception handling at underlying layers.For now swallowing + // exception. + logger.warn("Exception occurred while putting item onto heap cache", e); + } + } else { + if (ex != null) { + logger.warn("Exception occurred while trying to compute the value", ex); + } + } + completableFutureMap.remove(key);// Remove key from map as not needed anymore. + return null; + }; + V value = null; + if (future == null) { + future = completableFutureMap.get(key); + future.handle(handler); + try { + value = loader.load(key); + } catch (Exception ex) { + future.completeExceptionally(ex); + throw new ExecutionException(ex); + } + if (value == null) { + NullPointerException npe = new NullPointerException("Loader returned a null value"); + future.completeExceptionally(npe); + throw new ExecutionException(npe); + } else { + future.complete(new Tuple<>(key, value)); + } + } else { + try { + value = future.get().v2(); + } catch (InterruptedException ex) { + throw new IllegalStateException(ex); + } } return value; - } else { - // Handle stats for an initial hit from getValueFromTieredCache() - if (cacheValueTuple.v2().equals(TIER_DIMENSION_VALUE_ON_HEAP)) { - // A hit for the heap tier - statsHolder.incrementHits(heapDimensionValues); - } else if (cacheValueTuple.v2().equals(TIER_DIMENSION_VALUE_DISK)) { - // Miss for the heap tier, hit for the disk tier - statsHolder.incrementMisses(heapDimensionValues); - statsHolder.incrementHits(diskDimensionValues); + } + + @Override + public void invalidate(ICacheKey key) { + for (Map.Entry, TierInfo> cacheEntry : caches.entrySet()) { + if (key.getDropStatsForDimensions()) { + List dimensionValues = statsHolder.getDimensionsWithTierValue(key.dimensions, cacheEntry.getValue().tierName); + statsHolder.removeDimensions(dimensionValues); + } + if (key.key != null) { + try (ReleasableLock ignore = writeLock.acquire()) { + cacheEntry.getKey().invalidate(key); + } + } } } - return cacheValueTuple.v1(); - } - private V compute(ICacheKey key, LoadAwareCacheLoader, V> loader, CompletableFuture, V>> future) - throws Exception { - // Handler to handle results post processing. Takes a tuple or exception as an input and returns - // the value. Also before returning value, puts the value in cache. - BiFunction, V>, Throwable, Void> handler = (pair, ex) -> { - if (pair != null) { - try (ReleasableLock ignore = writeLock.acquire()) { - onHeapCache.put(pair.v1(), pair.v2()); - } catch (Exception e) { - // TODO: Catch specific exceptions to know whether this resulted from cache or underlying removal - // listeners/stats. Needs better exception handling at underlying layers.For now swallowing - // exception. - logger.warn("Exception occurred while putting item onto heap cache", e); + @Override + public void invalidateAll() { + try (ReleasableLock ignore = writeLock.acquire()) { + for (Map.Entry, TierInfo> cacheEntry : caches.entrySet()) { + cacheEntry.getKey().invalidateAll(); } - } else { - if (ex != null) { - logger.warn("Exception occurred while trying to compute the value", ex); + } + statsHolder.reset(); + } + + @SuppressWarnings({ "unchecked" }) + @Override + public Iterable> keys() { + List>> iterableList = new ArrayList<>(); + for (Map.Entry, TierInfo> cacheEntry : caches.entrySet()) { + iterableList.add(cacheEntry.getKey().keys()); + } + Iterable>[] iterables = (Iterable>[]) iterableList.toArray(new Iterable[0]); + return new ConcatenatedIterables<>(iterables); + } + + @Override + public long count() { + return onHeapCache.count() + diskCache.count(); + } + + @Override + public void refresh() { + try (ReleasableLock ignore = writeLock.acquire()) { + for (Map.Entry, TierInfo> cacheEntry : caches.entrySet()) { + cacheEntry.getKey().refresh(); } } - completableFutureMap.remove(key);// Remove key from map as not needed anymore. + } + + @Override + public ImmutableCacheStatsHolder stats(String[] levels) { return null; - }; - V value = null; - if (future == null) { - future = completableFutureMap.get(key); - future.handle(handler); - try { - value = loader.load(key); - } catch (Exception ex) { - future.completeExceptionally(ex); - throw new ExecutionException(ex); + } + + @Override + public void close() throws IOException { + for (Map.Entry, TierInfo> cacheEntry : caches.entrySet()) { + // Close all the caches here irrespective of whether they are enabled or not. + cacheEntry.getKey().close(); } - if (value == null) { - NullPointerException npe = new NullPointerException("Loader returned a null value"); - future.completeExceptionally(npe); - throw new ExecutionException(npe); - } else { - future.complete(new Tuple<>(key, value)); + } + + void handleRemovalFromHeapTier(RemovalNotification, V> notification) { + ICacheKey key = notification.getKey(); + boolean wasEvicted = SPILLOVER_REMOVAL_REASONS.contains(notification.getRemovalReason()); + boolean countEvictionTowardsTotal = false; // Don't count this eviction towards the cache's total if it ends up in the disk tier + boolean exceptionOccurredOnDiskCachePut = false; + boolean canCacheOnDisk = caches.get(diskCache).isEnabled() && wasEvicted && evaluatePolicies(notification.getValue()); + if (canCacheOnDisk) { + try (ReleasableLock ignore = writeLock.acquire()) { + diskCache.put(key, notification.getValue()); // spill over to the disk tier and increment its stats + } catch (Exception ex) { + // TODO: Catch specific exceptions. Needs better exception handling. We are just swallowing exception + // in this case as it shouldn't cause upstream request to fail. + logger.warn("Exception occurred while putting item to disk cache", ex); + exceptionOccurredOnDiskCachePut = true; + } + if (!exceptionOccurredOnDiskCachePut) { + updateStatsOnPut(TIER_DIMENSION_VALUE_DISK, key, notification.getValue()); + } } - } else { - try { - value = future.get().v2(); - } catch (InterruptedException ex) { - throw new IllegalStateException(ex); + if (!canCacheOnDisk || exceptionOccurredOnDiskCachePut) { + // If the value is not going to the disk cache, send this notification to the TSC's removal listener + // as the value is leaving the TSC entirely + removalListener.onRemoval(notification); + countEvictionTowardsTotal = true; } + updateStatsOnRemoval(TIER_DIMENSION_VALUE_ON_HEAP, wasEvicted, key, notification.getValue(), countEvictionTowardsTotal); } - return value; - } - @Override - public void invalidate(ICacheKey key) { - // We are trying to invalidate the key from all caches though it would be present in only of them. - // Doing this as we don't know where it is located. We could do a get from both and check that, but what will - // also trigger a hit/miss listener event, so ignoring it for now. - // We don't update stats here, as this is handled by the removal listeners for the tiers. - for (Map.Entry, TierInfo> cacheEntry : caches.entrySet()) { - if (key.getDropStatsForDimensions()) { - List dimensionValues = statsHolder.getDimensionsWithTierValue(key.dimensions, cacheEntry.getValue().tierName); - statsHolder.removeDimensions(dimensionValues); + boolean evaluatePolicies(V value) { + for (Predicate policy : policies) { + if (!policy.test(value)) { + return false; + } } - if (key.key != null) { - try (ReleasableLock ignore = writeLock.acquire()) { - cacheEntry.getKey().invalidate(key); + return true; + } + + /** + * Get a value from the tiered cache, and the name of the tier it was found in. + * @param captureStats Whether to record hits/misses for this call of the function + * @return A tuple of the value and the name of the tier it was found in. + */ + private Function, Tuple> getValueFromTieredCache(boolean captureStats) { + return key -> { + try (ReleasableLock ignore = readLock.acquire()) { + for (Map.Entry, TierInfo> cacheEntry : caches.entrySet()) { + if (cacheEntry.getValue().isEnabled()) { + V value = cacheEntry.getKey().get(key); + // Get the tier value corresponding to this cache + String tierValue = cacheEntry.getValue().tierName; + List dimensionValues = statsHolder.getDimensionsWithTierValue(key.dimensions, tierValue); + if (value != null) { + if (captureStats) { + statsHolder.incrementHits(dimensionValues); + } + return new Tuple<>(value, tierValue); + } else if (captureStats) { + statsHolder.incrementMisses(dimensionValues); + } + } + } + return null; } + }; + } + + void handleRemovalFromDiskTier(RemovalNotification, V> notification) { + // Values removed from the disk tier leave the TSC entirely + removalListener.onRemoval(notification); + boolean wasEvicted = SPILLOVER_REMOVAL_REASONS.contains(notification.getRemovalReason()); + updateStatsOnRemoval(TIER_DIMENSION_VALUE_DISK, wasEvicted, notification.getKey(), notification.getValue(), true); + } + + void updateStatsOnRemoval( + String removedFromTierValue, + boolean wasEvicted, + ICacheKey key, + V value, + boolean countEvictionTowardsTotal + ) { + List dimensionValues = statsHolder.getDimensionsWithTierValue(key.dimensions, removedFromTierValue); + if (wasEvicted) { + statsHolder.incrementEvictions(dimensionValues, countEvictionTowardsTotal); + } + statsHolder.decrementItems(dimensionValues); + statsHolder.decrementSizeInBytes(dimensionValues, weigher.applyAsLong(key, value)); + } + + void updateStatsOnPut(String destinationTierValue, ICacheKey key, V value) { + List dimensionValues = statsHolder.getDimensionsWithTierValue(key.dimensions, destinationTierValue); + statsHolder.incrementItems(dimensionValues); + statsHolder.incrementSizeInBytes(dimensionValues, weigher.applyAsLong(key, value)); + } + + /** + * A class which receives removal events from the heap tier. + */ + private class HeapTierRemovalListener implements RemovalListener, V> { + private final TieredSpilloverCacheSegment tsc; + + HeapTierRemovalListener(TieredSpilloverCacheSegment tsc) { + this.tsc = tsc; + } + + @Override + public void onRemoval(RemovalNotification, V> notification) { + tsc.handleRemovalFromHeapTier(notification); + } + } + + /** + * A class which receives removal events from the disk tier. + */ + private class DiskTierRemovalListener implements RemovalListener, V> { + private final TieredSpilloverCacheSegment tsc; + + DiskTierRemovalListener(TieredSpilloverCacheSegment tsc) { + this.tsc = tsc; + } + + @Override + public void onRemoval(RemovalNotification, V> notification) { + tsc.handleRemovalFromDiskTier(notification); } } } + // Package private for testing. + void enableDisableDiskCache(Boolean isDiskCacheEnabled) { + for (int iter = 0; iter < this.numberOfSegments; iter++) { + tieredSpilloverCacheSegments[iter].enableDisableDiskCache(isDiskCacheEnabled); + } + this.statsHolder.setDiskCacheEnabled(isDiskCacheEnabled); + } + + // Package private for testing. + TieredSpilloverCacheSegment getTieredCacheSegment(ICacheKey key) { + return tieredSpilloverCacheSegments[getSegmentNumber(key)]; + } + + int getSegmentNumber(ICacheKey key) { + return key.hashCode() & (this.numberOfSegments - 1); + } + + int getNumberOfSegments() { + return tieredSpilloverCacheSegments.length; + } + + @Override + public V get(ICacheKey key) { + TieredSpilloverCacheSegment tieredSpilloverCacheSegment = getTieredCacheSegment(key); + return tieredSpilloverCacheSegment.get(key); + } + + @Override + public void put(ICacheKey key, V value) { + TieredSpilloverCacheSegment tieredSpilloverCacheSegment = getTieredCacheSegment(key); + tieredSpilloverCacheSegment.put(key, value); + } + + @Override + public V computeIfAbsent(ICacheKey key, LoadAwareCacheLoader, V> loader) throws Exception { + TieredSpilloverCacheSegment tieredSpilloverCacheSegment = getTieredCacheSegment(key); + return tieredSpilloverCacheSegment.computeIfAbsent(key, loader); + } + + @Override + public void invalidate(ICacheKey key) { + TieredSpilloverCacheSegment tieredSpilloverCacheSegment = getTieredCacheSegment(key); + tieredSpilloverCacheSegment.invalidate(key); + } + @Override public void invalidateAll() { - try (ReleasableLock ignore = writeLock.acquire()) { - for (Map.Entry, TierInfo> cacheEntry : caches.entrySet()) { - cacheEntry.getKey().invalidateAll(); - } + for (int iter = 0; iter < this.numberOfSegments; iter++) { + tieredSpilloverCacheSegments[iter].invalidateAll(); } - statsHolder.reset(); } /** @@ -327,8 +621,8 @@ public void invalidateAll() { @Override public Iterable> keys() { List>> iterableList = new ArrayList<>(); - for (Map.Entry, TierInfo> cacheEntry : caches.entrySet()) { - iterableList.add(cacheEntry.getKey().keys()); + for (int iter = 0; iter < this.numberOfSegments; iter++) { + iterableList.add(tieredSpilloverCacheSegments[iter].keys()); } Iterable>[] iterables = (Iterable>[]) iterableList.toArray(new Iterable[0]); return new ConcatenatedIterables<>(iterables); @@ -343,18 +637,15 @@ public long count() { @Override public void refresh() { - try (ReleasableLock ignore = writeLock.acquire()) { - for (Map.Entry, TierInfo> cacheEntry : caches.entrySet()) { - cacheEntry.getKey().refresh(); - } + for (int iter = 0; iter < this.numberOfSegments; iter++) { + tieredSpilloverCacheSegments[iter].refresh(); } } @Override public void close() throws IOException { - for (Map.Entry, TierInfo> cacheEntry : caches.entrySet()) { - // Close all the caches here irrespective of whether they are enabled or not. - cacheEntry.getKey().close(); + for (int iter = 0; iter < this.numberOfSegments; iter++) { + tieredSpilloverCacheSegments[iter].close(); } } @@ -363,130 +654,44 @@ public ImmutableCacheStatsHolder stats(String[] levels) { return statsHolder.getImmutableCacheStatsHolder(levels); } - /** - * Get a value from the tiered cache, and the name of the tier it was found in. - * @param captureStats Whether to record hits/misses for this call of the function - * @return A tuple of the value and the name of the tier it was found in. - */ - private Function, Tuple> getValueFromTieredCache(boolean captureStats) { - return key -> { - try (ReleasableLock ignore = readLock.acquire()) { - for (Map.Entry, TierInfo> cacheEntry : caches.entrySet()) { - if (cacheEntry.getValue().isEnabled()) { - V value = cacheEntry.getKey().get(key); - // Get the tier value corresponding to this cache - String tierValue = cacheEntry.getValue().tierName; - List dimensionValues = statsHolder.getDimensionsWithTierValue(key.dimensions, tierValue); - if (value != null) { - if (captureStats) { - statsHolder.incrementHits(dimensionValues); - } - return new Tuple<>(value, tierValue); - } else if (captureStats) { - statsHolder.incrementMisses(dimensionValues); - } - } - } - return null; - } - }; - } - - void handleRemovalFromHeapTier(RemovalNotification, V> notification) { - ICacheKey key = notification.getKey(); - boolean wasEvicted = SPILLOVER_REMOVAL_REASONS.contains(notification.getRemovalReason()); - boolean countEvictionTowardsTotal = false; // Don't count this eviction towards the cache's total if it ends up in the disk tier - boolean exceptionOccurredOnDiskCachePut = false; - boolean canCacheOnDisk = caches.get(diskCache).isEnabled() && wasEvicted && evaluatePolicies(notification.getValue()); - if (canCacheOnDisk) { - try (ReleasableLock ignore = writeLock.acquire()) { - diskCache.put(key, notification.getValue()); // spill over to the disk tier and increment its stats - } catch (Exception ex) { - // TODO: Catch specific exceptions. Needs better exception handling. We are just swallowing exception - // in this case as it shouldn't cause upstream request to fail. - logger.warn("Exception occurred while putting item to disk cache", ex); - exceptionOccurredOnDiskCachePut = true; - } - if (!exceptionOccurredOnDiskCachePut) { - updateStatsOnPut(TIER_DIMENSION_VALUE_DISK, key, notification.getValue()); - } - } - if (!canCacheOnDisk || exceptionOccurredOnDiskCachePut) { - // If the value is not going to the disk cache, send this notification to the TSC's removal listener - // as the value is leaving the TSC entirely - removalListener.onRemoval(notification); - countEvictionTowardsTotal = true; + // Package private for testing. + @SuppressWarnings({ "unchecked" }) + Iterable> getOnHeapCacheKeys() { + List>> iterableList = new ArrayList<>(); + for (int iter = 0; iter < this.numberOfSegments; iter++) { + iterableList.add(tieredSpilloverCacheSegments[iter].onHeapCache.keys()); } - updateStatsOnRemoval(TIER_DIMENSION_VALUE_ON_HEAP, wasEvicted, key, notification.getValue(), countEvictionTowardsTotal); - } - - void handleRemovalFromDiskTier(RemovalNotification, V> notification) { - // Values removed from the disk tier leave the TSC entirely - removalListener.onRemoval(notification); - boolean wasEvicted = SPILLOVER_REMOVAL_REASONS.contains(notification.getRemovalReason()); - updateStatsOnRemoval(TIER_DIMENSION_VALUE_DISK, wasEvicted, notification.getKey(), notification.getValue(), true); - } - - void updateStatsOnRemoval( - String removedFromTierValue, - boolean wasEvicted, - ICacheKey key, - V value, - boolean countEvictionTowardsTotal - ) { - List dimensionValues = statsHolder.getDimensionsWithTierValue(key.dimensions, removedFromTierValue); - if (wasEvicted) { - statsHolder.incrementEvictions(dimensionValues, countEvictionTowardsTotal); - } - statsHolder.decrementItems(dimensionValues); - statsHolder.decrementSizeInBytes(dimensionValues, weigher.applyAsLong(key, value)); - } - - void updateStatsOnPut(String destinationTierValue, ICacheKey key, V value) { - List dimensionValues = statsHolder.getDimensionsWithTierValue(key.dimensions, destinationTierValue); - statsHolder.incrementItems(dimensionValues); - statsHolder.incrementSizeInBytes(dimensionValues, weigher.applyAsLong(key, value)); + Iterable>[] iterables = (Iterable>[]) iterableList.toArray(new Iterable[0]); + return new ConcatenatedIterables<>(iterables); } - boolean evaluatePolicies(V value) { - for (Predicate policy : policies) { - if (!policy.test(value)) { - return false; - } + // Package private for testing. + @SuppressWarnings({ "unchecked" }) + Iterable> getDiskCacheKeys() { + List>> iterableList = new ArrayList<>(); + for (int iter = 0; iter < this.numberOfSegments; iter++) { + iterableList.add(tieredSpilloverCacheSegments[iter].diskCache.keys()); } - return true; + Iterable>[] iterables = (Iterable>[]) iterableList.toArray(new Iterable[0]); + return new ConcatenatedIterables<>(iterables); } - /** - * A class which receives removal events from the heap tier. - */ - private class HeapTierRemovalListener implements RemovalListener, V> { - private final TieredSpilloverCache tsc; - - HeapTierRemovalListener(TieredSpilloverCache tsc) { - this.tsc = tsc; - } - - @Override - public void onRemoval(RemovalNotification, V> notification) { - tsc.handleRemovalFromHeapTier(notification); + // Package private for testing. + long onHeapCacheCount() { + long onHeapCacheEntries = 0; + for (int iter = 0; iter < this.numberOfSegments; iter++) { + onHeapCacheEntries += tieredSpilloverCacheSegments[iter].onHeapCache.count(); } + return onHeapCacheEntries; } - /** - * A class which receives removal events from the disk tier. - */ - private class DiskTierRemovalListener implements RemovalListener, V> { - private final TieredSpilloverCache tsc; - - DiskTierRemovalListener(TieredSpilloverCache tsc) { - this.tsc = tsc; - } - - @Override - public void onRemoval(RemovalNotification, V> notification) { - tsc.handleRemovalFromDiskTier(notification); + // Package private for testing. + long diskCacheCount() { + long diskCacheEntries = 0; + for (int iter = 0; iter < this.numberOfSegments; iter++) { + diskCacheEntries += tieredSpilloverCacheSegments[iter].diskCache.count(); } + return diskCacheEntries; } /** @@ -550,7 +755,7 @@ public void remove() { } } - private class TierInfo { + private static class TierInfo { AtomicBoolean isEnabled; final String tierName; @@ -611,12 +816,29 @@ public ICache create(CacheConfig config, CacheType cacheType, "Cached result parser fn can't be null" ); + int numberOfSegments = TIERED_SPILLOVER_SEGMENTS.getConcreteSettingForNamespace(cacheType.getSettingPrefix()).get(settings); + + if (!VALID_SEGMENT_COUNT_VALUES.contains(numberOfSegments)) { + throw new IllegalArgumentException( + String.format(Locale.ROOT, INVALID_SEGMENT_COUNT_EXCEPTION_MESSAGE, TIERED_SPILLOVER_CACHE_NAME) + ); + } + + long onHeapCacheSize = TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace(cacheType.getSettingPrefix()) + .get(settings) + .getBytes(); + long diskCacheSize = TIERED_SPILLOVER_DISK_STORE_SIZE.getConcreteSettingForNamespace(cacheType.getSettingPrefix()) + .get(settings); + return new Builder().setDiskCacheFactory(diskCacheFactory) .setOnHeapCacheFactory(onHeapCacheFactory) .setRemovalListener(config.getRemovalListener()) .setCacheConfig(config) .setCacheType(cacheType) + .setNumberOfSegments(numberOfSegments) .addPolicy(new TookTimePolicy(diskPolicyThreshold, cachedResultParser, config.getClusterSettings(), cacheType)) + .setOnHeapCacheSizeInBytes(onHeapCacheSize) + .setDiskCacheSize(diskCacheSize) .build(); } @@ -640,6 +862,10 @@ public static class Builder { private Map cacheFactories; private final ArrayList> policies = new ArrayList<>(); + private int numberOfSegments; + private long onHeapCacheSizeInBytes; + private long diskCacheSizeInBytes; + /** * Default constructor */ @@ -725,6 +951,36 @@ public Builder addPolicies(List> policies) { return this; } + /** + * Sets number of segments for tiered cache + * @param numberOfSegments number of segments + * @return builder + */ + public Builder setNumberOfSegments(int numberOfSegments) { + this.numberOfSegments = numberOfSegments; + return this; + } + + /** + * Sets onHeap cache size + * @param onHeapCacheSizeInBytes size of onHeap cache in bytes + * @return builder + */ + public Builder setOnHeapCacheSizeInBytes(long onHeapCacheSizeInBytes) { + this.onHeapCacheSizeInBytes = onHeapCacheSizeInBytes; + return this; + } + + /** + * Sets disk cache siz + * @param diskCacheSizeInBytes size of diskCache in bytes + * @return buider + */ + public Builder setDiskCacheSize(long diskCacheSizeInBytes) { + this.diskCacheSizeInBytes = diskCacheSizeInBytes; + return this; + } + /** * Build tiered spillover cache. * @return TieredSpilloverCache diff --git a/modules/cache-common/src/main/java/org/opensearch/cache/common/tier/TieredSpilloverCachePlugin.java b/modules/cache-common/src/main/java/org/opensearch/cache/common/tier/TieredSpilloverCachePlugin.java index 1c10e51630460..bf522b42b70ca 100644 --- a/modules/cache-common/src/main/java/org/opensearch/cache/common/tier/TieredSpilloverCachePlugin.java +++ b/modules/cache-common/src/main/java/org/opensearch/cache/common/tier/TieredSpilloverCachePlugin.java @@ -65,6 +65,15 @@ public List> getSettings() { if (FeatureFlags.PLUGGABLE_CACHE_SETTING.get(settings)) { settingList.add(DISK_CACHE_ENABLED_SETTING_MAP.get(cacheType)); } + settingList.add( + TieredSpilloverCacheSettings.TIERED_SPILLOVER_SEGMENTS.getConcreteSettingForNamespace(cacheType.getSettingPrefix()) + ); + settingList.add( + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace(cacheType.getSettingPrefix()) + ); + settingList.add( + TieredSpilloverCacheSettings.TIERED_SPILLOVER_DISK_STORE_SIZE.getConcreteSettingForNamespace(cacheType.getSettingPrefix()) + ); } return settingList; } diff --git a/modules/cache-common/src/main/java/org/opensearch/cache/common/tier/TieredSpilloverCacheSettings.java b/modules/cache-common/src/main/java/org/opensearch/cache/common/tier/TieredSpilloverCacheSettings.java index e8e441d6bd3a6..122d00af3bd1e 100644 --- a/modules/cache-common/src/main/java/org/opensearch/cache/common/tier/TieredSpilloverCacheSettings.java +++ b/modules/cache-common/src/main/java/org/opensearch/cache/common/tier/TieredSpilloverCacheSettings.java @@ -11,11 +11,16 @@ import org.opensearch.common.cache.CacheType; import org.opensearch.common.settings.Setting; import org.opensearch.common.unit.TimeValue; +import org.opensearch.core.common.unit.ByteSizeValue; +import org.opensearch.threadpool.ThreadPool; import java.util.HashMap; +import java.util.Locale; import java.util.Map; import java.util.concurrent.TimeUnit; +import static org.opensearch.common.cache.settings.CacheSettings.INVALID_SEGMENT_COUNT_EXCEPTION_MESSAGE; +import static org.opensearch.common.cache.settings.CacheSettings.VALID_SEGMENT_COUNT_VALUES; import static org.opensearch.common.settings.Setting.Property.NodeScope; /** @@ -23,6 +28,16 @@ */ public class TieredSpilloverCacheSettings { + /** + * Default cache size in bytes ie 1gb. + */ + public static final long DEFAULT_DISK_CACHE_SIZE_IN_BYTES = 1073741824L; + + /** + * Minimum disk cache size ie 10mb. May not make such sense to keep a value smaller than this. + */ + public static final long MIN_DISK_CACHE_SIZE_IN_BYTES = 10485760L; + /** * Setting which defines the onHeap cache store to be used in TieredSpilloverCache. * @@ -50,6 +65,43 @@ public class TieredSpilloverCacheSettings { (key) -> Setting.boolSetting(key, true, NodeScope, Setting.Property.Dynamic) ); + /** + * Setting defining the number of segments within Tiered cache + */ + public static final Setting.AffixSetting TIERED_SPILLOVER_SEGMENTS = Setting.suffixKeySetting( + TieredSpilloverCache.TieredSpilloverCacheFactory.TIERED_SPILLOVER_CACHE_NAME + ".segments", + (key) -> Setting.intSetting(key, defaultSegments(), 1, k -> { + if (!VALID_SEGMENT_COUNT_VALUES.contains(k)) { + throw new IllegalArgumentException( + String.format( + Locale.ROOT, + INVALID_SEGMENT_COUNT_EXCEPTION_MESSAGE, + TieredSpilloverCache.TieredSpilloverCacheFactory.TIERED_SPILLOVER_CACHE_NAME + ) + ); + } + }, NodeScope) + ); + + /** + * Setting which defines the onHeap cache size to be used within tiered cache. + * + * Pattern: {cache_type}.tiered_spillover.onheap.store.size + * Example: indices.request.cache.tiered_spillover.onheap.store.size + */ + public static final Setting.AffixSetting TIERED_SPILLOVER_ONHEAP_STORE_SIZE = Setting.suffixKeySetting( + TieredSpilloverCache.TieredSpilloverCacheFactory.TIERED_SPILLOVER_CACHE_NAME + ".onheap.store.size", + (key) -> Setting.memorySizeSetting(key, "1%", NodeScope) + ); + + /** + * Setting which defines the disk cache size to be used within tiered cache. + */ + public static final Setting.AffixSetting TIERED_SPILLOVER_DISK_STORE_SIZE = Setting.suffixKeySetting( + TieredSpilloverCache.TieredSpilloverCacheFactory.TIERED_SPILLOVER_CACHE_NAME + ".disk.store.size", + (key) -> Setting.longSetting(key, DEFAULT_DISK_CACHE_SIZE_IN_BYTES, MIN_DISK_CACHE_SIZE_IN_BYTES, NodeScope) + ); + /** * Setting defining the minimum took time for a query to be allowed into the disk cache. */ @@ -96,6 +148,23 @@ public class TieredSpilloverCacheSettings { DISK_CACHE_ENABLED_SETTING_MAP = diskCacheSettingMap; } + /** + * Returns the default segment count to be used within TieredCache. + * @return default segment count + */ + public static int defaultSegments() { + // For now, we use number of search threads as the default segment count. If needed each cache type can + // configure its own segmentCount via setting in the future. + int defaultSegmentCount = ThreadPool.searchThreadPoolSize(Runtime.getRuntime().availableProcessors()); + // Now round it off to the next power of 2 as we don't support any other values. + for (int segmentValue : VALID_SEGMENT_COUNT_VALUES) { + if (defaultSegmentCount <= segmentValue) { + return segmentValue; + } + } + return VALID_SEGMENT_COUNT_VALUES.get(VALID_SEGMENT_COUNT_VALUES.size() - 1); + } + /** * Default constructor */ diff --git a/modules/cache-common/src/test/java/org/opensearch/cache/common/tier/MockDiskCache.java b/modules/cache-common/src/test/java/org/opensearch/cache/common/tier/MockDiskCache.java index 69e2060f7ea2f..fcddd489a27aa 100644 --- a/modules/cache-common/src/test/java/org/opensearch/cache/common/tier/MockDiskCache.java +++ b/modules/cache-common/src/test/java/org/opensearch/cache/common/tier/MockDiskCache.java @@ -62,6 +62,7 @@ public void put(ICacheKey key, V value) { if (this.cache.size() >= maxSize) { // For simplification this.removalListener.onRemoval(new RemovalNotification<>(key, value, RemovalReason.EVICTED)); this.statsHolder.decrementItems(List.of()); + return; } try { Thread.sleep(delay); @@ -86,8 +87,10 @@ public V computeIfAbsent(ICacheKey key, LoadAwareCacheLoader, V> @Override public void invalidate(ICacheKey key) { - removalListener.onRemoval(new RemovalNotification<>(key, cache.get(key), RemovalReason.INVALIDATED)); - this.cache.remove(key); + V value = this.cache.remove(key); + if (value != null) { + removalListener.onRemoval(new RemovalNotification<>(key, cache.get(key), RemovalReason.INVALIDATED)); + } } @Override @@ -131,11 +134,13 @@ public static class MockDiskCacheFactory implements Factory { final long delay; final int maxSize; final boolean statsTrackingEnabled; + final int keyValueSize; - public MockDiskCacheFactory(long delay, int maxSize, boolean statsTrackingEnabled) { + public MockDiskCacheFactory(long delay, int maxSize, boolean statsTrackingEnabled, int keyValueSize) { this.delay = delay; this.maxSize = maxSize; this.statsTrackingEnabled = statsTrackingEnabled; + this.keyValueSize = keyValueSize; } @Override @@ -145,13 +150,21 @@ public ICache create(CacheConfig config, CacheType cacheType, // cache would require. assert config.getKeySerializer() != null; assert config.getValueSerializer() != null; - return new Builder().setKeySerializer((Serializer) config.getKeySerializer()) + MockDiskCache.Builder builder = (Builder) new Builder().setKeySerializer( + (Serializer) config.getKeySerializer() + ) .setValueSerializer((Serializer) config.getValueSerializer()) - .setMaxSize(maxSize) .setDeliberateDelay(delay) .setRemovalListener(config.getRemovalListener()) - .setStatsTrackingEnabled(config.getStatsTrackingEnabled()) - .build(); + .setStatsTrackingEnabled(config.getStatsTrackingEnabled()); + + // For mock disk cache, size refers to number of entries for simplicity. + if (config.getMaxSizeInBytes() > 0) { + builder.setMaxSize(Math.toIntExact(config.getMaxSizeInBytes())); + } else { + builder.setMaxSize(maxSize); + } + return builder.build(); } @Override diff --git a/modules/cache-common/src/test/java/org/opensearch/cache/common/tier/TieredSpilloverCacheTests.java b/modules/cache-common/src/test/java/org/opensearch/cache/common/tier/TieredSpilloverCacheTests.java index c6440a1e1797f..1215a2130ac2d 100644 --- a/modules/cache-common/src/test/java/org/opensearch/cache/common/tier/TieredSpilloverCacheTests.java +++ b/modules/cache-common/src/test/java/org/opensearch/cache/common/tier/TieredSpilloverCacheTests.java @@ -8,6 +8,7 @@ package org.opensearch.cache.common.tier; +import org.opensearch.OpenSearchException; import org.opensearch.common.Randomness; import org.opensearch.common.cache.CacheType; import org.opensearch.common.cache.ICache; @@ -29,6 +30,7 @@ import org.opensearch.common.settings.Settings; import org.opensearch.common.unit.TimeValue; import org.opensearch.common.util.FeatureFlags; +import org.opensearch.env.NodeEnvironment; import org.opensearch.test.OpenSearchTestCase; import org.junit.Before; @@ -54,11 +56,16 @@ import java.util.function.Function; import java.util.function.Predicate; +import static org.opensearch.cache.common.tier.TieredSpilloverCache.ZERO_SEGMENT_COUNT_EXCEPTION_MESSAGE; import static org.opensearch.cache.common.tier.TieredSpilloverCacheSettings.DISK_CACHE_ENABLED_SETTING_MAP; +import static org.opensearch.cache.common.tier.TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE; +import static org.opensearch.cache.common.tier.TieredSpilloverCacheSettings.TIERED_SPILLOVER_SEGMENTS; import static org.opensearch.cache.common.tier.TieredSpilloverCacheSettings.TOOK_TIME_POLICY_CONCRETE_SETTINGS_MAP; import static org.opensearch.cache.common.tier.TieredSpilloverCacheStatsHolder.TIER_DIMENSION_NAME; import static org.opensearch.cache.common.tier.TieredSpilloverCacheStatsHolder.TIER_DIMENSION_VALUE_DISK; import static org.opensearch.cache.common.tier.TieredSpilloverCacheStatsHolder.TIER_DIMENSION_VALUE_ON_HEAP; +import static org.opensearch.common.cache.settings.CacheSettings.INVALID_SEGMENT_COUNT_EXCEPTION_MESSAGE; +import static org.opensearch.common.cache.settings.CacheSettings.VALID_SEGMENT_COUNT_VALUES; import static org.opensearch.common.cache.store.settings.OpenSearchOnHeapCacheSettings.MAXIMUM_SIZE_IN_BYTES_KEY; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.doThrow; @@ -89,13 +96,14 @@ public void testComputeIfAbsentWithoutAnyOnHeapCacheEviction() throws Exception removalListener, Settings.builder() .put( - OpenSearchOnHeapCacheSettings.getSettingListForCacheType(CacheType.INDICES_REQUEST_CACHE) - .get(MAXIMUM_SIZE_IN_BYTES_KEY) - .getKey(), + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), onHeapCacheSize * keyValueSize + "b" ) .build(), - 0 + 0, + 1 ); int numOfItems1 = randomIntBetween(1, onHeapCacheSize / 2 - 1); List> keys = new ArrayList<>(); @@ -160,18 +168,19 @@ public void testComputeIfAbsentWithFactoryBasedCacheCreation() throws Exception MockDiskCache.MockDiskCacheFactory.NAME ) .put( - OpenSearchOnHeapCacheSettings.getSettingListForCacheType(CacheType.INDICES_REQUEST_CACHE) - .get(MAXIMUM_SIZE_IN_BYTES_KEY) - .getKey(), + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), onHeapCacheSize * keyValueSize + "b" ) .put( CacheSettings.getConcreteStoreNameSettingForCacheType(CacheType.INDICES_REQUEST_CACHE).getKey(), TieredSpilloverCache.TieredSpilloverCacheFactory.TIERED_SPILLOVER_CACHE_NAME ) + .put(TIERED_SPILLOVER_SEGMENTS.getConcreteSettingForNamespace(CacheType.INDICES_REQUEST_CACHE.getSettingPrefix()).getKey(), 1) .put(FeatureFlags.PLUGGABLE_CACHE, "true") .build(); - + String storagePath = getStoragePath(settings); ICache tieredSpilloverICache = new TieredSpilloverCache.TieredSpilloverCacheFactory().create( new CacheConfig.Builder().setKeyType(String.class) .setKeyType(String.class) @@ -182,15 +191,16 @@ public void testComputeIfAbsentWithFactoryBasedCacheCreation() throws Exception .setSettings(settings) .setDimensionNames(dimensionNames) .setCachedResultParser(s -> new CachedQueryResult.PolicyValues(20_000_000L)) // Values will always appear to have taken - // 20_000_000 ns = 20 ms to compute + // 20_000_000 ns = 20 ms to compute .setClusterSettings(clusterSettings) + .setStoragePath(storagePath) .build(), CacheType.INDICES_REQUEST_CACHE, Map.of( OpenSearchOnHeapCache.OpenSearchOnHeapCacheFactory.NAME, new OpenSearchOnHeapCache.OpenSearchOnHeapCacheFactory(), MockDiskCache.MockDiskCacheFactory.NAME, - new MockDiskCache.MockDiskCacheFactory(0, randomIntBetween(100, 300), false) + new MockDiskCache.MockDiskCacheFactory(0, randomIntBetween(100, 300), false, keyValueSize) ) ); @@ -206,10 +216,10 @@ public void testComputeIfAbsentWithFactoryBasedCacheCreation() throws Exception } int expectedDiskEntries = numOfItems1 - onHeapCacheSize; - tieredSpilloverCache.getOnHeapCache().keys().forEach(onHeapKeys::add); - tieredSpilloverCache.getDiskCache().keys().forEach(diskTierKeys::add); + tieredSpilloverCache.tieredSpilloverCacheSegments[0].getOnHeapCache().keys().forEach(onHeapKeys::add); + tieredSpilloverCache.tieredSpilloverCacheSegments[0].getDiskCache().keys().forEach(diskTierKeys::add); // Verify on heap cache stats. - assertEquals(onHeapCacheSize, tieredSpilloverCache.getOnHeapCache().count()); + assertEquals(onHeapCacheSize, tieredSpilloverCache.tieredSpilloverCacheSegments[0].getOnHeapCache().count()); assertEquals(onHeapCacheSize, getItemsForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_ON_HEAP)); assertEquals(numOfItems1, getMissesForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_ON_HEAP)); assertEquals(0, getHitsForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_ON_HEAP)); @@ -217,7 +227,7 @@ public void testComputeIfAbsentWithFactoryBasedCacheCreation() throws Exception assertEquals(onHeapCacheSize * keyValueSize, getSizeInBytesForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_ON_HEAP)); // Verify disk cache stats. - assertEquals(expectedDiskEntries, tieredSpilloverCache.getDiskCache().count()); + assertEquals(expectedDiskEntries, tieredSpilloverCache.tieredSpilloverCacheSegments[0].getDiskCache().count()); assertEquals(expectedDiskEntries, getItemsForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_DISK)); assertEquals(0, getHitsForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_DISK)); assertEquals(numOfItems1, getMissesForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_DISK)); @@ -225,6 +235,149 @@ public void testComputeIfAbsentWithFactoryBasedCacheCreation() throws Exception assertEquals(expectedDiskEntries * keyValueSize, getSizeInBytesForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_DISK)); } + public void testComputeIfAbsentWithSegmentedCache() throws Exception { + int onHeapCacheSize = randomIntBetween(300, 600); + int diskCacheSize = randomIntBetween(700, 1200); + int totalSize = onHeapCacheSize + diskCacheSize; + int numberOfSegments = getNumberOfSegments(); + int keyValueSize = 11; + MockCacheRemovalListener removalListener = new MockCacheRemovalListener<>(); + int onHeapCacheSizeInBytes = onHeapCacheSize * keyValueSize; + Map expectedSegmentOnHeapCacheSize = getSegmentOnHeapCacheSize( + numberOfSegments, + onHeapCacheSizeInBytes, + keyValueSize + ); + int totalOnHeapEntries = 0; + int totalOnDiskEntries = 0; + // Set the desired settings needed to create a TieredSpilloverCache object with INDICES_REQUEST_CACHE cacheType. + Settings settings = Settings.builder() + .put( + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_NAME.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), + OpenSearchOnHeapCache.OpenSearchOnHeapCacheFactory.NAME + ) + .put( + TIERED_SPILLOVER_SEGMENTS.getConcreteSettingForNamespace(CacheType.INDICES_REQUEST_CACHE.getSettingPrefix()).getKey(), + numberOfSegments + ) + .put( + TieredSpilloverCacheSettings.TIERED_SPILLOVER_DISK_STORE_NAME.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), + MockDiskCache.MockDiskCacheFactory.NAME + ) + .put( + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), + onHeapCacheSizeInBytes + "b" + ) + .put( + CacheSettings.getConcreteStoreNameSettingForCacheType(CacheType.INDICES_REQUEST_CACHE).getKey(), + TieredSpilloverCache.TieredSpilloverCacheFactory.TIERED_SPILLOVER_CACHE_NAME + ) + .put(FeatureFlags.PLUGGABLE_CACHE, "true") + .build(); + String storagePath = getStoragePath(settings); + ICache tieredSpilloverICache = new TieredSpilloverCache.TieredSpilloverCacheFactory().create( + new CacheConfig.Builder().setKeyType(String.class) + .setKeyType(String.class) + .setWeigher((k, v) -> keyValueSize) + .setRemovalListener(removalListener) + .setKeySerializer(new StringSerializer()) + .setValueSerializer(new StringSerializer()) + .setSettings(settings) + .setDimensionNames(dimensionNames) + .setCachedResultParser(s -> new CachedQueryResult.PolicyValues(20_000_000L)) // Values will always appear to have taken + // 20_000_000 ns = 20 ms to compute + .setClusterSettings(clusterSettings) + .setStoragePath(storagePath) + .setSegmentCount(numberOfSegments) + .build(), + CacheType.INDICES_REQUEST_CACHE, + Map.of( + OpenSearchOnHeapCache.OpenSearchOnHeapCacheFactory.NAME, + new OpenSearchOnHeapCache.OpenSearchOnHeapCacheFactory(), + MockDiskCache.MockDiskCacheFactory.NAME, + // Make disk cache big enough to hold all entries + new MockDiskCache.MockDiskCacheFactory(0, diskCacheSize * 500, false, keyValueSize) + ) + ); + TieredSpilloverCache tieredSpilloverCache = (TieredSpilloverCache) tieredSpilloverICache; + TieredSpilloverCache.TieredSpilloverCacheSegment[] tieredSpilloverCacheSegments = + tieredSpilloverCache.tieredSpilloverCacheSegments; + assertEquals(numberOfSegments, tieredSpilloverCacheSegments.length); + + Map, TieredSpilloverCache.TieredSpilloverCacheSegment> tieredSpilloverCacheSegmentMap = + new HashMap<>(); + + int numOfItems1 = randomIntBetween(onHeapCacheSize + 1, totalSize); + List> onHeapKeys = new ArrayList<>(); + List> diskTierKeys = new ArrayList<>(); + Map expectedNumberOfEntriesInSegment = new HashMap<>(); + for (int iter = 0; iter < numOfItems1; iter++) { + String key = UUID.randomUUID().toString(); + LoadAwareCacheLoader, String> tieredCacheLoader = getLoadAwareCacheLoader(); + ICacheKey iCacheKey = getICacheKey(key); + int keySegment = tieredSpilloverCache.getSegmentNumber(iCacheKey); + if (expectedNumberOfEntriesInSegment.get(keySegment) == null) { + expectedNumberOfEntriesInSegment.put(keySegment, Integer.valueOf(1)); + } else { + Integer updatedValue = expectedNumberOfEntriesInSegment.get(keySegment) + 1; + expectedNumberOfEntriesInSegment.put(keySegment, updatedValue); + } + tieredSpilloverCacheSegmentMap.put(iCacheKey, tieredSpilloverCache.getTieredCacheSegment(iCacheKey)); + tieredSpilloverCache.computeIfAbsent(iCacheKey, tieredCacheLoader); + } + + // We now calculate expected onHeap cache entries and then verify it later. + for (int i = 0; i < numberOfSegments; i++) { + if (expectedNumberOfEntriesInSegment.get(i) == null) { + continue; + } + if (expectedNumberOfEntriesInSegment.get(i) >= expectedSegmentOnHeapCacheSize.get(i)) { + totalOnHeapEntries += expectedSegmentOnHeapCacheSize.get(i); + totalOnDiskEntries += expectedNumberOfEntriesInSegment.get(i) - expectedSegmentOnHeapCacheSize.get(i); + } else { + // In this case onHeap cache wasn't utilized fully. + totalOnHeapEntries += expectedNumberOfEntriesInSegment.get(i); + } + } + + tieredSpilloverCache.getOnHeapCacheKeys().forEach(onHeapKeys::add); + tieredSpilloverCache.getDiskCacheKeys().forEach(diskTierKeys::add); + // Verify on heap cache stats. + assertEquals(totalOnHeapEntries, tieredSpilloverCache.onHeapCacheCount()); + assertEquals(totalOnHeapEntries, getItemsForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_ON_HEAP)); + assertEquals(numOfItems1, getMissesForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_ON_HEAP)); + assertEquals(0, getHitsForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_ON_HEAP)); + assertEquals(totalOnDiskEntries, getEvictionsForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_ON_HEAP)); + assertEquals(totalOnHeapEntries * keyValueSize, getSizeInBytesForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_ON_HEAP)); + + // Verify disk cache stats. + assertEquals(totalOnDiskEntries, tieredSpilloverCache.diskCacheCount()); + assertEquals(totalOnDiskEntries, getItemsForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_DISK)); + assertEquals(0, getHitsForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_DISK)); + assertEquals(numOfItems1, getMissesForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_DISK)); + assertEquals(0, getEvictionsForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_DISK)); + assertEquals(totalOnDiskEntries * keyValueSize, getSizeInBytesForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_DISK)); + + // Verify the keys for onHeap and disk cache + + for (ICacheKey key : onHeapKeys) { + assertNotNull(tieredSpilloverCache.get(key)); + } + assertEquals(numOfItems1, getMissesForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_ON_HEAP)); + assertEquals(onHeapKeys.size(), getHitsForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_ON_HEAP)); + for (ICacheKey key : diskTierKeys) { + assertNotNull(tieredSpilloverCache.get(key)); + } + assertEquals(numOfItems1, getMissesForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_DISK)); + assertEquals(diskTierKeys.size(), getHitsForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_DISK)); + } + public void testWithFactoryCreationWithOnHeapCacheNotPresent() { int onHeapCacheSize = randomIntBetween(10, 30); int keyValueSize = 50; @@ -239,9 +392,9 @@ public void testWithFactoryCreationWithOnHeapCacheNotPresent() { MockDiskCache.MockDiskCacheFactory.NAME ) .put( - OpenSearchOnHeapCacheSettings.getSettingListForCacheType(CacheType.INDICES_REQUEST_CACHE) - .get(MAXIMUM_SIZE_IN_BYTES_KEY) - .getKey(), + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), onHeapCacheSize * keyValueSize + "b" ) .put( @@ -265,7 +418,7 @@ public void testWithFactoryCreationWithOnHeapCacheNotPresent() { OpenSearchOnHeapCache.OpenSearchOnHeapCacheFactory.NAME, new OpenSearchOnHeapCache.OpenSearchOnHeapCacheFactory(), MockDiskCache.MockDiskCacheFactory.NAME, - new MockDiskCache.MockDiskCacheFactory(0, randomIntBetween(100, 300), false) + new MockDiskCache.MockDiskCacheFactory(0, randomIntBetween(100, 300), false, keyValueSize) ) ) ); @@ -289,9 +442,9 @@ public void testWithFactoryCreationWithDiskCacheNotPresent() { OpenSearchOnHeapCache.OpenSearchOnHeapCacheFactory.NAME ) .put( - OpenSearchOnHeapCacheSettings.getSettingListForCacheType(CacheType.INDICES_REQUEST_CACHE) - .get(MAXIMUM_SIZE_IN_BYTES_KEY) - .getKey(), + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), onHeapCacheSize * keyValueSize + "b" ) .build(); @@ -310,7 +463,7 @@ public void testWithFactoryCreationWithDiskCacheNotPresent() { OpenSearchOnHeapCache.OpenSearchOnHeapCacheFactory.NAME, new OpenSearchOnHeapCache.OpenSearchOnHeapCacheFactory(), MockDiskCache.MockDiskCacheFactory.NAME, - new MockDiskCache.MockDiskCacheFactory(0, randomIntBetween(100, 300), false) + new MockDiskCache.MockDiskCacheFactory(0, randomIntBetween(100, 300), false, keyValueSize) ) ) ); @@ -327,6 +480,20 @@ public void testComputeIfAbsentWithEvictionsFromOnHeapCache() throws Exception { int keyValueSize = 50; MockCacheRemovalListener removalListener = new MockCacheRemovalListener<>(); ICache.Factory onHeapCacheFactory = new OpenSearchOnHeapCache.OpenSearchOnHeapCacheFactory(); + + Settings settings = Settings.builder() + .put( + CacheSettings.getConcreteStoreNameSettingForCacheType(CacheType.INDICES_REQUEST_CACHE).getKey(), + TieredSpilloverCache.TieredSpilloverCacheFactory.TIERED_SPILLOVER_CACHE_NAME + ) + .put(FeatureFlags.PLUGGABLE_CACHE, "true") + .put( + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), + onHeapCacheSize * keyValueSize + "b" + ) + .build(); CacheConfig cacheConfig = new CacheConfig.Builder().setKeyType(String.class) .setKeyType(String.class) .setWeigher((k, v) -> keyValueSize) @@ -334,25 +501,12 @@ public void testComputeIfAbsentWithEvictionsFromOnHeapCache() throws Exception { .setKeySerializer(new StringSerializer()) .setValueSerializer(new StringSerializer()) .setDimensionNames(dimensionNames) - .setSettings( - Settings.builder() - .put( - CacheSettings.getConcreteStoreNameSettingForCacheType(CacheType.INDICES_REQUEST_CACHE).getKey(), - TieredSpilloverCache.TieredSpilloverCacheFactory.TIERED_SPILLOVER_CACHE_NAME - ) - .put(FeatureFlags.PLUGGABLE_CACHE, "true") - .put( - OpenSearchOnHeapCacheSettings.getSettingListForCacheType(CacheType.INDICES_REQUEST_CACHE) - .get(MAXIMUM_SIZE_IN_BYTES_KEY) - .getKey(), - onHeapCacheSize * keyValueSize + "b" - ) - .build() - ) + .setSettings(settings) + .setStoragePath(getStoragePath(settings)) .setClusterSettings(clusterSettings) .build(); - ICache.Factory mockDiskCacheFactory = new MockDiskCache.MockDiskCacheFactory(0, diskCacheSize, false); + ICache.Factory mockDiskCacheFactory = new MockDiskCache.MockDiskCacheFactory(0, diskCacheSize, false, keyValueSize); TieredSpilloverCache tieredSpilloverCache = new TieredSpilloverCache.Builder() .setOnHeapCacheFactory(onHeapCacheFactory) @@ -360,6 +514,9 @@ public void testComputeIfAbsentWithEvictionsFromOnHeapCache() throws Exception { .setCacheConfig(cacheConfig) .setRemovalListener(removalListener) .setCacheType(CacheType.INDICES_REQUEST_CACHE) + .setNumberOfSegments(1) + .setOnHeapCacheSizeInBytes(onHeapCacheSize * keyValueSize) + .setDiskCacheSize(diskCacheSize) .build(); // Put values in cache more than it's size and cause evictions from onHeap. @@ -372,7 +529,7 @@ public void testComputeIfAbsentWithEvictionsFromOnHeapCache() throws Exception { tieredSpilloverCache.computeIfAbsent(key, tieredCacheLoader); } - long actualDiskCacheSize = tieredSpilloverCache.getDiskCache().count(); + long actualDiskCacheSize = tieredSpilloverCache.diskCacheCount(); assertEquals(numOfItems1, getMissesForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_ON_HEAP)); assertEquals(0, getHitsForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_ON_HEAP)); @@ -381,8 +538,8 @@ public void testComputeIfAbsentWithEvictionsFromOnHeapCache() throws Exception { assertEquals(onHeapCacheSize * keyValueSize, getSizeInBytesForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_ON_HEAP)); assertEquals(actualDiskCacheSize, getItemsForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_DISK)); - tieredSpilloverCache.getOnHeapCache().keys().forEach(onHeapKeys::add); - tieredSpilloverCache.getDiskCache().keys().forEach(diskTierKeys::add); + tieredSpilloverCache.getOnHeapCacheKeys().forEach(onHeapKeys::add); + tieredSpilloverCache.getDiskCacheKeys().forEach(diskTierKeys::add); // Try to hit cache again with some randomization. int numOfItems2 = randomIntBetween(50, 200); @@ -420,11 +577,11 @@ public void testComputeIfAbsentWithEvictionsFromOnHeapCache() throws Exception { } public void testComputeIfAbsentWithEvictionsFromTieredCache() throws Exception { - int onHeapCacheSize = randomIntBetween(10, 30); - int diskCacheSize = randomIntBetween(onHeapCacheSize + 1, 100); + int onHeapCacheSize = randomIntBetween(300, 600); + int diskCacheSize = randomIntBetween(700, 1200); int totalSize = onHeapCacheSize + diskCacheSize; - int keyValueSize = 50; - + int numberOfSegments = getNumberOfSegments(); + int keyValueSize = 11; MockCacheRemovalListener removalListener = new MockCacheRemovalListener<>(); TieredSpilloverCache tieredSpilloverCache = initializeTieredSpilloverCache( keyValueSize, @@ -432,25 +589,53 @@ public void testComputeIfAbsentWithEvictionsFromTieredCache() throws Exception { removalListener, Settings.builder() .put( - OpenSearchOnHeapCacheSettings.getSettingListForCacheType(CacheType.INDICES_REQUEST_CACHE) - .get(MAXIMUM_SIZE_IN_BYTES_KEY) - .getKey(), + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), onHeapCacheSize * keyValueSize + "b" ) .build(), - 0 + 0, + numberOfSegments + ); + Map onHeapCacheExpectedSize = getSegmentOnHeapCacheSize( + numberOfSegments, + onHeapCacheSize * keyValueSize, + keyValueSize ); + Map mockDiskCacheExpectedSize = getSegmentMockDiskCacheSize(numberOfSegments, diskCacheSize); + Map perSegmentEntryCapacity = new HashMap<>(); + for (int i = 0; i < numberOfSegments; i++) { + int totalEntriesForSegment = onHeapCacheExpectedSize.get(i) + mockDiskCacheExpectedSize.get(i); + perSegmentEntryCapacity.put(i, totalEntriesForSegment); + } int numOfItems = randomIntBetween(totalSize + 1, totalSize * 3); + Map segmentSizeTracker = new HashMap<>(); + int expectedEvictions = 0; for (int iter = 0; iter < numOfItems; iter++) { + ICacheKey iCacheKey = getICacheKey(UUID.randomUUID().toString()); + int keySegment = tieredSpilloverCache.getSegmentNumber(iCacheKey); LoadAwareCacheLoader, String> tieredCacheLoader = getLoadAwareCacheLoader(); - tieredSpilloverCache.computeIfAbsent(getICacheKey(UUID.randomUUID().toString()), tieredCacheLoader); + tieredSpilloverCache.computeIfAbsent(iCacheKey, tieredCacheLoader); + if (segmentSizeTracker.get(keySegment) == null) { + segmentSizeTracker.put(keySegment, Integer.valueOf(1)); + } else { + Integer updatedValue = segmentSizeTracker.get(keySegment) + 1; + segmentSizeTracker.put(keySegment, updatedValue); + } } - - int evictions = numOfItems - (totalSize); // Evictions from the cache as a whole - assertEquals(evictions, removalListener.evictionsMetric.count()); - assertEquals(evictions, getEvictionsForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_DISK)); + for (int i = 0; i < numberOfSegments; i++) { + if (segmentSizeTracker.get(i) == null) { + continue; + } + if (segmentSizeTracker.get(i) > perSegmentEntryCapacity.get(i)) { + expectedEvictions += segmentSizeTracker.get(i) - perSegmentEntryCapacity.get(i); + } + } + assertEquals(expectedEvictions, removalListener.evictionsMetric.count()); + assertEquals(expectedEvictions, getEvictionsForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_DISK)); assertEquals( - evictions + getItemsForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_DISK), + expectedEvictions + getItemsForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_DISK), getEvictionsForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_ON_HEAP) ); } @@ -468,13 +653,14 @@ public void testGetAndCount() throws Exception { removalListener, Settings.builder() .put( - OpenSearchOnHeapCacheSettings.getSettingListForCacheType(CacheType.INDICES_REQUEST_CACHE) - .get(MAXIMUM_SIZE_IN_BYTES_KEY) - .getKey(), + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), onHeapCacheSize * keyValueSize + "b" ) .build(), - 0 + 0, + 1 ); int numOfItems1 = randomIntBetween(onHeapCacheSize + 1, totalSize); @@ -509,8 +695,9 @@ public void testGetAndCount() throws Exception { } public void testPut() throws Exception { - int onHeapCacheSize = randomIntBetween(10, 30); - int diskCacheSize = randomIntBetween(onHeapCacheSize + 1, 100); + int numberOfSegments = getNumberOfSegments(); + int onHeapCacheSize = randomIntBetween(10, 30) * numberOfSegments; + int diskCacheSize = randomIntBetween(onHeapCacheSize + 1, onHeapCacheSize * 2); int keyValueSize = 50; MockCacheRemovalListener removalListener = new MockCacheRemovalListener<>(); @@ -520,13 +707,14 @@ public void testPut() throws Exception { removalListener, Settings.builder() .put( - OpenSearchOnHeapCacheSettings.getSettingListForCacheType(CacheType.INDICES_REQUEST_CACHE) - .get(MAXIMUM_SIZE_IN_BYTES_KEY) - .getKey(), + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), onHeapCacheSize * keyValueSize + "b" ) .build(), - 0 + 0, + numberOfSegments ); ICacheKey key = getICacheKey(UUID.randomUUID().toString()); String value = UUID.randomUUID().toString(); @@ -541,7 +729,7 @@ public void testPutAndVerifyNewItemsArePresentOnHeapCache() throws Exception { int keyValueSize = 50; MockCacheRemovalListener removalListener = new MockCacheRemovalListener<>(); - + int numberOfSegments = 1; TieredSpilloverCache tieredSpilloverCache = initializeTieredSpilloverCache( keyValueSize, diskCacheSize, @@ -552,13 +740,14 @@ public void testPutAndVerifyNewItemsArePresentOnHeapCache() throws Exception { TieredSpilloverCache.TieredSpilloverCacheFactory.TIERED_SPILLOVER_CACHE_NAME ) .put( - OpenSearchOnHeapCacheSettings.getSettingListForCacheType(CacheType.INDICES_REQUEST_CACHE) - .get(MAXIMUM_SIZE_IN_BYTES_KEY) - .getKey(), + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), (onHeapCacheSize * keyValueSize) + "b" ) .build(), - 0 + 0, + numberOfSegments ); for (int i = 0; i < onHeapCacheSize; i++) { @@ -580,7 +769,7 @@ public void testPutAndVerifyNewItemsArePresentOnHeapCache() throws Exception { // Verify that new items are part of onHeap cache. List> actualOnHeapCacheKeys = new ArrayList<>(); - tieredSpilloverCache.getOnHeapCache().keys().forEach(actualOnHeapCacheKeys::add); + tieredSpilloverCache.getOnHeapCacheKeys().forEach(actualOnHeapCacheKeys::add); assertEquals(newKeyList.size(), actualOnHeapCacheKeys.size()); for (int i = 0; i < actualOnHeapCacheKeys.size(); i++) { @@ -594,7 +783,6 @@ public void testInvalidate() throws Exception { int onHeapCacheSize = 1; int diskCacheSize = 10; int keyValueSize = 20; - MockCacheRemovalListener removalListener = new MockCacheRemovalListener<>(); TieredSpilloverCache tieredSpilloverCache = initializeTieredSpilloverCache( keyValueSize, @@ -602,13 +790,14 @@ public void testInvalidate() throws Exception { removalListener, Settings.builder() .put( - OpenSearchOnHeapCacheSettings.getSettingListForCacheType(CacheType.INDICES_REQUEST_CACHE) - .get(MAXIMUM_SIZE_IN_BYTES_KEY) - .getKey(), + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), onHeapCacheSize * keyValueSize + "b" ) .build(), - 0 + 0, + 1 ); ICacheKey key = getICacheKey(UUID.randomUUID().toString()); String value = UUID.randomUUID().toString(); @@ -652,13 +841,14 @@ public void testCacheKeys() throws Exception { removalListener, Settings.builder() .put( - OpenSearchOnHeapCacheSettings.getSettingListForCacheType(CacheType.INDICES_REQUEST_CACHE) - .get(MAXIMUM_SIZE_IN_BYTES_KEY) - .getKey(), + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), onHeapCacheSize * keyValueSize + "b" ) .build(), - 0 + 0, + 1 ); List> onHeapKeys = new ArrayList<>(); List> diskTierKeys = new ArrayList<>(); @@ -678,8 +868,8 @@ public void testCacheKeys() throws Exception { List> actualOnHeapKeys = new ArrayList<>(); List> actualOnDiskKeys = new ArrayList<>(); - Iterable> onHeapiterable = tieredSpilloverCache.getOnHeapCache().keys(); - Iterable> onDiskiterable = tieredSpilloverCache.getDiskCache().keys(); + Iterable> onHeapiterable = tieredSpilloverCache.getOnHeapCacheKeys(); + Iterable> onDiskiterable = tieredSpilloverCache.getDiskCacheKeys(); onHeapiterable.iterator().forEachRemaining(actualOnHeapKeys::add); onDiskiterable.iterator().forEachRemaining(actualOnDiskKeys::add); for (ICacheKey onHeapKey : onHeapKeys) { @@ -713,16 +903,18 @@ public void testRefresh() { diskCacheSize, removalListener, Settings.EMPTY, - 0 + 0, + 1 ); tieredSpilloverCache.refresh(); } public void testInvalidateAll() throws Exception { - int onHeapCacheSize = randomIntBetween(10, 30); - int diskCacheSize = randomIntBetween(60, 100); - int keyValueSize = 50; + int onHeapCacheSize = randomIntBetween(300, 600); + int diskCacheSize = randomIntBetween(700, 1200); int totalSize = onHeapCacheSize + diskCacheSize; + int numberOfSegments = getNumberOfSegments(); + int keyValueSize = 50; MockCacheRemovalListener removalListener = new MockCacheRemovalListener<>(); TieredSpilloverCache tieredSpilloverCache = initializeTieredSpilloverCache( @@ -731,30 +923,29 @@ public void testInvalidateAll() throws Exception { removalListener, Settings.builder() .put( - OpenSearchOnHeapCacheSettings.getSettingListForCacheType(CacheType.INDICES_REQUEST_CACHE) - .get(MAXIMUM_SIZE_IN_BYTES_KEY) - .getKey(), + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), onHeapCacheSize * keyValueSize + "b" ) .build(), - 0 + 0, + numberOfSegments ); // Put values in cache more than it's size and cause evictions from onHeap. int numOfItems1 = randomIntBetween(onHeapCacheSize + 1, totalSize); - List> onHeapKeys = new ArrayList<>(); - List> diskTierKeys = new ArrayList<>(); for (int iter = 0; iter < numOfItems1; iter++) { ICacheKey key = getICacheKey(UUID.randomUUID().toString()); - if (iter > (onHeapCacheSize - 1)) { - // All these are bound to go to disk based cache. - diskTierKeys.add(key); - } else { - onHeapKeys.add(key); - } LoadAwareCacheLoader, String> tieredCacheLoader = getLoadAwareCacheLoader(); tieredSpilloverCache.computeIfAbsent(key, tieredCacheLoader); } - assertEquals(numOfItems1, tieredSpilloverCache.count()); + assertEquals( + getItemsForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_ON_HEAP) + getItemsForTier( + tieredSpilloverCache, + TIER_DIMENSION_VALUE_DISK + ), + tieredSpilloverCache.count() + ); tieredSpilloverCache.invalidateAll(); assertEquals(0, tieredSpilloverCache.count()); } @@ -767,9 +958,9 @@ public void testComputeIfAbsentConcurrently() throws Exception { MockCacheRemovalListener removalListener = new MockCacheRemovalListener<>(); Settings settings = Settings.builder() .put( - OpenSearchOnHeapCacheSettings.getSettingListForCacheType(CacheType.INDICES_REQUEST_CACHE) - .get(MAXIMUM_SIZE_IN_BYTES_KEY) - .getKey(), + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), onHeapCacheSize * keyValueSize + "b" ) .build(); @@ -779,7 +970,8 @@ public void testComputeIfAbsentConcurrently() throws Exception { diskCacheSize, removalListener, settings, - 0 + 0, + 1 ); int numberOfSameKeys = randomIntBetween(400, onHeapCacheSize - 1); @@ -837,16 +1029,17 @@ public String load(ICacheKey key) { } public void testComputIfAbsentConcurrentlyWithMultipleKeys() throws Exception { - int onHeapCacheSize = randomIntBetween(300, 500); + int numberOfSegments = getNumberOfSegments(); + int onHeapCacheSize = randomIntBetween(300, 500) * numberOfSegments; // Able to support all keys in case of + // skewness as well. int diskCacheSize = randomIntBetween(600, 700); int keyValueSize = 50; - MockCacheRemovalListener removalListener = new MockCacheRemovalListener<>(); Settings settings = Settings.builder() .put( - OpenSearchOnHeapCacheSettings.getSettingListForCacheType(CacheType.INDICES_REQUEST_CACHE) - .get(MAXIMUM_SIZE_IN_BYTES_KEY) - .getKey(), + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), onHeapCacheSize * keyValueSize + "b" ) .build(); @@ -856,7 +1049,8 @@ public void testComputIfAbsentConcurrentlyWithMultipleKeys() throws Exception { diskCacheSize, removalListener, settings, - 0 + 0, + numberOfSegments ); int iterations = 10; @@ -942,9 +1136,9 @@ public void testComputeIfAbsentWithOnHeapCacheThrowingExceptionOnPut() throws Ex MockCacheRemovalListener removalListener = new MockCacheRemovalListener<>(); Settings settings = Settings.builder() .put( - OpenSearchOnHeapCacheSettings.getSettingListForCacheType(CacheType.INDICES_REQUEST_CACHE) - .get(MAXIMUM_SIZE_IN_BYTES_KEY) - .getKey(), + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), onHeapCacheSize * keyValueSize + "b" ) .build(); @@ -952,15 +1146,18 @@ public void testComputeIfAbsentWithOnHeapCacheThrowingExceptionOnPut() throws Ex ICache mockOnHeapCache = mock(ICache.class); when(onHeapCacheFactory.create(any(), any(), any())).thenReturn(mockOnHeapCache); doThrow(new RuntimeException("Testing")).when(mockOnHeapCache).put(any(), any()); - CacheConfig cacheConfig = getCacheConfig(keyValueSize, settings, removalListener); - ICache.Factory mockDiskCacheFactory = new MockDiskCache.MockDiskCacheFactory(0, diskCacheSize, false); + CacheConfig cacheConfig = getCacheConfig(keyValueSize, settings, removalListener, 1); + ICache.Factory mockDiskCacheFactory = new MockDiskCache.MockDiskCacheFactory(0, diskCacheSize, false, keyValueSize); TieredSpilloverCache tieredSpilloverCache = getTieredSpilloverCache( onHeapCacheFactory, mockDiskCacheFactory, cacheConfig, null, - removalListener + removalListener, + 1, + onHeapCacheSize * keyValueSize, + diskCacheSize ); String value = ""; value = tieredSpilloverCache.computeIfAbsent(getICacheKey("test"), new LoadAwareCacheLoader<>() { @@ -980,20 +1177,20 @@ public String load(ICacheKey key) { @SuppressWarnings({ "rawtypes", "unchecked" }) public void testComputeIfAbsentWithDiskCacheThrowingExceptionOnPut() throws Exception { - int onHeapCacheSize = 0; + int onHeapCacheSize = 1; int keyValueSize = 50; MockCacheRemovalListener removalListener = new MockCacheRemovalListener<>(); Settings settings = Settings.builder() .put( - OpenSearchOnHeapCacheSettings.getSettingListForCacheType(CacheType.INDICES_REQUEST_CACHE) - .get(MAXIMUM_SIZE_IN_BYTES_KEY) - .getKey(), + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), onHeapCacheSize * keyValueSize + "b" ) .build(); ICache.Factory onHeapCacheFactory = new OpenSearchOnHeapCache.OpenSearchOnHeapCacheFactory(); - CacheConfig cacheConfig = getCacheConfig(keyValueSize, settings, removalListener); + CacheConfig cacheConfig = getCacheConfig(keyValueSize, settings, removalListener, 1); ICache.Factory mockDiskCacheFactory = mock(MockDiskCache.MockDiskCacheFactory.class); ICache mockDiskCache = mock(ICache.class); when(mockDiskCacheFactory.create(any(), any(), any())).thenReturn(mockDiskCache); @@ -1004,10 +1201,25 @@ public void testComputeIfAbsentWithDiskCacheThrowingExceptionOnPut() throws Exce mockDiskCacheFactory, cacheConfig, null, - removalListener + removalListener, + 1, + onHeapCacheSize * keyValueSize, + 200 ); String response = ""; + // This first computeIfAbsent ensures onHeap cache has 1 item present and rest will be evicted to disk. + tieredSpilloverCache.computeIfAbsent(getICacheKey("test1"), new LoadAwareCacheLoader<>() { + @Override + public boolean isLoaded() { + return false; + } + + @Override + public String load(ICacheKey key) { + return "test1"; + } + }); response = tieredSpilloverCache.computeIfAbsent(getICacheKey("test"), new LoadAwareCacheLoader<>() { @Override public boolean isLoaded() { @@ -1050,29 +1262,37 @@ public void testConcurrencyForEvictionFlowFromOnHeapToDiskTier() throws Exceptio MockCacheRemovalListener removalListener = new MockCacheRemovalListener<>(); ICache.Factory onHeapCacheFactory = new OpenSearchOnHeapCache.OpenSearchOnHeapCacheFactory(); - ICache.Factory diskCacheFactory = new MockDiskCache.MockDiskCacheFactory(500, diskCacheSize, false); + ICache.Factory diskCacheFactory = new MockDiskCache.MockDiskCacheFactory(500, diskCacheSize, false, 1); + + Settings settings = Settings.builder() + .put( + CacheSettings.getConcreteStoreNameSettingForCacheType(CacheType.INDICES_REQUEST_CACHE).getKey(), + TieredSpilloverCache.TieredSpilloverCacheFactory.TIERED_SPILLOVER_CACHE_NAME + ) + .put(FeatureFlags.PLUGGABLE_CACHE, "true") + .put( + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), + 200 + "b" + ) + .build(); + String storagePath; + try (NodeEnvironment environment = newNodeEnvironment(settings)) { + storagePath = environment.nodePaths()[0].path + "/test"; + } catch (IOException e) { + throw new OpenSearchException("Exception occurred", e); + } + CacheConfig cacheConfig = new CacheConfig.Builder().setKeyType(String.class) .setKeyType(String.class) .setWeigher((k, v) -> 150) .setRemovalListener(removalListener) .setKeySerializer(new StringSerializer()) .setValueSerializer(new StringSerializer()) - .setSettings( - Settings.builder() - .put( - CacheSettings.getConcreteStoreNameSettingForCacheType(CacheType.INDICES_REQUEST_CACHE).getKey(), - TieredSpilloverCache.TieredSpilloverCacheFactory.TIERED_SPILLOVER_CACHE_NAME - ) - .put(FeatureFlags.PLUGGABLE_CACHE, "true") - .put( - OpenSearchOnHeapCacheSettings.getSettingListForCacheType(CacheType.INDICES_REQUEST_CACHE) - .get(MAXIMUM_SIZE_IN_BYTES_KEY) - .getKey(), - 200 + "b" - ) - .build() - ) + .setSettings(settings) .setClusterSettings(clusterSettings) + .setStoragePath(storagePath) .setDimensionNames(dimensionNames) .build(); TieredSpilloverCache tieredSpilloverCache = new TieredSpilloverCache.Builder() @@ -1080,6 +1300,9 @@ public void testConcurrencyForEvictionFlowFromOnHeapToDiskTier() throws Exceptio .setDiskCacheFactory(diskCacheFactory) .setRemovalListener(removalListener) .setCacheConfig(cacheConfig) + .setOnHeapCacheSizeInBytes(200) + .setDiskCacheSize(diskCacheSize) + .setNumberOfSegments(1) .setCacheType(CacheType.INDICES_REQUEST_CACHE) .build(); @@ -1104,9 +1327,7 @@ public void testConcurrencyForEvictionFlowFromOnHeapToDiskTier() throws Exceptio }); thread.start(); assertBusy(() -> { assertTrue(loadAwareCacheLoader.isLoaded()); }, 100, TimeUnit.MILLISECONDS); // We wait for new key to be loaded - // after which it eviction flow is - // guaranteed to occur. - ICache onDiskCache = tieredSpilloverCache.getDiskCache(); + // after which it eviction flow is guaranteed to occur. // Now on a different thread, try to get key(above one which got evicted) from tiered cache. We expect this // should return not null value as it should be present on diskCache. @@ -1124,13 +1345,13 @@ public void testConcurrencyForEvictionFlowFromOnHeapToDiskTier() throws Exceptio assertNotNull(actualValue.get()); countDownLatch1.await(); - assertEquals(1, tieredSpilloverCache.getOnHeapCache().count()); - assertEquals(1, onDiskCache.count()); + assertEquals(1, tieredSpilloverCache.onHeapCacheCount()); + assertEquals(1, tieredSpilloverCache.diskCacheCount()); assertEquals(1, getEvictionsForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_ON_HEAP)); assertEquals(1, getItemsForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_ON_HEAP)); assertEquals(1, getItemsForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_DISK)); - assertNotNull(onDiskCache.get(keyToBeEvicted)); + assertNotNull(tieredSpilloverCache.getTieredCacheSegment(keyToBeEvicted).getDiskCache().get(keyToBeEvicted)); } public void testDiskTierPolicies() throws Exception { @@ -1140,7 +1361,6 @@ public void testDiskTierPolicies() throws Exception { policies.add(new AllowEvenLengths()); int keyValueSize = 50; - int onHeapCacheSize = 0; MockCacheRemovalListener removalListener = new MockCacheRemovalListener<>(); TieredSpilloverCache tieredSpilloverCache = intializeTieredSpilloverCache( keyValueSize, @@ -1148,14 +1368,15 @@ public void testDiskTierPolicies() throws Exception { removalListener, Settings.builder() .put( - OpenSearchOnHeapCacheSettings.getSettingListForCacheType(CacheType.INDICES_REQUEST_CACHE) - .get(MAXIMUM_SIZE_IN_BYTES_KEY) - .getKey(), - onHeapCacheSize * keyValueSize + "b" + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), + keyValueSize - 1 + "b" ) .build(), 0, - policies + policies, + 1 ); Map keyValuePairs = new HashMap<>(); @@ -1233,15 +1454,16 @@ public void testTookTimePolicyFromFactory() throws Exception { MockDiskCache.MockDiskCacheFactory.NAME ) .put( - OpenSearchOnHeapCacheSettings.getSettingListForCacheType(CacheType.INDICES_REQUEST_CACHE) - .get(MAXIMUM_SIZE_IN_BYTES_KEY) - .getKey(), + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), onHeapCacheSize * keyValueSize + "b" ) .put( TieredSpilloverCacheSettings.TOOK_TIME_POLICY_CONCRETE_SETTINGS_MAP.get(CacheType.INDICES_REQUEST_CACHE).getKey(), new TimeValue(timeValueThresholdNanos / 1_000_000) ) + .put(TIERED_SPILLOVER_SEGMENTS.getConcreteSettingForNamespace(CacheType.INDICES_REQUEST_CACHE.getSettingPrefix()).getKey(), 1) .build(); ICache tieredSpilloverICache = new TieredSpilloverCache.TieredSpilloverCacheFactory().create( @@ -1261,13 +1483,14 @@ public CachedQueryResult.PolicyValues apply(String s) { } }) .setClusterSettings(clusterSettings) + .setStoragePath(getStoragePath(settings)) .build(), CacheType.INDICES_REQUEST_CACHE, Map.of( OpenSearchOnHeapCache.OpenSearchOnHeapCacheFactory.NAME, new OpenSearchOnHeapCache.OpenSearchOnHeapCacheFactory(), MockDiskCache.MockDiskCacheFactory.NAME, - new MockDiskCache.MockDiskCacheFactory(0, randomIntBetween(100, 300), false) + new MockDiskCache.MockDiskCacheFactory(0, randomIntBetween(100, 300), false, keyValueSize) ) ); @@ -1279,13 +1502,14 @@ public CachedQueryResult.PolicyValues apply(String s) { } assertEquals(tookTimeMap.size(), tieredSpilloverCache.count()); - // Ensure all these keys get evicted from the on heap tier by adding > heap tier size worth of random keys + // Ensure all these keys get evicted from the on heap tier by adding > heap tier size worth of random keys (this works as we have 1 + // segment) for (int i = 0; i < onHeapCacheSize; i++) { tieredSpilloverCache.computeIfAbsent(getICacheKey(UUID.randomUUID().toString()), getLoadAwareCacheLoader(keyValueMap)); } - ICache onHeapCache = tieredSpilloverCache.getOnHeapCache(); for (String key : tookTimeMap.keySet()) { - assertNull(onHeapCache.get(getICacheKey(key))); + ICacheKey iCacheKey = getICacheKey(key); + assertNull(tieredSpilloverCache.getTieredCacheSegment(iCacheKey).getOnHeapCache().get(iCacheKey)); } // Now the original keys should be in the disk tier if the policy allows them, or misses if not @@ -1320,10 +1544,10 @@ public void testMinimumThresholdSettingValue() throws Exception { public void testPutWithDiskCacheDisabledSetting() throws Exception { int onHeapCacheSize = randomIntBetween(10, 30); - int diskCacheSize = randomIntBetween(onHeapCacheSize + 1, 100); + int diskCacheSize = randomIntBetween(300, 500); int keyValueSize = 50; int totalSize = onHeapCacheSize + diskCacheSize; - + int numberOfSegments = getNumberOfSegments(); MockCacheRemovalListener removalListener = new MockCacheRemovalListener<>(); TieredSpilloverCache tieredSpilloverCache = initializeTieredSpilloverCache( keyValueSize, @@ -1331,14 +1555,15 @@ public void testPutWithDiskCacheDisabledSetting() throws Exception { removalListener, Settings.builder() .put( - OpenSearchOnHeapCacheSettings.getSettingListForCacheType(CacheType.INDICES_REQUEST_CACHE) - .get(MAXIMUM_SIZE_IN_BYTES_KEY) - .getKey(), + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), onHeapCacheSize * keyValueSize + "b" ) .put(DISK_CACHE_ENABLED_SETTING_MAP.get(CacheType.INDICES_REQUEST_CACHE).getKey(), false) .build(), - 0 + 0, + numberOfSegments ); int numOfItems1 = randomIntBetween(onHeapCacheSize + 1, totalSize); // Create more items than onHeap cache. @@ -1347,59 +1572,95 @@ public void testPutWithDiskCacheDisabledSetting() throws Exception { LoadAwareCacheLoader, String> loadAwareCacheLoader = getLoadAwareCacheLoader(); tieredSpilloverCache.computeIfAbsent(key, loadAwareCacheLoader); } - ICache onHeapCache = tieredSpilloverCache.getOnHeapCache(); - ICache diskCache = tieredSpilloverCache.getDiskCache(); - assertEquals(onHeapCacheSize, onHeapCache.count()); - assertEquals(0, diskCache.count()); // Disk cache shouldn't have anything considering it is disabled. - assertEquals(numOfItems1 - onHeapCacheSize, removalListener.evictionsMetric.count()); + + assertEquals(getItemsForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_ON_HEAP), tieredSpilloverCache.onHeapCacheCount()); + assertEquals(0, tieredSpilloverCache.diskCacheCount()); // Disk cache shouldn't have anything considering it is + // disabled. + assertEquals(numOfItems1 - tieredSpilloverCache.onHeapCacheCount(), removalListener.evictionsMetric.count()); } public void testGetPutAndInvalidateWithDiskCacheDisabled() throws Exception { - int onHeapCacheSize = randomIntBetween(10, 30); - int diskCacheSize = randomIntBetween(onHeapCacheSize + 1, 100); - int keyValueSize = 50; + int onHeapCacheSize = randomIntBetween(300, 400); + int diskCacheSize = randomIntBetween(onHeapCacheSize + 1, 500); + int keyValueSize = 12; int totalSize = onHeapCacheSize + diskCacheSize; + MockCacheRemovalListener removalListener = new MockCacheRemovalListener<>(); + int numberOfSegments = getNumberOfSegments(); TieredSpilloverCache tieredSpilloverCache = initializeTieredSpilloverCache( keyValueSize, diskCacheSize, removalListener, Settings.builder() .put( - OpenSearchOnHeapCacheSettings.getSettingListForCacheType(CacheType.INDICES_REQUEST_CACHE) - .get(MAXIMUM_SIZE_IN_BYTES_KEY) - .getKey(), + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), onHeapCacheSize * keyValueSize + "b" ) .build(), - 0 + 0, + numberOfSegments ); - + Map onHeapCacheExpectedSize = getSegmentOnHeapCacheSize( + numberOfSegments, + onHeapCacheSize * keyValueSize, + keyValueSize + ); + Map mockDiskCacheExpectedSize = getSegmentMockDiskCacheSize(numberOfSegments, diskCacheSize); + Map perSegmentEntryCapacity = new HashMap<>(); + for (int i = 0; i < numberOfSegments; i++) { + int totalEntriesForSegment = onHeapCacheExpectedSize.get(i) + mockDiskCacheExpectedSize.get(i); + perSegmentEntryCapacity.put(i, totalEntriesForSegment); + } int numOfItems1 = randomIntBetween(onHeapCacheSize + 1, totalSize - 1); // Create more items than onHeap // cache to cause spillover. + Map segmentSizeTracker = new HashMap<>(); for (int iter = 0; iter < numOfItems1; iter++) { ICacheKey key = getICacheKey(UUID.randomUUID().toString()); + int keySegment = tieredSpilloverCache.getSegmentNumber(key); LoadAwareCacheLoader, String> loadAwareCacheLoader = getLoadAwareCacheLoader(); tieredSpilloverCache.computeIfAbsent(key, loadAwareCacheLoader); + if (segmentSizeTracker.get(keySegment) == null) { + segmentSizeTracker.put(keySegment, Integer.valueOf(1)); + } else { + Integer updatedValue = segmentSizeTracker.get(keySegment) + 1; + segmentSizeTracker.put(keySegment, updatedValue); + } + } + int expectedEvictions = 0; + for (int i = 0; i < numberOfSegments; i++) { + if (segmentSizeTracker.get(i) == null) { + continue; + } + if (segmentSizeTracker.get(i) > perSegmentEntryCapacity.get(i)) { + expectedEvictions += segmentSizeTracker.get(i) - perSegmentEntryCapacity.get(i); + } } - ICache onHeapCache = tieredSpilloverCache.getOnHeapCache(); - ICache diskCache = tieredSpilloverCache.getDiskCache(); List> diskCacheKeys = new ArrayList<>(); - tieredSpilloverCache.getDiskCache().keys().forEach(diskCacheKeys::add); - long actualDiskCacheCount = diskCache.count(); + tieredSpilloverCache.getDiskCacheKeys().forEach(diskCacheKeys::add); + long actualDiskCacheCount = tieredSpilloverCache.diskCacheCount(); long actualTieredCacheCount = tieredSpilloverCache.count(); - assertEquals(onHeapCacheSize, onHeapCache.count()); - assertEquals(numOfItems1 - onHeapCacheSize, actualDiskCacheCount); - assertEquals(0, removalListener.evictionsMetric.count()); - assertEquals(numOfItems1, actualTieredCacheCount); + assertEquals(getItemsForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_ON_HEAP), tieredSpilloverCache.onHeapCacheCount()); + assertEquals(getItemsForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_DISK), actualDiskCacheCount); + assertEquals(expectedEvictions, removalListener.evictionsMetric.count()); + assertEquals( + getItemsForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_ON_HEAP) + getItemsForTier( + tieredSpilloverCache, + TIER_DIMENSION_VALUE_DISK + ), + actualTieredCacheCount + ); for (ICacheKey diskKey : diskCacheKeys) { assertNotNull(tieredSpilloverCache.get(diskKey)); } - tieredSpilloverCache.enableDisableDiskCache(false); // Disable disk cache now. int numOfItems2 = totalSize - numOfItems1; for (int iter = 0; iter < numOfItems2; iter++) { ICacheKey key = getICacheKey(UUID.randomUUID().toString()); + int keySegment = tieredSpilloverCache.getSegmentNumber(key); + TieredSpilloverCache.TieredSpilloverCacheSegment segment = + tieredSpilloverCache.tieredSpilloverCacheSegments[keySegment]; LoadAwareCacheLoader, String> loadAwareCacheLoader = getLoadAwareCacheLoader(); tieredSpilloverCache.computeIfAbsent(key, loadAwareCacheLoader); } @@ -1407,20 +1668,21 @@ public void testGetPutAndInvalidateWithDiskCacheDisabled() throws Exception { assertNull(tieredSpilloverCache.get(diskKey)); // Considering disk cache is disabled, we shouldn't find // these keys. } - assertEquals(onHeapCacheSize, onHeapCache.count()); // Should remain same. - assertEquals(0, diskCache.count() - actualDiskCacheCount); // Considering it is disabled now, shouldn't cache + assertEquals(getItemsForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_ON_HEAP), tieredSpilloverCache.onHeapCacheCount()); // Should + // remain + // same. + assertEquals(0, tieredSpilloverCache.diskCacheCount() - actualDiskCacheCount); // Considering it is disabled now, shouldn't + // cache // any more items. - assertEquals(numOfItems2, removalListener.evictionsMetric.count()); // Considering onHeap cache was already - // full, we should all existing onHeap entries being evicted. - assertEquals(0, tieredSpilloverCache.count() - actualTieredCacheCount); // Count still returns disk cache - // entries count as they haven't been cleared yet. + assertTrue(removalListener.evictionsMetric.count() > 0); + // Considering onHeap cache was already full, we should have some onHeap entries being evicted. long lastKnownTieredCacheEntriesCount = tieredSpilloverCache.count(); // Clear up disk cache keys. for (ICacheKey diskKey : diskCacheKeys) { tieredSpilloverCache.invalidate(diskKey); } - assertEquals(0, diskCache.count()); + assertEquals(0, tieredSpilloverCache.diskCacheCount()); assertEquals(lastKnownTieredCacheEntriesCount - diskCacheKeys.size(), tieredSpilloverCache.count()); tieredSpilloverCache.invalidateAll(); // Clear up all the keys. @@ -1438,13 +1700,14 @@ public void testTiersDoNotTrackStats() throws Exception { removalListener, Settings.builder() .put( - OpenSearchOnHeapCacheSettings.getSettingListForCacheType(CacheType.INDICES_REQUEST_CACHE) - .get(MAXIMUM_SIZE_IN_BYTES_KEY) - .getKey(), + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), onHeapCacheSize * keyValueSize + "b" ) .build(), - 0 + 0, + 1 ); // do some gets to put entries in both tiers @@ -1455,8 +1718,11 @@ public void testTiersDoNotTrackStats() throws Exception { tieredSpilloverCache.computeIfAbsent(key, tieredCacheLoader); } assertNotEquals(new ImmutableCacheStats(0, 0, 0, 0, 0), tieredSpilloverCache.stats().getTotalStats()); - assertEquals(new ImmutableCacheStats(0, 0, 0, 0, 0), tieredSpilloverCache.getOnHeapCache().stats().getTotalStats()); - ImmutableCacheStats diskStats = tieredSpilloverCache.getDiskCache().stats().getTotalStats(); + assertEquals( + new ImmutableCacheStats(0, 0, 0, 0, 0), + tieredSpilloverCache.tieredSpilloverCacheSegments[0].getOnHeapCache().stats().getTotalStats() + ); + ImmutableCacheStats diskStats = tieredSpilloverCache.tieredSpilloverCacheSegments[0].getDiskCache().stats().getTotalStats(); assertEquals(new ImmutableCacheStats(0, 0, 0, 0, 0), diskStats); } @@ -1469,8 +1735,9 @@ public void testTierStatsAddCorrectly() throws Exception { * totalEntries = heapEntries + diskEntries */ - int onHeapCacheSize = randomIntBetween(10, 30); - int diskCacheSize = randomIntBetween(onHeapCacheSize + 1, 100); + int onHeapCacheSize = randomIntBetween(300, 600); + int diskCacheSize = randomIntBetween(700, 1200); + int numberOfSegments = getNumberOfSegments(); int keyValueSize = 50; MockCacheRemovalListener removalListener = new MockCacheRemovalListener<>(); TieredSpilloverCache tieredSpilloverCache = initializeTieredSpilloverCache( @@ -1479,13 +1746,14 @@ public void testTierStatsAddCorrectly() throws Exception { removalListener, Settings.builder() .put( - OpenSearchOnHeapCacheSettings.getSettingListForCacheType(CacheType.INDICES_REQUEST_CACHE) - .get(MAXIMUM_SIZE_IN_BYTES_KEY) - .getKey(), + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), onHeapCacheSize * keyValueSize + "b" ) .build(), - 0 + 0, + numberOfSegments ); List> usedKeys = new ArrayList<>(); @@ -1529,16 +1797,28 @@ public void testTierStatsAddCorrectly() throws Exception { clusterSettings.applySettings( Settings.builder().put(DISK_CACHE_ENABLED_SETTING_MAP.get(CacheType.INDICES_REQUEST_CACHE).getKey(), false).build() ); - - int newMisses = randomIntBetween(10, 30); - for (int i = 0; i < newMisses; i++) { + Map onHeapExpectedSize = getSegmentOnHeapCacheSize( + numberOfSegments, + onHeapCacheSize * keyValueSize, + keyValueSize + ); + int numOfItems = randomIntBetween(10, 30); + int newMisses = 0; + for (int i = 0; i < numOfItems; i++) { LoadAwareCacheLoader, String> tieredCacheLoader = getLoadAwareCacheLoader(); - tieredSpilloverCache.computeIfAbsent(getICacheKey(UUID.randomUUID().toString()), tieredCacheLoader); + ICacheKey iCacheKey = getICacheKey(UUID.randomUUID().toString()); + int keySegment = tieredSpilloverCache.getSegmentNumber(iCacheKey); + if (tieredSpilloverCache.tieredSpilloverCacheSegments[keySegment].getOnHeapCache().count() >= onHeapExpectedSize.get( + keySegment + )) { + newMisses++; + } + tieredSpilloverCache.computeIfAbsent(iCacheKey, tieredCacheLoader); } totalStats = tieredSpilloverCache.stats().getTotalStats(); heapStats = getStatsSnapshotForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_ON_HEAP); - assertEquals(missesBeforeDisablingDiskCache + newMisses, totalStats.getMisses()); + assertEquals(missesBeforeDisablingDiskCache + numOfItems, totalStats.getMisses()); assertEquals(heapTierEvictionsBeforeDisablingDiskCache + newMisses, heapStats.getEvictions()); assertEquals(evictionsBeforeDisablingDiskCache + newMisses, totalStats.getEvictions()); @@ -1546,7 +1826,290 @@ public void testTierStatsAddCorrectly() throws Exception { clusterSettings.applySettings( Settings.builder().put(DISK_CACHE_ENABLED_SETTING_MAP.get(CacheType.INDICES_REQUEST_CACHE).getKey(), true).build() ); + } + + public void testPutForAKeyWhichAlreadyExists() { + int onHeapCacheSize = 1; + int diskCacheSize = 3; + int keyValueSize = 1; + MockCacheRemovalListener removalListener = new MockCacheRemovalListener<>(); + TieredSpilloverCache tieredSpilloverCache = initializeTieredSpilloverCache( + keyValueSize, + diskCacheSize, + removalListener, + Settings.builder() + .put( + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), + onHeapCacheSize * keyValueSize + "b" + ) + .build(), + 0, + 1 + ); + ICacheKey key1 = getICacheKey("key1"); + ICacheKey key2 = getICacheKey("key2"); + tieredSpilloverCache.put(key1, "key1"); // Goes to onHeap cache. + tieredSpilloverCache.put(key2, "key2"); // Goes to onHeap cache. And key1 evicted to disk cache. + List> diskKeys = new ArrayList<>(); + List> onHeapCacheKeys = new ArrayList<>(); + tieredSpilloverCache.tieredSpilloverCacheSegments[0].getDiskCache().keys().forEach(diskKeys::add); + tieredSpilloverCache.tieredSpilloverCacheSegments[0].getOnHeapCache().keys().forEach(onHeapCacheKeys::add); + assertEquals(1, onHeapCacheKeys.size()); + assertEquals(1, diskKeys.size()); + assertTrue(onHeapCacheKeys.contains(key2)); + assertTrue(diskKeys.contains(key1)); + assertEquals("key1", tieredSpilloverCache.get(key1)); + + // Now try to put key1 again onto tiered cache with new value. + tieredSpilloverCache.put(key1, "dummy"); + diskKeys.clear(); + onHeapCacheKeys.clear(); + tieredSpilloverCache.tieredSpilloverCacheSegments[0].getDiskCache().keys().forEach(diskKeys::add); + tieredSpilloverCache.tieredSpilloverCacheSegments[0].getOnHeapCache().keys().forEach(onHeapCacheKeys::add); + assertEquals(1, onHeapCacheKeys.size()); + assertEquals(1, diskKeys.size()); + assertTrue(onHeapCacheKeys.contains(key2)); + assertTrue(diskKeys.contains(key1)); + assertEquals("dummy", tieredSpilloverCache.get(key1)); + } + public void testTieredCacheThrowingExceptionOnPerSegmentSizeBeingZero() { + int onHeapCacheSize = 10; + int diskCacheSize = randomIntBetween(700, 1200); + int keyValueSize = 1; + MockCacheRemovalListener removalListener = new MockCacheRemovalListener<>(); + assertThrows( + IllegalArgumentException.class, + () -> initializeTieredSpilloverCache( + keyValueSize, + diskCacheSize, + removalListener, + Settings.builder() + .put( + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), + onHeapCacheSize * keyValueSize + "b" + ) + .build(), + 0, + 256 + ) + ); + } + + public void testTieredCacheWithZeroNumberOfSegments() { + int onHeapCacheSize = 10; + int diskCacheSize = randomIntBetween(700, 1200); + int keyValueSize = 1; + MockCacheRemovalListener removalListener = new MockCacheRemovalListener<>(); + assertThrows( + ZERO_SEGMENT_COUNT_EXCEPTION_MESSAGE, + IllegalArgumentException.class, + () -> initializeTieredSpilloverCache( + keyValueSize, + diskCacheSize, + removalListener, + Settings.builder() + .put( + OpenSearchOnHeapCacheSettings.getSettingListForCacheType(CacheType.INDICES_REQUEST_CACHE) + .get(MAXIMUM_SIZE_IN_BYTES_KEY) + .getKey(), + onHeapCacheSize * keyValueSize + "b" + ) + .build(), + 0, + 0 + ) + ); + } + + public void testWithInvalidSegmentNumber() throws Exception { + int onHeapCacheSize = 10; + int diskCacheSize = randomIntBetween(700, 1200); + int keyValueSize = 1; + MockCacheRemovalListener removalListener = new MockCacheRemovalListener<>(); + Settings settings = Settings.builder() + .put( + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_NAME.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), + OpenSearchOnHeapCache.OpenSearchOnHeapCacheFactory.NAME + ) + .put( + TieredSpilloverCacheSettings.TIERED_SPILLOVER_DISK_STORE_NAME.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), + MockDiskCache.MockDiskCacheFactory.NAME + ) + .put( + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), + onHeapCacheSize * keyValueSize + "b" + ) + .put( + CacheSettings.getConcreteStoreNameSettingForCacheType(CacheType.INDICES_REQUEST_CACHE).getKey(), + TieredSpilloverCache.TieredSpilloverCacheFactory.TIERED_SPILLOVER_CACHE_NAME + ) + .put(TIERED_SPILLOVER_SEGMENTS.getConcreteSettingForNamespace(CacheType.INDICES_REQUEST_CACHE.getSettingPrefix()).getKey(), 1) + .put(FeatureFlags.PLUGGABLE_CACHE, "true") + .put(TIERED_SPILLOVER_SEGMENTS.getConcreteSettingForNamespace(CacheType.INDICES_REQUEST_CACHE.getSettingPrefix()).getKey(), 3) + .build(); + String storagePath = getStoragePath(settings); + assertThrows( + INVALID_SEGMENT_COUNT_EXCEPTION_MESSAGE, + IllegalArgumentException.class, + () -> new TieredSpilloverCache.TieredSpilloverCacheFactory().create( + new CacheConfig.Builder().setKeyType(String.class) + .setKeyType(String.class) + .setWeigher((k, v) -> keyValueSize) + .setRemovalListener(removalListener) + .setKeySerializer(new StringSerializer()) + .setValueSerializer(new StringSerializer()) + .setSettings(settings) + .setDimensionNames(dimensionNames) + .setCachedResultParser(s -> new CachedQueryResult.PolicyValues(20_000_000L)) // Values will always appear to have taken + // 20_000_000 ns = 20 ms to compute + .setClusterSettings(clusterSettings) + .setStoragePath(storagePath) + .build(), + CacheType.INDICES_REQUEST_CACHE, + Map.of( + OpenSearchOnHeapCache.OpenSearchOnHeapCacheFactory.NAME, + new OpenSearchOnHeapCache.OpenSearchOnHeapCacheFactory(), + MockDiskCache.MockDiskCacheFactory.NAME, + new MockDiskCache.MockDiskCacheFactory(0, randomIntBetween(100, 300), false, keyValueSize) + ) + ) + ); + } + + public void testWithVeryLowDiskCacheSize() throws Exception { + int onHeapCacheSize = 10; + int diskCacheSize = randomIntBetween(700, 1200); + int keyValueSize = 1; + MockCacheRemovalListener removalListener = new MockCacheRemovalListener<>(); + Settings settings = Settings.builder() + .put( + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_NAME.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), + OpenSearchOnHeapCache.OpenSearchOnHeapCacheFactory.NAME + ) + .put( + TieredSpilloverCacheSettings.TIERED_SPILLOVER_DISK_STORE_NAME.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), + MockDiskCache.MockDiskCacheFactory.NAME + ) + .put( + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), + onHeapCacheSize * keyValueSize + "b" + ) + .put( + CacheSettings.getConcreteStoreNameSettingForCacheType(CacheType.INDICES_REQUEST_CACHE).getKey(), + TieredSpilloverCache.TieredSpilloverCacheFactory.TIERED_SPILLOVER_CACHE_NAME + ) + .put( + TieredSpilloverCacheSettings.TIERED_SPILLOVER_DISK_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), + 1L + ) + .put(FeatureFlags.PLUGGABLE_CACHE, "true") + .put(TIERED_SPILLOVER_SEGMENTS.getConcreteSettingForNamespace(CacheType.INDICES_REQUEST_CACHE.getSettingPrefix()).getKey(), 2) + .build(); + String storagePath = getStoragePath(settings); + assertThrows( + IllegalArgumentException.class, + () -> new TieredSpilloverCache.TieredSpilloverCacheFactory().create( + new CacheConfig.Builder().setKeyType(String.class) + .setKeyType(String.class) + .setWeigher((k, v) -> keyValueSize) + .setRemovalListener(removalListener) + .setKeySerializer(new StringSerializer()) + .setValueSerializer(new StringSerializer()) + .setSettings(settings) + .setDimensionNames(dimensionNames) + .setCachedResultParser(s -> new CachedQueryResult.PolicyValues(20_000_000L)) // Values will always appear to have taken + // 20_000_000 ns = 20 ms to compute + .setClusterSettings(clusterSettings) + .setStoragePath(storagePath) + .build(), + CacheType.INDICES_REQUEST_CACHE, + Map.of( + OpenSearchOnHeapCache.OpenSearchOnHeapCacheFactory.NAME, + new OpenSearchOnHeapCache.OpenSearchOnHeapCacheFactory(), + MockDiskCache.MockDiskCacheFactory.NAME, + new MockDiskCache.MockDiskCacheFactory(0, randomIntBetween(100, 300), false, keyValueSize) + ) + ) + ); + } + + public void testTieredCacheDefaultSegmentCount() { + int onHeapCacheSize = 500; + int keyValueSize = 1; + + MockCacheRemovalListener removalListener = new MockCacheRemovalListener<>(); + Settings settings = Settings.builder() + .put( + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_NAME.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), + OpenSearchOnHeapCache.OpenSearchOnHeapCacheFactory.NAME + ) + .put( + TieredSpilloverCacheSettings.TIERED_SPILLOVER_DISK_STORE_NAME.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), + MockDiskCache.MockDiskCacheFactory.NAME + ) + .put( + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), + onHeapCacheSize * keyValueSize + "b" + ) + .put( + CacheSettings.getConcreteStoreNameSettingForCacheType(CacheType.INDICES_REQUEST_CACHE).getKey(), + TieredSpilloverCache.TieredSpilloverCacheFactory.TIERED_SPILLOVER_CACHE_NAME + ) + .put(FeatureFlags.PLUGGABLE_CACHE, "true") + .build(); + String storagePath = getStoragePath(settings); + + TieredSpilloverCache tieredSpilloverCache = (TieredSpilloverCache< + String, + String>) new TieredSpilloverCache.TieredSpilloverCacheFactory().create( + new CacheConfig.Builder().setKeyType(String.class) + .setKeyType(String.class) + .setWeigher((k, v) -> keyValueSize) + .setRemovalListener(removalListener) + .setKeySerializer(new StringSerializer()) + .setValueSerializer(new StringSerializer()) + .setSettings(settings) + .setDimensionNames(dimensionNames) + .setCachedResultParser(s -> new CachedQueryResult.PolicyValues(20_000_000L)) // Values will always appear to have taken + // 20_000_000 ns = 20 ms to compute + .setClusterSettings(clusterSettings) + .setStoragePath(storagePath) + .build(), + CacheType.INDICES_REQUEST_CACHE, + Map.of( + OpenSearchOnHeapCache.OpenSearchOnHeapCacheFactory.NAME, + new OpenSearchOnHeapCache.OpenSearchOnHeapCacheFactory(), + MockDiskCache.MockDiskCacheFactory.NAME, + new MockDiskCache.MockDiskCacheFactory(0, randomIntBetween(100, 300), false, keyValueSize) + ) + ); + assertEquals(TieredSpilloverCacheSettings.defaultSegments(), tieredSpilloverCache.getNumberOfSegments()); + assertTrue(VALID_SEGMENT_COUNT_VALUES.contains(tieredSpilloverCache.getNumberOfSegments())); } private List getMockDimensions() { @@ -1631,7 +2194,10 @@ private TieredSpilloverCache getTieredSpilloverCache( ICache.Factory mockDiskCacheFactory, CacheConfig cacheConfig, List> policies, - RemovalListener, String> removalListener + RemovalListener, String> removalListener, + int numberOfSegments, + long onHeapCacheSizeInBytes, + long diskCacheSize ) { TieredSpilloverCache.Builder builder = new TieredSpilloverCache.Builder().setCacheType( CacheType.INDICES_REQUEST_CACHE @@ -1639,6 +2205,9 @@ private TieredSpilloverCache getTieredSpilloverCache( .setRemovalListener(removalListener) .setOnHeapCacheFactory(onHeapCacheFactory) .setDiskCacheFactory(mockDiskCacheFactory) + .setNumberOfSegments(numberOfSegments) + .setDiskCacheSize(diskCacheSize) + .setOnHeapCacheSizeInBytes(onHeapCacheSizeInBytes) .setCacheConfig(cacheConfig); if (policies != null) { builder.addPolicies(policies); @@ -1654,7 +2223,27 @@ private TieredSpilloverCache initializeTieredSpilloverCache( long diskDeliberateDelay ) { - return intializeTieredSpilloverCache(keyValueSize, diskCacheSize, removalListener, settings, diskDeliberateDelay, null); + return intializeTieredSpilloverCache(keyValueSize, diskCacheSize, removalListener, settings, diskDeliberateDelay, null, 256); + } + + private TieredSpilloverCache initializeTieredSpilloverCache( + int keyValueSize, + int diskCacheSize, + RemovalListener, String> removalListener, + Settings settings, + long diskDeliberateDelay, + int numberOfSegments + + ) { + return intializeTieredSpilloverCache( + keyValueSize, + diskCacheSize, + removalListener, + settings, + diskDeliberateDelay, + null, + numberOfSegments + ); } private TieredSpilloverCache intializeTieredSpilloverCache( @@ -1663,9 +2252,19 @@ private TieredSpilloverCache intializeTieredSpilloverCache( RemovalListener, String> removalListener, Settings settings, long diskDeliberateDelay, - List> policies + List> policies, + int numberOfSegments ) { ICache.Factory onHeapCacheFactory = new OpenSearchOnHeapCache.OpenSearchOnHeapCacheFactory(); + String storagePath; + try (NodeEnvironment environment = newNodeEnvironment(settings)) { + storagePath = environment.nodePaths()[0].path + "/test"; + } catch (IOException e) { + throw new OpenSearchException("Exception occurred", e); + } + long onHeapCacheSizeInBytes = TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).get(settings).getBytes(); CacheConfig cacheConfig = new CacheConfig.Builder().setKeyType(String.class) .setKeyType(String.class) .setWeigher((k, v) -> keyValueSize) @@ -1685,17 +2284,39 @@ private TieredSpilloverCache intializeTieredSpilloverCache( .build() ) .setClusterSettings(clusterSettings) + .setStoragePath(storagePath) .build(); - ICache.Factory mockDiskCacheFactory = new MockDiskCache.MockDiskCacheFactory(diskDeliberateDelay, diskCacheSize, false); + ICache.Factory mockDiskCacheFactory = new MockDiskCache.MockDiskCacheFactory( + diskDeliberateDelay, + diskCacheSize, + false, + keyValueSize + ); - return getTieredSpilloverCache(onHeapCacheFactory, mockDiskCacheFactory, cacheConfig, policies, removalListener); + return getTieredSpilloverCache( + onHeapCacheFactory, + mockDiskCacheFactory, + cacheConfig, + policies, + removalListener, + numberOfSegments, + onHeapCacheSizeInBytes, + diskCacheSize + ); } private CacheConfig getCacheConfig( int keyValueSize, Settings settings, - RemovalListener, String> removalListener + RemovalListener, String> removalListener, + int numberOfSegments ) { + String storagePath; + try (NodeEnvironment environment = newNodeEnvironment(settings)) { + storagePath = environment.nodePaths()[0].path + "/test"; + } catch (IOException e) { + throw new OpenSearchException("Exception occurred", e); + } return new CacheConfig.Builder().setKeyType(String.class) .setKeyType(String.class) .setWeigher((k, v) -> keyValueSize) @@ -1715,6 +2336,8 @@ private CacheConfig getCacheConfig( .build() ) .setClusterSettings(clusterSettings) + .setStoragePath(storagePath) + .setSegmentCount(numberOfSegments) .build(); } @@ -1756,6 +2379,16 @@ private ImmutableCacheStats getStatsSnapshotForTier(TieredSpilloverCache t return snapshot; } + private String getStoragePath(Settings settings) { + String storagePath; + try (NodeEnvironment environment = newNodeEnvironment(settings)) { + storagePath = environment.nodePaths()[0].path + "/test"; + } catch (IOException e) { + throw new OpenSearchException("Exception occurred", e); + } + return storagePath; + } + private void verifyComputeIfAbsentThrowsException( Class expectedException, LoadAwareCacheLoader, String> loader, @@ -1780,7 +2413,8 @@ private void verifyComputeIfAbsentThrowsException( diskCacheSize, removalListener, settings, - 0 + 0, + 1 ); int numberOfSameKeys = randomIntBetween(10, onHeapCacheSize - 1); @@ -1816,6 +2450,29 @@ private void verifyComputeIfAbsentThrowsException( assertEquals(0, tieredSpilloverCache.completableFutureMap.size()); } + private int getNumberOfSegments() { + return randomFrom(2, 4, 8, 16, 32, 64, 128, 256); + } + + private Map getSegmentOnHeapCacheSize(int numberOfSegments, int onHeapCacheSizeInBytes, int keyValueSize) { + Map expectedSegmentOnHeapCacheSize = new HashMap<>(); + for (int i = 0; i < numberOfSegments; i++) { + int perSegmentOnHeapCacheSizeBytes = onHeapCacheSizeInBytes / numberOfSegments; + int perSegmentOnHeapCacheEntries = perSegmentOnHeapCacheSizeBytes / keyValueSize; + expectedSegmentOnHeapCacheSize.put(i, perSegmentOnHeapCacheEntries); + } + return expectedSegmentOnHeapCacheSize; + } + + private Map getSegmentMockDiskCacheSize(int numberOfSegments, int diskCacheSize) { + Map expectedSegmentDiskCacheSize = new HashMap<>(); + for (int i = 0; i < numberOfSegments; i++) { + int perSegmentDiskCacheEntries = diskCacheSize / numberOfSegments; + expectedSegmentDiskCacheSize.put(i, perSegmentDiskCacheEntries); + } + return expectedSegmentDiskCacheSize; + } + private ImmutableCacheStats getTotalStatsSnapshot(TieredSpilloverCache tsc) throws IOException { ImmutableCacheStatsHolder cacheStats = tsc.stats(new String[0]); return cacheStats.getStatsForDimensionValues(List.of()); diff --git a/plugins/cache-ehcache/src/main/java/org/opensearch/cache/EhcacheDiskCacheSettings.java b/plugins/cache-ehcache/src/main/java/org/opensearch/cache/EhcacheDiskCacheSettings.java index d173155080f6a..cbc104f2d0b00 100644 --- a/plugins/cache-ehcache/src/main/java/org/opensearch/cache/EhcacheDiskCacheSettings.java +++ b/plugins/cache-ehcache/src/main/java/org/opensearch/cache/EhcacheDiskCacheSettings.java @@ -88,7 +88,7 @@ public class EhcacheDiskCacheSettings { */ public static final Setting.AffixSetting DISK_CACHE_ALIAS_SETTING = Setting.suffixKeySetting( EhcacheDiskCache.EhcacheDiskCacheFactory.EHCACHE_DISK_CACHE_NAME + ".alias", - (key) -> Setting.simpleString(key, "", NodeScope) + (key) -> Setting.simpleString(key, "ehcache_disk", NodeScope) ); /** diff --git a/plugins/cache-ehcache/src/main/java/org/opensearch/cache/store/disk/EhcacheDiskCache.java b/plugins/cache-ehcache/src/main/java/org/opensearch/cache/store/disk/EhcacheDiskCache.java index 4a95b04de3952..0fa0f8162bb98 100644 --- a/plugins/cache-ehcache/src/main/java/org/opensearch/cache/store/disk/EhcacheDiskCache.java +++ b/plugins/cache-ehcache/src/main/java/org/opensearch/cache/store/disk/EhcacheDiskCache.java @@ -101,14 +101,18 @@ public class EhcacheDiskCache implements ICache { private static final Logger logger = LogManager.getLogger(EhcacheDiskCache.class); // Unique id associated with this cache. - private final static String UNIQUE_ID = UUID.randomUUID().toString(); - private final static String THREAD_POOL_ALIAS_PREFIX = "ehcachePool"; + final static String UNIQUE_ID = UUID.randomUUID().toString(); + final static String THREAD_POOL_ALIAS_PREFIX = "ehcachePool"; + final static int MINIMUM_MAX_SIZE_IN_BYTES = 1024 * 100; // 100KB + final static String CACHE_DATA_CLEANUP_DURING_INITIALIZATION_EXCEPTION = "Failed to delete ehcache disk cache under " + + "path: %s during initialization. Please clean this up manually and restart the process"; + // A Cache manager can create many caches. private final PersistentCacheManager cacheManager; // Disk cache. Using ByteArrayWrapper to compare two byte[] by values rather than the default reference checks @SuppressWarnings({ "rawtypes" }) // We have to use the raw type as there's no way to pass the "generic class" to ehcache - private Cache cache; + private final Cache cache; private final long maxWeightInBytes; private final String storagePath; private final Class keyType; @@ -124,10 +128,6 @@ public class EhcacheDiskCache implements ICache { private final Serializer keySerializer; private final Serializer valueSerializer; - final static int MINIMUM_MAX_SIZE_IN_BYTES = 1024 * 100; // 100KB - final static String CACHE_DATA_CLEANUP_DURING_INITIALIZATION_EXCEPTION = "Failed to delete ehcache disk cache under " - + "path: %s during initialization. Please clean this up manually and restart the process"; - /** * Used in computeIfAbsent to synchronize loading of a given key. This is needed as ehcache doesn't provide a * computeIfAbsent method. @@ -199,6 +199,12 @@ private Cache buildCache(Duration expireAfterAccess // Creating the cache requires permissions specified in plugin-security.policy return AccessController.doPrivileged((PrivilegedAction>) () -> { try { + int segmentCount = (Integer) EhcacheDiskCacheSettings.getSettingListForCacheType(cacheType) + .get(DISK_SEGMENT_KEY) + .get(settings); + if (builder.getNumberOfSegments() > 0) { + segmentCount = builder.getNumberOfSegments(); + } return this.cacheManager.createCache( this.diskCacheAlias, CacheConfigurationBuilder.newCacheConfigurationBuilder( @@ -232,7 +238,7 @@ public Duration getExpiryForUpdate( (Integer) EhcacheDiskCacheSettings.getSettingListForCacheType(cacheType) .get(DISK_WRITE_CONCURRENCY_KEY) .get(settings), - (Integer) EhcacheDiskCacheSettings.getSettingListForCacheType(cacheType).get(DISK_SEGMENT_KEY).get(settings) + segmentCount ) ) .withKeySerializer(new KeySerializerWrapper(keySerializer)) @@ -709,8 +715,19 @@ public ICache create(CacheConfig config, CacheType cacheType, throw new IllegalArgumentException("EhcacheDiskCache requires a value serializer of type Serializer"); } - return new Builder().setStoragePath((String) settingList.get(DISK_STORAGE_PATH_KEY).get(settings)) - .setDiskCacheAlias((String) settingList.get(DISK_CACHE_ALIAS_KEY).get(settings)) + String storagePath = (String) settingList.get(DISK_STORAGE_PATH_KEY).get(settings); + // If we read the storage path directly from the setting, we have to add the segment number at the end. + if (storagePath == null || storagePath.isBlank()) { + // In case storage path is not explicitly set by user, use default path. + // Since this comes from the TSC, it already has the segment number at the end. + storagePath = config.getStoragePath(); + } + String diskCacheAlias = (String) settingList.get(DISK_CACHE_ALIAS_KEY).get(settings); + if (config.getCacheAlias() != null && !config.getCacheAlias().isBlank()) { + diskCacheAlias = config.getCacheAlias(); + } + EhcacheDiskCache.Builder builder = (Builder) new Builder().setStoragePath(storagePath) + .setDiskCacheAlias(diskCacheAlias) .setIsEventListenerModeSync((Boolean) settingList.get(DISK_LISTENER_MODE_SYNC_KEY).get(settings)) .setCacheType(cacheType) .setKeyType((config.getKeyType())) @@ -721,9 +738,21 @@ public ICache create(CacheConfig config, CacheType cacheType, .setWeigher(config.getWeigher()) .setRemovalListener(config.getRemovalListener()) .setExpireAfterAccess((TimeValue) settingList.get(DISK_CACHE_EXPIRE_AFTER_ACCESS_KEY).get(settings)) - .setMaximumWeightInBytes((Long) settingList.get(DISK_MAX_SIZE_IN_BYTES_KEY).get(settings)) - .setSettings(settings) - .build(); + .setSettings(settings); + long maxSizeInBytes = (Long) settingList.get(DISK_MAX_SIZE_IN_BYTES_KEY).get(settings); + // If config value is set, use this instead. + if (config.getMaxSizeInBytes() > 0) { + builder.setMaximumWeightInBytes(config.getMaxSizeInBytes()); + } else { + builder.setMaximumWeightInBytes(maxSizeInBytes); + } + int segmentCount = (Integer) EhcacheDiskCacheSettings.getSettingListForCacheType(cacheType).get(DISK_SEGMENT_KEY).get(settings); + if (config.getSegmentCount() > 0) { + builder.setNumberOfSegments(config.getSegmentCount()); + } else { + builder.setNumberOfSegments(segmentCount); + } + return builder.build(); } @Override diff --git a/plugins/cache-ehcache/src/test/java/org/opensearch/cache/store/disk/EhCacheDiskCacheTests.java b/plugins/cache-ehcache/src/test/java/org/opensearch/cache/store/disk/EhCacheDiskCacheTests.java index 2bc24227bb513..a0d0aa4ec4914 100644 --- a/plugins/cache-ehcache/src/test/java/org/opensearch/cache/store/disk/EhCacheDiskCacheTests.java +++ b/plugins/cache-ehcache/src/test/java/org/opensearch/cache/store/disk/EhCacheDiskCacheTests.java @@ -433,6 +433,74 @@ public String load(ICacheKey key) { } } + public void testComputeIfAbsentConcurrentlyWithMultipleEhcacheDiskCache() throws IOException { + Settings settings = Settings.builder().build(); + MockRemovalListener removalListener = new MockRemovalListener<>(); + List> iCaches = new ArrayList<>(); + int segments = 4; + try (NodeEnvironment env = newNodeEnvironment(settings)) { + ICache.Factory ehcacheFactory = new EhcacheDiskCache.EhcacheDiskCacheFactory(); + for (int i = 1; i <= segments; i++) { + ICache ehcacheTest = ehcacheFactory.create( + new CacheConfig.Builder().setValueType(String.class) + .setKeyType(String.class) + .setRemovalListener(removalListener) + .setKeySerializer(new StringSerializer()) + .setValueSerializer(new StringSerializer()) + .setDimensionNames(List.of(dimensionName)) + .setWeigher(getWeigher()) + .setMaxSizeInBytes(CACHE_SIZE_IN_BYTES * 100) + .setSettings( + Settings.builder() + .put( + EhcacheDiskCacheSettings.getSettingListForCacheType(CacheType.INDICES_REQUEST_CACHE) + .get(DISK_MAX_SIZE_IN_BYTES_KEY) + .getKey(), + CACHE_SIZE_IN_BYTES + ) + .put( + EhcacheDiskCacheSettings.getSettingListForCacheType(CacheType.INDICES_REQUEST_CACHE) + .get(DISK_STORAGE_PATH_KEY) + .getKey(), + env.nodePaths()[0].indicesPath.toString() + "/request_cache/" + i + ) + .put( + EhcacheDiskCacheSettings.getSettingListForCacheType(CacheType.INDICES_REQUEST_CACHE) + .get(DISK_LISTENER_MODE_SYNC_KEY) + .getKey(), + true + ) + .build() + ) + .build(), + CacheType.INDICES_REQUEST_CACHE, + Map.of() + ); + iCaches.add(ehcacheTest); + } + int randomKeys = randomIntBetween(100, 300); + Map, String> keyValueMap = new HashMap<>(); + for (int i = 0; i < randomKeys; i++) { + keyValueMap.put(getICacheKey(UUID.randomUUID().toString()), UUID.randomUUID().toString()); + } + for (Map.Entry, String> entry : keyValueMap.entrySet()) { + ICache ehcacheTest = iCaches.get(entry.getKey().hashCode() & (segments - 1)); + ehcacheTest.put(entry.getKey(), entry.getValue()); + } + for (Map.Entry, String> entry : keyValueMap.entrySet()) { + ICache ehcacheTest = iCaches.get(entry.getKey().hashCode() & (segments - 1)); + String value = ehcacheTest.get(entry.getKey()); + assertEquals(entry.getValue(), value); + } + int count = 0; + for (int i = 0; i < segments; i++) { + count += iCaches.get(i).count(); + iCaches.get(i).close(); + } + assertEquals(randomKeys, count); + } + } + public void testComputeIfAbsentConcurrentlyAndThrowsException() throws Exception { Settings settings = Settings.builder().build(); MockRemovalListener removalListener = new MockRemovalListener<>(); diff --git a/server/src/main/java/org/opensearch/common/cache/Cache.java b/server/src/main/java/org/opensearch/common/cache/Cache.java index caae81e4387b4..da683bfff088d 100644 --- a/server/src/main/java/org/opensearch/common/cache/Cache.java +++ b/server/src/main/java/org/opensearch/common/cache/Cache.java @@ -115,8 +115,20 @@ public class Cache { // the removal callback private RemovalListener removalListener = notification -> {}; - // use CacheBuilder to construct - Cache() {} + private final int numberOfSegments; + private static final int NUMBER_OF_SEGMENTS_DEFAULT = 256; + + Cache(final int numberOfSegments) { + if (numberOfSegments != -1) { + this.numberOfSegments = numberOfSegments; + } else { + this.numberOfSegments = NUMBER_OF_SEGMENTS_DEFAULT; + } + this.segments = new CacheSegment[this.numberOfSegments]; + for (int i = 0; i < this.numberOfSegments; i++) { + this.segments[i] = new CacheSegment<>(); + } + } void setExpireAfterAccessNanos(long expireAfterAccessNanos) { if (expireAfterAccessNanos <= 0) { @@ -366,15 +378,8 @@ void eviction() { } } - public static final int NUMBER_OF_SEGMENTS = 256; @SuppressWarnings("unchecked") - private final CacheSegment[] segments = new CacheSegment[NUMBER_OF_SEGMENTS]; - - { - for (int i = 0; i < segments.length; i++) { - segments[i] = new CacheSegment<>(); - } - } + private final CacheSegment[] segments; Entry head; Entry tail; @@ -382,6 +387,10 @@ void eviction() { // lock protecting mutations to the LRU list private final ReleasableLock lruLock = new ReleasableLock(new ReentrantLock()); + int getNumberOfSegments() { + return numberOfSegments; + } + /** * Returns the value to which the specified key is mapped, or null if this map contains no mapping for the key. * @@ -588,9 +597,9 @@ public void invalidate(K key, V value) { public void invalidateAll() { Entry h; - boolean[] haveSegmentLock = new boolean[NUMBER_OF_SEGMENTS]; + boolean[] haveSegmentLock = new boolean[this.numberOfSegments]; try { - for (int i = 0; i < NUMBER_OF_SEGMENTS; i++) { + for (int i = 0; i < this.numberOfSegments; i++) { segments[i].segmentLock.writeLock().lock(); haveSegmentLock[i] = true; } @@ -607,7 +616,7 @@ public void invalidateAll() { weight = 0; } } finally { - for (int i = NUMBER_OF_SEGMENTS - 1; i >= 0; i--) { + for (int i = this.numberOfSegments - 1; i >= 0; i--) { if (haveSegmentLock[i]) { segments[i].segmentLock.writeLock().unlock(); } @@ -940,6 +949,6 @@ public ToLongBiFunction getWeigher() { } private CacheSegment getCacheSegment(K key) { - return segments[key.hashCode() & 0xff]; + return segments[key.hashCode() & (this.numberOfSegments - 1)]; } } diff --git a/server/src/main/java/org/opensearch/common/cache/CacheBuilder.java b/server/src/main/java/org/opensearch/common/cache/CacheBuilder.java index b6d7979aa4108..8f93e3358a375 100644 --- a/server/src/main/java/org/opensearch/common/cache/CacheBuilder.java +++ b/server/src/main/java/org/opensearch/common/cache/CacheBuilder.java @@ -34,9 +34,13 @@ import org.opensearch.common.unit.TimeValue; +import java.util.Locale; import java.util.Objects; import java.util.function.ToLongBiFunction; +import static org.opensearch.common.cache.settings.CacheSettings.INVALID_SEGMENT_COUNT_EXCEPTION_MESSAGE; +import static org.opensearch.common.cache.settings.CacheSettings.VALID_SEGMENT_COUNT_VALUES; + /** * The cache builder. * @@ -48,6 +52,7 @@ public class CacheBuilder { private long expireAfterWriteNanos = -1; private ToLongBiFunction weigher; private RemovalListener removalListener; + private int numberOfSegments = -1; public static CacheBuilder builder() { return new CacheBuilder<>(); @@ -55,6 +60,14 @@ public static CacheBuilder builder() { private CacheBuilder() {} + public CacheBuilder setNumberOfSegments(int numberOfSegments) { + if (!VALID_SEGMENT_COUNT_VALUES.contains(numberOfSegments)) { + throw new IllegalArgumentException(String.format(Locale.ROOT, INVALID_SEGMENT_COUNT_EXCEPTION_MESSAGE, "Cache")); + } + this.numberOfSegments = numberOfSegments; + return this; + } + public CacheBuilder setMaximumWeight(long maximumWeight) { if (maximumWeight < 0) { throw new IllegalArgumentException("maximumWeight < 0"); @@ -108,7 +121,7 @@ public CacheBuilder removalListener(RemovalListener removalListener) } public Cache build() { - Cache cache = new Cache<>(); + Cache cache = new Cache<>(numberOfSegments); if (maximumWeight != -1) { cache.setMaximumWeight(maximumWeight); } diff --git a/server/src/main/java/org/opensearch/common/cache/settings/CacheSettings.java b/server/src/main/java/org/opensearch/common/cache/settings/CacheSettings.java index 43a047f0f22c6..7a810d92e3d0d 100644 --- a/server/src/main/java/org/opensearch/common/cache/settings/CacheSettings.java +++ b/server/src/main/java/org/opensearch/common/cache/settings/CacheSettings.java @@ -12,12 +12,25 @@ import org.opensearch.common.cache.CacheType; import org.opensearch.common.settings.Setting; +import java.util.List; + /** * Settings related to cache. */ @ExperimentalApi public class CacheSettings { + /** + * Only includes values which is power of 2 as we use bitwise logic: (key AND (segmentCount -1)) to calculate + * segmentNumber which works well only with such values. Stores it in sorted order. + */ + public static final List VALID_SEGMENT_COUNT_VALUES = List.of(1, 2, 4, 8, 16, 32, 64, 128, 256); + + /** + * Exception message for invalid segment number. + */ + public static final String INVALID_SEGMENT_COUNT_EXCEPTION_MESSAGE = "Cache: %s segment count should be power of two up-to 256"; + /** * Used to store cache store name for desired cache types within OpenSearch. * Setting pattern: {cache_type}.store.name diff --git a/server/src/main/java/org/opensearch/common/cache/store/OpenSearchOnHeapCache.java b/server/src/main/java/org/opensearch/common/cache/store/OpenSearchOnHeapCache.java index 569653bec2a3d..571383a9fce6a 100644 --- a/server/src/main/java/org/opensearch/common/cache/store/OpenSearchOnHeapCache.java +++ b/server/src/main/java/org/opensearch/common/cache/store/OpenSearchOnHeapCache.java @@ -54,8 +54,10 @@ public class OpenSearchOnHeapCache implements ICache, RemovalListene private final List dimensionNames; private final ToLongBiFunction, V> weigher; private final boolean statsTrackingEnabled; + private final long maximumWeight; public OpenSearchOnHeapCache(Builder builder) { + this.maximumWeight = builder.getMaxWeightInBytes(); CacheBuilder, V> cacheBuilder = CacheBuilder., V>builder() .setMaximumWeight(builder.getMaxWeightInBytes()) .weigher(builder.getWeigher()) @@ -63,6 +65,9 @@ public OpenSearchOnHeapCache(Builder builder) { if (builder.getExpireAfterAcess() != null) { cacheBuilder.setExpireAfterAccess(builder.getExpireAfterAcess()); } + if (builder.getNumberOfSegments() > 0) { + cacheBuilder.setNumberOfSegments(builder.getNumberOfSegments()); + } cache = cacheBuilder.build(); this.dimensionNames = Objects.requireNonNull(builder.dimensionNames, "Dimension names can't be null"); this.statsTrackingEnabled = builder.getStatsTrackingEnabled(); @@ -75,6 +80,11 @@ public OpenSearchOnHeapCache(Builder builder) { this.weigher = builder.getWeigher(); } + // package private for testing + long getMaximumWeight() { + return this.maximumWeight; + } + @Override public V get(ICacheKey key) { V value = cache.get(key); @@ -174,18 +184,33 @@ public ICache create(CacheConfig config, CacheType cacheType, boolean statsTrackingEnabled = statsTrackingEnabled(config.getSettings(), config.getStatsTrackingEnabled()); ICacheBuilder builder = new Builder().setDimensionNames(config.getDimensionNames()) .setStatsTrackingEnabled(statsTrackingEnabled) - .setMaximumWeightInBytes(((ByteSizeValue) settingList.get(MAXIMUM_SIZE_IN_BYTES_KEY).get(settings)).getBytes()) .setExpireAfterAccess(((TimeValue) settingList.get(EXPIRE_AFTER_ACCESS_KEY).get(settings))) .setWeigher(config.getWeigher()) .setRemovalListener(config.getRemovalListener()); Setting cacheSettingForCacheType = CacheSettings.CACHE_TYPE_STORE_NAME.getConcreteSettingForNamespace( cacheType.getSettingPrefix() ); + long maxSizeInBytes = ((ByteSizeValue) settingList.get(MAXIMUM_SIZE_IN_BYTES_KEY).get(settings)).getBytes(); + + if (config.getMaxSizeInBytes() > 0) { // If this is passed from upstream(like tieredCache), then use this + // instead. + builder.setMaximumWeightInBytes(config.getMaxSizeInBytes()); + } else { + builder.setMaximumWeightInBytes(maxSizeInBytes); + } + if (config.getSegmentCount() > 0) { + builder.setNumberOfSegments(config.getSegmentCount()); + } else { + builder.setNumberOfSegments(-1); // By default it will use 256 segments. + } + String storeName = cacheSettingForCacheType.get(settings); if (!FeatureFlags.PLUGGABLE_CACHE_SETTING.get(settings) || (storeName == null || storeName.isBlank())) { // For backward compatibility as the user intent is to use older settings. builder.setMaximumWeightInBytes(config.getMaxSizeInBytes()); builder.setExpireAfterAccess(config.getExpireAfterAccess()); + builder.setNumberOfSegments(-1); // By default it will use 256 as we don't want to use this setting + // when user wants to use older default onHeap cache settings. } return builder.build(); } diff --git a/server/src/main/java/org/opensearch/common/cache/store/builders/ICacheBuilder.java b/server/src/main/java/org/opensearch/common/cache/store/builders/ICacheBuilder.java index a308d1db88258..d35e22ffdc978 100644 --- a/server/src/main/java/org/opensearch/common/cache/store/builders/ICacheBuilder.java +++ b/server/src/main/java/org/opensearch/common/cache/store/builders/ICacheBuilder.java @@ -39,6 +39,8 @@ public abstract class ICacheBuilder { private boolean statsTrackingEnabled = true; + private int numberOfSegments; + public ICacheBuilder() {} public ICacheBuilder setMaximumWeightInBytes(long sizeInBytes) { @@ -71,6 +73,11 @@ public ICacheBuilder setStatsTrackingEnabled(boolean statsTrackingEnabled) return this; } + public ICacheBuilder setNumberOfSegments(int numberOfSegments) { + this.numberOfSegments = numberOfSegments; + return this; + } + public long getMaxWeightInBytes() { return maxWeightInBytes; } @@ -79,6 +86,10 @@ public TimeValue getExpireAfterAcess() { return expireAfterAcess; } + public int getNumberOfSegments() { + return numberOfSegments; + } + public ToLongBiFunction, V> getWeigher() { return weigher; } diff --git a/server/src/main/java/org/opensearch/common/cache/store/config/CacheConfig.java b/server/src/main/java/org/opensearch/common/cache/store/config/CacheConfig.java index 0c54ac57a9b18..ddad416c251a9 100644 --- a/server/src/main/java/org/opensearch/common/cache/store/config/CacheConfig.java +++ b/server/src/main/java/org/opensearch/common/cache/store/config/CacheConfig.java @@ -70,6 +70,12 @@ public class CacheConfig { private final boolean statsTrackingEnabled; + private final String storagePath; + + private final int segmentCount; + + private final String cacheAlias; + private CacheConfig(Builder builder) { this.keyType = builder.keyType; this.valueType = builder.valueType; @@ -84,6 +90,9 @@ private CacheConfig(Builder builder) { this.expireAfterAccess = builder.expireAfterAccess; this.clusterSettings = builder.clusterSettings; this.statsTrackingEnabled = builder.statsTrackingEnabled; + this.storagePath = builder.storagePath; + this.segmentCount = builder.segmentCount; + this.cacheAlias = builder.cacheAlias; } public Class getKeyType() { @@ -138,6 +147,18 @@ public boolean getStatsTrackingEnabled() { return statsTrackingEnabled; } + public String getStoragePath() { + return storagePath; + } + + public int getSegmentCount() { + return segmentCount; + } + + public String getCacheAlias() { + return cacheAlias; + } + /** * Builder class to build Cache config related parameters. * @param Type of key. @@ -163,6 +184,9 @@ public static class Builder { private TimeValue expireAfterAccess; private ClusterSettings clusterSettings; private boolean statsTrackingEnabled = true; + private String storagePath; + private int segmentCount; + private String cacheAlias; public Builder() {} @@ -231,6 +255,21 @@ public Builder setStatsTrackingEnabled(boolean statsTrackingEnabled) { return this; } + public Builder setStoragePath(String storagePath) { + this.storagePath = storagePath; + return this; + } + + public Builder setSegmentCount(int segmentCount) { + this.segmentCount = segmentCount; + return this; + } + + public Builder setCacheAlias(String cacheAlias) { + this.cacheAlias = cacheAlias; + return this; + } + public CacheConfig build() { return new CacheConfig<>(this); } diff --git a/server/src/main/java/org/opensearch/indices/IndicesRequestCache.java b/server/src/main/java/org/opensearch/indices/IndicesRequestCache.java index 71f8cf5a78ec5..156fe32ff5809 100644 --- a/server/src/main/java/org/opensearch/indices/IndicesRequestCache.java +++ b/server/src/main/java/org/opensearch/indices/IndicesRequestCache.java @@ -68,6 +68,7 @@ import org.opensearch.core.common.io.stream.Writeable; import org.opensearch.core.common.unit.ByteSizeValue; import org.opensearch.core.index.shard.ShardId; +import org.opensearch.env.NodeEnvironment; import org.opensearch.index.shard.IndexShard; import org.opensearch.threadpool.ThreadPool; @@ -167,7 +168,8 @@ public final class IndicesRequestCache implements RemovalListener> cacheEntityFunction, CacheService cacheService, ThreadPool threadPool, - ClusterService clusterService + ClusterService clusterService, + NodeEnvironment nodeEnvironment ) { this.size = INDICES_CACHE_QUERY_SIZE.get(settings); this.expire = INDICES_CACHE_QUERY_EXPIRE.exists(settings) ? INDICES_CACHE_QUERY_EXPIRE.get(settings) : null; @@ -202,6 +204,7 @@ public final class IndicesRequestCache implements RemovalListener keys = new HashSet<>(); Cache cache = CacheBuilder.builder() .setMaximumWeight(numberOfEntries / 2) + .setNumberOfSegments(numberOfSegments) .removalListener(notification -> { keys.remove(notification.getKey()); evictions.incrementAndGet(); @@ -114,11 +116,13 @@ public void testCacheStats() { // check that the evicted entries were evicted in LRU order (first the odds in a batch, then the evens in a batch) // for each batch public void testCacheEvictions() { + int numberOfSegments = randomFrom(1, 2, 4, 8, 16, 64, 128, 256); int maximumWeight = randomIntBetween(1, numberOfEntries); AtomicLong evictions = new AtomicLong(); List evictedKeys = new ArrayList<>(); Cache cache = CacheBuilder.builder() .setMaximumWeight(maximumWeight) + .setNumberOfSegments(numberOfSegments) .removalListener(notification -> { evictions.incrementAndGet(); evictedKeys.add(notification.getKey()); @@ -173,11 +177,13 @@ public void testCacheEvictions() { // cache some entries and exceed the maximum weight, then check that the cache has the expected weight and the // expected evictions occurred public void testWeigher() { + int numberOfSegments = randomFrom(1, 2, 4, 8, 16, 64, 128, 256); int maximumWeight = 2 * numberOfEntries; int weight = randomIntBetween(2, 10); AtomicLong evictions = new AtomicLong(); Cache cache = CacheBuilder.builder() .setMaximumWeight(maximumWeight) + .setNumberOfSegments(numberOfSegments) .weigher((k, v) -> weight) .removalListener(notification -> evictions.incrementAndGet()) .build(); @@ -212,7 +218,8 @@ public void testWeight() { // cache some entries, randomly invalidate some of them, then check that the number of cached entries is correct public void testCount() { - Cache cache = CacheBuilder.builder().build(); + int numberOfSegments = randomFrom(1, 2, 4, 8, 16, 64, 128, 256); + Cache cache = CacheBuilder.builder().setNumberOfSegments(numberOfSegments).build(); int count = 0; for (int i = 0; i < numberOfEntries; i++) { count++; @@ -230,8 +237,9 @@ public void testCount() { // cache some entries, step the clock forward, cache some more entries, step the clock forward and then check that // the first batch of cached entries expired and were removed public void testExpirationAfterAccess() { + int numberOfSegments = randomFrom(1, 2, 4, 8, 16, 64, 128, 256); AtomicLong now = new AtomicLong(); - Cache cache = new Cache() { + Cache cache = new Cache(numberOfSegments) { @Override protected long now() { return now.get(); @@ -267,8 +275,9 @@ protected long now() { } public void testSimpleExpireAfterAccess() { + int numberOfSegments = randomFrom(1, 2, 4, 8, 16, 64, 128, 256); AtomicLong now = new AtomicLong(); - Cache cache = new Cache() { + Cache cache = new Cache(numberOfSegments) { @Override protected long now() { return now.get(); @@ -289,8 +298,9 @@ protected long now() { } public void testExpirationAfterWrite() { + int numberOfSegments = randomFrom(1, 2, 4, 8, 16, 64, 128, 256); AtomicLong now = new AtomicLong(); - Cache cache = new Cache() { + Cache cache = new Cache(numberOfSegments) { @Override protected long now() { return now.get(); @@ -329,8 +339,9 @@ protected long now() { } public void testComputeIfAbsentAfterExpiration() throws ExecutionException { + int numberOfSegments = randomFrom(1, 2, 4, 8, 16, 64, 128, 256); AtomicLong now = new AtomicLong(); - Cache cache = new Cache() { + Cache cache = new Cache(numberOfSegments) { @Override protected long now() { return now.get(); @@ -352,8 +363,10 @@ protected long now() { } public void testComputeIfAbsentDeadlock() throws BrokenBarrierException, InterruptedException { + int numberOfSegments = randomFrom(1, 2, 4, 8, 16, 64, 128, 256); final int numberOfThreads = randomIntBetween(2, 32); final Cache cache = CacheBuilder.builder() + .setNumberOfSegments(numberOfSegments) .setExpireAfterAccess(TimeValue.timeValueNanos(1)) .build(); @@ -386,8 +399,9 @@ public void testComputeIfAbsentDeadlock() throws BrokenBarrierException, Interru // randomly promote some entries, step the clock forward, then check that the promoted entries remain and the // non-promoted entries were removed public void testPromotion() { + int numberOfSegments = randomFrom(1, 2, 4, 8, 16, 64, 128, 256); AtomicLong now = new AtomicLong(); - Cache cache = new Cache() { + Cache cache = new Cache(numberOfSegments) { @Override protected long now() { return now.get(); @@ -420,7 +434,8 @@ protected long now() { // randomly invalidate some cached entries, then check that a lookup for each of those and only those keys is null public void testInvalidate() { - Cache cache = CacheBuilder.builder().build(); + int numberOfSegments = randomFrom(1, 2, 4, 8, 16, 64, 128, 256); + Cache cache = CacheBuilder.builder().setNumberOfSegments(numberOfSegments).build(); for (int i = 0; i < numberOfEntries; i++) { cache.put(i, Integer.toString(i)); } @@ -443,11 +458,15 @@ public void testInvalidate() { // randomly invalidate some cached entries, then check that we receive invalidate notifications for those and only // those entries public void testNotificationOnInvalidate() { + int numberOfSegments = randomFrom(1, 2, 4, 8, 16, 64, 128, 256); Set notifications = new HashSet<>(); - Cache cache = CacheBuilder.builder().removalListener(notification -> { - assertEquals(RemovalReason.INVALIDATED, notification.getRemovalReason()); - notifications.add(notification.getKey()); - }).build(); + Cache cache = CacheBuilder.builder() + .setNumberOfSegments(numberOfSegments) + .removalListener(notification -> { + assertEquals(RemovalReason.INVALIDATED, notification.getRemovalReason()); + notifications.add(notification.getKey()); + }) + .build(); for (int i = 0; i < numberOfEntries; i++) { cache.put(i, Integer.toString(i)); } @@ -491,11 +510,15 @@ public void testInvalidateWithValue() { // randomly invalidate some cached entries, then check that we receive invalidate notifications for those and only // those entries public void testNotificationOnInvalidateWithValue() { + int numberOfSegments = randomFrom(1, 2, 4, 8, 16, 64, 128, 256); Set notifications = new HashSet<>(); - Cache cache = CacheBuilder.builder().removalListener(notification -> { - assertEquals(RemovalReason.INVALIDATED, notification.getRemovalReason()); - notifications.add(notification.getKey()); - }).build(); + Cache cache = CacheBuilder.builder() + .setNumberOfSegments(numberOfSegments) + .removalListener(notification -> { + assertEquals(RemovalReason.INVALIDATED, notification.getRemovalReason()); + notifications.add(notification.getKey()); + }) + .build(); for (int i = 0; i < numberOfEntries; i++) { cache.put(i, Integer.toString(i)); } @@ -607,8 +630,9 @@ public void testNotificationOnReplace() { } public void testComputeIfAbsentLoadsSuccessfully() { + int numberOfSegments = randomFrom(1, 2, 4, 8, 16, 64, 128, 256); Map map = new HashMap<>(); - Cache cache = CacheBuilder.builder().build(); + Cache cache = CacheBuilder.builder().setNumberOfSegments(numberOfSegments).build(); for (int i = 0; i < numberOfEntries; i++) { try { cache.computeIfAbsent(i, k -> { @@ -931,4 +955,12 @@ public void testRemoveUsingValuesIterator() { assertEquals(RemovalReason.INVALIDATED, removalNotifications.get(i).getRemovalReason()); } } + + public void testWithInvalidSegmentNumber() { + assertThrows( + "Number of segments for cache should be a power of two up-to 256", + IllegalArgumentException.class, + () -> CacheBuilder.builder().setMaximumWeight(1000).setNumberOfSegments(21).build() + ); + } } diff --git a/server/src/test/java/org/opensearch/common/cache/store/OpenSearchOnHeapCacheTests.java b/server/src/test/java/org/opensearch/common/cache/store/OpenSearchOnHeapCacheTests.java index f227db6fee2d1..45a7b273eb41e 100644 --- a/server/src/test/java/org/opensearch/common/cache/store/OpenSearchOnHeapCacheTests.java +++ b/server/src/test/java/org/opensearch/common/cache/store/OpenSearchOnHeapCacheTests.java @@ -105,6 +105,37 @@ public void testStatsWithoutPluggableCaches() throws Exception { } } + public void testWithCacheConfigSettings() { + MockRemovalListener listener = new MockRemovalListener<>(); + int maxKeys = between(10, 50); + ICache.Factory onHeapCacheFactory = new OpenSearchOnHeapCache.OpenSearchOnHeapCacheFactory(); + Settings settings = Settings.builder() + .put( + OpenSearchOnHeapCacheSettings.getSettingListForCacheType(CacheType.INDICES_REQUEST_CACHE) + .get(MAXIMUM_SIZE_IN_BYTES_KEY) + .getKey(), + 1000 + "b" // Setting some random value which shouldn't be honored. + ) + .put(FeatureFlags.PLUGGABLE_CACHE, true) + .build(); + + CacheConfig cacheConfig = new CacheConfig.Builder().setKeyType(String.class) + .setValueType(String.class) + .setWeigher((k, v) -> keyValueSize) + .setRemovalListener(listener) + .setSettings(settings) + .setDimensionNames(dimensionNames) + .setMaxSizeInBytes(maxKeys * keyValueSize) // this should get honored + .setStatsTrackingEnabled(true) + .build(); + OpenSearchOnHeapCache onHeapCache = (OpenSearchOnHeapCache) onHeapCacheFactory.create( + cacheConfig, + CacheType.INDICES_REQUEST_CACHE, + null + ); + assertEquals(maxKeys * keyValueSize, onHeapCache.getMaximumWeight()); + } + private void assertZeroStats(ImmutableCacheStatsHolder stats) { assertEquals(new ImmutableCacheStats(0, 0, 0, 0, 0), stats.getTotalStats()); } diff --git a/server/src/test/java/org/opensearch/indices/IndicesRequestCacheTests.java b/server/src/test/java/org/opensearch/indices/IndicesRequestCacheTests.java index 10688de3ab0ae..1a3aece74b3e2 100644 --- a/server/src/test/java/org/opensearch/indices/IndicesRequestCacheTests.java +++ b/server/src/test/java/org/opensearch/indices/IndicesRequestCacheTests.java @@ -73,6 +73,7 @@ import org.opensearch.core.index.shard.ShardId; import org.opensearch.core.xcontent.MediaTypeRegistry; import org.opensearch.core.xcontent.XContentHelper; +import org.opensearch.env.NodeEnvironment; import org.opensearch.index.IndexNotFoundException; import org.opensearch.index.IndexService; import org.opensearch.index.cache.request.RequestCacheStats; @@ -851,15 +852,18 @@ public void testAddingToCleanupKeyToCountMapWorksAppropriatelyWithMultipleThread assertFalse(concurrentModificationExceptionDetected.get()); } - private IndicesRequestCache getIndicesRequestCache(Settings settings) { + private IndicesRequestCache getIndicesRequestCache(Settings settings) throws IOException { IndicesService indicesService = getInstanceFromNode(IndicesService.class); - return new IndicesRequestCache( - settings, - indicesService.indicesRequestCache.cacheEntityLookup, - new CacheModule(new ArrayList<>(), Settings.EMPTY).getCacheService(), - threadPool, - ClusterServiceUtils.createClusterService(threadPool) - ); + try (NodeEnvironment env = newNodeEnvironment(settings)) { + return new IndicesRequestCache( + settings, + indicesService.indicesRequestCache.cacheEntityLookup, + new CacheModule(new ArrayList<>(), Settings.EMPTY).getCacheService(), + threadPool, + ClusterServiceUtils.createClusterService(threadPool), + env + ); + } } private DirectoryReader getReader(IndexWriter writer, ShardId shardId) throws IOException { @@ -913,23 +917,26 @@ public void testClosingIndexWipesStats() throws Exception { .put(INDICES_REQUEST_CACHE_STALENESS_THRESHOLD_SETTING.getKey(), "0.001%") .put(FeatureFlags.PLUGGABLE_CACHE, true) .build(); - cache = new IndicesRequestCache(settings, (shardId -> { - IndexService indexService = null; - try { - indexService = indicesService.indexServiceSafe(shardId.getIndex()); - } catch (IndexNotFoundException ex) { - return Optional.empty(); - } - try { - return Optional.of(new IndicesService.IndexShardCacheEntity(indexService.getShard(shardId.id()))); - } catch (ShardNotFoundException ex) { - return Optional.empty(); - } - }), - new CacheModule(new ArrayList<>(), Settings.EMPTY).getCacheService(), - threadPool, - ClusterServiceUtils.createClusterService(threadPool) - ); + try (NodeEnvironment env = newNodeEnvironment(settings)) { + cache = new IndicesRequestCache(settings, (shardId -> { + IndexService indexService = null; + try { + indexService = indicesService.indexServiceSafe(shardId.getIndex()); + } catch (IndexNotFoundException ex) { + return Optional.empty(); + } + try { + return Optional.of(new IndicesService.IndexShardCacheEntity(indexService.getShard(shardId.id()))); + } catch (ShardNotFoundException ex) { + return Optional.empty(); + } + }), + new CacheModule(new ArrayList<>(), Settings.EMPTY).getCacheService(), + threadPool, + ClusterServiceUtils.createClusterService(threadPool), + env + ); + } writer.addDocument(newDoc(0, "foo")); TermQueryBuilder termQuery = new TermQueryBuilder("id", "0"); @@ -1058,6 +1065,7 @@ public void testEviction() throws Exception { IOUtils.close(reader, secondReader, writer, dir, cache); } indexShard = createIndex("test1").getShard(0); + NodeEnvironment environment = newNodeEnvironment(); IndicesRequestCache cache = new IndicesRequestCache( // TODO: Add wiggle room to max size to allow for overhead of ICacheKey. This can be removed once API PR goes in, as it updates // the old API to account for the ICacheKey overhead. @@ -1065,7 +1073,8 @@ public void testEviction() throws Exception { (shardId -> Optional.of(new IndicesService.IndexShardCacheEntity(indexShard))), new CacheModule(new ArrayList<>(), Settings.EMPTY).getCacheService(), threadPool, - ClusterServiceUtils.createClusterService(threadPool) + ClusterServiceUtils.createClusterService(threadPool), + environment ); dir = newDirectory(); writer = new IndexWriter(dir, newIndexWriterConfig()); @@ -1085,7 +1094,7 @@ public void testEviction() throws Exception { assertEquals("baz", value3.streamInput().readString()); assertEquals(2, cache.count()); assertEquals(1, indexShard.requestCache().stats().getEvictions()); - IOUtils.close(reader, secondReader, thirdReader); + IOUtils.close(reader, secondReader, thirdReader, environment); } public void testClearAllEntityIdentity() throws Exception { From 1e49aa8625d26846ee7bea663bcc1e9296eb219a Mon Sep 17 00:00:00 2001 From: Bharathwaj G Date: Tue, 8 Oct 2024 08:31:30 +0530 Subject: [PATCH 5/7] [Star tree] Add date field rounding support in star tree (#15249) --------- Signed-off-by: Bharathwaj G --- .../index/mapper/StarTreeMapperIT.java | 221 ++++++++-- .../java/org/opensearch/common/Rounding.java | 18 +- .../datacube/DataCubeDateTimeUnit.java | 77 ++++ .../datacube/DateDimension.java | 110 ++++- .../compositeindex/datacube/Dimension.java | 23 ++ .../datacube/DimensionFactory.java | 18 +- .../datacube/NumericDimension.java | 17 + .../datacube/ReadDimension.java | 17 + .../datacube/startree/StarTreeField.java | 32 ++ .../startree/StarTreeIndexSettings.java | 17 +- .../builder/AbstractDocumentsFileManager.java | 8 +- .../startree/builder/BaseStarTreeBuilder.java | 90 ++++- .../builder/OffHeapStarTreeBuilder.java | 40 +- .../builder/OnHeapStarTreeBuilder.java | 35 +- .../builder/SegmentDocsFileManager.java | 12 +- .../builder/StarTreeDocsFileManager.java | 17 +- .../meta/StarTreeMetadataWriter.java | 9 +- .../utils/date/DateTimeUnitAdapter.java | 62 +++ .../utils/date/DateTimeUnitRounding.java | 29 ++ .../startree/utils/date/package-info.java | 14 + .../index/mapper/DateFieldMapper.java | 6 + .../index/mapper/StarTreeMapper.java | 65 ++- .../datacube/startree/DateDimensionTests.java | 365 +++++++++++++++++ .../datacube/startree/StarTreeTestUtils.java | 14 +- .../builder/AbstractStarTreeBuilderTests.java | 262 +++++++++++- .../index/mapper/ObjectMapperTests.java | 14 +- .../index/mapper/StarTreeMapperTests.java | 377 ++++++++++++++---- 27 files changed, 1710 insertions(+), 259 deletions(-) create mode 100644 server/src/main/java/org/opensearch/index/compositeindex/datacube/DataCubeDateTimeUnit.java create mode 100644 server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/date/DateTimeUnitAdapter.java create mode 100644 server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/date/DateTimeUnitRounding.java create mode 100644 server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/date/package-info.java create mode 100644 server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/DateDimensionTests.java diff --git a/server/src/internalClusterTest/java/org/opensearch/index/mapper/StarTreeMapperIT.java b/server/src/internalClusterTest/java/org/opensearch/index/mapper/StarTreeMapperIT.java index 0d69b762ab4f2..5840884f5422a 100644 --- a/server/src/internalClusterTest/java/org/opensearch/index/mapper/StarTreeMapperIT.java +++ b/server/src/internalClusterTest/java/org/opensearch/index/mapper/StarTreeMapperIT.java @@ -12,6 +12,7 @@ import org.opensearch.action.index.IndexResponse; import org.opensearch.action.search.SearchResponse; import org.opensearch.action.support.master.AcknowledgedResponse; +import org.opensearch.common.Rounding; import org.opensearch.common.settings.Settings; import org.opensearch.common.util.FeatureFlags; import org.opensearch.core.common.unit.ByteSizeUnit; @@ -22,9 +23,13 @@ import org.opensearch.index.IndexService; import org.opensearch.index.IndexSettings; import org.opensearch.index.compositeindex.CompositeIndexSettings; +import org.opensearch.index.compositeindex.datacube.DataCubeDateTimeUnit; +import org.opensearch.index.compositeindex.datacube.DateDimension; import org.opensearch.index.compositeindex.datacube.MetricStat; import org.opensearch.index.compositeindex.datacube.startree.StarTreeFieldConfiguration; import org.opensearch.index.compositeindex.datacube.startree.StarTreeIndexSettings; +import org.opensearch.index.compositeindex.datacube.startree.utils.date.DateTimeUnitAdapter; +import org.opensearch.index.compositeindex.datacube.startree.utils.date.DateTimeUnitRounding; import org.opensearch.index.query.QueryBuilders; import org.opensearch.indices.IndicesService; import org.opensearch.search.SearchHit; @@ -58,13 +63,10 @@ private static XContentBuilder createMinimalTestMapping(boolean invalidDim, bool .startObject("startree-1") .field("type", "star_tree") .startObject("config") - .startArray("ordered_dimensions") - .startObject() - .field("name", "numeric_dv_1") - .endObject() - .startObject() - .field("name", "numeric_dv_2") + .startObject("date_dimension") + .field("name", "timestamp") .endObject() + .startArray("ordered_dimensions") .startObject() .field("name", getDim(invalidDim, keywordDim)) .endObject() @@ -85,11 +87,58 @@ private static XContentBuilder createMinimalTestMapping(boolean invalidDim, bool .field("type", "integer") .field("doc_values", true) .endObject() - .startObject("numeric_dv_1") + .startObject("numeric") .field("type", "integer") + .field("doc_values", false) + .endObject() + .startObject("keyword_dv") + .field("type", "keyword") .field("doc_values", true) .endObject() - .startObject("numeric_dv_2") + .startObject("keyword") + .field("type", "keyword") + .field("doc_values", false) + .endObject() + .endObject() + .endObject(); + } catch (IOException e) { + throw new IllegalStateException(e); + } + } + + private static XContentBuilder createDateTestMapping(boolean duplicate) { + try { + return jsonBuilder().startObject() + .startObject("composite") + .startObject("startree-1") + .field("type", "star_tree") + .startObject("config") + .startObject("date_dimension") + .field("name", "timestamp") + .startArray("calendar_intervals") + .value("day") + .value("quarter-hour") + .value(duplicate ? "quarter-hour" : "half-hour") + .endArray() + .endObject() + .startArray("ordered_dimensions") + .startObject() + .field("name", "numeric_dv") + .endObject() + .endArray() + .startArray("metrics") + .startObject() + .field("name", "numeric_dv") + .endObject() + .endArray() + .endObject() + .endObject() + .endObject() + .startObject("properties") + .startObject("timestamp") + .field("type", "date") + .endObject() + .startObject("numeric_dv") .field("type", "integer") .field("doc_values", true) .endObject() @@ -119,10 +168,15 @@ private static XContentBuilder createMaxDimTestMapping() { .startObject("startree-1") .field("type", "star_tree") .startObject("config") - .startArray("ordered_dimensions") - .startObject() - .field("name", "dim4") + .startObject("date_dimension") + .field("name", "timestamp") + .startArray("calendar_intervals") + .value("day") + .value("month") + .value("half-hour") + .endArray() .endObject() + .startArray("ordered_dimensions") .startObject() .field("name", "dim2") .endObject() @@ -167,7 +221,7 @@ private static XContentBuilder createMaxDimTestMapping() { } } - private static XContentBuilder createTestMappingWithoutStarTree(boolean invalidDim, boolean invalidMetric, boolean keywordDim) { + private static XContentBuilder createTestMappingWithoutStarTree() { try { return jsonBuilder().startObject() .startObject("properties") @@ -204,10 +258,10 @@ private static XContentBuilder createUpdateTestMapping(boolean changeDim, boolea .startObject(sameStarTree ? "startree-1" : "startree-2") .field("type", "star_tree") .startObject("config") - .startArray("ordered_dimensions") - .startObject() - .field("name", "numeric_dv1") + .startObject("date_dimension") + .field("name", "timestamp") .endObject() + .startArray("ordered_dimensions") .startObject() .field("name", changeDim ? "numeric_new" : getDim(false, false)) .endObject() @@ -228,10 +282,6 @@ private static XContentBuilder createUpdateTestMapping(boolean changeDim, boolea .field("type", "integer") .field("doc_values", true) .endObject() - .startObject("numeric_dv1") - .field("type", "integer") - .field("doc_values", true) - .endObject() .startObject("numeric") .field("type", "integer") .field("doc_values", false) @@ -263,10 +313,10 @@ private XContentBuilder getMappingWithDuplicateFields(boolean isDuplicateDim, bo .startObject("startree-1") .field("type", "star_tree") .startObject("config") - .startArray("ordered_dimensions") - .startObject() - .field("name", "numeric_dv2") + .startObject("date_dimension") + .field("name", "timestamp") .endObject() + .startArray("ordered_dimensions") .startObject() .field("name", "numeric_dv") .endObject() @@ -293,10 +343,6 @@ private XContentBuilder getMappingWithDuplicateFields(boolean isDuplicateDim, bo .field("type", "integer") .field("doc_values", true) .endObject() - .startObject("numeric_dv2") - .field("type", "integer") - .field("doc_values", true) - .endObject() .startObject("numeric_dv1") .field("type", "integer") .field("doc_values", true) @@ -341,8 +387,92 @@ public void testValidCompositeIndex() { for (CompositeMappedFieldType ft : fts) { assertTrue(ft instanceof StarTreeMapper.StarTreeFieldType); StarTreeMapper.StarTreeFieldType starTreeFieldType = (StarTreeMapper.StarTreeFieldType) ft; - assertEquals("numeric_dv_1", starTreeFieldType.getDimensions().get(0).getField()); - assertEquals("numeric_dv_2", starTreeFieldType.getDimensions().get(1).getField()); + assertEquals("timestamp", starTreeFieldType.getDimensions().get(0).getField()); + assertTrue(starTreeFieldType.getDimensions().get(0) instanceof DateDimension); + DateDimension dateDim = (DateDimension) starTreeFieldType.getDimensions().get(0); + List expectedTimeUnits = Arrays.asList( + new DateTimeUnitAdapter(Rounding.DateTimeUnit.MINUTES_OF_HOUR), + DataCubeDateTimeUnit.HALF_HOUR_OF_DAY + ); + for (int i = 0; i < dateDim.getSortedCalendarIntervals().size(); i++) { + assertEquals(expectedTimeUnits.get(i).shortName(), dateDim.getSortedCalendarIntervals().get(i).shortName()); + } + assertEquals("numeric_dv", starTreeFieldType.getDimensions().get(1).getField()); + assertEquals("numeric_dv", starTreeFieldType.getMetrics().get(0).getField()); + List expectedMetrics = Arrays.asList(MetricStat.VALUE_COUNT, MetricStat.SUM, MetricStat.AVG); + assertEquals(expectedMetrics, starTreeFieldType.getMetrics().get(0).getMetrics()); + assertEquals(10000, starTreeFieldType.getStarTreeConfig().maxLeafDocs()); + assertEquals( + StarTreeFieldConfiguration.StarTreeBuildMode.OFF_HEAP, + starTreeFieldType.getStarTreeConfig().getBuildMode() + ); + assertEquals(Collections.emptySet(), starTreeFieldType.getStarTreeConfig().getSkipStarNodeCreationInDims()); + } + } + } + } + + public void testValidCompositeIndexWithDates() { + prepareCreate(TEST_INDEX).setMapping(createDateTestMapping(false)).setSettings(settings).get(); + Iterable dataNodeInstances = internalCluster().getDataNodeInstances(IndicesService.class); + for (IndicesService service : dataNodeInstances) { + final Index index = resolveIndex("test"); + if (service.hasIndex(index)) { + IndexService indexService = service.indexService(index); + Set fts = indexService.mapperService().getCompositeFieldTypes(); + + for (CompositeMappedFieldType ft : fts) { + assertTrue(ft instanceof StarTreeMapper.StarTreeFieldType); + StarTreeMapper.StarTreeFieldType starTreeFieldType = (StarTreeMapper.StarTreeFieldType) ft; + assertEquals("timestamp", starTreeFieldType.getDimensions().get(0).getField()); + assertTrue(starTreeFieldType.getDimensions().get(0) instanceof DateDimension); + DateDimension dateDim = (DateDimension) starTreeFieldType.getDimensions().get(0); + List expectedTimeUnits = Arrays.asList( + DataCubeDateTimeUnit.QUARTER_HOUR_OF_DAY, + DataCubeDateTimeUnit.HALF_HOUR_OF_DAY, + new DateTimeUnitAdapter(Rounding.DateTimeUnit.DAY_OF_MONTH) + ); + for (int i = 0; i < dateDim.getIntervals().size(); i++) { + assertEquals(expectedTimeUnits.get(i).shortName(), dateDim.getSortedCalendarIntervals().get(i).shortName()); + } + assertEquals("numeric_dv", starTreeFieldType.getDimensions().get(1).getField()); + assertEquals("numeric_dv", starTreeFieldType.getMetrics().get(0).getField()); + List expectedMetrics = Arrays.asList(MetricStat.VALUE_COUNT, MetricStat.SUM, MetricStat.AVG); + assertEquals(expectedMetrics, starTreeFieldType.getMetrics().get(0).getMetrics()); + assertEquals(10000, starTreeFieldType.getStarTreeConfig().maxLeafDocs()); + assertEquals( + StarTreeFieldConfiguration.StarTreeBuildMode.OFF_HEAP, + starTreeFieldType.getStarTreeConfig().getBuildMode() + ); + assertEquals(Collections.emptySet(), starTreeFieldType.getStarTreeConfig().getSkipStarNodeCreationInDims()); + } + } + } + } + + public void testValidCompositeIndexWithDuplicateDates() { + prepareCreate(TEST_INDEX).setMapping(createDateTestMapping(true)).setSettings(settings).get(); + Iterable dataNodeInstances = internalCluster().getDataNodeInstances(IndicesService.class); + for (IndicesService service : dataNodeInstances) { + final Index index = resolveIndex("test"); + if (service.hasIndex(index)) { + IndexService indexService = service.indexService(index); + Set fts = indexService.mapperService().getCompositeFieldTypes(); + + for (CompositeMappedFieldType ft : fts) { + assertTrue(ft instanceof StarTreeMapper.StarTreeFieldType); + StarTreeMapper.StarTreeFieldType starTreeFieldType = (StarTreeMapper.StarTreeFieldType) ft; + assertEquals("timestamp", starTreeFieldType.getDimensions().get(0).getField()); + assertTrue(starTreeFieldType.getDimensions().get(0) instanceof DateDimension); + DateDimension dateDim = (DateDimension) starTreeFieldType.getDimensions().get(0); + List expectedTimeUnits = Arrays.asList( + DataCubeDateTimeUnit.QUARTER_HOUR_OF_DAY, + new DateTimeUnitAdapter(Rounding.DateTimeUnit.DAY_OF_MONTH) + ); + for (int i = 0; i < dateDim.getIntervals().size(); i++) { + assertEquals(expectedTimeUnits.get(i).shortName(), dateDim.getSortedCalendarIntervals().get(i).shortName()); + } + assertEquals("numeric_dv", starTreeFieldType.getDimensions().get(1).getField()); assertEquals(2, starTreeFieldType.getMetrics().size()); assertEquals("numeric_dv", starTreeFieldType.getMetrics().get(0).getField()); @@ -458,7 +588,7 @@ public void testUpdateIndexWithAdditionOfStarTree() { } public void testUpdateIndexWithNewerStarTree() { - prepareCreate(TEST_INDEX).setSettings(settings).setMapping(createTestMappingWithoutStarTree(false, false, false)).get(); + prepareCreate(TEST_INDEX).setSettings(settings).setMapping(createTestMappingWithoutStarTree()).get(); IllegalArgumentException ex = expectThrows( IllegalArgumentException.class, @@ -502,8 +632,18 @@ public void testUpdateIndexWhenMappingIsSame() { for (CompositeMappedFieldType ft : fts) { assertTrue(ft instanceof StarTreeMapper.StarTreeFieldType); StarTreeMapper.StarTreeFieldType starTreeFieldType = (StarTreeMapper.StarTreeFieldType) ft; - assertEquals("numeric_dv_1", starTreeFieldType.getDimensions().get(0).getField()); - assertEquals("numeric_dv_2", starTreeFieldType.getDimensions().get(1).getField()); + assertEquals("timestamp", starTreeFieldType.getDimensions().get(0).getField()); + assertTrue(starTreeFieldType.getDimensions().get(0) instanceof DateDimension); + DateDimension dateDim = (DateDimension) starTreeFieldType.getDimensions().get(0); + List expectedTimeUnits = Arrays.asList( + new DateTimeUnitAdapter(Rounding.DateTimeUnit.MINUTES_OF_HOUR), + DataCubeDateTimeUnit.HALF_HOUR_OF_DAY + ); + for (int i = 0; i < expectedTimeUnits.size(); i++) { + assertEquals(expectedTimeUnits.get(i).shortName(), dateDim.getIntervals().get(i).shortName()); + } + + assertEquals("numeric_dv", starTreeFieldType.getDimensions().get(1).getField()); assertEquals("numeric_dv", starTreeFieldType.getMetrics().get(0).getField()); // Assert default metrics @@ -536,6 +676,7 @@ public void testMaxDimsCompositeIndex() { MapperParsingException.class, () -> prepareCreate(TEST_INDEX).setSettings(settings) .setMapping(createMaxDimTestMapping()) + // Date dimension is considered as one dimension regardless of number of actual calendar intervals .setSettings( Settings.builder() .put(StarTreeIndexSettings.STAR_TREE_MAX_DIMENSIONS_SETTING.getKey(), 2) @@ -569,6 +710,24 @@ public void testMaxMetricsCompositeIndex() { ); } + public void testMaxCalendarIntervalsCompositeIndex() { + MapperParsingException ex = expectThrows( + MapperParsingException.class, + () -> prepareCreate(TEST_INDEX).setMapping(createMaxDimTestMapping()) + .setSettings( + Settings.builder() + .put(StarTreeIndexSettings.STAR_TREE_MAX_DATE_INTERVALS_SETTING.getKey(), 1) + .put(StarTreeIndexSettings.IS_COMPOSITE_INDEX_SETTING.getKey(), true) + .put(IndexSettings.INDEX_TRANSLOG_FLUSH_THRESHOLD_SIZE_SETTING.getKey(), new ByteSizeValue(512, ByteSizeUnit.MB)) + ) + .get() + ); + assertEquals( + "Failed to parse mapping [_doc]: At most [1] calendar intervals are allowed in dimension [timestamp]", + ex.getMessage() + ); + } + public void testUnsupportedDim() { MapperParsingException ex = expectThrows( MapperParsingException.class, diff --git a/server/src/main/java/org/opensearch/common/Rounding.java b/server/src/main/java/org/opensearch/common/Rounding.java index 6f5f1e4328758..12a399635046e 100644 --- a/server/src/main/java/org/opensearch/common/Rounding.java +++ b/server/src/main/java/org/opensearch/common/Rounding.java @@ -95,7 +95,7 @@ public enum DateTimeUnit { WEEK_OF_WEEKYEAR((byte) 1, "week", IsoFields.WEEK_OF_WEEK_BASED_YEAR, true, TimeUnit.DAYS.toMillis(7)) { private final long extraLocalOffsetLookup = TimeUnit.DAYS.toMillis(7); - long roundFloor(long utcMillis) { + public long roundFloor(long utcMillis) { return DateUtils.roundWeekOfWeekYear(utcMillis); } @@ -107,7 +107,7 @@ long extraLocalOffsetLookup() { YEAR_OF_CENTURY((byte) 2, "year", ChronoField.YEAR_OF_ERA, false, 12) { private final long extraLocalOffsetLookup = TimeUnit.DAYS.toMillis(366); - long roundFloor(long utcMillis) { + public long roundFloor(long utcMillis) { return DateUtils.roundYear(utcMillis); } @@ -118,7 +118,7 @@ long extraLocalOffsetLookup() { QUARTER_OF_YEAR((byte) 3, "quarter", IsoFields.QUARTER_OF_YEAR, false, 3) { private final long extraLocalOffsetLookup = TimeUnit.DAYS.toMillis(92); - long roundFloor(long utcMillis) { + public long roundFloor(long utcMillis) { return DateUtils.roundQuarterOfYear(utcMillis); } @@ -129,7 +129,7 @@ long extraLocalOffsetLookup() { MONTH_OF_YEAR((byte) 4, "month", ChronoField.MONTH_OF_YEAR, false, 1) { private final long extraLocalOffsetLookup = TimeUnit.DAYS.toMillis(31); - long roundFloor(long utcMillis) { + public long roundFloor(long utcMillis) { return DateUtils.roundMonthOfYear(utcMillis); } @@ -138,7 +138,7 @@ long extraLocalOffsetLookup() { } }, DAY_OF_MONTH((byte) 5, "day", ChronoField.DAY_OF_MONTH, true, ChronoField.DAY_OF_MONTH.getBaseUnit().getDuration().toMillis()) { - long roundFloor(long utcMillis) { + public long roundFloor(long utcMillis) { return DateUtils.roundFloor(utcMillis, this.ratio); } @@ -147,7 +147,7 @@ long extraLocalOffsetLookup() { } }, HOUR_OF_DAY((byte) 6, "hour", ChronoField.HOUR_OF_DAY, true, ChronoField.HOUR_OF_DAY.getBaseUnit().getDuration().toMillis()) { - long roundFloor(long utcMillis) { + public long roundFloor(long utcMillis) { return DateUtils.roundFloor(utcMillis, ratio); } @@ -162,7 +162,7 @@ long extraLocalOffsetLookup() { true, ChronoField.MINUTE_OF_HOUR.getBaseUnit().getDuration().toMillis() ) { - long roundFloor(long utcMillis) { + public long roundFloor(long utcMillis) { return DateUtils.roundFloor(utcMillis, ratio); } @@ -177,7 +177,7 @@ long extraLocalOffsetLookup() { true, ChronoField.SECOND_OF_MINUTE.getBaseUnit().getDuration().toMillis() ) { - long roundFloor(long utcMillis) { + public long roundFloor(long utcMillis) { return DateUtils.roundFloor(utcMillis, ratio); } @@ -210,7 +210,7 @@ public long extraLocalOffsetLookup() { * @param utcMillis the milliseconds since the epoch * @return the rounded down milliseconds since the epoch */ - abstract long roundFloor(long utcMillis); + public abstract long roundFloor(long utcMillis); /** * When looking up {@link LocalTimeOffset} go this many milliseconds diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/DataCubeDateTimeUnit.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/DataCubeDateTimeUnit.java new file mode 100644 index 0000000000000..56b704103dcae --- /dev/null +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/DataCubeDateTimeUnit.java @@ -0,0 +1,77 @@ +/* + * 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.index.compositeindex.datacube.startree.utils.date.DateTimeUnitRounding; + +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.TimeUnit; + +import static java.util.Collections.unmodifiableMap; + +/** + * Enum representing the extended date time units supported for star tree index as part of index mapping. + * The enum values are: + *

      + *
    • HALF_HOUR_OF_DAY: Represents half hour of day rounding
      • + *
    • QUARTER_HOUR_OF_DAY: Represents quarter hour of day rounding
      • + *
    + *

    + * The enum also provides a static map of date field units to their corresponding ExtendedDateTimeUnit instances. + * + * @see org.opensearch.common.Rounding.DateTimeUnit for more information on the dateTimeUnit enum and rounding logic. + * + * @opensearch.experimental + */ +public enum DataCubeDateTimeUnit implements DateTimeUnitRounding { + HALF_HOUR_OF_DAY("half-hour") { + @Override + public long roundFloor(long utcMillis) { + return utcMillis - (utcMillis % TimeUnit.MINUTES.toMillis(30)); + } + }, + QUARTER_HOUR_OF_DAY("quarter-hour") { + @Override + public long roundFloor(long utcMillis) { + return utcMillis - (utcMillis % TimeUnit.MINUTES.toMillis(15)); + } + }; + + public static final Map DATE_FIELD_UNITS; + static { + Map dateFieldUnits = new HashMap<>(); + dateFieldUnits.put("30m", DataCubeDateTimeUnit.HALF_HOUR_OF_DAY); + dateFieldUnits.put("half-hour", DataCubeDateTimeUnit.HALF_HOUR_OF_DAY); + dateFieldUnits.put("15m", DataCubeDateTimeUnit.QUARTER_HOUR_OF_DAY); + dateFieldUnits.put("quarter-hour", DataCubeDateTimeUnit.QUARTER_HOUR_OF_DAY); + DATE_FIELD_UNITS = unmodifiableMap(dateFieldUnits); + } + + private final String shortName; + + DataCubeDateTimeUnit(String shortName) { + this.shortName = shortName; + } + + /** + * This rounds down the supplied milliseconds since the epoch down to the next unit. In order to retain performance this method + * should be as fast as possible and not try to convert dates to java-time objects if possible + * + * @param utcMillis the milliseconds since the epoch + * @return the rounded down milliseconds since the epoch + */ + @Override + public abstract long roundFloor(long utcMillis); + + @Override + public String shortName() { + return shortName; + } +} diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/DateDimension.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/DateDimension.java index 074016db2aed7..ee6d5b4680c73 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/DateDimension.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/DateDimension.java @@ -10,12 +10,21 @@ import org.opensearch.common.Rounding; import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.common.time.DateUtils; import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.index.compositeindex.datacube.startree.utils.date.DateTimeUnitRounding; import org.opensearch.index.mapper.CompositeDataCubeFieldType; +import org.opensearch.index.mapper.DateFieldMapper; import java.io.IOException; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.HashMap; import java.util.List; +import java.util.Map; import java.util.Objects; +import java.util.function.Consumer; +import java.util.stream.Collectors; /** * Date dimension class @@ -24,27 +33,78 @@ */ @ExperimentalApi public class DateDimension implements Dimension { - private final List calendarIntervals; + private final List calendarIntervals; public static final String CALENDAR_INTERVALS = "calendar_intervals"; public static final String DATE = "date"; private final String field; + private final List sortedCalendarIntervals; + private final DateFieldMapper.Resolution resolution; - public DateDimension(String field, List calendarIntervals) { + public DateDimension(String field, List calendarIntervals, DateFieldMapper.Resolution resolution) { this.field = field; this.calendarIntervals = calendarIntervals; + // Sort from the lowest unit to the highest unit + this.sortedCalendarIntervals = getSortedDateTimeUnits(calendarIntervals); + if (resolution == null) { + this.resolution = DateFieldMapper.Resolution.MILLISECONDS; + } else { + this.resolution = resolution; + } } - public List getIntervals() { + public List getIntervals() { return calendarIntervals; } + public List getSortedCalendarIntervals() { + return sortedCalendarIntervals; + } + + /** + * Sets the dimension values in sorted order in the provided array starting from the given index. + * + * @param val The value to be set + * @param dimSetter Consumer which sets the dimensions + */ + @Override + public void setDimensionValues(final Long val, final Consumer dimSetter) { + for (DateTimeUnitRounding dateTimeUnit : sortedCalendarIntervals) { + if (val == null) { + dimSetter.accept(null); + } else { + Long roundedValue = dateTimeUnit.roundFloor(storedDurationSinceEpoch(val)); + dimSetter.accept(roundedValue); + } + } + } + + /** + * Converts nanoseconds to milliseconds based on the resolution of the field + */ + private long storedDurationSinceEpoch(long nanoSecondsSinceEpoch) { + if (resolution.equals(DateFieldMapper.Resolution.NANOSECONDS)) return DateUtils.toMilliSeconds(nanoSecondsSinceEpoch); + return nanoSecondsSinceEpoch; + } + + /** + * Returns the list of fields that represent the dimension + */ + @Override + public List getSubDimensionNames() { + List fields = new ArrayList<>(calendarIntervals.size()); + for (DateTimeUnitRounding interval : sortedCalendarIntervals) { + fields.add(field + "_" + interval.shortName()); + } + return fields; + } + @Override public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { - builder.startObject(); + builder.startObject("date_dimension"); builder.field(CompositeDataCubeFieldType.NAME, this.getField()); builder.field(CompositeDataCubeFieldType.TYPE, DATE); builder.startArray(CALENDAR_INTERVALS); - for (Rounding.DateTimeUnit interval : calendarIntervals) { + for (DateTimeUnitRounding interval : calendarIntervals) { builder.value(interval.shortName()); } builder.endArray(); @@ -69,4 +129,44 @@ public int hashCode() { public String getField() { return field; } + + @Override + public int getNumSubDimensions() { + return calendarIntervals.size(); + } + + /** + * DateTimeUnit Comparator which tracks dateTimeUnits in sorted order from second unit to year unit + */ + public static class DateTimeUnitComparator implements Comparator { + public static final Map ORDERED_DATE_TIME_UNIT = new HashMap<>(); + + static { + ORDERED_DATE_TIME_UNIT.put(Rounding.DateTimeUnit.SECOND_OF_MINUTE.shortName(), 1); + ORDERED_DATE_TIME_UNIT.put(Rounding.DateTimeUnit.MINUTES_OF_HOUR.shortName(), 2); + ORDERED_DATE_TIME_UNIT.put(DataCubeDateTimeUnit.QUARTER_HOUR_OF_DAY.shortName(), 3); + ORDERED_DATE_TIME_UNIT.put(DataCubeDateTimeUnit.HALF_HOUR_OF_DAY.shortName(), 4); + ORDERED_DATE_TIME_UNIT.put(Rounding.DateTimeUnit.HOUR_OF_DAY.shortName(), 5); + ORDERED_DATE_TIME_UNIT.put(Rounding.DateTimeUnit.DAY_OF_MONTH.shortName(), 6); + ORDERED_DATE_TIME_UNIT.put(Rounding.DateTimeUnit.WEEK_OF_WEEKYEAR.shortName(), 7); + ORDERED_DATE_TIME_UNIT.put(Rounding.DateTimeUnit.MONTH_OF_YEAR.shortName(), 8); + ORDERED_DATE_TIME_UNIT.put(Rounding.DateTimeUnit.QUARTER_OF_YEAR.shortName(), 9); + ORDERED_DATE_TIME_UNIT.put(Rounding.DateTimeUnit.YEAR_OF_CENTURY.shortName(), 10); + } + + @Override + public int compare(DateTimeUnitRounding unit1, DateTimeUnitRounding unit2) { + return Integer.compare( + ORDERED_DATE_TIME_UNIT.getOrDefault(unit1.shortName(), Integer.MAX_VALUE), + ORDERED_DATE_TIME_UNIT.getOrDefault(unit2.shortName(), Integer.MAX_VALUE) + ); + } + } + + /** + * Returns a sorted list of dateTimeUnits based on the DateTimeUnitComparator + */ + public static List getSortedDateTimeUnits(List dateTimeUnits) { + return dateTimeUnits.stream().sorted(new DateTimeUnitComparator()).collect(Collectors.toList()); + } } diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/Dimension.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/Dimension.java index 0151a474579be..cfa8d3a2a8164 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/Dimension.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/Dimension.java @@ -11,6 +11,9 @@ import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.core.xcontent.ToXContent; +import java.util.List; +import java.util.function.Consumer; + /** * Base interface for data-cube dimensions * @@ -18,5 +21,25 @@ */ @ExperimentalApi public interface Dimension extends ToXContent { + String getField(); + + /** + * Returns the number of dimension values that gets added to star tree document + * as part of this dimension + */ + int getNumSubDimensions(); + + /** + * Sets the dimension values with the consumer + * + * @param value The value to be set + * @param dimSetter Consumer which sets the dimensions + */ + void setDimensionValues(final Long value, final Consumer dimSetter); + + /** + * Returns the list of dimension fields that represent the dimension + */ + List getSubDimensionNames(); } 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 index 3c418c68fe8ad..7e72a3f0d9de6 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/DimensionFactory.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/DimensionFactory.java @@ -8,16 +8,19 @@ package org.opensearch.index.compositeindex.datacube; -import org.opensearch.common.Rounding; import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.common.xcontent.support.XContentMapValues; import org.opensearch.index.compositeindex.datacube.startree.StarTreeIndexSettings; +import org.opensearch.index.compositeindex.datacube.startree.utils.date.DateTimeUnitRounding; +import org.opensearch.index.mapper.DateFieldMapper; import org.opensearch.index.mapper.Mapper; import java.util.ArrayList; +import java.util.LinkedHashSet; import java.util.List; import java.util.Locale; import java.util.Map; +import java.util.Set; import java.util.stream.Collectors; import static org.opensearch.index.compositeindex.datacube.DateDimension.CALENDAR_INTERVALS; @@ -70,13 +73,13 @@ private static DateDimension parseAndCreateDateDimension( Map dimensionMap, Mapper.TypeParser.ParserContext c ) { - List calendarIntervals = new ArrayList<>(); + Set calendarIntervals; List intervalStrings = XContentMapValues.extractRawValues(CALENDAR_INTERVALS, dimensionMap) .stream() .map(Object::toString) .collect(Collectors.toList()); if (intervalStrings == null || intervalStrings.isEmpty()) { - calendarIntervals = StarTreeIndexSettings.DEFAULT_DATE_INTERVALS.get(c.getSettings()); + calendarIntervals = new LinkedHashSet<>(StarTreeIndexSettings.DEFAULT_DATE_INTERVALS.get(c.getSettings())); } else { if (intervalStrings.size() > StarTreeIndexSettings.STAR_TREE_MAX_DATE_INTERVALS_SETTING.get(c.getSettings())) { throw new IllegalArgumentException( @@ -88,12 +91,17 @@ private static DateDimension parseAndCreateDateDimension( ) ); } + calendarIntervals = new LinkedHashSet<>(); for (String interval : intervalStrings) { calendarIntervals.add(StarTreeIndexSettings.getTimeUnit(interval)); } - calendarIntervals = new ArrayList<>(calendarIntervals); } dimensionMap.remove(CALENDAR_INTERVALS); - return new DateDimension(name, calendarIntervals); + DateFieldMapper.Resolution resolution = null; + if (c != null && c.mapperService() != null && c.mapperService().fieldType(name) != null) { + resolution = ((DateFieldMapper.DateFieldType) c.mapperService().fieldType(name)).resolution(); + } + + return new DateDimension(name, new ArrayList<>(calendarIntervals), resolution); } } 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 index 9c25ef5b25503..acc14f5f05c68 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/NumericDimension.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/NumericDimension.java @@ -13,7 +13,9 @@ import org.opensearch.index.mapper.CompositeDataCubeFieldType; import java.io.IOException; +import java.util.List; import java.util.Objects; +import java.util.function.Consumer; /** * Composite index numeric dimension class @@ -33,6 +35,21 @@ public String getField() { return field; } + @Override + public int getNumSubDimensions() { + return 1; + } + + @Override + public void setDimensionValues(final Long val, final Consumer dimSetter) { + dimSetter.accept(val); + } + + @Override + public List getSubDimensionNames() { + return List.of(field); + } + @Override public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { builder.startObject(); diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/ReadDimension.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/ReadDimension.java index 4264ec87d2c74..be3667f10b6da 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/ReadDimension.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/ReadDimension.java @@ -12,7 +12,9 @@ import org.opensearch.index.mapper.CompositeDataCubeFieldType; import java.io.IOException; +import java.util.List; import java.util.Objects; +import java.util.function.Consumer; /** * Represents a dimension for reconstructing StarTreeField from file formats during searches and merges. @@ -31,6 +33,21 @@ public String getField() { return field; } + @Override + public int getNumSubDimensions() { + return 1; + } + + @Override + public void setDimensionValues(final Long val, final Consumer dimSetter) { + dimSetter.accept(val); + } + + @Override + public List getSubDimensionNames() { + return List.of(field); + } + @Override public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { builder.startObject(); diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/StarTreeField.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/StarTreeField.java index 922ddcbea4fe2..833bf63c04a18 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/StarTreeField.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/StarTreeField.java @@ -11,10 +11,13 @@ import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.core.xcontent.ToXContent; import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.index.compositeindex.datacube.DateDimension; import org.opensearch.index.compositeindex.datacube.Dimension; import org.opensearch.index.compositeindex.datacube.Metric; +import org.opensearch.index.compositeindex.datacube.MetricStat; import java.io.IOException; +import java.util.ArrayList; import java.util.List; import java.util.Objects; @@ -29,12 +32,24 @@ public class StarTreeField implements ToXContent { private final List dimensionsOrder; private final List metrics; private final StarTreeFieldConfiguration starTreeConfig; + private final List dimensionNames; + private final List metricNames; public StarTreeField(String name, List dimensions, List metrics, StarTreeFieldConfiguration starTreeConfig) { this.name = name; this.dimensionsOrder = dimensions; this.metrics = metrics; this.starTreeConfig = starTreeConfig; + dimensionNames = new ArrayList<>(); + for (Dimension dimension : dimensions) { + dimensionNames.addAll(dimension.getSubDimensionNames()); + } + metricNames = new ArrayList<>(); + for (Metric metric : metrics) { + for (MetricStat metricStat : metric.getMetrics()) { + metricNames.add(metric.getField() + "_" + metricStat.name()); + } + } } public String getName() { @@ -45,6 +60,14 @@ public List getDimensionsOrder() { return dimensionsOrder; } + public List getDimensionNames() { + return dimensionNames; + } + + public List getMetricNames() { + return metricNames; + } + public List getMetrics() { return metrics; } @@ -57,13 +80,22 @@ public StarTreeFieldConfiguration getStarTreeConfig() { public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { builder.startObject(); builder.field("name", name); + DateDimension dateDim = null; if (dimensionsOrder != null && !dimensionsOrder.isEmpty()) { builder.startArray("ordered_dimensions"); for (Dimension dimension : dimensionsOrder) { + // Handle dateDimension for later + if (dimension instanceof DateDimension) { + dateDim = (DateDimension) dimension; + continue; + } dimension.toXContent(builder, params); } builder.endArray(); } + if (dateDim != null) { + dateDim.toXContent(builder, params); + } if (metrics != null && !metrics.isEmpty()) { builder.startArray("metrics"); for (Metric metric : metrics) { diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/StarTreeIndexSettings.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/StarTreeIndexSettings.java index e665831b83d93..72be8ac44cbbc 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/StarTreeIndexSettings.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/StarTreeIndexSettings.java @@ -10,7 +10,10 @@ import org.opensearch.common.Rounding; import org.opensearch.common.settings.Setting; +import org.opensearch.index.compositeindex.datacube.DataCubeDateTimeUnit; import org.opensearch.index.compositeindex.datacube.MetricStat; +import org.opensearch.index.compositeindex.datacube.startree.utils.date.DateTimeUnitAdapter; +import org.opensearch.index.compositeindex.datacube.startree.utils.date.DateTimeUnitRounding; import org.opensearch.search.aggregations.bucket.histogram.DateHistogramAggregationBuilder; import java.util.Arrays; @@ -97,9 +100,9 @@ public class StarTreeIndexSettings { /** * Default intervals for date dimension as part of star tree fields */ - public static final Setting> DEFAULT_DATE_INTERVALS = Setting.listSetting( + public static final Setting> DEFAULT_DATE_INTERVALS = Setting.listSetting( "index.composite_index.star_tree.field.default.date_intervals", - Arrays.asList(Rounding.DateTimeUnit.MINUTES_OF_HOUR.shortName(), Rounding.DateTimeUnit.HOUR_OF_DAY.shortName()), + Arrays.asList(Rounding.DateTimeUnit.MINUTES_OF_HOUR.shortName(), DataCubeDateTimeUnit.HALF_HOUR_OF_DAY.shortName()), StarTreeIndexSettings::getTimeUnit, Setting.Property.IndexScope, Setting.Property.Final @@ -116,11 +119,13 @@ public class StarTreeIndexSettings { Setting.Property.Final ); - public static Rounding.DateTimeUnit getTimeUnit(String expression) { - if (!DateHistogramAggregationBuilder.DATE_FIELD_UNITS.containsKey(expression)) { - throw new IllegalArgumentException("unknown calendar intervals specified in star tree index mapping"); + public static DateTimeUnitRounding getTimeUnit(String expression) { + if (DateHistogramAggregationBuilder.DATE_FIELD_UNITS.containsKey(expression)) { + return new DateTimeUnitAdapter(DateHistogramAggregationBuilder.DATE_FIELD_UNITS.get(expression)); + } else if (DataCubeDateTimeUnit.DATE_FIELD_UNITS.containsKey(expression)) { + return DataCubeDateTimeUnit.DATE_FIELD_UNITS.get(expression); } - return DateHistogramAggregationBuilder.DATE_FIELD_UNITS.get(expression); + throw new IllegalArgumentException("unknown calendar intervals specified in star tree index mapping"); } public static final Setting IS_COMPOSITE_INDEX_SETTING = Setting.boolSetting( diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/AbstractDocumentsFileManager.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/AbstractDocumentsFileManager.java index 327fd26c00608..3f4e302e0a0f2 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/AbstractDocumentsFileManager.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/AbstractDocumentsFileManager.java @@ -43,17 +43,20 @@ public abstract class AbstractDocumentsFileManager implements Closeable { protected final TrackingDirectoryWrapper tmpDirectory; protected final SegmentWriteState state; protected int docSizeInBytes = -1; + protected final int numDimensions; public AbstractDocumentsFileManager( SegmentWriteState state, StarTreeField starTreeField, - List metricAggregatorInfos + List metricAggregatorInfos, + int numDimensions ) { this.starTreeField = starTreeField; this.tmpDirectory = new TrackingDirectoryWrapper(state.directory); this.metricAggregatorInfos = metricAggregatorInfos; this.state = state; numMetrics = metricAggregatorInfos.size(); + this.numDimensions = numDimensions; } private void setDocSizeInBytes(int numBytes) { @@ -124,8 +127,7 @@ protected int writeMetrics(StarTreeDocument starTreeDocument, IndexOutput output * @throws IOException IOException in case of I/O errors */ protected StarTreeDocument readStarTreeDocument(RandomAccessInput input, long offset, boolean isAggregatedDoc) throws IOException { - int dimSize = starTreeField.getDimensionsOrder().size(); - Long[] dimensions = new Long[dimSize]; + Long[] dimensions = new Long[numDimensions]; long initialOffset = offset; offset = readDimensions(dimensions, input, offset); 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 index 5e5814528fcd2..2d4938eeb45b3 100644 --- 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 @@ -18,6 +18,7 @@ import org.apache.lucene.index.SegmentWriteState; import org.apache.lucene.index.SortedNumericDocValues; import org.apache.lucene.index.SortedNumericDocValuesWriterWrapper; +import org.apache.lucene.search.DocIdSetIterator; import org.apache.lucene.store.IndexOutput; import org.apache.lucene.util.Counter; import org.apache.lucene.util.NumericUtils; @@ -52,6 +53,7 @@ import java.util.Set; import java.util.concurrent.atomic.AtomicInteger; +import static org.opensearch.index.compositeindex.CompositeIndexConstants.SEGMENT_DOCS_COUNT; import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils.ALL; import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils.fullyQualifiedFieldNameForStarTreeDimensionsDocValues; import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues; @@ -82,7 +84,7 @@ public abstract class BaseStarTreeBuilder implements StarTreeBuilder { protected int totalSegmentDocs; protected int numStarTreeNodes; protected final int maxLeafDocuments; - + List dimensionsSplitOrder = new ArrayList<>(); protected final InMemoryTreeNode rootNode = getNewNode(); protected final StarTreeField starTreeField; @@ -112,9 +114,12 @@ protected BaseStarTreeBuilder( this.starTreeField = starTreeField; StarTreeFieldConfiguration starTreeFieldSpec = starTreeField.getStarTreeConfig(); - - List dimensionsSplitOrder = starTreeField.getDimensionsOrder(); - this.numDimensions = dimensionsSplitOrder.size(); + int numDims = 0; + for (Dimension dim : starTreeField.getDimensionsOrder()) { + numDims += dim.getNumSubDimensions(); + dimensionsSplitOrder.add(dim); + } + this.numDimensions = numDims; this.skipStarNodeCreationForDimensions = new HashSet<>(); this.totalSegmentDocs = writeState.segmentInfo.maxDoc(); @@ -122,9 +127,12 @@ protected BaseStarTreeBuilder( Set skipStarNodeCreationForDimensions = starTreeFieldSpec.getSkipStarNodeCreationInDims(); - for (int i = 0; i < numDimensions; i++) { + for (int i = 0; i < dimensionsSplitOrder.size(); i++) { if (skipStarNodeCreationForDimensions.contains(dimensionsSplitOrder.get(i).getField())) { - this.skipStarNodeCreationForDimensions.add(i); + // add the dimension indices + for (int dimIndex = 0; dimIndex < dimensionsSplitOrder.get(i).getNumSubDimensions(); dimIndex++) { + this.skipStarNodeCreationForDimensions.add(i + dimIndex); + } } } @@ -224,7 +232,7 @@ public void build( List metricReaders = getMetricReaders(writeState, fieldProducerMap); List dimensionsSplitOrder = starTreeField.getDimensionsOrder(); SequentialDocValuesIterator[] dimensionReaders = new SequentialDocValuesIterator[dimensionsSplitOrder.size()]; - for (int i = 0; i < numDimensions; i++) { + for (int i = 0; i < dimensionReaders.length; i++) { String dimension = dimensionsSplitOrder.get(i).getField(); FieldInfo dimensionFieldInfo = writeState.fieldInfos.fieldInfo(dimension); if (dimensionFieldInfo == null) { @@ -313,19 +321,20 @@ private void createSortedDocValuesIndices(DocValuesConsumer docValuesConsumer, A throws IOException { List dimensionWriters = new ArrayList<>(); List metricWriters = new ArrayList<>(); - FieldInfo[] dimensionFieldInfoList = new FieldInfo[starTreeField.getDimensionsOrder().size()]; + FieldInfo[] dimensionFieldInfoList = new FieldInfo[numDimensions]; FieldInfo[] metricFieldInfoList = new FieldInfo[metricAggregatorInfos.size()]; - for (int i = 0; i < dimensionFieldInfoList.length; i++) { - final FieldInfo fi = getFieldInfo( - fullyQualifiedFieldNameForStarTreeDimensionsDocValues( - starTreeField.getName(), - starTreeField.getDimensionsOrder().get(i).getField() - ), - DocValuesType.SORTED_NUMERIC, - fieldNumberAcrossStarTrees.getAndIncrement() - ); - dimensionFieldInfoList[i] = fi; - dimensionWriters.add(new SortedNumericDocValuesWriterWrapper(fi, Counter.newCounter())); + int dimIndex = 0; + for (Dimension dim : dimensionsSplitOrder) { + for (String name : dim.getSubDimensionNames()) { + final FieldInfo fi = getFieldInfo( + fullyQualifiedFieldNameForStarTreeDimensionsDocValues(starTreeField.getName(), name), + DocValuesType.SORTED_NUMERIC, + fieldNumberAcrossStarTrees.getAndIncrement() + ); + dimensionFieldInfoList[dimIndex] = fi; + dimensionWriters.add(new SortedNumericDocValuesWriterWrapper(fi, Counter.newCounter())); + dimIndex++; + } } for (int i = 0; i < metricAggregatorInfos.size(); i++) { @@ -371,7 +380,7 @@ private void createSortedDocValuesIndices(DocValuesConsumer docValuesConsumer, A } } - addStarTreeDocValueFields(docValuesConsumer, dimensionWriters, dimensionFieldInfoList, starTreeField.getDimensionsOrder().size()); + addStarTreeDocValueFields(docValuesConsumer, dimensionWriters, dimensionFieldInfoList, numDimensions); addStarTreeDocValueFields(docValuesConsumer, metricWriters, metricFieldInfoList, metricAggregatorInfos.size()); } @@ -422,6 +431,36 @@ protected StarTreeDocument getStarTreeDocument( return new StarTreeDocument(dims, metrics); } + /** + * Sets dimensions / metric readers nnd numSegmentDocs + */ + protected void setReadersAndNumSegmentDocs( + SequentialDocValuesIterator[] dimensionReaders, + List metricReaders, + AtomicInteger numSegmentDocs, + StarTreeValues starTreeValues + ) { + List dimensionNames = starTreeValues.getStarTreeField().getDimensionNames(); + for (int i = 0; i < numDimensions; i++) { + dimensionReaders[i] = new SequentialDocValuesIterator(starTreeValues.getDimensionValuesIterator(dimensionNames.get(i))); + } + // get doc id set iterators for metrics + for (Metric metric : starTreeValues.getStarTreeField().getMetrics()) { + for (MetricStat metricStat : metric.getBaseMetrics()) { + String metricFullName = fullyQualifiedFieldNameForStarTreeMetricsDocValues( + starTreeValues.getStarTreeField().getName(), + metric.getField(), + metricStat.getTypeName() + ); + metricReaders.add(new SequentialDocValuesIterator(starTreeValues.getMetricValuesIterator(metricFullName))); + } + } + + numSegmentDocs.set( + Integer.parseInt(starTreeValues.getAttributes().getOrDefault(SEGMENT_DOCS_COUNT, String.valueOf(DocIdSetIterator.NO_MORE_DOCS))) + ); + } + /** * Adds a document to the star-tree. * @@ -500,7 +539,8 @@ protected StarTreeDocument getSegmentStarTreeDocument( */ Long[] getStarTreeDimensionsFromSegment(int currentDocId, SequentialDocValuesIterator[] dimensionReaders) throws IOException { Long[] dimensions = new Long[numDimensions]; - for (int i = 0; i < numDimensions; i++) { + AtomicInteger dimIndex = new AtomicInteger(0); + for (int i = 0; i < dimensionReaders.length; i++) { if (dimensionReaders[i] != null) { try { dimensionReaders[i].nextEntry(currentDocId); @@ -511,11 +551,16 @@ Long[] getStarTreeDimensionsFromSegment(int currentDocId, SequentialDocValuesIte logger.error("unable to read the dimension values from the segment", e); throw new IllegalStateException("unable to read the dimension values from the segment", e); } - dimensions[i] = dimensionReaders[i].value(currentDocId); + dimensionsSplitOrder.get(i).setDimensionValues(dimensionReaders[i].value(currentDocId), value -> { + dimensions[dimIndex.getAndIncrement()] = value; + }); } else { throw new IllegalStateException("dimension readers are empty"); } } + if (dimIndex.get() != numDimensions) { + throw new IllegalStateException("Values are not set for all dimensions"); + } return dimensions; } @@ -685,6 +730,7 @@ private SequentialDocValuesIterator getIteratorForNumericField( * @throws IOException throws an exception if we are unable to add the doc */ private void appendToStarTree(StarTreeDocument starTreeDocument) throws IOException { + appendStarTreeDocument(starTreeDocument); numStarTreeDocs++; } diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/OffHeapStarTreeBuilder.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/OffHeapStarTreeBuilder.java index 11cd9e0c94113..9c39555396397 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/OffHeapStarTreeBuilder.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/OffHeapStarTreeBuilder.java @@ -12,13 +12,9 @@ import org.apache.logging.log4j.Logger; import org.apache.lucene.codecs.DocValuesConsumer; import org.apache.lucene.index.SegmentWriteState; -import org.apache.lucene.search.DocIdSetIterator; import org.apache.lucene.store.IndexOutput; 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.Metric; -import org.opensearch.index.compositeindex.datacube.MetricStat; 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; @@ -36,9 +32,6 @@ import java.util.Objects; import java.util.concurrent.atomic.AtomicInteger; -import static org.opensearch.index.compositeindex.CompositeIndexConstants.SEGMENT_DOCS_COUNT; -import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues; - /** * Off-heap implementation of the star tree builder. * @@ -68,9 +61,9 @@ protected OffHeapStarTreeBuilder( MapperService mapperService ) throws IOException { super(metaOut, dataOut, starTreeField, state, mapperService); - segmentDocumentFileManager = new SegmentDocsFileManager(state, starTreeField, metricAggregatorInfos); + segmentDocumentFileManager = new SegmentDocsFileManager(state, starTreeField, metricAggregatorInfos, numDimensions); try { - starTreeDocumentFileManager = new StarTreeDocsFileManager(state, starTreeField, metricAggregatorInfos); + starTreeDocumentFileManager = new StarTreeDocsFileManager(state, starTreeField, metricAggregatorInfos, numDimensions); } catch (IOException e) { IOUtils.closeWhileHandlingException(segmentDocumentFileManager); throw e; @@ -147,31 +140,12 @@ Iterator mergeStarTrees(List starTreeValuesSub int[] docIds; try { for (StarTreeValues starTreeValues : starTreeValuesSubs) { - List dimensionsSplitOrder = starTreeValues.getStarTreeField().getDimensionsOrder(); - SequentialDocValuesIterator[] dimensionReaders = new SequentialDocValuesIterator[starTreeValues.getStarTreeField() - .getDimensionsOrder() - .size()]; - for (int i = 0; i < dimensionsSplitOrder.size(); i++) { - String dimension = dimensionsSplitOrder.get(i).getField(); - dimensionReaders[i] = new SequentialDocValuesIterator(starTreeValues.getDimensionValuesIterator(dimension)); - } + SequentialDocValuesIterator[] dimensionReaders = new SequentialDocValuesIterator[numDimensions]; List metricReaders = new ArrayList<>(); - // get doc id set iterators for metrics - for (Metric metric : starTreeValues.getStarTreeField().getMetrics()) { - for (MetricStat metricStat : metric.getBaseMetrics()) { - String metricFullName = fullyQualifiedFieldNameForStarTreeMetricsDocValues( - starTreeValues.getStarTreeField().getName(), - metric.getField(), - metricStat.getTypeName() - ); - metricReaders.add(new SequentialDocValuesIterator(starTreeValues.getMetricValuesIterator(metricFullName))); - } - } + AtomicInteger numSegmentDocs = new AtomicInteger(); + setReadersAndNumSegmentDocs(dimensionReaders, metricReaders, numSegmentDocs, starTreeValues); int currentDocId = 0; - int numSegmentDocs = Integer.parseInt( - starTreeValues.getAttributes().getOrDefault(SEGMENT_DOCS_COUNT, String.valueOf(DocIdSetIterator.NO_MORE_DOCS)) - ); - while (currentDocId < numSegmentDocs) { + while (currentDocId < numSegmentDocs.get()) { StarTreeDocument starTreeDocument = getStarTreeDocument(currentDocId, dimensionReaders, metricReaders); segmentDocumentFileManager.writeStarTreeDocument(starTreeDocument, true); numDocs++; @@ -317,7 +291,7 @@ public Iterator generateStarTreeDocumentsForStarNode(int start int docId = 1; private boolean hasSameDimensions(StarTreeDocument document1, StarTreeDocument document2) { - for (int i = dimensionId + 1; i < starTreeField.getDimensionsOrder().size(); i++) { + for (int i = dimensionId + 1; i < numDimensions; i++) { if (!Objects.equals(document1.dimensions[i], document2.dimensions[i])) { return false; } diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/OnHeapStarTreeBuilder.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/OnHeapStarTreeBuilder.java index 13c6d03c4dc3d..07142fc5c8be7 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/OnHeapStarTreeBuilder.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/OnHeapStarTreeBuilder.java @@ -9,12 +9,8 @@ import org.apache.lucene.codecs.DocValuesConsumer; import org.apache.lucene.index.SegmentWriteState; -import org.apache.lucene.search.DocIdSetIterator; import org.apache.lucene.store.IndexOutput; import org.opensearch.common.annotation.ExperimentalApi; -import org.opensearch.index.compositeindex.datacube.Dimension; -import org.opensearch.index.compositeindex.datacube.Metric; -import org.opensearch.index.compositeindex.datacube.MetricStat; 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; @@ -29,9 +25,6 @@ import java.util.Objects; import java.util.concurrent.atomic.AtomicInteger; -import static org.opensearch.index.compositeindex.CompositeIndexConstants.SEGMENT_DOCS_COUNT; -import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues; - /** * On heap single tree builder * @@ -133,33 +126,13 @@ Iterator mergeStarTrees(List starTreeValuesSub StarTreeDocument[] getSegmentsStarTreeDocuments(List starTreeValuesSubs) throws IOException { List starTreeDocuments = new ArrayList<>(); for (StarTreeValues starTreeValues : starTreeValuesSubs) { - List dimensionsSplitOrder = starTreeValues.getStarTreeField().getDimensionsOrder(); - SequentialDocValuesIterator[] dimensionReaders = new SequentialDocValuesIterator[dimensionsSplitOrder.size()]; - - for (int i = 0; i < dimensionsSplitOrder.size(); i++) { - String dimension = dimensionsSplitOrder.get(i).getField(); - dimensionReaders[i] = new SequentialDocValuesIterator(starTreeValues.getDimensionValuesIterator(dimension)); - } + SequentialDocValuesIterator[] dimensionReaders = new SequentialDocValuesIterator[numDimensions]; List metricReaders = new ArrayList<>(); - // get doc id set iterators for metrics - for (Metric metric : starTreeValues.getStarTreeField().getMetrics()) { - for (MetricStat metricStat : metric.getBaseMetrics()) { - String metricFullName = fullyQualifiedFieldNameForStarTreeMetricsDocValues( - starTreeValues.getStarTreeField().getName(), - metric.getField(), - metricStat.getTypeName() - ); - metricReaders.add(new SequentialDocValuesIterator(starTreeValues.getMetricValuesIterator(metricFullName))); - - } - } - + AtomicInteger numSegmentDocs = new AtomicInteger(); + setReadersAndNumSegmentDocs(dimensionReaders, metricReaders, numSegmentDocs, starTreeValues); int currentDocId = 0; - int numSegmentDocs = Integer.parseInt( - starTreeValues.getAttributes().getOrDefault(SEGMENT_DOCS_COUNT, String.valueOf(DocIdSetIterator.NO_MORE_DOCS)) - ); - while (currentDocId < numSegmentDocs) { + while (currentDocId < numSegmentDocs.get()) { starTreeDocuments.add(getStarTreeDocument(currentDocId, dimensionReaders, metricReaders)); currentDocId++; } diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/SegmentDocsFileManager.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/SegmentDocsFileManager.java index fe94df57d9535..363e02b52d5e3 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/SegmentDocsFileManager.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/SegmentDocsFileManager.java @@ -40,9 +40,13 @@ public class SegmentDocsFileManager extends AbstractDocumentsFileManager impleme private RandomAccessInput segmentRandomInput; final IndexOutput segmentDocsFileOutput; - public SegmentDocsFileManager(SegmentWriteState state, StarTreeField starTreeField, List metricAggregatorInfos) - throws IOException { - super(state, starTreeField, metricAggregatorInfos); + public SegmentDocsFileManager( + SegmentWriteState state, + StarTreeField starTreeField, + List metricAggregatorInfos, + int numDimensions + ) throws IOException { + super(state, starTreeField, metricAggregatorInfos, numDimensions); try { segmentDocsFileOutput = tmpDirectory.createTempOutput(SEGMENT_DOC_FILE_NAME, state.segmentSuffix, state.context); } catch (IOException e) { @@ -78,7 +82,7 @@ public StarTreeDocument readStarTreeDocument(int docId, boolean isAggregatedDoc) @Override public Long[] readDimensions(int docId) throws IOException { maybeInitializeSegmentInput(); - Long[] dims = new Long[starTreeField.getDimensionsOrder().size()]; + Long[] dims = new Long[numDimensions]; readDimensions(dims, segmentRandomInput, (long) docId * docSizeInBytes); return dims; } diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeDocsFileManager.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeDocsFileManager.java index 15ed153249243..7e920b912731d 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeDocsFileManager.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeDocsFileManager.java @@ -60,18 +60,23 @@ public class StarTreeDocsFileManager extends AbstractDocumentsFileManager implem private final int fileCountMergeThreshold; private int numStarTreeDocs = 0; - public StarTreeDocsFileManager(SegmentWriteState state, StarTreeField starTreeField, List metricAggregatorInfos) - throws IOException { - this(state, starTreeField, metricAggregatorInfos, DEFAULT_FILE_COUNT_MERGE_THRESHOLD); + public StarTreeDocsFileManager( + SegmentWriteState state, + StarTreeField starTreeField, + List metricAggregatorInfos, + int numDimensions + ) throws IOException { + this(state, starTreeField, metricAggregatorInfos, DEFAULT_FILE_COUNT_MERGE_THRESHOLD, numDimensions); } public StarTreeDocsFileManager( SegmentWriteState state, StarTreeField starTreeField, List metricAggregatorInfos, - int fileCountThreshold + int fileCountThreshold, + int numDimensions ) throws IOException { - super(state, starTreeField, metricAggregatorInfos); + super(state, starTreeField, metricAggregatorInfos, numDimensions); fileToEndDocIdMap = new LinkedHashMap<>(); try { starTreeDocsFileOutput = createStarTreeDocumentsFileOutput(); @@ -126,7 +131,7 @@ public Long getDimensionValue(int docId, int dimensionId) throws IOException { @Override public Long[] readDimensions(int docId) throws IOException { ensureDocumentReadable(docId); - Long[] dims = new Long[starTreeField.getDimensionsOrder().size()]; + Long[] dims = new Long[numDimensions]; readDimensions(dims, starTreeDocsFileRandomInput, getOffset(docId)); return dims; } diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/meta/StarTreeMetadataWriter.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/meta/StarTreeMetadataWriter.java index 1c04350e25047..42e6f3c59866a 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/meta/StarTreeMetadataWriter.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/fileformats/meta/StarTreeMetadataWriter.java @@ -11,7 +11,6 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.apache.lucene.store.IndexOutput; -import org.opensearch.index.compositeindex.datacube.Dimension; import org.opensearch.index.compositeindex.datacube.startree.StarTreeField; import org.opensearch.index.compositeindex.datacube.startree.aggregators.MetricAggregatorInfo; import org.opensearch.index.mapper.CompositeMappedFieldType; @@ -128,13 +127,11 @@ private static void writeMeta( metaOut.writeVInt(numNodes); // number of dimensions - // TODO: Revisit the number of dimensions for timestamps (as we will split timestamp into min, hour, etc.) - metaOut.writeVInt(starTreeField.getDimensionsOrder().size()); + metaOut.writeVInt(starTreeField.getDimensionNames().size()); // dimensions - // TODO: Add sub-dimensions for timestamps (as we will split timestamp into min, hour, etc.) - for (Dimension dimension : starTreeField.getDimensionsOrder()) { - metaOut.writeString(dimension.getField()); + for (String dim : starTreeField.getDimensionNames()) { + metaOut.writeString(dim); } // number of metrics diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/date/DateTimeUnitAdapter.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/date/DateTimeUnitAdapter.java new file mode 100644 index 0000000000000..b18ac33a8f657 --- /dev/null +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/date/DateTimeUnitAdapter.java @@ -0,0 +1,62 @@ +/* + * 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.startree.utils.date; + +import org.opensearch.common.Rounding; + +import java.util.Objects; + +/** + * Adapter class to convert {@link Rounding.DateTimeUnit} to {@link DateTimeUnitRounding} + * + * @opensearch.experimental + */ +public class DateTimeUnitAdapter implements DateTimeUnitRounding { + private final Rounding.DateTimeUnit dateTimeUnit; + + public DateTimeUnitAdapter(Rounding.DateTimeUnit dateTimeUnit) { + this.dateTimeUnit = dateTimeUnit; + } + + @Override + public String shortName() { + return dateTimeUnit.shortName(); + } + + @Override + public long roundFloor(long utcMillis) { + return dateTimeUnit.roundFloor(utcMillis); + } + + /** + * Checks if this DateTimeUnitRounding is equal to another object. + * Two DateTimeUnitRounding instances are considered equal if they have the same short name. + * + * @param obj The object to compare with + * @return true if the objects are equal, false otherwise + */ + @Override + public boolean equals(Object obj) { + if (this == obj) return true; + if (!(obj instanceof DateTimeUnitRounding)) return false; + DateTimeUnitRounding other = (DateTimeUnitRounding) obj; + return Objects.equals(shortName(), other.shortName()); + } + + /** + * Returns a hash code value for the object. + * This method is implemented to be consistent with equals. + * + * @return a hash code value for this object + */ + @Override + public int hashCode() { + return Objects.hash(shortName()); + } +} diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/date/DateTimeUnitRounding.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/date/DateTimeUnitRounding.java new file mode 100644 index 0000000000000..93212a2424e46 --- /dev/null +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/date/DateTimeUnitRounding.java @@ -0,0 +1,29 @@ +/* + * 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.startree.utils.date; + +import org.opensearch.common.annotation.ExperimentalApi; + +/** + * Interface for rounding time units in starTree + * + * @opensearch.experimental + */ +@ExperimentalApi +public interface DateTimeUnitRounding { + /** + * Returns the short name of the time unit + */ + String shortName(); + + /** + * rounds down the given utcMillis to the nearest unit of time + */ + long roundFloor(long utcMillis); +} diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/date/package-info.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/date/package-info.java new file mode 100644 index 0000000000000..bb285c894e3e0 --- /dev/null +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/date/package-info.java @@ -0,0 +1,14 @@ +/* + * 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. + */ + +/** + * Classes and utilities for working with dateTimeUnits + * + * @opensearch.experimental + */ +package org.opensearch.index.compositeindex.datacube.startree.utils.date; diff --git a/server/src/main/java/org/opensearch/index/mapper/DateFieldMapper.java b/server/src/main/java/org/opensearch/index/mapper/DateFieldMapper.java index bc50f507010c4..4698440f08036 100644 --- a/server/src/main/java/org/opensearch/index/mapper/DateFieldMapper.java +++ b/server/src/main/java/org/opensearch/index/mapper/DateFieldMapper.java @@ -54,6 +54,7 @@ import org.opensearch.common.unit.TimeValue; import org.opensearch.common.util.FeatureFlags; import org.opensearch.common.util.LocaleUtils; +import org.opensearch.index.compositeindex.datacube.DimensionType; import org.opensearch.index.fielddata.IndexFieldData; import org.opensearch.index.fielddata.IndexNumericFieldData.NumericType; import org.opensearch.index.fielddata.plain.SortedNumericIndexFieldData; @@ -76,6 +77,7 @@ import java.util.List; import java.util.Locale; import java.util.Map; +import java.util.Optional; import java.util.function.BiFunction; import java.util.function.Function; import java.util.function.LongSupplier; @@ -350,6 +352,10 @@ public DateFieldMapper build(BuilderContext context) { return new DateFieldMapper(name, ft, multiFieldsBuilder.build(this, context), copyTo.build(), nullTimestamp, resolution, this); } + @Override + public Optional getSupportedDataCubeDimensionType() { + return Optional.of(DimensionType.DATE); + } } public static final TypeParser MILLIS_PARSER = new TypeParser((n, c) -> { diff --git a/server/src/main/java/org/opensearch/index/mapper/StarTreeMapper.java b/server/src/main/java/org/opensearch/index/mapper/StarTreeMapper.java index 52dab17e0b0bb..40f05a8b76755 100644 --- a/server/src/main/java/org/opensearch/index/mapper/StarTreeMapper.java +++ b/server/src/main/java/org/opensearch/index/mapper/StarTreeMapper.java @@ -11,6 +11,7 @@ import org.apache.lucene.search.Query; import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.common.xcontent.support.XContentMapValues; +import org.opensearch.index.compositeindex.datacube.DateDimension; import org.opensearch.index.compositeindex.datacube.Dimension; import org.opensearch.index.compositeindex.datacube.DimensionFactory; import org.opensearch.index.compositeindex.datacube.Metric; @@ -44,8 +45,8 @@ public class StarTreeMapper extends ParametrizedFieldMapper { public static final String CONFIG = "config"; public static final String MAX_LEAF_DOCS = "max_leaf_docs"; public static final String SKIP_STAR_NODE_IN_DIMS = "skip_star_node_creation_for_dimensions"; - public static final String BUILD_MODE = "build_mode"; public static final String ORDERED_DIMENSIONS = "ordered_dimensions"; + public static final String DATE_DIMENSION = "date_dimension"; public static final String METRICS = "metrics"; public static final String STATS = "stats"; @@ -84,6 +85,7 @@ public static class Builder extends ParametrizedFieldMapper.Builder { StarTreeFieldConfiguration.StarTreeBuildMode buildMode = StarTreeFieldConfiguration.StarTreeBuildMode.OFF_HEAP; List dimensions = buildDimensions(name, paramMap, context); + paramMap.remove(DATE_DIMENSION); paramMap.remove(ORDERED_DIMENSIONS); List metrics = buildMetrics(name, paramMap, context); paramMap.remove(METRICS); @@ -119,16 +121,21 @@ public static class Builder extends ParametrizedFieldMapper.Builder { */ @SuppressWarnings("unchecked") private List buildDimensions(String fieldName, Map map, Mapper.TypeParser.ParserContext context) { + List dimensions = new LinkedList<>(); + DateDimension dateDim = buildDateDimension(fieldName, map, context); + if (dateDim != null) { + dimensions.add(dateDim); + } Object dims = XContentMapValues.extractValue("ordered_dimensions", map); if (dims == null) { throw new IllegalArgumentException( String.format(Locale.ROOT, "ordered_dimensions is required for star tree field [%s]", fieldName) ); } - List dimensions = new LinkedList<>(); + if (dims instanceof List) { - List dimList = (List) dims; - if (dimList.size() > context.getSettings() + List orderedDimensionsList = (List) dims; + if (orderedDimensionsList.size() + dimensions.size() > context.getSettings() .getAsInt( StarTreeIndexSettings.STAR_TREE_MAX_DIMENSIONS_SETTING.getKey(), StarTreeIndexSettings.STAR_TREE_MAX_DIMENSIONS_DEFAULT @@ -146,13 +153,13 @@ private List buildDimensions(String fieldName, Map ma ) ); } - if (dimList.size() < 2) { + if (dimensions.size() + orderedDimensionsList.size() < 2) { throw new IllegalArgumentException( String.format(Locale.ROOT, "Atleast two dimensions are required to build star tree index field [%s]", fieldName) ); } Set dimensionFieldNames = new HashSet<>(); - for (Object dim : dimList) { + for (Object dim : orderedDimensionsList) { Dimension dimension = getDimension(fieldName, dim, context); if (dimensionFieldNames.add(dimension.getField()) == false) { throw new IllegalArgumentException( @@ -174,6 +181,52 @@ private List buildDimensions(String fieldName, Map ma return dimensions; } + private DateDimension buildDateDimension(String fieldName, Map map, Mapper.TypeParser.ParserContext context) { + Object dims = XContentMapValues.extractValue("date_dimension", map); + if (dims == null) { + return null; + } + return getDateDimension(fieldName, dims, context); + } + + /** + * Get dimension based on mapping + */ + @SuppressWarnings("unchecked") + private DateDimension getDateDimension(String fieldName, Object dimensionMapping, Mapper.TypeParser.ParserContext context) { + DateDimension dimension; + Map dimensionMap = (Map) dimensionMapping; + String name = (String) XContentMapValues.extractValue(CompositeDataCubeFieldType.NAME, dimensionMap); + dimensionMap.remove(CompositeDataCubeFieldType.NAME); + if (this.objbuilder == null || this.objbuilder.mappersBuilders == null) { + String type = (String) XContentMapValues.extractValue(CompositeDataCubeFieldType.TYPE, dimensionMap); + dimensionMap.remove(CompositeDataCubeFieldType.TYPE); + if (type == null || type.equals(DateDimension.DATE) == false) { + throw new MapperParsingException( + String.format(Locale.ROOT, "unable to parse date dimension for star tree field [%s]", fieldName) + ); + } + return (DateDimension) DimensionFactory.parseAndCreateDimension(name, type, dimensionMap, context); + } else { + Optional dimBuilder = findMapperBuilderByName(name, this.objbuilder.mappersBuilders); + if (dimBuilder.isEmpty()) { + throw new IllegalArgumentException(String.format(Locale.ROOT, "unknown date dimension field [%s]", name)); + } + if (dimBuilder.get() instanceof DateFieldMapper.Builder == false) { + throw new IllegalArgumentException( + String.format(Locale.ROOT, "date_dimension [%s] should be of type date for star tree field [%s]", name, fieldName) + ); + } + dimension = (DateDimension) DimensionFactory.parseAndCreateDimension(name, dimBuilder.get(), dimensionMap, context); + } + DocumentMapperParser.checkNoRemainingFields( + dimensionMap, + context.indexVersionCreated(), + "Star tree mapping definition has unsupported parameters: " + ); + return dimension; + } + /** * Get dimension based on mapping */ diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/DateDimensionTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/DateDimensionTests.java new file mode 100644 index 0000000000000..0aa563ed0cafd --- /dev/null +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/DateDimensionTests.java @@ -0,0 +1,365 @@ +/* + * 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.startree; + +import org.opensearch.common.Rounding; +import org.opensearch.index.compositeindex.datacube.DataCubeDateTimeUnit; +import org.opensearch.index.compositeindex.datacube.DateDimension; +import org.opensearch.index.compositeindex.datacube.startree.utils.date.DateTimeUnitAdapter; +import org.opensearch.index.compositeindex.datacube.startree.utils.date.DateTimeUnitRounding; +import org.opensearch.index.mapper.DateFieldMapper; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.Arrays; +import java.util.Comparator; +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.opensearch.index.compositeindex.datacube.DateDimension.DateTimeUnitComparator.ORDERED_DATE_TIME_UNIT; + +public class DateDimensionTests extends OpenSearchTestCase { + public void testDateDimension() { + String field = "timestamp"; + List intervals = Arrays.asList( + DataCubeDateTimeUnit.HALF_HOUR_OF_DAY, + new DateTimeUnitAdapter(Rounding.DateTimeUnit.MONTH_OF_YEAR), + new DateTimeUnitAdapter(Rounding.DateTimeUnit.YEAR_OF_CENTURY) + ); + DateDimension dateDimension = new DateDimension(field, intervals, DateFieldMapper.Resolution.MILLISECONDS); + + assertEquals(field, dateDimension.getField()); + assertEquals(intervals, dateDimension.getIntervals()); + for (int i = 0; i < intervals.size(); i++) { + assertEquals(intervals.get(i).shortName(), dateDimension.getSortedCalendarIntervals().get(i).shortName()); + } + } + + public void testSetDimensionValuesWithMultipleIntervalsYear() { + List intervals = Arrays.asList( + new DateTimeUnitAdapter(Rounding.DateTimeUnit.YEAR_OF_CENTURY), + new DateTimeUnitAdapter(Rounding.DateTimeUnit.MONTH_OF_YEAR), + new DateTimeUnitAdapter(Rounding.DateTimeUnit.DAY_OF_MONTH), + new DateTimeUnitAdapter(Rounding.DateTimeUnit.HOUR_OF_DAY) + ); + DateDimension dateDimension = new DateDimension("timestamp", intervals, DateFieldMapper.Resolution.MILLISECONDS); + Long[] dims = new Long[4]; + Long testValue = 1609459200000L; // 2021-01-01 00:00:00 UTC + + AtomicInteger dimIndex = new AtomicInteger(0); + dateDimension.setDimensionValues(testValue, value -> { dims[dimIndex.getAndIncrement()] = value; }); + + assertEquals(4, dimIndex.get()); + assertEquals(1609459200000L, (long) dims[0]); // Hour rounded + assertEquals(1609459200000L, (long) dims[1]); // Day rounded + assertEquals(1609459200000L, (long) dims[2]); // Month rounded + assertEquals(1609459200000L, (long) dims[3]); // Year rounded + assertEquals(4, dateDimension.getNumSubDimensions()); + } + + public void testSetDimensionValuesForHalfAndQuarterHour() { + List intervals = Arrays.asList( + DataCubeDateTimeUnit.HALF_HOUR_OF_DAY, + DataCubeDateTimeUnit.QUARTER_HOUR_OF_DAY + ); + DateDimension dateDimension = new DateDimension("timestamp", intervals, DateFieldMapper.Resolution.MILLISECONDS); + Long[] dims = new Long[2]; + long testValue = 1724230620123L; // August 21, 2024 8:57:00.123 UTC + + AtomicInteger dimIndex = new AtomicInteger(0); + dateDimension.setDimensionValues(testValue, value -> { dims[dimIndex.getAndIncrement()] = value; }); + + assertEquals(2, dimIndex.get()); + assertEquals(1724229900000L, (long) dims[0]); // Quarter Hour rounded - Wed, 21 Aug 2024 08:45:00 UTC + assertEquals(1724229000000L, (long) dims[1]); // Half hour rounded - Wed, 21 Aug 2024 08:30:00 UTC + assertEquals(2, dateDimension.getNumSubDimensions()); + + Long[] dims1 = new Long[2]; + testValue = 1724229899234L; // Wed, 21 Aug 2024 08:44:59 GMT + + AtomicInteger dimIndex1 = new AtomicInteger(0); + dateDimension.setDimensionValues(testValue, value -> { dims1[dimIndex1.getAndIncrement()] = value; }); + + assertEquals(1724229000000L, (long) dims1[0]); // Quarter Hour rounded - Wed, 21 Aug 2024 08:30:00 UTC + assertEquals(1724229000000L, (long) dims1[1]); // Half hour rounded - Wed, 21 Aug 2024 08:30:00 UTC + assertEquals(2, dateDimension.getNumSubDimensions()); + + Long[] dims2 = new Long[2]; + testValue = 1724229000123L; // Wed, 21 Aug 2024 08:30:00 GMT + + AtomicInteger dimIndex2 = new AtomicInteger(0); + dateDimension.setDimensionValues(testValue, value -> { dims2[dimIndex2.getAndIncrement()] = value; }); + + assertEquals(2, dimIndex2.get()); + assertEquals(1724229000000L, (long) dims2[0]); // Quarter Hour rounded - Wed, 21 Aug 2024 08:30:00 UTC + assertEquals(1724229000000L, (long) dims2[1]); // Half hour rounded - Wed, 21 Aug 2024 08:30:00 UTC + assertEquals(2, dateDimension.getNumSubDimensions()); + + Long[] dims3 = new Long[2]; + testValue = 1724228940000L; // Wed, 21 Aug 2024 08:29:00 GMT + + AtomicInteger dimIndex3 = new AtomicInteger(0); + dateDimension.setDimensionValues(testValue, value -> { dims3[dimIndex3.getAndIncrement()] = value; }); + + assertEquals(2, dimIndex3.get()); + assertEquals(1724228100000L, (long) dims3[0]); // Quarter Hour rounded - Wed, 21 Aug 2024 08:15:00 UTC + assertEquals(1724227200000L, (long) dims3[1]); // Half hour rounded - Wed, 21 Aug 2024 08:00:00 UTC + assertEquals(2, dateDimension.getNumSubDimensions()); + } + + public void testRoundingAndSortingAllDateTimeUnitsNanos() { + List allUnits = getAllTimeUnits(); + + DateDimension dateDimension = new DateDimension("timestamp", allUnits, DateFieldMapper.Resolution.NANOSECONDS); + + // Test sorting + List sortedUnits = dateDimension.getSortedCalendarIntervals(); + assertEquals(allUnits.size(), sortedUnits.size()); + for (int i = 0; i < sortedUnits.size() - 1; i++) { + assertTrue( + "Units should be sorted in ascending order", + ORDERED_DATE_TIME_UNIT.get(sortedUnits.get(i).shortName()) + .compareTo(ORDERED_DATE_TIME_UNIT.get(sortedUnits.get(i + 1).shortName())) <= 0 + ); + } + + // Test rounding + long testValueNanos = 1655293505382719622L; // 2022-06-15T11:45:05.382719622Z + Long[] dims = new Long[allUnits.size()]; + + AtomicInteger dimIndex = new AtomicInteger(0); + dateDimension.setDimensionValues(testValueNanos, value -> { dims[dimIndex.getAndIncrement()] = value; }); + + // Expected rounded values (in nanoseconds) + long secondRounded = 1655293505000L; // 2022-06-15T11:45:05Z + long minuteRounded = 1655293500000L; // 2022-06-15T11:45:00Z + long quarterHourRounded = 1655293500000L; // 2022-06-15T11:45:00Z + long halfHourRounded = 1655292600000L; // 2022-06-15T11:30:00Z + long hourRounded = 1655290800000L; // 2022-06-15T11:00:00Z + long dayRounded = 1655251200000L; // 2022-06-15T00:00:00Z + long weekRounded = 1655078400000L; // 2022-06-13T00:00:00Z (Monday) + long monthRounded = 1654041600000L; // 2022-06-01T00:00:00Z + long quarterRounded = 1648771200000L; // 2022-04-01T00:00:00Z + long yearRounded = 1640995200000L; // 2022-01-01T00:00:00Z + + assertTimeUnits( + sortedUnits, + dims, + secondRounded, + minuteRounded, + hourRounded, + dayRounded, + weekRounded, + monthRounded, + quarterRounded, + yearRounded, + halfHourRounded, + quarterHourRounded + ); + } + + public void testRoundingAndSortingAllDateTimeUnitsMillis() { + List allUnits = getAllTimeUnits(); + + DateDimension dateDimension = new DateDimension("timestamp", allUnits, DateFieldMapper.Resolution.MILLISECONDS); + + // Test sorting + List sortedUnits = dateDimension.getSortedCalendarIntervals(); + assertEquals(allUnits.size(), sortedUnits.size()); + for (int i = 0; i < sortedUnits.size() - 1; i++) { + assertTrue( + "Units should be sorted in ascending order", + ORDERED_DATE_TIME_UNIT.get(sortedUnits.get(i).shortName()) + .compareTo(ORDERED_DATE_TIME_UNIT.get(sortedUnits.get(i + 1).shortName())) <= 0 + ); + } + + // Test rounding + long testValueNanos = 1724114825234L; // 2024-08-20T00:47:05.234Z + Long[] dims = new Long[allUnits.size()]; + + AtomicInteger dimIndex = new AtomicInteger(0); + dateDimension.setDimensionValues(testValueNanos, value -> { dims[dimIndex.getAndIncrement()] = value; }); + + // Expected rounded values (in millis) + long secondRounded = 1724114825000L; // 2024-08-20T00:47:05.000Z + long minuteRounded = 1724114820000L; // 2024-08-20T00:47:00.000Z + long quarterHourRounded = 1724114700000L; // 2024-08-20T00:45:00.000Z + long halfHourRounded = 1724113800000L; // 2024-08-20T00:30:00.000Z + long hourRounded = 1724112000000L; // 2024-08-20T00:00:00.000Z + long dayRounded = 1724112000000L; // 2024-08-20T00:00:00.000Z + long weekRounded = 1724025600000L; // 2024-08-15T00:00:00.000Z (Monday) + long monthRounded = 1722470400000L; // 2024-08-01T00:00:00.000Z + long quarterRounded = 1719792000000L; // 2024-07-01T00:00:00.000Z + long yearRounded = 1704067200000L; // 2024-01-01T00:00:00.000Z + + assertTimeUnits( + sortedUnits, + dims, + secondRounded, + minuteRounded, + hourRounded, + dayRounded, + weekRounded, + monthRounded, + quarterRounded, + yearRounded, + halfHourRounded, + quarterHourRounded + ); + } + + private static List getAllTimeUnits() { + return Arrays.asList( + new DateTimeUnitAdapter(Rounding.DateTimeUnit.WEEK_OF_WEEKYEAR), + new DateTimeUnitAdapter(Rounding.DateTimeUnit.MONTH_OF_YEAR), + new DateTimeUnitAdapter(Rounding.DateTimeUnit.QUARTER_OF_YEAR), + new DateTimeUnitAdapter(Rounding.DateTimeUnit.YEAR_OF_CENTURY), + new DateTimeUnitAdapter(Rounding.DateTimeUnit.SECOND_OF_MINUTE), + new DateTimeUnitAdapter(Rounding.DateTimeUnit.MINUTES_OF_HOUR), + new DateTimeUnitAdapter(Rounding.DateTimeUnit.HOUR_OF_DAY), + new DateTimeUnitAdapter(Rounding.DateTimeUnit.DAY_OF_MONTH), + DataCubeDateTimeUnit.HALF_HOUR_OF_DAY, + DataCubeDateTimeUnit.QUARTER_HOUR_OF_DAY + ); + } + + private static void assertTimeUnits( + List sortedUnits, + Long[] dims, + long secondRounded, + long minuteRounded, + long hourRounded, + long dayRounded, + long weekRounded, + long monthRounded, + long quarterRounded, + long yearRounded, + long halfHourRounded, + long quarterHourRounded + ) { + for (int i = 0; i < sortedUnits.size(); i++) { + DateTimeUnitRounding unit = sortedUnits.get(i); + String unitName = unit.shortName(); + switch (unitName) { + case "second": + assertEquals(secondRounded, (long) dims[i]); + break; + case "minute": + assertEquals(minuteRounded, (long) dims[i]); + break; + case "hour": + assertEquals(hourRounded, (long) dims[i]); + break; + case "day": + assertEquals(dayRounded, (long) dims[i]); + break; + case "week": + assertEquals(weekRounded, (long) dims[i]); + break; + case "month": + assertEquals(monthRounded, (long) dims[i]); + break; + case "quarter": + assertEquals(quarterRounded, (long) dims[i]); + break; + case "year": + assertEquals(yearRounded, (long) dims[i]); + break; + case "half-hour": + assertEquals(halfHourRounded, (long) dims[i]); + break; + case "quarter-hour": + assertEquals(quarterHourRounded, (long) dims[i]); + break; + default: + fail("Unexpected DateTimeUnit: " + unit); + } + } + } + + public void testGetSubDimensionNames() { + DateDimension dateDimension = new DateDimension( + "timestamp", + Arrays.asList( + new DateTimeUnitAdapter(Rounding.DateTimeUnit.HOUR_OF_DAY), + new DateTimeUnitAdapter(Rounding.DateTimeUnit.DAY_OF_MONTH) + ), + DateFieldMapper.Resolution.MILLISECONDS + ); + + List fields = dateDimension.getSubDimensionNames(); + + assertEquals(2, fields.size()); + assertEquals("timestamp_hour", fields.get(0)); + assertEquals("timestamp_day", fields.get(1)); + } + + public void testGetExtendedTimeUnitFieldsNames() { + DateDimension dateDimension = new DateDimension( + "timestamp", + Arrays.asList(DataCubeDateTimeUnit.HALF_HOUR_OF_DAY, DataCubeDateTimeUnit.QUARTER_HOUR_OF_DAY), + DateFieldMapper.Resolution.MILLISECONDS + ); + + List fields = dateDimension.getSubDimensionNames(); + + assertEquals(2, fields.size()); + assertEquals("timestamp_quarter-hour", fields.get(0)); + assertEquals("timestamp_half-hour", fields.get(1)); + } + + public void testSetDimensionValues() { + DateDimension dateDimension = new DateDimension( + "timestamp", + Arrays.asList( + new DateTimeUnitAdapter(Rounding.DateTimeUnit.YEAR_OF_CENTURY), + new DateTimeUnitAdapter(Rounding.DateTimeUnit.HOUR_OF_DAY) + ), + DateFieldMapper.Resolution.MILLISECONDS + ); + Long[] dims = new Long[2]; + Long testValue = 1609459200000L; // 2021-01-01 00:00:00 UTC + + AtomicInteger dimIndex = new AtomicInteger(0); + dateDimension.setDimensionValues(testValue, value -> { dims[dimIndex.getAndIncrement()] = value; }); + + assertEquals(2, dimIndex.get()); + assertEquals(1609459200000L, (long) dims[0]); // Hour rounded + assertEquals(1609459200000L, (long) dims[1]); // Year rounded + } + + public void testDateTimeUnitComparator() { + Comparator comparator = new DateDimension.DateTimeUnitComparator(); + assertTrue( + comparator.compare( + new DateTimeUnitAdapter(Rounding.DateTimeUnit.SECOND_OF_MINUTE), + new DateTimeUnitAdapter(Rounding.DateTimeUnit.MINUTES_OF_HOUR) + ) < 0 + ); + assertTrue( + comparator.compare( + new DateTimeUnitAdapter(Rounding.DateTimeUnit.HOUR_OF_DAY), + new DateTimeUnitAdapter(Rounding.DateTimeUnit.DAY_OF_MONTH) + ) < 0 + ); + assertTrue( + comparator.compare( + new DateTimeUnitAdapter(Rounding.DateTimeUnit.YEAR_OF_CENTURY), + new DateTimeUnitAdapter(Rounding.DateTimeUnit.MONTH_OF_YEAR) + ) > 0 + ); + assertEquals( + 0, + comparator.compare( + new DateTimeUnitAdapter(Rounding.DateTimeUnit.WEEK_OF_WEEKYEAR), + new DateTimeUnitAdapter(Rounding.DateTimeUnit.WEEK_OF_WEEKYEAR) + ) + ); + } +} diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/StarTreeTestUtils.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/StarTreeTestUtils.java index 7cae1cd25ee93..dc8b3320f3de2 100644 --- a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/StarTreeTestUtils.java +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/StarTreeTestUtils.java @@ -50,11 +50,19 @@ public static StarTreeDocument[] getSegmentsStarTreeDocuments( List starTreeDocuments = new ArrayList<>(); for (StarTreeValues starTreeValues : starTreeValuesSubs) { List dimensionsSplitOrder = starTreeValues.getStarTreeField().getDimensionsOrder(); - SequentialDocValuesIterator[] dimensionReaders = new SequentialDocValuesIterator[dimensionsSplitOrder.size()]; + int numDimensions = 0; + for (Dimension dimension : dimensionsSplitOrder) { + numDimensions += dimension.getNumSubDimensions(); + } + SequentialDocValuesIterator[] dimensionReaders = new SequentialDocValuesIterator[numDimensions]; + int dimIndex = 0; for (int i = 0; i < dimensionsSplitOrder.size(); i++) { - String dimension = dimensionsSplitOrder.get(i).getField(); - dimensionReaders[i] = new SequentialDocValuesIterator(starTreeValues.getDimensionValuesIterator(dimension)); + Dimension dimension = dimensionsSplitOrder.get(i); + for (String name : dimension.getSubDimensionNames()) { + dimensionReaders[dimIndex] = new SequentialDocValuesIterator(starTreeValues.getDimensionValuesIterator(name)); + dimIndex++; + } } List metricReaders = new ArrayList<>(); 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 index 95fd1579c2f63..0d4e4c37d6924 100644 --- 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 @@ -33,12 +33,15 @@ import org.apache.lucene.util.InfoStream; import org.apache.lucene.util.NumericUtils; import org.apache.lucene.util.Version; +import org.opensearch.common.Rounding; import org.opensearch.common.settings.Settings; import org.opensearch.index.codec.composite.LuceneDocValuesConsumerFactory; import org.opensearch.index.codec.composite.LuceneDocValuesProducerFactory; import org.opensearch.index.codec.composite.composite99.Composite99Codec; import org.opensearch.index.codec.composite.composite99.Composite99DocValuesFormat; import org.opensearch.index.compositeindex.CompositeIndexConstants; +import org.opensearch.index.compositeindex.datacube.DataCubeDateTimeUnit; +import org.opensearch.index.compositeindex.datacube.DateDimension; import org.opensearch.index.compositeindex.datacube.Dimension; import org.opensearch.index.compositeindex.datacube.Metric; import org.opensearch.index.compositeindex.datacube.MetricStat; @@ -53,9 +56,12 @@ import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNodeType; import org.opensearch.index.compositeindex.datacube.startree.utils.SequentialDocValuesIterator; import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils; +import org.opensearch.index.compositeindex.datacube.startree.utils.date.DateTimeUnitAdapter; +import org.opensearch.index.compositeindex.datacube.startree.utils.date.DateTimeUnitRounding; import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.SortedNumericStarTreeValuesIterator; import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.StarTreeValuesIterator; import org.opensearch.index.mapper.ContentPath; +import org.opensearch.index.mapper.DateFieldMapper; import org.opensearch.index.mapper.DocumentMapper; import org.opensearch.index.mapper.FieldValueConverter; import org.opensearch.index.mapper.Mapper; @@ -84,6 +90,7 @@ import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Supplier; +import static org.opensearch.index.compositeindex.CompositeIndexConstants.SEGMENT_DOCS_COUNT; import static org.opensearch.index.compositeindex.datacube.startree.StarTreeTestUtils.validateFileFormats; import static org.opensearch.index.compositeindex.datacube.startree.fileformats.StarTreeWriter.VERSION_CURRENT; import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils.fullyQualifiedFieldNameForStarTreeDimensionsDocValues; @@ -1864,7 +1871,7 @@ private List getStarTreeDimensionNames(List dimensionsOrder) List dimensionNames = new ArrayList<>(); for (Dimension dimension : dimensionsOrder) { - dimensionNames.add(dimension.getField()); + dimensionNames.addAll(dimension.getSubDimensionNames()); } return dimensionNames; @@ -3842,7 +3849,7 @@ public void testMergeFlowWithDuplicateDimensionValueWithMaxLeafDocs() throws IOE } public static long getLongFromDouble(double value) { - return Double.doubleToLongBits(value); + return NumericUtils.doubleToSortableLong(value); } public void testMergeFlowWithMaxLeafDocsAndStarTreeNodesAssertion() throws IOException { @@ -4199,9 +4206,11 @@ public void testMergeFlow() throws IOException { ... [999, 999, 999, 999] | [19980.0] */ + builder.appendDocumentsToStarTree(starTreeDocumentIterator); for (StarTreeDocument starTreeDocument : builder.getStarTreeDocuments()) { assertEquals(starTreeDocument.dimensions[0] * 20.0, starTreeDocument.metrics[0]); - assertEquals(2L, starTreeDocument.metrics[1]); + assertEquals(starTreeDocument.dimensions[0] * 2, starTreeDocument.metrics[1]); + assertEquals(2L, starTreeDocument.metrics[2]); } builder.build(starTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); @@ -4213,7 +4222,7 @@ public void testMergeFlow() throws IOException { docValuesConsumer.close(); StarTreeMetadata starTreeMetadata = new StarTreeMetadata( - "sf", + compositeField.getName(), STAR_TREE, mock(IndexInput.class), VERSION_CURRENT, @@ -4237,6 +4246,251 @@ public void testMergeFlow() throws IOException { ); } + public void testFlushFlowWithTimestamps() throws IOException { + List dimList = List.of(1655288152000L, 1655288092000L, 1655288032000L, 1655287972000L, 1655288092000L); + List docsWithField = List.of(0, 1, 3, 4, 5); + List dimList2 = List.of(0L, 1L, 2L, 3L, 4L, 5L); + List docsWithField2 = List.of(0, 1, 2, 3, 4, 5); + + List metricsList = List.of( + getLongFromDouble(0.0), + getLongFromDouble(10.0), + getLongFromDouble(20.0), + getLongFromDouble(30.0), + getLongFromDouble(40.0), + getLongFromDouble(50.0) + ); + List metricsWithField = List.of(0, 1, 2, 3, 4, 5); + + 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) + ); + this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( + writeState, + Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, + Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, + Composite99DocValuesFormat.META_DOC_VALUES_CODEC, + Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION + ); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, getWriteState(6, writeState.segmentInfo.getId()), mapperService); + SequentialDocValuesIterator[] dimDvs = { new SequentialDocValuesIterator(d1sndv), new SequentialDocValuesIterator(d2sndv) }; + Iterator starTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( + dimDvs, + List.of(new SequentialDocValuesIterator(m1sndv), new SequentialDocValuesIterator(m2sndv)) + ); + /** + * Asserting following dim / metrics [ dim1, dim2 / Sum [metric], count [metric] ] + [1655287920000, 1655287200000, 1655287200000, 4] | [40.0, 1] + [1655287980000, 1655287200000, 1655287200000, 3] | [30.0, 1] + [1655288040000, 1655287200000, 1655287200000, 1] | [10.0, 1] + [1655288040000, 1655287200000, 1655287200000, 5] | [50.0, 1] + [1655288100000, 1655287200000, 1655287200000, 0] | [0.0, 1] + [null, null, null, 2] | [20.0, 1] + */ + int count = 0; + List starTreeDocumentsList = new ArrayList<>(); + starTreeDocumentIterator.forEachRemaining(starTreeDocumentsList::add); + starTreeDocumentIterator = starTreeDocumentsList.iterator(); + while (starTreeDocumentIterator.hasNext()) { + count++; + StarTreeDocument starTreeDocument = starTreeDocumentIterator.next(); + assertEquals(starTreeDocument.dimensions[3] * 1 * 10.0, starTreeDocument.metrics[1]); + assertEquals(1L, starTreeDocument.metrics[0]); + } + assertEquals(6, count); + builder.build(starTreeDocumentsList.iterator(), new AtomicInteger(), docValuesConsumer); + validateStarTree(builder.getRootNode(), 3, 10, builder.getStarTreeDocuments()); + } + + public void testMergeFlowWithTimestamps() throws IOException { + List dimList = List.of(1655288152000L, 1655288092000L, 1655288032000L, 1655287972000L, 1655288092000L, 1655288092000L); + List docsWithField = List.of(0, 1, 2, 3, 4, 6); + List dimList2 = List.of(1655288152000L, 1655288092000L, 1655288032000L, 1655287972000L, 1655288092000L, 1655288092000L, -1L); + List docsWithField2 = List.of(0, 1, 2, 3, 4, 6); + List dimList7 = List.of(1655288152000L, 1655288092000L, 1655288032000L, 1655287972000L, 1655288092000L, 1655288092000L, -1L); + List docsWithField7 = List.of(0, 1, 2, 3, 4, 6); + + List dimList5 = List.of(0L, 1L, 2L, 3L, 4L, 5L, 6L); + List docsWithField5 = List.of(0, 1, 2, 3, 4, 5, 6); + List metricsList1 = List.of( + getLongFromDouble(0.0), + getLongFromDouble(10.0), + getLongFromDouble(20.0), + getLongFromDouble(30.0), + getLongFromDouble(40.0), + getLongFromDouble(50.0), + getLongFromDouble(60.0) + ); + List metricsWithField1 = List.of(0, 1, 2, 3, 4, 5, 6); + List metricsList = List.of(0L, 1L, 2L, 3L, 4L, 5L, 6L); + List metricsWithField = List.of(0, 1, 2, 3, 4, 5, 6); + + List dimList3 = List.of(1655288152000L, 1655288092000L, 1655288032000L, -1L); + List docsWithField3 = List.of(0, 1, 3, 4); + List dimList4 = List.of(1655288152000L, 1655288092000L, 1655288032000L, -1L); + List docsWithField4 = List.of(0, 1, 3, 4); + List dimList8 = List.of(1655288152000L, 1655288092000L, 1655288032000L, -1L); + List docsWithField8 = List.of(0, 1, 3, 4); + + List dimList6 = List.of(5L, 6L, 7L, 8L); + List docsWithField6 = List.of(0, 1, 2, 3); + List metricsList21 = List.of( + getLongFromDouble(50.0), + getLongFromDouble(60.0), + getLongFromDouble(70.0), + getLongFromDouble(80.0), + getLongFromDouble(90.0) + ); + List metricsWithField21 = List.of(0, 1, 2, 3, 4); + List metricsList2 = List.of(5L, 6L, 7L, 8L, 9L); + List metricsWithField2 = List.of(0, 1, 2, 3, 4); + + compositeField = getStarTreeFieldWithDateDimension(); + StarTreeValues starTreeValues = getStarTreeValuesWithDates( + getSortedNumericMock(dimList, docsWithField), + getSortedNumericMock(dimList2, docsWithField2), + getSortedNumericMock(dimList7, docsWithField7), + getSortedNumericMock(dimList5, docsWithField5), + getSortedNumericMock(metricsList, metricsWithField), + getSortedNumericMock(metricsList1, metricsWithField1), + compositeField, + "6" + ); + + StarTreeValues starTreeValues2 = getStarTreeValuesWithDates( + getSortedNumericMock(dimList3, docsWithField3), + getSortedNumericMock(dimList4, docsWithField4), + getSortedNumericMock(dimList8, docsWithField8), + getSortedNumericMock(dimList6, docsWithField6), + getSortedNumericMock(metricsList2, metricsWithField2), + getSortedNumericMock(metricsList21, metricsWithField21), + compositeField, + "4" + ); + this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( + writeState, + Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, + Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, + Composite99DocValuesFormat.META_DOC_VALUES_CODEC, + Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION + ); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, getWriteState(4, writeState.segmentInfo.getId()), mapperService); + Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); + /** + [1655287972000, 1655287972000, 1655287972000, 3] | [30.0, 3] + [1655288032000, 1655288032000, 1655288032000, 2] | [20.0, 2] + [1655288032000, 1655288032000, 1655288032000, 8] | [80.0, 8] + [1655288092000, 1655288092000, 1655288092000, 1] | [10.0, 1] + [1655288092000, 1655288092000, 1655288092000, 4] | [40.0, 4] + [1655288092000, 1655288092000, 1655288092000, 6] | [60.0, 6] + [1655288152000, 1655288152000, 1655288152000, 0] | [0.0, 0] + [1655288152000, 1655288152000, 1655288152000, 5] | [50.0, 5] + [null, null, null, 5] | [50.0, 5] + [null, null, null, 7] | [70.0, 7] + */ + int count = 0; + builder.appendDocumentsToStarTree(starTreeDocumentIterator); + for (StarTreeDocument starTreeDocument : builder.getStarTreeDocuments()) { + count++; + assertEquals(starTreeDocument.dimensions[3] * 10.0, (double) starTreeDocument.metrics[1], 0); + assertEquals(starTreeDocument.dimensions[3], starTreeDocument.metrics[0]); + } + assertEquals(10, count); + builder.build(starTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); + validateStarTree(builder.getRootNode(), 4, 10, builder.getStarTreeDocuments()); + metaOut.close(); + dataOut.close(); + docValuesConsumer.close(); + + StarTreeMetadata starTreeMetadata = new StarTreeMetadata( + "sf", + STAR_TREE, + mock(IndexInput.class), + VERSION_CURRENT, + builder.numStarTreeNodes, + getStarTreeDimensionNames(compositeField.getDimensionsOrder()), + compositeField.getMetrics(), + 10, + builder.numStarTreeDocs, + compositeField.getStarTreeConfig().maxLeafDocs(), + Set.of(), + getBuildMode(), + 0, + 231 + ); + + validateStarTreeFileFormats( + builder.getRootNode(), + builder.getStarTreeDocuments().size(), + starTreeMetadata, + builder.getStarTreeDocuments() + ); + } + + private StarTreeValues getStarTreeValuesWithDates( + SortedNumericDocValues dimList, + SortedNumericDocValues dimList2, + SortedNumericDocValues dimList4, + SortedNumericDocValues dimList3, + SortedNumericDocValues metricsList, + SortedNumericDocValues metricsList1, + StarTreeField sf, + String number + ) { + Map> dimDocIdSetIterators = Map.of( + "field1_minute", + () -> new SortedNumericStarTreeValuesIterator(dimList), + "field1_half-hour", + () -> new SortedNumericStarTreeValuesIterator(dimList4), + "field1_hour", + () -> new SortedNumericStarTreeValuesIterator(dimList2), + "field3", + () -> new SortedNumericStarTreeValuesIterator(dimList3) + ); + Map> metricDocIdSetIterators = new LinkedHashMap<>(); + + metricDocIdSetIterators.put( + fullyQualifiedFieldNameForStarTreeMetricsDocValues( + sf.getName(), + "field2", + sf.getMetrics().get(0).getMetrics().get(0).getTypeName() + ), + () -> new SortedNumericStarTreeValuesIterator(metricsList) + ); + metricDocIdSetIterators.put( + fullyQualifiedFieldNameForStarTreeMetricsDocValues( + sf.getName(), + "field2", + sf.getMetrics().get(0).getMetrics().get(1).getTypeName() + ), + () -> new SortedNumericStarTreeValuesIterator(metricsList1) + ); + return new StarTreeValues(sf, null, dimDocIdSetIterators, metricDocIdSetIterators, Map.of(SEGMENT_DOCS_COUNT, number), null); + } + + private StarTreeField getStarTreeFieldWithDateDimension() { + List intervals = new ArrayList<>(); + intervals.add(new DateTimeUnitAdapter(Rounding.DateTimeUnit.MINUTES_OF_HOUR)); + intervals.add(new DateTimeUnitAdapter(Rounding.DateTimeUnit.HOUR_OF_DAY)); + intervals.add(DataCubeDateTimeUnit.HALF_HOUR_OF_DAY); + Dimension d1 = new DateDimension("field1", intervals, DateFieldMapper.Resolution.MILLISECONDS); + Dimension d2 = new NumericDimension("field3"); + Metric m1 = new Metric("field2", List.of(MetricStat.VALUE_COUNT, MetricStat.SUM)); + List dims = List.of(d1, d2); + List metrics = List.of(m1); + StarTreeFieldConfiguration c = new StarTreeFieldConfiguration(10, new HashSet<>(), getBuildMode()); + StarTreeField sf = new StarTreeField("sf", dims, metrics, c); + return sf; + } + private void validateStarTree( InMemoryTreeNode root, int totalDimensions, diff --git a/server/src/test/java/org/opensearch/index/mapper/ObjectMapperTests.java b/server/src/test/java/org/opensearch/index/mapper/ObjectMapperTests.java index b415e1e657f7f..c35cf3fc1e591 100644 --- a/server/src/test/java/org/opensearch/index/mapper/ObjectMapperTests.java +++ b/server/src/test/java/org/opensearch/index/mapper/ObjectMapperTests.java @@ -506,10 +506,10 @@ public void testCompositeFields() throws Exception { .startObject("startree") .field("type", "star_tree") .startObject("config") - .startArray("ordered_dimensions") - .startObject() - .field("name", "node") + .startObject("date_dimension") + .field("name", "@timestamp") .endObject() + .startArray("ordered_dimensions") .startObject() .field("name", "status") .endObject() @@ -526,8 +526,8 @@ public void testCompositeFields() throws Exception { .endObject() .endObject() .startObject("properties") - .startObject("node") - .field("type", "integer") + .startObject("@timestamp") + .field("type", "date") .endObject() .startObject("status") .field("type", "integer") @@ -565,9 +565,9 @@ public void testCompositeFields() throws Exception { StarTreeMapper starTreeMapper = (StarTreeMapper) mapper; assertEquals("star_tree", starTreeMapper.fieldType().typeName()); // Check that field in properties was parsed correctly as well - mapper = documentMapper.root().getMapper("node"); + mapper = documentMapper.root().getMapper("@timestamp"); assertNotNull(mapper); - assertEquals("integer", mapper.typeName()); + assertEquals("date", mapper.typeName()); FeatureFlags.initializeFeatureFlags(Settings.EMPTY); } diff --git a/server/src/test/java/org/opensearch/index/mapper/StarTreeMapperTests.java b/server/src/test/java/org/opensearch/index/mapper/StarTreeMapperTests.java index daa9fda7a5a3b..aac460bd5e332 100644 --- a/server/src/test/java/org/opensearch/index/mapper/StarTreeMapperTests.java +++ b/server/src/test/java/org/opensearch/index/mapper/StarTreeMapperTests.java @@ -18,6 +18,7 @@ import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.index.compositeindex.CompositeIndexSettings; import org.opensearch.index.compositeindex.CompositeIndexValidator; +import org.opensearch.index.compositeindex.datacube.DataCubeDateTimeUnit; import org.opensearch.index.compositeindex.datacube.DateDimension; import org.opensearch.index.compositeindex.datacube.Dimension; import org.opensearch.index.compositeindex.datacube.Metric; @@ -27,6 +28,8 @@ import org.opensearch.index.compositeindex.datacube.startree.StarTreeField; import org.opensearch.index.compositeindex.datacube.startree.StarTreeFieldConfiguration; import org.opensearch.index.compositeindex.datacube.startree.StarTreeIndexSettings; +import org.opensearch.index.compositeindex.datacube.startree.utils.date.DateTimeUnitAdapter; +import org.opensearch.index.compositeindex.datacube.startree.utils.date.DateTimeUnitRounding; import org.junit.After; import org.junit.Before; @@ -42,6 +45,7 @@ import static org.opensearch.index.IndexSettings.INDEX_TRANSLOG_FLUSH_THRESHOLD_SIZE_SETTING; import static org.opensearch.index.compositeindex.CompositeIndexSettings.COMPOSITE_INDEX_MAX_TRANSLOG_FLUSH_THRESHOLD_SIZE_SETTING; import static org.hamcrest.Matchers.containsString; +import static com.carrotsearch.randomizedtesting.RandomizedTest.getRandom; /** * Tests for {@link StarTreeMapper}. @@ -74,7 +78,16 @@ public void testValidStarTree() throws IOException { for (CompositeMappedFieldType type : compositeFieldTypes) { StarTreeMapper.StarTreeFieldType starTreeFieldType = (StarTreeMapper.StarTreeFieldType) type; assertEquals(2, starTreeFieldType.getDimensions().size()); - assertEquals("node", starTreeFieldType.getDimensions().get(0).getField()); + assertEquals("@timestamp", starTreeFieldType.getDimensions().get(0).getField()); + assertTrue(starTreeFieldType.getDimensions().get(0) instanceof DateDimension); + DateDimension dateDim = (DateDimension) starTreeFieldType.getDimensions().get(0); + List expectedTimeUnits = Arrays.asList( + new DateTimeUnitAdapter(Rounding.DateTimeUnit.DAY_OF_MONTH), + new DateTimeUnitAdapter(Rounding.DateTimeUnit.MONTH_OF_YEAR) + ); + for (int i = 0; i < expectedTimeUnits.size(); i++) { + assertEquals(expectedTimeUnits.get(i).shortName(), dateDim.getIntervals().get(i).shortName()); + } assertEquals("status", starTreeFieldType.getDimensions().get(1).getField()); assertEquals(2, starTreeFieldType.getMetrics().size()); assertEquals("size", starTreeFieldType.getMetrics().get(0).getField()); @@ -85,7 +98,7 @@ public void testValidStarTree() throws IOException { assertEquals(100, starTreeFieldType.getStarTreeConfig().maxLeafDocs()); assertEquals(StarTreeFieldConfiguration.StarTreeBuildMode.OFF_HEAP, starTreeFieldType.getStarTreeConfig().getBuildMode()); assertEquals( - new HashSet<>(Arrays.asList("node", "status")), + new HashSet<>(Arrays.asList("@timestamp", "status")), starTreeFieldType.getStarTreeConfig().getSkipStarNodeCreationInDims() ); } @@ -130,7 +143,16 @@ public void testMetricsWithJustSum() throws IOException { Set compositeFieldTypes = mapperService.getCompositeFieldTypes(); for (CompositeMappedFieldType type : compositeFieldTypes) { StarTreeMapper.StarTreeFieldType starTreeFieldType = (StarTreeMapper.StarTreeFieldType) type; - assertEquals("node", starTreeFieldType.getDimensions().get(0).getField()); + assertEquals("@timestamp", starTreeFieldType.getDimensions().get(0).getField()); + assertTrue(starTreeFieldType.getDimensions().get(0) instanceof DateDimension); + DateDimension dateDim = (DateDimension) starTreeFieldType.getDimensions().get(0); + List expectedTimeUnits = Arrays.asList( + Rounding.DateTimeUnit.DAY_OF_MONTH, + Rounding.DateTimeUnit.MONTH_OF_YEAR + ); + for (int i = 0; i < expectedTimeUnits.size(); i++) { + assertEquals(new DateTimeUnitAdapter(expectedTimeUnits.get(i)), dateDim.getIntervals().get(i)); + } assertEquals("status", starTreeFieldType.getDimensions().get(1).getField()); assertEquals("size", starTreeFieldType.getMetrics().get(0).getField()); @@ -140,7 +162,7 @@ public void testMetricsWithJustSum() throws IOException { assertEquals(100, starTreeFieldType.getStarTreeConfig().maxLeafDocs()); assertEquals(StarTreeFieldConfiguration.StarTreeBuildMode.OFF_HEAP, starTreeFieldType.getStarTreeConfig().getBuildMode()); assertEquals( - new HashSet<>(Arrays.asList("node", "status")), + new HashSet<>(Arrays.asList("@timestamp", "status")), starTreeFieldType.getStarTreeConfig().getSkipStarNodeCreationInDims() ); } @@ -151,7 +173,16 @@ public void testMetricsWithCountAndSum() throws IOException { Set compositeFieldTypes = mapperService.getCompositeFieldTypes(); for (CompositeMappedFieldType type : compositeFieldTypes) { StarTreeMapper.StarTreeFieldType starTreeFieldType = (StarTreeMapper.StarTreeFieldType) type; - assertEquals("node", starTreeFieldType.getDimensions().get(0).getField()); + assertEquals("@timestamp", starTreeFieldType.getDimensions().get(0).getField()); + assertTrue(starTreeFieldType.getDimensions().get(0) instanceof DateDimension); + DateDimension dateDim = (DateDimension) starTreeFieldType.getDimensions().get(0); + List expectedTimeUnits = Arrays.asList( + Rounding.DateTimeUnit.DAY_OF_MONTH, + Rounding.DateTimeUnit.MONTH_OF_YEAR + ); + for (int i = 0; i < expectedTimeUnits.size(); i++) { + assertEquals(new DateTimeUnitAdapter(expectedTimeUnits.get(i)), dateDim.getIntervals().get(i)); + } assertEquals("status", starTreeFieldType.getDimensions().get(1).getField()); assertEquals("size", starTreeFieldType.getMetrics().get(0).getField()); @@ -166,18 +197,47 @@ public void testMetricsWithCountAndSum() throws IOException { assertEquals(100, starTreeFieldType.getStarTreeConfig().maxLeafDocs()); assertEquals(StarTreeFieldConfiguration.StarTreeBuildMode.OFF_HEAP, starTreeFieldType.getStarTreeConfig().getBuildMode()); assertEquals( - new HashSet<>(Arrays.asList("node", "status")), + new HashSet<>(Arrays.asList("@timestamp", "status")), starTreeFieldType.getStarTreeConfig().getSkipStarNodeCreationInDims() ); } } + public void testValidStarTreeWithNoDateDim() throws IOException { + MapperService mapperService = createMapperService(getMinMappingWithDateDims(false, true, true)); + Set compositeFieldTypes = mapperService.getCompositeFieldTypes(); + for (CompositeMappedFieldType type : compositeFieldTypes) { + StarTreeMapper.StarTreeFieldType starTreeFieldType = (StarTreeMapper.StarTreeFieldType) type; + assertEquals("status", starTreeFieldType.getDimensions().get(0).getField()); + assertTrue(starTreeFieldType.getDimensions().get(0) instanceof NumericDimension); + assertEquals("metric_field", starTreeFieldType.getDimensions().get(1).getField()); + assertTrue(starTreeFieldType.getDimensions().get(0) instanceof NumericDimension); + assertEquals("status", starTreeFieldType.getMetrics().get(0).getField()); + List expectedMetrics = Arrays.asList(MetricStat.VALUE_COUNT, MetricStat.SUM, MetricStat.AVG); + assertEquals(expectedMetrics, starTreeFieldType.getMetrics().get(0).getMetrics()); + assertEquals(10000, starTreeFieldType.getStarTreeConfig().maxLeafDocs()); + assertEquals(StarTreeFieldConfiguration.StarTreeBuildMode.OFF_HEAP, starTreeFieldType.getStarTreeConfig().getBuildMode()); + assertEquals(new HashSet<>(), starTreeFieldType.getStarTreeConfig().getSkipStarNodeCreationInDims()); + } + } + public void testValidStarTreeDefaults() throws IOException { MapperService mapperService = createMapperService(getMinMapping()); Set compositeFieldTypes = mapperService.getCompositeFieldTypes(); for (CompositeMappedFieldType type : compositeFieldTypes) { StarTreeMapper.StarTreeFieldType starTreeFieldType = (StarTreeMapper.StarTreeFieldType) type; - assertEquals("node", starTreeFieldType.getDimensions().get(0).getField()); + assertEquals("@timestamp", starTreeFieldType.getDimensions().get(0).getField()); + assertTrue(starTreeFieldType.getDimensions().get(0) instanceof DateDimension); + DateDimension dateDim = (DateDimension) starTreeFieldType.getDimensions().get(0); + List expectedDimensionFields = Arrays.asList("@timestamp_minute", "@timestamp_half-hour"); + assertEquals(expectedDimensionFields, dateDim.getSubDimensionNames()); + List expectedTimeUnits = Arrays.asList( + new DateTimeUnitAdapter(Rounding.DateTimeUnit.MINUTES_OF_HOUR), + DataCubeDateTimeUnit.HALF_HOUR_OF_DAY + ); + for (int i = 0; i < expectedTimeUnits.size(); i++) { + assertEquals(expectedTimeUnits.get(i).shortName(), dateDim.getIntervals().get(i).shortName()); + } assertEquals("status", starTreeFieldType.getDimensions().get(1).getField()); assertEquals(3, starTreeFieldType.getMetrics().size()); assertEquals("status", starTreeFieldType.getMetrics().get(0).getField()); @@ -196,6 +256,56 @@ public void testValidStarTreeDefaults() throws IOException { } } + public void testValidStarTreeDateDims() throws IOException { + MapperService mapperService = createMapperService(getMinMappingWithDateDims(false, false, false)); + Set compositeFieldTypes = mapperService.getCompositeFieldTypes(); + for (CompositeMappedFieldType type : compositeFieldTypes) { + StarTreeMapper.StarTreeFieldType starTreeFieldType = (StarTreeMapper.StarTreeFieldType) type; + assertEquals("@timestamp", starTreeFieldType.getDimensions().get(0).getField()); + assertTrue(starTreeFieldType.getDimensions().get(0) instanceof DateDimension); + DateDimension dateDim = (DateDimension) starTreeFieldType.getDimensions().get(0); + List expectedDimensionFields = Arrays.asList("@timestamp_half-hour", "@timestamp_week", "@timestamp_month"); + assertEquals(expectedDimensionFields, dateDim.getSubDimensionNames()); + List expectedTimeUnits = Arrays.asList( + DataCubeDateTimeUnit.HALF_HOUR_OF_DAY, + new DateTimeUnitAdapter(Rounding.DateTimeUnit.WEEK_OF_WEEKYEAR), + new DateTimeUnitAdapter(Rounding.DateTimeUnit.MONTH_OF_YEAR) + ); + for (int i = 0; i < expectedTimeUnits.size(); i++) { + assertEquals(expectedTimeUnits.get(i).shortName(), dateDim.getSortedCalendarIntervals().get(i).shortName()); + } + assertEquals("status", starTreeFieldType.getDimensions().get(1).getField()); + assertEquals("status", starTreeFieldType.getMetrics().get(0).getField()); + List expectedMetrics = Arrays.asList(MetricStat.VALUE_COUNT, MetricStat.SUM, MetricStat.AVG); + assertEquals(expectedMetrics, starTreeFieldType.getMetrics().get(0).getMetrics()); + assertEquals(10000, starTreeFieldType.getStarTreeConfig().maxLeafDocs()); + assertEquals(StarTreeFieldConfiguration.StarTreeBuildMode.OFF_HEAP, starTreeFieldType.getStarTreeConfig().getBuildMode()); + assertEquals(Collections.emptySet(), starTreeFieldType.getStarTreeConfig().getSkipStarNodeCreationInDims()); + } + } + + public void testInValidStarTreeMinDims() throws IOException { + MapperParsingException ex = expectThrows( + MapperParsingException.class, + () -> createMapperService(getMinMappingWithDateDims(false, true, false)) + ); + assertEquals( + "Failed to parse mapping [_doc]: Atleast two dimensions are required to build star tree index field [startree]", + ex.getMessage() + ); + } + + public void testInvalidStarTreeDateDims() throws IOException { + MapperParsingException ex = expectThrows( + MapperParsingException.class, + () -> createMapperService(getMinMappingWithDateDims(true, false, false)) + ); + assertEquals( + "Failed to parse mapping [_doc]: At most [3] calendar intervals are allowed in dimension [@timestamp]", + ex.getMessage() + ); + } + public void testInvalidDim() { MapperParsingException ex = expectThrows( MapperParsingException.class, @@ -215,7 +325,7 @@ public void testInvalidMetric() { public void testNoMetrics() { MapperParsingException ex = expectThrows( MapperParsingException.class, - () -> createMapperService(getMinMapping(false, true, false, false)) + () -> createMapperService(getMinMapping(false, true, false, false, false)) ); assertThat( ex.getMessage(), @@ -226,7 +336,7 @@ public void testNoMetrics() { public void testInvalidParam() { MapperParsingException ex = expectThrows( MapperParsingException.class, - () -> createMapperService(getInvalidMapping(false, false, false, false, true, false)) + () -> createMapperService(getInvalidMapping(false, false, false, false, true, false, false)) ); assertEquals( "Failed to parse mapping [_doc]: Star tree mapping definition has unsupported parameters: [invalid : {invalid=invalid}]", @@ -237,7 +347,7 @@ public void testInvalidParam() { public void testNoDims() { MapperParsingException ex = expectThrows( MapperParsingException.class, - () -> createMapperService(getMinMapping(true, false, false, false)) + () -> createMapperService(getMinMapping(true, false, false, false, false)) ); assertThat( ex.getMessage(), @@ -245,18 +355,26 @@ public void testNoDims() { ); } + public void testMissingDateDims() { + MapperParsingException ex = expectThrows( + MapperParsingException.class, + () -> createMapperService(getMinMapping(false, false, false, false, true)) + ); + assertThat(ex.getMessage(), containsString("Failed to parse mapping [_doc]: unknown date dimension field [@timestamp]")); + } + public void testMissingDims() { MapperParsingException ex = expectThrows( MapperParsingException.class, - () -> createMapperService(getMinMapping(false, false, true, false)) + () -> createMapperService(getMinMapping(false, false, true, false, false)) ); - assertThat(ex.getMessage(), containsString("Failed to parse mapping [_doc]: unknown dimension field [node]")); + assertThat(ex.getMessage(), containsString("Failed to parse mapping [_doc]: unknown dimension field [status]")); } public void testMissingMetrics() { MapperParsingException ex = expectThrows( MapperParsingException.class, - () -> createMapperService(getMinMapping(false, false, false, true)) + () -> createMapperService(getMinMapping(false, false, false, true, false)) ); assertThat(ex.getMessage(), containsString("Failed to parse mapping [_doc]: unknown metric field [metric_field]")); } @@ -275,7 +393,7 @@ public void testInvalidMetricType() { public void testInvalidMetricTypeWithDocCount() { MapperParsingException ex = expectThrows( MapperParsingException.class, - () -> createMapperService(getInvalidMapping(false, false, false, false, false, true)) + () -> createMapperService(getInvalidMapping(false, false, false, false, false, true, false)) ); assertEquals("Failed to parse mapping [_doc]: Invalid metric stat: _doc_count", ex.getMessage()); } @@ -286,7 +404,18 @@ public void testInvalidDimType() { () -> createMapperService(getInvalidMapping(false, false, true, false)) ); assertEquals( - "Failed to parse mapping [_doc]: unsupported field type associated with dimension [node] as part of star tree field [startree]", + "Failed to parse mapping [_doc]: unsupported field type associated with dimension [status] as part of star tree field [startree]", + ex.getMessage() + ); + } + + public void testInvalidDateDimType() { + MapperParsingException ex = expectThrows( + MapperParsingException.class, + () -> createMapperService(getInvalidMapping(false, false, true, false, false, false, true)) + ); + assertEquals( + "Failed to parse mapping [_doc]: date_dimension [@timestamp] should be of type date for star tree field [startree]", ex.getMessage() ); } @@ -356,17 +485,17 @@ public void testMetric() { } public void testDimensions() { - List d1CalendarIntervals = new ArrayList<>(); - d1CalendarIntervals.add(Rounding.DateTimeUnit.HOUR_OF_DAY); - DateDimension d1 = new DateDimension("name", d1CalendarIntervals); - DateDimension d2 = new DateDimension("name", d1CalendarIntervals); + List d1CalendarIntervals = new ArrayList<>(); + d1CalendarIntervals.add(new DateTimeUnitAdapter(Rounding.DateTimeUnit.HOUR_OF_DAY)); + DateDimension d1 = new DateDimension("name", d1CalendarIntervals, DateFieldMapper.Resolution.MILLISECONDS); + DateDimension d2 = new DateDimension("name", d1CalendarIntervals, DateFieldMapper.Resolution.MILLISECONDS); assertEquals(d1, d2); - d2 = new DateDimension("name1", d1CalendarIntervals); + d2 = new DateDimension("name1", d1CalendarIntervals, DateFieldMapper.Resolution.MILLISECONDS); assertNotEquals(d1, d2); - List d2CalendarIntervals = new ArrayList<>(); - d2CalendarIntervals.add(Rounding.DateTimeUnit.HOUR_OF_DAY); - d2CalendarIntervals.add(Rounding.DateTimeUnit.HOUR_OF_DAY); - d2 = new DateDimension("name", d2CalendarIntervals); + List d2CalendarIntervals = new ArrayList<>(); + d2CalendarIntervals.add(new DateTimeUnitAdapter(Rounding.DateTimeUnit.HOUR_OF_DAY)); + d2CalendarIntervals.add(new DateTimeUnitAdapter(Rounding.DateTimeUnit.HOUR_OF_DAY)); + d2 = new DateDimension("name", d2CalendarIntervals, DateFieldMapper.Resolution.MILLISECONDS); assertNotEquals(d1, d2); NumericDimension n1 = new NumericDimension("name"); NumericDimension n2 = new NumericDimension("name"); @@ -387,9 +516,9 @@ public void testStarTreeField() { List m1 = new ArrayList<>(); m1.add(MetricStat.MAX); Metric metric1 = new Metric("name", m1); - List d1CalendarIntervals = new ArrayList<>(); - d1CalendarIntervals.add(Rounding.DateTimeUnit.HOUR_OF_DAY); - DateDimension d1 = new DateDimension("name", d1CalendarIntervals); + List d1CalendarIntervals = new ArrayList<>(); + d1CalendarIntervals.add(new DateTimeUnitAdapter(Rounding.DateTimeUnit.HOUR_OF_DAY)); + DateDimension d1 = new DateDimension("name", d1CalendarIntervals, DateFieldMapper.Resolution.MILLISECONDS); NumericDimension n1 = new NumericDimension("numeric"); NumericDimension n2 = new NumericDimension("name1"); @@ -483,7 +612,7 @@ public void testValidations() throws IOException { ) ); assertEquals( - "Aggregations not supported for the dimension field [node] with field type [integer] as part of star tree field", + "Aggregations not supported for the dimension field [@timestamp] with field type [date] as part of star tree field", ex.getMessage() ); @@ -506,14 +635,18 @@ private XContentBuilder getExpandedMappingWithJustAvg(String dim, String metric) b.field("max_leaf_docs", 100); b.startArray("skip_star_node_creation_for_dimensions"); { - b.value("node"); + b.value("@timestamp"); b.value("status"); } b.endArray(); - b.startArray("ordered_dimensions"); - b.startObject(); - b.field("name", "node"); + b.startObject("date_dimension"); + b.field("name", "@timestamp"); + b.startArray("calendar_intervals"); + b.value("day"); + b.value("month"); + b.endArray(); b.endObject(); + b.startArray("ordered_dimensions"); b.startObject(); b.field("name", dim); b.endObject(); @@ -530,8 +663,8 @@ private XContentBuilder getExpandedMappingWithJustAvg(String dim, String metric) b.endObject(); b.endObject(); b.startObject("properties"); - b.startObject("node"); - b.field("type", "integer"); + b.startObject("@timestamp"); + b.field("type", "date"); b.endObject(); b.startObject("status"); b.field("type", "integer"); @@ -551,10 +684,10 @@ private XContentBuilder getMappingWithDuplicateFields(boolean isDuplicateDim, bo .startObject("startree-1") .field("type", "star_tree") .startObject("config") - .startArray("ordered_dimensions") - .startObject() - .field("name", "node") + .startObject("date_dimension") + .field("name", "timestamp") .endObject() + .startArray("ordered_dimensions") .startObject() .field("name", "numeric_dv") .endObject() @@ -574,8 +707,8 @@ private XContentBuilder getMappingWithDuplicateFields(boolean isDuplicateDim, bo .endObject() .endObject() .startObject("properties") - .startObject("node") - .field("type", "integer") + .startObject("timestamp") + .field("type", "date") .endObject() .startObject("numeric_dv") .field("type", "integer") @@ -602,14 +735,18 @@ private XContentBuilder getExpandedMappingWithJustSum(String dim, String metric) b.field("max_leaf_docs", 100); b.startArray("skip_star_node_creation_for_dimensions"); { - b.value("node"); + b.value("@timestamp"); b.value("status"); } b.endArray(); - b.startArray("ordered_dimensions"); - b.startObject(); - b.field("name", "node"); + b.startObject("date_dimension"); + b.field("name", "@timestamp"); + b.startArray("calendar_intervals"); + b.value("day"); + b.value("month"); + b.endArray(); b.endObject(); + b.startArray("ordered_dimensions"); b.startObject(); b.field("name", dim); b.endObject(); @@ -626,8 +763,8 @@ private XContentBuilder getExpandedMappingWithJustSum(String dim, String metric) b.endObject(); b.endObject(); b.startObject("properties"); - b.startObject("node"); - b.field("type", "integer"); + b.startObject("@timestamp"); + b.field("type", "date"); b.endObject(); b.startObject("status"); b.field("type", "integer"); @@ -648,14 +785,18 @@ private XContentBuilder getExpandedMappingWithSumAndCount(String dim, String met b.field("max_leaf_docs", 100); b.startArray("skip_star_node_creation_for_dimensions"); { - b.value("node"); + b.value("@timestamp"); b.value("status"); } b.endArray(); - b.startArray("ordered_dimensions"); - b.startObject(); - b.field("name", "node"); + b.startObject("date_dimension"); + b.field("name", "@timestamp"); + b.startArray("calendar_intervals"); + b.value("day"); + b.value("month"); + b.endArray(); b.endObject(); + b.startArray("ordered_dimensions"); b.startObject(); b.field("name", dim); b.endObject(); @@ -673,8 +814,8 @@ private XContentBuilder getExpandedMappingWithSumAndCount(String dim, String met b.endObject(); b.endObject(); b.startObject("properties"); - b.startObject("node"); - b.field("type", "integer"); + b.startObject("@timestamp"); + b.field("type", "date"); b.endObject(); b.startObject("status"); b.field("type", "integer"); @@ -687,21 +828,95 @@ private XContentBuilder getExpandedMappingWithSumAndCount(String dim, String met } private XContentBuilder getMinMapping() throws IOException { - return getMinMapping(false, false, false, false); + return getMinMapping(false, false, false, false, false); } - private XContentBuilder getMinMapping(boolean isEmptyDims, boolean isEmptyMetrics, boolean missingDim, boolean missingMetric) + private XContentBuilder getMinMappingWithDateDims(boolean calendarIntervalsExceeded, boolean dateDimsAbsent, boolean additionalDim) throws IOException { return topMapping(b -> { b.startObject("composite"); b.startObject("startree"); + b.field("type", "star_tree"); + b.startObject("config"); - if (!isEmptyDims) { - b.startArray("ordered_dimensions"); + if (!dateDimsAbsent) { + b.startObject("date_dimension"); + b.field("name", "@timestamp"); + b.startArray("calendar_intervals"); + b.value(getRandom().nextBoolean() ? "week" : "1w"); + if (calendarIntervalsExceeded) { + b.value(getRandom().nextBoolean() ? "day" : "1d"); + b.value(getRandom().nextBoolean() ? "second" : "1s"); + b.value(getRandom().nextBoolean() ? "hour" : "1h"); + b.value(getRandom().nextBoolean() ? "minute" : "1m"); + b.value(getRandom().nextBoolean() ? "year" : "1y"); + b.value(getRandom().nextBoolean() ? "quarter-hour" : "15m"); + } + b.value(getRandom().nextBoolean() ? "month" : "1M"); + b.value(getRandom().nextBoolean() ? "half-hour" : "30m"); + b.endArray(); + b.endObject(); + } + b.startArray("ordered_dimensions"); + b.startObject(); + b.field("name", "status"); + b.endObject(); + if (additionalDim) { b.startObject(); - b.field("name", "node"); + b.field("name", "metric_field"); b.endObject(); + } + b.endArray(); + + b.startArray("metrics"); + b.startObject(); + b.field("name", "status"); + b.endObject(); + b.startObject(); + b.field("name", "metric_field"); + b.endObject(); + b.endArray(); + + b.endObject(); + + b.endObject(); + b.endObject(); + b.startObject("properties"); + + b.startObject("@timestamp"); + b.field("type", "date"); + b.endObject(); + + b.startObject("status"); + b.field("type", "integer"); + b.endObject(); + + b.startObject("metric_field"); + b.field("type", "integer"); + b.endObject(); + + b.endObject(); + }); + } + + private XContentBuilder getMinMapping( + boolean isEmptyDims, + boolean isEmptyMetrics, + boolean missingDim, + boolean missingMetric, + boolean missingDateDim + ) throws IOException { + return topMapping(b -> { + b.startObject("composite"); + b.startObject("startree"); + b.field("type", "star_tree"); + b.startObject("config"); + if (!isEmptyDims) { + b.startObject("date_dimension"); + b.field("name", "@timestamp"); + b.endObject(); + b.startArray("ordered_dimensions"); b.startObject(); b.field("name", "status"); b.endObject(); @@ -721,14 +936,16 @@ private XContentBuilder getMinMapping(boolean isEmptyDims, boolean isEmptyMetric b.endObject(); b.endObject(); b.startObject("properties"); + if (!missingDateDim) { + b.startObject("@timestamp"); + b.field("type", "date"); + b.endObject(); + } if (!missingDim) { - b.startObject("node"); + b.startObject("status"); b.field("type", "integer"); b.endObject(); } - b.startObject("status"); - b.field("type", "integer"); - b.endObject(); if (!missingMetric) { b.startObject("metric_field"); b.field("type", "integer"); @@ -812,14 +1029,14 @@ private XContentBuilder getInvalidMapping( boolean invalidDimType, boolean invalidMetricType, boolean invalidParam, - boolean invalidDocCountMetricType + boolean invalidDocCountMetricType, + boolean invalidDate ) throws IOException { return topMapping(b -> { b.startObject("composite"); b.startObject("startree"); b.field("type", "star_tree"); b.startObject("config"); - b.startArray("skip_star_node_creation_for_dimensions"); { if (invalidSkipDims) { @@ -828,6 +1045,9 @@ private XContentBuilder getInvalidMapping( b.value("status"); } b.endArray(); + b.startObject("date_dimension"); + b.field("name", "@timestamp"); + b.endObject(); if (invalidParam) { b.startObject("invalid"); b.field("invalid", "invalid"); @@ -836,12 +1056,9 @@ private XContentBuilder getInvalidMapping( b.startArray("ordered_dimensions"); if (!singleDim) { b.startObject(); - b.field("name", "node"); + b.field("name", "status"); b.endObject(); } - b.startObject(); - b.field("name", "status"); - b.endObject(); b.endArray(); b.startArray("metrics"); b.startObject(); @@ -861,16 +1078,20 @@ private XContentBuilder getInvalidMapping( b.endObject(); b.endObject(); b.startObject("properties"); - b.startObject("node"); - if (!invalidDimType) { - b.field("type", "integer"); + b.startObject("@timestamp"); + if (!invalidDate) { + b.field("type", "date"); } else { b.field("type", "keyword"); } b.endObject(); b.startObject("status"); - b.field("type", "integer"); + if (!invalidDimType) { + b.field("type", "integer"); + } else { + b.field("type", "keyword"); + } b.endObject(); b.startObject("metric_field"); if (invalidMetricType) { @@ -903,15 +1124,15 @@ private XContentBuilder getInvalidMappingWithDv( b.value("status"); } b.endArray(); + b.startObject("date_dimension"); + b.field("name", "@timestamp"); + b.endObject(); b.startArray("ordered_dimensions"); if (!singleDim) { b.startObject(); - b.field("name", "node"); + b.field("name", "status"); b.endObject(); } - b.startObject(); - b.field("name", "status"); - b.endObject(); b.endArray(); b.startArray("metrics"); b.startObject(); @@ -925,12 +1146,12 @@ private XContentBuilder getInvalidMappingWithDv( b.endObject(); b.endObject(); b.startObject("properties"); - b.startObject("node"); + b.startObject("@timestamp"); if (!invalidDimType) { - b.field("type", "integer"); + b.field("type", "date"); b.field("doc_values", "true"); } else { - b.field("type", "integer"); + b.field("type", "date"); b.field("doc_values", "false"); } b.endObject(); @@ -953,7 +1174,7 @@ private XContentBuilder getInvalidMappingWithDv( private XContentBuilder getInvalidMapping(boolean singleDim, boolean invalidSkipDims, boolean invalidDimType, boolean invalidMetricType) throws IOException { - return getInvalidMapping(singleDim, invalidSkipDims, invalidDimType, invalidMetricType, false, false); + return getInvalidMapping(singleDim, invalidSkipDims, invalidDimType, invalidMetricType, false, false, false); } protected boolean supportsOrIgnoresBoost() { From 3696c296220acf4bf0ecb1bc10cb4a71675d4939 Mon Sep 17 00:00:00 2001 From: Bharathwaj G Date: Tue, 8 Oct 2024 10:03:53 +0530 Subject: [PATCH 6/7] [Star tree] Refactoring builder tests (#16036) * Refactoring builder tests Signed-off-by: Bharathwaj G * adding date tests Signed-off-by: Bharathwaj G --------- Signed-off-by: Bharathwaj G --- .../builder/AbstractStarTreeBuilderTests.java | 4740 ----------------- .../startree/builder/BuilderTestsUtils.java | 526 ++ .../builder/OffHeapStarTreeBuilderTests.java | 36 - .../builder/OnHeapStarTreeBuilderTests.java | 37 - .../builder/StarTreeBuildMetricTests.java | 954 ++++ .../StarTreeBuilderFlushFlowTests.java | 419 ++ .../StarTreeBuilderMergeFlowTests.java | 1921 +++++++ .../StarTreeBuilderSortAndAggregateTests.java | 695 +++ .../builder/StarTreeBuilderTestCase.java | 368 ++ 9 files changed, 4883 insertions(+), 4813 deletions(-) delete mode 100644 server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/AbstractStarTreeBuilderTests.java create mode 100644 server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/BuilderTestsUtils.java delete mode 100644 server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/OffHeapStarTreeBuilderTests.java delete mode 100644 server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/OnHeapStarTreeBuilderTests.java create mode 100644 server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuildMetricTests.java create mode 100644 server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuilderFlushFlowTests.java create mode 100644 server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuilderMergeFlowTests.java create mode 100644 server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuilderSortAndAggregateTests.java create mode 100644 server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuilderTestCase.java 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 deleted file mode 100644 index 0d4e4c37d6924..0000000000000 --- a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/AbstractStarTreeBuilderTests.java +++ /dev/null @@ -1,4740 +0,0 @@ -/* - * 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.startree.builder; - -import org.apache.lucene.codecs.DocValuesConsumer; -import org.apache.lucene.codecs.DocValuesProducer; -import org.apache.lucene.codecs.lucene912.Lucene912Codec; -import org.apache.lucene.index.DocValues; -import org.apache.lucene.index.DocValuesType; -import org.apache.lucene.index.EmptyDocValuesProducer; -import org.apache.lucene.index.FieldInfo; -import org.apache.lucene.index.FieldInfos; -import org.apache.lucene.index.IndexFileNames; -import org.apache.lucene.index.IndexOptions; -import org.apache.lucene.index.SegmentInfo; -import org.apache.lucene.index.SegmentReadState; -import org.apache.lucene.index.SegmentWriteState; -import org.apache.lucene.index.SortedNumericDocValues; -import org.apache.lucene.index.VectorEncoding; -import org.apache.lucene.index.VectorSimilarityFunction; -import org.apache.lucene.sandbox.document.HalfFloatPoint; -import org.apache.lucene.search.DocIdSetIterator; -import org.apache.lucene.store.Directory; -import org.apache.lucene.store.IOContext; -import org.apache.lucene.store.IndexInput; -import org.apache.lucene.store.IndexOutput; -import org.apache.lucene.util.InfoStream; -import org.apache.lucene.util.NumericUtils; -import org.apache.lucene.util.Version; -import org.opensearch.common.Rounding; -import org.opensearch.common.settings.Settings; -import org.opensearch.index.codec.composite.LuceneDocValuesConsumerFactory; -import org.opensearch.index.codec.composite.LuceneDocValuesProducerFactory; -import org.opensearch.index.codec.composite.composite99.Composite99Codec; -import org.opensearch.index.codec.composite.composite99.Composite99DocValuesFormat; -import org.opensearch.index.compositeindex.CompositeIndexConstants; -import org.opensearch.index.compositeindex.datacube.DataCubeDateTimeUnit; -import org.opensearch.index.compositeindex.datacube.DateDimension; -import org.opensearch.index.compositeindex.datacube.Dimension; -import org.opensearch.index.compositeindex.datacube.Metric; -import org.opensearch.index.compositeindex.datacube.MetricStat; -import org.opensearch.index.compositeindex.datacube.NumericDimension; -import org.opensearch.index.compositeindex.datacube.startree.StarTreeDocument; -import org.opensearch.index.compositeindex.datacube.startree.StarTreeField; -import org.opensearch.index.compositeindex.datacube.startree.StarTreeFieldConfiguration; -import org.opensearch.index.compositeindex.datacube.startree.StarTreeTestUtils; -import org.opensearch.index.compositeindex.datacube.startree.fileformats.meta.StarTreeMetadata; -import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues; -import org.opensearch.index.compositeindex.datacube.startree.node.InMemoryTreeNode; -import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNodeType; -import org.opensearch.index.compositeindex.datacube.startree.utils.SequentialDocValuesIterator; -import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils; -import org.opensearch.index.compositeindex.datacube.startree.utils.date.DateTimeUnitAdapter; -import org.opensearch.index.compositeindex.datacube.startree.utils.date.DateTimeUnitRounding; -import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.SortedNumericStarTreeValuesIterator; -import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.StarTreeValuesIterator; -import org.opensearch.index.mapper.ContentPath; -import org.opensearch.index.mapper.DateFieldMapper; -import org.opensearch.index.mapper.DocumentMapper; -import org.opensearch.index.mapper.FieldValueConverter; -import org.opensearch.index.mapper.Mapper; -import org.opensearch.index.mapper.MapperService; -import org.opensearch.index.mapper.MappingLookup; -import org.opensearch.index.mapper.NumberFieldMapper; -import org.opensearch.test.OpenSearchTestCase; -import org.junit.Before; - -import java.io.IOException; -import java.nio.charset.StandardCharsets; -import java.util.ArrayDeque; -import java.util.ArrayList; -import java.util.Collections; -import java.util.HashMap; -import java.util.HashSet; -import java.util.Iterator; -import java.util.LinkedHashMap; -import java.util.LinkedList; -import java.util.List; -import java.util.Map; -import java.util.Objects; -import java.util.Queue; -import java.util.Set; -import java.util.UUID; -import java.util.concurrent.atomic.AtomicInteger; -import java.util.function.Supplier; - -import static org.opensearch.index.compositeindex.CompositeIndexConstants.SEGMENT_DOCS_COUNT; -import static org.opensearch.index.compositeindex.datacube.startree.StarTreeTestUtils.validateFileFormats; -import static org.opensearch.index.compositeindex.datacube.startree.fileformats.StarTreeWriter.VERSION_CURRENT; -import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils.fullyQualifiedFieldNameForStarTreeDimensionsDocValues; -import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues; -import static org.opensearch.index.mapper.CompositeMappedFieldType.CompositeFieldType.STAR_TREE; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; - -public abstract class AbstractStarTreeBuilderTests extends OpenSearchTestCase { - protected MapperService mapperService; - protected List dimensionsOrder; - protected List fields = List.of(); - protected List metrics; - protected Directory directory; - protected FieldInfo[] fieldsInfo; - protected StarTreeField compositeField; - protected Map fieldProducerMap; - protected SegmentWriteState writeState; - protected BaseStarTreeBuilder builder; - protected IndexOutput dataOut; - protected IndexOutput metaOut; - protected DocValuesConsumer docValuesConsumer; - protected String dataFileName; - protected String metaFileName; - - @Before - public void setup() throws IOException { - fields = List.of("field1", "field2", "field3", "field4", "field5", "field6", "field7", "field8", "field9", "field10"); - - dimensionsOrder = List.of( - new NumericDimension("field1"), - new NumericDimension("field3"), - new NumericDimension("field5"), - new NumericDimension("field8") - ); - metrics = List.of( - new Metric("field2", List.of(MetricStat.SUM)), - new Metric("field4", List.of(MetricStat.SUM)), - new Metric("field6", List.of(MetricStat.VALUE_COUNT)), - new Metric("field9", List.of(MetricStat.MIN)), - new Metric("field10", List.of(MetricStat.MAX)), - new Metric("_doc_count", List.of(MetricStat.DOC_COUNT)) - ); - - DocValuesProducer docValuesProducer = mock(DocValuesProducer.class); - - compositeField = new StarTreeField( - "test", - dimensionsOrder, - metrics, - new StarTreeFieldConfiguration(1, Set.of("field8"), getBuildMode()) - ); - directory = newFSDirectory(createTempDir()); - - fieldsInfo = new FieldInfo[fields.size()]; - fieldProducerMap = new HashMap<>(); - for (int i = 0; i < fieldsInfo.length; i++) { - fieldsInfo[i] = new FieldInfo( - fields.get(i), - i, - false, - false, - true, - IndexOptions.DOCS_AND_FREQS_AND_POSITIONS_AND_OFFSETS, - DocValuesType.SORTED_NUMERIC, - -1, - Collections.emptyMap(), - 0, - 0, - 0, - 0, - VectorEncoding.FLOAT32, - VectorSimilarityFunction.EUCLIDEAN, - false, - false - ); - fieldProducerMap.put(fields.get(i), docValuesProducer); - } - writeState = getWriteState(5, UUID.randomUUID().toString().substring(0, 16).getBytes(StandardCharsets.UTF_8)); - - dataFileName = IndexFileNames.segmentFileName( - writeState.segmentInfo.name, - writeState.segmentSuffix, - Composite99DocValuesFormat.DATA_EXTENSION - ); - dataOut = writeState.directory.createOutput(dataFileName, writeState.context); - - metaFileName = IndexFileNames.segmentFileName( - writeState.segmentInfo.name, - writeState.segmentSuffix, - Composite99DocValuesFormat.META_EXTENSION - ); - metaOut = writeState.directory.createOutput(metaFileName, writeState.context); - - mapperService = mock(MapperService.class); - 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("field2", NumberFieldMapper.NumberType.DOUBLE, false, true) - .build(new Mapper.BuilderContext(settings, new ContentPath())); - NumberFieldMapper numberFieldMapper2 = new NumberFieldMapper.Builder("field4", NumberFieldMapper.NumberType.DOUBLE, false, true) - .build(new Mapper.BuilderContext(settings, new ContentPath())); - NumberFieldMapper numberFieldMapper3 = new NumberFieldMapper.Builder("field6", NumberFieldMapper.NumberType.DOUBLE, false, true) - .build(new Mapper.BuilderContext(settings, new ContentPath())); - NumberFieldMapper numberFieldMapper4 = new NumberFieldMapper.Builder("field9", NumberFieldMapper.NumberType.DOUBLE, false, true) - .build(new Mapper.BuilderContext(settings, new ContentPath())); - NumberFieldMapper numberFieldMapper5 = new NumberFieldMapper.Builder("field10", NumberFieldMapper.NumberType.DOUBLE, false, true) - .build(new Mapper.BuilderContext(settings, new ContentPath())); - MappingLookup fieldMappers = new MappingLookup( - Set.of(numberFieldMapper1, numberFieldMapper2, numberFieldMapper3, numberFieldMapper4, numberFieldMapper5), - Collections.emptyList(), - Collections.emptyList(), - 0, - null - ); - when(documentMapper.mappers()).thenReturn(fieldMappers); - docValuesConsumer = mock(DocValuesConsumer.class); - } - - private SegmentReadState getReadState(int numDocs, List dimensionFields, List metrics) { - - int numMetrics = 0; - for (Metric metric : metrics) { - numMetrics += metric.getBaseMetrics().size(); - } - - FieldInfo[] fields = new FieldInfo[dimensionFields.size() + numMetrics]; - - int i = 0; - for (String dimension : dimensionFields) { - fields[i] = new FieldInfo( - fullyQualifiedFieldNameForStarTreeDimensionsDocValues(compositeField.getName(), dimension), - i, - false, - false, - true, - IndexOptions.DOCS_AND_FREQS_AND_POSITIONS_AND_OFFSETS, - DocValuesType.SORTED_NUMERIC, - -1, - Collections.emptyMap(), - 0, - 0, - 0, - 0, - VectorEncoding.FLOAT32, - VectorSimilarityFunction.EUCLIDEAN, - false, - false - ); - i++; - } - - for (Metric metric : metrics) { - for (MetricStat metricStat : metric.getBaseMetrics()) { - fields[i] = new FieldInfo( - fullyQualifiedFieldNameForStarTreeMetricsDocValues( - compositeField.getName(), - metric.getField(), - metricStat.getTypeName() - ), - i, - false, - false, - true, - IndexOptions.DOCS_AND_FREQS_AND_POSITIONS_AND_OFFSETS, - DocValuesType.SORTED_NUMERIC, - -1, - Collections.emptyMap(), - 0, - 0, - 0, - 0, - VectorEncoding.FLOAT32, - VectorSimilarityFunction.EUCLIDEAN, - false, - false - ); - i++; - } - } - - SegmentInfo segmentInfo = new SegmentInfo( - directory, - Version.LATEST, - Version.LUCENE_9_11_0, - "test_segment", - numDocs, - false, - false, - new Lucene912Codec(), - new HashMap<>(), - writeState.segmentInfo.getId(), - new HashMap<>(), - null - ); - return new SegmentReadState(segmentInfo.dir, segmentInfo, new FieldInfos(fields), writeState.context); - } - - private SegmentWriteState getWriteState(int numDocs, byte[] id) { - FieldInfos fieldInfos = new FieldInfos(fieldsInfo); - SegmentInfo segmentInfo = new SegmentInfo( - directory, - Version.LATEST, - Version.LUCENE_9_12_0, - "test_segment", - numDocs, - false, - false, - new Lucene912Codec(), - new HashMap<>(), - id, - new HashMap<>(), - null - ); - return new SegmentWriteState(InfoStream.getDefault(), segmentInfo.dir, segmentInfo, fieldInfos, null, newIOContext(random())); - } - - public abstract BaseStarTreeBuilder getStarTreeBuilder( - IndexOutput metaOut, - IndexOutput dataOut, - StarTreeField starTreeField, - SegmentWriteState segmentWriteState, - MapperService mapperService - ) throws IOException; - - public void test_sortAndAggregateStarTreeDocuments() throws IOException { - - int noOfStarTreeDocuments = 5; - StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - - starTreeDocuments[0] = new StarTreeDocument( - new Long[] { 2L, 4L, 3L, 4L }, - new Object[] { 12.0, 10.0, randomDouble(), 8.0, 20.0, 10L } - ); - starTreeDocuments[1] = new StarTreeDocument( - new Long[] { 3L, 4L, 2L, 1L }, - new Object[] { 10.0, 6.0, randomDouble(), 12.0, 10.0, 10L } - ); - starTreeDocuments[2] = new StarTreeDocument( - new Long[] { 3L, 4L, 2L, 1L }, - new Object[] { 14.0, 12.0, randomDouble(), 6.0, 24.0, 10L } - ); - starTreeDocuments[3] = new StarTreeDocument( - new Long[] { 2L, 4L, 3L, 4L }, - new Object[] { 9.0, 4.0, randomDouble(), 9.0, 12.0, null } - ); - starTreeDocuments[4] = new StarTreeDocument( - new Long[] { 3L, 4L, 2L, 1L }, - new Object[] { 11.0, 16.0, randomDouble(), 8.0, 13.0, null } - ); - - StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - for (int i = 0; i < noOfStarTreeDocuments; i++) { - long metric1 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]); - long metric2 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[1]); - long metric3 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]); - long metric4 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[3]); - long metric5 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[4]); - Long metric6 = (Long) starTreeDocuments[i].metrics[5]; - segmentStarTreeDocuments[i] = new StarTreeDocument( - starTreeDocuments[i].dimensions, - new Long[] { metric1, metric2, metric3, metric4, metric5, metric6 } - ); - } - List inorderStarTreeDocuments = List.of( - new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Object[] { 21.0, 14.0, 2L, 8.0, 20.0, 11L }), - new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Object[] { 35.0, 34.0, 3L, 6.0, 24.0, 21L }) - ); - Iterator expectedStarTreeDocumentIterator = inorderStarTreeDocuments.iterator(); - - SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); - List metricsIterators = getMetricIterators(segmentStarTreeDocuments); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( - dimsIterators, - metricsIterators - ); - - int numOfAggregatedDocuments = 0; - while (segmentStarTreeDocumentIterator.hasNext() && expectedStarTreeDocumentIterator.hasNext()) { - StarTreeDocument resultStarTreeDocument = segmentStarTreeDocumentIterator.next(); - StarTreeDocument expectedStarTreeDocument = expectedStarTreeDocumentIterator.next(); - - assertEquals(expectedStarTreeDocument.dimensions[0], resultStarTreeDocument.dimensions[0]); - assertEquals(expectedStarTreeDocument.dimensions[1], resultStarTreeDocument.dimensions[1]); - assertEquals(expectedStarTreeDocument.dimensions[2], resultStarTreeDocument.dimensions[2]); - assertEquals(expectedStarTreeDocument.dimensions[3], resultStarTreeDocument.dimensions[3]); - assertEquals(expectedStarTreeDocument.metrics[0], resultStarTreeDocument.metrics[0]); - assertEquals(expectedStarTreeDocument.metrics[1], resultStarTreeDocument.metrics[1]); - assertEquals(expectedStarTreeDocument.metrics[2], resultStarTreeDocument.metrics[2]); - assertEquals(expectedStarTreeDocument.metrics[3], resultStarTreeDocument.metrics[3]); - assertEquals(expectedStarTreeDocument.metrics[4], resultStarTreeDocument.metrics[4]); - assertEquals(expectedStarTreeDocument.metrics[5], resultStarTreeDocument.metrics[5]); - - numOfAggregatedDocuments++; - } - - assertEquals(inorderStarTreeDocuments.size(), numOfAggregatedDocuments); - } - - SequentialDocValuesIterator[] getDimensionIterators(StarTreeDocument[] starTreeDocuments) { - SequentialDocValuesIterator[] sequentialDocValuesIterators = - new SequentialDocValuesIterator[starTreeDocuments[0].dimensions.length]; - for (int j = 0; j < starTreeDocuments[0].dimensions.length; j++) { - List dimList = new ArrayList<>(); - List docsWithField = new ArrayList<>(); - - for (int i = 0; i < starTreeDocuments.length; i++) { - if (starTreeDocuments[i].dimensions[j] != null) { - dimList.add(starTreeDocuments[i].dimensions[j]); - docsWithField.add(i); - } - } - sequentialDocValuesIterators[j] = new SequentialDocValuesIterator( - new SortedNumericStarTreeValuesIterator(getSortedNumericMock(dimList, docsWithField)) - ); - } - return sequentialDocValuesIterators; - } - - List getMetricIterators(StarTreeDocument[] starTreeDocuments) { - List sequentialDocValuesIterators = new ArrayList<>(); - for (int j = 0; j < starTreeDocuments[0].metrics.length; j++) { - List metricslist = new ArrayList<>(); - List docsWithField = new ArrayList<>(); - - for (int i = 0; i < starTreeDocuments.length; i++) { - if (starTreeDocuments[i].metrics[j] != null) { - metricslist.add((long) starTreeDocuments[i].metrics[j]); - docsWithField.add(i); - } - } - sequentialDocValuesIterators.add( - new SequentialDocValuesIterator(new SortedNumericStarTreeValuesIterator(getSortedNumericMock(metricslist, docsWithField))) - ); - } - return sequentialDocValuesIterators; - } - - public void test_sortAndAggregateStarTreeDocuments_nullMetric() throws IOException { - - int noOfStarTreeDocuments = 5; - StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - - starTreeDocuments[0] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Object[] { 12.0, 10.0, randomDouble(), 8.0, 20.0 }); - starTreeDocuments[1] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Object[] { 10.0, 6.0, randomDouble(), 12.0, 10.0 }); - starTreeDocuments[2] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Object[] { 14.0, 12.0, randomDouble(), 6.0, 24.0 }); - starTreeDocuments[3] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Object[] { 9.0, 4.0, randomDouble(), 9.0, 12.0 }); - starTreeDocuments[4] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Object[] { 11.0, null, randomDouble(), 8.0, 13.0 }); - - List inorderStarTreeDocuments = List.of( - new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Object[] { 21.0, 14.0, 2L, 8.0, 20.0, 2L }), - new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Object[] { 35.0, 18.0, 3L, 6.0, 24.0, 3L }) - ); - Iterator expectedStarTreeDocumentIterator = inorderStarTreeDocuments.iterator(); - - StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - for (int i = 0; i < noOfStarTreeDocuments; i++) { - long metric1 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]); - Long metric2 = starTreeDocuments[i].metrics[1] != null - ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[1]) - : null; - long metric3 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]); - long metric4 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[3]); - long metric5 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[4]); - segmentStarTreeDocuments[i] = new StarTreeDocument( - starTreeDocuments[i].dimensions, - new Object[] { metric1, metric2, metric3, metric4, metric5, null } - ); - } - SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); - List metricsIterators = getMetricIterators(segmentStarTreeDocuments); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( - dimsIterators, - metricsIterators - ); - - while (segmentStarTreeDocumentIterator.hasNext() && expectedStarTreeDocumentIterator.hasNext()) { - StarTreeDocument resultStarTreeDocument = segmentStarTreeDocumentIterator.next(); - StarTreeDocument expectedStarTreeDocument = expectedStarTreeDocumentIterator.next(); - assertEquals(expectedStarTreeDocument.dimensions[0], resultStarTreeDocument.dimensions[0]); - assertEquals(expectedStarTreeDocument.dimensions[1], resultStarTreeDocument.dimensions[1]); - assertEquals(expectedStarTreeDocument.dimensions[2], resultStarTreeDocument.dimensions[2]); - assertEquals(expectedStarTreeDocument.dimensions[3], resultStarTreeDocument.dimensions[3]); - assertEquals(expectedStarTreeDocument.metrics[0], resultStarTreeDocument.metrics[0]); - assertEquals(expectedStarTreeDocument.metrics[1], resultStarTreeDocument.metrics[1]); - assertEquals(expectedStarTreeDocument.metrics[2], resultStarTreeDocument.metrics[2]); - assertEquals(expectedStarTreeDocument.metrics[3], resultStarTreeDocument.metrics[3]); - assertEquals(expectedStarTreeDocument.metrics[4], resultStarTreeDocument.metrics[4]); - assertEquals(expectedStarTreeDocument.metrics[5], resultStarTreeDocument.metrics[5]); - } - } - - public void test_sortAndAggregateStarTreeDocuments_nullMetricField() throws IOException { - - int noOfStarTreeDocuments = 5; - StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - // Setting second metric iterator as empty sorted numeric , indicating a metric field is null - starTreeDocuments[0] = new StarTreeDocument( - new Long[] { 2L, 4L, 3L, 4L }, - new Object[] { 12.0, null, randomDouble(), 8.0, 20.0, null } - ); - starTreeDocuments[1] = new StarTreeDocument( - new Long[] { 3L, 4L, 2L, 1L }, - new Object[] { 10.0, null, randomDouble(), 12.0, 10.0, null } - ); - starTreeDocuments[2] = new StarTreeDocument( - new Long[] { 3L, 4L, 2L, 1L }, - new Object[] { 14.0, null, randomDouble(), 6.0, 24.0, null } - ); - starTreeDocuments[3] = new StarTreeDocument( - new Long[] { 2L, 4L, 3L, 4L }, - new Object[] { 9.0, null, randomDouble(), 9.0, 12.0, 10L } - ); - starTreeDocuments[4] = new StarTreeDocument( - new Long[] { 3L, 4L, 2L, 1L }, - new Object[] { 11.0, null, randomDouble(), 8.0, 13.0, null } - ); - - List inorderStarTreeDocuments = List.of( - new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Object[] { 21.0, 0.0, 2L, 8.0, 20.0, 11L }), - new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Object[] { 35.0, 0.0, 3L, 6.0, 24.0, 3L }) - ); - Iterator expectedStarTreeDocumentIterator = inorderStarTreeDocuments.iterator(); - - StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - for (int i = 0; i < noOfStarTreeDocuments; i++) { - long metric1 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]); - Long metric2 = starTreeDocuments[i].metrics[1] != null - ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[1]) - : null; - long metric3 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]); - long metric4 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[3]); - long metric5 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[4]); - Long metric6 = starTreeDocuments[i].metrics[5] != null ? (Long) starTreeDocuments[i].metrics[5] : null; - segmentStarTreeDocuments[i] = new StarTreeDocument( - starTreeDocuments[i].dimensions, - new Object[] { metric1, metric2, metric3, metric4, metric5, metric6 } - ); - } - SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); - List metricsIterators = getMetricIterators(segmentStarTreeDocuments); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( - dimsIterators, - metricsIterators - ); - - while (segmentStarTreeDocumentIterator.hasNext() && expectedStarTreeDocumentIterator.hasNext()) { - StarTreeDocument resultStarTreeDocument = segmentStarTreeDocumentIterator.next(); - StarTreeDocument expectedStarTreeDocument = expectedStarTreeDocumentIterator.next(); - assertEquals(expectedStarTreeDocument.dimensions[0], resultStarTreeDocument.dimensions[0]); - assertEquals(expectedStarTreeDocument.dimensions[1], resultStarTreeDocument.dimensions[1]); - assertEquals(expectedStarTreeDocument.dimensions[2], resultStarTreeDocument.dimensions[2]); - assertEquals(expectedStarTreeDocument.dimensions[3], resultStarTreeDocument.dimensions[3]); - assertEquals(expectedStarTreeDocument.metrics[0], resultStarTreeDocument.metrics[0]); - assertEquals(expectedStarTreeDocument.metrics[1], resultStarTreeDocument.metrics[1]); - assertEquals(expectedStarTreeDocument.metrics[2], resultStarTreeDocument.metrics[2]); - assertEquals(expectedStarTreeDocument.metrics[3], resultStarTreeDocument.metrics[3]); - assertEquals(expectedStarTreeDocument.metrics[4], resultStarTreeDocument.metrics[4]); - assertEquals(expectedStarTreeDocument.metrics[5], resultStarTreeDocument.metrics[5]); - } - } - - public void test_sortAndAggregateStarTreeDocuments_nullAndMinusOneInDimensionField() throws IOException { - int noOfStarTreeDocuments = 5; - StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - // Setting second metric iterator as empty sorted numeric , indicating a metric field is null - starTreeDocuments[0] = new StarTreeDocument( - new Long[] { 2L, null, 3L, 4L }, - new Object[] { 12.0, null, randomDouble(), 8.0, 20.0, null } - ); - starTreeDocuments[1] = new StarTreeDocument( - new Long[] { null, 4L, 2L, 1L }, - new Object[] { 10.0, null, randomDouble(), 12.0, 10.0, null } - ); - starTreeDocuments[2] = new StarTreeDocument( - new Long[] { null, 4L, 2L, 1L }, - new Object[] { 14.0, null, randomDouble(), 6.0, 24.0, null } - ); - starTreeDocuments[3] = new StarTreeDocument( - new Long[] { 2L, null, 3L, 4L }, - new Object[] { 9.0, null, randomDouble(), 9.0, 12.0, 10L } - ); - starTreeDocuments[4] = new StarTreeDocument( - new Long[] { -1L, 4L, 2L, 1L }, - new Object[] { 11.0, null, randomDouble(), 8.0, 13.0, null } - ); - - List inorderStarTreeDocuments = List.of( - new StarTreeDocument(new Long[] { -1L, 4L, 2L, 1L }, new Object[] { 11.0, 0.0, 1L, 8.0, 13.0, 1L }), - new StarTreeDocument(new Long[] { 2L, null, 3L, 4L }, new Object[] { 21.0, 0.0, 2L, 8.0, 20.0, 11L }), - new StarTreeDocument(new Long[] { null, 4L, 2L, 1L }, new Object[] { 24.0, 0.0, 2L, 6.0, 24.0, 2L }) - ); - Iterator expectedStarTreeDocumentIterator = inorderStarTreeDocuments.iterator(); - - StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - for (int i = 0; i < noOfStarTreeDocuments; i++) { - long metric1 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]); - Long metric2 = starTreeDocuments[i].metrics[1] != null - ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[1]) - : null; - long metric3 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]); - long metric4 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[3]); - long metric5 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[4]); - Long metric6 = starTreeDocuments[i].metrics[5] != null ? (long) starTreeDocuments[i].metrics[5] : null; - segmentStarTreeDocuments[i] = new StarTreeDocument( - starTreeDocuments[i].dimensions, - new Object[] { metric1, metric2, metric3, metric4, metric5, metric6 } - ); - } - SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); - List metricsIterators = getMetricIterators(segmentStarTreeDocuments); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( - dimsIterators, - metricsIterators - ); - - while (segmentStarTreeDocumentIterator.hasNext()) { - StarTreeDocument resultStarTreeDocument = segmentStarTreeDocumentIterator.next(); - StarTreeDocument expectedStarTreeDocument = expectedStarTreeDocumentIterator.next(); - assertEquals(expectedStarTreeDocument.dimensions[0], resultStarTreeDocument.dimensions[0]); - assertEquals(expectedStarTreeDocument.dimensions[1], resultStarTreeDocument.dimensions[1]); - assertEquals(expectedStarTreeDocument.dimensions[2], resultStarTreeDocument.dimensions[2]); - assertEquals(expectedStarTreeDocument.dimensions[3], resultStarTreeDocument.dimensions[3]); - assertEquals(expectedStarTreeDocument.metrics[0], resultStarTreeDocument.metrics[0]); - assertEquals(expectedStarTreeDocument.metrics[1], resultStarTreeDocument.metrics[1]); - assertEquals(expectedStarTreeDocument.metrics[2], resultStarTreeDocument.metrics[2]); - assertEquals(expectedStarTreeDocument.metrics[3], resultStarTreeDocument.metrics[3]); - assertEquals(expectedStarTreeDocument.metrics[4], resultStarTreeDocument.metrics[4]); - assertEquals(expectedStarTreeDocument.metrics[5], resultStarTreeDocument.metrics[5]); - } - - assertFalse(expectedStarTreeDocumentIterator.hasNext()); - - builder.build(segmentStarTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); - validateStarTree(builder.getRootNode(), 4, 1, builder.getStarTreeDocuments()); - } - - public void test_sortAndAggregateStarTreeDocuments_nullDimensionsAndNullMetrics() throws IOException { - int noOfStarTreeDocuments = 5; - StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - // Setting second metric iterator as empty sorted numeric , indicating a metric field is null - starTreeDocuments[0] = new StarTreeDocument( - new Long[] { null, null, null, null }, - new Object[] { null, null, null, null, null, null } - ); - starTreeDocuments[1] = new StarTreeDocument( - new Long[] { null, null, null, null }, - new Object[] { null, null, null, null, null, null } - ); - starTreeDocuments[2] = new StarTreeDocument( - new Long[] { null, null, null, null }, - new Object[] { null, null, null, null, null, null } - ); - starTreeDocuments[3] = new StarTreeDocument( - new Long[] { null, null, null, null }, - new Object[] { null, null, null, null, null, null } - ); - starTreeDocuments[4] = new StarTreeDocument( - new Long[] { null, null, null, null }, - new Object[] { null, null, null, null, null, null } - ); - - List inorderStarTreeDocuments = List.of( - new StarTreeDocument(new Long[] { null, null, null, null }, new Object[] { 0.0, 0.0, 0L, null, null, 5L }) - ); - Iterator expectedStarTreeDocumentIterator = inorderStarTreeDocuments.iterator(); - - StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - for (int i = 0; i < noOfStarTreeDocuments; i++) { - Long metric1 = starTreeDocuments[i].metrics[0] != null - ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]) - : null; - Long metric2 = starTreeDocuments[i].metrics[1] != null - ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[1]) - : null; - Long metric3 = starTreeDocuments[i].metrics[2] != null - ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]) - : null; - Long metric4 = starTreeDocuments[i].metrics[3] != null - ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[3]) - : null; - Long metric5 = starTreeDocuments[i].metrics[4] != null - ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[4]) - : null; - segmentStarTreeDocuments[i] = new StarTreeDocument( - starTreeDocuments[i].dimensions, - new Object[] { metric1, metric2, metric3, metric4, metric5, null } - ); - } - SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); - List metricsIterators = getMetricIterators(segmentStarTreeDocuments); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( - dimsIterators, - metricsIterators - ); - - while (segmentStarTreeDocumentIterator.hasNext() && expectedStarTreeDocumentIterator.hasNext()) { - StarTreeDocument resultStarTreeDocument = segmentStarTreeDocumentIterator.next(); - StarTreeDocument expectedStarTreeDocument = expectedStarTreeDocumentIterator.next(); - assertEquals(expectedStarTreeDocument.dimensions[0], resultStarTreeDocument.dimensions[0]); - assertEquals(expectedStarTreeDocument.dimensions[1], resultStarTreeDocument.dimensions[1]); - assertEquals(expectedStarTreeDocument.dimensions[2], resultStarTreeDocument.dimensions[2]); - assertEquals(expectedStarTreeDocument.dimensions[3], resultStarTreeDocument.dimensions[3]); - assertEquals(expectedStarTreeDocument.metrics[0], resultStarTreeDocument.metrics[0]); - assertEquals(expectedStarTreeDocument.metrics[1], resultStarTreeDocument.metrics[1]); - assertEquals(expectedStarTreeDocument.metrics[2], resultStarTreeDocument.metrics[2]); - assertEquals(expectedStarTreeDocument.metrics[3], resultStarTreeDocument.metrics[3]); - assertEquals(expectedStarTreeDocument.metrics[4], resultStarTreeDocument.metrics[4]); - assertEquals(expectedStarTreeDocument.metrics[5], resultStarTreeDocument.metrics[5]); - } - builder.build(segmentStarTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); - validateStarTree(builder.getRootNode(), 4, 1, builder.getStarTreeDocuments()); - } - - public void test_sortAndAggregateStarTreeDocuments_nullDimensionsAndFewNullMetrics() throws IOException { - int noOfStarTreeDocuments = 5; - StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - - double sumValue = randomDouble(); - double minValue = randomDouble(); - double maxValue = randomDouble(); - - // Setting second metric iterator as empty sorted numeric , indicating a metric field is null - starTreeDocuments[0] = new StarTreeDocument( - new Long[] { null, null, null, null }, - new Object[] { null, null, randomDouble(), null, maxValue } - ); - starTreeDocuments[1] = new StarTreeDocument(new Long[] { null, null, null, null }, new Object[] { null, null, null, null, null }); - starTreeDocuments[2] = new StarTreeDocument( - new Long[] { null, null, null, null }, - new Object[] { null, null, null, minValue, null } - ); - starTreeDocuments[3] = new StarTreeDocument(new Long[] { null, null, null, null }, new Object[] { null, null, null, null, null }); - starTreeDocuments[4] = new StarTreeDocument( - new Long[] { null, null, null, null }, - new Object[] { sumValue, null, randomDouble(), null, null } - ); - - List inorderStarTreeDocuments = List.of( - new StarTreeDocument(new Long[] { null, null, null, null }, new Object[] { sumValue, 0.0, 2L, minValue, maxValue, 5L }) - ); - Iterator expectedStarTreeDocumentIterator = inorderStarTreeDocuments.iterator(); - - StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - for (int i = 0; i < noOfStarTreeDocuments; i++) { - Long metric1 = starTreeDocuments[i].metrics[0] != null - ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]) - : null; - Long metric2 = starTreeDocuments[i].metrics[1] != null - ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[1]) - : null; - Long metric3 = starTreeDocuments[i].metrics[2] != null - ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]) - : null; - Long metric4 = starTreeDocuments[i].metrics[3] != null - ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[3]) - : null; - Long metric5 = starTreeDocuments[i].metrics[4] != null - ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[4]) - : null; - segmentStarTreeDocuments[i] = new StarTreeDocument( - starTreeDocuments[i].dimensions, - new Object[] { metric1, metric2, metric3, metric4, metric5, null } - ); - } - SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); - List metricsIterators = getMetricIterators(segmentStarTreeDocuments); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( - dimsIterators, - metricsIterators - ); - - while (segmentStarTreeDocumentIterator.hasNext() && expectedStarTreeDocumentIterator.hasNext()) { - StarTreeDocument resultStarTreeDocument = segmentStarTreeDocumentIterator.next(); - StarTreeDocument expectedStarTreeDocument = expectedStarTreeDocumentIterator.next(); - assertEquals(expectedStarTreeDocument.dimensions[0], resultStarTreeDocument.dimensions[0]); - assertEquals(expectedStarTreeDocument.dimensions[1], resultStarTreeDocument.dimensions[1]); - assertEquals(expectedStarTreeDocument.dimensions[2], resultStarTreeDocument.dimensions[2]); - assertEquals(expectedStarTreeDocument.dimensions[3], resultStarTreeDocument.dimensions[3]); - assertEquals(expectedStarTreeDocument.metrics[0], resultStarTreeDocument.metrics[0]); - assertEquals(expectedStarTreeDocument.metrics[1], resultStarTreeDocument.metrics[1]); - assertEquals(expectedStarTreeDocument.metrics[2], resultStarTreeDocument.metrics[2]); - assertEquals(expectedStarTreeDocument.metrics[3], resultStarTreeDocument.metrics[3]); - assertEquals(expectedStarTreeDocument.metrics[4], resultStarTreeDocument.metrics[4]); - assertEquals(expectedStarTreeDocument.metrics[5], resultStarTreeDocument.metrics[5]); - } - builder.build(segmentStarTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); - validateStarTree(builder.getRootNode(), 4, 1, builder.getStarTreeDocuments()); - } - - public void test_sortAndAggregateStarTreeDocuments_emptyDimensions() throws IOException { - - int noOfStarTreeDocuments = 5; - StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - // Setting second metric iterator as empty sorted numeric , indicating a metric field is null - starTreeDocuments[0] = new StarTreeDocument( - new Long[] { null, null, null, null }, - new Object[] { 12.0, null, randomDouble(), 8.0, 20.0, 10L } - ); - starTreeDocuments[1] = new StarTreeDocument( - new Long[] { null, null, null, null }, - new Object[] { 10.0, null, randomDouble(), 12.0, 10.0, 10L } - ); - starTreeDocuments[2] = new StarTreeDocument( - new Long[] { null, null, null, null }, - new Object[] { 14.0, null, randomDouble(), 6.0, 24.0, 10L } - ); - starTreeDocuments[3] = new StarTreeDocument( - new Long[] { null, null, null, null }, - new Object[] { 9.0, null, randomDouble(), 9.0, 12.0, 10L } - ); - starTreeDocuments[4] = new StarTreeDocument( - new Long[] { null, null, null, null }, - new Object[] { 11.0, null, randomDouble(), 8.0, 13.0, 10L } - ); - - List inorderStarTreeDocuments = List.of( - new StarTreeDocument(new Long[] { null, null, null, null }, new Object[] { 56.0, 0.0, 5L, 6.0, 24.0, 50L }) - ); - Iterator expectedStarTreeDocumentIterator = inorderStarTreeDocuments.iterator(); - - StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - for (int i = 0; i < noOfStarTreeDocuments; i++) { - Long metric1 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]); - Long metric2 = starTreeDocuments[i].metrics[1] != null - ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[1]) - : null; - Long metric3 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]); - Long metric4 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[3]); - Long metric5 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[4]); - Long metric6 = (Long) starTreeDocuments[i].metrics[5]; - segmentStarTreeDocuments[i] = new StarTreeDocument( - starTreeDocuments[i].dimensions, - new Object[] { metric1, metric2, metric3, metric4, metric5, metric6 } - ); - } - SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); - List metricsIterators = getMetricIterators(segmentStarTreeDocuments); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( - dimsIterators, - metricsIterators - ); - - while (segmentStarTreeDocumentIterator.hasNext() && expectedStarTreeDocumentIterator.hasNext()) { - StarTreeDocument resultStarTreeDocument = segmentStarTreeDocumentIterator.next(); - StarTreeDocument expectedStarTreeDocument = expectedStarTreeDocumentIterator.next(); - assertEquals(expectedStarTreeDocument.dimensions[0], resultStarTreeDocument.dimensions[0]); - assertEquals(expectedStarTreeDocument.dimensions[1], resultStarTreeDocument.dimensions[1]); - assertEquals(expectedStarTreeDocument.dimensions[2], resultStarTreeDocument.dimensions[2]); - assertEquals(expectedStarTreeDocument.dimensions[3], resultStarTreeDocument.dimensions[3]); - assertEquals(expectedStarTreeDocument.metrics[0], resultStarTreeDocument.metrics[0]); - assertEquals(expectedStarTreeDocument.metrics[1], resultStarTreeDocument.metrics[1]); - assertEquals(expectedStarTreeDocument.metrics[2], resultStarTreeDocument.metrics[2]); - assertEquals(expectedStarTreeDocument.metrics[3], resultStarTreeDocument.metrics[3]); - assertEquals(expectedStarTreeDocument.metrics[4], resultStarTreeDocument.metrics[4]); - assertEquals(expectedStarTreeDocument.metrics[5], resultStarTreeDocument.metrics[5]); - } - } - - public void test_sortAndAggregateStarTreeDocument_longMaxAndLongMinDimensions() throws IOException { - - int noOfStarTreeDocuments = 5; - StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - - starTreeDocuments[0] = new StarTreeDocument( - new Long[] { Long.MIN_VALUE, 4L, 3L, 4L }, - new Object[] { 12.0, 10.0, randomDouble(), 8.0, 20.0 } - ); - starTreeDocuments[1] = new StarTreeDocument( - new Long[] { 3L, 4L, 2L, Long.MAX_VALUE }, - new Object[] { 10.0, 6.0, randomDouble(), 12.0, 10.0 } - ); - starTreeDocuments[2] = new StarTreeDocument( - new Long[] { 3L, 4L, 2L, Long.MAX_VALUE }, - new Object[] { 14.0, 12.0, randomDouble(), 6.0, 24.0 } - ); - starTreeDocuments[3] = new StarTreeDocument( - new Long[] { Long.MIN_VALUE, 4L, 3L, 4L }, - new Object[] { 9.0, 4.0, randomDouble(), 9.0, 12.0 } - ); - starTreeDocuments[4] = new StarTreeDocument( - new Long[] { 3L, 4L, 2L, Long.MAX_VALUE }, - new Object[] { 11.0, 16.0, randomDouble(), 8.0, 13.0 } - ); - - List inorderStarTreeDocuments = List.of( - new StarTreeDocument(new Long[] { Long.MIN_VALUE, 4L, 3L, 4L }, new Object[] { 21.0, 14.0, 2L, 8.0, 20.0, 2L }), - new StarTreeDocument(new Long[] { 3L, 4L, 2L, Long.MAX_VALUE }, new Object[] { 35.0, 34.0, 3L, 6.0, 24.0, 3L }) - ); - Iterator expectedStarTreeDocumentIterator = inorderStarTreeDocuments.iterator(); - - StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - for (int i = 0; i < noOfStarTreeDocuments; i++) { - long metric1 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]); - long metric2 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[1]); - long metric3 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]); - long metric4 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[3]); - long metric5 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[4]); - segmentStarTreeDocuments[i] = new StarTreeDocument( - starTreeDocuments[i].dimensions, - new Long[] { metric1, metric2, metric3, metric4, metric5, null } - ); - } - - SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); - List metricsIterators = getMetricIterators(segmentStarTreeDocuments); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( - dimsIterators, - metricsIterators - ); - int numOfAggregatedDocuments = 0; - while (segmentStarTreeDocumentIterator.hasNext() && expectedStarTreeDocumentIterator.hasNext()) { - StarTreeDocument resultStarTreeDocument = segmentStarTreeDocumentIterator.next(); - StarTreeDocument expectedStarTreeDocument = expectedStarTreeDocumentIterator.next(); - - assertEquals(expectedStarTreeDocument.dimensions[0], resultStarTreeDocument.dimensions[0]); - assertEquals(expectedStarTreeDocument.dimensions[1], resultStarTreeDocument.dimensions[1]); - assertEquals(expectedStarTreeDocument.dimensions[2], resultStarTreeDocument.dimensions[2]); - assertEquals(expectedStarTreeDocument.dimensions[3], resultStarTreeDocument.dimensions[3]); - assertEquals(expectedStarTreeDocument.metrics[0], resultStarTreeDocument.metrics[0]); - assertEquals(expectedStarTreeDocument.metrics[1], resultStarTreeDocument.metrics[1]); - assertEquals(expectedStarTreeDocument.metrics[2], resultStarTreeDocument.metrics[2]); - assertEquals(expectedStarTreeDocument.metrics[3], resultStarTreeDocument.metrics[3]); - assertEquals(expectedStarTreeDocument.metrics[4], resultStarTreeDocument.metrics[4]); - assertEquals(expectedStarTreeDocument.metrics[5], resultStarTreeDocument.metrics[5]); - - numOfAggregatedDocuments++; - } - - assertEquals(inorderStarTreeDocuments.size(), numOfAggregatedDocuments); - - } - - public void test_sortAndAggregateStarTreeDocument_DoubleMaxAndDoubleMinMetrics() throws IOException { - - int noOfStarTreeDocuments = 5; - StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - - starTreeDocuments[0] = new StarTreeDocument( - new Long[] { 2L, 4L, 3L, 4L }, - new Object[] { Double.MAX_VALUE, 10.0, randomDouble(), 8.0, 20.0, 100L } - ); - starTreeDocuments[1] = new StarTreeDocument( - new Long[] { 3L, 4L, 2L, 1L }, - new Object[] { 10.0, 6.0, randomDouble(), 12.0, 10.0, 100L } - ); - starTreeDocuments[2] = new StarTreeDocument( - new Long[] { 3L, 4L, 2L, 1L }, - new Object[] { 14.0, Double.MIN_VALUE, randomDouble(), 6.0, 24.0, 100L } - ); - starTreeDocuments[3] = new StarTreeDocument( - new Long[] { 2L, 4L, 3L, 4L }, - new Object[] { 9.0, 4.0, randomDouble(), 9.0, 12.0, 100L } - ); - starTreeDocuments[4] = new StarTreeDocument( - new Long[] { 3L, 4L, 2L, 1L }, - new Object[] { 11.0, 16.0, randomDouble(), 8.0, 13.0, 100L } - ); - - List inorderStarTreeDocuments = List.of( - new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Object[] { Double.MAX_VALUE + 9, 14.0, 2L, 8.0, 20.0, 200L }), - new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Object[] { 35.0, Double.MIN_VALUE + 22, 3L, 6.0, 24.0, 300L }) - ); - Iterator expectedStarTreeDocumentIterator = inorderStarTreeDocuments.iterator(); - - StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - for (int i = 0; i < noOfStarTreeDocuments; i++) { - long metric1 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]); - long metric2 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[1]); - long metric3 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]); - long metric4 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[3]); - long metric5 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[4]); - Long metric6 = (Long) starTreeDocuments[i].metrics[5]; - segmentStarTreeDocuments[i] = new StarTreeDocument( - starTreeDocuments[i].dimensions, - new Long[] { metric1, metric2, metric3, metric4, metric5, metric6 } - ); - } - - SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); - List metricsIterators = getMetricIterators(segmentStarTreeDocuments); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( - dimsIterators, - metricsIterators - ); - int numOfAggregatedDocuments = 0; - while (segmentStarTreeDocumentIterator.hasNext() && expectedStarTreeDocumentIterator.hasNext()) { - StarTreeDocument resultStarTreeDocument = segmentStarTreeDocumentIterator.next(); - StarTreeDocument expectedStarTreeDocument = expectedStarTreeDocumentIterator.next(); - - assertEquals(expectedStarTreeDocument.dimensions[0], resultStarTreeDocument.dimensions[0]); - assertEquals(expectedStarTreeDocument.dimensions[1], resultStarTreeDocument.dimensions[1]); - assertEquals(expectedStarTreeDocument.dimensions[2], resultStarTreeDocument.dimensions[2]); - assertEquals(expectedStarTreeDocument.dimensions[3], resultStarTreeDocument.dimensions[3]); - assertEquals(expectedStarTreeDocument.metrics[0], resultStarTreeDocument.metrics[0]); - assertEquals(expectedStarTreeDocument.metrics[1], resultStarTreeDocument.metrics[1]); - assertEquals(expectedStarTreeDocument.metrics[2], resultStarTreeDocument.metrics[2]); - assertEquals(expectedStarTreeDocument.metrics[3], resultStarTreeDocument.metrics[3]); - assertEquals(expectedStarTreeDocument.metrics[4], resultStarTreeDocument.metrics[4]); - - numOfAggregatedDocuments++; - } - - assertEquals(inorderStarTreeDocuments.size(), numOfAggregatedDocuments); - builder.build(segmentStarTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); - validateStarTree(builder.getRootNode(), 3, 1, builder.getStarTreeDocuments()); - - } - - public void test_build_halfFloatMetrics() throws IOException { - - mapperService = mock(MapperService.class); - 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("field2", NumberFieldMapper.NumberType.HALF_FLOAT, false, true) - .build(new Mapper.BuilderContext(settings, new ContentPath())); - NumberFieldMapper numberFieldMapper2 = new NumberFieldMapper.Builder("field4", NumberFieldMapper.NumberType.HALF_FLOAT, false, true) - .build(new Mapper.BuilderContext(settings, new ContentPath())); - NumberFieldMapper numberFieldMapper3 = new NumberFieldMapper.Builder("field6", NumberFieldMapper.NumberType.HALF_FLOAT, false, true) - .build(new Mapper.BuilderContext(settings, new ContentPath())); - NumberFieldMapper numberFieldMapper4 = new NumberFieldMapper.Builder("field9", NumberFieldMapper.NumberType.HALF_FLOAT, false, true) - .build(new Mapper.BuilderContext(settings, new ContentPath())); - NumberFieldMapper numberFieldMapper5 = new NumberFieldMapper.Builder( - "field10", - NumberFieldMapper.NumberType.HALF_FLOAT, - false, - true - ).build(new Mapper.BuilderContext(settings, new ContentPath())); - MappingLookup fieldMappers = new MappingLookup( - Set.of(numberFieldMapper1, numberFieldMapper2, numberFieldMapper3, numberFieldMapper4, numberFieldMapper5), - Collections.emptyList(), - Collections.emptyList(), - 0, - null - ); - when(documentMapper.mappers()).thenReturn(fieldMappers); - - int noOfStarTreeDocuments = 5; - StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - - starTreeDocuments[0] = new StarTreeDocument( - new Long[] { 2L, 4L, 3L, 4L }, - new HalfFloatPoint[] { - new HalfFloatPoint("hf1", 12), - new HalfFloatPoint("hf6", 10), - new HalfFloatPoint("field6", 10), - new HalfFloatPoint("field9", 8), - new HalfFloatPoint("field10", 20) } - ); - starTreeDocuments[1] = new StarTreeDocument( - new Long[] { 3L, 4L, 2L, 1L }, - new HalfFloatPoint[] { - new HalfFloatPoint("hf2", 10), - new HalfFloatPoint("hf7", 6), - new HalfFloatPoint("field6", 10), - new HalfFloatPoint("field9", 12), - new HalfFloatPoint("field10", 10) } - ); - starTreeDocuments[2] = new StarTreeDocument( - new Long[] { 3L, 4L, 2L, 1L }, - new HalfFloatPoint[] { - new HalfFloatPoint("hf3", 14), - new HalfFloatPoint("hf8", 12), - new HalfFloatPoint("field6", 10), - new HalfFloatPoint("field9", 6), - new HalfFloatPoint("field10", 24) } - ); - starTreeDocuments[3] = new StarTreeDocument( - new Long[] { 2L, 4L, 3L, 4L }, - new HalfFloatPoint[] { - new HalfFloatPoint("hf4", 9), - new HalfFloatPoint("hf9", 4), - new HalfFloatPoint("field6", 10), - new HalfFloatPoint("field9", 9), - new HalfFloatPoint("field10", 12) } - ); - starTreeDocuments[4] = new StarTreeDocument( - new Long[] { 3L, 4L, 2L, 1L }, - new HalfFloatPoint[] { - new HalfFloatPoint("hf5", 11), - new HalfFloatPoint("hf10", 16), - new HalfFloatPoint("field6", 10), - new HalfFloatPoint("field9", 8), - new HalfFloatPoint("field10", 13) } - ); - - StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - for (int i = 0; i < noOfStarTreeDocuments; i++) { - long metric1 = HalfFloatPoint.halfFloatToSortableShort( - ((HalfFloatPoint) starTreeDocuments[i].metrics[0]).numericValue().floatValue() - ); - long metric2 = HalfFloatPoint.halfFloatToSortableShort( - ((HalfFloatPoint) starTreeDocuments[i].metrics[1]).numericValue().floatValue() - ); - long metric3 = HalfFloatPoint.halfFloatToSortableShort( - ((HalfFloatPoint) starTreeDocuments[i].metrics[2]).numericValue().floatValue() - ); - long metric4 = HalfFloatPoint.halfFloatToSortableShort( - ((HalfFloatPoint) starTreeDocuments[i].metrics[3]).numericValue().floatValue() - ); - long metric5 = HalfFloatPoint.halfFloatToSortableShort( - ((HalfFloatPoint) starTreeDocuments[i].metrics[4]).numericValue().floatValue() - ); - segmentStarTreeDocuments[i] = new StarTreeDocument( - starTreeDocuments[i].dimensions, - new Long[] { metric1, metric2, metric3, metric4, metric5, null } - ); - } - - SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); - List metricsIterators = getMetricIterators(segmentStarTreeDocuments); - this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( - writeState, - Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, - Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, - Composite99DocValuesFormat.META_DOC_VALUES_CODEC, - Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION - ); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - - Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( - dimsIterators, - metricsIterators - ); - builder.build(segmentStarTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); - List resultStarTreeDocuments = builder.getStarTreeDocuments(); - assertEquals(7, resultStarTreeDocuments.size()); - - Iterator expectedStarTreeDocumentIterator = getExpectedStarTreeDocumentIterator().iterator(); - assertStarTreeDocuments(resultStarTreeDocuments, expectedStarTreeDocumentIterator); - - metaOut.close(); - dataOut.close(); - docValuesConsumer.close(); - - StarTreeMetadata starTreeMetadata = new StarTreeMetadata( - "test", - STAR_TREE, - mock(IndexInput.class), - VERSION_CURRENT, - builder.numStarTreeNodes, - getStarTreeDimensionNames(compositeField.getDimensionsOrder()), - compositeField.getMetrics(), - 2, - getExpectedStarTreeDocumentIterator().size(), - 1, - Set.of("field8"), - getBuildMode(), - 0, - 330 - ); - - validateStarTreeFileFormats( - builder.getRootNode(), - getExpectedStarTreeDocumentIterator().size(), - starTreeMetadata, - getExpectedStarTreeDocumentIterator() - ); - } - - public void test_build_floatMetrics() throws IOException { - - mapperService = mock(MapperService.class); - 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("field2", NumberFieldMapper.NumberType.FLOAT, false, true) - .build(new Mapper.BuilderContext(settings, new ContentPath())); - NumberFieldMapper numberFieldMapper2 = new NumberFieldMapper.Builder("field4", NumberFieldMapper.NumberType.FLOAT, false, true) - .build(new Mapper.BuilderContext(settings, new ContentPath())); - NumberFieldMapper numberFieldMapper3 = new NumberFieldMapper.Builder("field6", NumberFieldMapper.NumberType.FLOAT, false, true) - .build(new Mapper.BuilderContext(settings, new ContentPath())); - NumberFieldMapper numberFieldMapper4 = new NumberFieldMapper.Builder("field9", NumberFieldMapper.NumberType.FLOAT, false, true) - .build(new Mapper.BuilderContext(settings, new ContentPath())); - NumberFieldMapper numberFieldMapper5 = new NumberFieldMapper.Builder("field10", NumberFieldMapper.NumberType.FLOAT, false, true) - .build(new Mapper.BuilderContext(settings, new ContentPath())); - MappingLookup fieldMappers = new MappingLookup( - Set.of(numberFieldMapper1, numberFieldMapper2, numberFieldMapper3, numberFieldMapper4, numberFieldMapper5), - Collections.emptyList(), - Collections.emptyList(), - 0, - null - ); - when(documentMapper.mappers()).thenReturn(fieldMappers); - - int noOfStarTreeDocuments = 5; - StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - - starTreeDocuments[0] = new StarTreeDocument( - new Long[] { 2L, 4L, 3L, 4L }, - new Object[] { 12.0F, 10.0F, randomFloat(), 8.0F, 20.0F, null } - ); - starTreeDocuments[1] = new StarTreeDocument( - new Long[] { 3L, 4L, 2L, 1L }, - new Object[] { 10.0F, 6.0F, randomFloat(), 12.0F, 10.0F, null } - ); - starTreeDocuments[2] = new StarTreeDocument( - new Long[] { 3L, 4L, 2L, 1L }, - new Object[] { 14.0F, 12.0F, randomFloat(), 6.0F, 24.0F, null } - ); - starTreeDocuments[3] = new StarTreeDocument( - new Long[] { 2L, 4L, 3L, 4L }, - new Object[] { 9.0F, 4.0F, randomFloat(), 9.0F, 12.0F, null } - ); - starTreeDocuments[4] = new StarTreeDocument( - new Long[] { 3L, 4L, 2L, 1L }, - new Object[] { 11.0F, 16.0F, randomFloat(), 8.0F, 13.0F, null } - ); - - StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - for (int i = 0; i < noOfStarTreeDocuments; i++) { - long metric1 = NumericUtils.floatToSortableInt((Float) starTreeDocuments[i].metrics[0]); - long metric2 = NumericUtils.floatToSortableInt((Float) starTreeDocuments[i].metrics[1]); - long metric3 = NumericUtils.floatToSortableInt((Float) starTreeDocuments[i].metrics[2]); - long metric4 = NumericUtils.floatToSortableInt((Float) starTreeDocuments[i].metrics[3]); - long metric5 = NumericUtils.floatToSortableInt((Float) starTreeDocuments[i].metrics[4]); - Long metric6 = (Long) starTreeDocuments[i].metrics[5]; - segmentStarTreeDocuments[i] = new StarTreeDocument( - starTreeDocuments[i].dimensions, - new Long[] { metric1, metric2, metric3, metric4, metric5, metric6 } - ); - } - - SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); - List metricsIterators = getMetricIterators(segmentStarTreeDocuments); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( - dimsIterators, - metricsIterators - ); - this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( - writeState, - Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, - Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, - Composite99DocValuesFormat.META_DOC_VALUES_CODEC, - Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION - ); - builder.build(segmentStarTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); - - List resultStarTreeDocuments = builder.getStarTreeDocuments(); - assertEquals(7, resultStarTreeDocuments.size()); - - Iterator expectedStarTreeDocumentIterator = getExpectedStarTreeDocumentIterator().iterator(); - assertStarTreeDocuments(resultStarTreeDocuments, expectedStarTreeDocumentIterator); - - metaOut.close(); - dataOut.close(); - docValuesConsumer.close(); - - StarTreeMetadata starTreeMetadata = new StarTreeMetadata( - "test", - STAR_TREE, - mock(IndexInput.class), - VERSION_CURRENT, - builder.numStarTreeNodes, - getStarTreeDimensionNames(compositeField.getDimensionsOrder()), - compositeField.getMetrics(), - 2, - getExpectedStarTreeDocumentIterator().size(), - 1, - Set.of("field8"), - getBuildMode(), - 0, - 330 - ); - - validateStarTreeFileFormats( - builder.getRootNode(), - getExpectedStarTreeDocumentIterator().size(), - starTreeMetadata, - getExpectedStarTreeDocumentIterator() - ); - } - - abstract StarTreeFieldConfiguration.StarTreeBuildMode getBuildMode(); - - public void test_build_longMetrics() throws IOException { - - mapperService = mock(MapperService.class); - 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("field2", NumberFieldMapper.NumberType.LONG, false, true) - .build(new Mapper.BuilderContext(settings, new ContentPath())); - NumberFieldMapper numberFieldMapper2 = new NumberFieldMapper.Builder("field4", NumberFieldMapper.NumberType.LONG, false, true) - .build(new Mapper.BuilderContext(settings, new ContentPath())); - NumberFieldMapper numberFieldMapper3 = new NumberFieldMapper.Builder("field6", NumberFieldMapper.NumberType.LONG, false, true) - .build(new Mapper.BuilderContext(settings, new ContentPath())); - NumberFieldMapper numberFieldMapper4 = new NumberFieldMapper.Builder("field9", NumberFieldMapper.NumberType.LONG, false, true) - .build(new Mapper.BuilderContext(settings, new ContentPath())); - NumberFieldMapper numberFieldMapper5 = new NumberFieldMapper.Builder("field10", NumberFieldMapper.NumberType.LONG, false, true) - .build(new Mapper.BuilderContext(settings, new ContentPath())); - MappingLookup fieldMappers = new MappingLookup( - Set.of(numberFieldMapper1, numberFieldMapper2, numberFieldMapper3, numberFieldMapper4, numberFieldMapper5), - Collections.emptyList(), - Collections.emptyList(), - 0, - null - ); - when(documentMapper.mappers()).thenReturn(fieldMappers); - - int noOfStarTreeDocuments = 5; - StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - - starTreeDocuments[0] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Long[] { 12L, 10L, randomLong(), 8L, 20L }); - starTreeDocuments[1] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Long[] { 10L, 6L, randomLong(), 12L, 10L }); - starTreeDocuments[2] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Long[] { 14L, 12L, randomLong(), 6L, 24L }); - starTreeDocuments[3] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Long[] { 9L, 4L, randomLong(), 9L, 12L }); - starTreeDocuments[4] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Long[] { 11L, 16L, randomLong(), 8L, 13L }); - - StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - for (int i = 0; i < noOfStarTreeDocuments; i++) { - long metric1 = (Long) starTreeDocuments[i].metrics[0]; - long metric2 = (Long) starTreeDocuments[i].metrics[1]; - long metric3 = (Long) starTreeDocuments[i].metrics[2]; - long metric4 = (Long) starTreeDocuments[i].metrics[3]; - long metric5 = (Long) starTreeDocuments[i].metrics[4]; - segmentStarTreeDocuments[i] = new StarTreeDocument( - starTreeDocuments[i].dimensions, - new Long[] { metric1, metric2, metric3, metric4, metric5, null } - ); - } - - SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); - List metricsIterators = getMetricIterators(segmentStarTreeDocuments); - this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( - writeState, - Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, - Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, - Composite99DocValuesFormat.META_DOC_VALUES_CODEC, - Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION - ); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( - dimsIterators, - metricsIterators - ); - builder.build(segmentStarTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); - - List resultStarTreeDocuments = builder.getStarTreeDocuments(); - assertEquals(7, resultStarTreeDocuments.size()); - - Iterator expectedStarTreeDocumentIterator = getExpectedStarTreeDocumentIterator().iterator(); - assertStarTreeDocuments(resultStarTreeDocuments, expectedStarTreeDocumentIterator); - - metaOut.close(); - dataOut.close(); - docValuesConsumer.close(); - - StarTreeMetadata starTreeMetadata = new StarTreeMetadata( - "test", - STAR_TREE, - mock(IndexInput.class), - VERSION_CURRENT, - builder.numStarTreeNodes, - getStarTreeDimensionNames(compositeField.getDimensionsOrder()), - compositeField.getMetrics(), - 2, - getExpectedStarTreeDocumentIterator().size(), - 1, - Set.of("field8"), - getBuildMode(), - 0, - 330 - ); - - validateStarTreeFileFormats( - builder.getRootNode(), - getExpectedStarTreeDocumentIterator().size(), - starTreeMetadata, - getExpectedStarTreeDocumentIterator() - ); - } - - private static List getExpectedStarTreeDocumentIterator() { - return List.of( - new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Object[] { 21.0, 14.0, 2L, 8.0, 20.0, 2L }), - new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Object[] { 35.0, 34.0, 3L, 6.0, 24.0, 3L }), - new StarTreeDocument(new Long[] { null, 4L, 2L, 1L }, new Object[] { 35.0, 34.0, 3L, 6.0, 24.0, 3L }), - new StarTreeDocument(new Long[] { null, 4L, 3L, 4L }, new Object[] { 21.0, 14.0, 2L, 8.0, 20.0, 2L }), - new StarTreeDocument(new Long[] { null, 4L, null, 1L }, new Object[] { 35.0, 34.0, 3L, 6.0, 24.0, 3L }), - new StarTreeDocument(new Long[] { null, 4L, null, 4L }, new Object[] { 21.0, 14.0, 2L, 8.0, 20.0, 2L }), - new StarTreeDocument(new Long[] { null, 4L, null, null }, new Object[] { 56.0, 48.0, 5L, 6.0, 24.0, 5L }) - ); - } - - public void test_build_multipleStarTrees() throws IOException { - - int noOfStarTreeDocuments = 5; - StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - - starTreeDocuments[0] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Double[] { 12.0, 10.0, randomDouble(), 8.0, 20.0 }); - starTreeDocuments[1] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 10.0, 6.0, randomDouble(), 12.0, 10.0 }); - starTreeDocuments[2] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 14.0, 12.0, randomDouble(), 6.0, 24.0 }); - starTreeDocuments[3] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Double[] { 9.0, 4.0, randomDouble(), 9.0, 12.0 }); - starTreeDocuments[4] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 11.0, 16.0, randomDouble(), 8.0, 13.0 }); - - StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - for (int i = 0; i < noOfStarTreeDocuments; i++) { - long metric1 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]); - long metric2 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[1]); - long metric3 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]); - long metric4 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[3]); - long metric5 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[4]); - segmentStarTreeDocuments[i] = new StarTreeDocument( - starTreeDocuments[i].dimensions, - new Long[] { metric1, metric2, metric3, metric4, metric5 } - ); - } - - SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); - List metricsIterators = getMetricIterators(segmentStarTreeDocuments); - - metrics = List.of( - new Metric("field2", List.of(MetricStat.SUM)), - new Metric("field4", List.of(MetricStat.SUM)), - new Metric("field6", List.of(MetricStat.VALUE_COUNT)), - new Metric("field9", List.of(MetricStat.MIN)), - new Metric("field10", List.of(MetricStat.MAX)) - ); - - compositeField = new StarTreeField( - "test", - dimensionsOrder, - metrics, - new StarTreeFieldConfiguration(1, Set.of("field8"), getBuildMode()) - ); - - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( - dimsIterators, - metricsIterators - ); - builder.build(segmentStarTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); - - List resultStarTreeDocuments = builder.getStarTreeDocuments(); - assertEquals(7, resultStarTreeDocuments.size()); - - Iterator expectedStarTreeDocumentIterator = getExpectedStarTreeDocumentIterator().iterator(); - assertStarTreeDocuments(resultStarTreeDocuments, expectedStarTreeDocumentIterator); - builder.close(); - - // building another tree in the same file - fields = List.of("fieldC", "fieldB", "fieldL", "fieldI"); - - dimensionsOrder = List.of(new NumericDimension("fieldC"), new NumericDimension("fieldB"), new NumericDimension("fieldL")); - metrics = List.of(new Metric("fieldI", List.of(MetricStat.SUM))); - - DocValuesProducer docValuesProducer = mock(DocValuesProducer.class); - - compositeField = new StarTreeField("test", dimensionsOrder, metrics, new StarTreeFieldConfiguration(1, Set.of(), getBuildMode())); - SegmentInfo segmentInfo = new SegmentInfo( - directory, - Version.LATEST, - Version.LUCENE_9_11_0, - "test_segment", - 7, - false, - false, - new Lucene912Codec(), - new HashMap<>(), - UUID.randomUUID().toString().substring(0, 16).getBytes(StandardCharsets.UTF_8), - new HashMap<>(), - null - ); - - fieldsInfo = new FieldInfo[fields.size()]; - fieldProducerMap = new HashMap<>(); - for (int i = 0; i < fieldsInfo.length; i++) { - fieldsInfo[i] = new FieldInfo( - fields.get(i), - i, - false, - false, - true, - IndexOptions.DOCS_AND_FREQS_AND_POSITIONS_AND_OFFSETS, - DocValuesType.SORTED_NUMERIC, - -1, - Collections.emptyMap(), - 0, - 0, - 0, - 0, - VectorEncoding.FLOAT32, - VectorSimilarityFunction.EUCLIDEAN, - false, - false - ); - fieldProducerMap.put(fields.get(i), docValuesProducer); - } - FieldInfos fieldInfos = new FieldInfos(fieldsInfo); - writeState = new SegmentWriteState(InfoStream.getDefault(), segmentInfo.dir, segmentInfo, fieldInfos, null, newIOContext(random())); - - mapperService = mock(MapperService.class); - 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("fieldI", NumberFieldMapper.NumberType.DOUBLE, false, true) - .build(new Mapper.BuilderContext(settings, new ContentPath())); - MappingLookup fieldMappers = new MappingLookup( - Set.of(numberFieldMapper1), - Collections.emptyList(), - Collections.emptyList(), - 0, - null - ); - when(documentMapper.mappers()).thenReturn(fieldMappers); - - InMemoryTreeNode rootNode1 = builder.getRootNode(); - - int noOfStarTreeDocuments2 = 7; - StarTreeDocument[] starTreeDocuments2 = new StarTreeDocument[noOfStarTreeDocuments2]; - starTreeDocuments2[0] = new StarTreeDocument(new Long[] { 1L, 11L, 21L }, new Double[] { 400.0 }); - starTreeDocuments2[1] = new StarTreeDocument(new Long[] { 1L, 12L, 22L }, new Double[] { 200.0 }); - starTreeDocuments2[2] = new StarTreeDocument(new Long[] { 2L, 13L, 23L }, new Double[] { 300.0 }); - starTreeDocuments2[3] = new StarTreeDocument(new Long[] { 2L, 13L, 21L }, new Double[] { 100.0 }); - starTreeDocuments2[4] = new StarTreeDocument(new Long[] { 3L, 11L, 21L }, new Double[] { 600.0 }); - starTreeDocuments2[5] = new StarTreeDocument(new Long[] { 3L, 12L, 23L }, new Double[] { 200.0 }); - starTreeDocuments2[6] = new StarTreeDocument(new Long[] { 3L, 12L, 21L }, new Double[] { 400.0 }); - - StarTreeDocument[] segmentStarTreeDocuments2 = new StarTreeDocument[noOfStarTreeDocuments2]; - for (int i = 0; i < noOfStarTreeDocuments2; i++) { - long metric1 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments2[i].metrics[0]); - segmentStarTreeDocuments2[i] = new StarTreeDocument(starTreeDocuments2[i].dimensions, new Long[] { metric1 }); - } - - SequentialDocValuesIterator[] dimsIterators2 = getDimensionIterators(segmentStarTreeDocuments2); - List metricsIterators2 = getMetricIterators(segmentStarTreeDocuments2); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator segmentStarTreeDocumentIterator2 = builder.sortAndAggregateSegmentDocuments( - dimsIterators2, - metricsIterators2 - ); - builder.build(segmentStarTreeDocumentIterator2, new AtomicInteger(), mock(DocValuesConsumer.class)); - InMemoryTreeNode rootNode2 = builder.getRootNode(); - - metaOut.close(); - dataOut.close(); - - StarTreeMetadata starTreeMetadata = new StarTreeMetadata( - "test", - STAR_TREE, - mock(IndexInput.class), - VERSION_CURRENT, - builder.numStarTreeNodes, - List.of("field1", "field3", "field5", "field8"), - List.of( - new Metric("field2", List.of(MetricStat.SUM)), - new Metric("field4", List.of(MetricStat.SUM)), - new Metric("field6", List.of(MetricStat.VALUE_COUNT)), - new Metric("field9", List.of(MetricStat.MIN)), - new Metric("field10", List.of(MetricStat.MAX)) - ), - 2, - getExpectedStarTreeDocumentIterator().size(), - 1, - Set.of("field8"), - getBuildMode(), - 0, - 330 - ); - - StarTreeMetadata starTreeMetadata2 = new StarTreeMetadata( - "test", - STAR_TREE, - mock(IndexInput.class), - VERSION_CURRENT, - builder.numStarTreeNodes, - List.of("fieldC", "fieldB", "fieldL"), - List.of(new Metric("fieldI", List.of(MetricStat.SUM))), - 7, - 27, - 1, - Set.of(), - getBuildMode(), - 330, - 1287 - ); - - List totalDimensionFields = new ArrayList<>(); - totalDimensionFields.addAll(starTreeMetadata.getDimensionFields()); - totalDimensionFields.addAll(starTreeMetadata2.getDimensionFields()); - - List metrics = new ArrayList<>(); - metrics.addAll(starTreeMetadata.getMetrics()); - metrics.addAll(starTreeMetadata2.getMetrics()); - - SegmentReadState readState = getReadState(3, totalDimensionFields, metrics); - - IndexInput dataIn = readState.directory.openInput(dataFileName, IOContext.DEFAULT); - IndexInput metaIn = readState.directory.openInput(metaFileName, IOContext.DEFAULT); - - validateFileFormats(dataIn, metaIn, rootNode1, starTreeMetadata); - validateFileFormats(dataIn, metaIn, rootNode2, starTreeMetadata2); - - dataIn.close(); - metaIn.close(); - - } - - public void test_build() throws IOException { - - int noOfStarTreeDocuments = 5; - StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - - starTreeDocuments[0] = new StarTreeDocument( - new Long[] { 2L, 4L, 3L, 4L }, - new Object[] { 12.0, 10.0, randomDouble(), 8.0, 20.0, 1L } - ); - starTreeDocuments[1] = new StarTreeDocument( - new Long[] { 3L, 4L, 2L, 1L }, - new Object[] { 10.0, 6.0, randomDouble(), 12.0, 10.0, null } - ); - starTreeDocuments[2] = new StarTreeDocument( - new Long[] { 3L, 4L, 2L, 1L }, - new Object[] { 14.0, 12.0, randomDouble(), 6.0, 24.0, null } - ); - starTreeDocuments[3] = new StarTreeDocument( - new Long[] { 2L, 4L, 3L, 4L }, - new Object[] { 9.0, 4.0, randomDouble(), 9.0, 12.0, null } - ); - starTreeDocuments[4] = new StarTreeDocument( - new Long[] { 3L, 4L, 2L, 1L }, - new Object[] { 11.0, 16.0, randomDouble(), 8.0, 13.0, null } - ); - - StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - for (int i = 0; i < noOfStarTreeDocuments; i++) { - long metric1 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]); - long metric2 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[1]); - long metric3 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]); - long metric4 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[3]); - long metric5 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[4]); - Long metric6 = (Long) starTreeDocuments[i].metrics[5]; - segmentStarTreeDocuments[i] = new StarTreeDocument( - starTreeDocuments[i].dimensions, - new Long[] { metric1, metric2, metric3, metric4, metric5, metric6 } - ); - } - - SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); - List metricsIterators = getMetricIterators(segmentStarTreeDocuments); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( - dimsIterators, - metricsIterators - ); - docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( - writeState, - Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, - Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, - Composite99DocValuesFormat.META_DOC_VALUES_CODEC, - Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION - ); - builder.build(segmentStarTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); - - List resultStarTreeDocuments = builder.getStarTreeDocuments(); - assertEquals(7, resultStarTreeDocuments.size()); - - Iterator expectedStarTreeDocumentIterator = getExpectedStarTreeDocumentIterator().iterator(); - assertStarTreeDocuments(resultStarTreeDocuments, expectedStarTreeDocumentIterator); - - metaOut.close(); - dataOut.close(); - docValuesConsumer.close(); - - StarTreeMetadata starTreeMetadata = new StarTreeMetadata( - "test", - STAR_TREE, - mock(IndexInput.class), - VERSION_CURRENT, - builder.numStarTreeNodes, - getStarTreeDimensionNames(compositeField.getDimensionsOrder()), - compositeField.getMetrics(), - 2, - getExpectedStarTreeDocumentIterator().size(), - 1, - Set.of("field8"), - getBuildMode(), - 0, - 330 - ); - - validateStarTreeFileFormats( - builder.getRootNode(), - getExpectedStarTreeDocumentIterator().size(), - starTreeMetadata, - getExpectedStarTreeDocumentIterator() - ); - } - - private void assertStarTreeDocuments( - List resultStarTreeDocuments, - Iterator expectedStarTreeDocumentIterator - ) { - Iterator resultStarTreeDocumentIterator = resultStarTreeDocuments.iterator(); - while (resultStarTreeDocumentIterator.hasNext() && expectedStarTreeDocumentIterator.hasNext()) { - StarTreeDocument resultStarTreeDocument = resultStarTreeDocumentIterator.next(); - StarTreeDocument expectedStarTreeDocument = expectedStarTreeDocumentIterator.next(); - - assertEquals(expectedStarTreeDocument.dimensions[0], resultStarTreeDocument.dimensions[0]); - assertEquals(expectedStarTreeDocument.dimensions[1], resultStarTreeDocument.dimensions[1]); - assertEquals(expectedStarTreeDocument.dimensions[2], resultStarTreeDocument.dimensions[2]); - assertEquals(expectedStarTreeDocument.dimensions[3], resultStarTreeDocument.dimensions[3]); - assertEquals(expectedStarTreeDocument.metrics[0], resultStarTreeDocument.metrics[0]); - assertEquals(expectedStarTreeDocument.metrics[1], resultStarTreeDocument.metrics[1]); - assertEquals(expectedStarTreeDocument.metrics[2], resultStarTreeDocument.metrics[2]); - assertEquals(expectedStarTreeDocument.metrics[3], resultStarTreeDocument.metrics[3]); - assertEquals(expectedStarTreeDocument.metrics[4], resultStarTreeDocument.metrics[4]); - } - } - - public void test_build_starTreeDataset() throws IOException { - - fields = List.of("fieldC", "fieldB", "fieldL", "fieldI"); - - dimensionsOrder = List.of(new NumericDimension("fieldC"), new NumericDimension("fieldB"), new NumericDimension("fieldL")); - metrics = List.of(new Metric("fieldI", List.of(MetricStat.SUM)), new Metric("_doc_count", List.of(MetricStat.DOC_COUNT))); - - DocValuesProducer docValuesProducer = mock(DocValuesProducer.class); - - compositeField = new StarTreeField("test", dimensionsOrder, metrics, new StarTreeFieldConfiguration(1, Set.of(), getBuildMode())); - SegmentInfo segmentInfo = new SegmentInfo( - directory, - Version.LATEST, - Version.LUCENE_9_11_0, - "test_segment", - 7, - false, - false, - new Lucene912Codec(), - new HashMap<>(), - UUID.randomUUID().toString().substring(0, 16).getBytes(StandardCharsets.UTF_8), - new HashMap<>(), - null - ); - - fieldsInfo = new FieldInfo[fields.size()]; - fieldProducerMap = new HashMap<>(); - for (int i = 0; i < fieldsInfo.length; i++) { - fieldsInfo[i] = new FieldInfo( - fields.get(i), - i, - false, - false, - true, - IndexOptions.DOCS_AND_FREQS_AND_POSITIONS_AND_OFFSETS, - DocValuesType.SORTED_NUMERIC, - -1, - Collections.emptyMap(), - 0, - 0, - 0, - 0, - VectorEncoding.FLOAT32, - VectorSimilarityFunction.EUCLIDEAN, - false, - false - ); - fieldProducerMap.put(fields.get(i), docValuesProducer); - } - FieldInfos fieldInfos = new FieldInfos(fieldsInfo); - writeState = new SegmentWriteState(InfoStream.getDefault(), segmentInfo.dir, segmentInfo, fieldInfos, null, newIOContext(random())); - this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( - writeState, - Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, - Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, - Composite99DocValuesFormat.META_DOC_VALUES_CODEC, - Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION - ); - mapperService = mock(MapperService.class); - 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("fieldI", NumberFieldMapper.NumberType.DOUBLE, false, true) - .build(new Mapper.BuilderContext(settings, new ContentPath())); - MappingLookup fieldMappers = new MappingLookup( - Set.of(numberFieldMapper1), - Collections.emptyList(), - Collections.emptyList(), - 0, - null - ); - when(documentMapper.mappers()).thenReturn(fieldMappers); - - int noOfStarTreeDocuments = 7; - StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - starTreeDocuments[0] = new StarTreeDocument(new Long[] { 1L, 11L, 21L }, new Object[] { 400.0, null }); - starTreeDocuments[1] = new StarTreeDocument(new Long[] { 1L, 12L, 22L }, new Object[] { 200.0, null }); - starTreeDocuments[2] = new StarTreeDocument(new Long[] { 2L, 13L, 23L }, new Object[] { 300.0, null }); - starTreeDocuments[3] = new StarTreeDocument(new Long[] { 2L, 13L, 21L }, new Object[] { 100.0, null }); - starTreeDocuments[4] = new StarTreeDocument(new Long[] { 3L, 11L, 21L }, new Object[] { 600.0, null }); - starTreeDocuments[5] = new StarTreeDocument(new Long[] { 3L, 12L, 23L }, new Object[] { 200.0, null }); - starTreeDocuments[6] = new StarTreeDocument(new Long[] { 3L, 12L, 21L }, new Object[] { 400.0, null }); - - StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; - for (int i = 0; i < noOfStarTreeDocuments; i++) { - long metric1 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]); - segmentStarTreeDocuments[i] = new StarTreeDocument(starTreeDocuments[i].dimensions, new Long[] { metric1, null }); - } - - SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); - List metricsIterators = getMetricIterators(segmentStarTreeDocuments); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( - dimsIterators, - metricsIterators - ); - builder.build(segmentStarTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); - - List resultStarTreeDocuments = builder.getStarTreeDocuments(); - Iterator expectedStarTreeDocumentIterator = expectedStarTreeDocuments().iterator(); - Iterator resultStarTreeDocumentIterator = resultStarTreeDocuments.iterator(); - Map> dimValueToDocIdMap = new HashMap<>(); - builder.rootNode.setNodeType(StarTreeNodeType.STAR.getValue()); - traverseStarTree(builder.rootNode, dimValueToDocIdMap, true); - - Map> expectedDimToValueMap = getExpectedDimToValueMap(); - for (Map.Entry> entry : dimValueToDocIdMap.entrySet()) { - int dimId = entry.getKey(); - if (dimId == -1) continue; - Map map = expectedDimToValueMap.get(dimId); - for (Map.Entry dimValueToDocIdEntry : entry.getValue().entrySet()) { - long dimValue = dimValueToDocIdEntry.getKey(); - int docId = dimValueToDocIdEntry.getValue(); - if (map.get(dimValue) != null) { - assertEquals(map.get(dimValue), resultStarTreeDocuments.get(docId).metrics[0]); - } - } - } - - while (resultStarTreeDocumentIterator.hasNext() && expectedStarTreeDocumentIterator.hasNext()) { - StarTreeDocument resultStarTreeDocument = resultStarTreeDocumentIterator.next(); - StarTreeDocument expectedStarTreeDocument = expectedStarTreeDocumentIterator.next(); - assertEquals(expectedStarTreeDocument.dimensions[0], resultStarTreeDocument.dimensions[0]); - assertEquals(expectedStarTreeDocument.dimensions[1], resultStarTreeDocument.dimensions[1]); - assertEquals(expectedStarTreeDocument.dimensions[2], resultStarTreeDocument.dimensions[2]); - assertEquals(expectedStarTreeDocument.metrics[0], resultStarTreeDocument.metrics[0]); - assertEquals(expectedStarTreeDocument.metrics[1], resultStarTreeDocument.metrics[1]); - } - - metaOut.close(); - dataOut.close(); - docValuesConsumer.close(); - - validateStarTree(builder.getRootNode(), 3, 1, builder.getStarTreeDocuments()); - - StarTreeMetadata starTreeMetadata = new StarTreeMetadata( - "test", - STAR_TREE, - mock(IndexInput.class), - VERSION_CURRENT, - builder.numStarTreeNodes, - getStarTreeDimensionNames(compositeField.getDimensionsOrder()), - compositeField.getMetrics(), - 7, - 27, - 1, - Set.of(), - getBuildMode(), - 0, - 1287 - ); - validateStarTreeFileFormats(builder.getRootNode(), 27, starTreeMetadata, expectedStarTreeDocuments()); - } - - private List getStarTreeDimensionNames(List dimensionsOrder) { - - List dimensionNames = new ArrayList<>(); - for (Dimension dimension : dimensionsOrder) { - dimensionNames.addAll(dimension.getSubDimensionNames()); - } - return dimensionNames; - - } - - private void validateStarTreeFileFormats( - InMemoryTreeNode rootNode, - int numDocs, - StarTreeMetadata expectedStarTreeMetadata, - List expectedStarTreeDocuments - ) throws IOException { - - assertNotNull(rootNode.getChildren()); - assertFalse(rootNode.getChildren().isEmpty()); - SegmentReadState readState = getReadState( - numDocs, - expectedStarTreeMetadata.getDimensionFields(), - expectedStarTreeMetadata.getMetrics() - ); - - DocValuesProducer compositeDocValuesProducer = LuceneDocValuesProducerFactory.getDocValuesProducerForCompositeCodec( - Composite99Codec.COMPOSITE_INDEX_CODEC_NAME, - readState, - Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, - Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, - Composite99DocValuesFormat.META_DOC_VALUES_CODEC, - Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION - ); - - IndexInput dataIn = readState.directory.openInput(dataFileName, IOContext.DEFAULT); - IndexInput metaIn = readState.directory.openInput(metaFileName, IOContext.DEFAULT); - - StarTreeValues starTreeValues = new StarTreeValues(expectedStarTreeMetadata, dataIn, compositeDocValuesProducer, readState); - assertEquals(expectedStarTreeMetadata.getStarTreeDocCount(), starTreeValues.getStarTreeDocumentCount()); - List fieldValueConverters = new ArrayList<>(); - builder.metricAggregatorInfos.forEach( - metricAggregatorInfo -> fieldValueConverters.add(metricAggregatorInfo.getValueAggregators().getAggregatedValueType()) - ); - StarTreeDocument[] starTreeDocuments = StarTreeTestUtils.getSegmentsStarTreeDocuments( - List.of(starTreeValues), - fieldValueConverters, - readState.segmentInfo.maxDoc() - ); - - StarTreeDocument[] expectedStarTreeDocumentsArray = expectedStarTreeDocuments.toArray(new StarTreeDocument[0]); - StarTreeTestUtils.assertStarTreeDocuments(starTreeDocuments, expectedStarTreeDocumentsArray); - - validateFileFormats(dataIn, metaIn, rootNode, expectedStarTreeMetadata); - - dataIn.close(); - metaIn.close(); - compositeDocValuesProducer.close(); - } - - private static Map> getExpectedDimToValueMap() { - Map> expectedDimToValueMap = new HashMap<>(); - Map dimValueMap = new HashMap<>(); - dimValueMap.put(1L, 600.0); - dimValueMap.put(2L, 400.0); - dimValueMap.put(3L, 1200.0); - expectedDimToValueMap.put(0, dimValueMap); - - dimValueMap = new HashMap<>(); - dimValueMap.put(11L, 1000.0); - dimValueMap.put(12L, 800.0); - dimValueMap.put(13L, 400.0); - expectedDimToValueMap.put(1, dimValueMap); - - dimValueMap = new HashMap<>(); - dimValueMap.put(21L, 1500.0); - dimValueMap.put(22L, 200.0); - dimValueMap.put(23L, 500.0); - expectedDimToValueMap.put(2, dimValueMap); - return expectedDimToValueMap; - } - - private List expectedStarTreeDocuments() { - return List.of( - new StarTreeDocument(new Long[] { 1L, 11L, 21L }, new Object[] { 400.0, 1L }), - new StarTreeDocument(new Long[] { 1L, 12L, 22L }, new Object[] { 200.0, 1L }), - new StarTreeDocument(new Long[] { 2L, 13L, 21L }, new Object[] { 100.0, 1L }), - new StarTreeDocument(new Long[] { 2L, 13L, 23L }, new Object[] { 300.0, 1L }), - new StarTreeDocument(new Long[] { 3L, 11L, 21L }, new Object[] { 600.0, 1L }), - new StarTreeDocument(new Long[] { 3L, 12L, 21L }, new Object[] { 400.0, 1L }), - new StarTreeDocument(new Long[] { 3L, 12L, 23L }, new Object[] { 200.0, 1L }), - new StarTreeDocument(new Long[] { null, 11L, 21L }, new Object[] { 1000.0, 2L }), - new StarTreeDocument(new Long[] { null, 12L, 21L }, new Object[] { 400.0, 1L }), - new StarTreeDocument(new Long[] { null, 12L, 22L }, new Object[] { 200.0, 1L }), - new StarTreeDocument(new Long[] { null, 12L, 23L }, new Object[] { 200.0, 1L }), - new StarTreeDocument(new Long[] { null, 13L, 21L }, new Object[] { 100.0, 1L }), - new StarTreeDocument(new Long[] { null, 13L, 23L }, new Object[] { 300.0, 1L }), - new StarTreeDocument(new Long[] { null, null, 21L }, new Object[] { 1500.0, 4L }), - new StarTreeDocument(new Long[] { null, null, 22L }, new Object[] { 200.0, 1L }), - new StarTreeDocument(new Long[] { null, null, 23L }, new Object[] { 500.0, 2L }), - new StarTreeDocument(new Long[] { null, null, null }, new Object[] { 2200.0, 7L }), - new StarTreeDocument(new Long[] { null, 12L, null }, new Object[] { 800.0, 3L }), - new StarTreeDocument(new Long[] { null, 13L, null }, new Object[] { 400.0, 2L }), - new StarTreeDocument(new Long[] { 1L, null, 21L }, new Object[] { 400.0, 1L }), - new StarTreeDocument(new Long[] { 1L, null, 22L }, new Object[] { 200.0, 1L }), - new StarTreeDocument(new Long[] { 1L, null, null }, new Object[] { 600.0, 2L }), - new StarTreeDocument(new Long[] { 2L, 13L, null }, new Object[] { 400.0, 2L }), - new StarTreeDocument(new Long[] { 3L, null, 21L }, new Object[] { 1000.0, 2L }), - new StarTreeDocument(new Long[] { 3L, null, 23L }, new Object[] { 200.0, 1L }), - new StarTreeDocument(new Long[] { 3L, null, null }, new Object[] { 1200.0, 3L }), - new StarTreeDocument(new Long[] { 3L, 12L, null }, new Object[] { 600.0, 2L }) - ); - - } - - public void testFlushFlow() throws IOException { - List dimList = List.of(0L, 1L, 3L, 4L, 5L); - List docsWithField = List.of(0, 1, 3, 4, 5); - List dimList2 = List.of(0L, 1L, 2L, 3L, 4L, 5L); - List docsWithField2 = List.of(0, 1, 2, 3, 4, 5); - - List metricsList = List.of( - getLongFromDouble(0.0), - getLongFromDouble(10.0), - getLongFromDouble(20.0), - getLongFromDouble(30.0), - getLongFromDouble(40.0), - getLongFromDouble(50.0) - ); - List metricsWithField = List.of(0, 1, 2, 3, 4, 5); - - compositeField = getStarTreeFieldWithMultipleMetrics(); - 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) - ); - - writeState = getWriteState(6, writeState.segmentInfo.getId()); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - SequentialDocValuesIterator[] dimDvs = { new SequentialDocValuesIterator(d1sndv), new SequentialDocValuesIterator(d2sndv) }; - Iterator starTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( - dimDvs, - List.of(new SequentialDocValuesIterator(m1sndv), new SequentialDocValuesIterator(m2sndv)) - ); - /** - * Asserting following dim / metrics [ dim1, dim2 / Sum [metric], count [metric] ] - [0, 0] | [0.0, 1] - [1, 1] | [10.0, 1] - [3, 3] | [30.0, 1] - [4, 4] | [40.0, 1] - [5, 5] | [50.0, 1] - [null, 2] | [20.0, 1] - */ - this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( - writeState, - Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, - Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, - Composite99DocValuesFormat.META_DOC_VALUES_CODEC, - Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION - ); - builder.build(starTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); - List starTreeDocuments = builder.getStarTreeDocuments(); - int count = 0; - for (StarTreeDocument starTreeDocument : starTreeDocuments) { - count++; - if (starTreeDocument.dimensions[1] != null) { - assertEquals( - starTreeDocument.dimensions[0] == null - ? starTreeDocument.dimensions[1] * 1 * 10.0 - : starTreeDocument.dimensions[0] * 10, - starTreeDocument.metrics[0] - ); - assertEquals(1L, starTreeDocument.metrics[1]); - } else { - assertEquals(150D, starTreeDocument.metrics[0]); - assertEquals(6L, starTreeDocument.metrics[1]); - } - } - assertEquals(13, count); - validateStarTree(builder.getRootNode(), 2, 1000, builder.getStarTreeDocuments()); - - metaOut.close(); - dataOut.close(); - docValuesConsumer.close(); - - StarTreeMetadata starTreeMetadata = new StarTreeMetadata( - "sf", - STAR_TREE, - mock(IndexInput.class), - VERSION_CURRENT, - builder.numStarTreeNodes, - List.of("field1", "field3"), - List.of(new Metric("field2", List.of(MetricStat.SUM, MetricStat.VALUE_COUNT, MetricStat.AVG))), - 6, - builder.numStarTreeDocs, - 1000, - Set.of(), - getBuildMode(), - 0, - 264 - ); - - validateStarTreeFileFormats( - builder.getRootNode(), - builder.getStarTreeDocuments().size(), - starTreeMetadata, - builder.getStarTreeDocuments() - ); - - } - - public void testFlushFlowDimsReverse() throws IOException { - List dimList = List.of(5L, 4L, 3L, 2L, 1L); - List docsWithField = List.of(0, 1, 2, 3, 4); - List dimList2 = List.of(5L, 4L, 3L, 2L, 1L, 0L); - List docsWithField2 = List.of(0, 1, 2, 3, 4, 5); - - List metricsList = List.of( - getLongFromDouble(50.0), - getLongFromDouble(40.0), - getLongFromDouble(30.0), - getLongFromDouble(20.0), - getLongFromDouble(10.0), - getLongFromDouble(0.0) - ); - List metricsWithField = List.of(0, 1, 2, 3, 4, 5); - - compositeField = getStarTreeFieldWithMultipleMetrics(); - 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) - ); - - writeState = getWriteState(6, writeState.segmentInfo.getId()); - this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( - writeState, - Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, - Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, - Composite99DocValuesFormat.META_DOC_VALUES_CODEC, - Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION - ); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - SequentialDocValuesIterator[] dimDvs = { new SequentialDocValuesIterator(d1sndv), new SequentialDocValuesIterator(d2sndv) }; - Iterator starTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( - dimDvs, - List.of(new SequentialDocValuesIterator(m1sndv), new SequentialDocValuesIterator(m2sndv)) - ); - /** - * Asserting following dim / metrics [ dim1, dim2 / Sum [metric], count [metric] ] - [1, 1] | [10.0, 1] - [2, 2] | [20.0, 1] - [3, 3] | [30.0, 1] - [4, 4] | [40.0, 1] - [5, 5] | [50.0, 1] - [null, 0] | [0.0, 1] - */ - builder.appendDocumentsToStarTree(starTreeDocumentIterator); - assertEquals(6, builder.getStarTreeDocuments().size()); - builder.build(starTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); - int count = 0; - for (StarTreeDocument starTreeDocument : builder.getStarTreeDocuments()) { - if (count <= 6) { - count++; - if (starTreeDocument.dimensions[0] != null) { - assertEquals(count, (long) starTreeDocument.dimensions[0]); - } - assertEquals(starTreeDocument.dimensions[1] * 10.0, starTreeDocument.metrics[0]); - assertEquals(1L, starTreeDocument.metrics[1]); - } - } - validateStarTree(builder.getRootNode(), 2, 1000, builder.getStarTreeDocuments()); - - metaOut.close(); - dataOut.close(); - docValuesConsumer.close(); - - StarTreeMetadata starTreeMetadata = new StarTreeMetadata( - "sf", - STAR_TREE, - mock(IndexInput.class), - VERSION_CURRENT, - builder.numStarTreeNodes, - List.of("field1", "field3"), - List.of(new Metric("field2", List.of(MetricStat.SUM, MetricStat.VALUE_COUNT, MetricStat.AVG))), - 6, - builder.numStarTreeDocs, - 1000, - Set.of(), - getBuildMode(), - 0, - 264 - ); - - validateStarTreeFileFormats( - builder.getRootNode(), - builder.getStarTreeDocuments().size(), - starTreeMetadata, - builder.getStarTreeDocuments() - ); - } - - public void testFlushFlowBuild() throws IOException { - List dimList = new ArrayList<>(100); - List docsWithField = new ArrayList<>(100); - for (int i = 0; i < 100; i++) { - dimList.add((long) i); - docsWithField.add(i); - } - - List dimList2 = new ArrayList<>(100); - List docsWithField2 = new ArrayList<>(100); - for (int i = 0; i < 100; i++) { - dimList2.add((long) i); - docsWithField2.add(i); - } - - List metricsList = new ArrayList<>(100); - List metricsWithField = new ArrayList<>(100); - for (int i = 0; i < 100; i++) { - metricsList.add(getLongFromDouble(i * 10.0)); - metricsWithField.add(i); - } - - Dimension d1 = new NumericDimension("field1"); - Dimension d2 = new NumericDimension("field3"); - Metric m1 = new Metric("field2", List.of(MetricStat.SUM)); - List dims = List.of(d1, d2); - List metrics = List.of(m1); - StarTreeFieldConfiguration c = new StarTreeFieldConfiguration(1, new HashSet<>(), getBuildMode()); - compositeField = new StarTreeField("sf", dims, metrics, c); - SortedNumericDocValues d1sndv = getSortedNumericMock(dimList, docsWithField); - SortedNumericDocValues d2sndv = getSortedNumericMock(dimList2, docsWithField2); - SortedNumericDocValues m1sndv = getSortedNumericMock(metricsList, metricsWithField); - - writeState = getWriteState(100, writeState.segmentInfo.getId()); - SegmentWriteState consumerWriteState = getWriteState(DocIdSetIterator.NO_MORE_DOCS, writeState.segmentInfo.getId()); - this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( - consumerWriteState, - Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, - Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, - Composite99DocValuesFormat.META_DOC_VALUES_CODEC, - Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION - ); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - - DocValuesProducer d1vp = getDocValuesProducer(d1sndv); - DocValuesProducer d2vp = getDocValuesProducer(d2sndv); - DocValuesProducer m1vp = getDocValuesProducer(m1sndv); - Map fieldProducerMap = Map.of("field1", d1vp, "field3", d2vp, "field2", m1vp); - builder.build(fieldProducerMap, new AtomicInteger(), docValuesConsumer); - /** - * Asserting following dim / metrics [ dim1, dim2 / Sum [ metric] ] - [0, 0] | [0.0] - [1, 1] | [10.0] - [2, 2] | [20.0] - [3, 3] | [30.0] - [4, 4] | [40.0] - .... - [null, 0] | [0.0] - [null, 1] | [10.0] - ... - [null, null] | [49500.0] - */ - List starTreeDocuments = builder.getStarTreeDocuments(); - for (StarTreeDocument starTreeDocument : starTreeDocuments) { - assertEquals( - starTreeDocument.dimensions[1] != null ? starTreeDocument.dimensions[1] * 10.0 : 49500.0, - starTreeDocument.metrics[0] - ); - } - validateStarTree(builder.getRootNode(), 2, 1, builder.getStarTreeDocuments()); - - metaOut.close(); - dataOut.close(); - docValuesConsumer.close(); - - StarTreeMetadata starTreeMetadata = new StarTreeMetadata( - "sf", - STAR_TREE, - mock(IndexInput.class), - VERSION_CURRENT, - builder.numStarTreeNodes, - List.of("field1", "field3"), - compositeField.getMetrics(), - 100, - builder.numStarTreeDocs, - 1, - Set.of(), - getBuildMode(), - 0, - 6699 - ); - - validateStarTreeFileFormats( - builder.getRootNode(), - builder.getStarTreeDocuments().size(), - starTreeMetadata, - builder.getStarTreeDocuments() - ); - } - - private static DocValuesProducer getDocValuesProducer(SortedNumericDocValues sndv) { - return new EmptyDocValuesProducer() { - @Override - public SortedNumericDocValues getSortedNumeric(FieldInfo field) throws IOException { - return sndv; - } - }; - } - - private StarTreeField getStarTreeFieldWithMultipleMetrics() { - Dimension d1 = new NumericDimension("field1"); - Dimension d2 = new NumericDimension("field3"); - Metric m1 = new Metric("field2", List.of(MetricStat.SUM)); - Metric m2 = new Metric("field2", List.of(MetricStat.VALUE_COUNT)); - Metric m3 = new Metric("field2", List.of(MetricStat.AVG)); - List dims = List.of(d1, d2); - List metrics = List.of(m1, m2, m3); - StarTreeFieldConfiguration c = new StarTreeFieldConfiguration(1000, new HashSet<>(), getBuildMode()); - return new StarTreeField("sf", dims, metrics, c); - } - - 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(); - } - - public void testMergeFlowWithSum() throws IOException { - List dimList = List.of(0L, 1L, 3L, 4L, 5L, 6L); - List docsWithField = List.of(0, 1, 3, 4, 5, 6); - List dimList2 = List.of(0L, 1L, 2L, 3L, 4L, 5L, -1L); - List docsWithField2 = List.of(0, 1, 2, 3, 4, 5, 6); - - List metricsList = List.of( - getLongFromDouble(0.0), - getLongFromDouble(10.0), - getLongFromDouble(20.0), - getLongFromDouble(30.0), - getLongFromDouble(40.0), - getLongFromDouble(50.0), - getLongFromDouble(60.0) - - ); - List metricsWithField = List.of(0, 1, 2, 3, 4, 5, 6); - - compositeField = getStarTreeField(MetricStat.SUM); - StarTreeValues starTreeValues = getStarTreeValues( - getSortedNumericMock(dimList, docsWithField), - getSortedNumericMock(dimList2, docsWithField2), - getSortedNumericMock(metricsList, metricsWithField), - compositeField, - "6" - ); - - StarTreeValues starTreeValues2 = getStarTreeValues( - getSortedNumericMock(dimList, docsWithField), - getSortedNumericMock(dimList2, docsWithField2), - getSortedNumericMock(metricsList, metricsWithField), - compositeField, - "6" - ); - writeState = getWriteState(6, writeState.segmentInfo.getId()); - this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( - writeState, - Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, - Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, - Composite99DocValuesFormat.META_DOC_VALUES_CODEC, - Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION - ); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); - /** - * Asserting following dim / metrics [ dim1, dim2 / Sum [ metric] ] - * [0, 0] | [0.0] - * [1, 1] | [20.0] - * [3, 3] | [60.0] - * [4, 4] | [80.0] - * [5, 5] | [100.0] - * [null, 2] | [40.0] - * ------------------ We only take non star docs - * [6,-1] | [120.0] - */ - builder.appendDocumentsToStarTree(starTreeDocumentIterator); - assertEquals(6, builder.getStarTreeDocuments().size()); - builder.build(starTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); - int count = 0; - for (StarTreeDocument starTreeDocument : builder.getStarTreeDocuments()) { - count++; - if (count <= 6) { - assertEquals( - starTreeDocument.dimensions[0] != null ? starTreeDocument.dimensions[0] * 2 * 10.0 : 40.0, - starTreeDocument.metrics[0] - ); - } - } - - validateStarTree(builder.getRootNode(), 2, 1000, builder.getStarTreeDocuments()); - - metaOut.close(); - dataOut.close(); - docValuesConsumer.close(); - - StarTreeMetadata starTreeMetadata = new StarTreeMetadata( - "sf", - STAR_TREE, - mock(IndexInput.class), - VERSION_CURRENT, - builder.numStarTreeNodes, - List.of("field1", "field3"), - compositeField.getMetrics(), - 6, - builder.numStarTreeDocs, - 1000, - Set.of(), - getBuildMode(), - 0, - 264 - ); - - validateStarTreeFileFormats( - builder.getRootNode(), - builder.getStarTreeDocuments().size(), - starTreeMetadata, - builder.getStarTreeDocuments() - ); - } - - public void testMergeFlowWithCount() throws IOException { - List dimList = List.of(0L, 1L, 3L, 4L, 5L, 6L); - List docsWithField = List.of(0, 1, 3, 4, 5, 6); - List dimList2 = List.of(0L, 1L, 2L, 3L, 4L, 5L, -1L); - List docsWithField2 = List.of(0, 1, 2, 3, 4, 5, 6); - - List metricsList = List.of(0L, 1L, 2L, 3L, 4L, 5L, 6L); - List metricsWithField = List.of(0, 1, 2, 3, 4, 5, 6); - - compositeField = getStarTreeField(MetricStat.VALUE_COUNT); - StarTreeValues starTreeValues = getStarTreeValues( - getSortedNumericMock(dimList, docsWithField), - getSortedNumericMock(dimList2, docsWithField2), - getSortedNumericMock(metricsList, metricsWithField), - compositeField, - "6" - ); - - StarTreeValues starTreeValues2 = getStarTreeValues( - getSortedNumericMock(dimList, docsWithField), - getSortedNumericMock(dimList2, docsWithField2), - getSortedNumericMock(metricsList, metricsWithField), - compositeField, - "6" - ); - writeState = getWriteState(6, writeState.segmentInfo.getId()); - this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( - writeState, - Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, - Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, - Composite99DocValuesFormat.META_DOC_VALUES_CODEC, - Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION - ); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); - /** - * Asserting following dim / metrics [ dim1, dim2 / Count [ metric] ] - [0, 0] | [0] - [1, 1] | [2] - [3, 3] | [6] - [4, 4] | [8] - [5, 5] | [10] - [null, 2] | [4] - --------------- - [6,-1] | [12] - */ - builder.appendDocumentsToStarTree(starTreeDocumentIterator); - assertEquals(6, builder.getStarTreeDocuments().size()); - builder.build(starTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); - int count = 0; - for (StarTreeDocument starTreeDocument : builder.getStarTreeDocuments()) { - count++; - if (count <= 6) { - assertEquals(starTreeDocument.dimensions[0] != null ? starTreeDocument.dimensions[0] * 2 : 4, starTreeDocument.metrics[0]); - } - } - - validateStarTree(builder.getRootNode(), 2, 1000, builder.getStarTreeDocuments()); - - metaOut.close(); - dataOut.close(); - docValuesConsumer.close(); - - StarTreeMetadata starTreeMetadata = new StarTreeMetadata( - "sf", - STAR_TREE, - mock(IndexInput.class), - VERSION_CURRENT, - builder.numStarTreeNodes, - List.of("field1", "field3"), - compositeField.getMetrics(), - 6, - builder.numStarTreeDocs, - 1000, - Set.of(), - getBuildMode(), - 0, - 264 - ); - - validateStarTreeFileFormats( - builder.getRootNode(), - builder.getStarTreeDocuments().size(), - starTreeMetadata, - builder.getStarTreeDocuments() - ); - - } - - private StarTreeValues getStarTreeValues( - SortedNumericDocValues dimList, - SortedNumericDocValues dimList2, - SortedNumericDocValues metricsList, - StarTreeField sf, - String number - ) { - SortedNumericDocValues d1sndv = dimList; - SortedNumericDocValues d2sndv = dimList2; - SortedNumericDocValues m1sndv = metricsList; - Map> dimDocIdSetIterators = Map.of( - "field1", - () -> new SortedNumericStarTreeValuesIterator(d1sndv), - "field3", - () -> new SortedNumericStarTreeValuesIterator(d2sndv) - ); - - Map> metricDocIdSetIterators = new LinkedHashMap<>(); - for (Metric metric : sf.getMetrics()) { - for (MetricStat metricStat : metric.getMetrics()) { - String metricFullName = fullyQualifiedFieldNameForStarTreeMetricsDocValues( - sf.getName(), - metric.getField(), - metricStat.getTypeName() - ); - metricDocIdSetIterators.put(metricFullName, () -> new SortedNumericStarTreeValuesIterator(m1sndv)); - } - } - - StarTreeValues starTreeValues = new StarTreeValues( - sf, - null, - dimDocIdSetIterators, - metricDocIdSetIterators, - Map.of(CompositeIndexConstants.SEGMENT_DOCS_COUNT, number), - null - ); - return starTreeValues; - } - - public void testMergeFlowWithDifferentDocsFromSegments() throws IOException { - List dimList = List.of(0L, 1L, 3L, 4L, 5L, 6L); - List docsWithField = List.of(0, 1, 3, 4, 5, 6); - List dimList2 = List.of(0L, 1L, 2L, 3L, 4L, 5L, -1L); - List docsWithField2 = List.of(0, 1, 2, 3, 4, 5, 6); - - List metricsList = List.of(0L, 1L, 2L, 3L, 4L, 5L, 6L); - List metricsWithField = List.of(0, 1, 2, 3, 4, 5, 6); - - List dimList3 = List.of(5L, 6L, 8L, -1L); - List docsWithField3 = List.of(0, 1, 3, 4); - List dimList4 = List.of(5L, 6L, 7L, 8L, -1L); - List docsWithField4 = List.of(0, 1, 2, 3, 4); - - List metricsList2 = List.of(5L, 6L, 7L, 8L, 9L); - List metricsWithField2 = List.of(0, 1, 2, 3, 4); - - compositeField = getStarTreeField(MetricStat.VALUE_COUNT); - StarTreeValues starTreeValues = getStarTreeValues( - getSortedNumericMock(dimList, docsWithField), - getSortedNumericMock(dimList2, docsWithField2), - getSortedNumericMock(metricsList, metricsWithField), - compositeField, - "6" - ); - - StarTreeValues starTreeValues2 = getStarTreeValues( - getSortedNumericMock(dimList3, docsWithField3), - getSortedNumericMock(dimList4, docsWithField4), - getSortedNumericMock(metricsList2, metricsWithField2), - compositeField, - "4" - ); - writeState = getWriteState(4, writeState.segmentInfo.getId()); - this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( - writeState, - Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, - Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, - Composite99DocValuesFormat.META_DOC_VALUES_CODEC, - Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION - ); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); - /** - * Asserting following dim / metrics [ dim1, dim2 / Count [ metric] ] - [0, 0] | [0] - [1, 1] | [1] - [3, 3] | [3] - [4, 4] | [4] - [5, 5] | [10] - [6, 6] | [6] - [8, 8] | [8] - [null, 2] | [2] - [null, 7] | [7] - */ - builder.appendDocumentsToStarTree(starTreeDocumentIterator); - assertEquals(9, builder.getStarTreeDocuments().size()); - builder.build(starTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); - int count = 0; - for (StarTreeDocument starTreeDocument : builder.getStarTreeDocuments()) { - count++; - if (count <= 9) { - if (Objects.equals(starTreeDocument.dimensions[0], 5L)) { - assertEquals(starTreeDocument.dimensions[0] * 2, starTreeDocument.metrics[0]); - } else { - assertEquals(starTreeDocument.dimensions[1], starTreeDocument.metrics[0]); - } - } - } - validateStarTree(builder.getRootNode(), 2, 1000, builder.getStarTreeDocuments()); - - metaOut.close(); - dataOut.close(); - docValuesConsumer.close(); - - StarTreeMetadata starTreeMetadata = new StarTreeMetadata( - "sf", - STAR_TREE, - mock(IndexInput.class), - VERSION_CURRENT, - builder.numStarTreeNodes, - List.of("field1", "field3"), - compositeField.getMetrics(), - 9, - builder.numStarTreeDocs, - 1000, - Set.of(), - getBuildMode(), - 0, - 330 - ); - - validateStarTreeFileFormats( - builder.getRootNode(), - builder.getStarTreeDocuments().size(), - starTreeMetadata, - builder.getStarTreeDocuments() - ); - } - - public void testMergeFlowNumSegmentsDocs() throws IOException { - List dimList = List.of(0L, 1L, 2L, 3L, 4L, 5L, 6L, -1L, -1L, -1L); - List docsWithField = List.of(0, 1, 2, 3, 4, 5, 6, 7, 8, 9); - List dimList2 = List.of(0L, 1L, 2L, 3L, 4L, 5L, 6L, -1L, -1L, -1L); - List docsWithField2 = List.of(0, 1, 2, 3, 4, 5, 6, 7, 8, 9); - - List metricsList = List.of(0L, 1L, 2L, 3L, 4L, 5L, 6L, -1L, -1L, -1L); - List metricsWithField = List.of(0, 1, 2, 3, 4, 5, 6, 7, 8, 9); - - List dimList3 = List.of(5L, 6L, 7L, 8L, -1L); - List docsWithField3 = List.of(0, 1, 2, 3, 4); - List dimList4 = List.of(5L, 6L, 7L, 8L, -1L); - List docsWithField4 = List.of(0, 1, 2, 3, 4); - - List metricsList2 = List.of(5L, 6L, 7L, 8L, 9L); - List metricsWithField2 = List.of(0, 1, 2, 3, 4); - - StarTreeField sf = getStarTreeField(MetricStat.VALUE_COUNT); - StarTreeValues starTreeValues = getStarTreeValues( - getSortedNumericMock(dimList, docsWithField), - getSortedNumericMock(dimList2, docsWithField2), - getSortedNumericMock(metricsList, metricsWithField), - sf, - "6" - ); - - StarTreeValues starTreeValues2 = getStarTreeValues( - getSortedNumericMock(dimList3, docsWithField3), - getSortedNumericMock(dimList4, docsWithField4), - getSortedNumericMock(metricsList2, metricsWithField2), - sf, - "4" - ); - builder = getStarTreeBuilder(metaOut, dataOut, sf, getWriteState(4, writeState.segmentInfo.getId()), mapperService); - Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); - /** - * Asserting following dim / metrics [ dim1, dim2 / Count [ metric] ] - [0, 0] | [0] - [1, 1] | [1] - [2, 2] | [2] - [3, 3] | [3] - [4, 4] | [4] - [5, 5] | [10] - [6, 6] | [6] - [7, 7] | [7] - [8, 8] | [8] - */ - int count = 0; - while (starTreeDocumentIterator.hasNext()) { - count++; - StarTreeDocument starTreeDocument = starTreeDocumentIterator.next(); - if (Objects.equals(starTreeDocument.dimensions[0], 5L)) { - assertEquals(starTreeDocument.dimensions[0] * 2, starTreeDocument.metrics[0]); - } else { - assertEquals(starTreeDocument.dimensions[1], starTreeDocument.metrics[0]); - } - } - assertEquals(9, count); - } - - public void testMergeFlowWithMissingDocs() throws IOException { - List dimList = List.of(0L, 1L, 2L, 3L, 4L, 6L); - List docsWithField = List.of(0, 1, 2, 3, 4, 6); - List dimList2 = List.of(0L, 1L, 2L, 3L, 4L, 5L, -1L); - List docsWithField2 = List.of(0, 1, 2, 3, 4, 5, 6); - - List metricsList = List.of(0L, 1L, 2L, 3L, 4L, 5L, 6L); - List metricsWithField = List.of(0, 1, 2, 3, 4, 5, 6); - - List dimList3 = List.of(5L, 6L, 8L, -1L); - List docsWithField3 = List.of(0, 1, 3, 4); - List dimList4 = List.of(5L, 6L, 7L, 8L, -1L); - List docsWithField4 = List.of(0, 1, 2, 3, 4); - - List metricsList2 = List.of(5L, 6L, 7L, 8L, 9L); - List metricsWithField2 = List.of(0, 1, 2, 3, 4); - - compositeField = getStarTreeField(MetricStat.VALUE_COUNT); - StarTreeValues starTreeValues = getStarTreeValues( - getSortedNumericMock(dimList, docsWithField), - getSortedNumericMock(dimList2, docsWithField2), - getSortedNumericMock(metricsList, metricsWithField), - compositeField, - "6" - ); - - StarTreeValues starTreeValues2 = getStarTreeValues( - getSortedNumericMock(dimList3, docsWithField3), - getSortedNumericMock(dimList4, docsWithField4), - getSortedNumericMock(metricsList2, metricsWithField2), - compositeField, - "4" - ); - writeState = getWriteState(4, writeState.segmentInfo.getId()); - this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( - writeState, - Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, - Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, - Composite99DocValuesFormat.META_DOC_VALUES_CODEC, - Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION - ); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); - /** - * Asserting following dim / metrics [ dim1, dim2 / Count [ metric] ] - [0, 0] | [0] - [1, 1] | [1] - [2, 2] | [2] - [3, 3] | [3] - [4, 4] | [4] - [5, 5] | [5] - [6, 6] | [6] - [8, 8] | [8] - [null, 5] | [5] - [null, 7] | [7] - */ - builder.appendDocumentsToStarTree(starTreeDocumentIterator); - assertEquals(10, builder.getStarTreeDocuments().size()); - builder.build(starTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); - int count = 0; - for (StarTreeDocument starTreeDocument : builder.getStarTreeDocuments()) { - count++; - if (count <= 10) { - if (starTreeDocument.dimensions[0] == null) { - assertTrue(List.of(5L, 7L).contains(starTreeDocument.dimensions[1])); - } - assertEquals(starTreeDocument.dimensions[1], starTreeDocument.metrics[0]); - } - } - - validateStarTree(builder.getRootNode(), 2, 1000, builder.getStarTreeDocuments()); - - metaOut.close(); - dataOut.close(); - docValuesConsumer.close(); - - StarTreeMetadata starTreeMetadata = new StarTreeMetadata( - "sf", - STAR_TREE, - mock(IndexInput.class), - VERSION_CURRENT, - builder.numStarTreeNodes, - List.of("field1", "field3"), - compositeField.getMetrics(), - 10, - builder.numStarTreeDocs, - 1000, - Set.of(), - getBuildMode(), - 0, - 363 - ); - - validateStarTreeFileFormats( - builder.getRootNode(), - builder.getStarTreeDocuments().size(), - starTreeMetadata, - builder.getStarTreeDocuments() - ); - } - - public void testMergeFlowWithMissingDocsWithZero() throws IOException { - List dimList = List.of(0L, 0L, 0L, 0L); - List docsWithField = List.of(0, 1, 2, 6); - List dimList2 = List.of(0L, 0L, 0L, 0L); - List docsWithField2 = List.of(0, 1, 2, 6); - - List metricsList = List.of(0L, 1L, 2L, 3L, 4L, 5L, 6L); - List metricsWithField = List.of(0, 1, 2, 3, 4, 5, 6); - - List dimList3 = List.of(5L, 6L, 8L, -1L); - List docsWithField3 = List.of(0, 1, 3, 4); - List dimList4 = List.of(5L, 6L, 7L, 8L, -1L); - List docsWithField4 = List.of(0, 1, 2, 3, 4); - - List metricsList2 = List.of(5L, 6L, 7L, 8L, 9L); - List metricsWithField2 = List.of(0, 1, 2, 3, 4); - - compositeField = getStarTreeField(MetricStat.VALUE_COUNT); - StarTreeValues starTreeValues = getStarTreeValues( - getSortedNumericMock(dimList, docsWithField), - getSortedNumericMock(dimList2, docsWithField2), - getSortedNumericMock(metricsList, metricsWithField), - compositeField, - "7" - ); - - StarTreeValues starTreeValues2 = getStarTreeValues( - getSortedNumericMock(dimList3, docsWithField3), - getSortedNumericMock(dimList4, docsWithField4), - getSortedNumericMock(metricsList2, metricsWithField2), - compositeField, - "4" - ); - writeState = getWriteState(4, writeState.segmentInfo.getId()); - SegmentWriteState consumerWriteState = getWriteState(DocIdSetIterator.NO_MORE_DOCS, writeState.segmentInfo.getId()); - this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( - consumerWriteState, - Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, - Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, - Composite99DocValuesFormat.META_DOC_VALUES_CODEC, - Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION - ); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); - /** - * Asserting following dim / metrics [ dim1, dim2 / Count [ metric] ] - [0, 0] | [9] - [5, 5] | [5] - [6, 6] | [6] - [8, 8] | [8] - [null, 7] | [7] - [null, null] | [12] - */ - builder.appendDocumentsToStarTree(starTreeDocumentIterator); - assertEquals(6, builder.getStarTreeDocuments().size()); - builder.build(starTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); - int count = 0; - for (StarTreeDocument starTreeDocument : builder.getStarTreeDocuments()) { - count++; - if (count <= 6) { - if (starTreeDocument.dimensions[0] == null && starTreeDocument.dimensions[1] == null) { - assertEquals(12L, (long) starTreeDocument.metrics[0]); - } else if (starTreeDocument.dimensions[0] == null) { - assertEquals(7L, starTreeDocument.metrics[0]); - } else if (starTreeDocument.dimensions[0] == 0) { - assertEquals(9L, starTreeDocument.metrics[0]); - } else { - assertEquals(starTreeDocument.dimensions[1], starTreeDocument.metrics[0]); - } - } - } - - validateStarTree(builder.getRootNode(), 2, 1000, builder.getStarTreeDocuments()); - - metaOut.close(); - dataOut.close(); - docValuesConsumer.close(); - - StarTreeMetadata starTreeMetadata = new StarTreeMetadata( - "sf", - STAR_TREE, - mock(IndexInput.class), - VERSION_CURRENT, - builder.numStarTreeNodes, - List.of("field1", "field3"), - compositeField.getMetrics(), - 6, - builder.numStarTreeDocs, - 1000, - Set.of(), - getBuildMode(), - 0, - 231 - ); - - validateStarTreeFileFormats( - builder.getRootNode(), - builder.getStarTreeDocuments().size(), - starTreeMetadata, - builder.getStarTreeDocuments() - ); - } - - public void testMergeFlowWithMissingDocsWithZeroComplexCase() throws IOException { - List dimList = List.of(0L, 0L, 0L, 0L, 0L); - List docsWithField = List.of(0, 1, 2, 6, 8); - List dimList2 = List.of(0L, 0L, 0L, 0L); - List docsWithField2 = List.of(0, 1, 2, 6); - - List metricsList = List.of(0L, 1L, 2L, 3L, 4L, 5L, 6L, 7L, 8L, 9L); - List metricsWithField = List.of(0, 1, 2, 3, 4, 5, 6, 7, 8, 9); - - List dimList3 = List.of(5L, 6L, 8L, -1L); - List docsWithField3 = List.of(0, 1, 3, 4); - List dimList4 = List.of(5L, 6L, 7L, 8L, -1L); - List docsWithField4 = List.of(0, 1, 2, 3, 4); - - List metricsList2 = List.of(5L, 6L, 7L, 8L, 9L); - List metricsWithField2 = List.of(0, 1, 2, 3, 4); - - compositeField = getStarTreeField(MetricStat.VALUE_COUNT); - StarTreeValues starTreeValues = getStarTreeValues( - getSortedNumericMock(dimList, docsWithField), - getSortedNumericMock(dimList2, docsWithField2), - getSortedNumericMock(metricsList, metricsWithField), - compositeField, - "9" - ); - - StarTreeValues starTreeValues2 = getStarTreeValues( - getSortedNumericMock(dimList3, docsWithField3), - getSortedNumericMock(dimList4, docsWithField4), - getSortedNumericMock(metricsList2, metricsWithField2), - compositeField, - "4" - ); - writeState = getWriteState(4, writeState.segmentInfo.getId()); - this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( - writeState, - Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, - Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, - Composite99DocValuesFormat.META_DOC_VALUES_CODEC, - Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION - ); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); - /** - * Asserting following dim / metrics [ dim1, dim2 / Count [ metric] ] - [0, 0] | [9] - [0, null] | [8] - [5, 5] | [5] - [6, 6] | [6] - [8, 8] | [8] - [null, 7] | [7] - [null, null] | [19] - */ - builder.appendDocumentsToStarTree(starTreeDocumentIterator); - assertEquals(7, builder.getStarTreeDocuments().size()); - builder.build(starTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); - int count = 0; - for (StarTreeDocument starTreeDocument : builder.getStarTreeDocuments()) { - count++; - if (count <= 7) { - if (starTreeDocument.dimensions[0] == null && starTreeDocument.dimensions[1] == null) { - assertEquals(19L, (long) starTreeDocument.metrics[0]); - assertEquals(7, count); - } else if (starTreeDocument.dimensions[0] == null) { - assertEquals(7L, starTreeDocument.metrics[0]); - } else if (starTreeDocument.dimensions[1] == null) { - assertEquals(8L, starTreeDocument.metrics[0]); - } else if (starTreeDocument.dimensions[0] == 0) { - assertEquals(9L, starTreeDocument.metrics[0]); - } else { - assertEquals(starTreeDocument.dimensions[1], starTreeDocument.metrics[0]); - } - } - } - - validateStarTree(builder.getRootNode(), 2, 1000, builder.getStarTreeDocuments()); - - metaOut.close(); - dataOut.close(); - docValuesConsumer.close(); - - StarTreeMetadata starTreeMetadata = new StarTreeMetadata( - "sf", - STAR_TREE, - mock(IndexInput.class), - VERSION_CURRENT, - builder.numStarTreeNodes, - List.of("field1", "field3"), - compositeField.getMetrics(), - 7, - builder.numStarTreeDocs, - 1000, - Set.of(), - getBuildMode(), - 0, - 231 - ); - - validateStarTreeFileFormats( - builder.getRootNode(), - builder.getStarTreeDocuments().size(), - starTreeMetadata, - builder.getStarTreeDocuments() - ); - } - - public void testMergeFlowWithMissingDocsInSecondDim() throws IOException { - List dimList2 = List.of(0L, 1L, 2L, 3L, 4L, 6L); - List docsWithField2 = List.of(0, 1, 2, 3, 4, 6); - List dimList = List.of(0L, 1L, 2L, 3L, 4L, 5L, -1L); - List docsWithField = List.of(0, 1, 2, 3, 4, 5, 6); - - List metricsList = List.of(0L, 1L, 2L, 3L, 4L, 5L, 6L); - List metricsWithField = List.of(0, 1, 2, 3, 4, 5, 6); - - List dimList3 = List.of(5L, 6L, 8L, -1L); - List docsWithField3 = List.of(0, 1, 3, 4); - List dimList4 = List.of(5L, 6L, 7L, 8L, -1L); - List docsWithField4 = List.of(0, 1, 2, 3, 4); - - List metricsList2 = List.of(5L, 6L, 7L, 8L, 9L); - List metricsWithField2 = List.of(0, 1, 2, 3, 4); - - compositeField = getStarTreeField(MetricStat.VALUE_COUNT); - StarTreeValues starTreeValues = getStarTreeValues( - getSortedNumericMock(dimList, docsWithField), - getSortedNumericMock(dimList2, docsWithField2), - getSortedNumericMock(metricsList, metricsWithField), - compositeField, - "6" - ); - - StarTreeValues starTreeValues2 = getStarTreeValues( - getSortedNumericMock(dimList3, docsWithField3), - getSortedNumericMock(dimList4, docsWithField4), - getSortedNumericMock(metricsList2, metricsWithField2), - compositeField, - "4" - ); - writeState = getWriteState(4, writeState.segmentInfo.getId()); - this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( - writeState, - Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, - Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, - Composite99DocValuesFormat.META_DOC_VALUES_CODEC, - Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION - ); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); - /** - * Asserting following dim / metrics [ dim1, dim2 / Count [ metric] ] - [0, 0] | [0] - [1, 1] | [1] - [2, 2] | [2] - [3, 3] | [3] - [4, 4] | [4] - [5, 5] | [5] - [5, null] | [5] - [6, 6] | [6] - [8, 8] | [8] - [null, 7] | [7] - */ - builder.appendDocumentsToStarTree(starTreeDocumentIterator); - assertEquals(10, builder.getStarTreeDocuments().size()); - builder.build(starTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); - int count = 0; - for (StarTreeDocument starTreeDocument : builder.getStarTreeDocuments()) { - count++; - if (count <= 10) { - if (starTreeDocument.dimensions[0] != null && starTreeDocument.dimensions[0] == 5) { - assertEquals(starTreeDocument.dimensions[0], starTreeDocument.metrics[0]); - } else { - assertEquals(starTreeDocument.dimensions[1], starTreeDocument.metrics[0]); - } - } - } - - validateStarTree(builder.getRootNode(), 2, 1000, builder.getStarTreeDocuments()); - - metaOut.close(); - dataOut.close(); - docValuesConsumer.close(); - - StarTreeMetadata starTreeMetadata = new StarTreeMetadata( - "sf", - STAR_TREE, - mock(IndexInput.class), - VERSION_CURRENT, - builder.numStarTreeNodes, - List.of("field1", "field3"), - compositeField.getMetrics(), - 10, - builder.numStarTreeDocs, - 1000, - Set.of(), - getBuildMode(), - 0, - 363 - ); - - validateStarTreeFileFormats( - builder.getRootNode(), - builder.getStarTreeDocuments().size(), - starTreeMetadata, - builder.getStarTreeDocuments() - ); - } - - public void testMergeFlowWithDocsMissingAtTheEnd() throws IOException { - List dimList = List.of(0L, 1L, 2L, 3L, 4L); - List docsWithField = List.of(0, 1, 2, 3, 4); - List dimList2 = List.of(0L, 1L, 2L, 3L, 4L, 5L, -1L); - List docsWithField2 = List.of(0, 1, 2, 3, 4, 5, 6); - - List metricsList = List.of(0L, 1L, 2L, 3L, 4L, 5L, 6L); - List metricsWithField = List.of(0, 1, 2, 3, 4, 5, 6); - - List dimList3 = List.of(5L, 6L, 8L, -1L); - List docsWithField3 = List.of(0, 1, 3, 4); - List dimList4 = List.of(5L, 6L, 7L, 8L, -1L); - List docsWithField4 = List.of(0, 1, 2, 3, 4); - - List metricsList2 = List.of(5L, 6L, 7L, 8L, 9L); - List metricsWithField2 = List.of(0, 1, 2, 3, 4); - - compositeField = getStarTreeField(MetricStat.VALUE_COUNT); - StarTreeValues starTreeValues = getStarTreeValues( - getSortedNumericMock(dimList, docsWithField), - getSortedNumericMock(dimList2, docsWithField2), - getSortedNumericMock(metricsList, metricsWithField), - compositeField, - "6" - ); - - StarTreeValues starTreeValues2 = getStarTreeValues( - getSortedNumericMock(dimList3, docsWithField3), - getSortedNumericMock(dimList4, docsWithField4), - getSortedNumericMock(metricsList2, metricsWithField2), - compositeField, - "4" - ); - this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( - writeState, - Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, - Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, - Composite99DocValuesFormat.META_DOC_VALUES_CODEC, - Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION - ); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); - /** - * Asserting following dim / metrics [ dim1, dim2 / Count [ metric] ] - [0, 0] | [0] - [1, 1] | [1] - [2, 2] | [2] - [3, 3] | [3] - [4, 4] | [4] - [5, 5] | [5] - [6, 6] | [6] - [8, 8] | [8] - [null, 5] | [5] - [null, 7] | [7] - */ - builder.appendDocumentsToStarTree(starTreeDocumentIterator); - assertEquals(10, builder.getStarTreeDocuments().size()); - builder.build(starTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); - int count = 0; - for (StarTreeDocument starTreeDocument : builder.getStarTreeDocuments()) { - count++; - if (count <= 10) { - if (starTreeDocument.dimensions[0] == null) { - assertTrue(List.of(5L, 7L).contains(starTreeDocument.dimensions[1])); - } - assertEquals(starTreeDocument.dimensions[1], starTreeDocument.metrics[0]); - } - } - - validateStarTree(builder.getRootNode(), 2, 1000, builder.getStarTreeDocuments()); - - metaOut.close(); - dataOut.close(); - docValuesConsumer.close(); - - StarTreeMetadata starTreeMetadata = new StarTreeMetadata( - "sf", - STAR_TREE, - mock(IndexInput.class), - VERSION_CURRENT, - builder.numStarTreeNodes, - List.of("field1", "field3"), - compositeField.getMetrics(), - 10, - builder.numStarTreeDocs, - 1000, - Set.of(), - getBuildMode(), - 0, - 363 - ); - - validateStarTreeFileFormats( - builder.getRootNode(), - builder.getStarTreeDocuments().size(), - starTreeMetadata, - builder.getStarTreeDocuments() - ); - } - - public void testMergeFlowWithEmptyFieldsInOneSegment() throws IOException { - List dimList = List.of(0L, 1L, 2L, 3L, 4L); - List docsWithField = List.of(0, 1, 2, 3, 4); - List dimList2 = List.of(0L, 1L, 2L, 3L, 4L, 5L, -1L); - List docsWithField2 = List.of(0, 1, 2, 3, 4, 5, 6); - - List metricsList = List.of(0L, 1L, 2L, 3L, 4L, 5L, 6L); - List metricsWithField = List.of(0, 1, 2, 3, 4, 5, 6); - - compositeField = getStarTreeField(MetricStat.VALUE_COUNT); - StarTreeValues starTreeValues = getStarTreeValues( - getSortedNumericMock(dimList, docsWithField), - getSortedNumericMock(dimList2, docsWithField2), - getSortedNumericMock(metricsList, metricsWithField), - compositeField, - "6" - ); - - StarTreeValues starTreeValues2 = getStarTreeValues( - DocValues.emptySortedNumeric(), - DocValues.emptySortedNumeric(), - DocValues.emptySortedNumeric(), - compositeField, - "0" - ); - writeState = getWriteState(0, writeState.segmentInfo.getId()); - this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( - writeState, - Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, - Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, - Composite99DocValuesFormat.META_DOC_VALUES_CODEC, - Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION - ); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); - /** - * Asserting following dim / metrics [ dim1, dim2 / Count [ metric] ] - [0, 0] | [0] - [1, 1] | [1] - [2, 2] | [2] - [3, 3] | [3] - [4, 4] | [4] - [null, 5] | [5] - */ - builder.appendDocumentsToStarTree(starTreeDocumentIterator); - assertEquals(6, builder.getStarTreeDocuments().size()); - builder.build(starTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); - int count = 0; - for (StarTreeDocument starTreeDocument : builder.getStarTreeDocuments()) { - count++; - if (count <= 6) { - if (starTreeDocument.dimensions[0] == null) { - assertEquals(5L, (long) starTreeDocument.dimensions[1]); - } - assertEquals(starTreeDocument.dimensions[1], starTreeDocument.metrics[0]); - } - } - validateStarTree(builder.getRootNode(), 2, 1000, builder.getStarTreeDocuments()); - - metaOut.close(); - dataOut.close(); - docValuesConsumer.close(); - - StarTreeMetadata starTreeMetadata = new StarTreeMetadata( - "sf", - STAR_TREE, - mock(IndexInput.class), - VERSION_CURRENT, - builder.numStarTreeNodes, - List.of("field1", "field3"), - compositeField.getMetrics(), - 6, - builder.numStarTreeDocs, - 1000, - Set.of(), - getBuildMode(), - 0, - 264 - ); - - validateStarTreeFileFormats( - builder.getRootNode(), - builder.getStarTreeDocuments().size(), - starTreeMetadata, - builder.getStarTreeDocuments() - ); - } - - public void testMergeFlowWithDuplicateDimensionValues() throws IOException { - List dimList1 = new ArrayList<>(500); - List docsWithField1 = new ArrayList<>(500); - for (int i = 0; i < 100; i++) { - for (int j = 0; j < 5; j++) { - dimList1.add((long) i); - docsWithField1.add(i * 5 + j); - } - } - - List dimList2 = new ArrayList<>(500); - List docsWithField2 = new ArrayList<>(500); - for (int i = 0; i < 100; i++) { - for (int j = 0; j < 5; j++) { - dimList2.add((long) i); - docsWithField2.add(i * 5 + j); - } - } - - List dimList3 = new ArrayList<>(500); - List docsWithField3 = new ArrayList<>(500); - for (int i = 0; i < 100; i++) { - for (int j = 0; j < 5; j++) { - dimList3.add((long) i); - docsWithField3.add(i * 5 + j); - } - } - - List dimList4 = new ArrayList<>(500); - List docsWithField4 = new ArrayList<>(500); - for (int i = 0; i < 100; i++) { - for (int j = 0; j < 5; j++) { - dimList4.add((long) i); - docsWithField4.add(i * 5 + j); - } - } - - List metricsList = new ArrayList<>(100); - List metricsWithField = new ArrayList<>(100); - for (int i = 0; i < 500; i++) { - metricsList.add(getLongFromDouble(i * 10.0)); - metricsWithField.add(i); - } - List docCountMetricsList = new ArrayList<>(100); - List docCountMetricsWithField = new ArrayList<>(100); - for (int i = 0; i < 500; i++) { - docCountMetricsList.add(i * 10L); - docCountMetricsWithField.add(i); - } - - compositeField = getStarTreeFieldWithDocCount(1, true); - StarTreeValues starTreeValues = getStarTreeValues( - dimList1, - docsWithField1, - dimList2, - docsWithField2, - dimList3, - docsWithField3, - dimList4, - docsWithField4, - metricsList, - metricsWithField, - docCountMetricsList, - docCountMetricsWithField, - compositeField - ); - - StarTreeValues starTreeValues2 = getStarTreeValues( - dimList1, - docsWithField1, - dimList2, - docsWithField2, - dimList3, - docsWithField3, - dimList4, - docsWithField4, - metricsList, - metricsWithField, - docCountMetricsList, - docCountMetricsWithField, - compositeField - ); - this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( - writeState, - Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, - Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, - Composite99DocValuesFormat.META_DOC_VALUES_CODEC, - Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION - ); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - builder.build(builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)), new AtomicInteger(), docValuesConsumer); - List starTreeDocuments = builder.getStarTreeDocuments(); - assertEquals(401, starTreeDocuments.size()); - int count = 0; - double sum = 0; - /** - 401 docs get generated - [0, 0, 0, 0] | [200.0, 10] - [1, 1, 1, 1] | [700.0, 10] - [2, 2, 2, 2] | [1200.0, 10] - [3, 3, 3, 3] | [1700.0, 10] - [4, 4, 4, 4] | [2200.0, 10] - ..... - [null, null, null, 99] | [49700.0, 10] - [null, null, null, null] | [2495000.0, 1000] - */ - for (StarTreeDocument starTreeDocument : starTreeDocuments) { - if (starTreeDocument.dimensions[3] == null) { - assertEquals(sum, starTreeDocument.metrics[0]); - assertEquals(2495000L, (long) starTreeDocument.metrics[1]); - } else { - if (starTreeDocument.dimensions[0] != null) { - sum += (double) starTreeDocument.metrics[0]; - } - assertEquals(starTreeDocument.dimensions[3] * 500 + 200.0, starTreeDocument.metrics[0]); - assertEquals(starTreeDocument.dimensions[3] * 500 + 200L, (long) starTreeDocument.metrics[1]); - - } - count++; - } - assertEquals(401, count); - validateStarTree(builder.getRootNode(), 4, compositeField.getStarTreeConfig().maxLeafDocs(), builder.getStarTreeDocuments()); - metaOut.close(); - dataOut.close(); - docValuesConsumer.close(); - - StarTreeMetadata starTreeMetadata = new StarTreeMetadata( - "sf", - STAR_TREE, - mock(IndexInput.class), - VERSION_CURRENT, - builder.numStarTreeNodes, - getStarTreeDimensionNames(compositeField.getDimensionsOrder()), - compositeField.getMetrics(), - 100, - builder.numStarTreeDocs, - 1, - Set.of(), - getBuildMode(), - 0, - 13365 - ); - - validateStarTreeFileFormats( - builder.getRootNode(), - builder.getStarTreeDocuments().size(), - starTreeMetadata, - builder.getStarTreeDocuments() - ); - } - - public void testMergeFlowWithMaxLeafDocs() throws IOException { - List dimList1 = new ArrayList<>(500); - List docsWithField1 = new ArrayList<>(500); - - for (int i = 0; i < 20; i++) { - for (int j = 0; j < 20; j++) { - dimList1.add((long) i); - docsWithField1.add(i * 20 + j); - } - } - for (int i = 80; i < 100; i++) { - for (int j = 0; j < 5; j++) { - dimList1.add((long) i); - docsWithField1.add(i * 5 + j); - } - } - List dimList3 = new ArrayList<>(500); - List docsWithField3 = new ArrayList<>(500); - for (int i = 0; i < 100; i++) { - for (int j = 0; j < 5; j++) { - dimList3.add((long) i); - docsWithField3.add(i * 5 + j); - } - } - List dimList2 = new ArrayList<>(500); - List docsWithField2 = new ArrayList<>(500); - for (int i = 0; i < 10; i++) { - for (int j = 0; j < 50; j++) { - dimList2.add((long) i); - docsWithField2.add(i * 50 + j); - } - } - - List dimList4 = new ArrayList<>(500); - List docsWithField4 = new ArrayList<>(500); - for (int i = 0; i < 100; i++) { - for (int j = 0; j < 5; j++) { - dimList4.add((long) i); - docsWithField4.add(i * 5 + j); - } - } - - List metricsList = new ArrayList<>(100); - List metricsWithField = new ArrayList<>(100); - for (int i = 0; i < 500; i++) { - metricsList.add(getLongFromDouble(i * 10.0)); - metricsWithField.add(i); - } - - List metricsList1 = new ArrayList<>(100); - List metricsWithField1 = new ArrayList<>(100); - for (int i = 0; i < 500; i++) { - metricsList1.add(1L); - metricsWithField1.add(i); - } - - compositeField = getStarTreeFieldWithDocCount(3, true); - StarTreeValues starTreeValues = getStarTreeValues( - dimList1, - docsWithField1, - dimList2, - docsWithField2, - dimList3, - docsWithField3, - dimList4, - docsWithField4, - metricsList, - metricsWithField, - metricsList1, - metricsWithField1, - compositeField - ); - - StarTreeValues starTreeValues2 = getStarTreeValues( - dimList1, - docsWithField1, - dimList2, - docsWithField2, - dimList3, - docsWithField3, - dimList4, - docsWithField4, - metricsList, - metricsWithField, - metricsList1, - metricsWithField1, - compositeField - ); - - this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( - writeState, - Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, - Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, - Composite99DocValuesFormat.META_DOC_VALUES_CODEC, - Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION - ); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - builder.build(builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)), new AtomicInteger(), docValuesConsumer); - List starTreeDocuments = builder.getStarTreeDocuments(); - /** - 635 docs get generated - [0, 0, 0, 0] | [200.0, 10] - [0, 0, 1, 1] | [700.0, 10] - [0, 0, 2, 2] | [1200.0, 10] - [0, 0, 3, 3] | [1700.0, 10] - [1, 0, 4, 4] | [2200.0, 10] - [1, 0, 5, 5] | [2700.0, 10] - [1, 0, 6, 6] | [3200.0, 10] - [1, 0, 7, 7] | [3700.0, 10] - [2, 0, 8, 8] | [4200.0, 10] - [2, 0, 9, 9] | [4700.0, 10] - [2, 1, 10, 10] | [5200.0, 10] - [2, 1, 11, 11] | [5700.0, 10] - ..... - [18, 7, null, null] | [147800.0, 40] - ... - [7, 2, null, null] | [28900.0, 20] - ... - [null, null, null, 99] | [49700.0, 10] - ..... - [null, null, null, null] | [2495000.0, 1000] - */ - assertEquals(635, starTreeDocuments.size()); - for (StarTreeDocument starTreeDocument : starTreeDocuments) { - if (starTreeDocument.dimensions[0] != null - && starTreeDocument.dimensions[1] != null - && starTreeDocument.dimensions[2] != null - && starTreeDocument.dimensions[3] != null) { - assertEquals(10L, starTreeDocument.metrics[1]); - } else if (starTreeDocument.dimensions[1] != null - && starTreeDocument.dimensions[2] != null - && starTreeDocument.dimensions[3] != null) { - assertEquals(10L, starTreeDocument.metrics[1]); - } else if (starTreeDocument.dimensions[0] != null - && starTreeDocument.dimensions[2] != null - && starTreeDocument.dimensions[3] != null) { - assertEquals(10L, starTreeDocument.metrics[1]); - } else if (starTreeDocument.dimensions[0] != null - && starTreeDocument.dimensions[1] != null - && starTreeDocument.dimensions[3] != null) { - assertEquals(10L, starTreeDocument.metrics[1]); - } else if (starTreeDocument.dimensions[0] != null && starTreeDocument.dimensions[3] != null) { - assertEquals(10L, starTreeDocument.metrics[1]); - } else if (starTreeDocument.dimensions[0] != null && starTreeDocument.dimensions[1] != null) { - assertTrue((long) starTreeDocument.metrics[1] == 20L || (long) starTreeDocument.metrics[1] == 40L); - } else if (starTreeDocument.dimensions[1] != null && starTreeDocument.dimensions[3] != null) { - assertEquals(10L, starTreeDocument.metrics[1]); - } else if (starTreeDocument.dimensions[1] != null) { - assertEquals(100L, starTreeDocument.metrics[1]); - } else if (starTreeDocument.dimensions[0] != null) { - assertEquals(40L, starTreeDocument.metrics[1]); - } - } - validateStarTree(builder.getRootNode(), 4, compositeField.getStarTreeConfig().maxLeafDocs(), builder.getStarTreeDocuments()); - metaOut.close(); - dataOut.close(); - docValuesConsumer.close(); - - StarTreeMetadata starTreeMetadata = new StarTreeMetadata( - "sf", - STAR_TREE, - mock(IndexInput.class), - VERSION_CURRENT, - builder.numStarTreeNodes, - getStarTreeDimensionNames(compositeField.getDimensionsOrder()), - compositeField.getMetrics(), - 100, - builder.numStarTreeDocs, - 3, - Set.of(), - getBuildMode(), - 0, - 23199 - ); - - validateStarTreeFileFormats( - builder.getRootNode(), - builder.getStarTreeDocuments().size(), - starTreeMetadata, - builder.getStarTreeDocuments() - ); - } - - private StarTreeValues getStarTreeValues( - List dimList1, - List docsWithField1, - List dimList2, - List docsWithField2, - List dimList3, - List docsWithField3, - List dimList4, - List docsWithField4, - List metricsList, - List metricsWithField, - List metricsList1, - List metricsWithField1, - StarTreeField sf - ) { - SortedNumericDocValues d1sndv = getSortedNumericMock(dimList1, docsWithField1); - SortedNumericDocValues d2sndv = getSortedNumericMock(dimList2, docsWithField2); - SortedNumericDocValues d3sndv = getSortedNumericMock(dimList3, docsWithField3); - SortedNumericDocValues d4sndv = getSortedNumericMock(dimList4, docsWithField4); - SortedNumericDocValues m1sndv = getSortedNumericMock(metricsList, metricsWithField); - SortedNumericDocValues m2sndv = getSortedNumericMock(metricsList1, metricsWithField1); - Map> dimDocIdSetIterators = Map.of( - "field1", - () -> new SortedNumericStarTreeValuesIterator(d1sndv), - "field3", - () -> new SortedNumericStarTreeValuesIterator(d2sndv), - "field5", - () -> new SortedNumericStarTreeValuesIterator(d3sndv), - "field8", - () -> new SortedNumericStarTreeValuesIterator(d4sndv) - ); - - Map> metricDocIdSetIterators = new LinkedHashMap<>(); - - metricDocIdSetIterators.put( - fullyQualifiedFieldNameForStarTreeMetricsDocValues( - sf.getName(), - "field2", - sf.getMetrics().get(0).getMetrics().get(0).getTypeName() - ), - () -> new SortedNumericStarTreeValuesIterator(m1sndv) - ); - metricDocIdSetIterators.put( - fullyQualifiedFieldNameForStarTreeMetricsDocValues( - sf.getName(), - "_doc_count", - sf.getMetrics().get(1).getMetrics().get(0).getTypeName() - ), - () -> new SortedNumericStarTreeValuesIterator(m2sndv) - ); - // metricDocIdSetIterators.put("field2", () -> m1sndv); - // metricDocIdSetIterators.put("_doc_count", () -> m2sndv); - StarTreeValues starTreeValues = new StarTreeValues( - sf, - null, - dimDocIdSetIterators, - metricDocIdSetIterators, - getAttributes(500), - null - ); - return starTreeValues; - } - - public void testMergeFlowWithDuplicateDimensionValueWithMaxLeafDocs() throws IOException { - List dimList1 = new ArrayList<>(500); - List docsWithField1 = new ArrayList<>(500); - - for (int i = 0; i < 20; i++) { - for (int j = 0; j < 20; j++) { - dimList1.add((long) i); - docsWithField1.add(i * 20 + j); - } - } - for (int i = 80; i < 100; i++) { - for (int j = 0; j < 5; j++) { - dimList1.add((long) i); - docsWithField1.add(i * 5 + j); - } - } - List dimList3 = new ArrayList<>(500); - List docsWithField3 = new ArrayList<>(500); - for (int i = 0; i < 100; i++) { - for (int j = 0; j < 5; j++) { - dimList3.add((long) i); - docsWithField3.add(i * 5 + j); - } - } - List dimList2 = new ArrayList<>(500); - List docsWithField2 = new ArrayList<>(500); - for (int i = 0; i < 500; i++) { - dimList2.add((long) 1); - docsWithField2.add(i); - } - - List dimList4 = new ArrayList<>(500); - List docsWithField4 = new ArrayList<>(500); - for (int i = 0; i < 100; i++) { - for (int j = 0; j < 5; j++) { - dimList4.add((long) i); - docsWithField4.add(i * 5 + j); - } - } - - List metricsList = new ArrayList<>(100); - List metricsWithField = new ArrayList<>(100); - for (int i = 0; i < 500; i++) { - metricsList.add(getLongFromDouble(i * 10.0)); - metricsWithField.add(i); - } - - List docCountMetricsList = new ArrayList<>(100); - List docCountMetricsWithField = new ArrayList<>(100); - for (int i = 0; i < 500; i++) { - metricsList.add(getLongFromDouble(i * 2)); - metricsWithField.add(i); - } - - compositeField = getStarTreeFieldWithDocCount(3, true); - StarTreeValues starTreeValues = getStarTreeValues( - dimList1, - docsWithField1, - dimList2, - docsWithField2, - dimList3, - docsWithField3, - dimList4, - docsWithField4, - metricsList, - metricsWithField, - docCountMetricsList, - docCountMetricsWithField, - compositeField - ); - - StarTreeValues starTreeValues2 = getStarTreeValues( - dimList1, - docsWithField1, - dimList2, - docsWithField2, - dimList3, - docsWithField3, - dimList4, - docsWithField4, - metricsList, - metricsWithField, - docCountMetricsList, - docCountMetricsWithField, - compositeField - ); - this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( - writeState, - Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, - Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, - Composite99DocValuesFormat.META_DOC_VALUES_CODEC, - Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION - ); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - builder.build(builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)), new AtomicInteger(), docValuesConsumer); - List starTreeDocuments = builder.getStarTreeDocuments(); - assertEquals(401, starTreeDocuments.size()); - validateStarTree(builder.getRootNode(), 4, compositeField.getStarTreeConfig().maxLeafDocs(), builder.getStarTreeDocuments()); - - metaOut.close(); - dataOut.close(); - docValuesConsumer.close(); - - StarTreeMetadata starTreeMetadata = new StarTreeMetadata( - "sf", - STAR_TREE, - mock(IndexInput.class), - VERSION_CURRENT, - builder.numStarTreeNodes, - getStarTreeDimensionNames(compositeField.getDimensionsOrder()), - compositeField.getMetrics(), - 100, - builder.numStarTreeDocs, - compositeField.getStarTreeConfig().maxLeafDocs(), - Set.of(), - getBuildMode(), - 0, - 15345 - ); - - validateStarTreeFileFormats( - builder.getRootNode(), - builder.getStarTreeDocuments().size(), - starTreeMetadata, - builder.getStarTreeDocuments() - ); - } - - public static long getLongFromDouble(double value) { - return NumericUtils.doubleToSortableLong(value); - } - - public void testMergeFlowWithMaxLeafDocsAndStarTreeNodesAssertion() throws IOException { - List dimList1 = new ArrayList<>(500); - List docsWithField1 = new ArrayList<>(500); - Map> expectedDimToValueMap = new HashMap<>(); - Map dimValueMap = new HashMap<>(); - for (int i = 0; i < 20; i++) { - for (int j = 0; j < 20; j++) { - dimList1.add((long) i); - docsWithField1.add(i * 20 + j); - } - // metric = no of docs * 10.0 - dimValueMap.put((long) i, 200.0); - } - for (int i = 80; i < 100; i++) { - for (int j = 0; j < 5; j++) { - dimList1.add((long) i); - docsWithField1.add(i * 5 + j); - } - // metric = no of docs * 10.0 - dimValueMap.put((long) i, 50.0); - } - dimValueMap.put(Long.MAX_VALUE, 5000.0); - expectedDimToValueMap.put(0, dimValueMap); - dimValueMap = new HashMap<>(); - List dimList3 = new ArrayList<>(500); - List docsWithField3 = new ArrayList<>(500); - for (int i = 0; i < 500; i++) { - dimList3.add((long) 1); - docsWithField3.add(i); - dimValueMap.put((long) i, 10.0); - } - dimValueMap.put(Long.MAX_VALUE, 5000.0); - expectedDimToValueMap.put(2, dimValueMap); - dimValueMap = new HashMap<>(); - List dimList2 = new ArrayList<>(500); - List docsWithField2 = new ArrayList<>(500); - for (int i = 0; i < 500; i++) { - dimList2.add((long) i); - docsWithField2.add(i); - dimValueMap.put((long) i, 10.0); - } - dimValueMap.put(Long.MAX_VALUE, 200.0); - expectedDimToValueMap.put(1, dimValueMap); - dimValueMap = new HashMap<>(); - List dimList4 = new ArrayList<>(500); - List docsWithField4 = new ArrayList<>(500); - for (int i = 0; i < 500; i++) { - dimList4.add((long) 1); - docsWithField4.add(i); - dimValueMap.put((long) i, 10.0); - } - dimValueMap.put(Long.MAX_VALUE, 5000.0); - expectedDimToValueMap.put(3, dimValueMap); - List metricsList = new ArrayList<>(100); - List metricsWithField = new ArrayList<>(100); - for (int i = 0; i < 500; i++) { - metricsList.add(getLongFromDouble(10.0)); - metricsWithField.add(i); - } - List metricsList1 = new ArrayList<>(100); - List metricsWithField1 = new ArrayList<>(100); - for (int i = 0; i < 500; i++) { - metricsList.add(1L); - metricsWithField.add(i); - } - compositeField = getStarTreeFieldWithDocCount(10, true); - StarTreeValues starTreeValues = getStarTreeValues( - dimList1, - docsWithField1, - dimList2, - docsWithField2, - dimList3, - docsWithField3, - dimList4, - docsWithField4, - metricsList, - metricsWithField, - metricsList1, - metricsWithField1, - compositeField - ); - - StarTreeValues starTreeValues2 = getStarTreeValues( - dimList1, - docsWithField1, - dimList2, - docsWithField2, - dimList3, - docsWithField3, - dimList4, - docsWithField4, - metricsList, - metricsWithField, - metricsList1, - metricsWithField1, - compositeField - ); - this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( - writeState, - Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, - Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, - Composite99DocValuesFormat.META_DOC_VALUES_CODEC, - Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION - ); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - builder.build(builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)), new AtomicInteger(), docValuesConsumer); - List starTreeDocuments = builder.getStarTreeDocuments(); - Map> dimValueToDocIdMap = new HashMap<>(); - traverseStarTree(builder.rootNode, dimValueToDocIdMap, true); - for (Map.Entry> entry : dimValueToDocIdMap.entrySet()) { - int dimId = entry.getKey(); - if (dimId == -1) continue; - Map map = expectedDimToValueMap.get(dimId); - for (Map.Entry dimValueToDocIdEntry : entry.getValue().entrySet()) { - long dimValue = dimValueToDocIdEntry.getKey(); - int docId = dimValueToDocIdEntry.getValue(); - assertEquals(map.get(dimValue) * 2, starTreeDocuments.get(docId).metrics[0]); - } - } - assertEquals(1041, starTreeDocuments.size()); - validateStarTree(builder.getRootNode(), 4, compositeField.getStarTreeConfig().maxLeafDocs(), builder.getStarTreeDocuments()); - - metaOut.close(); - dataOut.close(); - docValuesConsumer.close(); - - StarTreeMetadata starTreeMetadata = new StarTreeMetadata( - "sf", - STAR_TREE, - mock(IndexInput.class), - VERSION_CURRENT, - builder.numStarTreeNodes, - getStarTreeDimensionNames(compositeField.getDimensionsOrder()), - compositeField.getMetrics(), - 500, - builder.numStarTreeDocs, - compositeField.getStarTreeConfig().maxLeafDocs(), - Set.of(), - getBuildMode(), - 0, - 31779 - ); - - validateStarTreeFileFormats( - builder.getRootNode(), - builder.getStarTreeDocuments().size(), - starTreeMetadata, - builder.getStarTreeDocuments() - ); - } - - private StarTreeField getStarTreeFieldWithDocCount(int maxLeafDocs, boolean includeDocCountMetric) { - Dimension d1 = new NumericDimension("field1"); - Dimension d2 = new NumericDimension("field3"); - Dimension d3 = new NumericDimension("field5"); - Dimension d4 = new NumericDimension("field8"); - List dims = List.of(d1, d2, d3, d4); - Metric m1 = new Metric("field2", List.of(MetricStat.SUM)); - Metric m2 = null; - if (includeDocCountMetric) { - m2 = new Metric("_doc_count", List.of(MetricStat.DOC_COUNT)); - } - List metrics = m2 == null ? List.of(m1) : List.of(m1, m2); - StarTreeFieldConfiguration c = new StarTreeFieldConfiguration(maxLeafDocs, new HashSet<>(), getBuildMode()); - StarTreeField sf = new StarTreeField("sf", dims, metrics, c); - return sf; - } - - private void traverseStarTree(InMemoryTreeNode root, Map> dimValueToDocIdMap, boolean traverStarNodes) { - InMemoryTreeNode starTree = root; - // Use BFS to traverse the star tree - Queue queue = new ArrayDeque<>(); - queue.add(starTree); - int currentDimensionId = -1; - InMemoryTreeNode starTreeNode; - List docIds = new ArrayList<>(); - while ((starTreeNode = queue.poll()) != null) { - int dimensionId = starTreeNode.getDimensionId(); - if (dimensionId > currentDimensionId) { - currentDimensionId = dimensionId; - } - - // store aggregated document of the node - int docId = starTreeNode.getAggregatedDocId(); - Map map = dimValueToDocIdMap.getOrDefault(dimensionId, new HashMap<>()); - if (starTreeNode.getNodeType() == StarTreeNodeType.STAR.getValue()) { - map.put(Long.MAX_VALUE, docId); - } else { - map.put(starTreeNode.getDimensionValue(), docId); - } - dimValueToDocIdMap.put(dimensionId, map); - - if (starTreeNode.getChildren() != null - && (!traverStarNodes || starTreeNode.getNodeType() == StarTreeNodeType.STAR.getValue())) { - Iterator childrenIterator = starTreeNode.getChildren().values().iterator(); - while (childrenIterator.hasNext()) { - InMemoryTreeNode childNode = childrenIterator.next(); - queue.add(childNode); - } - } - } - } - - public void testMergeFlow() throws IOException { - List dimList1 = new ArrayList<>(1000); - List docsWithField1 = new ArrayList<>(1000); - for (int i = 0; i < 1000; i++) { - dimList1.add((long) i); - docsWithField1.add(i); - } - - List dimList2 = new ArrayList<>(1000); - List docsWithField2 = new ArrayList<>(1000); - for (int i = 0; i < 1000; i++) { - dimList2.add((long) i); - docsWithField2.add(i); - } - - List dimList3 = new ArrayList<>(1000); - List docsWithField3 = new ArrayList<>(1000); - for (int i = 0; i < 1000; i++) { - dimList3.add((long) i); - docsWithField3.add(i); - } - - List dimList4 = new ArrayList<>(1000); - List docsWithField4 = new ArrayList<>(1000); - for (int i = 0; i < 1000; i++) { - dimList4.add((long) i); - docsWithField4.add(i); - } - - List dimList5 = new ArrayList<>(1000); - List docsWithField5 = new ArrayList<>(1000); - for (int i = 0; i < 1000; i++) { - dimList5.add((long) i); - docsWithField5.add(i); - } - - List metricsList = new ArrayList<>(1000); - List metricsWithField = new ArrayList<>(1000); - for (int i = 0; i < 1000; i++) { - metricsList.add(getLongFromDouble(i * 10.0)); - metricsWithField.add(i); - } - - List metricsListValueCount = new ArrayList<>(1000); - List metricsWithFieldValueCount = new ArrayList<>(1000); - for (int i = 0; i < 1000; i++) { - metricsListValueCount.add((long) i); - metricsWithFieldValueCount.add(i); - } - - Dimension d1 = new NumericDimension("field1"); - Dimension d2 = new NumericDimension("field3"); - Dimension d3 = new NumericDimension("field5"); - Dimension d4 = new NumericDimension("field8"); - // Dimension d5 = new NumericDimension("field5"); - Metric m1 = new Metric("field2", List.of(MetricStat.SUM, MetricStat.AVG, MetricStat.VALUE_COUNT)); - Metric m2 = new Metric("_doc_count", List.of(MetricStat.DOC_COUNT)); - List dims = List.of(d1, d2, d3, d4); - List metrics = List.of(m1, m2); - StarTreeFieldConfiguration c = new StarTreeFieldConfiguration(1, new HashSet<>(), getBuildMode()); - compositeField = new StarTreeField("sf", dims, metrics, c); - SortedNumericDocValues d1sndv = getSortedNumericMock(dimList1, docsWithField1); - SortedNumericDocValues d2sndv = getSortedNumericMock(dimList2, docsWithField2); - SortedNumericDocValues d3sndv = getSortedNumericMock(dimList3, docsWithField3); - SortedNumericDocValues d4sndv = getSortedNumericMock(dimList4, docsWithField4); - SortedNumericDocValues m1sndv = getSortedNumericMock(metricsList, metricsWithField); - SortedNumericDocValues valucountsndv = getSortedNumericMock(metricsListValueCount, metricsWithFieldValueCount); - SortedNumericDocValues m2sndv = DocValues.emptySortedNumeric(); - Map> dimDocIdSetIterators = Map.of( - "field1", - () -> new SortedNumericStarTreeValuesIterator(d1sndv), - "field3", - () -> new SortedNumericStarTreeValuesIterator(d2sndv), - "field5", - () -> new SortedNumericStarTreeValuesIterator(d3sndv), - "field8", - () -> new SortedNumericStarTreeValuesIterator(d4sndv) - ); - - Map> metricDocIdSetIterators = Map.of( - "sf_field2_sum_metric", - () -> new SortedNumericStarTreeValuesIterator(m1sndv), - "sf_field2_value_count_metric", - () -> new SortedNumericStarTreeValuesIterator(valucountsndv), - "sf__doc_count_doc_count_metric", - () -> new SortedNumericStarTreeValuesIterator(m2sndv) - ); - - StarTreeValues starTreeValues = new StarTreeValues( - compositeField, - null, - dimDocIdSetIterators, - metricDocIdSetIterators, - getAttributes(1000), - null - ); - - SortedNumericDocValues f2d1sndv = getSortedNumericMock(dimList1, docsWithField1); - SortedNumericDocValues f2d2sndv = getSortedNumericMock(dimList2, docsWithField2); - SortedNumericDocValues f2d3sndv = getSortedNumericMock(dimList3, docsWithField3); - SortedNumericDocValues f2d4sndv = getSortedNumericMock(dimList4, docsWithField4); - SortedNumericDocValues f2m1sndv = getSortedNumericMock(metricsList, metricsWithField); - SortedNumericDocValues f2valucountsndv = getSortedNumericMock(metricsListValueCount, metricsWithFieldValueCount); - SortedNumericDocValues f2m2sndv = DocValues.emptySortedNumeric(); - Map> f2dimDocIdSetIterators = Map.of( - "field1", - () -> new SortedNumericStarTreeValuesIterator(f2d1sndv), - "field3", - () -> new SortedNumericStarTreeValuesIterator(f2d2sndv), - "field5", - () -> new SortedNumericStarTreeValuesIterator(f2d3sndv), - "field8", - () -> new SortedNumericStarTreeValuesIterator(f2d4sndv) - ); - - Map> f2metricDocIdSetIterators = Map.of( - "sf_field2_sum_metric", - () -> new SortedNumericStarTreeValuesIterator(f2m1sndv), - "sf_field2_value_count_metric", - () -> new SortedNumericStarTreeValuesIterator(f2valucountsndv), - "sf__doc_count_doc_count_metric", - () -> new SortedNumericStarTreeValuesIterator(f2m2sndv) - ); - StarTreeValues starTreeValues2 = new StarTreeValues( - compositeField, - null, - f2dimDocIdSetIterators, - f2metricDocIdSetIterators, - getAttributes(1000), - null - ); - - this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( - writeState, - Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, - Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, - Composite99DocValuesFormat.META_DOC_VALUES_CODEC, - Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION - ); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); - Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); - /** - [0, 0, 0, 0] | [0.0, 2] - [1, 1, 1, 1] | [20.0, 2] - [2, 2, 2, 2] | [40.0, 2] - [3, 3, 3, 3] | [60.0, 2] - [4, 4, 4, 4] | [80.0, 2] - [5, 5, 5, 5] | [100.0, 2] - ... - [999, 999, 999, 999] | [19980.0] - */ - builder.appendDocumentsToStarTree(starTreeDocumentIterator); - for (StarTreeDocument starTreeDocument : builder.getStarTreeDocuments()) { - assertEquals(starTreeDocument.dimensions[0] * 20.0, starTreeDocument.metrics[0]); - assertEquals(starTreeDocument.dimensions[0] * 2, starTreeDocument.metrics[1]); - assertEquals(2L, starTreeDocument.metrics[2]); - } - builder.build(starTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); - - // Validate the star tree structure - validateStarTree(builder.getRootNode(), 4, 1, builder.getStarTreeDocuments()); - - metaOut.close(); - dataOut.close(); - docValuesConsumer.close(); - - StarTreeMetadata starTreeMetadata = new StarTreeMetadata( - compositeField.getName(), - STAR_TREE, - mock(IndexInput.class), - VERSION_CURRENT, - builder.numStarTreeNodes, - getStarTreeDimensionNames(compositeField.getDimensionsOrder()), - compositeField.getMetrics(), - 1000, - builder.numStarTreeDocs, - compositeField.getStarTreeConfig().maxLeafDocs(), - Set.of(), - getBuildMode(), - 0, - 132165 - ); - - validateStarTreeFileFormats( - builder.getRootNode(), - builder.getStarTreeDocuments().size(), - starTreeMetadata, - builder.getStarTreeDocuments() - ); - } - - public void testFlushFlowWithTimestamps() throws IOException { - List dimList = List.of(1655288152000L, 1655288092000L, 1655288032000L, 1655287972000L, 1655288092000L); - List docsWithField = List.of(0, 1, 3, 4, 5); - List dimList2 = List.of(0L, 1L, 2L, 3L, 4L, 5L); - List docsWithField2 = List.of(0, 1, 2, 3, 4, 5); - - List metricsList = List.of( - getLongFromDouble(0.0), - getLongFromDouble(10.0), - getLongFromDouble(20.0), - getLongFromDouble(30.0), - getLongFromDouble(40.0), - getLongFromDouble(50.0) - ); - List metricsWithField = List.of(0, 1, 2, 3, 4, 5); - - 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) - ); - this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( - writeState, - Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, - Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, - Composite99DocValuesFormat.META_DOC_VALUES_CODEC, - Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION - ); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, getWriteState(6, writeState.segmentInfo.getId()), mapperService); - SequentialDocValuesIterator[] dimDvs = { new SequentialDocValuesIterator(d1sndv), new SequentialDocValuesIterator(d2sndv) }; - Iterator starTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( - dimDvs, - List.of(new SequentialDocValuesIterator(m1sndv), new SequentialDocValuesIterator(m2sndv)) - ); - /** - * Asserting following dim / metrics [ dim1, dim2 / Sum [metric], count [metric] ] - [1655287920000, 1655287200000, 1655287200000, 4] | [40.0, 1] - [1655287980000, 1655287200000, 1655287200000, 3] | [30.0, 1] - [1655288040000, 1655287200000, 1655287200000, 1] | [10.0, 1] - [1655288040000, 1655287200000, 1655287200000, 5] | [50.0, 1] - [1655288100000, 1655287200000, 1655287200000, 0] | [0.0, 1] - [null, null, null, 2] | [20.0, 1] - */ - int count = 0; - List starTreeDocumentsList = new ArrayList<>(); - starTreeDocumentIterator.forEachRemaining(starTreeDocumentsList::add); - starTreeDocumentIterator = starTreeDocumentsList.iterator(); - while (starTreeDocumentIterator.hasNext()) { - count++; - StarTreeDocument starTreeDocument = starTreeDocumentIterator.next(); - assertEquals(starTreeDocument.dimensions[3] * 1 * 10.0, starTreeDocument.metrics[1]); - assertEquals(1L, starTreeDocument.metrics[0]); - } - assertEquals(6, count); - builder.build(starTreeDocumentsList.iterator(), new AtomicInteger(), docValuesConsumer); - validateStarTree(builder.getRootNode(), 3, 10, builder.getStarTreeDocuments()); - } - - public void testMergeFlowWithTimestamps() throws IOException { - List dimList = List.of(1655288152000L, 1655288092000L, 1655288032000L, 1655287972000L, 1655288092000L, 1655288092000L); - List docsWithField = List.of(0, 1, 2, 3, 4, 6); - List dimList2 = List.of(1655288152000L, 1655288092000L, 1655288032000L, 1655287972000L, 1655288092000L, 1655288092000L, -1L); - List docsWithField2 = List.of(0, 1, 2, 3, 4, 6); - List dimList7 = List.of(1655288152000L, 1655288092000L, 1655288032000L, 1655287972000L, 1655288092000L, 1655288092000L, -1L); - List docsWithField7 = List.of(0, 1, 2, 3, 4, 6); - - List dimList5 = List.of(0L, 1L, 2L, 3L, 4L, 5L, 6L); - List docsWithField5 = List.of(0, 1, 2, 3, 4, 5, 6); - List metricsList1 = List.of( - getLongFromDouble(0.0), - getLongFromDouble(10.0), - getLongFromDouble(20.0), - getLongFromDouble(30.0), - getLongFromDouble(40.0), - getLongFromDouble(50.0), - getLongFromDouble(60.0) - ); - List metricsWithField1 = List.of(0, 1, 2, 3, 4, 5, 6); - List metricsList = List.of(0L, 1L, 2L, 3L, 4L, 5L, 6L); - List metricsWithField = List.of(0, 1, 2, 3, 4, 5, 6); - - List dimList3 = List.of(1655288152000L, 1655288092000L, 1655288032000L, -1L); - List docsWithField3 = List.of(0, 1, 3, 4); - List dimList4 = List.of(1655288152000L, 1655288092000L, 1655288032000L, -1L); - List docsWithField4 = List.of(0, 1, 3, 4); - List dimList8 = List.of(1655288152000L, 1655288092000L, 1655288032000L, -1L); - List docsWithField8 = List.of(0, 1, 3, 4); - - List dimList6 = List.of(5L, 6L, 7L, 8L); - List docsWithField6 = List.of(0, 1, 2, 3); - List metricsList21 = List.of( - getLongFromDouble(50.0), - getLongFromDouble(60.0), - getLongFromDouble(70.0), - getLongFromDouble(80.0), - getLongFromDouble(90.0) - ); - List metricsWithField21 = List.of(0, 1, 2, 3, 4); - List metricsList2 = List.of(5L, 6L, 7L, 8L, 9L); - List metricsWithField2 = List.of(0, 1, 2, 3, 4); - - compositeField = getStarTreeFieldWithDateDimension(); - StarTreeValues starTreeValues = getStarTreeValuesWithDates( - getSortedNumericMock(dimList, docsWithField), - getSortedNumericMock(dimList2, docsWithField2), - getSortedNumericMock(dimList7, docsWithField7), - getSortedNumericMock(dimList5, docsWithField5), - getSortedNumericMock(metricsList, metricsWithField), - getSortedNumericMock(metricsList1, metricsWithField1), - compositeField, - "6" - ); - - StarTreeValues starTreeValues2 = getStarTreeValuesWithDates( - getSortedNumericMock(dimList3, docsWithField3), - getSortedNumericMock(dimList4, docsWithField4), - getSortedNumericMock(dimList8, docsWithField8), - getSortedNumericMock(dimList6, docsWithField6), - getSortedNumericMock(metricsList2, metricsWithField2), - getSortedNumericMock(metricsList21, metricsWithField21), - compositeField, - "4" - ); - this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( - writeState, - Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, - Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, - Composite99DocValuesFormat.META_DOC_VALUES_CODEC, - Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION - ); - builder = getStarTreeBuilder(metaOut, dataOut, compositeField, getWriteState(4, writeState.segmentInfo.getId()), mapperService); - Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); - /** - [1655287972000, 1655287972000, 1655287972000, 3] | [30.0, 3] - [1655288032000, 1655288032000, 1655288032000, 2] | [20.0, 2] - [1655288032000, 1655288032000, 1655288032000, 8] | [80.0, 8] - [1655288092000, 1655288092000, 1655288092000, 1] | [10.0, 1] - [1655288092000, 1655288092000, 1655288092000, 4] | [40.0, 4] - [1655288092000, 1655288092000, 1655288092000, 6] | [60.0, 6] - [1655288152000, 1655288152000, 1655288152000, 0] | [0.0, 0] - [1655288152000, 1655288152000, 1655288152000, 5] | [50.0, 5] - [null, null, null, 5] | [50.0, 5] - [null, null, null, 7] | [70.0, 7] - */ - int count = 0; - builder.appendDocumentsToStarTree(starTreeDocumentIterator); - for (StarTreeDocument starTreeDocument : builder.getStarTreeDocuments()) { - count++; - assertEquals(starTreeDocument.dimensions[3] * 10.0, (double) starTreeDocument.metrics[1], 0); - assertEquals(starTreeDocument.dimensions[3], starTreeDocument.metrics[0]); - } - assertEquals(10, count); - builder.build(starTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); - validateStarTree(builder.getRootNode(), 4, 10, builder.getStarTreeDocuments()); - metaOut.close(); - dataOut.close(); - docValuesConsumer.close(); - - StarTreeMetadata starTreeMetadata = new StarTreeMetadata( - "sf", - STAR_TREE, - mock(IndexInput.class), - VERSION_CURRENT, - builder.numStarTreeNodes, - getStarTreeDimensionNames(compositeField.getDimensionsOrder()), - compositeField.getMetrics(), - 10, - builder.numStarTreeDocs, - compositeField.getStarTreeConfig().maxLeafDocs(), - Set.of(), - getBuildMode(), - 0, - 231 - ); - - validateStarTreeFileFormats( - builder.getRootNode(), - builder.getStarTreeDocuments().size(), - starTreeMetadata, - builder.getStarTreeDocuments() - ); - } - - private StarTreeValues getStarTreeValuesWithDates( - SortedNumericDocValues dimList, - SortedNumericDocValues dimList2, - SortedNumericDocValues dimList4, - SortedNumericDocValues dimList3, - SortedNumericDocValues metricsList, - SortedNumericDocValues metricsList1, - StarTreeField sf, - String number - ) { - Map> dimDocIdSetIterators = Map.of( - "field1_minute", - () -> new SortedNumericStarTreeValuesIterator(dimList), - "field1_half-hour", - () -> new SortedNumericStarTreeValuesIterator(dimList4), - "field1_hour", - () -> new SortedNumericStarTreeValuesIterator(dimList2), - "field3", - () -> new SortedNumericStarTreeValuesIterator(dimList3) - ); - Map> metricDocIdSetIterators = new LinkedHashMap<>(); - - metricDocIdSetIterators.put( - fullyQualifiedFieldNameForStarTreeMetricsDocValues( - sf.getName(), - "field2", - sf.getMetrics().get(0).getMetrics().get(0).getTypeName() - ), - () -> new SortedNumericStarTreeValuesIterator(metricsList) - ); - metricDocIdSetIterators.put( - fullyQualifiedFieldNameForStarTreeMetricsDocValues( - sf.getName(), - "field2", - sf.getMetrics().get(0).getMetrics().get(1).getTypeName() - ), - () -> new SortedNumericStarTreeValuesIterator(metricsList1) - ); - return new StarTreeValues(sf, null, dimDocIdSetIterators, metricDocIdSetIterators, Map.of(SEGMENT_DOCS_COUNT, number), null); - } - - private StarTreeField getStarTreeFieldWithDateDimension() { - List intervals = new ArrayList<>(); - intervals.add(new DateTimeUnitAdapter(Rounding.DateTimeUnit.MINUTES_OF_HOUR)); - intervals.add(new DateTimeUnitAdapter(Rounding.DateTimeUnit.HOUR_OF_DAY)); - intervals.add(DataCubeDateTimeUnit.HALF_HOUR_OF_DAY); - Dimension d1 = new DateDimension("field1", intervals, DateFieldMapper.Resolution.MILLISECONDS); - Dimension d2 = new NumericDimension("field3"); - Metric m1 = new Metric("field2", List.of(MetricStat.VALUE_COUNT, MetricStat.SUM)); - List dims = List.of(d1, d2); - List metrics = List.of(m1); - StarTreeFieldConfiguration c = new StarTreeFieldConfiguration(10, new HashSet<>(), getBuildMode()); - StarTreeField sf = new StarTreeField("sf", dims, metrics, c); - return sf; - } - - private void validateStarTree( - InMemoryTreeNode root, - int totalDimensions, - int maxLeafDocuments, - List starTreeDocuments - ) { - Queue queue = new LinkedList<>(); - queue.offer(new Object[] { root, false }); - while (!queue.isEmpty()) { - Object[] current = queue.poll(); - InMemoryTreeNode node = (InMemoryTreeNode) current[0]; - boolean currentIsStarNode = (boolean) current[1]; - - assertNotNull(node); - - // assert dimensions - if (node.getDimensionId() != StarTreeUtils.ALL) { - assertTrue(node.getDimensionId() >= 0 && node.getDimensionId() < totalDimensions); - } - - if (node.getChildren() != null && !node.getChildren().isEmpty()) { - assertEquals(node.getDimensionId() + 1, node.getChildDimensionId()); - assertTrue(node.getChildDimensionId() < totalDimensions); - InMemoryTreeNode starNode = null; - Object[] nonStarNodeCumulativeMetrics = getMetrics(starTreeDocuments); - for (Map.Entry entry : node.getChildren().entrySet()) { - Long childDimensionValue = entry.getKey(); - InMemoryTreeNode child = entry.getValue(); - Object[] currMetrics = getMetrics(starTreeDocuments); - if (child.getNodeType() != StarTreeNodeType.STAR.getValue()) { - // Validate dimension values in documents - for (int i = child.getStartDocId(); i < child.getEndDocId(); i++) { - StarTreeDocument doc = starTreeDocuments.get(i); - int j = 0; - addMetrics(doc, currMetrics, j); - if (child.getNodeType() != StarTreeNodeType.STAR.getValue()) { - Long dimension = doc.dimensions[child.getDimensionId()]; - assertEquals(childDimensionValue, dimension); - if (dimension != null) { - assertEquals(child.getDimensionValue(), (long) dimension); - } else { - // TODO : fix this ? - assertEquals(child.getDimensionValue(), StarTreeUtils.ALL); - } - } - } - Object[] aggregatedMetrics = starTreeDocuments.get(child.getAggregatedDocId()).metrics; - int j = 0; - for (Object metric : currMetrics) { - /* - * TODO : refactor this to handle any data type - */ - if (metric instanceof Double) { - nonStarNodeCumulativeMetrics[j] = (double) nonStarNodeCumulativeMetrics[j] + (double) metric; - assertEquals((Double) metric, (Double) aggregatedMetrics[j], 0); - } else if (metric instanceof Long) { - nonStarNodeCumulativeMetrics[j] = (long) nonStarNodeCumulativeMetrics[j] + (long) metric; - assertEquals((long) metric, (long) aggregatedMetrics[j]); - } else if (metric instanceof Float) { - nonStarNodeCumulativeMetrics[j] = (float) nonStarNodeCumulativeMetrics[j] + (float) metric; - assertEquals((float) metric, (float) aggregatedMetrics[j], 0); - } - j++; - } - queue.offer(new Object[] { child, false }); - } else { - starNode = child; - } - } - // Add star node to queue - if (starNode != null) { - Object[] starNodeMetrics = getMetrics(starTreeDocuments); - for (int i = starNode.getStartDocId(); i < starNode.getEndDocId(); i++) { - StarTreeDocument doc = starTreeDocuments.get(i); - int j = 0; - addMetrics(doc, starNodeMetrics, j); - } - int j = 0; - Object[] aggregatedMetrics = starTreeDocuments.get(starNode.getAggregatedDocId()).metrics; - for (Object nonStarNodeCumulativeMetric : nonStarNodeCumulativeMetrics) { - assertEquals(nonStarNodeCumulativeMetric, starNodeMetrics[j]); - assertEquals(starNodeMetrics[j], aggregatedMetrics[j]); - /* - * TODO : refactor this to handle any data type - */ - if (nonStarNodeCumulativeMetric instanceof Double) { - assertEquals((double) nonStarNodeCumulativeMetric, (double) starNodeMetrics[j], 0); - assertEquals((double) nonStarNodeCumulativeMetric, (double) aggregatedMetrics[j], 0); - } else if (nonStarNodeCumulativeMetric instanceof Long) { - assertEquals((long) nonStarNodeCumulativeMetric, (long) starNodeMetrics[j]); - assertEquals((long) nonStarNodeCumulativeMetric, (long) aggregatedMetrics[j]); - } else if (nonStarNodeCumulativeMetric instanceof Float) { - assertEquals((float) nonStarNodeCumulativeMetric, (float) starNodeMetrics[j], 0); - assertEquals((float) nonStarNodeCumulativeMetric, (float) aggregatedMetrics[j], 0); - } - - j++; - } - assertEquals(-1L, starNode.getDimensionValue()); - queue.offer(new Object[] { starNode, true }); - } - } else { - assertTrue(node.getEndDocId() - node.getStartDocId() <= maxLeafDocuments); - } - - if (currentIsStarNode) { - StarTreeDocument prevDoc = null; - int docCount = 0; - int docId = node.getStartDocId(); - int dimensionId = node.getDimensionId(); - - while (docId < node.getEndDocId()) { - StarTreeDocument currentDoc = starTreeDocuments.get(docId); - docCount++; - - // Verify that the dimension at 'dimensionId' is set to STAR_IN_DOC_VALUES_INDEX - assertNull(currentDoc.dimensions[dimensionId]); - - // Verify sorting of documents - if (prevDoc != null) { - assertTrue(compareDocuments(prevDoc, currentDoc, dimensionId + 1, totalDimensions) <= 0); - } - prevDoc = currentDoc; - docId++; - } - - // Verify that the number of generated star documents matches the range in the star node - assertEquals(node.getEndDocId() - node.getStartDocId(), docCount); - } - } - } - - /** - * TODO : refactor this to handle any data type - */ - private static void addMetrics(StarTreeDocument doc, Object[] currMetrics, int j) { - for (Object metric : doc.metrics) { - if (metric instanceof Double) { - currMetrics[j] = (double) currMetrics[j] + (double) metric; - } else if (metric instanceof Long) { - currMetrics[j] = (long) currMetrics[j] + (long) metric; - } else if (metric instanceof Float) { - currMetrics[j] = (float) currMetrics[j] + (float) metric; - } - j++; - } - } - - private static Object[] getMetrics(List starTreeDocuments) { - Object[] nonStarNodeCumulativeMetrics = new Object[starTreeDocuments.get(0).metrics.length]; - for (int i = 0; i < nonStarNodeCumulativeMetrics.length; i++) { - if (starTreeDocuments.get(0).metrics[i] instanceof Long) { - nonStarNodeCumulativeMetrics[i] = 0L; - } else if (starTreeDocuments.get(0).metrics[i] instanceof Double) { - nonStarNodeCumulativeMetrics[i] = 0.0; - } else if (starTreeDocuments.get(0).metrics[i] instanceof Float) { - nonStarNodeCumulativeMetrics[i] = 0.0f; - } - } - return nonStarNodeCumulativeMetrics; - } - - private int compareDocuments(StarTreeDocument doc1, StarTreeDocument doc2, int startDim, int endDim) { - for (int i = startDim; i < endDim; i++) { - Long val1 = doc1.dimensions[i]; - Long val2 = doc2.dimensions[i]; - - if (!Objects.equals(val1, val2)) { - if (val1 == null) return 1; - if (val2 == null) return -1; - return Long.compare(val1, val2); - } - } - return 0; - } - - Map getAttributes(int numSegmentDocs) { - return Map.of(CompositeIndexConstants.SEGMENT_DOCS_COUNT, String.valueOf(numSegmentDocs)); - } - - private StarTreeField getStarTreeField(MetricStat count) { - Dimension d1 = new NumericDimension("field1"); - Dimension d2 = new NumericDimension("field3"); - Metric m1 = new Metric("field2", List.of(count)); - List dims = List.of(d1, d2); - List metrics = List.of(m1); - StarTreeFieldConfiguration c = new StarTreeFieldConfiguration(1000, new HashSet<>(), getBuildMode()); - return new StarTreeField("sf", dims, metrics, c); - } - - SortedNumericDocValues getSortedNumericMock(List dimList, List docsWithField) { - return new SortedNumericDocValues() { - int index = -1; - - @Override - public long nextValue() { - return dimList.get(index); - } - - @Override - public int docValueCount() { - return 0; - } - - @Override - public boolean advanceExact(int target) { - return false; - } - - @Override - public int docID() { - return index; - } - - @Override - public int nextDoc() { - if (index == docsWithField.size() - 1) { - return NO_MORE_DOCS; - } - index++; - return docsWithField.get(index); - } - - @Override - public int advance(int target) { - return 0; - } - - @Override - public long cost() { - return 0; - } - }; - } - - @Override - public void tearDown() throws Exception { - super.tearDown(); - if (builder != null) { - builder.close(); - } - docValuesConsumer.close(); - metaOut.close(); - dataOut.close(); - directory.close(); - } -} diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/BuilderTestsUtils.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/BuilderTestsUtils.java new file mode 100644 index 0000000000000..bb31bd6a7cc27 --- /dev/null +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/BuilderTestsUtils.java @@ -0,0 +1,526 @@ +/* + * 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.startree.builder; + +import org.apache.lucene.codecs.DocValuesProducer; +import org.apache.lucene.codecs.lucene912.Lucene912Codec; +import org.apache.lucene.index.DocValuesType; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.FieldInfos; +import org.apache.lucene.index.IndexOptions; +import org.apache.lucene.index.SegmentInfo; +import org.apache.lucene.index.SegmentReadState; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.index.SortedNumericDocValues; +import org.apache.lucene.index.VectorEncoding; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.IOContext; +import org.apache.lucene.store.IndexInput; +import org.apache.lucene.util.InfoStream; +import org.apache.lucene.util.Version; +import org.opensearch.index.codec.composite.LuceneDocValuesProducerFactory; +import org.opensearch.index.codec.composite.composite99.Composite99Codec; +import org.opensearch.index.codec.composite.composite99.Composite99DocValuesFormat; +import org.opensearch.index.compositeindex.datacube.Metric; +import org.opensearch.index.compositeindex.datacube.MetricStat; +import org.opensearch.index.compositeindex.datacube.startree.StarTreeDocument; +import org.opensearch.index.compositeindex.datacube.startree.StarTreeField; +import org.opensearch.index.compositeindex.datacube.startree.StarTreeTestUtils; +import org.opensearch.index.compositeindex.datacube.startree.fileformats.meta.StarTreeMetadata; +import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues; +import org.opensearch.index.compositeindex.datacube.startree.node.InMemoryTreeNode; +import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNodeType; +import org.opensearch.index.compositeindex.datacube.startree.utils.SequentialDocValuesIterator; +import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils; +import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.SortedNumericStarTreeValuesIterator; +import org.opensearch.index.mapper.FieldValueConverter; + +import java.io.IOException; +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.Iterator; +import java.util.LinkedList; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Queue; + +import static org.opensearch.index.compositeindex.datacube.startree.StarTreeTestUtils.validateFileFormats; +import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils.fullyQualifiedFieldNameForStarTreeDimensionsDocValues; +import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues; +import static org.apache.lucene.tests.util.LuceneTestCase.newIOContext; +import static org.apache.lucene.tests.util.LuceneTestCase.random; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; + +public class BuilderTestsUtils { + public static SequentialDocValuesIterator[] getDimensionIterators(StarTreeDocument[] starTreeDocuments) { + SequentialDocValuesIterator[] sequentialDocValuesIterators = + new SequentialDocValuesIterator[starTreeDocuments[0].dimensions.length]; + for (int j = 0; j < starTreeDocuments[0].dimensions.length; j++) { + List dimList = new ArrayList<>(); + List docsWithField = new ArrayList<>(); + + for (int i = 0; i < starTreeDocuments.length; i++) { + if (starTreeDocuments[i].dimensions[j] != null) { + dimList.add(starTreeDocuments[i].dimensions[j]); + docsWithField.add(i); + } + } + sequentialDocValuesIterators[j] = new SequentialDocValuesIterator( + new SortedNumericStarTreeValuesIterator(getSortedNumericMock(dimList, docsWithField)) + ); + } + return sequentialDocValuesIterators; + } + + public static List getMetricIterators(StarTreeDocument[] starTreeDocuments) { + List sequentialDocValuesIterators = new ArrayList<>(); + for (int j = 0; j < starTreeDocuments[0].metrics.length; j++) { + List metricslist = new ArrayList<>(); + List docsWithField = new ArrayList<>(); + + for (int i = 0; i < starTreeDocuments.length; i++) { + if (starTreeDocuments[i].metrics[j] != null) { + metricslist.add((long) starTreeDocuments[i].metrics[j]); + docsWithField.add(i); + } + } + sequentialDocValuesIterators.add( + new SequentialDocValuesIterator(new SortedNumericStarTreeValuesIterator(getSortedNumericMock(metricslist, docsWithField))) + ); + } + return sequentialDocValuesIterators; + } + + public static SortedNumericDocValues getSortedNumericMock(List dimList, List docsWithField) { + return new SortedNumericDocValues() { + int index = -1; + + @Override + public long nextValue() { + return dimList.get(index); + } + + @Override + public int docValueCount() { + return 0; + } + + @Override + public boolean advanceExact(int target) { + return false; + } + + @Override + public int docID() { + return index; + } + + @Override + public int nextDoc() { + if (index == docsWithField.size() - 1) { + return NO_MORE_DOCS; + } + index++; + return docsWithField.get(index); + } + + @Override + public int advance(int target) { + return 0; + } + + @Override + public long cost() { + return 0; + } + }; + } + + public static void validateStarTree( + InMemoryTreeNode root, + int totalDimensions, + int maxLeafDocuments, + List starTreeDocuments + ) { + Queue queue = new LinkedList<>(); + queue.offer(new Object[] { root, false }); + while (!queue.isEmpty()) { + Object[] current = queue.poll(); + InMemoryTreeNode node = (InMemoryTreeNode) current[0]; + boolean currentIsStarNode = (boolean) current[1]; + + assertNotNull(node); + + // assert dimensions + if (node.getDimensionId() != StarTreeUtils.ALL) { + assertTrue(node.getDimensionId() >= 0 && node.getDimensionId() < totalDimensions); + } + + if (node.getChildren() != null && !node.getChildren().isEmpty()) { + assertEquals(node.getDimensionId() + 1, node.getChildDimensionId()); + assertTrue(node.getChildDimensionId() < totalDimensions); + InMemoryTreeNode starNode = null; + Object[] nonStarNodeCumulativeMetrics = getMetrics(starTreeDocuments); + for (Map.Entry entry : node.getChildren().entrySet()) { + Long childDimensionValue = entry.getKey(); + InMemoryTreeNode child = entry.getValue(); + Object[] currMetrics = getMetrics(starTreeDocuments); + if (child.getNodeType() != StarTreeNodeType.STAR.getValue()) { + // Validate dimension values in documents + for (int i = child.getStartDocId(); i < child.getEndDocId(); i++) { + StarTreeDocument doc = starTreeDocuments.get(i); + int j = 0; + addMetrics(doc, currMetrics, j); + if (child.getNodeType() != StarTreeNodeType.STAR.getValue()) { + Long dimension = doc.dimensions[child.getDimensionId()]; + assertEquals(childDimensionValue, dimension); + if (dimension != null) { + assertEquals(child.getDimensionValue(), (long) dimension); + } else { + // TODO : fix this ? + assertEquals(child.getDimensionValue(), StarTreeUtils.ALL); + } + } + } + Object[] aggregatedMetrics = starTreeDocuments.get(child.getAggregatedDocId()).metrics; + int j = 0; + for (Object metric : currMetrics) { + /* + * TODO : refactor this to handle any data type + */ + if (metric instanceof Double) { + nonStarNodeCumulativeMetrics[j] = (double) nonStarNodeCumulativeMetrics[j] + (double) metric; + assertEquals((Double) metric, (Double) aggregatedMetrics[j], 0); + } else if (metric instanceof Long) { + nonStarNodeCumulativeMetrics[j] = (long) nonStarNodeCumulativeMetrics[j] + (long) metric; + assertEquals((long) metric, (long) aggregatedMetrics[j]); + } else if (metric instanceof Float) { + nonStarNodeCumulativeMetrics[j] = (float) nonStarNodeCumulativeMetrics[j] + (float) metric; + assertEquals((float) metric, (float) aggregatedMetrics[j], 0); + } + j++; + } + queue.offer(new Object[] { child, false }); + } else { + starNode = child; + } + } + // Add star node to queue + if (starNode != null) { + Object[] starNodeMetrics = getMetrics(starTreeDocuments); + for (int i = starNode.getStartDocId(); i < starNode.getEndDocId(); i++) { + StarTreeDocument doc = starTreeDocuments.get(i); + int j = 0; + addMetrics(doc, starNodeMetrics, j); + } + int j = 0; + Object[] aggregatedMetrics = starTreeDocuments.get(starNode.getAggregatedDocId()).metrics; + for (Object nonStarNodeCumulativeMetric : nonStarNodeCumulativeMetrics) { + assertEquals(nonStarNodeCumulativeMetric, starNodeMetrics[j]); + assertEquals(starNodeMetrics[j], aggregatedMetrics[j]); + /* + * TODO : refactor this to handle any data type + */ + if (nonStarNodeCumulativeMetric instanceof Double) { + assertEquals((double) nonStarNodeCumulativeMetric, (double) starNodeMetrics[j], 0); + assertEquals((double) nonStarNodeCumulativeMetric, (double) aggregatedMetrics[j], 0); + } else if (nonStarNodeCumulativeMetric instanceof Long) { + assertEquals((long) nonStarNodeCumulativeMetric, (long) starNodeMetrics[j]); + assertEquals((long) nonStarNodeCumulativeMetric, (long) aggregatedMetrics[j]); + } else if (nonStarNodeCumulativeMetric instanceof Float) { + assertEquals((float) nonStarNodeCumulativeMetric, (float) starNodeMetrics[j], 0); + assertEquals((float) nonStarNodeCumulativeMetric, (float) aggregatedMetrics[j], 0); + } + + j++; + } + assertEquals(-1L, starNode.getDimensionValue()); + queue.offer(new Object[] { starNode, true }); + } + } else { + assertTrue(node.getEndDocId() - node.getStartDocId() <= maxLeafDocuments); + } + + if (currentIsStarNode) { + StarTreeDocument prevDoc = null; + int docCount = 0; + int docId = node.getStartDocId(); + int dimensionId = node.getDimensionId(); + + while (docId < node.getEndDocId()) { + StarTreeDocument currentDoc = starTreeDocuments.get(docId); + docCount++; + + // Verify that the dimension at 'dimensionId' is set to STAR_IN_DOC_VALUES_INDEX + assertNull(currentDoc.dimensions[dimensionId]); + + // Verify sorting of documents + if (prevDoc != null) { + assertTrue(compareDocuments(prevDoc, currentDoc, dimensionId + 1, totalDimensions) <= 0); + } + prevDoc = currentDoc; + docId++; + } + + // Verify that the number of generated star documents matches the range in the star node + assertEquals(node.getEndDocId() - node.getStartDocId(), docCount); + } + } + } + + /** + * TODO : refactor this to handle any data type + */ + private static void addMetrics(StarTreeDocument doc, Object[] currMetrics, int j) { + for (Object metric : doc.metrics) { + if (metric instanceof Double) { + currMetrics[j] = (double) currMetrics[j] + (double) metric; + } else if (metric instanceof Long) { + currMetrics[j] = (long) currMetrics[j] + (long) metric; + } else if (metric instanceof Float) { + currMetrics[j] = (float) currMetrics[j] + (float) metric; + } + j++; + } + } + + private static Object[] getMetrics(List starTreeDocuments) { + Object[] nonStarNodeCumulativeMetrics = new Object[starTreeDocuments.get(0).metrics.length]; + for (int i = 0; i < nonStarNodeCumulativeMetrics.length; i++) { + if (starTreeDocuments.get(0).metrics[i] instanceof Long) { + nonStarNodeCumulativeMetrics[i] = 0L; + } else if (starTreeDocuments.get(0).metrics[i] instanceof Double) { + nonStarNodeCumulativeMetrics[i] = 0.0; + } else if (starTreeDocuments.get(0).metrics[i] instanceof Float) { + nonStarNodeCumulativeMetrics[i] = 0.0f; + } + } + return nonStarNodeCumulativeMetrics; + } + + private static int compareDocuments(StarTreeDocument doc1, StarTreeDocument doc2, int startDim, int endDim) { + for (int i = startDim; i < endDim; i++) { + Long val1 = doc1.dimensions[i]; + Long val2 = doc2.dimensions[i]; + + if (!Objects.equals(val1, val2)) { + if (val1 == null) return 1; + if (val2 == null) return -1; + return Long.compare(val1, val2); + } + } + return 0; + } + + public static void validateStarTreeFileFormats( + InMemoryTreeNode rootNode, + int numDocs, + StarTreeMetadata expectedStarTreeMetadata, + List expectedStarTreeDocuments, + String dataFileName, + String metaFileName, + BaseStarTreeBuilder builder, + StarTreeField starTreeField, + SegmentWriteState writeState, + Directory directory + ) throws IOException { + + assertNotNull(rootNode.getChildren()); + assertFalse(rootNode.getChildren().isEmpty()); + SegmentReadState readState = getReadState( + numDocs, + expectedStarTreeMetadata.getDimensionFields(), + expectedStarTreeMetadata.getMetrics(), + starTreeField, + writeState, + directory + ); + + DocValuesProducer compositeDocValuesProducer = LuceneDocValuesProducerFactory.getDocValuesProducerForCompositeCodec( + Composite99Codec.COMPOSITE_INDEX_CODEC_NAME, + readState, + Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, + Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, + Composite99DocValuesFormat.META_DOC_VALUES_CODEC, + Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION + ); + + IndexInput dataIn = readState.directory.openInput(dataFileName, IOContext.DEFAULT); + IndexInput metaIn = readState.directory.openInput(metaFileName, IOContext.DEFAULT); + + StarTreeValues starTreeValues = new StarTreeValues(expectedStarTreeMetadata, dataIn, compositeDocValuesProducer, readState); + assertEquals(expectedStarTreeMetadata.getStarTreeDocCount(), starTreeValues.getStarTreeDocumentCount()); + List fieldValueConverters = new ArrayList<>(); + builder.metricAggregatorInfos.forEach( + metricAggregatorInfo -> fieldValueConverters.add(metricAggregatorInfo.getValueAggregators().getAggregatedValueType()) + ); + StarTreeDocument[] starTreeDocuments = StarTreeTestUtils.getSegmentsStarTreeDocuments( + List.of(starTreeValues), + fieldValueConverters, + readState.segmentInfo.maxDoc() + ); + + StarTreeDocument[] expectedStarTreeDocumentsArray = expectedStarTreeDocuments.toArray(new StarTreeDocument[0]); + StarTreeTestUtils.assertStarTreeDocuments(starTreeDocuments, expectedStarTreeDocumentsArray); + + validateFileFormats(dataIn, metaIn, rootNode, expectedStarTreeMetadata); + + dataIn.close(); + metaIn.close(); + compositeDocValuesProducer.close(); + } + + public static SegmentReadState getReadState( + int numDocs, + List dimensionFields, + List metrics, + StarTreeField compositeField, + SegmentWriteState writeState, + Directory directory + ) { + + int numMetrics = 0; + for (Metric metric : metrics) { + numMetrics += metric.getBaseMetrics().size(); + } + + FieldInfo[] fields = new FieldInfo[dimensionFields.size() + numMetrics]; + + int i = 0; + for (String dimension : dimensionFields) { + fields[i] = new FieldInfo( + fullyQualifiedFieldNameForStarTreeDimensionsDocValues(compositeField.getName(), dimension), + i, + false, + false, + true, + IndexOptions.DOCS_AND_FREQS_AND_POSITIONS_AND_OFFSETS, + DocValuesType.SORTED_NUMERIC, + -1, + Collections.emptyMap(), + 0, + 0, + 0, + 0, + VectorEncoding.FLOAT32, + VectorSimilarityFunction.EUCLIDEAN, + false, + false + ); + i++; + } + + for (Metric metric : metrics) { + for (MetricStat metricStat : metric.getBaseMetrics()) { + fields[i] = new FieldInfo( + fullyQualifiedFieldNameForStarTreeMetricsDocValues( + compositeField.getName(), + metric.getField(), + metricStat.getTypeName() + ), + i, + false, + false, + true, + IndexOptions.DOCS_AND_FREQS_AND_POSITIONS_AND_OFFSETS, + DocValuesType.SORTED_NUMERIC, + -1, + Collections.emptyMap(), + 0, + 0, + 0, + 0, + VectorEncoding.FLOAT32, + VectorSimilarityFunction.EUCLIDEAN, + false, + false + ); + i++; + } + } + + SegmentInfo segmentInfo = new SegmentInfo( + directory, + Version.LATEST, + Version.LUCENE_9_11_0, + "test_segment", + numDocs, + false, + false, + new Lucene912Codec(), + new HashMap<>(), + writeState.segmentInfo.getId(), + new HashMap<>(), + null + ); + return new SegmentReadState(segmentInfo.dir, segmentInfo, new FieldInfos(fields), writeState.context); + } + + static void traverseStarTree(InMemoryTreeNode root, Map> dimValueToDocIdMap, boolean traverStarNodes) { + InMemoryTreeNode starTree = root; + // Use BFS to traverse the star tree + Queue queue = new ArrayDeque<>(); + queue.add(starTree); + int currentDimensionId = -1; + InMemoryTreeNode starTreeNode; + List docIds = new ArrayList<>(); + while ((starTreeNode = queue.poll()) != null) { + int dimensionId = starTreeNode.getDimensionId(); + if (dimensionId > currentDimensionId) { + currentDimensionId = dimensionId; + } + + // store aggregated document of the node + int docId = starTreeNode.getAggregatedDocId(); + Map map = dimValueToDocIdMap.getOrDefault(dimensionId, new HashMap<>()); + if (starTreeNode.getNodeType() == StarTreeNodeType.STAR.getValue()) { + map.put(Long.MAX_VALUE, docId); + } else { + map.put(starTreeNode.getDimensionValue(), docId); + } + dimValueToDocIdMap.put(dimensionId, map); + + if (starTreeNode.getChildren() != null + && (!traverStarNodes || starTreeNode.getNodeType() == StarTreeNodeType.STAR.getValue())) { + Iterator childrenIterator = starTreeNode.getChildren().values().iterator(); + while (childrenIterator.hasNext()) { + InMemoryTreeNode childNode = childrenIterator.next(); + queue.add(childNode); + } + } + } + } + + public static SegmentWriteState getWriteState(int numDocs, byte[] id, FieldInfo[] fieldsInfo, Directory directory) { + FieldInfos fieldInfos = new FieldInfos(fieldsInfo); + SegmentInfo segmentInfo = new SegmentInfo( + directory, + Version.LATEST, + Version.LUCENE_9_11_0, + "test_segment", + numDocs, + false, + false, + new Lucene912Codec(), + new HashMap<>(), + id, + new HashMap<>(), + null + ); + return new SegmentWriteState(InfoStream.getDefault(), segmentInfo.dir, segmentInfo, fieldInfos, null, newIOContext(random())); + } +} diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/OffHeapStarTreeBuilderTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/OffHeapStarTreeBuilderTests.java deleted file mode 100644 index 496558dbc2e83..0000000000000 --- a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/OffHeapStarTreeBuilderTests.java +++ /dev/null @@ -1,36 +0,0 @@ -/* - * 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.startree.builder; - -import org.apache.lucene.index.SegmentWriteState; -import org.apache.lucene.store.IndexOutput; -import org.opensearch.index.compositeindex.datacube.startree.StarTreeField; -import org.opensearch.index.compositeindex.datacube.startree.StarTreeFieldConfiguration; -import org.opensearch.index.mapper.MapperService; - -import java.io.IOException; - -public class OffHeapStarTreeBuilderTests extends AbstractStarTreeBuilderTests { - @Override - public BaseStarTreeBuilder getStarTreeBuilder( - IndexOutput metaOut, - IndexOutput dataOut, - StarTreeField starTreeField, - SegmentWriteState segmentWriteState, - MapperService mapperService - ) throws IOException { - return new OffHeapStarTreeBuilder(metaOut, dataOut, starTreeField, segmentWriteState, mapperService); - } - - @Override - StarTreeFieldConfiguration.StarTreeBuildMode getBuildMode() { - return StarTreeFieldConfiguration.StarTreeBuildMode.OFF_HEAP; - } - -} diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/OnHeapStarTreeBuilderTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/OnHeapStarTreeBuilderTests.java deleted file mode 100644 index 55cf3bde3cea7..0000000000000 --- a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/OnHeapStarTreeBuilderTests.java +++ /dev/null @@ -1,37 +0,0 @@ -/* - * 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.startree.builder; - -import org.apache.lucene.index.SegmentWriteState; -import org.apache.lucene.store.IndexOutput; -import org.opensearch.index.compositeindex.datacube.startree.StarTreeField; -import org.opensearch.index.compositeindex.datacube.startree.StarTreeFieldConfiguration; -import org.opensearch.index.mapper.MapperService; - -import java.io.IOException; - -public class OnHeapStarTreeBuilderTests extends AbstractStarTreeBuilderTests { - - @Override - public BaseStarTreeBuilder getStarTreeBuilder( - IndexOutput metaOut, - IndexOutput dataOut, - StarTreeField starTreeField, - SegmentWriteState segmentWriteState, - MapperService mapperService - ) throws IOException { - return new OnHeapStarTreeBuilder(metaOut, dataOut, starTreeField, segmentWriteState, mapperService); - } - - @Override - StarTreeFieldConfiguration.StarTreeBuildMode getBuildMode() { - return StarTreeFieldConfiguration.StarTreeBuildMode.ON_HEAP; - } - -} diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuildMetricTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuildMetricTests.java new file mode 100644 index 0000000000000..095eda2986b3a --- /dev/null +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuildMetricTests.java @@ -0,0 +1,954 @@ +/* + * 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.startree.builder; + +import org.apache.lucene.codecs.DocValuesConsumer; +import org.apache.lucene.codecs.DocValuesProducer; +import org.apache.lucene.codecs.lucene912.Lucene912Codec; +import org.apache.lucene.index.DocValuesType; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.FieldInfos; +import org.apache.lucene.index.IndexOptions; +import org.apache.lucene.index.SegmentInfo; +import org.apache.lucene.index.SegmentReadState; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.index.VectorEncoding; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.sandbox.document.HalfFloatPoint; +import org.apache.lucene.store.IOContext; +import org.apache.lucene.store.IndexInput; +import org.apache.lucene.util.InfoStream; +import org.apache.lucene.util.NumericUtils; +import org.apache.lucene.util.Version; +import org.opensearch.common.settings.Settings; +import org.opensearch.index.codec.composite.LuceneDocValuesConsumerFactory; +import org.opensearch.index.codec.composite.composite99.Composite99DocValuesFormat; +import org.opensearch.index.compositeindex.datacube.Metric; +import org.opensearch.index.compositeindex.datacube.MetricStat; +import org.opensearch.index.compositeindex.datacube.NumericDimension; +import org.opensearch.index.compositeindex.datacube.startree.StarTreeDocument; +import org.opensearch.index.compositeindex.datacube.startree.StarTreeField; +import org.opensearch.index.compositeindex.datacube.startree.StarTreeFieldConfiguration; +import org.opensearch.index.compositeindex.datacube.startree.fileformats.meta.StarTreeMetadata; +import org.opensearch.index.compositeindex.datacube.startree.node.InMemoryTreeNode; +import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNodeType; +import org.opensearch.index.compositeindex.datacube.startree.utils.SequentialDocValuesIterator; +import org.opensearch.index.mapper.ContentPath; +import org.opensearch.index.mapper.DocumentMapper; +import org.opensearch.index.mapper.Mapper; +import org.opensearch.index.mapper.MapperService; +import org.opensearch.index.mapper.MappingLookup; +import org.opensearch.index.mapper.NumberFieldMapper; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.UUID; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.opensearch.index.compositeindex.datacube.startree.StarTreeTestUtils.validateFileFormats; +import static org.opensearch.index.compositeindex.datacube.startree.builder.BuilderTestsUtils.getDimensionIterators; +import static org.opensearch.index.compositeindex.datacube.startree.builder.BuilderTestsUtils.getMetricIterators; +import static org.opensearch.index.compositeindex.datacube.startree.builder.BuilderTestsUtils.traverseStarTree; +import static org.opensearch.index.compositeindex.datacube.startree.builder.BuilderTestsUtils.validateStarTree; +import static org.opensearch.index.compositeindex.datacube.startree.fileformats.StarTreeWriter.VERSION_CURRENT; +import static org.opensearch.index.mapper.CompositeMappedFieldType.CompositeFieldType.STAR_TREE; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class StarTreeBuildMetricTests extends StarTreeBuilderTestCase { + + public StarTreeBuildMetricTests(StarTreeFieldConfiguration.StarTreeBuildMode buildMode) { + super(buildMode); + } + + public void test_build_halfFloatMetrics() throws IOException { + + mapperService = mock(MapperService.class); + 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("field2", NumberFieldMapper.NumberType.HALF_FLOAT, false, true) + .build(new Mapper.BuilderContext(settings, new ContentPath())); + NumberFieldMapper numberFieldMapper2 = new NumberFieldMapper.Builder("field4", NumberFieldMapper.NumberType.HALF_FLOAT, false, true) + .build(new Mapper.BuilderContext(settings, new ContentPath())); + NumberFieldMapper numberFieldMapper3 = new NumberFieldMapper.Builder("field6", NumberFieldMapper.NumberType.HALF_FLOAT, false, true) + .build(new Mapper.BuilderContext(settings, new ContentPath())); + NumberFieldMapper numberFieldMapper4 = new NumberFieldMapper.Builder("field9", NumberFieldMapper.NumberType.HALF_FLOAT, false, true) + .build(new Mapper.BuilderContext(settings, new ContentPath())); + NumberFieldMapper numberFieldMapper5 = new NumberFieldMapper.Builder( + "field10", + NumberFieldMapper.NumberType.HALF_FLOAT, + false, + true + ).build(new Mapper.BuilderContext(settings, new ContentPath())); + MappingLookup fieldMappers = new MappingLookup( + Set.of(numberFieldMapper1, numberFieldMapper2, numberFieldMapper3, numberFieldMapper4, numberFieldMapper5), + Collections.emptyList(), + Collections.emptyList(), + 0, + null + ); + when(documentMapper.mappers()).thenReturn(fieldMappers); + + int noOfStarTreeDocuments = 5; + StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; + + starTreeDocuments[0] = new StarTreeDocument( + new Long[] { 2L, 4L, 3L, 4L }, + new HalfFloatPoint[] { + new HalfFloatPoint("hf1", 12), + new HalfFloatPoint("hf6", 10), + new HalfFloatPoint("field6", 10), + new HalfFloatPoint("field9", 8), + new HalfFloatPoint("field10", 20) } + ); + starTreeDocuments[1] = new StarTreeDocument( + new Long[] { 3L, 4L, 2L, 1L }, + new HalfFloatPoint[] { + new HalfFloatPoint("hf2", 10), + new HalfFloatPoint("hf7", 6), + new HalfFloatPoint("field6", 10), + new HalfFloatPoint("field9", 12), + new HalfFloatPoint("field10", 10) } + ); + starTreeDocuments[2] = new StarTreeDocument( + new Long[] { 3L, 4L, 2L, 1L }, + new HalfFloatPoint[] { + new HalfFloatPoint("hf3", 14), + new HalfFloatPoint("hf8", 12), + new HalfFloatPoint("field6", 10), + new HalfFloatPoint("field9", 6), + new HalfFloatPoint("field10", 24) } + ); + starTreeDocuments[3] = new StarTreeDocument( + new Long[] { 2L, 4L, 3L, 4L }, + new HalfFloatPoint[] { + new HalfFloatPoint("hf4", 9), + new HalfFloatPoint("hf9", 4), + new HalfFloatPoint("field6", 10), + new HalfFloatPoint("field9", 9), + new HalfFloatPoint("field10", 12) } + ); + starTreeDocuments[4] = new StarTreeDocument( + new Long[] { 3L, 4L, 2L, 1L }, + new HalfFloatPoint[] { + new HalfFloatPoint("hf5", 11), + new HalfFloatPoint("hf10", 16), + new HalfFloatPoint("field6", 10), + new HalfFloatPoint("field9", 8), + new HalfFloatPoint("field10", 13) } + ); + + StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; + for (int i = 0; i < noOfStarTreeDocuments; i++) { + long metric1 = HalfFloatPoint.halfFloatToSortableShort( + ((HalfFloatPoint) starTreeDocuments[i].metrics[0]).numericValue().floatValue() + ); + long metric2 = HalfFloatPoint.halfFloatToSortableShort( + ((HalfFloatPoint) starTreeDocuments[i].metrics[1]).numericValue().floatValue() + ); + long metric3 = HalfFloatPoint.halfFloatToSortableShort( + ((HalfFloatPoint) starTreeDocuments[i].metrics[2]).numericValue().floatValue() + ); + long metric4 = HalfFloatPoint.halfFloatToSortableShort( + ((HalfFloatPoint) starTreeDocuments[i].metrics[3]).numericValue().floatValue() + ); + long metric5 = HalfFloatPoint.halfFloatToSortableShort( + ((HalfFloatPoint) starTreeDocuments[i].metrics[4]).numericValue().floatValue() + ); + segmentStarTreeDocuments[i] = new StarTreeDocument( + starTreeDocuments[i].dimensions, + new Long[] { metric1, metric2, metric3, metric4, metric5, null } + ); + } + + SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); + List metricsIterators = getMetricIterators(segmentStarTreeDocuments); + this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( + writeState, + Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, + Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, + Composite99DocValuesFormat.META_DOC_VALUES_CODEC, + Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION + ); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + + Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( + dimsIterators, + metricsIterators + ); + builder.build(segmentStarTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); + List resultStarTreeDocuments = builder.getStarTreeDocuments(); + assertEquals(7, resultStarTreeDocuments.size()); + + Iterator expectedStarTreeDocumentIterator = getExpectedStarTreeDocumentIterator().iterator(); + assertStarTreeDocuments(resultStarTreeDocuments, expectedStarTreeDocumentIterator); + + metaOut.close(); + dataOut.close(); + docValuesConsumer.close(); + + StarTreeMetadata starTreeMetadata = new StarTreeMetadata( + "test", + STAR_TREE, + mock(IndexInput.class), + VERSION_CURRENT, + builder.numStarTreeNodes, + getStarTreeDimensionNames(compositeField.getDimensionsOrder()), + compositeField.getMetrics(), + 2, + getExpectedStarTreeDocumentIterator().size(), + 1, + Set.of("field8"), + getBuildMode(), + 0, + 330 + ); + + validateStarTreeFileFormats( + builder.getRootNode(), + getExpectedStarTreeDocumentIterator().size(), + starTreeMetadata, + getExpectedStarTreeDocumentIterator() + ); + } + + public void test_build_floatMetrics() throws IOException { + + mapperService = mock(MapperService.class); + 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("field2", NumberFieldMapper.NumberType.FLOAT, false, true) + .build(new Mapper.BuilderContext(settings, new ContentPath())); + NumberFieldMapper numberFieldMapper2 = new NumberFieldMapper.Builder("field4", NumberFieldMapper.NumberType.FLOAT, false, true) + .build(new Mapper.BuilderContext(settings, new ContentPath())); + NumberFieldMapper numberFieldMapper3 = new NumberFieldMapper.Builder("field6", NumberFieldMapper.NumberType.FLOAT, false, true) + .build(new Mapper.BuilderContext(settings, new ContentPath())); + NumberFieldMapper numberFieldMapper4 = new NumberFieldMapper.Builder("field9", NumberFieldMapper.NumberType.FLOAT, false, true) + .build(new Mapper.BuilderContext(settings, new ContentPath())); + NumberFieldMapper numberFieldMapper5 = new NumberFieldMapper.Builder("field10", NumberFieldMapper.NumberType.FLOAT, false, true) + .build(new Mapper.BuilderContext(settings, new ContentPath())); + MappingLookup fieldMappers = new MappingLookup( + Set.of(numberFieldMapper1, numberFieldMapper2, numberFieldMapper3, numberFieldMapper4, numberFieldMapper5), + Collections.emptyList(), + Collections.emptyList(), + 0, + null + ); + when(documentMapper.mappers()).thenReturn(fieldMappers); + + int noOfStarTreeDocuments = 5; + StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; + + starTreeDocuments[0] = new StarTreeDocument( + new Long[] { 2L, 4L, 3L, 4L }, + new Object[] { 12.0F, 10.0F, randomFloat(), 8.0F, 20.0F, null } + ); + starTreeDocuments[1] = new StarTreeDocument( + new Long[] { 3L, 4L, 2L, 1L }, + new Object[] { 10.0F, 6.0F, randomFloat(), 12.0F, 10.0F, null } + ); + starTreeDocuments[2] = new StarTreeDocument( + new Long[] { 3L, 4L, 2L, 1L }, + new Object[] { 14.0F, 12.0F, randomFloat(), 6.0F, 24.0F, null } + ); + starTreeDocuments[3] = new StarTreeDocument( + new Long[] { 2L, 4L, 3L, 4L }, + new Object[] { 9.0F, 4.0F, randomFloat(), 9.0F, 12.0F, null } + ); + starTreeDocuments[4] = new StarTreeDocument( + new Long[] { 3L, 4L, 2L, 1L }, + new Object[] { 11.0F, 16.0F, randomFloat(), 8.0F, 13.0F, null } + ); + + StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; + for (int i = 0; i < noOfStarTreeDocuments; i++) { + long metric1 = NumericUtils.floatToSortableInt((Float) starTreeDocuments[i].metrics[0]); + long metric2 = NumericUtils.floatToSortableInt((Float) starTreeDocuments[i].metrics[1]); + long metric3 = NumericUtils.floatToSortableInt((Float) starTreeDocuments[i].metrics[2]); + long metric4 = NumericUtils.floatToSortableInt((Float) starTreeDocuments[i].metrics[3]); + long metric5 = NumericUtils.floatToSortableInt((Float) starTreeDocuments[i].metrics[4]); + Long metric6 = (Long) starTreeDocuments[i].metrics[5]; + segmentStarTreeDocuments[i] = new StarTreeDocument( + starTreeDocuments[i].dimensions, + new Long[] { metric1, metric2, metric3, metric4, metric5, metric6 } + ); + } + + SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); + List metricsIterators = getMetricIterators(segmentStarTreeDocuments); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( + dimsIterators, + metricsIterators + ); + this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( + writeState, + Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, + Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, + Composite99DocValuesFormat.META_DOC_VALUES_CODEC, + Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION + ); + builder.build(segmentStarTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); + + List resultStarTreeDocuments = builder.getStarTreeDocuments(); + assertEquals(7, resultStarTreeDocuments.size()); + + Iterator expectedStarTreeDocumentIterator = getExpectedStarTreeDocumentIterator().iterator(); + assertStarTreeDocuments(resultStarTreeDocuments, expectedStarTreeDocumentIterator); + + metaOut.close(); + dataOut.close(); + docValuesConsumer.close(); + + StarTreeMetadata starTreeMetadata = new StarTreeMetadata( + "test", + STAR_TREE, + mock(IndexInput.class), + VERSION_CURRENT, + builder.numStarTreeNodes, + getStarTreeDimensionNames(compositeField.getDimensionsOrder()), + compositeField.getMetrics(), + 2, + getExpectedStarTreeDocumentIterator().size(), + 1, + Set.of("field8"), + getBuildMode(), + 0, + 330 + ); + + validateStarTreeFileFormats( + builder.getRootNode(), + getExpectedStarTreeDocumentIterator().size(), + starTreeMetadata, + getExpectedStarTreeDocumentIterator() + ); + } + + public void test_build_longMetrics() throws IOException { + + mapperService = mock(MapperService.class); + 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("field2", NumberFieldMapper.NumberType.LONG, false, true) + .build(new Mapper.BuilderContext(settings, new ContentPath())); + NumberFieldMapper numberFieldMapper2 = new NumberFieldMapper.Builder("field4", NumberFieldMapper.NumberType.LONG, false, true) + .build(new Mapper.BuilderContext(settings, new ContentPath())); + NumberFieldMapper numberFieldMapper3 = new NumberFieldMapper.Builder("field6", NumberFieldMapper.NumberType.LONG, false, true) + .build(new Mapper.BuilderContext(settings, new ContentPath())); + NumberFieldMapper numberFieldMapper4 = new NumberFieldMapper.Builder("field9", NumberFieldMapper.NumberType.LONG, false, true) + .build(new Mapper.BuilderContext(settings, new ContentPath())); + NumberFieldMapper numberFieldMapper5 = new NumberFieldMapper.Builder("field10", NumberFieldMapper.NumberType.LONG, false, true) + .build(new Mapper.BuilderContext(settings, new ContentPath())); + MappingLookup fieldMappers = new MappingLookup( + Set.of(numberFieldMapper1, numberFieldMapper2, numberFieldMapper3, numberFieldMapper4, numberFieldMapper5), + Collections.emptyList(), + Collections.emptyList(), + 0, + null + ); + when(documentMapper.mappers()).thenReturn(fieldMappers); + + int noOfStarTreeDocuments = 5; + StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; + + starTreeDocuments[0] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Long[] { 12L, 10L, randomLong(), 8L, 20L }); + starTreeDocuments[1] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Long[] { 10L, 6L, randomLong(), 12L, 10L }); + starTreeDocuments[2] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Long[] { 14L, 12L, randomLong(), 6L, 24L }); + starTreeDocuments[3] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Long[] { 9L, 4L, randomLong(), 9L, 12L }); + starTreeDocuments[4] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Long[] { 11L, 16L, randomLong(), 8L, 13L }); + + StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; + for (int i = 0; i < noOfStarTreeDocuments; i++) { + long metric1 = (Long) starTreeDocuments[i].metrics[0]; + long metric2 = (Long) starTreeDocuments[i].metrics[1]; + long metric3 = (Long) starTreeDocuments[i].metrics[2]; + long metric4 = (Long) starTreeDocuments[i].metrics[3]; + long metric5 = (Long) starTreeDocuments[i].metrics[4]; + segmentStarTreeDocuments[i] = new StarTreeDocument( + starTreeDocuments[i].dimensions, + new Long[] { metric1, metric2, metric3, metric4, metric5, null } + ); + } + + SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); + List metricsIterators = getMetricIterators(segmentStarTreeDocuments); + this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( + writeState, + Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, + Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, + Composite99DocValuesFormat.META_DOC_VALUES_CODEC, + Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION + ); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( + dimsIterators, + metricsIterators + ); + builder.build(segmentStarTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); + + List resultStarTreeDocuments = builder.getStarTreeDocuments(); + assertEquals(7, resultStarTreeDocuments.size()); + + Iterator expectedStarTreeDocumentIterator = getExpectedStarTreeDocumentIterator().iterator(); + assertStarTreeDocuments(resultStarTreeDocuments, expectedStarTreeDocumentIterator); + + metaOut.close(); + dataOut.close(); + docValuesConsumer.close(); + + StarTreeMetadata starTreeMetadata = new StarTreeMetadata( + "test", + STAR_TREE, + mock(IndexInput.class), + VERSION_CURRENT, + builder.numStarTreeNodes, + getStarTreeDimensionNames(compositeField.getDimensionsOrder()), + compositeField.getMetrics(), + 2, + getExpectedStarTreeDocumentIterator().size(), + 1, + Set.of("field8"), + getBuildMode(), + 0, + 330 + ); + + validateStarTreeFileFormats( + builder.getRootNode(), + getExpectedStarTreeDocumentIterator().size(), + starTreeMetadata, + getExpectedStarTreeDocumentIterator() + ); + } + + public void test_build_multipleStarTrees() throws IOException { + + int noOfStarTreeDocuments = 5; + StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; + + starTreeDocuments[0] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Double[] { 12.0, 10.0, randomDouble(), 8.0, 20.0 }); + starTreeDocuments[1] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 10.0, 6.0, randomDouble(), 12.0, 10.0 }); + starTreeDocuments[2] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 14.0, 12.0, randomDouble(), 6.0, 24.0 }); + starTreeDocuments[3] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Double[] { 9.0, 4.0, randomDouble(), 9.0, 12.0 }); + starTreeDocuments[4] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Double[] { 11.0, 16.0, randomDouble(), 8.0, 13.0 }); + + StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; + for (int i = 0; i < noOfStarTreeDocuments; i++) { + long metric1 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]); + long metric2 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[1]); + long metric3 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]); + long metric4 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[3]); + long metric5 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[4]); + segmentStarTreeDocuments[i] = new StarTreeDocument( + starTreeDocuments[i].dimensions, + new Long[] { metric1, metric2, metric3, metric4, metric5 } + ); + } + + SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); + List metricsIterators = getMetricIterators(segmentStarTreeDocuments); + + metrics = List.of( + new Metric("field2", List.of(MetricStat.SUM)), + new Metric("field4", List.of(MetricStat.SUM)), + new Metric("field6", List.of(MetricStat.VALUE_COUNT)), + new Metric("field9", List.of(MetricStat.MIN)), + new Metric("field10", List.of(MetricStat.MAX)) + ); + + compositeField = new StarTreeField( + "test", + dimensionsOrder, + metrics, + new StarTreeFieldConfiguration(1, Set.of("field8"), getBuildMode()) + ); + + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( + dimsIterators, + metricsIterators + ); + builder.build(segmentStarTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); + + List resultStarTreeDocuments = builder.getStarTreeDocuments(); + assertEquals(7, resultStarTreeDocuments.size()); + + Iterator expectedStarTreeDocumentIterator = getExpectedStarTreeDocumentIterator().iterator(); + assertStarTreeDocuments(resultStarTreeDocuments, expectedStarTreeDocumentIterator); + builder.close(); + + // building another tree in the same file + fields = List.of("fieldC", "fieldB", "fieldL", "fieldI"); + + dimensionsOrder = List.of(new NumericDimension("fieldC"), new NumericDimension("fieldB"), new NumericDimension("fieldL")); + metrics = List.of(new Metric("fieldI", List.of(MetricStat.SUM))); + + DocValuesProducer docValuesProducer = mock(DocValuesProducer.class); + + compositeField = new StarTreeField("test", dimensionsOrder, metrics, new StarTreeFieldConfiguration(1, Set.of(), getBuildMode())); + SegmentInfo segmentInfo = new SegmentInfo( + directory, + Version.LATEST, + Version.LUCENE_9_11_0, + "test_segment", + 7, + false, + false, + new Lucene912Codec(), + new HashMap<>(), + UUID.randomUUID().toString().substring(0, 16).getBytes(StandardCharsets.UTF_8), + new HashMap<>(), + null + ); + + fieldsInfo = new FieldInfo[fields.size()]; + fieldProducerMap = new HashMap<>(); + for (int i = 0; i < fieldsInfo.length; i++) { + fieldsInfo[i] = new FieldInfo( + fields.get(i), + i, + false, + false, + true, + IndexOptions.DOCS_AND_FREQS_AND_POSITIONS_AND_OFFSETS, + DocValuesType.SORTED_NUMERIC, + -1, + Collections.emptyMap(), + 0, + 0, + 0, + 0, + VectorEncoding.FLOAT32, + VectorSimilarityFunction.EUCLIDEAN, + false, + false + ); + fieldProducerMap.put(fields.get(i), docValuesProducer); + } + FieldInfos fieldInfos = new FieldInfos(fieldsInfo); + writeState = new SegmentWriteState(InfoStream.getDefault(), segmentInfo.dir, segmentInfo, fieldInfos, null, newIOContext(random())); + + mapperService = mock(MapperService.class); + 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("fieldI", NumberFieldMapper.NumberType.DOUBLE, false, true) + .build(new Mapper.BuilderContext(settings, new ContentPath())); + MappingLookup fieldMappers = new MappingLookup( + Set.of(numberFieldMapper1), + Collections.emptyList(), + Collections.emptyList(), + 0, + null + ); + when(documentMapper.mappers()).thenReturn(fieldMappers); + + InMemoryTreeNode rootNode1 = builder.getRootNode(); + + int noOfStarTreeDocuments2 = 7; + StarTreeDocument[] starTreeDocuments2 = new StarTreeDocument[noOfStarTreeDocuments2]; + starTreeDocuments2[0] = new StarTreeDocument(new Long[] { 1L, 11L, 21L }, new Double[] { 400.0 }); + starTreeDocuments2[1] = new StarTreeDocument(new Long[] { 1L, 12L, 22L }, new Double[] { 200.0 }); + starTreeDocuments2[2] = new StarTreeDocument(new Long[] { 2L, 13L, 23L }, new Double[] { 300.0 }); + starTreeDocuments2[3] = new StarTreeDocument(new Long[] { 2L, 13L, 21L }, new Double[] { 100.0 }); + starTreeDocuments2[4] = new StarTreeDocument(new Long[] { 3L, 11L, 21L }, new Double[] { 600.0 }); + starTreeDocuments2[5] = new StarTreeDocument(new Long[] { 3L, 12L, 23L }, new Double[] { 200.0 }); + starTreeDocuments2[6] = new StarTreeDocument(new Long[] { 3L, 12L, 21L }, new Double[] { 400.0 }); + + StarTreeDocument[] segmentStarTreeDocuments2 = new StarTreeDocument[noOfStarTreeDocuments2]; + for (int i = 0; i < noOfStarTreeDocuments2; i++) { + long metric1 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments2[i].metrics[0]); + segmentStarTreeDocuments2[i] = new StarTreeDocument(starTreeDocuments2[i].dimensions, new Long[] { metric1 }); + } + + SequentialDocValuesIterator[] dimsIterators2 = getDimensionIterators(segmentStarTreeDocuments2); + List metricsIterators2 = getMetricIterators(segmentStarTreeDocuments2); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + Iterator segmentStarTreeDocumentIterator2 = builder.sortAndAggregateSegmentDocuments( + dimsIterators2, + metricsIterators2 + ); + builder.build(segmentStarTreeDocumentIterator2, new AtomicInteger(), mock(DocValuesConsumer.class)); + InMemoryTreeNode rootNode2 = builder.getRootNode(); + + metaOut.close(); + dataOut.close(); + + StarTreeMetadata starTreeMetadata = new StarTreeMetadata( + "test", + STAR_TREE, + mock(IndexInput.class), + VERSION_CURRENT, + builder.numStarTreeNodes, + List.of("field1", "field3", "field5", "field8"), + List.of( + new Metric("field2", List.of(MetricStat.SUM)), + new Metric("field4", List.of(MetricStat.SUM)), + new Metric("field6", List.of(MetricStat.VALUE_COUNT)), + new Metric("field9", List.of(MetricStat.MIN)), + new Metric("field10", List.of(MetricStat.MAX)) + ), + 2, + getExpectedStarTreeDocumentIterator().size(), + 1, + Set.of("field8"), + getBuildMode(), + 0, + 330 + ); + + StarTreeMetadata starTreeMetadata2 = new StarTreeMetadata( + "test", + STAR_TREE, + mock(IndexInput.class), + VERSION_CURRENT, + builder.numStarTreeNodes, + List.of("fieldC", "fieldB", "fieldL"), + List.of(new Metric("fieldI", List.of(MetricStat.SUM))), + 7, + 27, + 1, + Set.of(), + getBuildMode(), + 330, + 1287 + ); + + List totalDimensionFields = new ArrayList<>(); + totalDimensionFields.addAll(starTreeMetadata.getDimensionFields()); + totalDimensionFields.addAll(starTreeMetadata2.getDimensionFields()); + + List metrics = new ArrayList<>(); + metrics.addAll(starTreeMetadata.getMetrics()); + metrics.addAll(starTreeMetadata2.getMetrics()); + + SegmentReadState readState = getReadState(3, totalDimensionFields, metrics); + + IndexInput dataIn = readState.directory.openInput(dataFileName, IOContext.DEFAULT); + IndexInput metaIn = readState.directory.openInput(metaFileName, IOContext.DEFAULT); + + validateFileFormats(dataIn, metaIn, rootNode1, starTreeMetadata); + validateFileFormats(dataIn, metaIn, rootNode2, starTreeMetadata2); + + dataIn.close(); + metaIn.close(); + + } + + public void test_build() throws IOException { + + int noOfStarTreeDocuments = 5; + StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; + + starTreeDocuments[0] = new StarTreeDocument( + new Long[] { 2L, 4L, 3L, 4L }, + new Object[] { 12.0, 10.0, randomDouble(), 8.0, 20.0, 1L } + ); + starTreeDocuments[1] = new StarTreeDocument( + new Long[] { 3L, 4L, 2L, 1L }, + new Object[] { 10.0, 6.0, randomDouble(), 12.0, 10.0, null } + ); + starTreeDocuments[2] = new StarTreeDocument( + new Long[] { 3L, 4L, 2L, 1L }, + new Object[] { 14.0, 12.0, randomDouble(), 6.0, 24.0, null } + ); + starTreeDocuments[3] = new StarTreeDocument( + new Long[] { 2L, 4L, 3L, 4L }, + new Object[] { 9.0, 4.0, randomDouble(), 9.0, 12.0, null } + ); + starTreeDocuments[4] = new StarTreeDocument( + new Long[] { 3L, 4L, 2L, 1L }, + new Object[] { 11.0, 16.0, randomDouble(), 8.0, 13.0, null } + ); + + StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; + for (int i = 0; i < noOfStarTreeDocuments; i++) { + long metric1 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]); + long metric2 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[1]); + long metric3 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]); + long metric4 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[3]); + long metric5 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[4]); + Long metric6 = (Long) starTreeDocuments[i].metrics[5]; + segmentStarTreeDocuments[i] = new StarTreeDocument( + starTreeDocuments[i].dimensions, + new Long[] { metric1, metric2, metric3, metric4, metric5, metric6 } + ); + } + + SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); + List metricsIterators = getMetricIterators(segmentStarTreeDocuments); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( + dimsIterators, + metricsIterators + ); + docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( + writeState, + Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, + Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, + Composite99DocValuesFormat.META_DOC_VALUES_CODEC, + Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION + ); + builder.build(segmentStarTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); + + List resultStarTreeDocuments = builder.getStarTreeDocuments(); + assertEquals(7, resultStarTreeDocuments.size()); + + Iterator expectedStarTreeDocumentIterator = getExpectedStarTreeDocumentIterator().iterator(); + assertStarTreeDocuments(resultStarTreeDocuments, expectedStarTreeDocumentIterator); + + metaOut.close(); + dataOut.close(); + docValuesConsumer.close(); + + StarTreeMetadata starTreeMetadata = new StarTreeMetadata( + "test", + STAR_TREE, + mock(IndexInput.class), + VERSION_CURRENT, + builder.numStarTreeNodes, + getStarTreeDimensionNames(compositeField.getDimensionsOrder()), + compositeField.getMetrics(), + 2, + getExpectedStarTreeDocumentIterator().size(), + 1, + Set.of("field8"), + getBuildMode(), + 0, + 330 + ); + + validateStarTreeFileFormats( + builder.getRootNode(), + getExpectedStarTreeDocumentIterator().size(), + starTreeMetadata, + getExpectedStarTreeDocumentIterator() + ); + } + + public void test_build_starTreeDataset() throws IOException { + + fields = List.of("fieldC", "fieldB", "fieldL", "fieldI"); + + dimensionsOrder = List.of(new NumericDimension("fieldC"), new NumericDimension("fieldB"), new NumericDimension("fieldL")); + metrics = List.of(new Metric("fieldI", List.of(MetricStat.SUM)), new Metric("_doc_count", List.of(MetricStat.DOC_COUNT))); + + DocValuesProducer docValuesProducer = mock(DocValuesProducer.class); + + compositeField = new StarTreeField("test", dimensionsOrder, metrics, new StarTreeFieldConfiguration(1, Set.of(), getBuildMode())); + SegmentInfo segmentInfo = new SegmentInfo( + directory, + Version.LATEST, + Version.LUCENE_9_11_0, + "test_segment", + 7, + false, + false, + new Lucene912Codec(), + new HashMap<>(), + UUID.randomUUID().toString().substring(0, 16).getBytes(StandardCharsets.UTF_8), + new HashMap<>(), + null + ); + + fieldsInfo = new FieldInfo[fields.size()]; + fieldProducerMap = new HashMap<>(); + for (int i = 0; i < fieldsInfo.length; i++) { + fieldsInfo[i] = new FieldInfo( + fields.get(i), + i, + false, + false, + true, + IndexOptions.DOCS_AND_FREQS_AND_POSITIONS_AND_OFFSETS, + DocValuesType.SORTED_NUMERIC, + -1, + Collections.emptyMap(), + 0, + 0, + 0, + 0, + VectorEncoding.FLOAT32, + VectorSimilarityFunction.EUCLIDEAN, + false, + false + ); + fieldProducerMap.put(fields.get(i), docValuesProducer); + } + FieldInfos fieldInfos = new FieldInfos(fieldsInfo); + writeState = new SegmentWriteState(InfoStream.getDefault(), segmentInfo.dir, segmentInfo, fieldInfos, null, newIOContext(random())); + this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( + writeState, + Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, + Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, + Composite99DocValuesFormat.META_DOC_VALUES_CODEC, + Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION + ); + mapperService = mock(MapperService.class); + 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("fieldI", NumberFieldMapper.NumberType.DOUBLE, false, true) + .build(new Mapper.BuilderContext(settings, new ContentPath())); + MappingLookup fieldMappers = new MappingLookup( + Set.of(numberFieldMapper1), + Collections.emptyList(), + Collections.emptyList(), + 0, + null + ); + when(documentMapper.mappers()).thenReturn(fieldMappers); + + int noOfStarTreeDocuments = 7; + StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; + starTreeDocuments[0] = new StarTreeDocument(new Long[] { 1L, 11L, 21L }, new Object[] { 400.0, null }); + starTreeDocuments[1] = new StarTreeDocument(new Long[] { 1L, 12L, 22L }, new Object[] { 200.0, null }); + starTreeDocuments[2] = new StarTreeDocument(new Long[] { 2L, 13L, 23L }, new Object[] { 300.0, null }); + starTreeDocuments[3] = new StarTreeDocument(new Long[] { 2L, 13L, 21L }, new Object[] { 100.0, null }); + starTreeDocuments[4] = new StarTreeDocument(new Long[] { 3L, 11L, 21L }, new Object[] { 600.0, null }); + starTreeDocuments[5] = new StarTreeDocument(new Long[] { 3L, 12L, 23L }, new Object[] { 200.0, null }); + starTreeDocuments[6] = new StarTreeDocument(new Long[] { 3L, 12L, 21L }, new Object[] { 400.0, null }); + + StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; + for (int i = 0; i < noOfStarTreeDocuments; i++) { + long metric1 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]); + segmentStarTreeDocuments[i] = new StarTreeDocument(starTreeDocuments[i].dimensions, new Long[] { metric1, null }); + } + + SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); + List metricsIterators = getMetricIterators(segmentStarTreeDocuments); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( + dimsIterators, + metricsIterators + ); + builder.build(segmentStarTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); + + List resultStarTreeDocuments = builder.getStarTreeDocuments(); + Iterator expectedStarTreeDocumentIterator = expectedStarTreeDocuments().iterator(); + Iterator resultStarTreeDocumentIterator = resultStarTreeDocuments.iterator(); + Map> dimValueToDocIdMap = new HashMap<>(); + builder.rootNode.setNodeType(StarTreeNodeType.STAR.getValue()); + traverseStarTree(builder.rootNode, dimValueToDocIdMap, true); + + Map> expectedDimToValueMap = getExpectedDimToValueMap(); + for (Map.Entry> entry : dimValueToDocIdMap.entrySet()) { + int dimId = entry.getKey(); + if (dimId == -1) continue; + Map map = expectedDimToValueMap.get(dimId); + for (Map.Entry dimValueToDocIdEntry : entry.getValue().entrySet()) { + long dimValue = dimValueToDocIdEntry.getKey(); + int docId = dimValueToDocIdEntry.getValue(); + if (map.get(dimValue) != null) { + assertEquals(map.get(dimValue), resultStarTreeDocuments.get(docId).metrics[0]); + } + } + } + + while (resultStarTreeDocumentIterator.hasNext() && expectedStarTreeDocumentIterator.hasNext()) { + StarTreeDocument resultStarTreeDocument = resultStarTreeDocumentIterator.next(); + StarTreeDocument expectedStarTreeDocument = expectedStarTreeDocumentIterator.next(); + assertEquals(expectedStarTreeDocument.dimensions[0], resultStarTreeDocument.dimensions[0]); + assertEquals(expectedStarTreeDocument.dimensions[1], resultStarTreeDocument.dimensions[1]); + assertEquals(expectedStarTreeDocument.dimensions[2], resultStarTreeDocument.dimensions[2]); + assertEquals(expectedStarTreeDocument.metrics[0], resultStarTreeDocument.metrics[0]); + assertEquals(expectedStarTreeDocument.metrics[1], resultStarTreeDocument.metrics[1]); + } + + metaOut.close(); + dataOut.close(); + docValuesConsumer.close(); + + validateStarTree(builder.getRootNode(), 3, 1, builder.getStarTreeDocuments()); + + StarTreeMetadata starTreeMetadata = new StarTreeMetadata( + "test", + STAR_TREE, + mock(IndexInput.class), + VERSION_CURRENT, + builder.numStarTreeNodes, + getStarTreeDimensionNames(compositeField.getDimensionsOrder()), + compositeField.getMetrics(), + 7, + 27, + 1, + Set.of(), + getBuildMode(), + 0, + 1287 + ); + validateStarTreeFileFormats(builder.getRootNode(), 27, starTreeMetadata, expectedStarTreeDocuments()); + } + + private static Map> getExpectedDimToValueMap() { + Map> expectedDimToValueMap = new HashMap<>(); + Map dimValueMap = new HashMap<>(); + dimValueMap.put(1L, 600.0); + dimValueMap.put(2L, 400.0); + dimValueMap.put(3L, 1200.0); + expectedDimToValueMap.put(0, dimValueMap); + + dimValueMap = new HashMap<>(); + dimValueMap.put(11L, 1000.0); + dimValueMap.put(12L, 800.0); + dimValueMap.put(13L, 400.0); + expectedDimToValueMap.put(1, dimValueMap); + + dimValueMap = new HashMap<>(); + dimValueMap.put(21L, 1500.0); + dimValueMap.put(22L, 200.0); + dimValueMap.put(23L, 500.0); + expectedDimToValueMap.put(2, dimValueMap); + return expectedDimToValueMap; + } + + private List expectedStarTreeDocuments() { + return List.of( + new StarTreeDocument(new Long[] { 1L, 11L, 21L }, new Object[] { 400.0, 1L }), + new StarTreeDocument(new Long[] { 1L, 12L, 22L }, new Object[] { 200.0, 1L }), + new StarTreeDocument(new Long[] { 2L, 13L, 21L }, new Object[] { 100.0, 1L }), + new StarTreeDocument(new Long[] { 2L, 13L, 23L }, new Object[] { 300.0, 1L }), + new StarTreeDocument(new Long[] { 3L, 11L, 21L }, new Object[] { 600.0, 1L }), + new StarTreeDocument(new Long[] { 3L, 12L, 21L }, new Object[] { 400.0, 1L }), + new StarTreeDocument(new Long[] { 3L, 12L, 23L }, new Object[] { 200.0, 1L }), + new StarTreeDocument(new Long[] { null, 11L, 21L }, new Object[] { 1000.0, 2L }), + new StarTreeDocument(new Long[] { null, 12L, 21L }, new Object[] { 400.0, 1L }), + new StarTreeDocument(new Long[] { null, 12L, 22L }, new Object[] { 200.0, 1L }), + new StarTreeDocument(new Long[] { null, 12L, 23L }, new Object[] { 200.0, 1L }), + new StarTreeDocument(new Long[] { null, 13L, 21L }, new Object[] { 100.0, 1L }), + new StarTreeDocument(new Long[] { null, 13L, 23L }, new Object[] { 300.0, 1L }), + new StarTreeDocument(new Long[] { null, null, 21L }, new Object[] { 1500.0, 4L }), + new StarTreeDocument(new Long[] { null, null, 22L }, new Object[] { 200.0, 1L }), + new StarTreeDocument(new Long[] { null, null, 23L }, new Object[] { 500.0, 2L }), + new StarTreeDocument(new Long[] { null, null, null }, new Object[] { 2200.0, 7L }), + new StarTreeDocument(new Long[] { null, 12L, null }, new Object[] { 800.0, 3L }), + new StarTreeDocument(new Long[] { null, 13L, null }, new Object[] { 400.0, 2L }), + new StarTreeDocument(new Long[] { 1L, null, 21L }, new Object[] { 400.0, 1L }), + new StarTreeDocument(new Long[] { 1L, null, 22L }, new Object[] { 200.0, 1L }), + new StarTreeDocument(new Long[] { 1L, null, null }, new Object[] { 600.0, 2L }), + new StarTreeDocument(new Long[] { 2L, 13L, null }, new Object[] { 400.0, 2L }), + new StarTreeDocument(new Long[] { 3L, null, 21L }, new Object[] { 1000.0, 2L }), + new StarTreeDocument(new Long[] { 3L, null, 23L }, new Object[] { 200.0, 1L }), + new StarTreeDocument(new Long[] { 3L, null, null }, new Object[] { 1200.0, 3L }), + new StarTreeDocument(new Long[] { 3L, 12L, null }, new Object[] { 600.0, 2L }) + ); + } + +} diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuilderFlushFlowTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuilderFlushFlowTests.java new file mode 100644 index 0000000000000..1aa830e3587df --- /dev/null +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuilderFlushFlowTests.java @@ -0,0 +1,419 @@ +/* + * 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.startree.builder; + +import org.apache.lucene.codecs.DocValuesProducer; +import org.apache.lucene.index.EmptyDocValuesProducer; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.index.SortedNumericDocValues; +import org.apache.lucene.search.DocIdSetIterator; +import org.apache.lucene.store.IndexInput; +import org.opensearch.index.codec.composite.LuceneDocValuesConsumerFactory; +import org.opensearch.index.codec.composite.composite99.Composite99DocValuesFormat; +import org.opensearch.index.compositeindex.datacube.Dimension; +import org.opensearch.index.compositeindex.datacube.Metric; +import org.opensearch.index.compositeindex.datacube.MetricStat; +import org.opensearch.index.compositeindex.datacube.NumericDimension; +import org.opensearch.index.compositeindex.datacube.startree.StarTreeDocument; +import org.opensearch.index.compositeindex.datacube.startree.StarTreeField; +import org.opensearch.index.compositeindex.datacube.startree.StarTreeFieldConfiguration; +import org.opensearch.index.compositeindex.datacube.startree.fileformats.meta.StarTreeMetadata; +import org.opensearch.index.compositeindex.datacube.startree.utils.SequentialDocValuesIterator; +import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.SortedNumericStarTreeValuesIterator; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.opensearch.index.compositeindex.datacube.startree.builder.BuilderTestsUtils.getSortedNumericMock; +import static org.opensearch.index.compositeindex.datacube.startree.builder.BuilderTestsUtils.validateStarTree; +import static org.opensearch.index.compositeindex.datacube.startree.fileformats.StarTreeWriter.VERSION_CURRENT; +import static org.opensearch.index.mapper.CompositeMappedFieldType.CompositeFieldType.STAR_TREE; +import static org.mockito.Mockito.mock; + +public class StarTreeBuilderFlushFlowTests extends StarTreeBuilderTestCase { + + public StarTreeBuilderFlushFlowTests(StarTreeFieldConfiguration.StarTreeBuildMode buildMode) { + super(buildMode); + } + + public void testFlushFlow() throws IOException { + List dimList = List.of(0L, 1L, 3L, 4L, 5L); + List docsWithField = List.of(0, 1, 3, 4, 5); + List dimList2 = List.of(0L, 1L, 2L, 3L, 4L, 5L); + List docsWithField2 = List.of(0, 1, 2, 3, 4, 5); + + List metricsList = List.of( + getLongFromDouble(0.0), + getLongFromDouble(10.0), + getLongFromDouble(20.0), + getLongFromDouble(30.0), + getLongFromDouble(40.0), + getLongFromDouble(50.0) + ); + List metricsWithField = List.of(0, 1, 2, 3, 4, 5); + + compositeField = getStarTreeFieldWithMultipleMetrics(); + 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) + ); + + writeState = getWriteState(6, writeState.segmentInfo.getId()); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + SequentialDocValuesIterator[] dimDvs = { new SequentialDocValuesIterator(d1sndv), new SequentialDocValuesIterator(d2sndv) }; + Iterator starTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( + dimDvs, + List.of(new SequentialDocValuesIterator(m1sndv), new SequentialDocValuesIterator(m2sndv)) + ); + /** + * Asserting following dim / metrics [ dim1, dim2 / Sum [metric], count [metric] ] + [0, 0] | [0.0, 1] + [1, 1] | [10.0, 1] + [3, 3] | [30.0, 1] + [4, 4] | [40.0, 1] + [5, 5] | [50.0, 1] + [null, 2] | [20.0, 1] + */ + this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( + writeState, + Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, + Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, + Composite99DocValuesFormat.META_DOC_VALUES_CODEC, + Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION + ); + builder.build(starTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); + List starTreeDocuments = builder.getStarTreeDocuments(); + int count = 0; + for (StarTreeDocument starTreeDocument : starTreeDocuments) { + count++; + if (starTreeDocument.dimensions[1] != null) { + assertEquals( + starTreeDocument.dimensions[0] == null + ? starTreeDocument.dimensions[1] * 1 * 10.0 + : starTreeDocument.dimensions[0] * 10, + starTreeDocument.metrics[0] + ); + assertEquals(1L, starTreeDocument.metrics[1]); + } else { + assertEquals(150D, starTreeDocument.metrics[0]); + assertEquals(6L, starTreeDocument.metrics[1]); + } + } + assertEquals(13, count); + validateStarTree(builder.getRootNode(), 2, 1000, builder.getStarTreeDocuments()); + + metaOut.close(); + dataOut.close(); + docValuesConsumer.close(); + + StarTreeMetadata starTreeMetadata = new StarTreeMetadata( + "sf", + STAR_TREE, + mock(IndexInput.class), + VERSION_CURRENT, + builder.numStarTreeNodes, + List.of("field1", "field3"), + List.of(new Metric("field2", List.of(MetricStat.SUM, MetricStat.VALUE_COUNT, MetricStat.AVG))), + 6, + builder.numStarTreeDocs, + 1000, + Set.of(), + getBuildMode(), + 0, + 264 + ); + + validateStarTreeFileFormats( + builder.getRootNode(), + builder.getStarTreeDocuments().size(), + starTreeMetadata, + builder.getStarTreeDocuments() + ); + + } + + public void testFlushFlowDimsReverse() throws IOException { + List dimList = List.of(5L, 4L, 3L, 2L, 1L); + List docsWithField = List.of(0, 1, 2, 3, 4); + List dimList2 = List.of(5L, 4L, 3L, 2L, 1L, 0L); + List docsWithField2 = List.of(0, 1, 2, 3, 4, 5); + + List metricsList = List.of( + getLongFromDouble(50.0), + getLongFromDouble(40.0), + getLongFromDouble(30.0), + getLongFromDouble(20.0), + getLongFromDouble(10.0), + getLongFromDouble(0.0) + ); + List metricsWithField = List.of(0, 1, 2, 3, 4, 5); + + compositeField = getStarTreeFieldWithMultipleMetrics(); + 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) + ); + + writeState = getWriteState(6, writeState.segmentInfo.getId()); + this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( + writeState, + Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, + Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, + Composite99DocValuesFormat.META_DOC_VALUES_CODEC, + Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION + ); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + SequentialDocValuesIterator[] dimDvs = { new SequentialDocValuesIterator(d1sndv), new SequentialDocValuesIterator(d2sndv) }; + Iterator starTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( + dimDvs, + List.of(new SequentialDocValuesIterator(m1sndv), new SequentialDocValuesIterator(m2sndv)) + ); + /** + * Asserting following dim / metrics [ dim1, dim2 / Sum [metric], count [metric] ] + [1, 1] | [10.0, 1] + [2, 2] | [20.0, 1] + [3, 3] | [30.0, 1] + [4, 4] | [40.0, 1] + [5, 5] | [50.0, 1] + [null, 0] | [0.0, 1] + */ + builder.appendDocumentsToStarTree(starTreeDocumentIterator); + assertEquals(6, builder.getStarTreeDocuments().size()); + builder.build(starTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); + int count = 0; + for (StarTreeDocument starTreeDocument : builder.getStarTreeDocuments()) { + if (count <= 6) { + count++; + if (starTreeDocument.dimensions[0] != null) { + assertEquals(count, (long) starTreeDocument.dimensions[0]); + } + assertEquals(starTreeDocument.dimensions[1] * 10.0, starTreeDocument.metrics[0]); + assertEquals(1L, starTreeDocument.metrics[1]); + } + } + validateStarTree(builder.getRootNode(), 2, 1000, builder.getStarTreeDocuments()); + + metaOut.close(); + dataOut.close(); + docValuesConsumer.close(); + + StarTreeMetadata starTreeMetadata = new StarTreeMetadata( + "sf", + STAR_TREE, + mock(IndexInput.class), + VERSION_CURRENT, + builder.numStarTreeNodes, + List.of("field1", "field3"), + List.of(new Metric("field2", List.of(MetricStat.SUM, MetricStat.VALUE_COUNT, MetricStat.AVG))), + 6, + builder.numStarTreeDocs, + 1000, + Set.of(), + getBuildMode(), + 0, + 264 + ); + + validateStarTreeFileFormats( + builder.getRootNode(), + builder.getStarTreeDocuments().size(), + starTreeMetadata, + builder.getStarTreeDocuments() + ); + } + + public void testFlushFlowBuild() throws IOException { + List dimList = new ArrayList<>(100); + List docsWithField = new ArrayList<>(100); + for (int i = 0; i < 100; i++) { + dimList.add((long) i); + docsWithField.add(i); + } + + List dimList2 = new ArrayList<>(100); + List docsWithField2 = new ArrayList<>(100); + for (int i = 0; i < 100; i++) { + dimList2.add((long) i); + docsWithField2.add(i); + } + + List metricsList = new ArrayList<>(100); + List metricsWithField = new ArrayList<>(100); + for (int i = 0; i < 100; i++) { + metricsList.add(getLongFromDouble(i * 10.0)); + metricsWithField.add(i); + } + + Dimension d1 = new NumericDimension("field1"); + Dimension d2 = new NumericDimension("field3"); + Metric m1 = new Metric("field2", List.of(MetricStat.SUM)); + List dims = List.of(d1, d2); + List metrics = List.of(m1); + StarTreeFieldConfiguration c = new StarTreeFieldConfiguration(1, new HashSet<>(), getBuildMode()); + compositeField = new StarTreeField("sf", dims, metrics, c); + SortedNumericDocValues d1sndv = getSortedNumericMock(dimList, docsWithField); + SortedNumericDocValues d2sndv = getSortedNumericMock(dimList2, docsWithField2); + SortedNumericDocValues m1sndv = getSortedNumericMock(metricsList, metricsWithField); + + writeState = getWriteState(100, writeState.segmentInfo.getId()); + SegmentWriteState consumerWriteState = getWriteState(DocIdSetIterator.NO_MORE_DOCS, writeState.segmentInfo.getId()); + this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( + consumerWriteState, + Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, + Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, + Composite99DocValuesFormat.META_DOC_VALUES_CODEC, + Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION + ); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + + DocValuesProducer d1vp = getDocValuesProducer(d1sndv); + DocValuesProducer d2vp = getDocValuesProducer(d2sndv); + DocValuesProducer m1vp = getDocValuesProducer(m1sndv); + Map fieldProducerMap = Map.of("field1", d1vp, "field3", d2vp, "field2", m1vp); + builder.build(fieldProducerMap, new AtomicInteger(), docValuesConsumer); + /** + * Asserting following dim / metrics [ dim1, dim2 / Sum [ metric] ] + [0, 0] | [0.0] + [1, 1] | [10.0] + [2, 2] | [20.0] + [3, 3] | [30.0] + [4, 4] | [40.0] + .... + [null, 0] | [0.0] + [null, 1] | [10.0] + ... + [null, null] | [49500.0] + */ + List starTreeDocuments = builder.getStarTreeDocuments(); + for (StarTreeDocument starTreeDocument : starTreeDocuments) { + assertEquals( + starTreeDocument.dimensions[1] != null ? starTreeDocument.dimensions[1] * 10.0 : 49500.0, + starTreeDocument.metrics[0] + ); + } + validateStarTree(builder.getRootNode(), 2, 1, builder.getStarTreeDocuments()); + + metaOut.close(); + dataOut.close(); + docValuesConsumer.close(); + + StarTreeMetadata starTreeMetadata = getStarTreeMetadata(List.of("field1", "field3"), 100, 1, 6699); + + validateStarTreeFileFormats( + builder.getRootNode(), + builder.getStarTreeDocuments().size(), + starTreeMetadata, + builder.getStarTreeDocuments() + ); + } + + public void testFlushFlowWithTimestamps() throws IOException { + List dimList = List.of(1655288152000L, 1655288092000L, 1655288032000L, 1655287972000L, 1655288092000L); + List docsWithField = List.of(0, 1, 3, 4, 5); + List dimList2 = List.of(0L, 1L, 2L, 3L, 4L, 5L); + List docsWithField2 = List.of(0, 1, 2, 3, 4, 5); + + List metricsList = List.of( + getLongFromDouble(0.0), + getLongFromDouble(10.0), + getLongFromDouble(20.0), + getLongFromDouble(30.0), + getLongFromDouble(40.0), + getLongFromDouble(50.0) + ); + List metricsWithField = List.of(0, 1, 2, 3, 4, 5); + + 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) + ); + this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( + writeState, + Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, + Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, + Composite99DocValuesFormat.META_DOC_VALUES_CODEC, + Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION + ); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, getWriteState(6, writeState.segmentInfo.getId()), mapperService); + SequentialDocValuesIterator[] dimDvs = { new SequentialDocValuesIterator(d1sndv), new SequentialDocValuesIterator(d2sndv) }; + Iterator starTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( + dimDvs, + List.of(new SequentialDocValuesIterator(m1sndv), new SequentialDocValuesIterator(m2sndv)) + ); + /** + * Asserting following dim / metrics [ dim1, dim2 / Sum [metric], count [metric] ] + [1655287920000, 1655287200000, 1655287200000, 4] | [40.0, 1] + [1655287980000, 1655287200000, 1655287200000, 3] | [30.0, 1] + [1655288040000, 1655287200000, 1655287200000, 1] | [10.0, 1] + [1655288040000, 1655287200000, 1655287200000, 5] | [50.0, 1] + [1655288100000, 1655287200000, 1655287200000, 0] | [0.0, 1] + [null, null, null, 2] | [20.0, 1] + */ + int count = 0; + List starTreeDocumentsList = new ArrayList<>(); + starTreeDocumentIterator.forEachRemaining(starTreeDocumentsList::add); + starTreeDocumentIterator = starTreeDocumentsList.iterator(); + while (starTreeDocumentIterator.hasNext()) { + count++; + StarTreeDocument starTreeDocument = starTreeDocumentIterator.next(); + assertEquals(starTreeDocument.dimensions[3] * 1 * 10.0, starTreeDocument.metrics[1]); + assertEquals(1L, starTreeDocument.metrics[0]); + } + assertEquals(6, count); + builder.build(starTreeDocumentsList.iterator(), new AtomicInteger(), docValuesConsumer); + validateStarTree(builder.getRootNode(), 3, 10, builder.getStarTreeDocuments()); + } + + private StarTreeField getStarTreeFieldWithMultipleMetrics() { + Dimension d1 = new NumericDimension("field1"); + Dimension d2 = new NumericDimension("field3"); + Metric m1 = new Metric("field2", List.of(MetricStat.SUM)); + Metric m2 = new Metric("field2", List.of(MetricStat.VALUE_COUNT)); + Metric m3 = new Metric("field2", List.of(MetricStat.AVG)); + List dims = List.of(d1, d2); + List metrics = List.of(m1, m2, m3); + StarTreeFieldConfiguration c = new StarTreeFieldConfiguration(1000, new HashSet<>(), getBuildMode()); + return new StarTreeField("sf", dims, metrics, c); + } + + private static DocValuesProducer getDocValuesProducer(SortedNumericDocValues sndv) { + return new EmptyDocValuesProducer() { + @Override + public SortedNumericDocValues getSortedNumeric(FieldInfo field) throws IOException { + return sndv; + } + }; + } +} diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuilderMergeFlowTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuilderMergeFlowTests.java new file mode 100644 index 0000000000000..f983365dfec30 --- /dev/null +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuilderMergeFlowTests.java @@ -0,0 +1,1921 @@ +/* + * 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.startree.builder; + +import org.apache.lucene.index.DocValues; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.index.SortedNumericDocValues; +import org.apache.lucene.search.DocIdSetIterator; +import org.opensearch.common.settings.Settings; +import org.opensearch.index.codec.composite.LuceneDocValuesConsumerFactory; +import org.opensearch.index.codec.composite.composite99.Composite99DocValuesFormat; +import org.opensearch.index.compositeindex.datacube.Dimension; +import org.opensearch.index.compositeindex.datacube.Metric; +import org.opensearch.index.compositeindex.datacube.MetricStat; +import org.opensearch.index.compositeindex.datacube.NumericDimension; +import org.opensearch.index.compositeindex.datacube.startree.StarTreeDocument; +import org.opensearch.index.compositeindex.datacube.startree.StarTreeField; +import org.opensearch.index.compositeindex.datacube.startree.StarTreeFieldConfiguration; +import org.opensearch.index.compositeindex.datacube.startree.fileformats.meta.StarTreeMetadata; +import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues; +import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.SortedNumericStarTreeValuesIterator; +import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.StarTreeValuesIterator; +import org.opensearch.index.mapper.ContentPath; +import org.opensearch.index.mapper.DocumentMapper; +import org.opensearch.index.mapper.Mapper; +import org.opensearch.index.mapper.MappingLookup; +import org.opensearch.index.mapper.NumberFieldMapper; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Iterator; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Set; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Supplier; + +import static org.opensearch.index.compositeindex.CompositeIndexConstants.SEGMENT_DOCS_COUNT; +import static org.opensearch.index.compositeindex.datacube.startree.builder.BuilderTestsUtils.getSortedNumericMock; +import static org.opensearch.index.compositeindex.datacube.startree.builder.BuilderTestsUtils.traverseStarTree; +import static org.opensearch.index.compositeindex.datacube.startree.builder.BuilderTestsUtils.validateStarTree; +import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class StarTreeBuilderMergeFlowTests extends StarTreeBuilderTestCase { + + public StarTreeBuilderMergeFlowTests(StarTreeFieldConfiguration.StarTreeBuildMode buildMode) { + super(buildMode); + } + + public void testMergeFlow() throws IOException { + List dimList1 = new ArrayList<>(1000); + List docsWithField1 = new ArrayList<>(1000); + for (int i = 0; i < 1000; i++) { + dimList1.add((long) i); + docsWithField1.add(i); + } + + List dimList2 = new ArrayList<>(1000); + List docsWithField2 = new ArrayList<>(1000); + for (int i = 0; i < 1000; i++) { + dimList2.add((long) i); + docsWithField2.add(i); + } + + List dimList3 = new ArrayList<>(1000); + List docsWithField3 = new ArrayList<>(1000); + for (int i = 0; i < 1000; i++) { + dimList3.add((long) i); + docsWithField3.add(i); + } + + List dimList4 = new ArrayList<>(1000); + List docsWithField4 = new ArrayList<>(1000); + for (int i = 0; i < 1000; i++) { + dimList4.add((long) i); + docsWithField4.add(i); + } + + List dimList5 = new ArrayList<>(1000); + List docsWithField5 = new ArrayList<>(1000); + for (int i = 0; i < 1000; i++) { + dimList5.add((long) i); + docsWithField5.add(i); + } + + List metricsList = new ArrayList<>(1000); + List metricsWithField = new ArrayList<>(1000); + for (int i = 0; i < 1000; i++) { + metricsList.add(getLongFromDouble(i * 10.0)); + metricsWithField.add(i); + } + + List metricsListValueCount = new ArrayList<>(1000); + List metricsWithFieldValueCount = new ArrayList<>(1000); + for (int i = 0; i < 1000; i++) { + metricsListValueCount.add((long) i); + metricsWithFieldValueCount.add(i); + } + + Dimension d1 = new NumericDimension("field1"); + Dimension d2 = new NumericDimension("field3"); + Dimension d3 = new NumericDimension("field5"); + Dimension d4 = new NumericDimension("field8"); + // Dimension d5 = new NumericDimension("field5"); + Metric m1 = new Metric("field2", List.of(MetricStat.SUM, MetricStat.AVG, MetricStat.VALUE_COUNT)); + Metric m2 = new Metric("_doc_count", List.of(MetricStat.DOC_COUNT)); + List dims = List.of(d1, d2, d3, d4); + List metrics = List.of(m1, m2); + StarTreeFieldConfiguration c = new StarTreeFieldConfiguration(1, new HashSet<>(), getBuildMode()); + compositeField = new StarTreeField("sf", dims, metrics, c); + SortedNumericDocValues d1sndv = getSortedNumericMock(dimList1, docsWithField1); + SortedNumericDocValues d2sndv = getSortedNumericMock(dimList2, docsWithField2); + SortedNumericDocValues d3sndv = getSortedNumericMock(dimList3, docsWithField3); + SortedNumericDocValues d4sndv = getSortedNumericMock(dimList4, docsWithField4); + SortedNumericDocValues m1sndv = getSortedNumericMock(metricsList, metricsWithField); + SortedNumericDocValues valucountsndv = getSortedNumericMock(metricsListValueCount, metricsWithFieldValueCount); + SortedNumericDocValues m2sndv = DocValues.emptySortedNumeric(); + Map> dimDocIdSetIterators = Map.of( + "field1", + () -> new SortedNumericStarTreeValuesIterator(d1sndv), + "field3", + () -> new SortedNumericStarTreeValuesIterator(d2sndv), + "field5", + () -> new SortedNumericStarTreeValuesIterator(d3sndv), + "field8", + () -> new SortedNumericStarTreeValuesIterator(d4sndv) + ); + + Map> metricDocIdSetIterators = Map.of( + "sf_field2_sum_metric", + () -> new SortedNumericStarTreeValuesIterator(m1sndv), + "sf_field2_value_count_metric", + () -> new SortedNumericStarTreeValuesIterator(valucountsndv), + "sf__doc_count_doc_count_metric", + () -> new SortedNumericStarTreeValuesIterator(m2sndv) + ); + + StarTreeValues starTreeValues = new StarTreeValues( + compositeField, + null, + dimDocIdSetIterators, + metricDocIdSetIterators, + getAttributes(1000), + null + ); + + SortedNumericDocValues f2d1sndv = getSortedNumericMock(dimList1, docsWithField1); + SortedNumericDocValues f2d2sndv = getSortedNumericMock(dimList2, docsWithField2); + SortedNumericDocValues f2d3sndv = getSortedNumericMock(dimList3, docsWithField3); + SortedNumericDocValues f2d4sndv = getSortedNumericMock(dimList4, docsWithField4); + SortedNumericDocValues f2m1sndv = getSortedNumericMock(metricsList, metricsWithField); + SortedNumericDocValues f2valucountsndv = getSortedNumericMock(metricsListValueCount, metricsWithFieldValueCount); + SortedNumericDocValues f2m2sndv = DocValues.emptySortedNumeric(); + Map> f2dimDocIdSetIterators = Map.of( + "field1", + () -> new SortedNumericStarTreeValuesIterator(f2d1sndv), + "field3", + () -> new SortedNumericStarTreeValuesIterator(f2d2sndv), + "field5", + () -> new SortedNumericStarTreeValuesIterator(f2d3sndv), + "field8", + () -> new SortedNumericStarTreeValuesIterator(f2d4sndv) + ); + + Map> f2metricDocIdSetIterators = Map.of( + "sf_field2_sum_metric", + () -> new SortedNumericStarTreeValuesIterator(f2m1sndv), + "sf_field2_value_count_metric", + () -> new SortedNumericStarTreeValuesIterator(f2valucountsndv), + "sf__doc_count_doc_count_metric", + () -> new SortedNumericStarTreeValuesIterator(f2m2sndv) + ); + StarTreeValues starTreeValues2 = new StarTreeValues( + compositeField, + null, + f2dimDocIdSetIterators, + f2metricDocIdSetIterators, + getAttributes(1000), + null + ); + + this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( + writeState, + Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, + Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, + Composite99DocValuesFormat.META_DOC_VALUES_CODEC, + Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION + ); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); + /** + [0, 0, 0, 0] | [0.0, 2] + [1, 1, 1, 1] | [20.0, 2] + [2, 2, 2, 2] | [40.0, 2] + [3, 3, 3, 3] | [60.0, 2] + [4, 4, 4, 4] | [80.0, 2] + [5, 5, 5, 5] | [100.0, 2] + ... + [999, 999, 999, 999] | [19980.0] + */ + for (StarTreeDocument starTreeDocument : builder.getStarTreeDocuments()) { + assertEquals(starTreeDocument.dimensions[0] * 20.0, starTreeDocument.metrics[0]); + assertEquals(2L, starTreeDocument.metrics[1]); + } + builder.build(starTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); + + // Validate the star tree structure + validateStarTree(builder.getRootNode(), 4, 1, builder.getStarTreeDocuments()); + + metaOut.close(); + dataOut.close(); + docValuesConsumer.close(); + + StarTreeMetadata starTreeMetadata = getStarTreeMetadata( + getStarTreeDimensionNames(compositeField.getDimensionsOrder()), + 1000, + compositeField.getStarTreeConfig().maxLeafDocs(), + 132165 + ); + + validateStarTreeFileFormats( + builder.getRootNode(), + builder.getStarTreeDocuments().size(), + starTreeMetadata, + builder.getStarTreeDocuments() + ); + } + + 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(); + } + + public void testMergeFlowWithSum() throws IOException { + List dimList = List.of(0L, 1L, 3L, 4L, 5L, 6L); + List docsWithField = List.of(0, 1, 3, 4, 5, 6); + List dimList2 = List.of(0L, 1L, 2L, 3L, 4L, 5L, -1L); + List docsWithField2 = List.of(0, 1, 2, 3, 4, 5, 6); + + List metricsList = List.of( + getLongFromDouble(0.0), + getLongFromDouble(10.0), + getLongFromDouble(20.0), + getLongFromDouble(30.0), + getLongFromDouble(40.0), + getLongFromDouble(50.0), + getLongFromDouble(60.0) + + ); + List metricsWithField = List.of(0, 1, 2, 3, 4, 5, 6); + + compositeField = getStarTreeField(MetricStat.SUM); + StarTreeValues starTreeValues = getStarTreeValues( + getSortedNumericMock(dimList, docsWithField), + getSortedNumericMock(dimList2, docsWithField2), + getSortedNumericMock(metricsList, metricsWithField), + compositeField, + "6" + ); + + StarTreeValues starTreeValues2 = getStarTreeValues( + getSortedNumericMock(dimList, docsWithField), + getSortedNumericMock(dimList2, docsWithField2), + getSortedNumericMock(metricsList, metricsWithField), + compositeField, + "6" + ); + writeState = getWriteState(6, writeState.segmentInfo.getId()); + this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( + writeState, + Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, + Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, + Composite99DocValuesFormat.META_DOC_VALUES_CODEC, + Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION + ); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); + /** + * Asserting following dim / metrics [ dim1, dim2 / Sum [ metric] ] + * [0, 0] | [0.0] + * [1, 1] | [20.0] + * [3, 3] | [60.0] + * [4, 4] | [80.0] + * [5, 5] | [100.0] + * [null, 2] | [40.0] + * ------------------ We only take non star docs + * [6,-1] | [120.0] + */ + builder.appendDocumentsToStarTree(starTreeDocumentIterator); + assertEquals(6, builder.getStarTreeDocuments().size()); + builder.build(starTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); + int count = 0; + for (StarTreeDocument starTreeDocument : builder.getStarTreeDocuments()) { + count++; + if (count <= 6) { + assertEquals( + starTreeDocument.dimensions[0] != null ? starTreeDocument.dimensions[0] * 2 * 10.0 : 40.0, + starTreeDocument.metrics[0] + ); + } + } + + validateStarTree(builder.getRootNode(), 2, 1000, builder.getStarTreeDocuments()); + + metaOut.close(); + dataOut.close(); + docValuesConsumer.close(); + + StarTreeMetadata starTreeMetadata = getStarTreeMetadata(List.of("field1", "field3"), 6, 1000, 264); + + validateStarTreeFileFormats( + builder.getRootNode(), + builder.getStarTreeDocuments().size(), + starTreeMetadata, + builder.getStarTreeDocuments() + ); + } + + public void testMergeFlowWithCount() throws IOException { + List dimList = List.of(0L, 1L, 3L, 4L, 5L, 6L); + List docsWithField = List.of(0, 1, 3, 4, 5, 6); + List dimList2 = List.of(0L, 1L, 2L, 3L, 4L, 5L, -1L); + List docsWithField2 = List.of(0, 1, 2, 3, 4, 5, 6); + + List metricsList = List.of(0L, 1L, 2L, 3L, 4L, 5L, 6L); + List metricsWithField = List.of(0, 1, 2, 3, 4, 5, 6); + + compositeField = getStarTreeField(MetricStat.VALUE_COUNT); + StarTreeValues starTreeValues = getStarTreeValues( + getSortedNumericMock(dimList, docsWithField), + getSortedNumericMock(dimList2, docsWithField2), + getSortedNumericMock(metricsList, metricsWithField), + compositeField, + "6" + ); + + StarTreeValues starTreeValues2 = getStarTreeValues( + getSortedNumericMock(dimList, docsWithField), + getSortedNumericMock(dimList2, docsWithField2), + getSortedNumericMock(metricsList, metricsWithField), + compositeField, + "6" + ); + writeState = getWriteState(6, writeState.segmentInfo.getId()); + this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( + writeState, + Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, + Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, + Composite99DocValuesFormat.META_DOC_VALUES_CODEC, + Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION + ); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); + /** + * Asserting following dim / metrics [ dim1, dim2 / Count [ metric] ] + [0, 0] | [0] + [1, 1] | [2] + [3, 3] | [6] + [4, 4] | [8] + [5, 5] | [10] + [null, 2] | [4] + --------------- + [6,-1] | [12] + */ + builder.appendDocumentsToStarTree(starTreeDocumentIterator); + assertEquals(6, builder.getStarTreeDocuments().size()); + builder.build(starTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); + int count = 0; + for (StarTreeDocument starTreeDocument : builder.getStarTreeDocuments()) { + count++; + if (count <= 6) { + assertEquals(starTreeDocument.dimensions[0] != null ? starTreeDocument.dimensions[0] * 2 : 4, starTreeDocument.metrics[0]); + } + } + + validateStarTree(builder.getRootNode(), 2, 1000, builder.getStarTreeDocuments()); + + metaOut.close(); + dataOut.close(); + docValuesConsumer.close(); + + StarTreeMetadata starTreeMetadata = getStarTreeMetadata(List.of("field1", "field3"), 6, 1000, 264); + + validateStarTreeFileFormats( + builder.getRootNode(), + builder.getStarTreeDocuments().size(), + starTreeMetadata, + builder.getStarTreeDocuments() + ); + + } + + public void testMergeFlowNumSegmentsDocs() throws IOException { + List dimList = List.of(0L, 1L, 2L, 3L, 4L, 5L, 6L, -1L, -1L, -1L); + List docsWithField = List.of(0, 1, 2, 3, 4, 5, 6, 7, 8, 9); + List dimList2 = List.of(0L, 1L, 2L, 3L, 4L, 5L, 6L, -1L, -1L, -1L); + List docsWithField2 = List.of(0, 1, 2, 3, 4, 5, 6, 7, 8, 9); + + List metricsList = List.of(0L, 1L, 2L, 3L, 4L, 5L, 6L, -1L, -1L, -1L); + List metricsWithField = List.of(0, 1, 2, 3, 4, 5, 6, 7, 8, 9); + + List dimList3 = List.of(5L, 6L, 7L, 8L, -1L); + List docsWithField3 = List.of(0, 1, 2, 3, 4); + List dimList4 = List.of(5L, 6L, 7L, 8L, -1L); + List docsWithField4 = List.of(0, 1, 2, 3, 4); + + List metricsList2 = List.of(5L, 6L, 7L, 8L, 9L); + List metricsWithField2 = List.of(0, 1, 2, 3, 4); + + StarTreeField sf = getStarTreeField(MetricStat.VALUE_COUNT); + StarTreeValues starTreeValues = getStarTreeValues( + getSortedNumericMock(dimList, docsWithField), + getSortedNumericMock(dimList2, docsWithField2), + getSortedNumericMock(metricsList, metricsWithField), + sf, + "6" + ); + + StarTreeValues starTreeValues2 = getStarTreeValues( + getSortedNumericMock(dimList3, docsWithField3), + getSortedNumericMock(dimList4, docsWithField4), + getSortedNumericMock(metricsList2, metricsWithField2), + sf, + "4" + ); + builder = getStarTreeBuilder(metaOut, dataOut, sf, getWriteState(4, writeState.segmentInfo.getId()), mapperService); + Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); + /** + * Asserting following dim / metrics [ dim1, dim2 / Count [ metric] ] + [0, 0] | [0] + [1, 1] | [1] + [2, 2] | [2] + [3, 3] | [3] + [4, 4] | [4] + [5, 5] | [10] + [6, 6] | [6] + [7, 7] | [7] + [8, 8] | [8] + */ + int count = 0; + while (starTreeDocumentIterator.hasNext()) { + count++; + StarTreeDocument starTreeDocument = starTreeDocumentIterator.next(); + if (Objects.equals(starTreeDocument.dimensions[0], 5L)) { + assertEquals(starTreeDocument.dimensions[0] * 2, starTreeDocument.metrics[0]); + } else { + assertEquals(starTreeDocument.dimensions[1], starTreeDocument.metrics[0]); + } + } + assertEquals(9, count); + } + + public void testMergeFlowWithMissingDocs() throws IOException { + List dimList = List.of(0L, 1L, 2L, 3L, 4L, 6L); + List docsWithField = List.of(0, 1, 2, 3, 4, 6); + List dimList2 = List.of(0L, 1L, 2L, 3L, 4L, 5L, -1L); + List docsWithField2 = List.of(0, 1, 2, 3, 4, 5, 6); + + List metricsList = List.of(0L, 1L, 2L, 3L, 4L, 5L, 6L); + List metricsWithField = List.of(0, 1, 2, 3, 4, 5, 6); + + List dimList3 = List.of(5L, 6L, 8L, -1L); + List docsWithField3 = List.of(0, 1, 3, 4); + List dimList4 = List.of(5L, 6L, 7L, 8L, -1L); + List docsWithField4 = List.of(0, 1, 2, 3, 4); + + List metricsList2 = List.of(5L, 6L, 7L, 8L, 9L); + List metricsWithField2 = List.of(0, 1, 2, 3, 4); + + compositeField = getStarTreeField(MetricStat.VALUE_COUNT); + StarTreeValues starTreeValues = getStarTreeValues( + getSortedNumericMock(dimList, docsWithField), + getSortedNumericMock(dimList2, docsWithField2), + getSortedNumericMock(metricsList, metricsWithField), + compositeField, + "6" + ); + + StarTreeValues starTreeValues2 = getStarTreeValues( + getSortedNumericMock(dimList3, docsWithField3), + getSortedNumericMock(dimList4, docsWithField4), + getSortedNumericMock(metricsList2, metricsWithField2), + compositeField, + "4" + ); + writeState = getWriteState(4, writeState.segmentInfo.getId()); + this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( + writeState, + Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, + Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, + Composite99DocValuesFormat.META_DOC_VALUES_CODEC, + Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION + ); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); + /** + * Asserting following dim / metrics [ dim1, dim2 / Count [ metric] ] + [0, 0] | [0] + [1, 1] | [1] + [2, 2] | [2] + [3, 3] | [3] + [4, 4] | [4] + [5, 5] | [5] + [6, 6] | [6] + [8, 8] | [8] + [null, 5] | [5] + [null, 7] | [7] + */ + builder.appendDocumentsToStarTree(starTreeDocumentIterator); + assertEquals(10, builder.getStarTreeDocuments().size()); + builder.build(starTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); + int count = 0; + for (StarTreeDocument starTreeDocument : builder.getStarTreeDocuments()) { + count++; + if (count <= 10) { + if (starTreeDocument.dimensions[0] == null) { + assertTrue(List.of(5L, 7L).contains(starTreeDocument.dimensions[1])); + } + assertEquals(starTreeDocument.dimensions[1], starTreeDocument.metrics[0]); + } + } + + validateStarTree(builder.getRootNode(), 2, 1000, builder.getStarTreeDocuments()); + + metaOut.close(); + dataOut.close(); + docValuesConsumer.close(); + + StarTreeMetadata starTreeMetadata = getStarTreeMetadata(List.of("field1", "field3"), 10, 1000, 363); + + validateStarTreeFileFormats( + builder.getRootNode(), + builder.getStarTreeDocuments().size(), + starTreeMetadata, + builder.getStarTreeDocuments() + ); + } + + public void testMergeFlowWithMissingDocsWithZero() throws IOException { + List dimList = List.of(0L, 0L, 0L, 0L); + List docsWithField = List.of(0, 1, 2, 6); + List dimList2 = List.of(0L, 0L, 0L, 0L); + List docsWithField2 = List.of(0, 1, 2, 6); + + List metricsList = List.of(0L, 1L, 2L, 3L, 4L, 5L, 6L); + List metricsWithField = List.of(0, 1, 2, 3, 4, 5, 6); + + List dimList3 = List.of(5L, 6L, 8L, -1L); + List docsWithField3 = List.of(0, 1, 3, 4); + List dimList4 = List.of(5L, 6L, 7L, 8L, -1L); + List docsWithField4 = List.of(0, 1, 2, 3, 4); + + List metricsList2 = List.of(5L, 6L, 7L, 8L, 9L); + List metricsWithField2 = List.of(0, 1, 2, 3, 4); + + compositeField = getStarTreeField(MetricStat.VALUE_COUNT); + StarTreeValues starTreeValues = getStarTreeValues( + getSortedNumericMock(dimList, docsWithField), + getSortedNumericMock(dimList2, docsWithField2), + getSortedNumericMock(metricsList, metricsWithField), + compositeField, + "7" + ); + + StarTreeValues starTreeValues2 = getStarTreeValues( + getSortedNumericMock(dimList3, docsWithField3), + getSortedNumericMock(dimList4, docsWithField4), + getSortedNumericMock(metricsList2, metricsWithField2), + compositeField, + "4" + ); + writeState = getWriteState(4, writeState.segmentInfo.getId()); + SegmentWriteState consumerWriteState = getWriteState(DocIdSetIterator.NO_MORE_DOCS, writeState.segmentInfo.getId()); + this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( + consumerWriteState, + Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, + Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, + Composite99DocValuesFormat.META_DOC_VALUES_CODEC, + Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION + ); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); + /** + * Asserting following dim / metrics [ dim1, dim2 / Count [ metric] ] + [0, 0] | [9] + [5, 5] | [5] + [6, 6] | [6] + [8, 8] | [8] + [null, 7] | [7] + [null, null] | [12] + */ + builder.appendDocumentsToStarTree(starTreeDocumentIterator); + assertEquals(6, builder.getStarTreeDocuments().size()); + builder.build(starTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); + int count = 0; + for (StarTreeDocument starTreeDocument : builder.getStarTreeDocuments()) { + count++; + if (count <= 6) { + if (starTreeDocument.dimensions[0] == null && starTreeDocument.dimensions[1] == null) { + assertEquals(12L, (long) starTreeDocument.metrics[0]); + } else if (starTreeDocument.dimensions[0] == null) { + assertEquals(7L, starTreeDocument.metrics[0]); + } else if (starTreeDocument.dimensions[0] == 0) { + assertEquals(9L, starTreeDocument.metrics[0]); + } else { + assertEquals(starTreeDocument.dimensions[1], starTreeDocument.metrics[0]); + } + } + } + + validateStarTree(builder.getRootNode(), 2, 1000, builder.getStarTreeDocuments()); + + metaOut.close(); + dataOut.close(); + docValuesConsumer.close(); + + StarTreeMetadata starTreeMetadata = getStarTreeMetadata(List.of("field1", "field3"), 6, 1000, 231); + + validateStarTreeFileFormats( + builder.getRootNode(), + builder.getStarTreeDocuments().size(), + starTreeMetadata, + builder.getStarTreeDocuments() + ); + } + + public void testMergeFlowWithMissingDocsWithZeroComplexCase() throws IOException { + List dimList = List.of(0L, 0L, 0L, 0L, 0L); + List docsWithField = List.of(0, 1, 2, 6, 8); + List dimList2 = List.of(0L, 0L, 0L, 0L); + List docsWithField2 = List.of(0, 1, 2, 6); + + List metricsList = List.of(0L, 1L, 2L, 3L, 4L, 5L, 6L, 7L, 8L, 9L); + List metricsWithField = List.of(0, 1, 2, 3, 4, 5, 6, 7, 8, 9); + + List dimList3 = List.of(5L, 6L, 8L, -1L); + List docsWithField3 = List.of(0, 1, 3, 4); + List dimList4 = List.of(5L, 6L, 7L, 8L, -1L); + List docsWithField4 = List.of(0, 1, 2, 3, 4); + + List metricsList2 = List.of(5L, 6L, 7L, 8L, 9L); + List metricsWithField2 = List.of(0, 1, 2, 3, 4); + + compositeField = getStarTreeField(MetricStat.VALUE_COUNT); + StarTreeValues starTreeValues = getStarTreeValues( + getSortedNumericMock(dimList, docsWithField), + getSortedNumericMock(dimList2, docsWithField2), + getSortedNumericMock(metricsList, metricsWithField), + compositeField, + "9" + ); + + StarTreeValues starTreeValues2 = getStarTreeValues( + getSortedNumericMock(dimList3, docsWithField3), + getSortedNumericMock(dimList4, docsWithField4), + getSortedNumericMock(metricsList2, metricsWithField2), + compositeField, + "4" + ); + writeState = getWriteState(4, writeState.segmentInfo.getId()); + this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( + writeState, + Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, + Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, + Composite99DocValuesFormat.META_DOC_VALUES_CODEC, + Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION + ); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); + /** + * Asserting following dim / metrics [ dim1, dim2 / Count [ metric] ] + [0, 0] | [9] + [0, null] | [8] + [5, 5] | [5] + [6, 6] | [6] + [8, 8] | [8] + [null, 7] | [7] + [null, null] | [19] + */ + builder.appendDocumentsToStarTree(starTreeDocumentIterator); + assertEquals(7, builder.getStarTreeDocuments().size()); + builder.build(starTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); + int count = 0; + for (StarTreeDocument starTreeDocument : builder.getStarTreeDocuments()) { + count++; + if (count <= 7) { + if (starTreeDocument.dimensions[0] == null && starTreeDocument.dimensions[1] == null) { + assertEquals(19L, (long) starTreeDocument.metrics[0]); + assertEquals(7, count); + } else if (starTreeDocument.dimensions[0] == null) { + assertEquals(7L, starTreeDocument.metrics[0]); + } else if (starTreeDocument.dimensions[1] == null) { + assertEquals(8L, starTreeDocument.metrics[0]); + } else if (starTreeDocument.dimensions[0] == 0) { + assertEquals(9L, starTreeDocument.metrics[0]); + } else { + assertEquals(starTreeDocument.dimensions[1], starTreeDocument.metrics[0]); + } + } + } + + validateStarTree(builder.getRootNode(), 2, 1000, builder.getStarTreeDocuments()); + + metaOut.close(); + dataOut.close(); + docValuesConsumer.close(); + + StarTreeMetadata starTreeMetadata = getStarTreeMetadata(List.of("field1", "field3"), 7, 1000, 231); + + validateStarTreeFileFormats( + builder.getRootNode(), + builder.getStarTreeDocuments().size(), + starTreeMetadata, + builder.getStarTreeDocuments() + ); + } + + public void testMergeFlowWithMissingDocsInSecondDim() throws IOException { + List dimList2 = List.of(0L, 1L, 2L, 3L, 4L, 6L); + List docsWithField2 = List.of(0, 1, 2, 3, 4, 6); + List dimList = List.of(0L, 1L, 2L, 3L, 4L, 5L, -1L); + List docsWithField = List.of(0, 1, 2, 3, 4, 5, 6); + + List metricsList = List.of(0L, 1L, 2L, 3L, 4L, 5L, 6L); + List metricsWithField = List.of(0, 1, 2, 3, 4, 5, 6); + + List dimList3 = List.of(5L, 6L, 8L, -1L); + List docsWithField3 = List.of(0, 1, 3, 4); + List dimList4 = List.of(5L, 6L, 7L, 8L, -1L); + List docsWithField4 = List.of(0, 1, 2, 3, 4); + + List metricsList2 = List.of(5L, 6L, 7L, 8L, 9L); + List metricsWithField2 = List.of(0, 1, 2, 3, 4); + + compositeField = getStarTreeField(MetricStat.VALUE_COUNT); + StarTreeValues starTreeValues = getStarTreeValues( + getSortedNumericMock(dimList, docsWithField), + getSortedNumericMock(dimList2, docsWithField2), + getSortedNumericMock(metricsList, metricsWithField), + compositeField, + "6" + ); + + StarTreeValues starTreeValues2 = getStarTreeValues( + getSortedNumericMock(dimList3, docsWithField3), + getSortedNumericMock(dimList4, docsWithField4), + getSortedNumericMock(metricsList2, metricsWithField2), + compositeField, + "4" + ); + writeState = getWriteState(4, writeState.segmentInfo.getId()); + this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( + writeState, + Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, + Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, + Composite99DocValuesFormat.META_DOC_VALUES_CODEC, + Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION + ); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); + /** + * Asserting following dim / metrics [ dim1, dim2 / Count [ metric] ] + [0, 0] | [0] + [1, 1] | [1] + [2, 2] | [2] + [3, 3] | [3] + [4, 4] | [4] + [5, 5] | [5] + [5, null] | [5] + [6, 6] | [6] + [8, 8] | [8] + [null, 7] | [7] + */ + builder.appendDocumentsToStarTree(starTreeDocumentIterator); + assertEquals(10, builder.getStarTreeDocuments().size()); + builder.build(starTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); + int count = 0; + for (StarTreeDocument starTreeDocument : builder.getStarTreeDocuments()) { + count++; + if (count <= 10) { + if (starTreeDocument.dimensions[0] != null && starTreeDocument.dimensions[0] == 5) { + assertEquals(starTreeDocument.dimensions[0], starTreeDocument.metrics[0]); + } else { + assertEquals(starTreeDocument.dimensions[1], starTreeDocument.metrics[0]); + } + } + } + + validateStarTree(builder.getRootNode(), 2, 1000, builder.getStarTreeDocuments()); + + metaOut.close(); + dataOut.close(); + docValuesConsumer.close(); + + StarTreeMetadata starTreeMetadata = getStarTreeMetadata(List.of("field1", "field3"), 10, 1000, 363); + + validateStarTreeFileFormats( + builder.getRootNode(), + builder.getStarTreeDocuments().size(), + starTreeMetadata, + builder.getStarTreeDocuments() + ); + } + + public void testMergeFlowWithDocsMissingAtTheEnd() throws IOException { + List dimList = List.of(0L, 1L, 2L, 3L, 4L); + List docsWithField = List.of(0, 1, 2, 3, 4); + List dimList2 = List.of(0L, 1L, 2L, 3L, 4L, 5L, -1L); + List docsWithField2 = List.of(0, 1, 2, 3, 4, 5, 6); + + List metricsList = List.of(0L, 1L, 2L, 3L, 4L, 5L, 6L); + List metricsWithField = List.of(0, 1, 2, 3, 4, 5, 6); + + List dimList3 = List.of(5L, 6L, 8L, -1L); + List docsWithField3 = List.of(0, 1, 3, 4); + List dimList4 = List.of(5L, 6L, 7L, 8L, -1L); + List docsWithField4 = List.of(0, 1, 2, 3, 4); + + List metricsList2 = List.of(5L, 6L, 7L, 8L, 9L); + List metricsWithField2 = List.of(0, 1, 2, 3, 4); + + compositeField = getStarTreeField(MetricStat.VALUE_COUNT); + StarTreeValues starTreeValues = getStarTreeValues( + getSortedNumericMock(dimList, docsWithField), + getSortedNumericMock(dimList2, docsWithField2), + getSortedNumericMock(metricsList, metricsWithField), + compositeField, + "6" + ); + + StarTreeValues starTreeValues2 = getStarTreeValues( + getSortedNumericMock(dimList3, docsWithField3), + getSortedNumericMock(dimList4, docsWithField4), + getSortedNumericMock(metricsList2, metricsWithField2), + compositeField, + "4" + ); + this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( + writeState, + Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, + Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, + Composite99DocValuesFormat.META_DOC_VALUES_CODEC, + Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION + ); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); + /** + * Asserting following dim / metrics [ dim1, dim2 / Count [ metric] ] + [0, 0] | [0] + [1, 1] | [1] + [2, 2] | [2] + [3, 3] | [3] + [4, 4] | [4] + [5, 5] | [5] + [6, 6] | [6] + [8, 8] | [8] + [null, 5] | [5] + [null, 7] | [7] + */ + builder.appendDocumentsToStarTree(starTreeDocumentIterator); + assertEquals(10, builder.getStarTreeDocuments().size()); + builder.build(starTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); + int count = 0; + for (StarTreeDocument starTreeDocument : builder.getStarTreeDocuments()) { + count++; + if (count <= 10) { + if (starTreeDocument.dimensions[0] == null) { + assertTrue(List.of(5L, 7L).contains(starTreeDocument.dimensions[1])); + } + assertEquals(starTreeDocument.dimensions[1], starTreeDocument.metrics[0]); + } + } + + validateStarTree(builder.getRootNode(), 2, 1000, builder.getStarTreeDocuments()); + + metaOut.close(); + dataOut.close(); + docValuesConsumer.close(); + + StarTreeMetadata starTreeMetadata = getStarTreeMetadata(List.of("field1", "field3"), 10, 1000, 363); + + validateStarTreeFileFormats( + builder.getRootNode(), + builder.getStarTreeDocuments().size(), + starTreeMetadata, + builder.getStarTreeDocuments() + ); + } + + public void testMergeFlowWithEmptyFieldsInOneSegment() throws IOException { + List dimList = List.of(0L, 1L, 2L, 3L, 4L); + List docsWithField = List.of(0, 1, 2, 3, 4); + List dimList2 = List.of(0L, 1L, 2L, 3L, 4L, 5L, -1L); + List docsWithField2 = List.of(0, 1, 2, 3, 4, 5, 6); + + List metricsList = List.of(0L, 1L, 2L, 3L, 4L, 5L, 6L); + List metricsWithField = List.of(0, 1, 2, 3, 4, 5, 6); + + compositeField = getStarTreeField(MetricStat.VALUE_COUNT); + StarTreeValues starTreeValues = getStarTreeValues( + getSortedNumericMock(dimList, docsWithField), + getSortedNumericMock(dimList2, docsWithField2), + getSortedNumericMock(metricsList, metricsWithField), + compositeField, + "6" + ); + + StarTreeValues starTreeValues2 = getStarTreeValues( + DocValues.emptySortedNumeric(), + DocValues.emptySortedNumeric(), + DocValues.emptySortedNumeric(), + compositeField, + "0" + ); + writeState = getWriteState(0, writeState.segmentInfo.getId()); + this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( + writeState, + Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, + Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, + Composite99DocValuesFormat.META_DOC_VALUES_CODEC, + Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION + ); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); + /** + * Asserting following dim / metrics [ dim1, dim2 / Count [ metric] ] + [0, 0] | [0] + [1, 1] | [1] + [2, 2] | [2] + [3, 3] | [3] + [4, 4] | [4] + [null, 5] | [5] + */ + builder.appendDocumentsToStarTree(starTreeDocumentIterator); + assertEquals(6, builder.getStarTreeDocuments().size()); + builder.build(starTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); + int count = 0; + for (StarTreeDocument starTreeDocument : builder.getStarTreeDocuments()) { + count++; + if (count <= 6) { + if (starTreeDocument.dimensions[0] == null) { + assertEquals(5L, (long) starTreeDocument.dimensions[1]); + } + assertEquals(starTreeDocument.dimensions[1], starTreeDocument.metrics[0]); + } + } + validateStarTree(builder.getRootNode(), 2, 1000, builder.getStarTreeDocuments()); + + metaOut.close(); + dataOut.close(); + docValuesConsumer.close(); + + StarTreeMetadata starTreeMetadata = getStarTreeMetadata(List.of("field1", "field3"), 6, 1000, 264); + + validateStarTreeFileFormats( + builder.getRootNode(), + builder.getStarTreeDocuments().size(), + starTreeMetadata, + builder.getStarTreeDocuments() + ); + } + + public void testMergeFlowWithDuplicateDimensionValues() throws IOException { + List dimList1 = new ArrayList<>(500); + List docsWithField1 = new ArrayList<>(500); + for (int i = 0; i < 100; i++) { + for (int j = 0; j < 5; j++) { + dimList1.add((long) i); + docsWithField1.add(i * 5 + j); + } + } + + List dimList2 = new ArrayList<>(500); + List docsWithField2 = new ArrayList<>(500); + for (int i = 0; i < 100; i++) { + for (int j = 0; j < 5; j++) { + dimList2.add((long) i); + docsWithField2.add(i * 5 + j); + } + } + + List dimList3 = new ArrayList<>(500); + List docsWithField3 = new ArrayList<>(500); + for (int i = 0; i < 100; i++) { + for (int j = 0; j < 5; j++) { + dimList3.add((long) i); + docsWithField3.add(i * 5 + j); + } + } + + List dimList4 = new ArrayList<>(500); + List docsWithField4 = new ArrayList<>(500); + for (int i = 0; i < 100; i++) { + for (int j = 0; j < 5; j++) { + dimList4.add((long) i); + docsWithField4.add(i * 5 + j); + } + } + + List metricsList = new ArrayList<>(100); + List metricsWithField = new ArrayList<>(100); + for (int i = 0; i < 500; i++) { + metricsList.add(getLongFromDouble(i * 10.0)); + metricsWithField.add(i); + } + List docCountMetricsList = new ArrayList<>(100); + List docCountMetricsWithField = new ArrayList<>(100); + for (int i = 0; i < 500; i++) { + docCountMetricsList.add(i * 10L); + docCountMetricsWithField.add(i); + } + + compositeField = getStarTreeFieldWithDocCount(1, true); + StarTreeValues starTreeValues = getStarTreeValues( + dimList1, + docsWithField1, + dimList2, + docsWithField2, + dimList3, + docsWithField3, + dimList4, + docsWithField4, + metricsList, + metricsWithField, + docCountMetricsList, + docCountMetricsWithField, + compositeField + ); + + StarTreeValues starTreeValues2 = getStarTreeValues( + dimList1, + docsWithField1, + dimList2, + docsWithField2, + dimList3, + docsWithField3, + dimList4, + docsWithField4, + metricsList, + metricsWithField, + docCountMetricsList, + docCountMetricsWithField, + compositeField + ); + this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( + writeState, + Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, + Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, + Composite99DocValuesFormat.META_DOC_VALUES_CODEC, + Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION + ); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + builder.build(builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)), new AtomicInteger(), docValuesConsumer); + List starTreeDocuments = builder.getStarTreeDocuments(); + assertEquals(401, starTreeDocuments.size()); + int count = 0; + double sum = 0; + /** + 401 docs get generated + [0, 0, 0, 0] | [200.0, 10] + [1, 1, 1, 1] | [700.0, 10] + [2, 2, 2, 2] | [1200.0, 10] + [3, 3, 3, 3] | [1700.0, 10] + [4, 4, 4, 4] | [2200.0, 10] + ..... + [null, null, null, 99] | [49700.0, 10] + [null, null, null, null] | [2495000.0, 1000] + */ + for (StarTreeDocument starTreeDocument : starTreeDocuments) { + if (starTreeDocument.dimensions[3] == null) { + assertEquals(sum, starTreeDocument.metrics[0]); + assertEquals(2495000L, (long) starTreeDocument.metrics[1]); + } else { + if (starTreeDocument.dimensions[0] != null) { + sum += (double) starTreeDocument.metrics[0]; + } + assertEquals(starTreeDocument.dimensions[3] * 500 + 200.0, starTreeDocument.metrics[0]); + assertEquals(starTreeDocument.dimensions[3] * 500 + 200L, (long) starTreeDocument.metrics[1]); + + } + count++; + } + assertEquals(401, count); + validateStarTree(builder.getRootNode(), 4, compositeField.getStarTreeConfig().maxLeafDocs(), builder.getStarTreeDocuments()); + metaOut.close(); + dataOut.close(); + docValuesConsumer.close(); + + StarTreeMetadata starTreeMetadata = getStarTreeMetadata( + getStarTreeDimensionNames(compositeField.getDimensionsOrder()), + 100, + 1, + 13365 + ); + + validateStarTreeFileFormats( + builder.getRootNode(), + builder.getStarTreeDocuments().size(), + starTreeMetadata, + builder.getStarTreeDocuments() + ); + } + + public void testMergeFlowWithMaxLeafDocs() throws IOException { + List dimList1 = new ArrayList<>(500); + List docsWithField1 = new ArrayList<>(500); + + for (int i = 0; i < 20; i++) { + for (int j = 0; j < 20; j++) { + dimList1.add((long) i); + docsWithField1.add(i * 20 + j); + } + } + for (int i = 80; i < 100; i++) { + for (int j = 0; j < 5; j++) { + dimList1.add((long) i); + docsWithField1.add(i * 5 + j); + } + } + List dimList3 = new ArrayList<>(500); + List docsWithField3 = new ArrayList<>(500); + for (int i = 0; i < 100; i++) { + for (int j = 0; j < 5; j++) { + dimList3.add((long) i); + docsWithField3.add(i * 5 + j); + } + } + List dimList2 = new ArrayList<>(500); + List docsWithField2 = new ArrayList<>(500); + for (int i = 0; i < 10; i++) { + for (int j = 0; j < 50; j++) { + dimList2.add((long) i); + docsWithField2.add(i * 50 + j); + } + } + + List dimList4 = new ArrayList<>(500); + List docsWithField4 = new ArrayList<>(500); + for (int i = 0; i < 100; i++) { + for (int j = 0; j < 5; j++) { + dimList4.add((long) i); + docsWithField4.add(i * 5 + j); + } + } + + List metricsList = new ArrayList<>(100); + List metricsWithField = new ArrayList<>(100); + for (int i = 0; i < 500; i++) { + metricsList.add(getLongFromDouble(i * 10.0)); + metricsWithField.add(i); + } + + List metricsList1 = new ArrayList<>(100); + List metricsWithField1 = new ArrayList<>(100); + for (int i = 0; i < 500; i++) { + metricsList1.add(1L); + metricsWithField1.add(i); + } + + compositeField = getStarTreeFieldWithDocCount(3, true); + StarTreeValues starTreeValues = getStarTreeValues( + dimList1, + docsWithField1, + dimList2, + docsWithField2, + dimList3, + docsWithField3, + dimList4, + docsWithField4, + metricsList, + metricsWithField, + metricsList1, + metricsWithField1, + compositeField + ); + + StarTreeValues starTreeValues2 = getStarTreeValues( + dimList1, + docsWithField1, + dimList2, + docsWithField2, + dimList3, + docsWithField3, + dimList4, + docsWithField4, + metricsList, + metricsWithField, + metricsList1, + metricsWithField1, + compositeField + ); + + this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( + writeState, + Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, + Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, + Composite99DocValuesFormat.META_DOC_VALUES_CODEC, + Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION + ); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + builder.build(builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)), new AtomicInteger(), docValuesConsumer); + List starTreeDocuments = builder.getStarTreeDocuments(); + /** + 635 docs get generated + [0, 0, 0, 0] | [200.0, 10] + [0, 0, 1, 1] | [700.0, 10] + [0, 0, 2, 2] | [1200.0, 10] + [0, 0, 3, 3] | [1700.0, 10] + [1, 0, 4, 4] | [2200.0, 10] + [1, 0, 5, 5] | [2700.0, 10] + [1, 0, 6, 6] | [3200.0, 10] + [1, 0, 7, 7] | [3700.0, 10] + [2, 0, 8, 8] | [4200.0, 10] + [2, 0, 9, 9] | [4700.0, 10] + [2, 1, 10, 10] | [5200.0, 10] + [2, 1, 11, 11] | [5700.0, 10] + ..... + [18, 7, null, null] | [147800.0, 40] + ... + [7, 2, null, null] | [28900.0, 20] + ... + [null, null, null, 99] | [49700.0, 10] + ..... + [null, null, null, null] | [2495000.0, 1000] + */ + assertEquals(635, starTreeDocuments.size()); + for (StarTreeDocument starTreeDocument : starTreeDocuments) { + if (starTreeDocument.dimensions[0] != null + && starTreeDocument.dimensions[1] != null + && starTreeDocument.dimensions[2] != null + && starTreeDocument.dimensions[3] != null) { + assertEquals(10L, starTreeDocument.metrics[1]); + } else if (starTreeDocument.dimensions[1] != null + && starTreeDocument.dimensions[2] != null + && starTreeDocument.dimensions[3] != null) { + assertEquals(10L, starTreeDocument.metrics[1]); + } else if (starTreeDocument.dimensions[0] != null + && starTreeDocument.dimensions[2] != null + && starTreeDocument.dimensions[3] != null) { + assertEquals(10L, starTreeDocument.metrics[1]); + } else if (starTreeDocument.dimensions[0] != null + && starTreeDocument.dimensions[1] != null + && starTreeDocument.dimensions[3] != null) { + assertEquals(10L, starTreeDocument.metrics[1]); + } else if (starTreeDocument.dimensions[0] != null && starTreeDocument.dimensions[3] != null) { + assertEquals(10L, starTreeDocument.metrics[1]); + } else if (starTreeDocument.dimensions[0] != null && starTreeDocument.dimensions[1] != null) { + assertTrue((long) starTreeDocument.metrics[1] == 20L || (long) starTreeDocument.metrics[1] == 40L); + } else if (starTreeDocument.dimensions[1] != null && starTreeDocument.dimensions[3] != null) { + assertEquals(10L, starTreeDocument.metrics[1]); + } else if (starTreeDocument.dimensions[1] != null) { + assertEquals(100L, starTreeDocument.metrics[1]); + } else if (starTreeDocument.dimensions[0] != null) { + assertEquals(40L, starTreeDocument.metrics[1]); + } + } + validateStarTree(builder.getRootNode(), 4, compositeField.getStarTreeConfig().maxLeafDocs(), builder.getStarTreeDocuments()); + metaOut.close(); + dataOut.close(); + docValuesConsumer.close(); + + StarTreeMetadata starTreeMetadata = getStarTreeMetadata( + getStarTreeDimensionNames(compositeField.getDimensionsOrder()), + 100, + 3, + 23199 + ); + + validateStarTreeFileFormats( + builder.getRootNode(), + builder.getStarTreeDocuments().size(), + starTreeMetadata, + builder.getStarTreeDocuments() + ); + } + + public void testMergeFlowWithDifferentDocsFromSegments() throws IOException { + List dimList = List.of(0L, 1L, 3L, 4L, 5L, 6L); + List docsWithField = List.of(0, 1, 3, 4, 5, 6); + List dimList2 = List.of(0L, 1L, 2L, 3L, 4L, 5L, -1L); + List docsWithField2 = List.of(0, 1, 2, 3, 4, 5, 6); + + List metricsList = List.of(0L, 1L, 2L, 3L, 4L, 5L, 6L); + List metricsWithField = List.of(0, 1, 2, 3, 4, 5, 6); + + List dimList3 = List.of(5L, 6L, 8L, -1L); + List docsWithField3 = List.of(0, 1, 3, 4); + List dimList4 = List.of(5L, 6L, 7L, 8L, -1L); + List docsWithField4 = List.of(0, 1, 2, 3, 4); + + List metricsList2 = List.of(5L, 6L, 7L, 8L, 9L); + List metricsWithField2 = List.of(0, 1, 2, 3, 4); + + compositeField = getStarTreeField(MetricStat.VALUE_COUNT); + StarTreeValues starTreeValues = getStarTreeValues( + getSortedNumericMock(dimList, docsWithField), + getSortedNumericMock(dimList2, docsWithField2), + getSortedNumericMock(metricsList, metricsWithField), + compositeField, + "6" + ); + + StarTreeValues starTreeValues2 = getStarTreeValues( + getSortedNumericMock(dimList3, docsWithField3), + getSortedNumericMock(dimList4, docsWithField4), + getSortedNumericMock(metricsList2, metricsWithField2), + compositeField, + "4" + ); + writeState = getWriteState(4, writeState.segmentInfo.getId()); + this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( + writeState, + Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, + Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, + Composite99DocValuesFormat.META_DOC_VALUES_CODEC, + Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION + ); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); + /** + * Asserting following dim / metrics [ dim1, dim2 / Count [ metric] ] + [0, 0] | [0] + [1, 1] | [1] + [3, 3] | [3] + [4, 4] | [4] + [5, 5] | [10] + [6, 6] | [6] + [8, 8] | [8] + [null, 2] | [2] + [null, 7] | [7] + */ + builder.appendDocumentsToStarTree(starTreeDocumentIterator); + assertEquals(9, builder.getStarTreeDocuments().size()); + builder.build(starTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); + int count = 0; + for (StarTreeDocument starTreeDocument : builder.getStarTreeDocuments()) { + count++; + if (count <= 9) { + if (Objects.equals(starTreeDocument.dimensions[0], 5L)) { + assertEquals(starTreeDocument.dimensions[0] * 2, starTreeDocument.metrics[0]); + } else { + assertEquals(starTreeDocument.dimensions[1], starTreeDocument.metrics[0]); + } + } + } + validateStarTree(builder.getRootNode(), 2, 1000, builder.getStarTreeDocuments()); + + metaOut.close(); + dataOut.close(); + docValuesConsumer.close(); + + StarTreeMetadata starTreeMetadata = getStarTreeMetadata(List.of("field1", "field3"), 9, 1000, 330); + + validateStarTreeFileFormats( + builder.getRootNode(), + builder.getStarTreeDocuments().size(), + starTreeMetadata, + builder.getStarTreeDocuments() + ); + } + + public void testMergeFlowWithDuplicateDimensionValueWithMaxLeafDocs() throws IOException { + List dimList1 = new ArrayList<>(500); + List docsWithField1 = new ArrayList<>(500); + + for (int i = 0; i < 20; i++) { + for (int j = 0; j < 20; j++) { + dimList1.add((long) i); + docsWithField1.add(i * 20 + j); + } + } + for (int i = 80; i < 100; i++) { + for (int j = 0; j < 5; j++) { + dimList1.add((long) i); + docsWithField1.add(i * 5 + j); + } + } + List dimList3 = new ArrayList<>(500); + List docsWithField3 = new ArrayList<>(500); + for (int i = 0; i < 100; i++) { + for (int j = 0; j < 5; j++) { + dimList3.add((long) i); + docsWithField3.add(i * 5 + j); + } + } + List dimList2 = new ArrayList<>(500); + List docsWithField2 = new ArrayList<>(500); + for (int i = 0; i < 500; i++) { + dimList2.add((long) 1); + docsWithField2.add(i); + } + + List dimList4 = new ArrayList<>(500); + List docsWithField4 = new ArrayList<>(500); + for (int i = 0; i < 100; i++) { + for (int j = 0; j < 5; j++) { + dimList4.add((long) i); + docsWithField4.add(i * 5 + j); + } + } + + List metricsList = new ArrayList<>(100); + List metricsWithField = new ArrayList<>(100); + for (int i = 0; i < 500; i++) { + metricsList.add(getLongFromDouble(i * 10.0)); + metricsWithField.add(i); + } + + List docCountMetricsList = new ArrayList<>(100); + List docCountMetricsWithField = new ArrayList<>(100); + for (int i = 0; i < 500; i++) { + metricsList.add(getLongFromDouble(i * 2)); + metricsWithField.add(i); + } + + compositeField = getStarTreeFieldWithDocCount(3, true); + StarTreeValues starTreeValues = getStarTreeValues( + dimList1, + docsWithField1, + dimList2, + docsWithField2, + dimList3, + docsWithField3, + dimList4, + docsWithField4, + metricsList, + metricsWithField, + docCountMetricsList, + docCountMetricsWithField, + compositeField + ); + + StarTreeValues starTreeValues2 = getStarTreeValues( + dimList1, + docsWithField1, + dimList2, + docsWithField2, + dimList3, + docsWithField3, + dimList4, + docsWithField4, + metricsList, + metricsWithField, + docCountMetricsList, + docCountMetricsWithField, + compositeField + ); + this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( + writeState, + Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, + Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, + Composite99DocValuesFormat.META_DOC_VALUES_CODEC, + Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION + ); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + builder.build(builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)), new AtomicInteger(), docValuesConsumer); + List starTreeDocuments = builder.getStarTreeDocuments(); + assertEquals(401, starTreeDocuments.size()); + validateStarTree(builder.getRootNode(), 4, compositeField.getStarTreeConfig().maxLeafDocs(), builder.getStarTreeDocuments()); + + metaOut.close(); + dataOut.close(); + docValuesConsumer.close(); + + StarTreeMetadata starTreeMetadata = getStarTreeMetadata( + getStarTreeDimensionNames(compositeField.getDimensionsOrder()), + 100, + compositeField.getStarTreeConfig().maxLeafDocs(), + 15345 + ); + + validateStarTreeFileFormats( + builder.getRootNode(), + builder.getStarTreeDocuments().size(), + starTreeMetadata, + builder.getStarTreeDocuments() + ); + } + + public void testMergeFlowWithMaxLeafDocsAndStarTreeNodesAssertion() throws IOException { + List dimList1 = new ArrayList<>(500); + List docsWithField1 = new ArrayList<>(500); + Map> expectedDimToValueMap = new HashMap<>(); + Map dimValueMap = new HashMap<>(); + for (int i = 0; i < 20; i++) { + for (int j = 0; j < 20; j++) { + dimList1.add((long) i); + docsWithField1.add(i * 20 + j); + } + // metric = no of docs * 10.0 + dimValueMap.put((long) i, 200.0); + } + for (int i = 80; i < 100; i++) { + for (int j = 0; j < 5; j++) { + dimList1.add((long) i); + docsWithField1.add(i * 5 + j); + } + // metric = no of docs * 10.0 + dimValueMap.put((long) i, 50.0); + } + dimValueMap.put(Long.MAX_VALUE, 5000.0); + expectedDimToValueMap.put(0, dimValueMap); + dimValueMap = new HashMap<>(); + List dimList3 = new ArrayList<>(500); + List docsWithField3 = new ArrayList<>(500); + for (int i = 0; i < 500; i++) { + dimList3.add((long) 1); + docsWithField3.add(i); + dimValueMap.put((long) i, 10.0); + } + dimValueMap.put(Long.MAX_VALUE, 5000.0); + expectedDimToValueMap.put(2, dimValueMap); + dimValueMap = new HashMap<>(); + List dimList2 = new ArrayList<>(500); + List docsWithField2 = new ArrayList<>(500); + for (int i = 0; i < 500; i++) { + dimList2.add((long) i); + docsWithField2.add(i); + dimValueMap.put((long) i, 10.0); + } + dimValueMap.put(Long.MAX_VALUE, 200.0); + expectedDimToValueMap.put(1, dimValueMap); + dimValueMap = new HashMap<>(); + List dimList4 = new ArrayList<>(500); + List docsWithField4 = new ArrayList<>(500); + for (int i = 0; i < 500; i++) { + dimList4.add((long) 1); + docsWithField4.add(i); + dimValueMap.put((long) i, 10.0); + } + dimValueMap.put(Long.MAX_VALUE, 5000.0); + expectedDimToValueMap.put(3, dimValueMap); + List metricsList = new ArrayList<>(100); + List metricsWithField = new ArrayList<>(100); + for (int i = 0; i < 500; i++) { + metricsList.add(getLongFromDouble(10.0)); + metricsWithField.add(i); + } + List metricsList1 = new ArrayList<>(100); + List metricsWithField1 = new ArrayList<>(100); + for (int i = 0; i < 500; i++) { + metricsList.add(1L); + metricsWithField.add(i); + } + compositeField = getStarTreeFieldWithDocCount(10, true); + StarTreeValues starTreeValues = getStarTreeValues( + dimList1, + docsWithField1, + dimList2, + docsWithField2, + dimList3, + docsWithField3, + dimList4, + docsWithField4, + metricsList, + metricsWithField, + metricsList1, + metricsWithField1, + compositeField + ); + + StarTreeValues starTreeValues2 = getStarTreeValues( + dimList1, + docsWithField1, + dimList2, + docsWithField2, + dimList3, + docsWithField3, + dimList4, + docsWithField4, + metricsList, + metricsWithField, + metricsList1, + metricsWithField1, + compositeField + ); + this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( + writeState, + Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, + Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, + Composite99DocValuesFormat.META_DOC_VALUES_CODEC, + Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION + ); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + builder.build(builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)), new AtomicInteger(), docValuesConsumer); + List starTreeDocuments = builder.getStarTreeDocuments(); + Map> dimValueToDocIdMap = new HashMap<>(); + traverseStarTree(builder.rootNode, dimValueToDocIdMap, true); + for (Map.Entry> entry : dimValueToDocIdMap.entrySet()) { + int dimId = entry.getKey(); + if (dimId == -1) continue; + Map map = expectedDimToValueMap.get(dimId); + for (Map.Entry dimValueToDocIdEntry : entry.getValue().entrySet()) { + long dimValue = dimValueToDocIdEntry.getKey(); + int docId = dimValueToDocIdEntry.getValue(); + assertEquals(map.get(dimValue) * 2, starTreeDocuments.get(docId).metrics[0]); + } + } + assertEquals(1041, starTreeDocuments.size()); + validateStarTree(builder.getRootNode(), 4, compositeField.getStarTreeConfig().maxLeafDocs(), builder.getStarTreeDocuments()); + + metaOut.close(); + dataOut.close(); + docValuesConsumer.close(); + + StarTreeMetadata starTreeMetadata = getStarTreeMetadata( + getStarTreeDimensionNames(compositeField.getDimensionsOrder()), + 500, + compositeField.getStarTreeConfig().maxLeafDocs(), + 31779 + ); + + validateStarTreeFileFormats( + builder.getRootNode(), + builder.getStarTreeDocuments().size(), + starTreeMetadata, + builder.getStarTreeDocuments() + ); + } + + public void testMergeFlowWithTimestamps() throws IOException { + List dimList = List.of(1655288152000L, 1655288092000L, 1655288032000L, 1655287972000L, 1655288092000L, 1655288092000L); + List docsWithField = List.of(0, 1, 2, 3, 4, 6); + List dimList2 = List.of(1655288152000L, 1655288092000L, 1655288032000L, 1655287972000L, 1655288092000L, 1655288092000L, -1L); + List docsWithField2 = List.of(0, 1, 2, 3, 4, 6); + List dimList7 = List.of(1655288152000L, 1655288092000L, 1655288032000L, 1655287972000L, 1655288092000L, 1655288092000L, -1L); + List docsWithField7 = List.of(0, 1, 2, 3, 4, 6); + + List dimList5 = List.of(0L, 1L, 2L, 3L, 4L, 5L, 6L); + List docsWithField5 = List.of(0, 1, 2, 3, 4, 5, 6); + List metricsList1 = List.of( + getLongFromDouble(0.0), + getLongFromDouble(10.0), + getLongFromDouble(20.0), + getLongFromDouble(30.0), + getLongFromDouble(40.0), + getLongFromDouble(50.0), + getLongFromDouble(60.0) + ); + List metricsWithField1 = List.of(0, 1, 2, 3, 4, 5, 6); + List metricsList = List.of(0L, 1L, 2L, 3L, 4L, 5L, 6L); + List metricsWithField = List.of(0, 1, 2, 3, 4, 5, 6); + + List dimList3 = List.of(1655288152000L, 1655288092000L, 1655288032000L, -1L); + List docsWithField3 = List.of(0, 1, 3, 4); + List dimList4 = List.of(1655288152000L, 1655288092000L, 1655288032000L, -1L); + List docsWithField4 = List.of(0, 1, 3, 4); + List dimList8 = List.of(1655288152000L, 1655288092000L, 1655288032000L, -1L); + List docsWithField8 = List.of(0, 1, 3, 4); + + List dimList6 = List.of(5L, 6L, 7L, 8L); + List docsWithField6 = List.of(0, 1, 2, 3); + List metricsList21 = List.of( + getLongFromDouble(50.0), + getLongFromDouble(60.0), + getLongFromDouble(70.0), + getLongFromDouble(80.0), + getLongFromDouble(90.0) + ); + List metricsWithField21 = List.of(0, 1, 2, 3, 4); + List metricsList2 = List.of(5L, 6L, 7L, 8L, 9L); + List metricsWithField2 = List.of(0, 1, 2, 3, 4); + + compositeField = getStarTreeFieldWithDateDimension(); + StarTreeValues starTreeValues = getStarTreeValuesWithDates( + getSortedNumericMock(dimList, docsWithField), + getSortedNumericMock(dimList2, docsWithField2), + getSortedNumericMock(dimList7, docsWithField7), + getSortedNumericMock(dimList5, docsWithField5), + getSortedNumericMock(metricsList, metricsWithField), + getSortedNumericMock(metricsList1, metricsWithField1), + compositeField, + "6" + ); + + StarTreeValues starTreeValues2 = getStarTreeValuesWithDates( + getSortedNumericMock(dimList3, docsWithField3), + getSortedNumericMock(dimList4, docsWithField4), + getSortedNumericMock(dimList8, docsWithField8), + getSortedNumericMock(dimList6, docsWithField6), + getSortedNumericMock(metricsList2, metricsWithField2), + getSortedNumericMock(metricsList21, metricsWithField21), + compositeField, + "4" + ); + this.docValuesConsumer = LuceneDocValuesConsumerFactory.getDocValuesConsumerForCompositeCodec( + writeState, + Composite99DocValuesFormat.DATA_DOC_VALUES_CODEC, + Composite99DocValuesFormat.DATA_DOC_VALUES_EXTENSION, + Composite99DocValuesFormat.META_DOC_VALUES_CODEC, + Composite99DocValuesFormat.META_DOC_VALUES_EXTENSION + ); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, getWriteState(4, writeState.segmentInfo.getId()), mapperService); + Iterator starTreeDocumentIterator = builder.mergeStarTrees(List.of(starTreeValues, starTreeValues2)); + /** + [1655287972000, 1655287972000, 1655287972000, 3] | [30.0, 3] + [1655288032000, 1655288032000, 1655288032000, 2] | [20.0, 2] + [1655288032000, 1655288032000, 1655288032000, 8] | [80.0, 8] + [1655288092000, 1655288092000, 1655288092000, 1] | [10.0, 1] + [1655288092000, 1655288092000, 1655288092000, 4] | [40.0, 4] + [1655288092000, 1655288092000, 1655288092000, 6] | [60.0, 6] + [1655288152000, 1655288152000, 1655288152000, 0] | [0.0, 0] + [1655288152000, 1655288152000, 1655288152000, 5] | [50.0, 5] + [null, null, null, 5] | [50.0, 5] + [null, null, null, 7] | [70.0, 7] + */ + int count = 0; + builder.appendDocumentsToStarTree(starTreeDocumentIterator); + for (StarTreeDocument starTreeDocument : builder.getStarTreeDocuments()) { + count++; + assertEquals(starTreeDocument.dimensions[3] * 10.0, (double) starTreeDocument.metrics[1], 0); + assertEquals(starTreeDocument.dimensions[3], starTreeDocument.metrics[0]); + } + assertEquals(10, count); + builder.build(starTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); + validateStarTree(builder.getRootNode(), 4, 10, builder.getStarTreeDocuments()); + metaOut.close(); + dataOut.close(); + docValuesConsumer.close(); + + StarTreeMetadata starTreeMetadata = getStarTreeMetadata( + getStarTreeDimensionNames(compositeField.getDimensionsOrder()), + 10, + compositeField.getStarTreeConfig().maxLeafDocs(), + 231 + ); + + validateStarTreeFileFormats( + builder.getRootNode(), + builder.getStarTreeDocuments().size(), + starTreeMetadata, + builder.getStarTreeDocuments() + ); + } + + private StarTreeValues getStarTreeValuesWithDates( + SortedNumericDocValues dimList, + SortedNumericDocValues dimList2, + SortedNumericDocValues dimList4, + SortedNumericDocValues dimList3, + SortedNumericDocValues metricsList, + SortedNumericDocValues metricsList1, + StarTreeField sf, + String number + ) { + Map> dimDocIdSetIterators = Map.of( + "field1_minute", + () -> new SortedNumericStarTreeValuesIterator(dimList), + "field1_half-hour", + () -> new SortedNumericStarTreeValuesIterator(dimList4), + "field1_hour", + () -> new SortedNumericStarTreeValuesIterator(dimList2), + "field3", + () -> new SortedNumericStarTreeValuesIterator(dimList3) + ); + Map> metricDocIdSetIterators = new LinkedHashMap<>(); + + metricDocIdSetIterators.put( + fullyQualifiedFieldNameForStarTreeMetricsDocValues( + sf.getName(), + "field2", + sf.getMetrics().get(0).getMetrics().get(0).getTypeName() + ), + () -> new SortedNumericStarTreeValuesIterator(metricsList) + ); + metricDocIdSetIterators.put( + fullyQualifiedFieldNameForStarTreeMetricsDocValues( + sf.getName(), + "field2", + sf.getMetrics().get(0).getMetrics().get(1).getTypeName() + ), + () -> new SortedNumericStarTreeValuesIterator(metricsList1) + ); + return new StarTreeValues(sf, null, dimDocIdSetIterators, metricDocIdSetIterators, Map.of(SEGMENT_DOCS_COUNT, number), null); + } + + private StarTreeValues getStarTreeValues( + SortedNumericDocValues dimList, + SortedNumericDocValues dimList2, + SortedNumericDocValues metricsList, + StarTreeField sf, + String number + ) { + SortedNumericDocValues d1sndv = dimList; + SortedNumericDocValues d2sndv = dimList2; + SortedNumericDocValues m1sndv = metricsList; + Map> dimDocIdSetIterators = Map.of( + "field1", + () -> new SortedNumericStarTreeValuesIterator(d1sndv), + "field3", + () -> new SortedNumericStarTreeValuesIterator(d2sndv) + ); + + Map> metricDocIdSetIterators = new LinkedHashMap<>(); + for (Metric metric : sf.getMetrics()) { + for (MetricStat metricStat : metric.getMetrics()) { + String metricFullName = fullyQualifiedFieldNameForStarTreeMetricsDocValues( + sf.getName(), + metric.getField(), + metricStat.getTypeName() + ); + metricDocIdSetIterators.put(metricFullName, () -> new SortedNumericStarTreeValuesIterator(m1sndv)); + } + } + + StarTreeValues starTreeValues = new StarTreeValues( + sf, + null, + dimDocIdSetIterators, + metricDocIdSetIterators, + Map.of(SEGMENT_DOCS_COUNT, number), + null + ); + return starTreeValues; + } + + private StarTreeValues getStarTreeValues( + List dimList1, + List docsWithField1, + List dimList2, + List docsWithField2, + List dimList3, + List docsWithField3, + List dimList4, + List docsWithField4, + List metricsList, + List metricsWithField, + List metricsList1, + List metricsWithField1, + StarTreeField sf + ) { + SortedNumericDocValues d1sndv = getSortedNumericMock(dimList1, docsWithField1); + SortedNumericDocValues d2sndv = getSortedNumericMock(dimList2, docsWithField2); + SortedNumericDocValues d3sndv = getSortedNumericMock(dimList3, docsWithField3); + SortedNumericDocValues d4sndv = getSortedNumericMock(dimList4, docsWithField4); + SortedNumericDocValues m1sndv = getSortedNumericMock(metricsList, metricsWithField); + SortedNumericDocValues m2sndv = getSortedNumericMock(metricsList1, metricsWithField1); + Map> dimDocIdSetIterators = Map.of( + "field1", + () -> new SortedNumericStarTreeValuesIterator(d1sndv), + "field3", + () -> new SortedNumericStarTreeValuesIterator(d2sndv), + "field5", + () -> new SortedNumericStarTreeValuesIterator(d3sndv), + "field8", + () -> new SortedNumericStarTreeValuesIterator(d4sndv) + ); + + Map> metricDocIdSetIterators = new LinkedHashMap<>(); + + metricDocIdSetIterators.put( + fullyQualifiedFieldNameForStarTreeMetricsDocValues( + sf.getName(), + "field2", + sf.getMetrics().get(0).getMetrics().get(0).getTypeName() + ), + () -> new SortedNumericStarTreeValuesIterator(m1sndv) + ); + metricDocIdSetIterators.put( + fullyQualifiedFieldNameForStarTreeMetricsDocValues( + sf.getName(), + "_doc_count", + sf.getMetrics().get(1).getMetrics().get(0).getTypeName() + ), + () -> new SortedNumericStarTreeValuesIterator(m2sndv) + ); + // metricDocIdSetIterators.put("field2", () -> m1sndv); + // metricDocIdSetIterators.put("_doc_count", () -> m2sndv); + StarTreeValues starTreeValues = new StarTreeValues( + sf, + null, + dimDocIdSetIterators, + metricDocIdSetIterators, + getAttributes(500), + null + ); + return starTreeValues; + } +} diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuilderSortAndAggregateTests.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuilderSortAndAggregateTests.java new file mode 100644 index 0000000000000..851535bcc2b26 --- /dev/null +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuilderSortAndAggregateTests.java @@ -0,0 +1,695 @@ +/* + * 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.startree.builder; + +import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; + +import org.apache.lucene.util.NumericUtils; +import org.opensearch.index.compositeindex.datacube.startree.StarTreeDocument; +import org.opensearch.index.compositeindex.datacube.startree.StarTreeFieldConfiguration; +import org.opensearch.index.compositeindex.datacube.startree.utils.SequentialDocValuesIterator; + +import java.io.IOException; +import java.util.Arrays; +import java.util.Collection; +import java.util.Iterator; +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.opensearch.index.compositeindex.datacube.startree.builder.BuilderTestsUtils.getDimensionIterators; +import static org.opensearch.index.compositeindex.datacube.startree.builder.BuilderTestsUtils.getMetricIterators; +import static org.opensearch.index.compositeindex.datacube.startree.builder.BuilderTestsUtils.validateStarTree; + +public class StarTreeBuilderSortAndAggregateTests extends StarTreeBuilderTestCase { + + public StarTreeBuilderSortAndAggregateTests(StarTreeFieldConfiguration.StarTreeBuildMode buildMode) { + super(buildMode); + } + + @ParametersFactory + public static Collection parameters() { + return Arrays.asList( + new Object[] { StarTreeFieldConfiguration.StarTreeBuildMode.ON_HEAP }, + new Object[] { StarTreeFieldConfiguration.StarTreeBuildMode.OFF_HEAP } + ); + } + + public void test_sortAndAggregateStarTreeDocuments() throws IOException { + + int noOfStarTreeDocuments = 5; + StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; + + starTreeDocuments[0] = new StarTreeDocument( + new Long[] { 2L, 4L, 3L, 4L }, + new Object[] { 12.0, 10.0, randomDouble(), 8.0, 20.0, 10L } + ); + starTreeDocuments[1] = new StarTreeDocument( + new Long[] { 3L, 4L, 2L, 1L }, + new Object[] { 10.0, 6.0, randomDouble(), 12.0, 10.0, 10L } + ); + starTreeDocuments[2] = new StarTreeDocument( + new Long[] { 3L, 4L, 2L, 1L }, + new Object[] { 14.0, 12.0, randomDouble(), 6.0, 24.0, 10L } + ); + starTreeDocuments[3] = new StarTreeDocument( + new Long[] { 2L, 4L, 3L, 4L }, + new Object[] { 9.0, 4.0, randomDouble(), 9.0, 12.0, null } + ); + starTreeDocuments[4] = new StarTreeDocument( + new Long[] { 3L, 4L, 2L, 1L }, + new Object[] { 11.0, 16.0, randomDouble(), 8.0, 13.0, null } + ); + + StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; + for (int i = 0; i < noOfStarTreeDocuments; i++) { + long metric1 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]); + long metric2 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[1]); + long metric3 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]); + long metric4 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[3]); + long metric5 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[4]); + Long metric6 = (Long) starTreeDocuments[i].metrics[5]; + segmentStarTreeDocuments[i] = new StarTreeDocument( + starTreeDocuments[i].dimensions, + new Long[] { metric1, metric2, metric3, metric4, metric5, metric6 } + ); + } + List inorderStarTreeDocuments = List.of( + new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Object[] { 21.0, 14.0, 2L, 8.0, 20.0, 11L }), + new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Object[] { 35.0, 34.0, 3L, 6.0, 24.0, 21L }) + ); + Iterator expectedStarTreeDocumentIterator = inorderStarTreeDocuments.iterator(); + + SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); + List metricsIterators = getMetricIterators(segmentStarTreeDocuments); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( + dimsIterators, + metricsIterators + ); + + int numOfAggregatedDocuments = 0; + while (segmentStarTreeDocumentIterator.hasNext() && expectedStarTreeDocumentIterator.hasNext()) { + StarTreeDocument resultStarTreeDocument = segmentStarTreeDocumentIterator.next(); + StarTreeDocument expectedStarTreeDocument = expectedStarTreeDocumentIterator.next(); + + assertEquals(expectedStarTreeDocument.dimensions[0], resultStarTreeDocument.dimensions[0]); + assertEquals(expectedStarTreeDocument.dimensions[1], resultStarTreeDocument.dimensions[1]); + assertEquals(expectedStarTreeDocument.dimensions[2], resultStarTreeDocument.dimensions[2]); + assertEquals(expectedStarTreeDocument.dimensions[3], resultStarTreeDocument.dimensions[3]); + assertEquals(expectedStarTreeDocument.metrics[0], resultStarTreeDocument.metrics[0]); + assertEquals(expectedStarTreeDocument.metrics[1], resultStarTreeDocument.metrics[1]); + assertEquals(expectedStarTreeDocument.metrics[2], resultStarTreeDocument.metrics[2]); + assertEquals(expectedStarTreeDocument.metrics[3], resultStarTreeDocument.metrics[3]); + assertEquals(expectedStarTreeDocument.metrics[4], resultStarTreeDocument.metrics[4]); + assertEquals(expectedStarTreeDocument.metrics[5], resultStarTreeDocument.metrics[5]); + + numOfAggregatedDocuments++; + } + + assertEquals(inorderStarTreeDocuments.size(), numOfAggregatedDocuments); + } + + public void test_sortAndAggregateStarTreeDocument_DoubleMaxAndDoubleMinMetrics() throws IOException { + + int noOfStarTreeDocuments = 5; + StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; + + starTreeDocuments[0] = new StarTreeDocument( + new Long[] { 2L, 4L, 3L, 4L }, + new Object[] { Double.MAX_VALUE, 10.0, randomDouble(), 8.0, 20.0, 100L } + ); + starTreeDocuments[1] = new StarTreeDocument( + new Long[] { 3L, 4L, 2L, 1L }, + new Object[] { 10.0, 6.0, randomDouble(), 12.0, 10.0, 100L } + ); + starTreeDocuments[2] = new StarTreeDocument( + new Long[] { 3L, 4L, 2L, 1L }, + new Object[] { 14.0, Double.MIN_VALUE, randomDouble(), 6.0, 24.0, 100L } + ); + starTreeDocuments[3] = new StarTreeDocument( + new Long[] { 2L, 4L, 3L, 4L }, + new Object[] { 9.0, 4.0, randomDouble(), 9.0, 12.0, 100L } + ); + starTreeDocuments[4] = new StarTreeDocument( + new Long[] { 3L, 4L, 2L, 1L }, + new Object[] { 11.0, 16.0, randomDouble(), 8.0, 13.0, 100L } + ); + + List inorderStarTreeDocuments = List.of( + new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Object[] { Double.MAX_VALUE + 9, 14.0, 2L, 8.0, 20.0, 200L }), + new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Object[] { 35.0, Double.MIN_VALUE + 22, 3L, 6.0, 24.0, 300L }) + ); + Iterator expectedStarTreeDocumentIterator = inorderStarTreeDocuments.iterator(); + + StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; + for (int i = 0; i < noOfStarTreeDocuments; i++) { + long metric1 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]); + long metric2 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[1]); + long metric3 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]); + long metric4 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[3]); + long metric5 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[4]); + Long metric6 = (Long) starTreeDocuments[i].metrics[5]; + segmentStarTreeDocuments[i] = new StarTreeDocument( + starTreeDocuments[i].dimensions, + new Long[] { metric1, metric2, metric3, metric4, metric5, metric6 } + ); + } + + SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); + List metricsIterators = getMetricIterators(segmentStarTreeDocuments); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( + dimsIterators, + metricsIterators + ); + int numOfAggregatedDocuments = 0; + while (segmentStarTreeDocumentIterator.hasNext() && expectedStarTreeDocumentIterator.hasNext()) { + StarTreeDocument resultStarTreeDocument = segmentStarTreeDocumentIterator.next(); + StarTreeDocument expectedStarTreeDocument = expectedStarTreeDocumentIterator.next(); + + assertEquals(expectedStarTreeDocument.dimensions[0], resultStarTreeDocument.dimensions[0]); + assertEquals(expectedStarTreeDocument.dimensions[1], resultStarTreeDocument.dimensions[1]); + assertEquals(expectedStarTreeDocument.dimensions[2], resultStarTreeDocument.dimensions[2]); + assertEquals(expectedStarTreeDocument.dimensions[3], resultStarTreeDocument.dimensions[3]); + assertEquals(expectedStarTreeDocument.metrics[0], resultStarTreeDocument.metrics[0]); + assertEquals(expectedStarTreeDocument.metrics[1], resultStarTreeDocument.metrics[1]); + assertEquals(expectedStarTreeDocument.metrics[2], resultStarTreeDocument.metrics[2]); + assertEquals(expectedStarTreeDocument.metrics[3], resultStarTreeDocument.metrics[3]); + assertEquals(expectedStarTreeDocument.metrics[4], resultStarTreeDocument.metrics[4]); + + numOfAggregatedDocuments++; + } + + assertEquals(inorderStarTreeDocuments.size(), numOfAggregatedDocuments); + builder.build(segmentStarTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); + validateStarTree(builder.getRootNode(), 3, 1, builder.getStarTreeDocuments()); + + } + + public void test_sortAndAggregateStarTreeDocument_longMaxAndLongMinDimensions() throws IOException { + + int noOfStarTreeDocuments = 5; + StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; + + starTreeDocuments[0] = new StarTreeDocument( + new Long[] { Long.MIN_VALUE, 4L, 3L, 4L }, + new Object[] { 12.0, 10.0, randomDouble(), 8.0, 20.0 } + ); + starTreeDocuments[1] = new StarTreeDocument( + new Long[] { 3L, 4L, 2L, Long.MAX_VALUE }, + new Object[] { 10.0, 6.0, randomDouble(), 12.0, 10.0 } + ); + starTreeDocuments[2] = new StarTreeDocument( + new Long[] { 3L, 4L, 2L, Long.MAX_VALUE }, + new Object[] { 14.0, 12.0, randomDouble(), 6.0, 24.0 } + ); + starTreeDocuments[3] = new StarTreeDocument( + new Long[] { Long.MIN_VALUE, 4L, 3L, 4L }, + new Object[] { 9.0, 4.0, randomDouble(), 9.0, 12.0 } + ); + starTreeDocuments[4] = new StarTreeDocument( + new Long[] { 3L, 4L, 2L, Long.MAX_VALUE }, + new Object[] { 11.0, 16.0, randomDouble(), 8.0, 13.0 } + ); + + List inorderStarTreeDocuments = List.of( + new StarTreeDocument(new Long[] { Long.MIN_VALUE, 4L, 3L, 4L }, new Object[] { 21.0, 14.0, 2L, 8.0, 20.0, 2L }), + new StarTreeDocument(new Long[] { 3L, 4L, 2L, Long.MAX_VALUE }, new Object[] { 35.0, 34.0, 3L, 6.0, 24.0, 3L }) + ); + Iterator expectedStarTreeDocumentIterator = inorderStarTreeDocuments.iterator(); + + StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; + for (int i = 0; i < noOfStarTreeDocuments; i++) { + long metric1 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]); + long metric2 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[1]); + long metric3 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]); + long metric4 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[3]); + long metric5 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[4]); + segmentStarTreeDocuments[i] = new StarTreeDocument( + starTreeDocuments[i].dimensions, + new Long[] { metric1, metric2, metric3, metric4, metric5, null } + ); + } + + SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); + List metricsIterators = getMetricIterators(segmentStarTreeDocuments); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( + dimsIterators, + metricsIterators + ); + int numOfAggregatedDocuments = 0; + while (segmentStarTreeDocumentIterator.hasNext() && expectedStarTreeDocumentIterator.hasNext()) { + StarTreeDocument resultStarTreeDocument = segmentStarTreeDocumentIterator.next(); + StarTreeDocument expectedStarTreeDocument = expectedStarTreeDocumentIterator.next(); + + assertEquals(expectedStarTreeDocument.dimensions[0], resultStarTreeDocument.dimensions[0]); + assertEquals(expectedStarTreeDocument.dimensions[1], resultStarTreeDocument.dimensions[1]); + assertEquals(expectedStarTreeDocument.dimensions[2], resultStarTreeDocument.dimensions[2]); + assertEquals(expectedStarTreeDocument.dimensions[3], resultStarTreeDocument.dimensions[3]); + assertEquals(expectedStarTreeDocument.metrics[0], resultStarTreeDocument.metrics[0]); + assertEquals(expectedStarTreeDocument.metrics[1], resultStarTreeDocument.metrics[1]); + assertEquals(expectedStarTreeDocument.metrics[2], resultStarTreeDocument.metrics[2]); + assertEquals(expectedStarTreeDocument.metrics[3], resultStarTreeDocument.metrics[3]); + assertEquals(expectedStarTreeDocument.metrics[4], resultStarTreeDocument.metrics[4]); + assertEquals(expectedStarTreeDocument.metrics[5], resultStarTreeDocument.metrics[5]); + + numOfAggregatedDocuments++; + } + + assertEquals(inorderStarTreeDocuments.size(), numOfAggregatedDocuments); + } + + public void test_sortAndAggregateStarTreeDocuments_nullMetric() throws IOException { + + int noOfStarTreeDocuments = 5; + StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; + + starTreeDocuments[0] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Object[] { 12.0, 10.0, randomDouble(), 8.0, 20.0 }); + starTreeDocuments[1] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Object[] { 10.0, 6.0, randomDouble(), 12.0, 10.0 }); + starTreeDocuments[2] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Object[] { 14.0, 12.0, randomDouble(), 6.0, 24.0 }); + starTreeDocuments[3] = new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Object[] { 9.0, 4.0, randomDouble(), 9.0, 12.0 }); + starTreeDocuments[4] = new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Object[] { 11.0, null, randomDouble(), 8.0, 13.0 }); + + List inorderStarTreeDocuments = List.of( + new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Object[] { 21.0, 14.0, 2L, 8.0, 20.0, 2L }), + new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Object[] { 35.0, 18.0, 3L, 6.0, 24.0, 3L }) + ); + Iterator expectedStarTreeDocumentIterator = inorderStarTreeDocuments.iterator(); + + StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; + for (int i = 0; i < noOfStarTreeDocuments; i++) { + long metric1 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]); + Long metric2 = starTreeDocuments[i].metrics[1] != null + ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[1]) + : null; + long metric3 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]); + long metric4 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[3]); + long metric5 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[4]); + segmentStarTreeDocuments[i] = new StarTreeDocument( + starTreeDocuments[i].dimensions, + new Object[] { metric1, metric2, metric3, metric4, metric5, null } + ); + } + SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); + List metricsIterators = getMetricIterators(segmentStarTreeDocuments); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( + dimsIterators, + metricsIterators + ); + + while (segmentStarTreeDocumentIterator.hasNext() && expectedStarTreeDocumentIterator.hasNext()) { + StarTreeDocument resultStarTreeDocument = segmentStarTreeDocumentIterator.next(); + StarTreeDocument expectedStarTreeDocument = expectedStarTreeDocumentIterator.next(); + assertEquals(expectedStarTreeDocument.dimensions[0], resultStarTreeDocument.dimensions[0]); + assertEquals(expectedStarTreeDocument.dimensions[1], resultStarTreeDocument.dimensions[1]); + assertEquals(expectedStarTreeDocument.dimensions[2], resultStarTreeDocument.dimensions[2]); + assertEquals(expectedStarTreeDocument.dimensions[3], resultStarTreeDocument.dimensions[3]); + assertEquals(expectedStarTreeDocument.metrics[0], resultStarTreeDocument.metrics[0]); + assertEquals(expectedStarTreeDocument.metrics[1], resultStarTreeDocument.metrics[1]); + assertEquals(expectedStarTreeDocument.metrics[2], resultStarTreeDocument.metrics[2]); + assertEquals(expectedStarTreeDocument.metrics[3], resultStarTreeDocument.metrics[3]); + assertEquals(expectedStarTreeDocument.metrics[4], resultStarTreeDocument.metrics[4]); + assertEquals(expectedStarTreeDocument.metrics[5], resultStarTreeDocument.metrics[5]); + } + } + + public void test_sortAndAggregateStarTreeDocuments_nullMetricField() throws IOException { + + int noOfStarTreeDocuments = 5; + StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; + // Setting second metric iterator as empty sorted numeric , indicating a metric field is null + starTreeDocuments[0] = new StarTreeDocument( + new Long[] { 2L, 4L, 3L, 4L }, + new Object[] { 12.0, null, randomDouble(), 8.0, 20.0, null } + ); + starTreeDocuments[1] = new StarTreeDocument( + new Long[] { 3L, 4L, 2L, 1L }, + new Object[] { 10.0, null, randomDouble(), 12.0, 10.0, null } + ); + starTreeDocuments[2] = new StarTreeDocument( + new Long[] { 3L, 4L, 2L, 1L }, + new Object[] { 14.0, null, randomDouble(), 6.0, 24.0, null } + ); + starTreeDocuments[3] = new StarTreeDocument( + new Long[] { 2L, 4L, 3L, 4L }, + new Object[] { 9.0, null, randomDouble(), 9.0, 12.0, 10L } + ); + starTreeDocuments[4] = new StarTreeDocument( + new Long[] { 3L, 4L, 2L, 1L }, + new Object[] { 11.0, null, randomDouble(), 8.0, 13.0, null } + ); + + List inorderStarTreeDocuments = List.of( + new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Object[] { 21.0, 0.0, 2L, 8.0, 20.0, 11L }), + new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Object[] { 35.0, 0.0, 3L, 6.0, 24.0, 3L }) + ); + Iterator expectedStarTreeDocumentIterator = inorderStarTreeDocuments.iterator(); + + StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; + for (int i = 0; i < noOfStarTreeDocuments; i++) { + long metric1 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]); + Long metric2 = starTreeDocuments[i].metrics[1] != null + ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[1]) + : null; + long metric3 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]); + long metric4 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[3]); + long metric5 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[4]); + Long metric6 = starTreeDocuments[i].metrics[5] != null ? (Long) starTreeDocuments[i].metrics[5] : null; + segmentStarTreeDocuments[i] = new StarTreeDocument( + starTreeDocuments[i].dimensions, + new Object[] { metric1, metric2, metric3, metric4, metric5, metric6 } + ); + } + SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); + List metricsIterators = getMetricIterators(segmentStarTreeDocuments); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( + dimsIterators, + metricsIterators + ); + + while (segmentStarTreeDocumentIterator.hasNext() && expectedStarTreeDocumentIterator.hasNext()) { + StarTreeDocument resultStarTreeDocument = segmentStarTreeDocumentIterator.next(); + StarTreeDocument expectedStarTreeDocument = expectedStarTreeDocumentIterator.next(); + assertEquals(expectedStarTreeDocument.dimensions[0], resultStarTreeDocument.dimensions[0]); + assertEquals(expectedStarTreeDocument.dimensions[1], resultStarTreeDocument.dimensions[1]); + assertEquals(expectedStarTreeDocument.dimensions[2], resultStarTreeDocument.dimensions[2]); + assertEquals(expectedStarTreeDocument.dimensions[3], resultStarTreeDocument.dimensions[3]); + assertEquals(expectedStarTreeDocument.metrics[0], resultStarTreeDocument.metrics[0]); + assertEquals(expectedStarTreeDocument.metrics[1], resultStarTreeDocument.metrics[1]); + assertEquals(expectedStarTreeDocument.metrics[2], resultStarTreeDocument.metrics[2]); + assertEquals(expectedStarTreeDocument.metrics[3], resultStarTreeDocument.metrics[3]); + assertEquals(expectedStarTreeDocument.metrics[4], resultStarTreeDocument.metrics[4]); + assertEquals(expectedStarTreeDocument.metrics[5], resultStarTreeDocument.metrics[5]); + } + } + + public void test_sortAndAggregateStarTreeDocuments_nullAndMinusOneInDimensionField() throws IOException { + int noOfStarTreeDocuments = 5; + StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; + // Setting second metric iterator as empty sorted numeric , indicating a metric field is null + starTreeDocuments[0] = new StarTreeDocument( + new Long[] { 2L, null, 3L, 4L }, + new Object[] { 12.0, null, randomDouble(), 8.0, 20.0, null } + ); + starTreeDocuments[1] = new StarTreeDocument( + new Long[] { null, 4L, 2L, 1L }, + new Object[] { 10.0, null, randomDouble(), 12.0, 10.0, null } + ); + starTreeDocuments[2] = new StarTreeDocument( + new Long[] { null, 4L, 2L, 1L }, + new Object[] { 14.0, null, randomDouble(), 6.0, 24.0, null } + ); + starTreeDocuments[3] = new StarTreeDocument( + new Long[] { 2L, null, 3L, 4L }, + new Object[] { 9.0, null, randomDouble(), 9.0, 12.0, 10L } + ); + starTreeDocuments[4] = new StarTreeDocument( + new Long[] { -1L, 4L, 2L, 1L }, + new Object[] { 11.0, null, randomDouble(), 8.0, 13.0, null } + ); + + List inorderStarTreeDocuments = List.of( + new StarTreeDocument(new Long[] { -1L, 4L, 2L, 1L }, new Object[] { 11.0, 0.0, 1L, 8.0, 13.0, 1L }), + new StarTreeDocument(new Long[] { 2L, null, 3L, 4L }, new Object[] { 21.0, 0.0, 2L, 8.0, 20.0, 11L }), + new StarTreeDocument(new Long[] { null, 4L, 2L, 1L }, new Object[] { 24.0, 0.0, 2L, 6.0, 24.0, 2L }) + ); + Iterator expectedStarTreeDocumentIterator = inorderStarTreeDocuments.iterator(); + + StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; + for (int i = 0; i < noOfStarTreeDocuments; i++) { + long metric1 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]); + Long metric2 = starTreeDocuments[i].metrics[1] != null + ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[1]) + : null; + long metric3 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]); + long metric4 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[3]); + long metric5 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[4]); + Long metric6 = starTreeDocuments[i].metrics[5] != null ? (long) starTreeDocuments[i].metrics[5] : null; + segmentStarTreeDocuments[i] = new StarTreeDocument( + starTreeDocuments[i].dimensions, + new Object[] { metric1, metric2, metric3, metric4, metric5, metric6 } + ); + } + SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); + List metricsIterators = getMetricIterators(segmentStarTreeDocuments); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( + dimsIterators, + metricsIterators + ); + + while (segmentStarTreeDocumentIterator.hasNext()) { + StarTreeDocument resultStarTreeDocument = segmentStarTreeDocumentIterator.next(); + StarTreeDocument expectedStarTreeDocument = expectedStarTreeDocumentIterator.next(); + assertEquals(expectedStarTreeDocument.dimensions[0], resultStarTreeDocument.dimensions[0]); + assertEquals(expectedStarTreeDocument.dimensions[1], resultStarTreeDocument.dimensions[1]); + assertEquals(expectedStarTreeDocument.dimensions[2], resultStarTreeDocument.dimensions[2]); + assertEquals(expectedStarTreeDocument.dimensions[3], resultStarTreeDocument.dimensions[3]); + assertEquals(expectedStarTreeDocument.metrics[0], resultStarTreeDocument.metrics[0]); + assertEquals(expectedStarTreeDocument.metrics[1], resultStarTreeDocument.metrics[1]); + assertEquals(expectedStarTreeDocument.metrics[2], resultStarTreeDocument.metrics[2]); + assertEquals(expectedStarTreeDocument.metrics[3], resultStarTreeDocument.metrics[3]); + assertEquals(expectedStarTreeDocument.metrics[4], resultStarTreeDocument.metrics[4]); + assertEquals(expectedStarTreeDocument.metrics[5], resultStarTreeDocument.metrics[5]); + } + + assertFalse(expectedStarTreeDocumentIterator.hasNext()); + + builder.build(segmentStarTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); + validateStarTree(builder.getRootNode(), 4, 1, builder.getStarTreeDocuments()); + } + + public void test_sortAndAggregateStarTreeDocuments_nullDimensionsAndNullMetrics() throws IOException { + int noOfStarTreeDocuments = 5; + StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; + // Setting second metric iterator as empty sorted numeric , indicating a metric field is null + starTreeDocuments[0] = new StarTreeDocument( + new Long[] { null, null, null, null }, + new Object[] { null, null, null, null, null, null } + ); + starTreeDocuments[1] = new StarTreeDocument( + new Long[] { null, null, null, null }, + new Object[] { null, null, null, null, null, null } + ); + starTreeDocuments[2] = new StarTreeDocument( + new Long[] { null, null, null, null }, + new Object[] { null, null, null, null, null, null } + ); + starTreeDocuments[3] = new StarTreeDocument( + new Long[] { null, null, null, null }, + new Object[] { null, null, null, null, null, null } + ); + starTreeDocuments[4] = new StarTreeDocument( + new Long[] { null, null, null, null }, + new Object[] { null, null, null, null, null, null } + ); + + List inorderStarTreeDocuments = List.of( + new StarTreeDocument(new Long[] { null, null, null, null }, new Object[] { 0.0, 0.0, 0L, null, null, 5L }) + ); + Iterator expectedStarTreeDocumentIterator = inorderStarTreeDocuments.iterator(); + + StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; + for (int i = 0; i < noOfStarTreeDocuments; i++) { + Long metric1 = starTreeDocuments[i].metrics[0] != null + ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]) + : null; + Long metric2 = starTreeDocuments[i].metrics[1] != null + ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[1]) + : null; + Long metric3 = starTreeDocuments[i].metrics[2] != null + ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]) + : null; + Long metric4 = starTreeDocuments[i].metrics[3] != null + ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[3]) + : null; + Long metric5 = starTreeDocuments[i].metrics[4] != null + ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[4]) + : null; + segmentStarTreeDocuments[i] = new StarTreeDocument( + starTreeDocuments[i].dimensions, + new Object[] { metric1, metric2, metric3, metric4, metric5, null } + ); + } + SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); + List metricsIterators = getMetricIterators(segmentStarTreeDocuments); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( + dimsIterators, + metricsIterators + ); + + while (segmentStarTreeDocumentIterator.hasNext() && expectedStarTreeDocumentIterator.hasNext()) { + StarTreeDocument resultStarTreeDocument = segmentStarTreeDocumentIterator.next(); + StarTreeDocument expectedStarTreeDocument = expectedStarTreeDocumentIterator.next(); + assertEquals(expectedStarTreeDocument.dimensions[0], resultStarTreeDocument.dimensions[0]); + assertEquals(expectedStarTreeDocument.dimensions[1], resultStarTreeDocument.dimensions[1]); + assertEquals(expectedStarTreeDocument.dimensions[2], resultStarTreeDocument.dimensions[2]); + assertEquals(expectedStarTreeDocument.dimensions[3], resultStarTreeDocument.dimensions[3]); + assertEquals(expectedStarTreeDocument.metrics[0], resultStarTreeDocument.metrics[0]); + assertEquals(expectedStarTreeDocument.metrics[1], resultStarTreeDocument.metrics[1]); + assertEquals(expectedStarTreeDocument.metrics[2], resultStarTreeDocument.metrics[2]); + assertEquals(expectedStarTreeDocument.metrics[3], resultStarTreeDocument.metrics[3]); + assertEquals(expectedStarTreeDocument.metrics[4], resultStarTreeDocument.metrics[4]); + assertEquals(expectedStarTreeDocument.metrics[5], resultStarTreeDocument.metrics[5]); + } + builder.build(segmentStarTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); + validateStarTree(builder.getRootNode(), 4, 1, builder.getStarTreeDocuments()); + } + + public void test_sortAndAggregateStarTreeDocuments_nullDimensionsAndFewNullMetrics() throws IOException { + int noOfStarTreeDocuments = 5; + StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; + + double sumValue = randomDouble(); + double minValue = randomDouble(); + double maxValue = randomDouble(); + + // Setting second metric iterator as empty sorted numeric , indicating a metric field is null + starTreeDocuments[0] = new StarTreeDocument( + new Long[] { null, null, null, null }, + new Object[] { null, null, randomDouble(), null, maxValue } + ); + starTreeDocuments[1] = new StarTreeDocument(new Long[] { null, null, null, null }, new Object[] { null, null, null, null, null }); + starTreeDocuments[2] = new StarTreeDocument( + new Long[] { null, null, null, null }, + new Object[] { null, null, null, minValue, null } + ); + starTreeDocuments[3] = new StarTreeDocument(new Long[] { null, null, null, null }, new Object[] { null, null, null, null, null }); + starTreeDocuments[4] = new StarTreeDocument( + new Long[] { null, null, null, null }, + new Object[] { sumValue, null, randomDouble(), null, null } + ); + + List inorderStarTreeDocuments = List.of( + new StarTreeDocument(new Long[] { null, null, null, null }, new Object[] { sumValue, 0.0, 2L, minValue, maxValue, 5L }) + ); + Iterator expectedStarTreeDocumentIterator = inorderStarTreeDocuments.iterator(); + + StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; + for (int i = 0; i < noOfStarTreeDocuments; i++) { + Long metric1 = starTreeDocuments[i].metrics[0] != null + ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]) + : null; + Long metric2 = starTreeDocuments[i].metrics[1] != null + ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[1]) + : null; + Long metric3 = starTreeDocuments[i].metrics[2] != null + ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]) + : null; + Long metric4 = starTreeDocuments[i].metrics[3] != null + ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[3]) + : null; + Long metric5 = starTreeDocuments[i].metrics[4] != null + ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[4]) + : null; + segmentStarTreeDocuments[i] = new StarTreeDocument( + starTreeDocuments[i].dimensions, + new Object[] { metric1, metric2, metric3, metric4, metric5, null } + ); + } + SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); + List metricsIterators = getMetricIterators(segmentStarTreeDocuments); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( + dimsIterators, + metricsIterators + ); + + while (segmentStarTreeDocumentIterator.hasNext() && expectedStarTreeDocumentIterator.hasNext()) { + StarTreeDocument resultStarTreeDocument = segmentStarTreeDocumentIterator.next(); + StarTreeDocument expectedStarTreeDocument = expectedStarTreeDocumentIterator.next(); + assertEquals(expectedStarTreeDocument.dimensions[0], resultStarTreeDocument.dimensions[0]); + assertEquals(expectedStarTreeDocument.dimensions[1], resultStarTreeDocument.dimensions[1]); + assertEquals(expectedStarTreeDocument.dimensions[2], resultStarTreeDocument.dimensions[2]); + assertEquals(expectedStarTreeDocument.dimensions[3], resultStarTreeDocument.dimensions[3]); + assertEquals(expectedStarTreeDocument.metrics[0], resultStarTreeDocument.metrics[0]); + assertEquals(expectedStarTreeDocument.metrics[1], resultStarTreeDocument.metrics[1]); + assertEquals(expectedStarTreeDocument.metrics[2], resultStarTreeDocument.metrics[2]); + assertEquals(expectedStarTreeDocument.metrics[3], resultStarTreeDocument.metrics[3]); + assertEquals(expectedStarTreeDocument.metrics[4], resultStarTreeDocument.metrics[4]); + assertEquals(expectedStarTreeDocument.metrics[5], resultStarTreeDocument.metrics[5]); + } + builder.build(segmentStarTreeDocumentIterator, new AtomicInteger(), docValuesConsumer); + validateStarTree(builder.getRootNode(), 4, 1, builder.getStarTreeDocuments()); + } + + public void test_sortAndAggregateStarTreeDocuments_emptyDimensions() throws IOException { + + int noOfStarTreeDocuments = 5; + StarTreeDocument[] starTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; + // Setting second metric iterator as empty sorted numeric , indicating a metric field is null + starTreeDocuments[0] = new StarTreeDocument( + new Long[] { null, null, null, null }, + new Object[] { 12.0, null, randomDouble(), 8.0, 20.0, 10L } + ); + starTreeDocuments[1] = new StarTreeDocument( + new Long[] { null, null, null, null }, + new Object[] { 10.0, null, randomDouble(), 12.0, 10.0, 10L } + ); + starTreeDocuments[2] = new StarTreeDocument( + new Long[] { null, null, null, null }, + new Object[] { 14.0, null, randomDouble(), 6.0, 24.0, 10L } + ); + starTreeDocuments[3] = new StarTreeDocument( + new Long[] { null, null, null, null }, + new Object[] { 9.0, null, randomDouble(), 9.0, 12.0, 10L } + ); + starTreeDocuments[4] = new StarTreeDocument( + new Long[] { null, null, null, null }, + new Object[] { 11.0, null, randomDouble(), 8.0, 13.0, 10L } + ); + + List inorderStarTreeDocuments = List.of( + new StarTreeDocument(new Long[] { null, null, null, null }, new Object[] { 56.0, 0.0, 5L, 6.0, 24.0, 50L }) + ); + Iterator expectedStarTreeDocumentIterator = inorderStarTreeDocuments.iterator(); + + StarTreeDocument[] segmentStarTreeDocuments = new StarTreeDocument[noOfStarTreeDocuments]; + for (int i = 0; i < noOfStarTreeDocuments; i++) { + Long metric1 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[0]); + Long metric2 = starTreeDocuments[i].metrics[1] != null + ? NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[1]) + : null; + Long metric3 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[2]); + Long metric4 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[3]); + Long metric5 = NumericUtils.doubleToSortableLong((Double) starTreeDocuments[i].metrics[4]); + Long metric6 = (Long) starTreeDocuments[i].metrics[5]; + segmentStarTreeDocuments[i] = new StarTreeDocument( + starTreeDocuments[i].dimensions, + new Object[] { metric1, metric2, metric3, metric4, metric5, metric6 } + ); + } + SequentialDocValuesIterator[] dimsIterators = getDimensionIterators(segmentStarTreeDocuments); + List metricsIterators = getMetricIterators(segmentStarTreeDocuments); + builder = getStarTreeBuilder(metaOut, dataOut, compositeField, writeState, mapperService); + Iterator segmentStarTreeDocumentIterator = builder.sortAndAggregateSegmentDocuments( + dimsIterators, + metricsIterators + ); + + while (segmentStarTreeDocumentIterator.hasNext() && expectedStarTreeDocumentIterator.hasNext()) { + StarTreeDocument resultStarTreeDocument = segmentStarTreeDocumentIterator.next(); + StarTreeDocument expectedStarTreeDocument = expectedStarTreeDocumentIterator.next(); + assertEquals(expectedStarTreeDocument.dimensions[0], resultStarTreeDocument.dimensions[0]); + assertEquals(expectedStarTreeDocument.dimensions[1], resultStarTreeDocument.dimensions[1]); + assertEquals(expectedStarTreeDocument.dimensions[2], resultStarTreeDocument.dimensions[2]); + assertEquals(expectedStarTreeDocument.dimensions[3], resultStarTreeDocument.dimensions[3]); + assertEquals(expectedStarTreeDocument.metrics[0], resultStarTreeDocument.metrics[0]); + assertEquals(expectedStarTreeDocument.metrics[1], resultStarTreeDocument.metrics[1]); + assertEquals(expectedStarTreeDocument.metrics[2], resultStarTreeDocument.metrics[2]); + assertEquals(expectedStarTreeDocument.metrics[3], resultStarTreeDocument.metrics[3]); + assertEquals(expectedStarTreeDocument.metrics[4], resultStarTreeDocument.metrics[4]); + assertEquals(expectedStarTreeDocument.metrics[5], resultStarTreeDocument.metrics[5]); + } + } +} diff --git a/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuilderTestCase.java b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuilderTestCase.java new file mode 100644 index 0000000000000..4c854f7546197 --- /dev/null +++ b/server/src/test/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeBuilderTestCase.java @@ -0,0 +1,368 @@ +/* + * 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.startree.builder; + +import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; + +import org.apache.lucene.codecs.DocValuesConsumer; +import org.apache.lucene.codecs.DocValuesProducer; +import org.apache.lucene.index.DocValuesType; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.IndexFileNames; +import org.apache.lucene.index.IndexOptions; +import org.apache.lucene.index.SegmentReadState; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.index.VectorEncoding; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.IndexInput; +import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.util.NumericUtils; +import org.opensearch.common.Rounding; +import org.opensearch.common.settings.Settings; +import org.opensearch.index.codec.composite.composite99.Composite99DocValuesFormat; +import org.opensearch.index.compositeindex.CompositeIndexConstants; +import org.opensearch.index.compositeindex.datacube.DataCubeDateTimeUnit; +import org.opensearch.index.compositeindex.datacube.DateDimension; +import org.opensearch.index.compositeindex.datacube.Dimension; +import org.opensearch.index.compositeindex.datacube.Metric; +import org.opensearch.index.compositeindex.datacube.MetricStat; +import org.opensearch.index.compositeindex.datacube.NumericDimension; +import org.opensearch.index.compositeindex.datacube.startree.StarTreeDocument; +import org.opensearch.index.compositeindex.datacube.startree.StarTreeField; +import org.opensearch.index.compositeindex.datacube.startree.StarTreeFieldConfiguration; +import org.opensearch.index.compositeindex.datacube.startree.fileformats.meta.StarTreeMetadata; +import org.opensearch.index.compositeindex.datacube.startree.node.InMemoryTreeNode; +import org.opensearch.index.compositeindex.datacube.startree.utils.date.DateTimeUnitAdapter; +import org.opensearch.index.compositeindex.datacube.startree.utils.date.DateTimeUnitRounding; +import org.opensearch.index.mapper.ContentPath; +import org.opensearch.index.mapper.DateFieldMapper; +import org.opensearch.index.mapper.DocumentMapper; +import org.opensearch.index.mapper.Mapper; +import org.opensearch.index.mapper.MapperService; +import org.opensearch.index.mapper.MappingLookup; +import org.opensearch.index.mapper.NumberFieldMapper; +import org.opensearch.test.OpenSearchTestCase; +import org.junit.Before; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collection; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.UUID; + +import static org.opensearch.index.compositeindex.datacube.startree.fileformats.StarTreeWriter.VERSION_CURRENT; +import static org.opensearch.index.mapper.CompositeMappedFieldType.CompositeFieldType.STAR_TREE; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public abstract class StarTreeBuilderTestCase extends OpenSearchTestCase { + private final StarTreeFieldConfiguration.StarTreeBuildMode buildMode; + protected MapperService mapperService; + protected List fields = List.of(); + protected List metrics; + protected Directory directory; + protected FieldInfo[] fieldsInfo; + protected StarTreeField compositeField; + protected Map fieldProducerMap; + protected SegmentWriteState writeState; + protected BaseStarTreeBuilder builder; + protected IndexOutput dataOut; + protected IndexOutput metaOut; + protected DocValuesConsumer docValuesConsumer; + protected String dataFileName; + protected String metaFileName; + protected List dimensionsOrder; + + public StarTreeBuilderTestCase(StarTreeFieldConfiguration.StarTreeBuildMode buildMode) { + this.buildMode = buildMode; + } + + @ParametersFactory + public static Collection parameters() { + return Arrays.asList( + new Object[] { StarTreeFieldConfiguration.StarTreeBuildMode.ON_HEAP }, + new Object[] { StarTreeFieldConfiguration.StarTreeBuildMode.OFF_HEAP } + ); + } + + @Before + public void setup() throws IOException { + fields = List.of("field1", "field2", "field3", "field4", "field5", "field6", "field7", "field8", "field9", "field10"); + + dimensionsOrder = List.of( + new NumericDimension("field1"), + new NumericDimension("field3"), + new NumericDimension("field5"), + new NumericDimension("field8") + ); + metrics = List.of( + new Metric("field2", List.of(MetricStat.SUM)), + new Metric("field4", List.of(MetricStat.SUM)), + new Metric("field6", List.of(MetricStat.VALUE_COUNT)), + new Metric("field9", List.of(MetricStat.MIN)), + new Metric("field10", List.of(MetricStat.MAX)), + new Metric("_doc_count", List.of(MetricStat.DOC_COUNT)) + ); + + DocValuesProducer docValuesProducer = mock(DocValuesProducer.class); + + compositeField = new StarTreeField( + "test", + dimensionsOrder, + metrics, + new StarTreeFieldConfiguration(1, Set.of("field8"), getBuildMode()) + ); + directory = newFSDirectory(createTempDir()); + + fieldsInfo = new FieldInfo[fields.size()]; + fieldProducerMap = new HashMap<>(); + for (int i = 0; i < fieldsInfo.length; i++) { + fieldsInfo[i] = new FieldInfo( + fields.get(i), + i, + false, + false, + true, + IndexOptions.DOCS_AND_FREQS_AND_POSITIONS_AND_OFFSETS, + DocValuesType.SORTED_NUMERIC, + -1, + Collections.emptyMap(), + 0, + 0, + 0, + 0, + VectorEncoding.FLOAT32, + VectorSimilarityFunction.EUCLIDEAN, + false, + false + ); + fieldProducerMap.put(fields.get(i), docValuesProducer); + } + writeState = getWriteState(5, UUID.randomUUID().toString().substring(0, 16).getBytes(StandardCharsets.UTF_8)); + + dataFileName = IndexFileNames.segmentFileName( + writeState.segmentInfo.name, + writeState.segmentSuffix, + Composite99DocValuesFormat.DATA_EXTENSION + ); + dataOut = writeState.directory.createOutput(dataFileName, writeState.context); + + metaFileName = IndexFileNames.segmentFileName( + writeState.segmentInfo.name, + writeState.segmentSuffix, + Composite99DocValuesFormat.META_EXTENSION + ); + metaOut = writeState.directory.createOutput(metaFileName, writeState.context); + + mapperService = mock(MapperService.class); + 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("field2", NumberFieldMapper.NumberType.DOUBLE, false, true) + .build(new Mapper.BuilderContext(settings, new ContentPath())); + NumberFieldMapper numberFieldMapper2 = new NumberFieldMapper.Builder("field4", NumberFieldMapper.NumberType.DOUBLE, false, true) + .build(new Mapper.BuilderContext(settings, new ContentPath())); + NumberFieldMapper numberFieldMapper3 = new NumberFieldMapper.Builder("field6", NumberFieldMapper.NumberType.DOUBLE, false, true) + .build(new Mapper.BuilderContext(settings, new ContentPath())); + NumberFieldMapper numberFieldMapper4 = new NumberFieldMapper.Builder("field9", NumberFieldMapper.NumberType.DOUBLE, false, true) + .build(new Mapper.BuilderContext(settings, new ContentPath())); + NumberFieldMapper numberFieldMapper5 = new NumberFieldMapper.Builder("field10", NumberFieldMapper.NumberType.DOUBLE, false, true) + .build(new Mapper.BuilderContext(settings, new ContentPath())); + MappingLookup fieldMappers = new MappingLookup( + Set.of(numberFieldMapper1, numberFieldMapper2, numberFieldMapper3, numberFieldMapper4, numberFieldMapper5), + Collections.emptyList(), + Collections.emptyList(), + 0, + null + ); + when(documentMapper.mappers()).thenReturn(fieldMappers); + docValuesConsumer = mock(DocValuesConsumer.class); + } + + protected BaseStarTreeBuilder getStarTreeBuilder( + IndexOutput metaOut, + IndexOutput dataOut, + StarTreeField starTreeField, + SegmentWriteState segmentWriteState, + MapperService mapperService + ) throws IOException { + switch (buildMode) { + case ON_HEAP: + return new OnHeapStarTreeBuilder(metaOut, dataOut, starTreeField, segmentWriteState, mapperService); + case OFF_HEAP: + return new OffHeapStarTreeBuilder(metaOut, dataOut, starTreeField, segmentWriteState, mapperService); + default: + throw new IllegalArgumentException("Invalid build mode: " + buildMode); + } + } + + protected StarTreeFieldConfiguration.StarTreeBuildMode getBuildMode() { + return buildMode; + } + + protected void validateStarTreeFileFormats( + InMemoryTreeNode rootNode, + int numDocs, + StarTreeMetadata expectedStarTreeMetadata, + List expectedStarTreeDocuments + ) throws IOException { + BuilderTestsUtils.validateStarTreeFileFormats( + rootNode, + numDocs, + expectedStarTreeMetadata, + expectedStarTreeDocuments, + dataFileName, + metaFileName, + builder, + compositeField, + writeState, + directory + ); + + } + + SegmentWriteState getWriteState(int numDocs, byte[] id) { + return BuilderTestsUtils.getWriteState(numDocs, id, fieldsInfo, directory); + } + + SegmentReadState getReadState(int numDocs, List dimensionFields, List metrics) { + return BuilderTestsUtils.getReadState(numDocs, dimensionFields, metrics, compositeField, writeState, directory); + } + + protected Map getAttributes(int numSegmentDocs) { + return Map.of(CompositeIndexConstants.SEGMENT_DOCS_COUNT, String.valueOf(numSegmentDocs)); + } + + protected List getStarTreeDimensionNames(List dimensionsOrder) { + List dimensionNames = new ArrayList<>(); + for (Dimension dimension : dimensionsOrder) { + dimensionNames.addAll(dimension.getSubDimensionNames()); + } + return dimensionNames; + } + + protected StarTreeField getStarTreeField(MetricStat count) { + Dimension d1 = new NumericDimension("field1"); + Dimension d2 = new NumericDimension("field3"); + Metric m1 = new Metric("field2", List.of(count)); + List dims = List.of(d1, d2); + List metrics = List.of(m1); + StarTreeFieldConfiguration c = new StarTreeFieldConfiguration(1000, new HashSet<>(), getBuildMode()); + return new StarTreeField("sf", dims, metrics, c); + } + + protected StarTreeField getStarTreeFieldWithDocCount(int maxLeafDocs, boolean includeDocCountMetric) { + Dimension d1 = new NumericDimension("field1"); + Dimension d2 = new NumericDimension("field3"); + Dimension d3 = new NumericDimension("field5"); + Dimension d4 = new NumericDimension("field8"); + List dims = List.of(d1, d2, d3, d4); + Metric m1 = new Metric("field2", List.of(MetricStat.SUM)); + Metric m2 = null; + if (includeDocCountMetric) { + m2 = new Metric("_doc_count", List.of(MetricStat.DOC_COUNT)); + } + List metrics = m2 == null ? List.of(m1) : List.of(m1, m2); + StarTreeFieldConfiguration c = new StarTreeFieldConfiguration(maxLeafDocs, new HashSet<>(), getBuildMode()); + StarTreeField sf = new StarTreeField("sf", dims, metrics, c); + return sf; + } + + protected void assertStarTreeDocuments( + List resultStarTreeDocuments, + Iterator expectedStarTreeDocumentIterator + ) { + Iterator resultStarTreeDocumentIterator = resultStarTreeDocuments.iterator(); + while (resultStarTreeDocumentIterator.hasNext() && expectedStarTreeDocumentIterator.hasNext()) { + StarTreeDocument resultStarTreeDocument = resultStarTreeDocumentIterator.next(); + StarTreeDocument expectedStarTreeDocument = expectedStarTreeDocumentIterator.next(); + + assertEquals(expectedStarTreeDocument.dimensions[0], resultStarTreeDocument.dimensions[0]); + assertEquals(expectedStarTreeDocument.dimensions[1], resultStarTreeDocument.dimensions[1]); + assertEquals(expectedStarTreeDocument.dimensions[2], resultStarTreeDocument.dimensions[2]); + assertEquals(expectedStarTreeDocument.dimensions[3], resultStarTreeDocument.dimensions[3]); + assertEquals(expectedStarTreeDocument.metrics[0], resultStarTreeDocument.metrics[0]); + assertEquals(expectedStarTreeDocument.metrics[1], resultStarTreeDocument.metrics[1]); + assertEquals(expectedStarTreeDocument.metrics[2], resultStarTreeDocument.metrics[2]); + assertEquals(expectedStarTreeDocument.metrics[3], resultStarTreeDocument.metrics[3]); + assertEquals(expectedStarTreeDocument.metrics[4], resultStarTreeDocument.metrics[4]); + } + } + + protected static List getExpectedStarTreeDocumentIterator() { + return List.of( + new StarTreeDocument(new Long[] { 2L, 4L, 3L, 4L }, new Object[] { 21.0, 14.0, 2L, 8.0, 20.0, 2L }), + new StarTreeDocument(new Long[] { 3L, 4L, 2L, 1L }, new Object[] { 35.0, 34.0, 3L, 6.0, 24.0, 3L }), + new StarTreeDocument(new Long[] { null, 4L, 2L, 1L }, new Object[] { 35.0, 34.0, 3L, 6.0, 24.0, 3L }), + new StarTreeDocument(new Long[] { null, 4L, 3L, 4L }, new Object[] { 21.0, 14.0, 2L, 8.0, 20.0, 2L }), + new StarTreeDocument(new Long[] { null, 4L, null, 1L }, new Object[] { 35.0, 34.0, 3L, 6.0, 24.0, 3L }), + new StarTreeDocument(new Long[] { null, 4L, null, 4L }, new Object[] { 21.0, 14.0, 2L, 8.0, 20.0, 2L }), + new StarTreeDocument(new Long[] { null, 4L, null, null }, new Object[] { 56.0, 48.0, 5L, 6.0, 24.0, 5L }) + ); + } + + protected long getLongFromDouble(double value) { + return NumericUtils.doubleToSortableLong(value); + } + + protected StarTreeMetadata getStarTreeMetadata(List fields, int segmentAggregatedDocCount, int maxLeafDocs, int dataLength) { + return new StarTreeMetadata( + "sf", + STAR_TREE, + mock(IndexInput.class), + VERSION_CURRENT, + builder.numStarTreeNodes, + fields, + compositeField.getMetrics(), + segmentAggregatedDocCount, + builder.numStarTreeDocs, + maxLeafDocs, + Set.of(), + getBuildMode(), + 0, + dataLength + ); + } + + protected StarTreeField getStarTreeFieldWithDateDimension() { + List intervals = new ArrayList<>(); + intervals.add(new DateTimeUnitAdapter(Rounding.DateTimeUnit.MINUTES_OF_HOUR)); + intervals.add(new DateTimeUnitAdapter(Rounding.DateTimeUnit.HOUR_OF_DAY)); + intervals.add(DataCubeDateTimeUnit.HALF_HOUR_OF_DAY); + Dimension d1 = new DateDimension("field1", intervals, DateFieldMapper.Resolution.MILLISECONDS); + Dimension d2 = new NumericDimension("field3"); + Metric m1 = new Metric("field2", List.of(MetricStat.VALUE_COUNT, MetricStat.SUM)); + List dims = List.of(d1, d2); + List metrics = List.of(m1); + StarTreeFieldConfiguration c = new StarTreeFieldConfiguration(10, new HashSet<>(), getBuildMode()); + StarTreeField sf = new StarTreeField("sf", dims, metrics, c); + return sf; + } + + @Override + public void tearDown() throws Exception { + super.tearDown(); + if (builder != null) { + builder.close(); + } + docValuesConsumer.close(); + metaOut.close(); + dataOut.close(); + directory.close(); + } +} From 5279d21f8c85e78a91159af6e17e52789355f2a4 Mon Sep 17 00:00:00 2001 From: David Zane <38449481+dzane17@users.noreply.github.com> Date: Mon, 7 Oct 2024 21:59:54 -0700 Subject: [PATCH 7/7] Change successfulSearchShardIndices to Set (#16110) * Change successfulSearchShardIndices to Set Signed-off-by: David Zane * Update CHANGELOG.md Signed-off-by: Ankit Jain --------- Signed-off-by: David Zane Signed-off-by: Ankit Jain Co-authored-by: Ankit Jain --- CHANGELOG.md | 2 +- .../action/search/AbstractSearchAsyncAction.java | 4 +++- .../opensearch/action/search/SearchRequestContext.java | 9 +++++---- 3 files changed, 9 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 06ac7740eea6e..d8732d1ffa5cc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,7 +11,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - [Workload Management] Add orchestrator for wlm resiliency (QueryGroupService) ([#15925](https://github.com/opensearch-project/OpenSearch/pull/15925)) - Fallback to Remote cluster-state on Term-Version check mismatch - ([#15424](https://github.com/opensearch-project/OpenSearch/pull/15424)) - Implement WithFieldName interface in ValuesSourceAggregationBuilder & FieldSortBuilder ([#15916](https://github.com/opensearch-project/OpenSearch/pull/15916)) -- Add successfulSearchShardIndices in searchRequestContext ([#15967](https://github.com/opensearch-project/OpenSearch/pull/15967)) +- Add successfulSearchShardIndices in searchRequestContext ([#15967](https://github.com/opensearch-project/OpenSearch/pull/15967), [#16110](https://github.com/opensearch-project/OpenSearch/pull/16110)) - [Tiered Caching] Segmented cache changes ([#16047](https://github.com/opensearch-project/OpenSearch/pull/16047)) - Add support for msearch API to pass search pipeline name - ([#15923](https://github.com/opensearch-project/OpenSearch/pull/15923)) - Add _list/indices API as paginated alternate to _cat/indices ([#14718](https://github.com/opensearch-project/OpenSearch/pull/14718)) diff --git a/server/src/main/java/org/opensearch/action/search/AbstractSearchAsyncAction.java b/server/src/main/java/org/opensearch/action/search/AbstractSearchAsyncAction.java index 836083f91b995..85ea34e442c8f 100644 --- a/server/src/main/java/org/opensearch/action/search/AbstractSearchAsyncAction.java +++ b/server/src/main/java/org/opensearch/action/search/AbstractSearchAsyncAction.java @@ -755,7 +755,9 @@ public void sendSearchResponse(InternalSearchResponse internalSearchResponse, At searchRequestContext.setTotalHits(internalSearchResponse.hits().getTotalHits()); searchRequestContext.setShardStats(results.getNumShards(), successfulOps.get(), skippedOps.get(), failures.length); searchRequestContext.setSuccessfulSearchShardIndices( - results.getSuccessfulResults().map(result -> result.getSearchShardTarget().getIndex()).collect(Collectors.toSet()) + results.getSuccessfulResults() + .map(result -> result.getSearchShardTarget().getShardId().getIndex()) + .collect(Collectors.toSet()) ); onPhaseEnd(searchRequestContext); onRequestEnd(searchRequestContext); diff --git a/server/src/main/java/org/opensearch/action/search/SearchRequestContext.java b/server/src/main/java/org/opensearch/action/search/SearchRequestContext.java index 376cf71448d5c..398896734280e 100644 --- a/server/src/main/java/org/opensearch/action/search/SearchRequestContext.java +++ b/server/src/main/java/org/opensearch/action/search/SearchRequestContext.java @@ -12,6 +12,7 @@ import org.apache.logging.log4j.Logger; import org.apache.lucene.search.TotalHits; import org.opensearch.common.annotation.InternalApi; +import org.opensearch.core.index.Index; import org.opensearch.core.tasks.resourcetracker.TaskResourceInfo; import java.util.ArrayList; @@ -37,7 +38,7 @@ public class SearchRequestContext { private final Map phaseTookMap; private TotalHits totalHits; private final EnumMap shardStats; - private Set successfulSearchShardIndices; + private Set successfulSearchShardIndices; private final SearchRequest searchRequest; private final LinkedBlockingQueue phaseResourceUsage; @@ -144,15 +145,15 @@ public SearchRequest getRequest() { return searchRequest; } - void setSuccessfulSearchShardIndices(Set successfulSearchShardIndices) { + void setSuccessfulSearchShardIndices(Set successfulSearchShardIndices) { this.successfulSearchShardIndices = successfulSearchShardIndices; } /** - * @return A {@link List} of {@link String} representing the names of the indices that were + * @return A {@link Set} of {@link Index} representing the names of the indices that were * successfully queried at the shard level. */ - public Set getSuccessfulSearchShardIndices() { + public Set getSuccessfulSearchShardIndices() { return successfulSearchShardIndices; } }