From 701cc1ef4e447ba6a100d8a27d34bf88a26b5f48 Mon Sep 17 00:00:00 2001 From: Bolek Ziobrowski <26925920+bziobrowski@users.noreply.github.com> Date: Tue, 21 Jan 2025 10:21:17 +0100 Subject: [PATCH] Fix flaky tests (#14345) --- pinot-common/pom.xml | 4 +- ...DropwizardBrokerPrometheusMetricsTest.java | 20 +- ...wizardControllerPrometheusMetricsTest.java | 20 +- ...DropwizardMinionPrometheusMetricsTest.java | 20 +- ...DropwizardServerPrometheusMetricsTest.java | 20 +- .../utils/RoundRobinURIProviderTest.java | 238 ++++----- .../FixedByteSingleValueMultiColWriter.java | 2 +- .../forward/BaseChunkForwardIndexReader.java | 2 +- .../readers/text/NativeTextIndexReader.java | 10 + .../segment/store/FilePerIndexDirectory.java | 1 + .../utils/nativefst/NativeFSTIndexReader.java | 11 + .../PinotBuffersAfterClassCheckRule.java | 33 ++ .../PinotBuffersAfterMethodCheckRule.java | 70 +++ .../mutable/IndexingFailureTest.java | 9 +- .../mutable/MutableSegmentDedupeTest.java | 47 +- ...tableSegmentEntriesAboveThresholdTest.java | 83 +-- ...MutableSegmentImplNullValueVectorTest.java | 9 +- .../mutable/MutableSegmentImplRawMVTest.java | 5 +- .../mutable/MutableSegmentImplTest.java | 2 + .../mutable/MutableSegmentImplTestUtils.java | 18 +- ...bleSegmentImplUpsertComparisonColTest.java | 79 ++- .../mutable/MutableSegmentImplUpsertTest.java | 80 +-- .../io/reader/impl/FixedBitIntReaderTest.java | 3 +- .../local/io/util/PinotDataBitSetTest.java | 3 +- .../local/io/util/PinotDataBitSetV2Test.java | 488 +++++++++--------- .../util/VarLengthValueReaderWriterTest.java | 3 +- .../MmapMemoryManagerFileCleanupTest.java | 3 +- .../io/writer/impl/MmapMemoryManagerTest.java | 199 ++++--- .../MutableOffHeapByteArrayStoreTest.java | 121 ++--- .../VarByteChunkSVForwardIndexWriterTest.java | 3 +- .../RealtimeSegmentConverterTest.java | 445 ++++++++-------- .../dictionary/MultiValueDictionaryTest.java | 49 +- .../dictionary/MutableDictionaryTest.java | 3 +- .../OffHeapMutableBytesStoreTest.java | 11 +- .../BitmapInvertedIndexCreatorTest.java | 8 +- .../creator/BitmapInvertedIndexTest.java | 3 +- .../segment/creator/DictionariesTest.java | 161 +++--- .../creator/DictionaryOptimiserTest.java | 44 +- .../DictionaryOptimizerCardinalityTest.java | 46 +- .../creator/OnHeapDictionariesTest.java | 5 +- .../segment/creator/SegmentTestUtils.java | 8 +- .../impl/SegmentColumnarIndexCreatorTest.java | 10 +- .../local/segment/index/H3IndexTest.java | 3 +- .../local/segment/index/JsonIndexTest.java | 3 +- .../creator/BitSlicedIndexCreatorTest.java | 3 +- .../index/creator/BloomFilterCreatorTest.java | 3 +- .../creator/CLPForwardIndexCreatorTest.java | 27 +- .../creator/CLPForwardIndexCreatorV2Test.java | 89 ++-- .../creator/LuceneFSTIndexCreatorTest.java | 21 +- ...ultiValueFixedByteRawIndexCreatorTest.java | 41 +- .../MultiValueVarByteRawIndexCreatorTest.java | 3 +- .../creator/NativeFSTIndexCreatorTest.java | 20 +- .../creator/NativeTextIndexCreatorTest.java | 38 +- .../index/creator/RangeIndexCreatorTest.java | 3 +- .../index/creator/RawIndexCreatorTest.java | 6 +- .../SegmentGenerationWithBytesTypeTest.java | 50 +- ...egmentGenerationWithFilterRecordsTest.java | 14 +- .../SegmentGenerationWithMinMaxTest.java | 3 +- .../SegmentGenerationWithNoRecordsTest.java | 8 +- .../SegmentGenerationWithTimeColumnTest.java | 3 +- .../index/creator/SegmentPartitionTest.java | 4 +- .../index/creator/VarByteChunkV4Test.java | 10 +- .../index/creator/VarByteChunkV5Test.java | 3 +- .../inv/BitmapInvertedIndexWriterTest.java | 3 +- .../FixedBitMVEntryDictForwardIndexTest.java | 3 +- .../forward/FixedBitMVForwardIndexTest.java | 3 +- .../FixedBitSVForwardIndexReaderTest.java | 4 +- .../FixedByteChunkSVForwardIndexTest.java | 94 ++-- .../forward/SortedForwardIndexReaderTest.java | 41 +- .../VarByteChunkSVForwardIndexTest.java | 39 +- .../mutable/CLPMutableForwardIndexTest.java | 3 +- .../mutable/CLPMutableForwardIndexV2Test.java | 3 +- .../FixedByteMVMutableForwardIndexTest.java | 3 +- .../FixedByteSVMutableForwardIndexTest.java | 3 +- .../VarByteSVMutableForwardIndexTest.java | 3 +- .../index/loader/SegmentPreProcessorTest.java | 3 +- ...ixedByteWidthRowColDataFileReaderTest.java | 13 +- .../readers/ImmutableDictionaryTest.java | 102 ++-- ...ImmutableDictionaryTypeConversionTest.java | 105 ++-- .../NullValueVectorReaderImplTest.java | 12 +- .../FixedBitSVForwardIndexReaderV2Test.java | 3 +- .../text/LuceneTextIndexCompatibleTest.java | 9 +- ...ingleValueMultiColumnReaderWriterTest.java | 3 +- .../FixedByteValueReaderWriterTest.java | 3 +- .../ValueReaderComparisonTest.java | 11 +- ...ByteWidthRowColForwardIndexWriterTest.java | 107 ++-- .../store/ColumnIndexDirectoryTestHelper.java | 3 +- .../store/FilePerIndexDirectoryTest.java | 28 +- .../store/SegmentLocalFSDirectoryTest.java | 3 +- .../store/SingleFileIndexDirectoryTest.java | 46 +- .../store/StarTreeIndexReaderTest.java | 3 +- .../local/utils/fst/FSTBuilderTest.java | 22 +- .../nativefst/FSTRegexpWithWeirdTest.java | 12 +- .../local/utils/nativefst/FSTSanityTest.java | 12 +- .../utils/nativefst/FSTTraversalTest.java | 153 ++++-- .../ImmutableFSTDeserializedTest.java | 17 +- .../utils/nativefst/ImmutableFSTTest.java | 52 +- .../utils/nativefst/SerializerTestBase.java | 33 +- .../segment/spi/memory/PinotDataBuffer.java | 17 + 99 files changed, 2190 insertions(+), 1538 deletions(-) create mode 100644 pinot-segment-local/src/test/java/org/apache/pinot/segment/local/PinotBuffersAfterClassCheckRule.java create mode 100644 pinot-segment-local/src/test/java/org/apache/pinot/segment/local/PinotBuffersAfterMethodCheckRule.java diff --git a/pinot-common/pom.xml b/pinot-common/pom.xml index 59dc5dd7a9f0..7b44a94375bb 100644 --- a/pinot-common/pom.xml +++ b/pinot-common/pom.xml @@ -45,7 +45,9 @@ org.apache.maven.plugins maven-surefire-plugin - + + 1 + true usedefaultlisteners diff --git a/pinot-common/src/test/java/org/apache/pinot/common/metrics/prometheus/dropwizard/DropwizardBrokerPrometheusMetricsTest.java b/pinot-common/src/test/java/org/apache/pinot/common/metrics/prometheus/dropwizard/DropwizardBrokerPrometheusMetricsTest.java index 8d1abd44b286..6ba25fed7920 100644 --- a/pinot-common/src/test/java/org/apache/pinot/common/metrics/prometheus/dropwizard/DropwizardBrokerPrometheusMetricsTest.java +++ b/pinot-common/src/test/java/org/apache/pinot/common/metrics/prometheus/dropwizard/DropwizardBrokerPrometheusMetricsTest.java @@ -19,6 +19,9 @@ package org.apache.pinot.common.metrics.prometheus.dropwizard; +import org.apache.pinot.common.metrics.BrokerGauge; +import org.apache.pinot.common.metrics.BrokerMeter; +import org.apache.pinot.common.metrics.BrokerTimer; import org.apache.pinot.common.metrics.prometheus.BrokerPrometheusMetricsTest; import org.apache.pinot.plugin.metrics.dropwizard.DropwizardMetricsFactory; import org.apache.pinot.spi.annotations.metrics.PinotMetricsFactory; @@ -28,7 +31,7 @@ /** * Disabling tests as Pinot currently uses Yammer and these tests fail for for {@link DropwizardMetricsFactory} */ -@Test(enabled = false) +@Test(enabled = false) // enabled=false on class level doesn't seem to work in intellij public class DropwizardBrokerPrometheusMetricsTest extends BrokerPrometheusMetricsTest { @Override protected PinotMetricsFactory getPinotMetricsFactory() { @@ -40,4 +43,19 @@ protected String getConfigFile() { //todo: return the correct dir once this test is enabled return null; } + + @Test(dataProvider = "brokerGauges", enabled = false) + public void timerTest(BrokerTimer timer) { + super.timerTest(timer); + } + + @Test(dataProvider = "brokerMeters", enabled = false) + public void meterTest(BrokerMeter meter) { + super.meterTest(meter); + } + + @Test(dataProvider = "brokerGauges", enabled = false) + public void gaugeTest(BrokerGauge gauge) { + super.gaugeTest(gauge); + } } diff --git a/pinot-common/src/test/java/org/apache/pinot/common/metrics/prometheus/dropwizard/DropwizardControllerPrometheusMetricsTest.java b/pinot-common/src/test/java/org/apache/pinot/common/metrics/prometheus/dropwizard/DropwizardControllerPrometheusMetricsTest.java index 005d83466b91..9c4bc7b45dbb 100644 --- a/pinot-common/src/test/java/org/apache/pinot/common/metrics/prometheus/dropwizard/DropwizardControllerPrometheusMetricsTest.java +++ b/pinot-common/src/test/java/org/apache/pinot/common/metrics/prometheus/dropwizard/DropwizardControllerPrometheusMetricsTest.java @@ -19,6 +19,9 @@ package org.apache.pinot.common.metrics.prometheus.dropwizard; +import org.apache.pinot.common.metrics.ControllerGauge; +import org.apache.pinot.common.metrics.ControllerMeter; +import org.apache.pinot.common.metrics.ControllerTimer; import org.apache.pinot.common.metrics.prometheus.ControllerPrometheusMetricsTest; import org.apache.pinot.plugin.metrics.dropwizard.DropwizardMetricsFactory; import org.apache.pinot.spi.annotations.metrics.PinotMetricsFactory; @@ -28,7 +31,7 @@ /** * Disabling tests as Pinot currently uses Yammer and these tests fail for for {@link DropwizardMetricsFactory} */ -@Test(enabled = false) +@Test(enabled = false) // enabled=false on class level doesn't seem to work in intellij public class DropwizardControllerPrometheusMetricsTest extends ControllerPrometheusMetricsTest { @Override @@ -41,4 +44,19 @@ protected String getConfigFile() { //todo: return the correct dir once this test is enabled return null; } + + @Test(dataProvider = "controllerTimers", enabled = false) + public void timerTest(ControllerTimer controllerTimer) { + super.timerTest(controllerTimer); + } + + @Test(dataProvider = "controllerMeters", enabled = false) + public void meterTest(ControllerMeter meter) { + super.meterTest(meter); + } + + @Test(dataProvider = "controllerGauges", enabled = false) + public void gaugeTest(ControllerGauge controllerGauge) { + super.gaugeTest(controllerGauge); + } } diff --git a/pinot-common/src/test/java/org/apache/pinot/common/metrics/prometheus/dropwizard/DropwizardMinionPrometheusMetricsTest.java b/pinot-common/src/test/java/org/apache/pinot/common/metrics/prometheus/dropwizard/DropwizardMinionPrometheusMetricsTest.java index 08183428a8eb..3bfa51f52664 100644 --- a/pinot-common/src/test/java/org/apache/pinot/common/metrics/prometheus/dropwizard/DropwizardMinionPrometheusMetricsTest.java +++ b/pinot-common/src/test/java/org/apache/pinot/common/metrics/prometheus/dropwizard/DropwizardMinionPrometheusMetricsTest.java @@ -19,6 +19,9 @@ package org.apache.pinot.common.metrics.prometheus.dropwizard; +import org.apache.pinot.common.metrics.MinionGauge; +import org.apache.pinot.common.metrics.MinionMeter; +import org.apache.pinot.common.metrics.MinionTimer; import org.apache.pinot.common.metrics.prometheus.MinionPrometheusMetricsTest; import org.apache.pinot.plugin.metrics.dropwizard.DropwizardMetricsFactory; import org.apache.pinot.spi.annotations.metrics.PinotMetricsFactory; @@ -28,7 +31,7 @@ /** * Disabling tests as Pinot currently uses Yammer and these tests fail for for {@link DropwizardMetricsFactory} */ -@Test(enabled = false) +@Test(enabled = false) // enabled=false on class level doesn't seem to work in intellij public class DropwizardMinionPrometheusMetricsTest extends MinionPrometheusMetricsTest { @Override @@ -41,4 +44,19 @@ protected String getConfigFile() { //todo: return the correct dir once this test is enabled return null; } + + @Test(dataProvider = "minionTimers", enabled = false) + public void timerTest(MinionTimer timer) { + super.timerTest(timer); + } + + @Test(dataProvider = "minionMeters", enabled = false) + public void meterTest(MinionMeter meter) { + super.meterTest(meter); + } + + @Test(dataProvider = "minionGauges", enabled = false) + public void gaugeTest(MinionGauge gauge) { + super.gaugeTest(gauge); + } } diff --git a/pinot-common/src/test/java/org/apache/pinot/common/metrics/prometheus/dropwizard/DropwizardServerPrometheusMetricsTest.java b/pinot-common/src/test/java/org/apache/pinot/common/metrics/prometheus/dropwizard/DropwizardServerPrometheusMetricsTest.java index 9fc5bf4b9690..70768e995455 100644 --- a/pinot-common/src/test/java/org/apache/pinot/common/metrics/prometheus/dropwizard/DropwizardServerPrometheusMetricsTest.java +++ b/pinot-common/src/test/java/org/apache/pinot/common/metrics/prometheus/dropwizard/DropwizardServerPrometheusMetricsTest.java @@ -19,6 +19,9 @@ package org.apache.pinot.common.metrics.prometheus.dropwizard; +import org.apache.pinot.common.metrics.ServerGauge; +import org.apache.pinot.common.metrics.ServerMeter; +import org.apache.pinot.common.metrics.ServerTimer; import org.apache.pinot.common.metrics.prometheus.ServerPrometheusMetricsTest; import org.apache.pinot.plugin.metrics.dropwizard.DropwizardMetricsFactory; import org.apache.pinot.plugin.metrics.yammer.YammerMetricsFactory; @@ -29,7 +32,7 @@ /** * Disabling tests as Pinot currently uses Yammer and these tests fail for for {@link DropwizardMetricsFactory} */ -@Test(enabled = false) +@Test(enabled = false) // enabled=false on class level doesn't seem to work in intellij public class DropwizardServerPrometheusMetricsTest extends ServerPrometheusMetricsTest { @Override @@ -42,4 +45,19 @@ protected String getConfigFile() { //todo: return the correct dir once this test is enabled return null; } + + @Test(dataProvider = "serverTimers", enabled = false) + public void timerTest(ServerTimer serverTimer) { + super.timerTest(serverTimer); + } + + @Test(dataProvider = "serverMeters", enabled = false) + public void meterTest(ServerMeter serverMeter) { + super.meterTest(serverMeter); + } + + @Test(dataProvider = "serverGauges", enabled = false) + public void gaugeTest(ServerGauge serverGauge) { + super.gaugeTest(serverGauge); + } } diff --git a/pinot-common/src/test/java/org/apache/pinot/common/utils/RoundRobinURIProviderTest.java b/pinot-common/src/test/java/org/apache/pinot/common/utils/RoundRobinURIProviderTest.java index 05dfff8a6f45..f5dba93bf2f9 100644 --- a/pinot-common/src/test/java/org/apache/pinot/common/utils/RoundRobinURIProviderTest.java +++ b/pinot-common/src/test/java/org/apache/pinot/common/utils/RoundRobinURIProviderTest.java @@ -48,132 +48,134 @@ public void testHostAddressRoundRobin() InetAddress.getByAddress("localhost", InetAddresses.forString("0:0:0:0:0:0:0:1").getAddress()) }; - MockedStatic mock = Mockito.mockStatic(InetAddress.class); - mock.when(() -> InetAddress.getAllByName("localhost")).thenReturn(localHostAddresses); - mock.when(() -> InetAddress.getAllByName("testweb.com")).thenReturn(testWebAddresses); + try (MockedStatic mock = Mockito.mockStatic(InetAddress.class)) { + mock.when(() -> InetAddress.getAllByName("localhost")).thenReturn(localHostAddresses); + mock.when(() -> InetAddress.getAllByName("testweb.com")).thenReturn(testWebAddresses); - TestCase[] testCases = new TestCase[]{ - new TestCase("http://127.0.0.1", Collections.singletonList("http://127.0.0.1")), - new TestCase("http://127.0.0.1/", Collections.singletonList("http://127.0.0.1/")), - new TestCase("http://127.0.0.1/?", Collections.singletonList("http://127.0.0.1/?")), - new TestCase("http://127.0.0.1/?it=5", Collections.singletonList("http://127.0.0.1/?it=5")), - new TestCase("http://127.0.0.1/me/out?it=5", Collections.singletonList("http://127.0.0.1/me/out?it=5")), - new TestCase("http://127.0.0.1:20000", Collections.singletonList("http://127.0.0.1:20000")), - new TestCase("http://127.0.0.1:20000/", Collections.singletonList("http://127.0.0.1:20000/")), - new TestCase("http://127.0.0.1:20000/?", Collections.singletonList("http://127.0.0.1:20000/?")), - new TestCase("http://127.0.0.1:20000/?it=5", Collections.singletonList("http://127.0.0.1:20000/?it=5")), - new TestCase("http://127.0.0.1:20000/me/out?it=5", - Collections.singletonList("http://127.0.0.1:20000/me/out?it=5")), + TestCase[] testCases = new TestCase[]{ + new TestCase("http://127.0.0.1", Collections.singletonList("http://127.0.0.1")), + new TestCase("http://127.0.0.1/", Collections.singletonList("http://127.0.0.1/")), + new TestCase("http://127.0.0.1/?", Collections.singletonList("http://127.0.0.1/?")), + new TestCase("http://127.0.0.1/?it=5", Collections.singletonList("http://127.0.0.1/?it=5")), + new TestCase("http://127.0.0.1/me/out?it=5", Collections.singletonList("http://127.0.0.1/me/out?it=5")), + new TestCase("http://127.0.0.1:20000", Collections.singletonList("http://127.0.0.1:20000")), + new TestCase("http://127.0.0.1:20000/", Collections.singletonList("http://127.0.0.1:20000/")), + new TestCase("http://127.0.0.1:20000/?", Collections.singletonList("http://127.0.0.1:20000/?")), + new TestCase("http://127.0.0.1:20000/?it=5", Collections.singletonList("http://127.0.0.1:20000/?it=5")), + new TestCase("http://127.0.0.1:20000/me/out?it=5", + Collections.singletonList("http://127.0.0.1:20000/me/out?it=5")), - new TestCase("http://localhost", Arrays.asList("http://127.0.0.1", "http://[0:0:0:0:0:0:0:1]")), - new TestCase("http://localhost/", Arrays.asList("http://127.0.0.1/", "http://[0:0:0:0:0:0:0:1]/")), - new TestCase("http://localhost/?", Arrays.asList("http://127.0.0.1/?", "http://[0:0:0:0:0:0:0:1]/?")), - new TestCase("http://localhost/?it=5", - Arrays.asList("http://127.0.0.1/?it=5", "http://[0:0:0:0:0:0:0:1]/?it=5")), - new TestCase("http://localhost/me/out?it=5", - Arrays.asList("http://127.0.0.1/me/out?it=5", "http://[0:0:0:0:0:0:0:1]/me/out?it=5")), - new TestCase("http://localhost:20000", - Arrays.asList("http://127.0.0.1:20000", "http://[0:0:0:0:0:0:0:1]:20000")), - new TestCase("http://localhost:20000/", - Arrays.asList("http://127.0.0.1:20000/", "http://[0:0:0:0:0:0:0:1]:20000/")), - new TestCase("http://localhost:20000/?", - Arrays.asList("http://127.0.0.1:20000/?", "http://[0:0:0:0:0:0:0:1]:20000/?")), - new TestCase("http://localhost:20000/?it=5", - Arrays.asList("http://127.0.0.1:20000/?it=5", "http://[0:0:0:0:0:0:0:1]:20000/?it=5")), - new TestCase("http://localhost:20000/me/out?it=5", - Arrays.asList("http://127.0.0.1:20000/me/out?it=5", "http://[0:0:0:0:0:0:0:1]:20000/me/out?it=5")), + new TestCase("http://localhost", Arrays.asList("http://127.0.0.1", "http://[0:0:0:0:0:0:0:1]")), + new TestCase("http://localhost/", Arrays.asList("http://127.0.0.1/", "http://[0:0:0:0:0:0:0:1]/")), + new TestCase("http://localhost/?", Arrays.asList("http://127.0.0.1/?", "http://[0:0:0:0:0:0:0:1]/?")), + new TestCase("http://localhost/?it=5", + Arrays.asList("http://127.0.0.1/?it=5", "http://[0:0:0:0:0:0:0:1]/?it=5")), + new TestCase("http://localhost/me/out?it=5", + Arrays.asList("http://127.0.0.1/me/out?it=5", "http://[0:0:0:0:0:0:0:1]/me/out?it=5")), + new TestCase("http://localhost:20000", + Arrays.asList("http://127.0.0.1:20000", "http://[0:0:0:0:0:0:0:1]:20000")), + new TestCase("http://localhost:20000/", + Arrays.asList("http://127.0.0.1:20000/", "http://[0:0:0:0:0:0:0:1]:20000/")), + new TestCase("http://localhost:20000/?", + Arrays.asList("http://127.0.0.1:20000/?", "http://[0:0:0:0:0:0:0:1]:20000/?")), + new TestCase("http://localhost:20000/?it=5", + Arrays.asList("http://127.0.0.1:20000/?it=5", "http://[0:0:0:0:0:0:0:1]:20000/?it=5")), + new TestCase("http://localhost:20000/me/out?it=5", + Arrays.asList("http://127.0.0.1:20000/me/out?it=5", "http://[0:0:0:0:0:0:0:1]:20000/me/out?it=5")), - new TestCase("http://testweb.com", - Arrays.asList("http://192.168.3.1", "http://192.168.3.2", "http://192.168.3.3")), - new TestCase("http://testweb.com/", - Arrays.asList("http://192.168.3.1/", "http://192.168.3.2/", "http://192.168.3.3/")), - new TestCase("http://testweb.com/?", - Arrays.asList("http://192.168.3.1/?", "http://192.168.3.2/?", "http://192.168.3.3/?")), - new TestCase("http://testweb.com/?it=5", - Arrays.asList("http://192.168.3.1/?it=5", "http://192.168.3.2/?it=5", "http://192.168.3.3/?it=5")), - new TestCase("http://testweb.com/me/out?it=5", - Arrays.asList("http://192.168.3.1/me/out?it=5", "http://192.168.3.2/me/out?it=5", - "http://192.168.3.3/me/out?it=5")), - new TestCase("http://testweb.com:20000", - Arrays.asList("http://192.168.3.1:20000", "http://192.168.3.2:20000", "http://192.168.3.3:20000")), - new TestCase("http://testweb.com:20000/", - Arrays.asList("http://192.168.3.1:20000/", "http://192.168.3.2:20000/", "http://192.168.3.3:20000/")), - new TestCase("http://testweb.com:20000/?", - Arrays.asList("http://192.168.3.1:20000/?", "http://192.168.3.2:20000/?", "http://192.168.3.3:20000/?")), - new TestCase("http://testweb.com:20000/?it=5", - Arrays.asList("http://192.168.3.1:20000/?it=5", "http://192.168.3.2:20000/?it=5", - "http://192.168.3.3:20000/?it=5")), - new TestCase("http://testweb.com:20000/me/out?it=5", - Arrays.asList("http://192.168.3.1:20000/me/out?it=5", "http://192.168.3.2:20000/me/out?it=5", - "http://192.168.3.3:20000/me/out?it=5")), + new TestCase("http://testweb.com", + Arrays.asList("http://192.168.3.1", "http://192.168.3.2", "http://192.168.3.3")), + new TestCase("http://testweb.com/", + Arrays.asList("http://192.168.3.1/", "http://192.168.3.2/", "http://192.168.3.3/")), + new TestCase("http://testweb.com/?", + Arrays.asList("http://192.168.3.1/?", "http://192.168.3.2/?", "http://192.168.3.3/?")), + new TestCase("http://testweb.com/?it=5", + Arrays.asList("http://192.168.3.1/?it=5", "http://192.168.3.2/?it=5", "http://192.168.3.3/?it=5")), + new TestCase("http://testweb.com/me/out?it=5", + Arrays.asList("http://192.168.3.1/me/out?it=5", "http://192.168.3.2/me/out?it=5", + "http://192.168.3.3/me/out?it=5")), + new TestCase("http://testweb.com:20000", + Arrays.asList("http://192.168.3.1:20000", "http://192.168.3.2:20000", "http://192.168.3.3:20000")), + new TestCase("http://testweb.com:20000/", + Arrays.asList("http://192.168.3.1:20000/", "http://192.168.3.2:20000/", "http://192.168.3.3:20000/")), + new TestCase("http://testweb.com:20000/?", + Arrays.asList("http://192.168.3.1:20000/?", "http://192.168.3.2:20000/?", "http://192.168.3.3:20000/?")), + new TestCase("http://testweb.com:20000/?it=5", + Arrays.asList("http://192.168.3.1:20000/?it=5", "http://192.168.3.2:20000/?it=5", + "http://192.168.3.3:20000/?it=5")), + new TestCase("http://testweb.com:20000/me/out?it=5", + Arrays.asList("http://192.168.3.1:20000/me/out?it=5", "http://192.168.3.2:20000/me/out?it=5", + "http://192.168.3.3:20000/me/out?it=5")), - new TestCase("https://127.0.0.1", Collections.singletonList("https://127.0.0.1")), - new TestCase("https://127.0.0.1/", Collections.singletonList("https://127.0.0.1/")), - new TestCase("https://127.0.0.1/?", Collections.singletonList("https://127.0.0.1/?")), - new TestCase("https://127.0.0.1/?it=5", Collections.singletonList("https://127.0.0.1/?it=5")), - new TestCase("https://127.0.0.1/me/out?it=5", Collections.singletonList("https://127.0.0.1/me/out?it=5")), - new TestCase("https://127.0.0.1:20000", Collections.singletonList("https://127.0.0.1:20000")), - new TestCase("https://127.0.0.1:20000/", Collections.singletonList("https://127.0.0.1:20000/")), - new TestCase("https://127.0.0.1:20000/?", Collections.singletonList("https://127.0.0.1:20000/?")), - new TestCase("https://127.0.0.1:20000/?it=5", Collections.singletonList("https://127.0.0.1:20000/?it=5")), - new TestCase("https://127.0.0.1:20000/me/out?it=5", - Collections.singletonList("https://127.0.0.1:20000/me/out?it=5")), + new TestCase("https://127.0.0.1", Collections.singletonList("https://127.0.0.1")), + new TestCase("https://127.0.0.1/", Collections.singletonList("https://127.0.0.1/")), + new TestCase("https://127.0.0.1/?", Collections.singletonList("https://127.0.0.1/?")), + new TestCase("https://127.0.0.1/?it=5", Collections.singletonList("https://127.0.0.1/?it=5")), + new TestCase("https://127.0.0.1/me/out?it=5", Collections.singletonList("https://127.0.0.1/me/out?it=5")), + new TestCase("https://127.0.0.1:20000", Collections.singletonList("https://127.0.0.1:20000")), + new TestCase("https://127.0.0.1:20000/", Collections.singletonList("https://127.0.0.1:20000/")), + new TestCase("https://127.0.0.1:20000/?", Collections.singletonList("https://127.0.0.1:20000/?")), + new TestCase("https://127.0.0.1:20000/?it=5", Collections.singletonList("https://127.0.0.1:20000/?it=5")), + new TestCase("https://127.0.0.1:20000/me/out?it=5", + Collections.singletonList("https://127.0.0.1:20000/me/out?it=5")), - new TestCase("https://localhost", Arrays.asList("https://127.0.0.1", "https://[0:0:0:0:0:0:0:1]")), - new TestCase("https://localhost/", Arrays.asList("https://127.0.0.1/", "https://[0:0:0:0:0:0:0:1]/")), - new TestCase("https://localhost/?", Arrays.asList("https://127.0.0.1/?", "https://[0:0:0:0:0:0:0:1]/?")), - new TestCase("https://localhost/?it=5", - Arrays.asList("https://127.0.0.1/?it=5", "https://[0:0:0:0:0:0:0:1]/?it=5")), - new TestCase("https://localhost/me/out?it=5", - Arrays.asList("https://127.0.0.1/me/out?it=5", "https://[0:0:0:0:0:0:0:1]/me/out?it=5")), - new TestCase("https://localhost:20000", - Arrays.asList("https://127.0.0.1:20000", "https://[0:0:0:0:0:0:0:1]:20000")), - new TestCase("https://localhost:20000/", - Arrays.asList("https://127.0.0.1:20000/", "https://[0:0:0:0:0:0:0:1]:20000/")), - new TestCase("https://localhost:20000/?", - Arrays.asList("https://127.0.0.1:20000/?", "https://[0:0:0:0:0:0:0:1]:20000/?")), - new TestCase("https://localhost:20000/?it=5", - Arrays.asList("https://127.0.0.1:20000/?it=5", "https://[0:0:0:0:0:0:0:1]:20000/?it=5")), + new TestCase("https://localhost", Arrays.asList("https://127.0.0.1", "https://[0:0:0:0:0:0:0:1]")), + new TestCase("https://localhost/", Arrays.asList("https://127.0.0.1/", "https://[0:0:0:0:0:0:0:1]/")), + new TestCase("https://localhost/?", Arrays.asList("https://127.0.0.1/?", "https://[0:0:0:0:0:0:0:1]/?")), + new TestCase("https://localhost/?it=5", + Arrays.asList("https://127.0.0.1/?it=5", "https://[0:0:0:0:0:0:0:1]/?it=5")), + new TestCase("https://localhost/me/out?it=5", + Arrays.asList("https://127.0.0.1/me/out?it=5", "https://[0:0:0:0:0:0:0:1]/me/out?it=5")), + new TestCase("https://localhost:20000", + Arrays.asList("https://127.0.0.1:20000", "https://[0:0:0:0:0:0:0:1]:20000")), + new TestCase("https://localhost:20000/", + Arrays.asList("https://127.0.0.1:20000/", "https://[0:0:0:0:0:0:0:1]:20000/")), + new TestCase("https://localhost:20000/?", + Arrays.asList("https://127.0.0.1:20000/?", "https://[0:0:0:0:0:0:0:1]:20000/?")), + new TestCase("https://localhost:20000/?it=5", + Arrays.asList("https://127.0.0.1:20000/?it=5", "https://[0:0:0:0:0:0:0:1]:20000/?it=5")), - new TestCase("https://testweb.com", - Arrays.asList("https://192.168.3.1", "https://192.168.3.2", "https://192.168.3.3")), - new TestCase("https://testweb.com/", - Arrays.asList("https://192.168.3.1/", "https://192.168.3.2/", "https://192.168.3.3/")), - new TestCase("https://testweb.com/?", - Arrays.asList("https://192.168.3.1/?", "https://192.168.3.2/?", "https://192.168.3.3/?")), - new TestCase("https://testweb.com/?it=5", - Arrays.asList("https://192.168.3.1/?it=5", "https://192.168.3.2/?it=5", "https://192.168.3.3/?it=5")), - new TestCase("https://testweb.com/me/out?it=5", - Arrays.asList("https://192.168.3.1/me/out?it=5", "https://192.168.3.2/me/out?it=5", - "https://192.168.3.3/me/out?it=5")), - new TestCase("https://testweb.com:20000", - Arrays.asList("https://192.168.3.1:20000", "https://192.168.3.2:20000", "https://192.168.3.3:20000")), - new TestCase("https://testweb.com:20000/", - Arrays.asList("https://192.168.3.1:20000/", "https://192.168.3.2:20000/", "https://192.168.3.3:20000/")), - new TestCase("https://testweb.com:20000/?", - Arrays.asList("https://192.168.3.1:20000/?", "https://192.168.3.2:20000/?", "https://192.168.3.3:20000/?")), - new TestCase("https://testweb.com:20000/?it=5", - Arrays.asList("https://192.168.3.1:20000/?it=5", "https://192.168.3.2:20000/?it=5", - "https://192.168.3.3:20000/?it=5")), - new TestCase("https://testweb.com:20000/me/out?it=5", - Arrays.asList("https://192.168.3.1:20000/me/out?it=5", "https://192.168.3.2:20000/me/out?it=5", - "https://192.168.3.3:20000/me/out?it=5")), - }; + new TestCase("https://testweb.com", + Arrays.asList("https://192.168.3.1", "https://192.168.3.2", "https://192.168.3.3")), + new TestCase("https://testweb.com/", + Arrays.asList("https://192.168.3.1/", "https://192.168.3.2/", "https://192.168.3.3/")), + new TestCase("https://testweb.com/?", + Arrays.asList("https://192.168.3.1/?", "https://192.168.3.2/?", "https://192.168.3.3/?")), + new TestCase("https://testweb.com/?it=5", + Arrays.asList("https://192.168.3.1/?it=5", "https://192.168.3.2/?it=5", "https://192.168.3.3/?it=5")), + new TestCase("https://testweb.com/me/out?it=5", + Arrays.asList("https://192.168.3.1/me/out?it=5", "https://192.168.3.2/me/out?it=5", + "https://192.168.3.3/me/out?it=5")), + new TestCase("https://testweb.com:20000", + Arrays.asList("https://192.168.3.1:20000", "https://192.168.3.2:20000", "https://192.168.3.3:20000")), + new TestCase("https://testweb.com:20000/", + Arrays.asList("https://192.168.3.1:20000/", "https://192.168.3.2:20000/", "https://192.168.3.3:20000/")), + new TestCase("https://testweb.com:20000/?", + Arrays.asList("https://192.168.3.1:20000/?", "https://192.168.3.2:20000/?", + "https://192.168.3.3:20000/?")), + new TestCase("https://testweb.com:20000/?it=5", + Arrays.asList("https://192.168.3.1:20000/?it=5", "https://192.168.3.2:20000/?it=5", + "https://192.168.3.3:20000/?it=5")), + new TestCase("https://testweb.com:20000/me/out?it=5", + Arrays.asList("https://192.168.3.1:20000/me/out?it=5", "https://192.168.3.2:20000/me/out?it=5", + "https://192.168.3.3:20000/me/out?it=5")), + }; - for (TestCase testCase : testCases) { - String uri = testCase._originalUri; - RoundRobinURIProvider uriProvider = new RoundRobinURIProvider(List.of(new URI(uri)), true); - int n = testCase._expectedUris.size(); - int previousIndex = -1; - int currentIndex; - for (int i = 0; i < 2 * n; i++) { - String actualUri = uriProvider.next().toString(); - currentIndex = testCase._expectedUris.indexOf(actualUri); - Assert.assertTrue(currentIndex != -1); - if (previousIndex != -1) { - Assert.assertEquals((previousIndex + 1) % n, currentIndex); + for (TestCase testCase : testCases) { + String uri = testCase._originalUri; + RoundRobinURIProvider uriProvider = new RoundRobinURIProvider(List.of(new URI(uri)), true); + int n = testCase._expectedUris.size(); + int previousIndex = -1; + int currentIndex; + for (int i = 0; i < 2 * n; i++) { + String actualUri = uriProvider.next().toString(); + currentIndex = testCase._expectedUris.indexOf(actualUri); + Assert.assertTrue(currentIndex != -1); + if (previousIndex != -1) { + Assert.assertEquals((previousIndex + 1) % n, currentIndex); + } + previousIndex = currentIndex; } - previousIndex = currentIndex; } } } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/writer/impl/FixedByteSingleValueMultiColWriter.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/writer/impl/FixedByteSingleValueMultiColWriter.java index cafae34189bb..bc86b468f8c8 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/writer/impl/FixedByteSingleValueMultiColWriter.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/writer/impl/FixedByteSingleValueMultiColWriter.java @@ -59,7 +59,7 @@ public FixedByteSingleValueMultiColWriter(PinotDataBuffer dataBuffer, int cols, } _rowSizeInBytes = rowSizeInBytes; _dataBuffer = dataBuffer; - // For passed in PinotDataBuffer, the caller is responsible of closing the PinotDataBuffer. + // For passed in PinotDataBuffer, the caller is responsible for closing the PinotDataBuffer. _shouldCloseDataBuffer = false; } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/forward/BaseChunkForwardIndexReader.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/forward/BaseChunkForwardIndexReader.java index 80ed58a4162f..b11a991d33aa 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/forward/BaseChunkForwardIndexReader.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/forward/BaseChunkForwardIndexReader.java @@ -447,7 +447,7 @@ public void readValuesSV(int[] docIds, int length, double[] values, ChunkReaderC public void close() throws IOException { // NOTE: DO NOT close the PinotDataBuffer here because it is tracked by the caller and might be reused later. The - // caller is responsible of closing the PinotDataBuffer. + // caller is responsible for closing the PinotDataBuffer. _chunkDecompressor.close(); } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/text/NativeTextIndexReader.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/text/NativeTextIndexReader.java index a9cd64c91b9d..3131f60e314e 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/text/NativeTextIndexReader.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/text/NativeTextIndexReader.java @@ -18,6 +18,7 @@ */ package org.apache.pinot.segment.local.segment.index.readers.text; +import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; import java.io.File; import java.io.IOException; @@ -119,6 +120,15 @@ public MutableRoaringBitmap getDocIds(String searchQuery) { @Override public void close() throws IOException { + // TODO: this method doesn't release native buffers held by _fst _buffer.close(); } + + @VisibleForTesting + public void closeInTest() + throws IOException { + if (_fst instanceof ImmutableFST) { + ((ImmutableFST) _fst)._mutableBytesStore.close(); + } + } } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/store/FilePerIndexDirectory.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/store/FilePerIndexDirectory.java index 21fa74950acf..fe05e2e49605 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/store/FilePerIndexDirectory.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/store/FilePerIndexDirectory.java @@ -99,6 +99,7 @@ public void close() @Override public void removeIndex(String columnName, IndexType indexType) { + // TODO: this leaks the removed data buffer (it's not going to be freed in close() method) _indexBuffers.remove(new IndexKey(columnName, indexType)); if (indexType == StandardIndexes.text()) { TextIndexUtils.cleanupTextIndex(_segmentDirectory, columnName); diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/nativefst/NativeFSTIndexReader.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/nativefst/NativeFSTIndexReader.java index 0a03389336f2..5db6760f101b 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/nativefst/NativeFSTIndexReader.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/nativefst/NativeFSTIndexReader.java @@ -18,6 +18,7 @@ */ package org.apache.pinot.segment.local.utils.nativefst; +import com.google.common.annotations.VisibleForTesting; import java.io.IOException; import java.nio.ByteBuffer; import java.util.Collections; @@ -71,5 +72,15 @@ public ImmutableRoaringBitmap getDictIds(String searchQuery) { @Override public void close() throws IOException { + //TODO: why does this class not close FST ? + } + + @VisibleForTesting + public void closeInTest() + throws IOException { + // immutable fst contains native data buffers that need to be closed + if (_fst instanceof ImmutableFST) { + ((ImmutableFST) _fst)._mutableBytesStore.close(); + } } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/PinotBuffersAfterClassCheckRule.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/PinotBuffersAfterClassCheckRule.java new file mode 100644 index 000000000000..ef0894d178bf --- /dev/null +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/PinotBuffersAfterClassCheckRule.java @@ -0,0 +1,33 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.segment.local; + +import org.testng.ITestContext; +import org.testng.annotations.AfterClass; + + +/** Checks that tests don't leak buffers after executing all test methods. + * Meant for tests that share fixture that contains buffers and can't be verified with an 'after' method. */ +public interface PinotBuffersAfterClassCheckRule { + + @AfterClass + default void checkPinotBuffersAfterClass(ITestContext context) { + PinotBuffersAfterMethodCheckRule.assertPinotBuffers(context.getAllTestMethods()[0].getRealClass().getName()); + } +} diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/PinotBuffersAfterMethodCheckRule.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/PinotBuffersAfterMethodCheckRule.java new file mode 100644 index 000000000000..8d9f79ab922a --- /dev/null +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/PinotBuffersAfterMethodCheckRule.java @@ -0,0 +1,70 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.segment.local; + +import java.util.List; +import org.apache.pinot.segment.spi.memory.PinotDataBuffer; +import org.testng.ITestResult; +import org.testng.annotations.AfterMethod; +import org.testng.annotations.Test; + + +/** Check for pinot buffer leaks after each test method */ +public interface PinotBuffersAfterMethodCheckRule { + + @AfterMethod + default void checkPinotBuffers(ITestResult result) { + assertPinotBuffers(result); + } + + @Test(enabled = false) + static void assertPinotBuffers(Object target) { + List bufferInfos = PinotDataBuffer.getBufferInfo(); + if (bufferInfos.size() > 0) { + + StringBuilder builder; + if (target instanceof ITestResult) { + ITestResult result = (ITestResult) target; + if (result.getStatus() != ITestResult.SUCCESS) { + // don't override result of failed test + PinotDataBuffer.closeOpenBuffers(); + return; + } + + builder = new StringBuilder("Test method: ").append(result.getTestClass().getRealClass().getSimpleName()) + .append('.').append(result.getMethod().getMethodName()); + } else { + builder = new StringBuilder("Test class: ").append(target); + } + + StringBuilder message = + builder.append(target) + .append(" did not release ") + .append(bufferInfos.size()) + .append(" pinot buffer(s): \n"); + for (String bufferInfo : bufferInfos) { + message.append(bufferInfo).append('\n'); + } + //close all open buffers otherwise one test will fail all following tests + PinotDataBuffer.closeOpenBuffers(); + + throw new AssertionError(message.toString()); + } + } +} diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/IndexingFailureTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/IndexingFailureTest.java index 30be1374650b..86dfcbf27643 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/IndexingFailureTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/IndexingFailureTest.java @@ -24,12 +24,14 @@ import java.util.HashSet; import org.apache.pinot.common.metrics.ServerMeter; import org.apache.pinot.common.metrics.ServerMetrics; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.spi.config.table.JsonIndexConfig; import org.apache.pinot.spi.data.FieldSpec; import org.apache.pinot.spi.data.Schema; import org.apache.pinot.spi.data.readers.GenericRow; import org.apache.pinot.spi.stream.StreamMessageMetadata; import org.roaringbitmap.buffer.ImmutableRoaringBitmap; +import org.testng.annotations.AfterMethod; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; @@ -40,7 +42,7 @@ import static org.testng.Assert.assertTrue; -public class IndexingFailureTest { +public class IndexingFailureTest implements PinotBuffersAfterMethodCheckRule { private static final String TABLE_NAME = "testTable"; private static final String INT_COL = "int_col"; private static final String STRING_COL = "string_col"; @@ -61,6 +63,11 @@ public void setup() { Collections.singletonMap(JSON_COL, new JsonIndexConfig()), _serverMetrics); } + @AfterMethod + public void tearDown() { + _mutableSegment.destroy(); + } + @Test public void testIndexingFailures() throws IOException { diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentDedupeTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentDedupeTest.java index bb21b7b11cea..743fffa8d486 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentDedupeTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentDedupeTest.java @@ -28,6 +28,7 @@ import java.util.Map; import org.apache.commons.io.FileUtils; import org.apache.pinot.common.metrics.ServerMetrics; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.data.manager.TableDataManager; import org.apache.pinot.segment.local.dedup.PartitionDedupMetadataManager; import org.apache.pinot.segment.local.dedup.TableDedupMetadataManager; @@ -46,10 +47,11 @@ import org.apache.pinot.spi.utils.builder.TableConfigBuilder; import org.mockito.Mockito; import org.testng.Assert; +import org.testng.annotations.AfterMethod; import org.testng.annotations.Test; -public class MutableSegmentDedupeTest { +public class MutableSegmentDedupeTest implements PinotBuffersAfterMethodCheckRule { private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(), MutableSegmentDedupeTest.class.getSimpleName()); private static final String SCHEMA_FILE_PATH = "data/test_dedup_schema.json"; @@ -68,22 +70,37 @@ private void setup(boolean dedupEnabled, double metadataTTL, String dedupTimeCol DedupConfig dedupConfig = new DedupConfig(true, HashFunction.NONE, null, null, metadataTTL, dedupTimeColumn, false); PartitionDedupMetadataManager partitionDedupMetadataManager = (dedupEnabled) ? getTableDedupMetadataManager(schema, dedupConfig).getOrCreatePartitionManager(0) : null; - _mutableSegmentImpl = - MutableSegmentImplTestUtils.createMutableSegmentImpl(schema, Collections.emptySet(), Collections.emptySet(), - Collections.emptySet(), false, true, null, "secondsSinceEpoch", null, dedupConfig, - partitionDedupMetadataManager); - GenericRow reuse = new GenericRow(); - try (RecordReader recordReader = RecordReaderFactory.getRecordReader(FileFormat.JSON, jsonFile, - schema.getColumnNames(), null)) { - while (recordReader.hasNext()) { - recordReader.next(reuse); - GenericRow transformedRow = recordTransformer.transform(reuse); - _mutableSegmentImpl.index(transformedRow, null); - if (dedupEnabled) { - partitionDedupMetadataManager.removeExpiredPrimaryKeys(); + try { + _mutableSegmentImpl = + MutableSegmentImplTestUtils.createMutableSegmentImpl(schema, Collections.emptySet(), Collections.emptySet(), + Collections.emptySet(), false, true, null, "secondsSinceEpoch", null, dedupConfig, + partitionDedupMetadataManager); + GenericRow reuse = new GenericRow(); + try (RecordReader recordReader = RecordReaderFactory.getRecordReader(FileFormat.JSON, jsonFile, + schema.getColumnNames(), null)) { + while (recordReader.hasNext()) { + recordReader.next(reuse); + GenericRow transformedRow = recordTransformer.transform(reuse); + _mutableSegmentImpl.index(transformedRow, null); + if (dedupEnabled) { + partitionDedupMetadataManager.removeExpiredPrimaryKeys(); + } + reuse.clear(); } - reuse.clear(); } + } finally { + if (partitionDedupMetadataManager != null) { + partitionDedupMetadataManager.stop(); + partitionDedupMetadataManager.close(); + } + } + } + + @AfterMethod + public void tearDown() { + if (_mutableSegmentImpl != null) { + _mutableSegmentImpl.destroy(); + _mutableSegmentImpl = null; } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentEntriesAboveThresholdTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentEntriesAboveThresholdTest.java index 1eaaab657d21..97ad899a9c2a 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentEntriesAboveThresholdTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentEntriesAboveThresholdTest.java @@ -25,6 +25,7 @@ import java.util.Collections; import java.util.Map; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterClassCheckRule; import org.apache.pinot.segment.local.segment.creator.SegmentTestUtils; import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; import org.apache.pinot.segment.local.segment.index.forward.ForwardIndexPlugin; @@ -46,7 +47,7 @@ import org.testng.annotations.Test; -public class MutableSegmentEntriesAboveThresholdTest { +public class MutableSegmentEntriesAboveThresholdTest implements PinotBuffersAfterClassCheckRule { private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(), MutableSegmentEntriesAboveThresholdTest.class.getSimpleName()); private static final String AVRO_FILE = "data/test_data-mv.avro"; @@ -140,15 +141,19 @@ public void testNoLimitBreached() throws Exception { File avroFile = getAvroFile(); MutableSegmentImpl mutableSegment = getMutableSegment(avroFile); - StreamMessageMetadata defaultMetadata = new StreamMessageMetadata(System.currentTimeMillis(), new GenericRow()); - try (RecordReader recordReader = RecordReaderFactory - .getRecordReader(FileFormat.AVRO, avroFile, _schema.getColumnNames(), null)) { - GenericRow reuse = new GenericRow(); - while (recordReader.hasNext()) { - mutableSegment.index(recordReader.next(reuse), defaultMetadata); + try { + StreamMessageMetadata defaultMetadata = new StreamMessageMetadata(System.currentTimeMillis(), new GenericRow()); + try (RecordReader recordReader = RecordReaderFactory + .getRecordReader(FileFormat.AVRO, avroFile, _schema.getColumnNames(), null)) { + GenericRow reuse = new GenericRow(); + while (recordReader.hasNext()) { + mutableSegment.index(recordReader.next(reuse), defaultMetadata); + } } + assert mutableSegment.canAddMore(); + } finally { + mutableSegment.destroy(); } - assert mutableSegment.canAddMore(); } @Test @@ -156,39 +161,43 @@ public void testLimitBreached() throws Exception { File avroFile = getAvroFile(); MutableSegmentImpl mutableSegment = getMutableSegment(avroFile); - - Field indexContainerMapField = MutableSegmentImpl.class.getDeclaredField("_indexContainerMap"); - indexContainerMapField.setAccessible(true); - Map colVsIndexContainer = (Map) indexContainerMapField.get(mutableSegment); - - for (Map.Entry entry : colVsIndexContainer.entrySet()) { - Object indexContainer = entry.getValue(); - Field mutableIndexesField = indexContainer.getClass().getDeclaredField("_mutableIndexes"); - mutableIndexesField.setAccessible(true); - Map indexTypeVsMutableIndex = - (Map) mutableIndexesField.get(indexContainer); - - MutableForwardIndex mutableForwardIndex = null; - for (IndexType indexType : indexTypeVsMutableIndex.keySet()) { - if (indexType.getId().equals(StandardIndexes.FORWARD_ID)) { - mutableForwardIndex = (MutableForwardIndex) indexTypeVsMutableIndex.get(indexType); + try { + + Field indexContainerMapField = MutableSegmentImpl.class.getDeclaredField("_indexContainerMap"); + indexContainerMapField.setAccessible(true); + Map colVsIndexContainer = (Map) indexContainerMapField.get(mutableSegment); + + for (Map.Entry entry : colVsIndexContainer.entrySet()) { + Object indexContainer = entry.getValue(); + Field mutableIndexesField = indexContainer.getClass().getDeclaredField("_mutableIndexes"); + mutableIndexesField.setAccessible(true); + Map indexTypeVsMutableIndex = + (Map) mutableIndexesField.get(indexContainer); + + MutableForwardIndex mutableForwardIndex = null; + for (IndexType indexType : indexTypeVsMutableIndex.keySet()) { + if (indexType.getId().equals(StandardIndexes.FORWARD_ID)) { + mutableForwardIndex = (MutableForwardIndex) indexTypeVsMutableIndex.get(indexType); + } } - } - assert mutableForwardIndex != null; + assert mutableForwardIndex != null; - indexTypeVsMutableIndex.put(new ForwardIndexPlugin().getIndexType(), - new FakeMutableForwardIndex(mutableForwardIndex)); - } - StreamMessageMetadata defaultMetadata = new StreamMessageMetadata(System.currentTimeMillis(), new GenericRow()); - try (RecordReader recordReader = RecordReaderFactory - .getRecordReader(FileFormat.AVRO, avroFile, _schema.getColumnNames(), null)) { - GenericRow reuse = new GenericRow(); - while (recordReader.hasNext()) { - mutableSegment.index(recordReader.next(reuse), defaultMetadata); + indexTypeVsMutableIndex.put(new ForwardIndexPlugin().getIndexType(), + new FakeMutableForwardIndex(mutableForwardIndex)); + } + StreamMessageMetadata defaultMetadata = new StreamMessageMetadata(System.currentTimeMillis(), new GenericRow()); + try (RecordReader recordReader = RecordReaderFactory + .getRecordReader(FileFormat.AVRO, avroFile, _schema.getColumnNames(), null)) { + GenericRow reuse = new GenericRow(); + while (recordReader.hasNext()) { + mutableSegment.index(recordReader.next(reuse), defaultMetadata); + } } - } - assert !mutableSegment.canAddMore(); + assert !mutableSegment.canAddMore(); + } finally { + mutableSegment.destroy(); + } } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplNullValueVectorTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplNullValueVectorTest.java index 8500d0da2fbe..21b5640b34f6 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplNullValueVectorTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplNullValueVectorTest.java @@ -23,6 +23,7 @@ import java.util.Arrays; import java.util.Collections; import java.util.List; +import org.apache.pinot.segment.local.PinotBuffersAfterClassCheckRule; import org.apache.pinot.segment.local.recordtransformer.CompositeTransformer; import org.apache.pinot.segment.spi.datasource.DataSource; import org.apache.pinot.segment.spi.index.reader.NullValueVectorReader; @@ -35,11 +36,12 @@ import org.apache.pinot.spi.data.readers.RecordReaderFactory; import org.apache.pinot.spi.utils.builder.TableConfigBuilder; import org.testng.Assert; +import org.testng.annotations.AfterClass; import org.testng.annotations.BeforeClass; import org.testng.annotations.Test; -public class MutableSegmentImplNullValueVectorTest { +public class MutableSegmentImplNullValueVectorTest implements PinotBuffersAfterClassCheckRule { private static final String PINOT_SCHEMA_FILE_PATH = "data/test_null_value_vector_pinot_schema.json"; private static final String DATA_FILE = "data/test_null_value_vector_data.json"; private static CompositeTransformer _recordTransformer; @@ -73,6 +75,11 @@ public void setup() _finalNullColumns = Arrays.asList("signup_email", "cityid"); } + @AfterClass + public void tearDown() { + _mutableSegmentImpl.destroy(); + } + @Test public void testNullValueVector() throws Exception { diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplRawMVTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplRawMVTest.java index cab06c453108..b6da06bc63d5 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplRawMVTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplRawMVTest.java @@ -26,6 +26,7 @@ import java.util.List; import java.util.Set; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterClassCheckRule; import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; import org.apache.pinot.segment.local.segment.creator.SegmentTestUtils; import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; @@ -60,7 +61,7 @@ import static org.testng.Assert.assertNull; -public class MutableSegmentImplRawMVTest { +public class MutableSegmentImplRawMVTest implements PinotBuffersAfterClassCheckRule { private static final String AVRO_FILE = "data/test_data-mv.avro"; private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(), "MutableSegmentImplRawMVTest"); @@ -262,6 +263,8 @@ public void testDataSourceForMVColumns() @AfterClass public void tearDown() { + _mutableSegmentImpl.destroy(); + _immutableSegment.destroy(); FileUtils.deleteQuietly(TEMP_DIR); } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplTest.java index 5a58c2dddcb0..9010b3a6f5c8 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplTest.java @@ -204,6 +204,8 @@ public void testDataSourceForMVColumns() @AfterClass public void tearDown() { + _mutableSegmentImpl.destroy(); + _immutableSegment.destroy(); FileUtils.deleteQuietly(TEMP_DIR); } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplTestUtils.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplTestUtils.java index b23f203ec7a1..af2ec5dd2e0b 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplTestUtils.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplTestUtils.java @@ -109,16 +109,24 @@ public static MutableSegmentImpl createMutableSegmentImpl(Schema schema, Set entry : jsonIndexConfigs.entrySet()) { segmentConfBuilder.setIndex(entry.getKey(), StandardIndexes.json(), entry.getValue()); } + RealtimeSegmentConfig realtimeSegmentConfig = segmentConfBuilder.build(); return new MutableSegmentImpl(realtimeSegmentConfig, serverMetrics); } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplUpsertComparisonColTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplUpsertComparisonColTest.java index 73ea23360761..54e9019e6c49 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplUpsertComparisonColTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplUpsertComparisonColTest.java @@ -19,9 +19,11 @@ package org.apache.pinot.segment.local.indexsegment.mutable; import java.io.File; +import java.io.IOException; import java.net.URL; import java.util.Collections; import org.apache.pinot.common.metrics.ServerMetrics; +import org.apache.pinot.segment.local.PinotBuffersAfterClassCheckRule; import org.apache.pinot.segment.local.data.manager.TableDataManager; import org.apache.pinot.segment.local.recordtransformer.CompositeTransformer; import org.apache.pinot.segment.local.upsert.PartitionUpsertMetadataManager; @@ -48,7 +50,7 @@ import static org.mockito.Mockito.when; -public class MutableSegmentImplUpsertComparisonColTest { +public class MutableSegmentImplUpsertComparisonColTest implements PinotBuffersAfterClassCheckRule { private static final String SCHEMA_FILE_PATH = "data/test_upsert_comparison_col_schema.json"; private static final String DATA_FILE_PATH = "data/test_upsert_comparison_col_data.json"; private static final String RAW_TABLE_NAME = "testTable"; @@ -104,6 +106,19 @@ public void setup(UpsertConfig upsertConfig) } } + private void tearDown() + throws IOException { + if (_mutableSegmentImpl != null) { + _mutableSegmentImpl.destroy(); + _mutableSegmentImpl = null; + } + if (_partitionUpsertMetadataManager != null) { + _partitionUpsertMetadataManager.stop(); + _partitionUpsertMetadataManager.close(); + _partitionUpsertMetadataManager = null; + } + } + @Test public void testHashFunctions() throws Exception { @@ -131,25 +146,33 @@ public void testUpsertOutOfOrderRecordColumn() public void testUpsertIngestion(UpsertConfig upsertConfig) throws Exception { setup(upsertConfig); - ImmutableRoaringBitmap bitmap = _mutableSegmentImpl.getValidDocIds().getMutableRoaringBitmap(); - // note offset column is used for determining sequence but not time column - Assert.assertEquals(_mutableSegmentImpl.getNumDocsIndexed(), 4); - Assert.assertFalse(bitmap.contains(0)); - Assert.assertTrue(bitmap.contains(1)); - Assert.assertTrue(bitmap.contains(2)); - Assert.assertFalse(bitmap.contains(3)); + try { + ImmutableRoaringBitmap bitmap = _mutableSegmentImpl.getValidDocIds().getMutableRoaringBitmap(); + // note offset column is used for determining sequence but not time column + Assert.assertEquals(_mutableSegmentImpl.getNumDocsIndexed(), 4); + Assert.assertFalse(bitmap.contains(0)); + Assert.assertTrue(bitmap.contains(1)); + Assert.assertTrue(bitmap.contains(2)); + Assert.assertFalse(bitmap.contains(3)); + } finally { + tearDown(); + } } public void testUpsertDropOfOrderRecordIngestion(UpsertConfig upsertConfig) throws Exception { upsertConfig.setDropOutOfOrderRecord(true); setup(upsertConfig); - ImmutableRoaringBitmap bitmap = _mutableSegmentImpl.getValidDocIds().getMutableRoaringBitmap(); - // note offset column is used for determining sequence but not time column - Assert.assertEquals(_mutableSegmentImpl.getNumDocsIndexed(), 3); - Assert.assertFalse(bitmap.contains(0)); - Assert.assertTrue(bitmap.contains(1)); - Assert.assertTrue(bitmap.contains(2)); + try { + ImmutableRoaringBitmap bitmap = _mutableSegmentImpl.getValidDocIds().getMutableRoaringBitmap(); + // note offset column is used for determining sequence but not time column + Assert.assertEquals(_mutableSegmentImpl.getNumDocsIndexed(), 3); + Assert.assertFalse(bitmap.contains(0)); + Assert.assertTrue(bitmap.contains(1)); + Assert.assertTrue(bitmap.contains(2)); + } finally { + tearDown(); + } } public void testUpsertOutOfOrderRecordColumnIngestion(UpsertConfig upsertConfig) @@ -157,17 +180,21 @@ public void testUpsertOutOfOrderRecordColumnIngestion(UpsertConfig upsertConfig) String outOfOrderRecordColumn = "outOfOrderRecordColumn"; upsertConfig.setOutOfOrderRecordColumn(outOfOrderRecordColumn); setup(upsertConfig); - ImmutableRoaringBitmap bitmap = _mutableSegmentImpl.getValidDocIds().getMutableRoaringBitmap(); - // note offset column is used for determining sequence but not time column - Assert.assertEquals(_mutableSegmentImpl.getNumDocsIndexed(), 4); - Assert.assertFalse(bitmap.contains(0)); - Assert.assertTrue(bitmap.contains(1)); - Assert.assertTrue(bitmap.contains(2)); - Assert.assertFalse(bitmap.contains(3)); - - Assert.assertFalse(BooleanUtils.toBoolean(_mutableSegmentImpl.getValue(0, outOfOrderRecordColumn))); - Assert.assertFalse(BooleanUtils.toBoolean(_mutableSegmentImpl.getValue(1, outOfOrderRecordColumn))); - Assert.assertFalse(BooleanUtils.toBoolean(_mutableSegmentImpl.getValue(2, outOfOrderRecordColumn))); - Assert.assertTrue(BooleanUtils.toBoolean(_mutableSegmentImpl.getValue(3, outOfOrderRecordColumn))); + try { + ImmutableRoaringBitmap bitmap = _mutableSegmentImpl.getValidDocIds().getMutableRoaringBitmap(); + // note offset column is used for determining sequence but not time column + Assert.assertEquals(_mutableSegmentImpl.getNumDocsIndexed(), 4); + Assert.assertFalse(bitmap.contains(0)); + Assert.assertTrue(bitmap.contains(1)); + Assert.assertTrue(bitmap.contains(2)); + Assert.assertFalse(bitmap.contains(3)); + + Assert.assertFalse(BooleanUtils.toBoolean(_mutableSegmentImpl.getValue(0, outOfOrderRecordColumn))); + Assert.assertFalse(BooleanUtils.toBoolean(_mutableSegmentImpl.getValue(1, outOfOrderRecordColumn))); + Assert.assertFalse(BooleanUtils.toBoolean(_mutableSegmentImpl.getValue(2, outOfOrderRecordColumn))); + Assert.assertTrue(BooleanUtils.toBoolean(_mutableSegmentImpl.getValue(3, outOfOrderRecordColumn))); + } finally { + tearDown(); + } } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplUpsertTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplUpsertTest.java index 4f46bd6d5906..090c35e8f71e 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplUpsertTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplUpsertTest.java @@ -19,6 +19,7 @@ package org.apache.pinot.segment.local.indexsegment.mutable; import java.io.File; +import java.io.IOException; import java.net.URL; import java.util.Arrays; import java.util.Collections; @@ -115,6 +116,19 @@ private void setup(UpsertConfig upsertConfigWithHash) } } + private void tearDown() + throws IOException { + if (_mutableSegmentImpl != null) { + _mutableSegmentImpl.destroy(); + _mutableSegmentImpl = null; + } + if (_partitionUpsertMetadataManager != null) { + _partitionUpsertMetadataManager.stop(); + _partitionUpsertMetadataManager.close(); + _partitionUpsertMetadataManager = null; + } + } + @Test public void testHashFunctions() throws Exception { @@ -134,38 +148,42 @@ public void testMultipleComparisonColumns() private void testUpsertIngestion(UpsertConfig upsertConfig) throws Exception { setup(upsertConfig); - ImmutableRoaringBitmap bitmap = _mutableSegmentImpl.getValidDocIds().getMutableRoaringBitmap(); - if (upsertConfig.getComparisonColumns() == null) { - // aa - Assert.assertFalse(bitmap.contains(0)); - Assert.assertTrue(bitmap.contains(1)); - Assert.assertFalse(bitmap.contains(2)); - Assert.assertFalse(bitmap.contains(3)); - // bb - Assert.assertFalse(bitmap.contains(4)); - Assert.assertTrue(bitmap.contains(5)); - Assert.assertFalse(bitmap.contains(6)); - } else { - // aa - Assert.assertFalse(bitmap.contains(0)); - Assert.assertFalse(bitmap.contains(1)); - Assert.assertTrue(bitmap.contains(2)); - Assert.assertFalse(bitmap.contains(3)); - // Confirm that both comparison column values have made it into the persisted upserted doc - Assert.assertEquals(1567205397L, _mutableSegmentImpl.getValue(2, "secondsSinceEpoch")); - Assert.assertEquals(1567205395L, _mutableSegmentImpl.getValue(2, "otherComparisonColumn")); - Assert.assertFalse(_mutableSegmentImpl.getDataSource("secondsSinceEpoch").getNullValueVector().isNull(2)); - Assert.assertFalse(_mutableSegmentImpl.getDataSource("otherComparisonColumn").getNullValueVector().isNull(2)); + try { + ImmutableRoaringBitmap bitmap = _mutableSegmentImpl.getValidDocIds().getMutableRoaringBitmap(); + if (upsertConfig.getComparisonColumns() == null) { + // aa + Assert.assertFalse(bitmap.contains(0)); + Assert.assertTrue(bitmap.contains(1)); + Assert.assertFalse(bitmap.contains(2)); + Assert.assertFalse(bitmap.contains(3)); + // bb + Assert.assertFalse(bitmap.contains(4)); + Assert.assertTrue(bitmap.contains(5)); + Assert.assertFalse(bitmap.contains(6)); + } else { + // aa + Assert.assertFalse(bitmap.contains(0)); + Assert.assertFalse(bitmap.contains(1)); + Assert.assertTrue(bitmap.contains(2)); + Assert.assertFalse(bitmap.contains(3)); + // Confirm that both comparison column values have made it into the persisted upserted doc + Assert.assertEquals(1567205397L, _mutableSegmentImpl.getValue(2, "secondsSinceEpoch")); + Assert.assertEquals(1567205395L, _mutableSegmentImpl.getValue(2, "otherComparisonColumn")); + Assert.assertFalse(_mutableSegmentImpl.getDataSource("secondsSinceEpoch").getNullValueVector().isNull(2)); + Assert.assertFalse(_mutableSegmentImpl.getDataSource("otherComparisonColumn").getNullValueVector().isNull(2)); - // bb - Assert.assertFalse(bitmap.contains(4)); - Assert.assertTrue(bitmap.contains(5)); - Assert.assertFalse(bitmap.contains(6)); - // Confirm that comparison column values have made it into the persisted upserted doc - Assert.assertEquals(1567205396L, _mutableSegmentImpl.getValue(5, "secondsSinceEpoch")); - Assert.assertEquals(Long.MIN_VALUE, _mutableSegmentImpl.getValue(5, "otherComparisonColumn")); - Assert.assertTrue(_mutableSegmentImpl.getDataSource("otherComparisonColumn").getNullValueVector().isNull(5)); - Assert.assertFalse(_mutableSegmentImpl.getDataSource("secondsSinceEpoch").getNullValueVector().isNull(5)); + // bb + Assert.assertFalse(bitmap.contains(4)); + Assert.assertTrue(bitmap.contains(5)); + Assert.assertFalse(bitmap.contains(6)); + // Confirm that comparison column values have made it into the persisted upserted doc + Assert.assertEquals(1567205396L, _mutableSegmentImpl.getValue(5, "secondsSinceEpoch")); + Assert.assertEquals(Long.MIN_VALUE, _mutableSegmentImpl.getValue(5, "otherComparisonColumn")); + Assert.assertTrue(_mutableSegmentImpl.getDataSource("otherComparisonColumn").getNullValueVector().isNull(5)); + Assert.assertFalse(_mutableSegmentImpl.getDataSource("secondsSinceEpoch").getNullValueVector().isNull(5)); + } + } finally { + tearDown(); } } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/reader/impl/FixedBitIntReaderTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/reader/impl/FixedBitIntReaderTest.java index f30f3405760a..f957d4509966 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/reader/impl/FixedBitIntReaderTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/reader/impl/FixedBitIntReaderTest.java @@ -22,6 +22,7 @@ import java.io.IOException; import java.util.Random; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.io.writer.impl.FixedBitSVForwardIndexWriter; import org.apache.pinot.segment.spi.memory.PinotDataBuffer; import org.testng.annotations.AfterClass; @@ -31,7 +32,7 @@ import static org.testng.Assert.assertEquals; -public class FixedBitIntReaderTest { +public class FixedBitIntReaderTest implements PinotBuffersAfterMethodCheckRule { private static final File INDEX_DIR = new File(FileUtils.getTempDirectory(), "FixedBitIntReaderTest"); private static final int NUM_VALUES = 95; private static final Random RANDOM = new Random(); diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/util/PinotDataBitSetTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/util/PinotDataBitSetTest.java index b84d41d0b301..7b1101942542 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/util/PinotDataBitSetTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/util/PinotDataBitSetTest.java @@ -21,13 +21,14 @@ import java.io.IOException; import java.nio.ByteOrder; import java.util.Random; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.spi.memory.PinotDataBuffer; import org.testng.annotations.Test; import static org.testng.Assert.assertEquals; -public class PinotDataBitSetTest { +public class PinotDataBitSetTest implements PinotBuffersAfterMethodCheckRule { private static final Random RANDOM = new Random(); private static final int NUM_ITERATIONS = 1000; diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/util/PinotDataBitSetV2Test.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/util/PinotDataBitSetV2Test.java index 442249cd7f46..f916aa6bf63a 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/util/PinotDataBitSetV2Test.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/util/PinotDataBitSetV2Test.java @@ -18,14 +18,18 @@ */ package org.apache.pinot.segment.local.io.util; +import java.io.IOException; import java.nio.ByteOrder; import java.util.Random; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.spi.memory.PinotDataBuffer; import org.testng.Assert; import org.testng.annotations.Test; -public class PinotDataBitSetV2Test { +public class PinotDataBitSetV2Test implements PinotBuffersAfterMethodCheckRule { + + PinotDataBuffer _bitPackedBuffer; private void batchRead(PinotDataBitSetV2 bitset, int startDocId, int batchLength, int[] unpacked, int[] forwardIndex) { @@ -49,128 +53,129 @@ public void testBit2Encoded() int numBitsPerValue = PinotDataBitSet.getNumBitsPerValue(cardinality - 1); int bitPackedBufferSize = (rows * numBitsPerValue + Byte.SIZE - 1) / Byte.SIZE; - PinotDataBitSetV2 bitSet = getEmptyBitSet(bitPackedBufferSize, numBitsPerValue); + try (PinotDataBitSetV2 bitSet = getEmptyBitSet(bitPackedBufferSize, numBitsPerValue)) { - Assert.assertEquals(2, numBitsPerValue); - Assert.assertTrue(bitSet instanceof PinotDataBitSetV2.Bit2Encoded); + Assert.assertEquals(2, numBitsPerValue); + Assert.assertTrue(bitSet instanceof PinotDataBitSetV2.Bit2Encoded); - for (int i = 0; i < rows; i++) { - bitSet.writeInt(i, forwardIndex[i]); - } + for (int i = 0; i < rows; i++) { + bitSet.writeInt(i, forwardIndex[i]); + } - // test single read API for sequential consecutive - for (int i = 0; i < rows; i++) { - int unpacked = bitSet.readInt(i); - Assert.assertEquals(forwardIndex[i], unpacked); - } + // test single read API for sequential consecutive + for (int i = 0; i < rows; i++) { + int unpacked = bitSet.readInt(i); + Assert.assertEquals(forwardIndex[i], unpacked); + } - // for each batch: - // 3 aligned reads at 4-byte boundary to unpack 16 integers after each read -- 48 integers unpacked - // followed by reading the next byte to unpack 2 integers from first 4 bits - int batchLength = 50; - int[] unpacked = new int[batchLength]; - int startDocId; - for (startDocId = 0; startDocId < rows; startDocId += 50) { + // for each batch: + // 3 aligned reads at 4-byte boundary to unpack 16 integers after each read -- 48 integers unpacked + // followed by reading the next byte to unpack 2 integers from first 4 bits + int batchLength = 50; + int[] unpacked = new int[batchLength]; + int startDocId; + for (startDocId = 0; startDocId < rows; startDocId += 50) { + bitSet.readInt(startDocId, batchLength, unpacked); + for (int i = 0; i < batchLength; i++) { + Assert.assertEquals(forwardIndex[startDocId + i], unpacked[i]); + } + } + + // 3 aligned reads at 4-byte boundary to unpack 16 integers after each read -- 48 integers unpacked + // followed by reading the next 2 bytes to unpack 8 integers + batchLength = 56; + unpacked = new int[batchLength]; + startDocId = 1; + batchRead(bitSet, startDocId, batchLength, unpacked, forwardIndex); + + // 3 aligned reads at 4-byte boundary to unpack 16 integers after each read -- 48 integers unpacked + // followed by reading the next 2 bytes to unpack 8 integers + // followed by reading the next byte to unpack 4 integers + batchLength = 60; + unpacked = new int[batchLength]; + startDocId = 20; + batchRead(bitSet, startDocId, batchLength, unpacked, forwardIndex); + + // 3 aligned reads at 4-byte boundary to unpack 16 integers after each read -- 48 integers unpacked + // followed by reading the next 2 bytes to unpack 8 integers + // followed by reading the next byte to unpack 4 integers + // followed by reading the next byte to unpack 1 integer from first 2 bits + batchLength = 61; + unpacked = new int[batchLength]; + startDocId = 20; + batchRead(bitSet, startDocId, batchLength, unpacked, forwardIndex); + + // 3 aligned reads at 4-byte boundary to unpack 16 integers after each read -- 48 integers unpacked + // followed by reading the next 2 bytes to unpack 8 integers + // followed by reading the next byte to unpack 4 integers + // followed by reading the next byte to unpack 2 integers from first 4 bits + batchLength = 62; + unpacked = new int[batchLength]; + startDocId = 20; + batchRead(bitSet, startDocId, batchLength, unpacked, forwardIndex); + + // 3 aligned reads at 4-byte boundary to unpack 16 integers after each read -- 48 integers unpacked + // followed by reading the next 2 bytes to unpack 8 integers + // followed by reading the next byte to unpack 4 integers + // followed by reading the next byte to unpack 6 integers from first 6 bits + batchLength = 63; + unpacked = new int[batchLength]; + startDocId = 20; + batchRead(bitSet, startDocId, batchLength, unpacked, forwardIndex); + + // for each batch: + // unaligned read on the first byte to unpack 3 integers + // followed by 3 aligned reads at byte boundary to unpack 4 integers after each read -- 12 integers unpacked + // followed by reading the next byte to unpack 2 integer from first 4 bits + // 3 + 12 + 2 = 17 unpacked integers + batchLength = 17; + unpacked = new int[batchLength]; + startDocId = 1; bitSet.readInt(startDocId, batchLength, unpacked); for (int i = 0; i < batchLength; i++) { Assert.assertEquals(forwardIndex[startDocId + i], unpacked[i]); } - } - // 3 aligned reads at 4-byte boundary to unpack 16 integers after each read -- 48 integers unpacked - // followed by reading the next 2 bytes to unpack 8 integers - batchLength = 56; - unpacked = new int[batchLength]; - startDocId = 1; - batchRead(bitSet, startDocId, batchLength, unpacked, forwardIndex); - - // 3 aligned reads at 4-byte boundary to unpack 16 integers after each read -- 48 integers unpacked - // followed by reading the next 2 bytes to unpack 8 integers - // followed by reading the next byte to unpack 4 integers - batchLength = 60; - unpacked = new int[batchLength]; - startDocId = 20; - batchRead(bitSet, startDocId, batchLength, unpacked, forwardIndex); - - // 3 aligned reads at 4-byte boundary to unpack 16 integers after each read -- 48 integers unpacked - // followed by reading the next 2 bytes to unpack 8 integers - // followed by reading the next byte to unpack 4 integers - // followed by reading the next byte to unpack 1 integer from first 2 bits - batchLength = 61; - unpacked = new int[batchLength]; - startDocId = 20; - batchRead(bitSet, startDocId, batchLength, unpacked, forwardIndex); - - // 3 aligned reads at 4-byte boundary to unpack 16 integers after each read -- 48 integers unpacked - // followed by reading the next 2 bytes to unpack 8 integers - // followed by reading the next byte to unpack 4 integers - // followed by reading the next byte to unpack 2 integers from first 4 bits - batchLength = 62; - unpacked = new int[batchLength]; - startDocId = 20; - batchRead(bitSet, startDocId, batchLength, unpacked, forwardIndex); - - // 3 aligned reads at 4-byte boundary to unpack 16 integers after each read -- 48 integers unpacked - // followed by reading the next 2 bytes to unpack 8 integers - // followed by reading the next byte to unpack 4 integers - // followed by reading the next byte to unpack 6 integers from first 6 bits - batchLength = 63; - unpacked = new int[batchLength]; - startDocId = 20; - batchRead(bitSet, startDocId, batchLength, unpacked, forwardIndex); - - // for each batch: - // unaligned read on the first byte to unpack 3 integers - // followed by 3 aligned reads at byte boundary to unpack 4 integers after each read -- 12 integers unpacked - // followed by reading the next byte to unpack 2 integer from first 4 bits - // 3 + 12 + 2 = 17 unpacked integers - batchLength = 17; - unpacked = new int[batchLength]; - startDocId = 1; - bitSet.readInt(startDocId, batchLength, unpacked); - for (int i = 0; i < batchLength; i++) { - Assert.assertEquals(forwardIndex[startDocId + i], unpacked[i]); - } - - // unaligned read on the first byte to unpack 3 integers (bits 2 to 7) - // followed by 3 aligned reads at 4-byte boundary to unpack 16 integers after each read -- 48 integers unpacked - // followed by 1 aligned read at byte boundary to unpack 4 integers -- 4 integers unpacked - // followed by reading the next byte to unpack 3 integers from first 6 bits - // 3 + 48 + 4 + 3 = 58 unpacked integers - batchLength = 58; - unpacked = new int[batchLength]; - startDocId = 1; - bitSet.readInt(startDocId, batchLength, unpacked); - for (int i = 0; i < batchLength; i++) { - Assert.assertEquals(forwardIndex[startDocId + i], unpacked[i]); - } + // unaligned read on the first byte to unpack 3 integers (bits 2 to 7) + // followed by 3 aligned reads at 4-byte boundary to unpack 16 integers after each read -- 48 integers unpacked + // followed by 1 aligned read at byte boundary to unpack 4 integers -- 4 integers unpacked + // followed by reading the next byte to unpack 3 integers from first 6 bits + // 3 + 48 + 4 + 3 = 58 unpacked integers + batchLength = 58; + unpacked = new int[batchLength]; + startDocId = 1; + bitSet.readInt(startDocId, batchLength, unpacked); + for (int i = 0; i < batchLength; i++) { + Assert.assertEquals(forwardIndex[startDocId + i], unpacked[i]); + } - // for each batch: - // unaligned read on the first byte to unpack 2 integers (bits 4 to 7) - // followed by 3 aligned reads at 4-byte boundary to unpack 16 integers after each read -- 48 integers unpacked - // followed by 2 aligned reads at byte boundary to unpack 4 integers after each read -- 8 integers unpacked - // 3 + 48 + 8 = 58 unpacked integers - startDocId = 2; - unpacked = new int[batchLength]; - bitSet.readInt(startDocId, batchLength, unpacked); - for (int i = 0; i < batchLength; i++) { - Assert.assertEquals(forwardIndex[startDocId + i], unpacked[i]); - } + // for each batch: + // unaligned read on the first byte to unpack 2 integers (bits 4 to 7) + // followed by 3 aligned reads at 4-byte boundary to unpack 16 integers after each read -- 48 integers unpacked + // followed by 2 aligned reads at byte boundary to unpack 4 integers after each read -- 8 integers unpacked + // 3 + 48 + 8 = 58 unpacked integers + startDocId = 2; + unpacked = new int[batchLength]; + bitSet.readInt(startDocId, batchLength, unpacked); + for (int i = 0; i < batchLength; i++) { + Assert.assertEquals(forwardIndex[startDocId + i], unpacked[i]); + } - // for each batch: - // unaligned read on the first byte to unpack 1 integers (bits 6 to 7) - // followed by 3 aligned reads at 4-byte boundary to unpack 16 integers after each read -- 48 integers unpacked - // followed by 2 aligned reads at byte boundary to unpack 4 integers after each read -- 8 integers unpacked - // followed by reading the next byte to unpack 1 integer from first 2 bits - // 1 + 48 + 8 + 1 = 58 unpacked integers - startDocId = 3; - unpacked = new int[batchLength]; - bitSet.readInt(startDocId, batchLength, unpacked); - for (int i = 0; i < batchLength; i++) { - Assert.assertEquals(forwardIndex[startDocId + i], unpacked[i]); + // for each batch: + // unaligned read on the first byte to unpack 1 integers (bits 6 to 7) + // followed by 3 aligned reads at 4-byte boundary to unpack 16 integers after each read -- 48 integers unpacked + // followed by 2 aligned reads at byte boundary to unpack 4 integers after each read -- 8 integers unpacked + // followed by reading the next byte to unpack 1 integer from first 2 bits + // 1 + 48 + 8 + 1 = 58 unpacked integers + startDocId = 3; + unpacked = new int[batchLength]; + bitSet.readInt(startDocId, batchLength, unpacked); + for (int i = 0; i < batchLength; i++) { + Assert.assertEquals(forwardIndex[startDocId + i], unpacked[i]); + } + } finally { + closeBuffer(); } - - bitSet.close(); } @Test @@ -187,97 +192,98 @@ public void testBit4Encoded() int numBitsPerValue = PinotDataBitSet.getNumBitsPerValue(cardinality - 1); int bitPackedBufferSize = (rows * numBitsPerValue + Byte.SIZE - 1) / Byte.SIZE; - PinotDataBitSetV2 bitSet = getEmptyBitSet(bitPackedBufferSize, numBitsPerValue); + try (PinotDataBitSetV2 bitSet = getEmptyBitSet(bitPackedBufferSize, numBitsPerValue)) { - Assert.assertEquals(4, numBitsPerValue); - Assert.assertTrue(bitSet instanceof PinotDataBitSetV2.Bit4Encoded); + Assert.assertEquals(4, numBitsPerValue); + Assert.assertTrue(bitSet instanceof PinotDataBitSetV2.Bit4Encoded); - for (int i = 0; i < rows; i++) { - bitSet.writeInt(i, forwardIndex[i]); - } + for (int i = 0; i < rows; i++) { + bitSet.writeInt(i, forwardIndex[i]); + } - // test single read API for sequential consecutive - for (int i = 0; i < rows; i++) { - int unpacked = bitSet.readInt(i); - Assert.assertEquals(forwardIndex[i], unpacked); - } + // test single read API for sequential consecutive + for (int i = 0; i < rows; i++) { + int unpacked = bitSet.readInt(i); + Assert.assertEquals(forwardIndex[i], unpacked); + } - // test array API for sequential consecutive + // test array API for sequential consecutive + + // for each batch: do a combination of aligned and unaligned reads + // 6 aligned reads at 4-byte boundary to unpack 8 integers after each read -- 48 integers unpacked + // followed by reading the next byte to unpack 2 integers + int batchLength = 50; + int[] unpacked = new int[batchLength]; + int startDocId; + for (startDocId = 0; startDocId < rows; startDocId += batchLength) { + bitSet.readInt(startDocId, batchLength, unpacked); + for (int i = 0; i < batchLength; i++) { + Assert.assertEquals(forwardIndex[startDocId + i], unpacked[i]); + } + } - // for each batch: do a combination of aligned and unaligned reads - // 6 aligned reads at 4-byte boundary to unpack 8 integers after each read -- 48 integers unpacked - // followed by reading the next byte to unpack 2 integers - int batchLength = 50; - int[] unpacked = new int[batchLength]; - int startDocId; - for (startDocId = 0; startDocId < rows; startDocId += batchLength) { + // 12 aligned reads at 4-byte boundary to unpack 8 integers after each read -- 96 integers unpacked + batchLength = 96; + unpacked = new int[batchLength]; + startDocId = 19; bitSet.readInt(startDocId, batchLength, unpacked); for (int i = 0; i < batchLength; i++) { Assert.assertEquals(forwardIndex[startDocId + i], unpacked[i]); } - } - // 12 aligned reads at 4-byte boundary to unpack 8 integers after each read -- 96 integers unpacked - batchLength = 96; - unpacked = new int[batchLength]; - startDocId = 19; - bitSet.readInt(startDocId, batchLength, unpacked); - for (int i = 0; i < batchLength; i++) { - Assert.assertEquals(forwardIndex[startDocId + i], unpacked[i]); - } - - // only a single unaligned read from the middle of a byte (44th bit) - batchLength = 1; - startDocId = 21; - bitSet.readInt(startDocId, batchLength, unpacked); - Assert.assertEquals(forwardIndex[startDocId], unpacked[0]); + // only a single unaligned read from the middle of a byte (44th bit) + batchLength = 1; + startDocId = 21; + bitSet.readInt(startDocId, batchLength, unpacked); + Assert.assertEquals(forwardIndex[startDocId], unpacked[0]); + + // unaligned read within a byte to unpack an integer from bits 4 to 7 + // followed by 2 aligned reads at 4-byte boundary to unpack 8 integers after each read -- unpacked 16 integers + // followed by 1 aligned read at 2-byte boundary to unpack 4 integers + // followed by 1 aligned read at byte boundary to unpack 2 integers + // followed by reading the next byte to unpack integer from first 4 bits + // 1 + 16 + 4 + 2 + 1 = 24 unpacked integers + startDocId = 1; + batchLength = 24; + unpacked = new int[batchLength]; + bitSet.readInt(startDocId, batchLength, unpacked); + for (int i = 0; i < batchLength; i++) { + Assert.assertEquals(forwardIndex[startDocId + i], unpacked[i]); + } - // unaligned read within a byte to unpack an integer from bits 4 to 7 - // followed by 2 aligned reads at 4-byte boundary to unpack 8 integers after each read -- unpacked 16 integers - // followed by 1 aligned read at 2-byte boundary to unpack 4 integers - // followed by 1 aligned read at byte boundary to unpack 2 integers - // followed by reading the next byte to unpack integer from first 4 bits - // 1 + 16 + 4 + 2 + 1 = 24 unpacked integers - startDocId = 1; - batchLength = 24; - unpacked = new int[batchLength]; - bitSet.readInt(startDocId, batchLength, unpacked); - for (int i = 0; i < batchLength; i++) { - Assert.assertEquals(forwardIndex[startDocId + i], unpacked[i]); - } + // unaligned read within a byte to unpack an integer from bits 4 to 7 + // 1 aligned read at 2-byte boundary to unpack 4 integers + // followed by 1 aligned read at byte boundary to unpack 2 integers + // followed by reading the next byte to unpack integer from first 4 bits + // 1 + 4 + 2 + 1 = 8 unpacked integers + startDocId = 1; + batchLength = 8; + unpacked = new int[batchLength]; + bitSet.readInt(startDocId, batchLength, unpacked); + for (int i = 0; i < batchLength; i++) { + Assert.assertEquals(forwardIndex[startDocId + i], unpacked[i]); + } - // unaligned read within a byte to unpack an integer from bits 4 to 7 - // 1 aligned read at 2-byte boundary to unpack 4 integers - // followed by 1 aligned read at byte boundary to unpack 2 integers - // followed by reading the next byte to unpack integer from first 4 bits - // 1 + 4 + 2 + 1 = 8 unpacked integers - startDocId = 1; - batchLength = 8; - unpacked = new int[batchLength]; - bitSet.readInt(startDocId, batchLength, unpacked); - for (int i = 0; i < batchLength; i++) { - Assert.assertEquals(forwardIndex[startDocId + i], unpacked[i]); - } + // 1 aligned read at 2-byte boundary to unpack 4 integers + // followed by 1 aligned read at byte boundary to unpack 2 integers + // followed by reading the next byte to unpack integer from first 4 bits + // 4 + 2 + 1 = 7 unpacked integers + startDocId = 4; + batchLength = 7; + unpacked = new int[batchLength]; + bitSet.readInt(startDocId, batchLength, unpacked); + for (int i = 0; i < batchLength; i++) { + Assert.assertEquals(forwardIndex[startDocId + i], unpacked[i]); + } - // 1 aligned read at 2-byte boundary to unpack 4 integers - // followed by 1 aligned read at byte boundary to unpack 2 integers - // followed by reading the next byte to unpack integer from first 4 bits - // 4 + 2 + 1 = 7 unpacked integers - startDocId = 4; - batchLength = 7; - unpacked = new int[batchLength]; - bitSet.readInt(startDocId, batchLength, unpacked); - for (int i = 0; i < batchLength; i++) { - Assert.assertEquals(forwardIndex[startDocId + i], unpacked[i]); + // test bulk API for sequential but not necessarily consecutive + testBulkSequentialWithGaps(bitSet, 1, 50, -1, forwardIndex); + testBulkSequentialWithGaps(bitSet, 5, 57, 4, forwardIndex); + testBulkSequentialWithGaps(bitSet, 17, 109, 19, forwardIndex); + testBulkSequentialWithGaps(bitSet, 17, 1, 19, forwardIndex); + } finally { + closeBuffer(); } - - // test bulk API for sequential but not necessarily consecutive - testBulkSequentialWithGaps(bitSet, 1, 50, -1, forwardIndex); - testBulkSequentialWithGaps(bitSet, 5, 57, 4, forwardIndex); - testBulkSequentialWithGaps(bitSet, 17, 109, 19, forwardIndex); - testBulkSequentialWithGaps(bitSet, 17, 1, 19, forwardIndex); - - bitSet.close(); } @Test @@ -294,59 +300,60 @@ public void testBit8Encoded() int numBitsPerValue = PinotDataBitSet.getNumBitsPerValue(cardinality - 1); int bitPackedBufferSize = (rows * numBitsPerValue + Byte.SIZE - 1) / Byte.SIZE; - PinotDataBitSetV2 bitSet = getEmptyBitSet(bitPackedBufferSize, numBitsPerValue); + try (PinotDataBitSetV2 bitSet = getEmptyBitSet(bitPackedBufferSize, numBitsPerValue)) { - Assert.assertEquals(8, numBitsPerValue); - Assert.assertTrue(bitSet instanceof PinotDataBitSetV2.Bit8Encoded); + Assert.assertEquals(8, numBitsPerValue); + Assert.assertTrue(bitSet instanceof PinotDataBitSetV2.Bit8Encoded); - for (int i = 0; i < rows; i++) { - bitSet.writeInt(i, forwardIndex[i]); - } + for (int i = 0; i < rows; i++) { + bitSet.writeInt(i, forwardIndex[i]); + } - // test single read API for sequential consecutive - for (int i = 0; i < rows; i++) { - int unpacked = bitSet.readInt(i); - Assert.assertEquals(forwardIndex[i], unpacked); - } + // test single read API for sequential consecutive + for (int i = 0; i < rows; i++) { + int unpacked = bitSet.readInt(i); + Assert.assertEquals(forwardIndex[i], unpacked); + } - // test array API for sequential consecutive + // test array API for sequential consecutive + + // for each batch: + // 12 aligned reads at 4-byte boundary to unpack 4 integers after each read -- 48 integers unpacked + // followed by reading the next 2 bytes to unpack 2 integers + int batchLength = 50; + int[] unpacked = new int[batchLength]; + int startDocId; + for (startDocId = 0; startDocId < rows; startDocId += batchLength) { + bitSet.readInt(startDocId, batchLength, unpacked); + for (int i = 0; i < batchLength; i++) { + Assert.assertEquals(forwardIndex[startDocId + i], unpacked[i]); + } + } - // for each batch: - // 12 aligned reads at 4-byte boundary to unpack 4 integers after each read -- 48 integers unpacked - // followed by reading the next 2 bytes to unpack 2 integers - int batchLength = 50; - int[] unpacked = new int[batchLength]; - int startDocId; - for (startDocId = 0; startDocId < rows; startDocId += batchLength) { + // for each batch: + // 24 aligned reads at 4-byte boundary to unpack 4 integers after each read -- 96 integers unpacked + batchLength = 96; + unpacked = new int[batchLength]; + startDocId = 7; bitSet.readInt(startDocId, batchLength, unpacked); for (int i = 0; i < batchLength; i++) { Assert.assertEquals(forwardIndex[startDocId + i], unpacked[i]); } - } - // for each batch: - // 24 aligned reads at 4-byte boundary to unpack 4 integers after each read -- 96 integers unpacked - batchLength = 96; - unpacked = new int[batchLength]; - startDocId = 7; - bitSet.readInt(startDocId, batchLength, unpacked); - for (int i = 0; i < batchLength; i++) { - Assert.assertEquals(forwardIndex[startDocId + i], unpacked[i]); + // unaligned spill over + startDocId = 19; + batchLength = 3; + bitSet.readInt(startDocId, batchLength, unpacked); + Assert.assertEquals(forwardIndex[startDocId], unpacked[0]); + + // test bulk API for sequential but not necessarily consecutive + testBulkSequentialWithGaps(bitSet, 1, 50, -1, forwardIndex); + testBulkSequentialWithGaps(bitSet, 5, 57, 4, forwardIndex); + testBulkSequentialWithGaps(bitSet, 17, 109, 19, forwardIndex); + testBulkSequentialWithGaps(bitSet, 17, 1, 19, forwardIndex); + } finally { + closeBuffer(); } - - // unaligned spill over - startDocId = 19; - batchLength = 3; - bitSet.readInt(startDocId, batchLength, unpacked); - Assert.assertEquals(forwardIndex[startDocId], unpacked[0]); - - // test bulk API for sequential but not necessarily consecutive - testBulkSequentialWithGaps(bitSet, 1, 50, -1, forwardIndex); - testBulkSequentialWithGaps(bitSet, 5, 57, 4, forwardIndex); - testBulkSequentialWithGaps(bitSet, 17, 109, 19, forwardIndex); - testBulkSequentialWithGaps(bitSet, 17, 1, 19, forwardIndex); - - bitSet.close(); } @Test @@ -363,7 +370,7 @@ public void testBit16Encoded() int numBitsPerValue = PinotDataBitSet.getNumBitsPerValue(cardinality - 1); int bitPackedBufferSize = (rows * numBitsPerValue + Byte.SIZE - 1) / Byte.SIZE; - PinotDataBitSetV2 bitSet = getEmptyBitSet(bitPackedBufferSize, numBitsPerValue); + try (PinotDataBitSetV2 bitSet = getEmptyBitSet(bitPackedBufferSize, numBitsPerValue)) { Assert.assertEquals(16, numBitsPerValue); Assert.assertTrue(bitSet instanceof PinotDataBitSetV2.Bit16Encoded); @@ -413,8 +420,9 @@ public void testBit16Encoded() testBulkSequentialWithGaps(bitSet, 5, 57, 4, forwardIndex); testBulkSequentialWithGaps(bitSet, 17, 109, 19, forwardIndex); testBulkSequentialWithGaps(bitSet, 17, 1, 19, forwardIndex); - - bitSet.close(); + } finally { + closeBuffer(); + } } private void testBulkSequentialWithGaps(PinotDataBitSetV2 bitset, int gaps, int batchLength, int startDocId, @@ -434,10 +442,18 @@ private void testBulkSequentialWithGaps(PinotDataBitSetV2 bitset, int gaps, int } private PinotDataBitSetV2 getEmptyBitSet(int size, int numBitsPerValue) { - PinotDataBuffer bitPackedBuffer = PinotDataBuffer.allocateDirect(size, ByteOrder.BIG_ENDIAN, null); + _bitPackedBuffer = PinotDataBuffer.allocateDirect(size, ByteOrder.BIG_ENDIAN, null); for (int i = 0; i < size; i++) { - bitPackedBuffer.readFrom(0, new byte[size]); + _bitPackedBuffer.readFrom(0, new byte[size]); + } + return PinotDataBitSetV2.createBitSet(_bitPackedBuffer, numBitsPerValue); + } + + private void closeBuffer() + throws IOException { + if (_bitPackedBuffer != null) { + _bitPackedBuffer.close(); + _bitPackedBuffer = null; } - return PinotDataBitSetV2.createBitSet(bitPackedBuffer, numBitsPerValue); } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/util/VarLengthValueReaderWriterTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/util/VarLengthValueReaderWriterTest.java index 23f15001b62b..a76295690461 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/util/VarLengthValueReaderWriterTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/util/VarLengthValueReaderWriterTest.java @@ -22,6 +22,7 @@ import java.io.IOException; import org.apache.commons.io.FileUtils; import org.apache.commons.lang3.RandomStringUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.spi.memory.PinotDataBuffer; import org.testng.annotations.AfterClass; import org.testng.annotations.BeforeClass; @@ -35,7 +36,7 @@ /** * Unit test for {@link VarLengthValueReader} and {@link VarLengthValueWriter}. */ -public class VarLengthValueReaderWriterTest { +public class VarLengthValueReaderWriterTest implements PinotBuffersAfterMethodCheckRule { private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(), "VarLengthValueReaderWriterTest"); private static final int MAX_STRING_LENGTH = 200; private static final int NUM_VALUES = 1000; diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/writer/impl/MmapMemoryManagerFileCleanupTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/writer/impl/MmapMemoryManagerFileCleanupTest.java index b70f133a2c60..313d3ee67a74 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/writer/impl/MmapMemoryManagerFileCleanupTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/writer/impl/MmapMemoryManagerFileCleanupTest.java @@ -20,6 +20,7 @@ import java.io.File; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.spi.memory.PinotDataBufferMemoryManager; import org.testng.Assert; import org.testng.annotations.AfterClass; @@ -27,7 +28,7 @@ import org.testng.annotations.Test; -public class MmapMemoryManagerFileCleanupTest { +public class MmapMemoryManagerFileCleanupTest implements PinotBuffersAfterMethodCheckRule { private String _tmpDir; @BeforeClass diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/writer/impl/MmapMemoryManagerTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/writer/impl/MmapMemoryManagerTest.java index acc5b17d12b6..ae007b77ff77 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/writer/impl/MmapMemoryManagerTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/writer/impl/MmapMemoryManagerTest.java @@ -23,7 +23,9 @@ import java.util.Arrays; import java.util.Comparator; import java.util.List; +import java.util.stream.Collectors; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.spi.memory.PinotDataBuffer; import org.apache.pinot.segment.spi.memory.PinotDataBufferMemoryManager; import org.testng.Assert; @@ -34,7 +36,7 @@ import org.testng.annotations.Test; -public class MmapMemoryManagerTest { +public class MmapMemoryManagerTest implements PinotBuffersAfterMethodCheckRule { private String _tmpDir; @@ -66,47 +68,44 @@ public void afterMethodAssertions() { public void testLargeBlocks() throws Exception { final String segmentName = "someSegment"; - PinotDataBufferMemoryManager memoryManager = new MmapMemoryManager(_tmpDir, segmentName); - final long s1 = 2 * MmapMemoryManager.getDefaultFileLength(); - final long s2 = 1000; - final String col1 = "col1"; - final String col2 = "col2"; - final byte value = 34; - PinotDataBuffer buf1 = memoryManager.allocate(s1, col1); - - // Verify that we can write to and read from the buffer - ByteBuffer b1 = buf1.toDirectByteBuffer(0, (int) s1); - b1.putLong(0, s1); - Assert.assertEquals(b1.getLong(0), s1); - b1.put((int) s1 - 1, value); - Assert.assertEquals(b1.get((int) s1 - 1), value); - - PinotDataBuffer buf2 = memoryManager.allocate(s2, col2); - ByteBuffer b2 = buf2.toDirectByteBuffer(0, (int) s2); - b2.putLong(0, s2); - Assert.assertEquals(b2.getLong(0), s2); - - File dir = new File(_tmpDir); - File[] files = dir.listFiles(); - Assert.assertEquals(files.length, 2); - - Arrays.sort(files, new Comparator() { - @Override - public int compare(File o1, File o2) { - return o1.getName().compareTo(o2.getName()); - } - }); - - String fileName = files[0].getName(); - Assert.assertTrue(fileName.contains(segmentName)); - - fileName = files[1].getName(); - Assert.assertTrue(fileName.contains(segmentName)); - - buf1.close(); - buf2.close(); - - memoryManager.close(); + try ( + PinotDataBufferMemoryManager memoryManager = new MmapMemoryManager(_tmpDir, segmentName)) { + final long s1 = 2 * MmapMemoryManager.getDefaultFileLength(); + final long s2 = 1000; + final String col1 = "col1"; + final String col2 = "col2"; + final byte value = 34; + PinotDataBuffer buf1 = memoryManager.allocate(s1, col1); + + // Verify that we can write to and read from the buffer + ByteBuffer b1 = buf1.toDirectByteBuffer(0, (int) s1); + b1.putLong(0, s1); + Assert.assertEquals(b1.getLong(0), s1); + b1.put((int) s1 - 1, value); + Assert.assertEquals(b1.get((int) s1 - 1), value); + + PinotDataBuffer buf2 = memoryManager.allocate(s2, col2); + ByteBuffer b2 = buf2.toDirectByteBuffer(0, (int) s2); + b2.putLong(0, s2); + Assert.assertEquals(b2.getLong(0), s2); + + File dir = new File(_tmpDir); + File[] files = dir.listFiles(); + Assert.assertEquals(files.length, 2); + + Arrays.sort(files, new Comparator() { + @Override + public int compare(File o1, File o2) { + return o1.getName().compareTo(o2.getName()); + } + }); + + String fileName = files[0].getName(); + Assert.assertTrue(fileName.contains(segmentName)); + + fileName = files[1].getName(); + Assert.assertTrue(fileName.contains(segmentName)); + } List allocationContexts = PinotDataBuffer.getBufferInfo(); Assert.assertEquals(allocationContexts.size(), 0); @@ -117,34 +116,28 @@ public int compare(File o1, File o2) { public void testSmallBlocksForSameColumn() throws Exception { final String segmentName = "someSegment"; - PinotDataBufferMemoryManager memoryManager = new MmapMemoryManager(_tmpDir, segmentName); - final long s1 = 500; - final long s2 = 1000; - final String col1 = "col1"; - PinotDataBuffer buf1 = memoryManager.allocate(s1, col1); - - PinotDataBuffer buf2 = memoryManager.allocate(s2, col1); - - ByteBuffer b1 = buf1.toDirectByteBuffer(0, (int) s1); - b1.putLong(0, s1); - - ByteBuffer b2 = buf2.toDirectByteBuffer(0, (int) s2); - b2.putLong(0, s2); + File dir = new File(_tmpDir); + try (PinotDataBufferMemoryManager memoryManager = new MmapMemoryManager(_tmpDir, segmentName)) { + final long s1 = 500; + final long s2 = 1000; + final String col1 = "col1"; + PinotDataBuffer buf1 = memoryManager.allocate(s1, col1); - Assert.assertEquals(b1.getLong(0), s1); - Assert.assertEquals(b2.getLong(0), s2); + PinotDataBuffer buf2 = memoryManager.allocate(s2, col1); - File dir = new File(_tmpDir); - Assert.assertEquals(dir.listFiles().length, 1); + ByteBuffer b1 = buf1.toDirectByteBuffer(0, (int) s1); + b1.putLong(0, s1); - buf1.close(); - buf2.close(); + ByteBuffer b2 = buf2.toDirectByteBuffer(0, (int) s2); + b2.putLong(0, s2); - memoryManager.close(); + Assert.assertEquals(b1.getLong(0), s1); + Assert.assertEquals(b2.getLong(0), s2); + Assert.assertEquals(dir.listFiles().length, 1); + } List allocationContexts = PinotDataBuffer.getBufferInfo(); Assert.assertEquals(allocationContexts.size(), 0); - Assert.assertEquals(dir.listFiles().length, 0); } @@ -152,46 +145,48 @@ public void testSmallBlocksForSameColumn() public void testCornerConditions() throws Exception { final String segmentName = "someSegment"; - PinotDataBufferMemoryManager memoryManager = new MmapMemoryManager(_tmpDir, segmentName); - final long s1 = MmapMemoryManager.getDefaultFileLength() - 1; - final long s2 = 1; - final long s3 = 100 * 1024 * 1024; - final String colName = "col"; - final byte v1 = 56; - final byte v2 = 11; - final byte v3 = 32; - - // Allocate a buffer 1 less than the default file length, and write the last byte of the buffer. - PinotDataBuffer b1 = memoryManager.allocate(s1, colName); - ByteBuffer bb1 = b1.toDirectByteBuffer(0, (int) s1); - bb1.put((int) s1 - 1, v1); - - // Allocate another buffer that is 1 byte in size, should be in the same file. - // Write a value in the byte. - PinotDataBuffer b2 = memoryManager.allocate(s2, colName); - ByteBuffer bb2 = b2.toDirectByteBuffer(0, (int) s2); - bb2.put((int) s2 - 1, v2); - - // Verify that there is only one file. - File dir = new File(_tmpDir); - Assert.assertEquals(dir.listFiles().length, 1); - - // Now allocate another buffer that will open a second file, write a value in the first byte of the buffer. - PinotDataBuffer b3 = memoryManager.allocate(s3, colName); - ByteBuffer bb3 = b3.toDirectByteBuffer(0, (int) s3); - bb3.put(0, v3); - - // Ensure that there are 2 files. - Assert.assertEquals(dir.listFiles().length, 2); - - // Make sure that the values written are preserved. - Assert.assertEquals(bb1.get((int) s1 - 1), v1); - Assert.assertEquals(bb2.get((int) s2 - 1), v2); - Assert.assertEquals(bb3.get(0), v3); - - memoryManager.close(); + try (PinotDataBufferMemoryManager memoryManager = new MmapMemoryManager(_tmpDir, segmentName)) { + final long s1 = MmapMemoryManager.getDefaultFileLength() - 1; + final long s2 = 1; + final long s3 = 100 * 1024 * 1024; + final String colName = "col"; + final byte v1 = 56; + final byte v2 = 11; + final byte v3 = 32; + + // Allocate a buffer 1 less than the default file length, and write the last byte of the buffer. + PinotDataBuffer b1 = memoryManager.allocate(s1, colName); + ByteBuffer bb1 = b1.toDirectByteBuffer(0, (int) s1); + bb1.put((int) s1 - 1, v1); + + // Allocate another buffer that is 1 byte in size, should be in the same file. + // Write a value in the byte. + PinotDataBuffer b2 = memoryManager.allocate(s2, colName); + ByteBuffer bb2 = b2.toDirectByteBuffer(0, (int) s2); + bb2.put((int) s2 - 1, v2); + + // Verify that there is only one file. + File dir = new File(_tmpDir); + Assert.assertEquals(dir.listFiles().length, 1); + + // Now allocate another buffer that will open a second file, write a value in the first byte of the buffer. + PinotDataBuffer b3 = memoryManager.allocate(s3, colName); + ByteBuffer bb3 = b3.toDirectByteBuffer(0, (int) s3); + bb3.put(0, v3); + + // Ensure that there are 2 files. + Assert.assertEquals(dir.listFiles().length, 2); + + // Make sure that the values written are preserved. + Assert.assertEquals(bb1.get((int) s1 - 1), v1); + Assert.assertEquals(bb2.get((int) s2 - 1), v2); + Assert.assertEquals(bb3.get(0), v3); + } List allocationContexts = PinotDataBuffer.getBufferInfo(); - Assert.assertEquals(allocationContexts.size(), 0); + if (allocationContexts.size() != 0) { + throw new AssertionError("Expected no allocation contexts but got: \n" + + allocationContexts.stream().collect(Collectors.joining("\n"))); + } } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/writer/impl/MutableOffHeapByteArrayStoreTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/writer/impl/MutableOffHeapByteArrayStoreTest.java index f74757943144..c316992c0342 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/writer/impl/MutableOffHeapByteArrayStoreTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/writer/impl/MutableOffHeapByteArrayStoreTest.java @@ -19,6 +19,7 @@ package org.apache.pinot.segment.local.io.writer.impl; import java.util.Arrays; +import org.apache.pinot.segment.local.PinotBuffersAfterClassCheckRule; import org.apache.pinot.segment.spi.memory.PinotDataBufferMemoryManager; import org.testng.Assert; import org.testng.annotations.AfterClass; @@ -26,7 +27,7 @@ import org.testng.annotations.Test; -public class MutableOffHeapByteArrayStoreTest { +public class MutableOffHeapByteArrayStoreTest implements PinotBuffersAfterClassCheckRule { private PinotDataBufferMemoryManager _memoryManager; private static final int ONE_GB = 1024 * 1024 * 1024; @@ -47,17 +48,17 @@ public void maxValueTest() throws Exception { int numArrays = 1024; int avgArrayLen = 32; - MutableOffHeapByteArrayStore store = - new MutableOffHeapByteArrayStore(_memoryManager, "stringColumn", numArrays, avgArrayLen); - final int arrSize = MutableOffHeapByteArrayStore.getStartSize(numArrays, avgArrayLen); - byte[] dataIn = new byte[arrSize - 4]; - for (int i = 0; i < dataIn.length; i++) { - dataIn[i] = (byte) (i % Byte.MAX_VALUE); + try (MutableOffHeapByteArrayStore store = + new MutableOffHeapByteArrayStore(_memoryManager, "stringColumn", numArrays, avgArrayLen)) { + final int arrSize = MutableOffHeapByteArrayStore.getStartSize(numArrays, avgArrayLen); + byte[] dataIn = new byte[arrSize - 4]; + for (int i = 0; i < dataIn.length; i++) { + dataIn[i] = (byte) (i % Byte.MAX_VALUE); + } + int index = store.add(dataIn); + byte[] dataOut = store.get(index); + Assert.assertTrue(Arrays.equals(dataIn, dataOut)); } - int index = store.add(dataIn); - byte[] dataOut = store.get(index); - Assert.assertTrue(Arrays.equals(dataIn, dataOut)); - store.close(); } @Test @@ -72,66 +73,66 @@ public void overflowTest() throws Exception { int numArrays = 1024; int avgArrayLen = 32; - MutableOffHeapByteArrayStore store = - new MutableOffHeapByteArrayStore(_memoryManager, "stringColumn", numArrays, avgArrayLen); - final int maxSize = MutableOffHeapByteArrayStore.getStartSize(numArrays, avgArrayLen) - 4; + try (MutableOffHeapByteArrayStore store = + new MutableOffHeapByteArrayStore(_memoryManager, "stringColumn", numArrays, avgArrayLen)) { + final int maxSize = MutableOffHeapByteArrayStore.getStartSize(numArrays, avgArrayLen) - 4; - byte[] b1 = new byte[3]; - for (int i = 0; i < b1.length; i++) { - b1[i] = (byte) i; - } + byte[] b1 = new byte[3]; + for (int i = 0; i < b1.length; i++) { + b1[i] = (byte) i; + } - byte[] b2 = new byte[maxSize]; - for (int i = 0; i < b2.length; i++) { - b2[i] = (byte) (i % Byte.MAX_VALUE); - } + byte[] b2 = new byte[maxSize]; + for (int i = 0; i < b2.length; i++) { + b2[i] = (byte) (i % Byte.MAX_VALUE); + } - // Add small array - final int i1 = store.add(b1); - Assert.assertTrue(Arrays.equals(store.get(i1), b1)); + // Add small array + final int i1 = store.add(b1); + Assert.assertTrue(Arrays.equals(store.get(i1), b1)); - // And now the larger one, should result in a new buffer - final int i2 = store.add(b2); - Assert.assertTrue(Arrays.equals(store.get(i2), b2)); + // And now the larger one, should result in a new buffer + final int i2 = store.add(b2); + Assert.assertTrue(Arrays.equals(store.get(i2), b2)); - // And now one more, should result in a new buffer but exact fit. - final int i3 = store.add(b2); - Assert.assertTrue(Arrays.equals(store.get(i3), b2)); + // And now one more, should result in a new buffer but exact fit. + final int i3 = store.add(b2); + Assert.assertTrue(Arrays.equals(store.get(i3), b2)); - // One more buffer when we add the small one again. - final int i4 = store.add(b1); - Assert.assertTrue(Arrays.equals(store.get(i4), b1)); + // One more buffer when we add the small one again. + final int i4 = store.add(b1); + Assert.assertTrue(Arrays.equals(store.get(i4), b1)); - // Test with one more 'get' to ensure that things have not changed. - Assert.assertTrue(Arrays.equals(store.get(i1), b1)); - Assert.assertTrue(Arrays.equals(store.get(i2), b2)); - Assert.assertTrue(Arrays.equals(store.get(i3), b2)); - Assert.assertTrue(Arrays.equals(store.get(i4), b1)); + // Test with one more 'get' to ensure that things have not changed. + Assert.assertTrue(Arrays.equals(store.get(i1), b1)); + Assert.assertTrue(Arrays.equals(store.get(i2), b2)); + Assert.assertTrue(Arrays.equals(store.get(i3), b2)); + Assert.assertTrue(Arrays.equals(store.get(i4), b1)); - byte[] b3 = new byte[5]; - for (int i = 0; i < b3.length; i++) { - b3[i] = (byte) (i + 1); - } + byte[] b3 = new byte[5]; + for (int i = 0; i < b3.length; i++) { + b3[i] = (byte) (i + 1); + } - int ix = -1; - final int iters = 1_000_000; + int ix = -1; + final int iters = 1_000_000; - // Now add the small one multiple times causing many additions. - for (int i = 0; i < iters; i++) { - if (ix == -1) { - ix = store.add(b3); + // Now add the small one multiple times causing many additions. + for (int i = 0; i < iters; i++) { + if (ix == -1) { + ix = store.add(b3); + } + store.add(b3); + } + for (int i = 0; i < iters; i++) { + Assert.assertTrue(Arrays.equals(store.get(ix++), b3)); } - store.add(b3); - } - for (int i = 0; i < iters; i++) { - Assert.assertTrue(Arrays.equals(store.get(ix++), b3)); - } - // Original values should still be good. - Assert.assertTrue(Arrays.equals(store.get(i1), b1)); - Assert.assertTrue(Arrays.equals(store.get(i2), b2)); - Assert.assertTrue(Arrays.equals(store.get(i3), b2)); - Assert.assertTrue(Arrays.equals(store.get(i4), b1)); - store.close(); + // Original values should still be good. + Assert.assertTrue(Arrays.equals(store.get(i1), b1)); + Assert.assertTrue(Arrays.equals(store.get(i2), b2)); + Assert.assertTrue(Arrays.equals(store.get(i3), b2)); + Assert.assertTrue(Arrays.equals(store.get(i4), b1)); + } } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/writer/impl/VarByteChunkSVForwardIndexWriterTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/writer/impl/VarByteChunkSVForwardIndexWriterTest.java index a9b8c34bb538..b2385edc2702 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/writer/impl/VarByteChunkSVForwardIndexWriterTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/writer/impl/VarByteChunkSVForwardIndexWriterTest.java @@ -30,6 +30,7 @@ import java.util.UUID; import java.util.stream.IntStream; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.segment.index.readers.forward.ChunkReaderContext; import org.apache.pinot.segment.local.segment.index.readers.forward.VarByteChunkSVForwardIndexReader; import org.apache.pinot.segment.spi.V1Constants; @@ -45,7 +46,7 @@ import static org.testng.Assert.assertEquals; -public class VarByteChunkSVForwardIndexWriterTest { +public class VarByteChunkSVForwardIndexWriterTest implements PinotBuffersAfterMethodCheckRule { private static final File OUTPUT_DIR = new File(FileUtils.getTempDirectory(), VarByteChunkSVForwardIndexWriterTest.class.getSimpleName()); diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/realtime/converter/RealtimeSegmentConverterTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/realtime/converter/RealtimeSegmentConverterTest.java index 06581baabdb6..a3ad1c177106 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/realtime/converter/RealtimeSegmentConverterTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/realtime/converter/RealtimeSegmentConverterTest.java @@ -35,6 +35,7 @@ import org.apache.commons.io.FileUtils; import org.apache.pinot.common.metadata.segment.SegmentZKMetadata; import org.apache.pinot.common.metrics.ServerMetrics; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentImpl; import org.apache.pinot.segment.local.indexsegment.mutable.MutableSegmentImpl; import org.apache.pinot.segment.local.io.writer.impl.DirectMemoryManager; @@ -79,7 +80,7 @@ import static org.testng.Assert.assertTrue; -public class RealtimeSegmentConverterTest { +public class RealtimeSegmentConverterTest implements PinotBuffersAfterMethodCheckRule { private static final String STRING_COLUMN1 = "string_col1"; private static final String STRING_COLUMN2 = "string_col2"; @@ -170,25 +171,29 @@ public void testNoRecordsIndexedRowMajorSegmentBuilder() // create mutable segment impl MutableSegmentImpl mutableSegmentImpl = new MutableSegmentImpl(realtimeSegmentConfigBuilder.build(), null); - - File outputDir = new File(tmpDir, "outputDir"); - SegmentZKPropsConfig segmentZKPropsConfig = new SegmentZKPropsConfig(); - segmentZKPropsConfig.setStartOffset("1"); - segmentZKPropsConfig.setEndOffset("100"); - RealtimeSegmentConverter converter = - new RealtimeSegmentConverter(mutableSegmentImpl, segmentZKPropsConfig, outputDir.getAbsolutePath(), schema, - tableNameWithType, tableConfig, segmentName, false); - converter.build(SegmentVersion.v3, null); - - File indexDir = new File(outputDir, segmentName); - SegmentMetadataImpl segmentMetadata = new SegmentMetadataImpl(indexDir); - assertEquals(segmentMetadata.getTotalDocs(), 0); - assertEquals(segmentMetadata.getTimeColumn(), DATE_TIME_COLUMN); - assertEquals(segmentMetadata.getTimeUnit(), TimeUnit.MILLISECONDS); - assertEquals(segmentMetadata.getStartTime(), segmentMetadata.getEndTime()); - assertTrue(segmentMetadata.getAllColumns().containsAll(schema.getColumnNames())); - assertEquals(segmentMetadata.getStartOffset(), "1"); - assertEquals(segmentMetadata.getEndOffset(), "100"); + try { + + File outputDir = new File(tmpDir, "outputDir"); + SegmentZKPropsConfig segmentZKPropsConfig = new SegmentZKPropsConfig(); + segmentZKPropsConfig.setStartOffset("1"); + segmentZKPropsConfig.setEndOffset("100"); + RealtimeSegmentConverter converter = + new RealtimeSegmentConverter(mutableSegmentImpl, segmentZKPropsConfig, outputDir.getAbsolutePath(), schema, + tableNameWithType, tableConfig, segmentName, false); + converter.build(SegmentVersion.v3, null); + + File indexDir = new File(outputDir, segmentName); + SegmentMetadataImpl segmentMetadata = new SegmentMetadataImpl(indexDir); + assertEquals(segmentMetadata.getTotalDocs(), 0); + assertEquals(segmentMetadata.getTimeColumn(), DATE_TIME_COLUMN); + assertEquals(segmentMetadata.getTimeUnit(), TimeUnit.MILLISECONDS); + assertEquals(segmentMetadata.getStartTime(), segmentMetadata.getEndTime()); + assertTrue(segmentMetadata.getAllColumns().containsAll(schema.getColumnNames())); + assertEquals(segmentMetadata.getStartOffset(), "1"); + assertEquals(segmentMetadata.getEndOffset(), "100"); + } finally { + mutableSegmentImpl.destroy(); + } } @Test @@ -238,38 +243,42 @@ public void test10RecordsIndexedRowMajorSegmentBuilder() // create mutable segment impl MutableSegmentImpl mutableSegmentImpl = new MutableSegmentImpl(realtimeSegmentConfigBuilder.build(), null); - List rows = generateTestData(); + try { + List rows = generateTestData(); - for (GenericRow row : rows) { - mutableSegmentImpl.index(row, null); - } + for (GenericRow row : rows) { + mutableSegmentImpl.index(row, null); + } - File outputDir = new File(tmpDir, "outputDir"); - SegmentZKPropsConfig segmentZKPropsConfig = new SegmentZKPropsConfig(); - segmentZKPropsConfig.setStartOffset("1"); - segmentZKPropsConfig.setEndOffset("100"); - RealtimeSegmentConverter converter = - new RealtimeSegmentConverter(mutableSegmentImpl, segmentZKPropsConfig, outputDir.getAbsolutePath(), schema, - tableNameWithType, tableConfig, segmentName, false); - converter.build(SegmentVersion.v3, null); - - File indexDir = new File(outputDir, segmentName); - SegmentMetadataImpl segmentMetadata = new SegmentMetadataImpl(indexDir); - assertEquals(segmentMetadata.getVersion(), SegmentVersion.v3); - assertEquals(segmentMetadata.getTotalDocs(), rows.size()); - assertEquals(segmentMetadata.getTimeColumn(), DATE_TIME_COLUMN); - assertEquals(segmentMetadata.getTimeUnit(), TimeUnit.MILLISECONDS); - - long expectedStartTime = (long) rows.get(0).getValue(DATE_TIME_COLUMN); - assertEquals(segmentMetadata.getStartTime(), expectedStartTime); - long expectedEndTime = (long) rows.get(rows.size() - 1).getValue(DATE_TIME_COLUMN); - assertEquals(segmentMetadata.getEndTime(), expectedEndTime); - - assertTrue(segmentMetadata.getAllColumns().containsAll(schema.getColumnNames())); - assertEquals(segmentMetadata.getStartOffset(), "1"); - assertEquals(segmentMetadata.getEndOffset(), "100"); - - testSegment(rows, indexDir, tableConfig, segmentMetadata); + File outputDir = new File(tmpDir, "outputDir"); + SegmentZKPropsConfig segmentZKPropsConfig = new SegmentZKPropsConfig(); + segmentZKPropsConfig.setStartOffset("1"); + segmentZKPropsConfig.setEndOffset("100"); + RealtimeSegmentConverter converter = + new RealtimeSegmentConverter(mutableSegmentImpl, segmentZKPropsConfig, outputDir.getAbsolutePath(), schema, + tableNameWithType, tableConfig, segmentName, false); + converter.build(SegmentVersion.v3, null); + + File indexDir = new File(outputDir, segmentName); + SegmentMetadataImpl segmentMetadata = new SegmentMetadataImpl(indexDir); + assertEquals(segmentMetadata.getVersion(), SegmentVersion.v3); + assertEquals(segmentMetadata.getTotalDocs(), rows.size()); + assertEquals(segmentMetadata.getTimeColumn(), DATE_TIME_COLUMN); + assertEquals(segmentMetadata.getTimeUnit(), TimeUnit.MILLISECONDS); + + long expectedStartTime = (long) rows.get(0).getValue(DATE_TIME_COLUMN); + assertEquals(segmentMetadata.getStartTime(), expectedStartTime); + long expectedEndTime = (long) rows.get(rows.size() - 1).getValue(DATE_TIME_COLUMN); + assertEquals(segmentMetadata.getEndTime(), expectedEndTime); + + assertTrue(segmentMetadata.getAllColumns().containsAll(schema.getColumnNames())); + assertEquals(segmentMetadata.getStartOffset(), "1"); + assertEquals(segmentMetadata.getEndOffset(), "100"); + + testSegment(rows, indexDir, tableConfig, segmentMetadata); + } finally { + mutableSegmentImpl.destroy(); + } } @Test @@ -314,25 +323,29 @@ public void testNoRecordsIndexedColumnMajorSegmentBuilder() // create mutable segment impl MutableSegmentImpl mutableSegmentImpl = new MutableSegmentImpl(realtimeSegmentConfigBuilder.build(), null); - - File outputDir = new File(tmpDir, "outputDir"); - SegmentZKPropsConfig segmentZKPropsConfig = new SegmentZKPropsConfig(); - segmentZKPropsConfig.setStartOffset("1"); - segmentZKPropsConfig.setEndOffset("100"); - RealtimeSegmentConverter converter = - new RealtimeSegmentConverter(mutableSegmentImpl, segmentZKPropsConfig, outputDir.getAbsolutePath(), schema, - tableNameWithType, tableConfig, segmentName, false); - converter.build(SegmentVersion.v3, null); - - File indexDir = new File(outputDir, segmentName); - SegmentMetadataImpl segmentMetadata = new SegmentMetadataImpl(indexDir); - assertEquals(segmentMetadata.getTotalDocs(), 0); - assertEquals(segmentMetadata.getTimeColumn(), DATE_TIME_COLUMN); - assertEquals(segmentMetadata.getTimeUnit(), TimeUnit.MILLISECONDS); - assertEquals(segmentMetadata.getStartTime(), segmentMetadata.getEndTime()); - assertTrue(segmentMetadata.getAllColumns().containsAll(schema.getColumnNames())); - assertEquals(segmentMetadata.getStartOffset(), "1"); - assertEquals(segmentMetadata.getEndOffset(), "100"); + try { + + File outputDir = new File(tmpDir, "outputDir"); + SegmentZKPropsConfig segmentZKPropsConfig = new SegmentZKPropsConfig(); + segmentZKPropsConfig.setStartOffset("1"); + segmentZKPropsConfig.setEndOffset("100"); + RealtimeSegmentConverter converter = + new RealtimeSegmentConverter(mutableSegmentImpl, segmentZKPropsConfig, outputDir.getAbsolutePath(), schema, + tableNameWithType, tableConfig, segmentName, false); + converter.build(SegmentVersion.v3, null); + + File indexDir = new File(outputDir, segmentName); + SegmentMetadataImpl segmentMetadata = new SegmentMetadataImpl(indexDir); + assertEquals(segmentMetadata.getTotalDocs(), 0); + assertEquals(segmentMetadata.getTimeColumn(), DATE_TIME_COLUMN); + assertEquals(segmentMetadata.getTimeUnit(), TimeUnit.MILLISECONDS); + assertEquals(segmentMetadata.getStartTime(), segmentMetadata.getEndTime()); + assertTrue(segmentMetadata.getAllColumns().containsAll(schema.getColumnNames())); + assertEquals(segmentMetadata.getStartOffset(), "1"); + assertEquals(segmentMetadata.getEndOffset(), "100"); + } finally { + mutableSegmentImpl.destroy(); + } } @Test @@ -382,38 +395,42 @@ public void test10RecordsIndexedColumnMajorSegmentBuilder() // create mutable segment impl MutableSegmentImpl mutableSegmentImpl = new MutableSegmentImpl(realtimeSegmentConfigBuilder.build(), null); - List rows = generateTestData(); + try { + List rows = generateTestData(); - for (GenericRow row : rows) { - mutableSegmentImpl.index(row, null); - } + for (GenericRow row : rows) { + mutableSegmentImpl.index(row, null); + } - File outputDir = new File(tmpDir, "outputDir"); - SegmentZKPropsConfig segmentZKPropsConfig = new SegmentZKPropsConfig(); - segmentZKPropsConfig.setStartOffset("1"); - segmentZKPropsConfig.setEndOffset("100"); - RealtimeSegmentConverter converter = - new RealtimeSegmentConverter(mutableSegmentImpl, segmentZKPropsConfig, outputDir.getAbsolutePath(), schema, - tableNameWithType, tableConfig, segmentName, false); - converter.build(SegmentVersion.v3, null); - - File indexDir = new File(outputDir, segmentName); - SegmentMetadataImpl segmentMetadata = new SegmentMetadataImpl(indexDir); - assertEquals(segmentMetadata.getVersion(), SegmentVersion.v3); - assertEquals(segmentMetadata.getTotalDocs(), rows.size()); - assertEquals(segmentMetadata.getTimeColumn(), DATE_TIME_COLUMN); - assertEquals(segmentMetadata.getTimeUnit(), TimeUnit.MILLISECONDS); - - long expectedStartTime = (long) rows.get(0).getValue(DATE_TIME_COLUMN); - assertEquals(segmentMetadata.getStartTime(), expectedStartTime); - long expectedEndTime = (long) rows.get(rows.size() - 1).getValue(DATE_TIME_COLUMN); - assertEquals(segmentMetadata.getEndTime(), expectedEndTime); - - assertTrue(segmentMetadata.getAllColumns().containsAll(schema.getColumnNames())); - assertEquals(segmentMetadata.getStartOffset(), "1"); - assertEquals(segmentMetadata.getEndOffset(), "100"); - - testSegment(rows, indexDir, tableConfig, segmentMetadata); + File outputDir = new File(tmpDir, "outputDir"); + SegmentZKPropsConfig segmentZKPropsConfig = new SegmentZKPropsConfig(); + segmentZKPropsConfig.setStartOffset("1"); + segmentZKPropsConfig.setEndOffset("100"); + RealtimeSegmentConverter converter = + new RealtimeSegmentConverter(mutableSegmentImpl, segmentZKPropsConfig, outputDir.getAbsolutePath(), schema, + tableNameWithType, tableConfig, segmentName, false); + converter.build(SegmentVersion.v3, null); + + File indexDir = new File(outputDir, segmentName); + SegmentMetadataImpl segmentMetadata = new SegmentMetadataImpl(indexDir); + assertEquals(segmentMetadata.getVersion(), SegmentVersion.v3); + assertEquals(segmentMetadata.getTotalDocs(), rows.size()); + assertEquals(segmentMetadata.getTimeColumn(), DATE_TIME_COLUMN); + assertEquals(segmentMetadata.getTimeUnit(), TimeUnit.MILLISECONDS); + + long expectedStartTime = (long) rows.get(0).getValue(DATE_TIME_COLUMN); + assertEquals(segmentMetadata.getStartTime(), expectedStartTime); + long expectedEndTime = (long) rows.get(rows.size() - 1).getValue(DATE_TIME_COLUMN); + assertEquals(segmentMetadata.getEndTime(), expectedEndTime); + + assertTrue(segmentMetadata.getAllColumns().containsAll(schema.getColumnNames())); + assertEquals(segmentMetadata.getStartOffset(), "1"); + assertEquals(segmentMetadata.getEndOffset(), "100"); + + testSegment(rows, indexDir, tableConfig, segmentMetadata); + } finally { + mutableSegmentImpl.destroy(); + } } private void testSegment(List rows, File indexDir, @@ -492,45 +509,48 @@ public void testOptimizeDictionaryTypeConversion(Object[] params) // create mutable segment impl MutableSegmentImpl mutableSegmentImpl = new MutableSegmentImpl(realtimeSegmentConfigBuilder.build(), null); - List rows = new ArrayList<>(); - for (int i = 0; i < 10; i++) { - GenericRow row = new GenericRow(); - row.putValue(STRING_COLUMN1, "string" + i * 10); // var length - row.putValue(STRING_COLUMN2, "string" + i % 10); // fixed length - row.putValue(DATE_TIME_COLUMN, 1697814309L + i); - rows.add(row); - } - for (GenericRow row : rows) { - mutableSegmentImpl.index(row, null); - } + try { + List rows = new ArrayList<>(); + for (int i = 0; i < 10; i++) { + GenericRow row = new GenericRow(); + row.putValue(STRING_COLUMN1, "string" + i * 10); // var length + row.putValue(STRING_COLUMN2, "string" + i % 10); // fixed length + row.putValue(DATE_TIME_COLUMN, 1697814309L + i); + rows.add(row); + } + for (GenericRow row : rows) { + mutableSegmentImpl.index(row, null); + } - File outputDir = new File(tmpDir, "outputDir"); - SegmentZKPropsConfig segmentZKPropsConfig = new SegmentZKPropsConfig(); - segmentZKPropsConfig.setStartOffset("1"); - segmentZKPropsConfig.setEndOffset("100"); - RealtimeSegmentConverter converter = - new RealtimeSegmentConverter(mutableSegmentImpl, segmentZKPropsConfig, outputDir.getAbsolutePath(), schema, - tableNameWithType, tableConfig, segmentName, false); - converter.build(SegmentVersion.v3, null); - - - File indexDir = new File(outputDir, segmentName); - SegmentMetadataImpl segmentMetadata = new SegmentMetadataImpl(indexDir); - assertEquals(segmentMetadata.getCrc(), params[1]); - - assertEquals(segmentMetadata.getVersion(), SegmentVersion.v3); - assertEquals(segmentMetadata.getTotalDocs(), rows.size()); - assertEquals(segmentMetadata.getTimeColumn(), DATE_TIME_COLUMN); - assertEquals(segmentMetadata.getTimeUnit(), TimeUnit.MILLISECONDS); - - long expectedStartTime = (long) rows.get(0).getValue(DATE_TIME_COLUMN); - assertEquals(segmentMetadata.getStartTime(), expectedStartTime); - long expectedEndTime = (long) rows.get(rows.size() - 1).getValue(DATE_TIME_COLUMN); - assertEquals(segmentMetadata.getEndTime(), expectedEndTime); - - assertTrue(segmentMetadata.getAllColumns().containsAll(schema.getColumnNames())); - assertEquals(segmentMetadata.getStartOffset(), "1"); - assertEquals(segmentMetadata.getEndOffset(), "100"); + File outputDir = new File(tmpDir, "outputDir"); + SegmentZKPropsConfig segmentZKPropsConfig = new SegmentZKPropsConfig(); + segmentZKPropsConfig.setStartOffset("1"); + segmentZKPropsConfig.setEndOffset("100"); + RealtimeSegmentConverter converter = + new RealtimeSegmentConverter(mutableSegmentImpl, segmentZKPropsConfig, outputDir.getAbsolutePath(), schema, + tableNameWithType, tableConfig, segmentName, false); + converter.build(SegmentVersion.v3, null); + + File indexDir = new File(outputDir, segmentName); + SegmentMetadataImpl segmentMetadata = new SegmentMetadataImpl(indexDir); + assertEquals(segmentMetadata.getCrc(), params[1]); + + assertEquals(segmentMetadata.getVersion(), SegmentVersion.v3); + assertEquals(segmentMetadata.getTotalDocs(), rows.size()); + assertEquals(segmentMetadata.getTimeColumn(), DATE_TIME_COLUMN); + assertEquals(segmentMetadata.getTimeUnit(), TimeUnit.MILLISECONDS); + + long expectedStartTime = (long) rows.get(0).getValue(DATE_TIME_COLUMN); + assertEquals(segmentMetadata.getStartTime(), expectedStartTime); + long expectedEndTime = (long) rows.get(rows.size() - 1).getValue(DATE_TIME_COLUMN); + assertEquals(segmentMetadata.getEndTime(), expectedEndTime); + + assertTrue(segmentMetadata.getAllColumns().containsAll(schema.getColumnNames())); + assertEquals(segmentMetadata.getStartOffset(), "1"); + assertEquals(segmentMetadata.getEndOffset(), "100"); + } finally { + mutableSegmentImpl.destroy(); + } } @DataProvider @@ -606,88 +626,93 @@ public void testSegmentBuilderWithReuse(boolean columnMajorSegmentBuilder, Strin // create mutable segment impl RealtimeLuceneTextIndexSearcherPool.init(1); RealtimeLuceneIndexRefreshManager.init(1, 10); + ImmutableSegmentImpl segmentFile = null; MutableSegmentImpl mutableSegmentImpl = new MutableSegmentImpl(realtimeSegmentConfigBuilder.build(), null); - List rows = generateTestDataForReusePath(); - - for (GenericRow row : rows) { - mutableSegmentImpl.index(row, null); - } + try { + List rows = generateTestDataForReusePath(); - // build converted segment - File outputDir = new File(new File(tmpDir, segmentName), "tmp-" + segmentName + "-" + System.currentTimeMillis()); - SegmentZKPropsConfig segmentZKPropsConfig = new SegmentZKPropsConfig(); - segmentZKPropsConfig.setStartOffset("1"); - segmentZKPropsConfig.setEndOffset("100"); - RealtimeSegmentConverter converter = - new RealtimeSegmentConverter(mutableSegmentImpl, segmentZKPropsConfig, outputDir.getAbsolutePath(), schema, - tableNameWithType, tableConfig, segmentName, false); - converter.build(SegmentVersion.v3, null); - - File indexDir = new File(outputDir, segmentName); - SegmentMetadataImpl segmentMetadata = new SegmentMetadataImpl(indexDir); - assertEquals(segmentMetadata.getVersion(), SegmentVersion.v3); - assertEquals(segmentMetadata.getTotalDocs(), rows.size()); - assertEquals(segmentMetadata.getTimeColumn(), DATE_TIME_COLUMN); - assertEquals(segmentMetadata.getTimeUnit(), TimeUnit.MILLISECONDS); - - long expectedStartTime = (long) rows.get(0).getValue(DATE_TIME_COLUMN); - assertEquals(segmentMetadata.getStartTime(), expectedStartTime); - long expectedEndTime = (long) rows.get(rows.size() - 1).getValue(DATE_TIME_COLUMN); - assertEquals(segmentMetadata.getEndTime(), expectedEndTime); - - assertTrue(segmentMetadata.getAllColumns().containsAll(schema.getColumnNames())); - assertEquals(segmentMetadata.getStartOffset(), "1"); - assertEquals(segmentMetadata.getEndOffset(), "100"); - - // read converted segment - SegmentLocalFSDirectory segmentDir = new SegmentLocalFSDirectory(indexDir, segmentMetadata, ReadMode.mmap); - SegmentDirectory.Reader segmentReader = segmentDir.createReader(); + for (GenericRow row : rows) { + mutableSegmentImpl.index(row, null); + } - Map indexContainerMap = new HashMap<>(); - Map columnMetadataMap = segmentMetadata.getColumnMetadataMap(); - IndexLoadingConfig indexLoadingConfig = new IndexLoadingConfig(null, tableConfig); - for (Map.Entry entry : columnMetadataMap.entrySet()) { - indexContainerMap.put(entry.getKey(), - new PhysicalColumnIndexContainer(segmentReader, entry.getValue(), indexLoadingConfig)); - } - ImmutableSegmentImpl segmentFile = new ImmutableSegmentImpl(segmentDir, segmentMetadata, indexContainerMap, null); + // build converted segment + File outputDir = new File(new File(tmpDir, segmentName), "tmp-" + segmentName + "-" + System.currentTimeMillis()); + SegmentZKPropsConfig segmentZKPropsConfig = new SegmentZKPropsConfig(); + segmentZKPropsConfig.setStartOffset("1"); + segmentZKPropsConfig.setEndOffset("100"); + RealtimeSegmentConverter converter = + new RealtimeSegmentConverter(mutableSegmentImpl, segmentZKPropsConfig, outputDir.getAbsolutePath(), schema, + tableNameWithType, tableConfig, segmentName, false); + converter.build(SegmentVersion.v3, null); + + File indexDir = new File(outputDir, segmentName); + SegmentMetadataImpl segmentMetadata = new SegmentMetadataImpl(indexDir); + assertEquals(segmentMetadata.getVersion(), SegmentVersion.v3); + assertEquals(segmentMetadata.getTotalDocs(), rows.size()); + assertEquals(segmentMetadata.getTimeColumn(), DATE_TIME_COLUMN); + assertEquals(segmentMetadata.getTimeUnit(), TimeUnit.MILLISECONDS); + + long expectedStartTime = (long) rows.get(0).getValue(DATE_TIME_COLUMN); + assertEquals(segmentMetadata.getStartTime(), expectedStartTime); + long expectedEndTime = (long) rows.get(rows.size() - 1).getValue(DATE_TIME_COLUMN); + assertEquals(segmentMetadata.getEndTime(), expectedEndTime); + + assertTrue(segmentMetadata.getAllColumns().containsAll(schema.getColumnNames())); + assertEquals(segmentMetadata.getStartOffset(), "1"); + assertEquals(segmentMetadata.getEndOffset(), "100"); + + // read converted segment + SegmentLocalFSDirectory segmentDir = new SegmentLocalFSDirectory(indexDir, segmentMetadata, ReadMode.mmap); + SegmentDirectory.Reader segmentReader = segmentDir.createReader(); + + Map indexContainerMap = new HashMap<>(); + Map columnMetadataMap = segmentMetadata.getColumnMetadataMap(); + IndexLoadingConfig indexLoadingConfig = new IndexLoadingConfig(null, tableConfig); + for (Map.Entry entry : columnMetadataMap.entrySet()) { + indexContainerMap.put(entry.getKey(), + new PhysicalColumnIndexContainer(segmentReader, entry.getValue(), indexLoadingConfig)); + } + segmentFile = new ImmutableSegmentImpl(segmentDir, segmentMetadata, indexContainerMap, null); + + // test forward index contents + GenericRow readRow = new GenericRow(); + int docId = 0; + for (int i = 0; i < rows.size(); i++) { + GenericRow row; + if (sortedColumn == null) { + row = rows.get(i); + } else { + row = rows.get(rows.size() - i - 1); + } + + segmentFile.getRecord(docId, readRow); + + // if rawValueForTextIndex is set and mutable index is reused, the forward index should return the dummy value + if (rawValueForTextIndex != null && reuseMutableIndex) { + assertEquals(readRow.getValue(STRING_COLUMN1), rawValueForTextIndex); + assertEquals(readRow.getValue(DATE_TIME_COLUMN), row.getValue(DATE_TIME_COLUMN)); + } else { + assertEquals(readRow.getValue(STRING_COLUMN1), row.getValue(STRING_COLUMN1)); + assertEquals(readRow.getValue(DATE_TIME_COLUMN), row.getValue(DATE_TIME_COLUMN)); + } + docId += 1; + } - // test forward index contents - GenericRow readRow = new GenericRow(); - int docId = 0; - for (int i = 0; i < rows.size(); i++) { - GenericRow row; + // test docId conversion + TextIndexReader textIndexReader = segmentFile.getIndex(STRING_COLUMN1, StandardIndexes.text()); if (sortedColumn == null) { - row = rows.get(i); + assertEquals(textIndexReader.getDocIds("str-8"), ImmutableRoaringBitmap.bitmapOf(0)); + assertEquals(textIndexReader.getDocIds("str-4"), ImmutableRoaringBitmap.bitmapOf(4)); } else { - row = rows.get(rows.size() - i - 1); + assertEquals(textIndexReader.getDocIds("str-8"), ImmutableRoaringBitmap.bitmapOf(7)); + assertEquals(textIndexReader.getDocIds("str-4"), ImmutableRoaringBitmap.bitmapOf(3)); } - - segmentFile.getRecord(docId, readRow); - - // if rawValueForTextIndex is set and mutable index is reused, the forward index should return the dummy value - if (rawValueForTextIndex != null && reuseMutableIndex) { - assertEquals(readRow.getValue(STRING_COLUMN1), rawValueForTextIndex); - assertEquals(readRow.getValue(DATE_TIME_COLUMN), row.getValue(DATE_TIME_COLUMN)); - } else { - assertEquals(readRow.getValue(STRING_COLUMN1), row.getValue(STRING_COLUMN1)); - assertEquals(readRow.getValue(DATE_TIME_COLUMN), row.getValue(DATE_TIME_COLUMN)); + } finally { + mutableSegmentImpl.destroy(); + if (segmentFile != null) { + segmentFile.destroy(); } - docId += 1; - } - - // test docId conversion - TextIndexReader textIndexReader = segmentFile.getIndex(STRING_COLUMN1, StandardIndexes.text()); - if (sortedColumn == null) { - assertEquals(textIndexReader.getDocIds("str-8"), ImmutableRoaringBitmap.bitmapOf(0)); - assertEquals(textIndexReader.getDocIds("str-4"), ImmutableRoaringBitmap.bitmapOf(4)); - } else { - assertEquals(textIndexReader.getDocIds("str-8"), ImmutableRoaringBitmap.bitmapOf(7)); - assertEquals(textIndexReader.getDocIds("str-4"), ImmutableRoaringBitmap.bitmapOf(3)); } - - mutableSegmentImpl.destroy(); - segmentFile.destroy(); } private List generateTestData() { diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/realtime/impl/dictionary/MultiValueDictionaryTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/realtime/impl/dictionary/MultiValueDictionaryTest.java index 9d95e81cad3a..e58449613110 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/realtime/impl/dictionary/MultiValueDictionaryTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/realtime/impl/dictionary/MultiValueDictionaryTest.java @@ -20,6 +20,7 @@ import java.nio.charset.StandardCharsets; import java.util.Random; +import org.apache.pinot.segment.local.PinotBuffersAfterClassCheckRule; import org.apache.pinot.segment.local.io.writer.impl.DirectMemoryManager; import org.apache.pinot.segment.local.realtime.impl.forward.FixedByteMVMutableForwardIndex; import org.apache.pinot.segment.spi.memory.PinotDataBufferMemoryManager; @@ -33,7 +34,7 @@ import static org.testng.Assert.fail; -public class MultiValueDictionaryTest { +public class MultiValueDictionaryTest implements PinotBuffersAfterClassCheckRule { private static final int NROWS = 1000; private static final int MAX_N_VALUES = 1000; private PinotDataBufferMemoryManager _memoryManager; @@ -52,10 +53,11 @@ public void tearDown() @Test public void testMultiValueIndexingWithDictionary() { long seed = System.nanoTime(); + try (LongOnHeapMutableDictionary dict = new LongOnHeapMutableDictionary(); + DirectMemoryManager memManager = new DirectMemoryManager("test"); FixedByteMVMutableForwardIndex indexer = new FixedByteMVMutableForwardIndex(MAX_N_VALUES, MAX_N_VALUES / 2, - NROWS / 3, Integer.BYTES, new DirectMemoryManager("test"), "indexer", - true, FieldSpec.DataType.INT)) { + NROWS / 3, Integer.BYTES, memManager, "indexer", true, FieldSpec.DataType.INT)) { // Insert rows into the indexer and dictionary Random random = new Random(seed); for (int row = 0; row < NROWS; row++) { @@ -89,9 +91,10 @@ public void testMultiValueIndexingWithDictionary() { @Test public void testMultiValueIndexingWithRawInt() { long seed = System.nanoTime(); - try (FixedByteMVMutableForwardIndex indexer = new FixedByteMVMutableForwardIndex(MAX_N_VALUES, MAX_N_VALUES / 2, - NROWS / 3, Integer.BYTES, new DirectMemoryManager("test"), "indexer", - false, FieldSpec.DataType.INT)) { + + try (DirectMemoryManager memManager = new DirectMemoryManager("test"); + FixedByteMVMutableForwardIndex indexer = new FixedByteMVMutableForwardIndex(MAX_N_VALUES, MAX_N_VALUES / 2, + NROWS / 3, Integer.BYTES, memManager, "indexer", false, FieldSpec.DataType.INT)) { // Insert rows into the indexer Random random = new Random(seed); for (int row = 0; row < NROWS; row++) { @@ -122,9 +125,10 @@ public void testMultiValueIndexingWithRawInt() { @Test public void testMultiValueIndexingWithRawLong() { long seed = System.nanoTime(); - try (FixedByteMVMutableForwardIndex indexer = new FixedByteMVMutableForwardIndex(MAX_N_VALUES, MAX_N_VALUES / 2, - NROWS / 3, Long.BYTES, new DirectMemoryManager("test"), "indexer", - false, FieldSpec.DataType.LONG)) { + + try (DirectMemoryManager memManager = new DirectMemoryManager("test"); + FixedByteMVMutableForwardIndex indexer = new FixedByteMVMutableForwardIndex(MAX_N_VALUES, MAX_N_VALUES / 2, + NROWS / 3, Long.BYTES, memManager, "indexer", false, FieldSpec.DataType.LONG)) { // Insert rows into the indexer Random random = new Random(seed); for (int row = 0; row < NROWS; row++) { @@ -155,9 +159,10 @@ public void testMultiValueIndexingWithRawLong() { @Test public void testMultiValueIndexingWithRawFloat() { long seed = System.nanoTime(); - try (FixedByteMVMutableForwardIndex indexer = new FixedByteMVMutableForwardIndex(MAX_N_VALUES, MAX_N_VALUES / 2, - NROWS / 3, Float.BYTES, new DirectMemoryManager("test"), "indexer", - false, FieldSpec.DataType.FLOAT)) { + + try (DirectMemoryManager memManager = new DirectMemoryManager("test"); + FixedByteMVMutableForwardIndex indexer = new FixedByteMVMutableForwardIndex(MAX_N_VALUES, MAX_N_VALUES / 2, + NROWS / 3, Float.BYTES, memManager, "indexer", false, FieldSpec.DataType.FLOAT)) { // Insert rows into the indexer Random random = new Random(seed); for (int row = 0; row < NROWS; row++) { @@ -188,9 +193,10 @@ public void testMultiValueIndexingWithRawFloat() { @Test public void testMultiValueIndexingWithRawDouble() { long seed = System.nanoTime(); - try (FixedByteMVMutableForwardIndex indexer = new FixedByteMVMutableForwardIndex(MAX_N_VALUES, MAX_N_VALUES / 2, - NROWS / 3, Double.BYTES, new DirectMemoryManager("test"), "indexer", - false, FieldSpec.DataType.DOUBLE)) { + + try (DirectMemoryManager memManager = new DirectMemoryManager("test"); + FixedByteMVMutableForwardIndex indexer = new FixedByteMVMutableForwardIndex(MAX_N_VALUES, MAX_N_VALUES / 2, + NROWS / 3, Double.BYTES, memManager, "indexer", false, FieldSpec.DataType.DOUBLE)) { // Insert rows into the indexer Random random = new Random(seed); for (int row = 0; row < NROWS; row++) { @@ -221,9 +227,10 @@ public void testMultiValueIndexingWithRawDouble() { @Test public void testMultiValueIndexingWithRawString() { long seed = System.nanoTime(); - try (FixedByteMVMutableForwardIndex indexer = new FixedByteMVMutableForwardIndex(MAX_N_VALUES, MAX_N_VALUES / 2, - NROWS / 3, 24, new DirectMemoryManager("test"), "indexer", - false, FieldSpec.DataType.STRING)) { + + try (DirectMemoryManager memManager = new DirectMemoryManager("test"); + FixedByteMVMutableForwardIndex indexer = new FixedByteMVMutableForwardIndex(MAX_N_VALUES, MAX_N_VALUES / 2, + NROWS / 3, 24, memManager, "indexer", false, FieldSpec.DataType.STRING)) { // Insert rows into the indexer Random random = new Random(seed); for (int row = 0; row < NROWS; row++) { @@ -251,8 +258,10 @@ public void testMultiValueIndexingWithRawString() { @Test public void testMultiValueIndexingWithRawByte() { long seed = System.nanoTime(); - try (FixedByteMVMutableForwardIndex indexer = new FixedByteMVMutableForwardIndex(MAX_N_VALUES, MAX_N_VALUES / 2, - NROWS / 3, 24, new DirectMemoryManager("test"), "indexer", + + try (DirectMemoryManager memManager = new DirectMemoryManager("test"); + FixedByteMVMutableForwardIndex indexer = new FixedByteMVMutableForwardIndex(MAX_N_VALUES, MAX_N_VALUES / 2, + NROWS / 3, 24, memManager, "indexer", false, FieldSpec.DataType.BYTES)) { // Insert rows into the indexer Random random = new Random(seed); diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/realtime/impl/dictionary/MutableDictionaryTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/realtime/impl/dictionary/MutableDictionaryTest.java index 349bbe2e7dda..e26eb753f112 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/realtime/impl/dictionary/MutableDictionaryTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/realtime/impl/dictionary/MutableDictionaryTest.java @@ -32,6 +32,7 @@ import java.util.concurrent.Future; import java.util.stream.Collectors; import org.apache.commons.lang3.RandomStringUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterClassCheckRule; import org.apache.pinot.segment.local.io.writer.impl.DirectMemoryManager; import org.apache.pinot.segment.spi.index.mutable.MutableDictionary; import org.apache.pinot.segment.spi.memory.PinotDataBufferMemoryManager; @@ -46,7 +47,7 @@ * Tests for functionality and concurrent read/write against mutable dictionaries. *

Index contiguous integers from 1 so that the index for each value is deterministic. */ -public class MutableDictionaryTest { +public class MutableDictionaryTest implements PinotBuffersAfterClassCheckRule { private static final int NUM_ENTRIES = 100_000; private static final int EST_CARDINALITY = NUM_ENTRIES / 3; private static final int NUM_READERS = 3; diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/realtime/impl/dictionary/OffHeapMutableBytesStoreTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/realtime/impl/dictionary/OffHeapMutableBytesStoreTest.java index f705f9e91c57..efa8e863f4e7 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/realtime/impl/dictionary/OffHeapMutableBytesStoreTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/realtime/impl/dictionary/OffHeapMutableBytesStoreTest.java @@ -18,6 +18,7 @@ */ package org.apache.pinot.segment.local.realtime.impl.dictionary; +import java.io.IOException; import java.util.ArrayList; import java.util.Arrays; import java.util.List; @@ -25,9 +26,11 @@ import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Future; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.io.writer.impl.DirectMemoryManager; import org.apache.pinot.segment.spi.memory.PinotDataBufferMemoryManager; import org.testng.Assert; +import org.testng.annotations.AfterClass; import org.testng.annotations.BeforeClass; import org.testng.annotations.Test; @@ -35,7 +38,7 @@ import static org.testng.Assert.assertTrue; -public class OffHeapMutableBytesStoreTest { +public class OffHeapMutableBytesStoreTest implements PinotBuffersAfterMethodCheckRule { private static final int NUM_VALUES = 100_000; private static final int MAX_NUM_BYTES_PER_VALUE = 512; private static final long RANDOM_SEED = System.currentTimeMillis(); @@ -56,6 +59,12 @@ public void setUp() { } } + @AfterClass + public void tearDown() + throws IOException { + _memoryManager.close(); + } + @Test public void testAdd() throws Exception { diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/BitmapInvertedIndexCreatorTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/BitmapInvertedIndexCreatorTest.java index 790492d62b5c..52c591d25949 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/BitmapInvertedIndexCreatorTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/BitmapInvertedIndexCreatorTest.java @@ -24,6 +24,7 @@ import java.util.Random; import java.util.Set; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.segment.creator.impl.inv.OffHeapBitmapInvertedIndexCreator; import org.apache.pinot.segment.local.segment.creator.impl.inv.OnHeapBitmapInvertedIndexCreator; import org.apache.pinot.segment.local.segment.index.readers.BitmapInvertedIndexReader; @@ -39,7 +40,7 @@ import org.testng.annotations.Test; -public class BitmapInvertedIndexCreatorTest { +public class BitmapInvertedIndexCreatorTest implements PinotBuffersAfterMethodCheckRule { private static final File TEMP_DIR = FileUtils.getTempDirectory(); private static final File ON_HEAP_INDEX_DIR = new File(TEMP_DIR, "onHeap"); private static final File OFF_HEAP_INDEX_DIR = new File(TEMP_DIR, "offHeap"); @@ -148,8 +149,9 @@ public void testMultiValue() private void validate(File invertedIndex, Set[] postingLists) throws IOException { - try (BitmapInvertedIndexReader reader = new BitmapInvertedIndexReader( - PinotDataBuffer.mapReadOnlyBigEndianFile(invertedIndex), CARDINALITY)) { + + try (PinotDataBuffer dataBuffer = PinotDataBuffer.mapReadOnlyBigEndianFile(invertedIndex); + BitmapInvertedIndexReader reader = new BitmapInvertedIndexReader(dataBuffer, CARDINALITY)) { for (int dictId = 0; dictId < CARDINALITY; dictId++) { ImmutableRoaringBitmap bitmap = reader.getDocIds(dictId); Set expected = postingLists[dictId]; diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/BitmapInvertedIndexTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/BitmapInvertedIndexTest.java index 9f818c8c8543..52b034049fda 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/BitmapInvertedIndexTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/BitmapInvertedIndexTest.java @@ -27,6 +27,7 @@ import org.apache.avro.generic.GenericDatumReader; import org.apache.avro.generic.GenericRecord; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterClassCheckRule; import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; import org.apache.pinot.segment.local.segment.index.loader.IndexLoadingConfig; @@ -48,7 +49,7 @@ import org.testng.annotations.Test; -public class BitmapInvertedIndexTest { +public class BitmapInvertedIndexTest implements PinotBuffersAfterClassCheckRule { private static final String AVRO_FILE_PATH = "data" + File.separator + "test_sample_data.avro"; private static final File INDEX_DIR = new File(FileUtils.getTempDirectory(), BitmapInvertedIndexTest.class.getSimpleName()); diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/DictionariesTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/DictionariesTest.java index a678be426984..8514571826e7 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/DictionariesTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/DictionariesTest.java @@ -35,6 +35,7 @@ import org.apache.avro.util.Utf8; import org.apache.commons.io.FileUtils; import org.apache.pinot.plugin.inputformat.avro.AvroUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; import org.apache.pinot.segment.local.segment.creator.impl.SegmentCreationDriverFactory; import org.apache.pinot.segment.local.segment.creator.impl.SegmentDictionaryCreator; @@ -73,7 +74,7 @@ import org.testng.annotations.Test; -public class DictionariesTest { +public class DictionariesTest implements PinotBuffersAfterMethodCheckRule { private static final String AVRO_DATA = "data/test_sample_data.avro"; private static final File INDEX_DIR = new File(DictionariesTest.class.toString()); private static final Map> UNIQUE_ENTRIES = new HashMap<>(); @@ -136,64 +137,71 @@ public void test1() throws Exception { ImmutableSegment heapSegment = ImmutableSegmentLoader.load(_segmentDirectory, ReadMode.heap); ImmutableSegment mmapSegment = ImmutableSegmentLoader.load(_segmentDirectory, ReadMode.mmap); + try { - Schema schema = heapSegment.getSegmentMetadata().getSchema(); - for (FieldSpec fieldSpec : schema.getAllFieldSpecs()) { - // Skip virtual columns - if (fieldSpec.isVirtualColumn()) { - continue; - } + Schema schema = heapSegment.getSegmentMetadata().getSchema(); + for (FieldSpec fieldSpec : schema.getAllFieldSpecs()) { + // Skip virtual columns + if (fieldSpec.isVirtualColumn()) { + continue; + } - String columnName = fieldSpec.getName(); - Dictionary heapDictionary = heapSegment.getDictionary(columnName); - Dictionary mmapDictionary = mmapSegment.getDictionary(columnName); - - switch (fieldSpec.getDataType()) { - case INT: - Assert.assertTrue(heapDictionary instanceof IntDictionary); - Assert.assertTrue(mmapDictionary instanceof IntDictionary); - int firstInt = heapDictionary.getIntValue(0); - Assert.assertEquals(heapDictionary.indexOf(firstInt), heapDictionary.indexOf(String.valueOf(firstInt))); - Assert.assertEquals(mmapDictionary.indexOf(firstInt), mmapDictionary.indexOf(String.valueOf(firstInt))); - break; - case LONG: - Assert.assertTrue(heapDictionary instanceof LongDictionary); - Assert.assertTrue(mmapDictionary instanceof LongDictionary); - long firstLong = heapDictionary.getLongValue(0); - Assert.assertEquals(heapDictionary.indexOf(firstLong), heapDictionary.indexOf(String.valueOf(firstLong))); - Assert.assertEquals(mmapDictionary.indexOf(firstLong), mmapDictionary.indexOf(String.valueOf(firstLong))); - break; - case FLOAT: - Assert.assertTrue(heapDictionary instanceof FloatDictionary); - Assert.assertTrue(mmapDictionary instanceof FloatDictionary); - float firstFloat = heapDictionary.getFloatValue(0); - Assert.assertEquals(heapDictionary.indexOf(firstFloat), heapDictionary.indexOf(String.valueOf(firstFloat))); - Assert.assertEquals(mmapDictionary.indexOf(firstFloat), mmapDictionary.indexOf(String.valueOf(firstFloat))); - break; - case DOUBLE: - Assert.assertTrue(heapDictionary instanceof DoubleDictionary); - Assert.assertTrue(mmapDictionary instanceof DoubleDictionary); - double firstDouble = heapDictionary.getDoubleValue(0); - Assert.assertEquals(heapDictionary.indexOf(firstDouble), heapDictionary.indexOf(String.valueOf(firstDouble))); - Assert.assertEquals(mmapDictionary.indexOf(firstDouble), mmapDictionary.indexOf(String.valueOf(firstDouble))); - break; - case BIG_DECIMAL: - Assert.assertTrue(heapDictionary instanceof BigDecimalDictionary); - Assert.assertTrue(mmapDictionary instanceof BigDecimalDictionary); - break; - case STRING: - Assert.assertTrue(heapDictionary instanceof StringDictionary); - Assert.assertTrue(mmapDictionary instanceof StringDictionary); - break; - default: - Assert.fail(); - break; - } + String columnName = fieldSpec.getName(); + Dictionary heapDictionary = heapSegment.getDictionary(columnName); + Dictionary mmapDictionary = mmapSegment.getDictionary(columnName); + + switch (fieldSpec.getDataType()) { + case INT: + Assert.assertTrue(heapDictionary instanceof IntDictionary); + Assert.assertTrue(mmapDictionary instanceof IntDictionary); + int firstInt = heapDictionary.getIntValue(0); + Assert.assertEquals(heapDictionary.indexOf(firstInt), heapDictionary.indexOf(String.valueOf(firstInt))); + Assert.assertEquals(mmapDictionary.indexOf(firstInt), mmapDictionary.indexOf(String.valueOf(firstInt))); + break; + case LONG: + Assert.assertTrue(heapDictionary instanceof LongDictionary); + Assert.assertTrue(mmapDictionary instanceof LongDictionary); + long firstLong = heapDictionary.getLongValue(0); + Assert.assertEquals(heapDictionary.indexOf(firstLong), heapDictionary.indexOf(String.valueOf(firstLong))); + Assert.assertEquals(mmapDictionary.indexOf(firstLong), mmapDictionary.indexOf(String.valueOf(firstLong))); + break; + case FLOAT: + Assert.assertTrue(heapDictionary instanceof FloatDictionary); + Assert.assertTrue(mmapDictionary instanceof FloatDictionary); + float firstFloat = heapDictionary.getFloatValue(0); + Assert.assertEquals(heapDictionary.indexOf(firstFloat), heapDictionary.indexOf(String.valueOf(firstFloat))); + Assert.assertEquals(mmapDictionary.indexOf(firstFloat), mmapDictionary.indexOf(String.valueOf(firstFloat))); + break; + case DOUBLE: + Assert.assertTrue(heapDictionary instanceof DoubleDictionary); + Assert.assertTrue(mmapDictionary instanceof DoubleDictionary); + double firstDouble = heapDictionary.getDoubleValue(0); + Assert.assertEquals(heapDictionary.indexOf(firstDouble), + heapDictionary.indexOf(String.valueOf(firstDouble))); + Assert.assertEquals(mmapDictionary.indexOf(firstDouble), + mmapDictionary.indexOf(String.valueOf(firstDouble))); + break; + case BIG_DECIMAL: + Assert.assertTrue(heapDictionary instanceof BigDecimalDictionary); + Assert.assertTrue(mmapDictionary instanceof BigDecimalDictionary); + break; + case STRING: + Assert.assertTrue(heapDictionary instanceof StringDictionary); + Assert.assertTrue(mmapDictionary instanceof StringDictionary); + break; + default: + Assert.fail(); + break; + } - Assert.assertEquals(mmapDictionary.length(), heapDictionary.length()); - for (int i = 0; i < heapDictionary.length(); i++) { - Assert.assertEquals(mmapDictionary.get(i), heapDictionary.get(i)); + Assert.assertEquals(mmapDictionary.length(), heapDictionary.length()); + for (int i = 0; i < heapDictionary.length(); i++) { + Assert.assertEquals(mmapDictionary.get(i), heapDictionary.get(i)); + } } + } finally { + heapSegment.destroy(); + mmapSegment.destroy(); } } @@ -202,27 +210,32 @@ public void test2() throws Exception { ImmutableSegment heapSegment = ImmutableSegmentLoader.load(_segmentDirectory, ReadMode.heap); ImmutableSegment mmapSegment = ImmutableSegmentLoader.load(_segmentDirectory, ReadMode.mmap); - - Schema schema = heapSegment.getSegmentMetadata().getSchema(); - for (String columnName : schema.getPhysicalColumnNames()) { - Dictionary heapDictionary = heapSegment.getDictionary(columnName); - Dictionary mmapDictionary = mmapSegment.getDictionary(columnName); - - for (Object entry : UNIQUE_ENTRIES.get(columnName)) { - String stringValue = entry.toString(); - Assert.assertEquals(mmapDictionary.indexOf(stringValue), heapDictionary.indexOf(stringValue)); - if (!columnName.equals("pageKey")) { - Assert.assertFalse(heapDictionary.indexOf(stringValue) < 0); - Assert.assertFalse(mmapDictionary.indexOf(stringValue) < 0); - } - if (entry instanceof Integer) { - Assert.assertEquals(mmapDictionary.indexOf((int) entry), mmapDictionary.indexOf(stringValue)); - Assert.assertEquals(heapDictionary.indexOf((int) entry), heapDictionary.indexOf(stringValue)); - } else if (entry instanceof Long) { - Assert.assertEquals(mmapDictionary.indexOf((long) entry), mmapDictionary.indexOf(stringValue)); - Assert.assertEquals(heapDictionary.indexOf((long) entry), heapDictionary.indexOf(stringValue)); + try { + + Schema schema = heapSegment.getSegmentMetadata().getSchema(); + for (String columnName : schema.getPhysicalColumnNames()) { + Dictionary heapDictionary = heapSegment.getDictionary(columnName); + Dictionary mmapDictionary = mmapSegment.getDictionary(columnName); + + for (Object entry : UNIQUE_ENTRIES.get(columnName)) { + String stringValue = entry.toString(); + Assert.assertEquals(mmapDictionary.indexOf(stringValue), heapDictionary.indexOf(stringValue)); + if (!columnName.equals("pageKey")) { + Assert.assertFalse(heapDictionary.indexOf(stringValue) < 0); + Assert.assertFalse(mmapDictionary.indexOf(stringValue) < 0); + } + if (entry instanceof Integer) { + Assert.assertEquals(mmapDictionary.indexOf((int) entry), mmapDictionary.indexOf(stringValue)); + Assert.assertEquals(heapDictionary.indexOf((int) entry), heapDictionary.indexOf(stringValue)); + } else if (entry instanceof Long) { + Assert.assertEquals(mmapDictionary.indexOf((long) entry), mmapDictionary.indexOf(stringValue)); + Assert.assertEquals(heapDictionary.indexOf((long) entry), heapDictionary.indexOf(stringValue)); + } } } + } finally { + heapSegment.destroy(); + mmapSegment.destroy(); } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/DictionaryOptimiserTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/DictionaryOptimiserTest.java index aeb3feee4f5b..527712486dad 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/DictionaryOptimiserTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/DictionaryOptimiserTest.java @@ -31,6 +31,7 @@ import org.apache.avro.generic.GenericRecord; import org.apache.commons.io.FileUtils; import org.apache.pinot.plugin.inputformat.avro.AvroUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; import org.apache.pinot.segment.local.segment.creator.impl.SegmentCreationDriverFactory; import org.apache.pinot.segment.spi.ImmutableSegment; @@ -56,7 +57,7 @@ import org.testng.annotations.Test; -public class DictionaryOptimiserTest { +public class DictionaryOptimiserTest implements PinotBuffersAfterMethodCheckRule { private static final Logger LOGGER = LoggerFactory.getLogger(DictionaryOptimiserTest.class); private static final String AVRO_DATA = "data/mixed_cardinality_data.avro"; @@ -100,29 +101,32 @@ public static void before() public void testDictionaryForMixedCardinalities() throws Exception { ImmutableSegment heapSegment = ImmutableSegmentLoader.load(_segmentDirectory, ReadMode.heap); + try { + Schema schema = heapSegment.getSegmentMetadata().getSchema(); + for (FieldSpec fieldSpec : schema.getAllFieldSpecs()) { + // Skip virtual columns + if (fieldSpec.isVirtualColumn()) { + continue; + } - Schema schema = heapSegment.getSegmentMetadata().getSchema(); - for (FieldSpec fieldSpec : schema.getAllFieldSpecs()) { - // Skip virtual columns - if (fieldSpec.isVirtualColumn()) { - continue; - } - - String columnName = fieldSpec.getName(); - if (columnName.contains("low_cardinality")) { - Assert.assertTrue(heapSegment.getForwardIndex(columnName).isDictionaryEncoded(), - "No dictionary found for low cardinality columns"); - } + String columnName = fieldSpec.getName(); + if (columnName.contains("low_cardinality")) { + Assert.assertTrue(heapSegment.getForwardIndex(columnName).isDictionaryEncoded(), + "No dictionary found for low cardinality columns"); + } - if (columnName.contains("high_cardinality")) { - Assert.assertFalse(heapSegment.getForwardIndex(columnName).isDictionaryEncoded(), - "No Raw index for high cardinality columns"); - } + if (columnName.contains("high_cardinality")) { + Assert.assertFalse(heapSegment.getForwardIndex(columnName).isDictionaryEncoded(), + "No Raw index for high cardinality columns"); + } - if (columnName.contains("key")) { - Assert.assertFalse(heapSegment.getSegmentMetadata().getColumnMetadataFor(columnName).hasDictionary(), - "Dictionary found for text index column"); + if (columnName.contains("key")) { + Assert.assertFalse(heapSegment.getSegmentMetadata().getColumnMetadataFor(columnName).hasDictionary(), + "Dictionary found for text index column"); + } } + } finally { + heapSegment.destroy(); } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/DictionaryOptimizerCardinalityTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/DictionaryOptimizerCardinalityTest.java index 85f57beae7ac..ddf9158e88e7 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/DictionaryOptimizerCardinalityTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/DictionaryOptimizerCardinalityTest.java @@ -28,6 +28,7 @@ import java.util.concurrent.TimeUnit; import java.util.stream.Collectors; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterClassCheckRule; import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; import org.apache.pinot.segment.local.segment.creator.impl.SegmentCreationDriverFactory; import org.apache.pinot.segment.spi.ImmutableSegment; @@ -51,7 +52,7 @@ import org.testng.annotations.Test; -public class DictionaryOptimizerCardinalityTest { +public class DictionaryOptimizerCardinalityTest implements PinotBuffersAfterClassCheckRule { private static final Logger LOGGER = LoggerFactory.getLogger(DictionaryOptimizerCardinalityTest.class); private static final File INDEX_DIR = new File(DictionaryOptimizerCardinalityTest.class.toString()); @@ -64,24 +65,28 @@ public void testDictionaryForMixedCardinalitiesStringType() throws Exception { ImmutableSegment heapSegment = ImmutableSegmentLoader.load(_segmentDirectory, ReadMode.heap); + try { - Schema schema = heapSegment.getSegmentMetadata().getSchema(); - for (FieldSpec fieldSpec : schema.getAllFieldSpecs()) { - // Skip virtual columns - if (fieldSpec.isVirtualColumn()) { - continue; - } + Schema schema = heapSegment.getSegmentMetadata().getSchema(); + for (FieldSpec fieldSpec : schema.getAllFieldSpecs()) { + // Skip virtual columns + if (fieldSpec.isVirtualColumn()) { + continue; + } - String columnName = fieldSpec.getName(); - if ("low_cardinality_strings".equals(columnName)) { - Assert.assertTrue(heapSegment.getForwardIndex(columnName).isDictionaryEncoded(), - "Low cardinality columns should be dictionary encoded"); - } + String columnName = fieldSpec.getName(); + if ("low_cardinality_strings".equals(columnName)) { + Assert.assertTrue(heapSegment.getForwardIndex(columnName).isDictionaryEncoded(), + "Low cardinality columns should be dictionary encoded"); + } - if ("high_cardinality_strings".equals(columnName)) { - Assert.assertFalse(heapSegment.getForwardIndex(columnName).isDictionaryEncoded(), - "High cardinality columns should be raw encoded"); + if ("high_cardinality_strings".equals(columnName)) { + Assert.assertFalse(heapSegment.getForwardIndex(columnName).isDictionaryEncoded(), + "High cardinality columns should be raw encoded"); + } } + } finally { + heapSegment.destroy(); } } @@ -100,12 +105,15 @@ private void setup() ingestionConfig.setRowTimeValueCheck(false); ingestionConfig.setSegmentTimeValueCheck(false); Schema schema = - new Schema.SchemaBuilder().addSingleValueDimension("low_cardinality_strings", FieldSpec.DataType.STRING) + new Schema.SchemaBuilder() + .addSingleValueDimension("low_cardinality_strings", FieldSpec.DataType.STRING) .addSingleValueDimension("high_cardinality_strings", FieldSpec.DataType.STRING) - .addDateTimeField("ts", FieldSpec.DataType.LONG, "1:MILLISECONDS:EPOCH", "1:MILLISECONDS").build(); + .addDateTimeField("ts", FieldSpec.DataType.LONG, "1:MILLISECONDS:EPOCH", "1:MILLISECONDS") + .build(); + List stringColumns = - schema.getDimensionFieldSpecs().stream().filter(x -> x.getDataType() == FieldSpec.DataType.STRING).collect( - Collectors.toList()); + schema.getDimensionFieldSpecs().stream().filter(x -> x.getDataType() == FieldSpec.DataType.STRING) + .collect(Collectors.toList()); List fieldConfigList = stringColumns.stream().map( x -> new FieldConfig(x.getName(), FieldConfig.EncodingType.DICTIONARY, Collections.emptyList(), null, null)) diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/OnHeapDictionariesTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/OnHeapDictionariesTest.java index e455cf1041ba..ac34f4b9b43a 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/OnHeapDictionariesTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/OnHeapDictionariesTest.java @@ -26,6 +26,7 @@ import java.util.Random; import org.apache.commons.io.FileUtils; import org.apache.commons.lang3.RandomStringUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterClassCheckRule; import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; import org.apache.pinot.segment.local.segment.index.loader.IndexLoadingConfig; @@ -55,7 +56,7 @@ /** * Unit tests for On-Heap dictionary implementations. */ -public class OnHeapDictionariesTest { +public class OnHeapDictionariesTest implements PinotBuffersAfterClassCheckRule { private static final Logger LOGGER = LoggerFactory.getLogger(OnHeapDictionariesTest.class); private static final String SEGMENT_DIR_NAME = System.getProperty("java.io.tmpdir") + File.separator + "onHeapDict"; @@ -91,6 +92,8 @@ public void setup() @AfterClass public void tearDown() throws IOException { + _onHeapSegment.destroy(); + _offHeapSegment.destroy(); FileUtils.deleteDirectory(new File(SEGMENT_DIR_NAME)); } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/SegmentTestUtils.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/SegmentTestUtils.java index a0936c94104c..a9462d9f8402 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/SegmentTestUtils.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/SegmentTestUtils.java @@ -84,10 +84,13 @@ public static SegmentGeneratorConfig getSegmentGenSpecWithSchemAndProjectedColum String timeColumn, TimeUnit timeUnit, String tableName) throws IOException { Schema schema = extractSchemaFromAvroWithoutTime(inputAvro); - TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName(tableName) + TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE) + .setTableName(tableName) .setInvertedIndexColumns(new ArrayList<>(schema.getColumnNames())) - .setCreateInvertedIndexDuringSegmentGeneration(true).setIngestionConfig(getSkipTimeCheckIngestionConfig()) + .setCreateInvertedIndexDuringSegmentGeneration(true) + .setIngestionConfig(getSkipTimeCheckIngestionConfig()) .build(); + SegmentGeneratorConfig segmentGenSpec = new SegmentGeneratorConfig(tableConfig, schema); segmentGenSpec.setInputFilePath(inputAvro.getAbsolutePath()); segmentGenSpec.setTimeColumnName(timeColumn); @@ -95,6 +98,7 @@ public static SegmentGeneratorConfig getSegmentGenSpecWithSchemAndProjectedColum segmentGenSpec.setFormat(FileFormat.AVRO); segmentGenSpec.setSegmentVersion(SegmentVersion.v1); segmentGenSpec.setOutDir(outputDir.getAbsolutePath()); + return segmentGenSpec; } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentColumnarIndexCreatorTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentColumnarIndexCreatorTest.java index ca86b86166c1..71bd9b2621b7 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentColumnarIndexCreatorTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentColumnarIndexCreatorTest.java @@ -24,6 +24,7 @@ import java.util.List; import org.apache.commons.configuration2.PropertiesConfiguration; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterClassCheckRule; import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader; import org.apache.pinot.segment.spi.IndexSegment; @@ -47,7 +48,7 @@ import static org.testng.Assert.assertTrue; -public class SegmentColumnarIndexCreatorTest { +public class SegmentColumnarIndexCreatorTest implements PinotBuffersAfterClassCheckRule { private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(), "SegmentColumnarIndexCreatorTest"); private static final File CONFIG_FILE = new File(TEMP_DIR, "config"); private static final String COLUMN_NAME = "testColumn"; @@ -109,6 +110,8 @@ private static long getStartTimeInSegmentMetadata(String testDateTimeFormat, Str SegmentGeneratorConfig config = new SegmentGeneratorConfig(tableConfig, schema); config.setOutDir(indexDirPath); config.setSegmentName(segmentName); + IndexSegment indexSegment = null; + try { FileUtils.deleteQuietly(new File(indexDirPath)); @@ -119,10 +122,13 @@ private static long getStartTimeInSegmentMetadata(String testDateTimeFormat, Str SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl(); driver.init(config, new GenericRowRecordReader(rows)); driver.build(); - IndexSegment indexSegment = ImmutableSegmentLoader.load(new File(indexDirPath, segmentName), ReadMode.heap); + indexSegment = ImmutableSegmentLoader.load(new File(indexDirPath, segmentName), ReadMode.heap); SegmentMetadata md = indexSegment.getSegmentMetadata(); return md.getStartTime(); } finally { + if (indexSegment != null) { + indexSegment.destroy(); + } FileUtils.deleteQuietly(new File(indexDirPath)); } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/H3IndexTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/H3IndexTest.java index 39ba18a90413..9ca006fdaf13 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/H3IndexTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/H3IndexTest.java @@ -28,6 +28,7 @@ import java.util.Random; import java.util.stream.Collectors; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.realtime.impl.geospatial.MutableH3Index; import org.apache.pinot.segment.local.segment.creator.impl.inv.geospatial.OffHeapH3IndexCreator; import org.apache.pinot.segment.local.segment.creator.impl.inv.geospatial.OnHeapH3IndexCreator; @@ -56,7 +57,7 @@ import static org.testng.Assert.assertTrue; -public class H3IndexTest { +public class H3IndexTest implements PinotBuffersAfterMethodCheckRule { private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(), "H3IndexCreatorTest"); private static final Random RANDOM = new Random(); diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/JsonIndexTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/JsonIndexTest.java index 8a3ac7c7b710..1412214381b0 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/JsonIndexTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/JsonIndexTest.java @@ -28,6 +28,7 @@ import java.util.Map; import java.util.stream.Collectors; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.realtime.impl.json.MutableJsonIndexImpl; import org.apache.pinot.segment.local.segment.creator.impl.inv.json.OffHeapJsonIndexCreator; import org.apache.pinot.segment.local.segment.creator.impl.inv.json.OnHeapJsonIndexCreator; @@ -55,7 +56,7 @@ /** * Unit test for {@link JsonIndexCreator} and {@link JsonIndexReader}. */ -public class JsonIndexTest { +public class JsonIndexTest implements PinotBuffersAfterMethodCheckRule { private static final File INDEX_DIR = new File(FileUtils.getTempDirectory(), "JsonIndexTest"); private static final String ON_HEAP_COLUMN_NAME = "onHeap"; private static final String OFF_HEAP_COLUMN_NAME = "offHeap"; diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/BitSlicedIndexCreatorTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/BitSlicedIndexCreatorTest.java index d80723c35246..7e4571a29045 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/BitSlicedIndexCreatorTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/BitSlicedIndexCreatorTest.java @@ -26,6 +26,7 @@ import java.util.function.LongSupplier; import java.util.stream.IntStream; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.segment.creator.impl.inv.BitSlicedRangeIndexCreator; import org.apache.pinot.segment.local.segment.index.readers.BitSlicedRangeIndexReader; import org.apache.pinot.segment.spi.ColumnMetadata; @@ -44,7 +45,7 @@ import static org.testng.Assert.assertEquals; -public class BitSlicedIndexCreatorTest { +public class BitSlicedIndexCreatorTest implements PinotBuffersAfterMethodCheckRule { private static final File INDEX_DIR = new File(FileUtils.getTempDirectory(), "BitSlicedIndexCreatorTest"); private static final String COLUMN_NAME = "testColumn"; diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/BloomFilterCreatorTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/BloomFilterCreatorTest.java index 48a691f3a65d..589016d7bc86 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/BloomFilterCreatorTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/BloomFilterCreatorTest.java @@ -20,6 +20,7 @@ import java.io.File; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.segment.creator.impl.bloom.OnHeapGuavaBloomFilterCreator; import org.apache.pinot.segment.local.segment.index.readers.bloom.BloomFilterReaderFactory; import org.apache.pinot.segment.spi.V1Constants; @@ -35,7 +36,7 @@ import org.testng.annotations.Test; -public class BloomFilterCreatorTest { +public class BloomFilterCreatorTest implements PinotBuffersAfterMethodCheckRule { private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(), "BloomFilterCreatorTest"); @BeforeClass diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/CLPForwardIndexCreatorTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/CLPForwardIndexCreatorTest.java index c97a1d318fb3..140da3d22c53 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/CLPForwardIndexCreatorTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/CLPForwardIndexCreatorTest.java @@ -23,6 +23,7 @@ import java.util.Collections; import java.util.List; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.segment.creator.impl.fwd.CLPForwardIndexCreatorV1; import org.apache.pinot.segment.local.segment.creator.impl.stats.StringColumnPreIndexStatsCollector; import org.apache.pinot.segment.local.segment.index.readers.forward.CLPForwardIndexReaderV1; @@ -45,7 +46,8 @@ import org.testng.annotations.BeforeClass; import org.testng.annotations.Test; -public class CLPForwardIndexCreatorTest { + +public class CLPForwardIndexCreatorTest implements PinotBuffersAfterMethodCheckRule { private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(), "CLPForwardIndexCreatorTest"); @BeforeClass @@ -92,20 +94,21 @@ public void testCLPWriter() statsCollector.seal(); File indexFile = new File(TEMP_DIR, "column1" + V1Constants.Indexes.RAW_SV_FORWARD_INDEX_FILE_EXTENSION); - CLPForwardIndexCreatorV1 clpForwardIndexCreatorV1 = - new CLPForwardIndexCreatorV1(TEMP_DIR, "column1", logLines.size(), statsCollector); + try (CLPForwardIndexCreatorV1 clpForwardIndexCreatorV1 = + new CLPForwardIndexCreatorV1(TEMP_DIR, "column1", logLines.size(), statsCollector)) { - for (String logLine : logLines) { - clpForwardIndexCreatorV1.putString(logLine); + for (String logLine : logLines) { + clpForwardIndexCreatorV1.putString(logLine); + } + clpForwardIndexCreatorV1.seal(); } - clpForwardIndexCreatorV1.seal(); - clpForwardIndexCreatorV1.close(); - PinotDataBuffer pinotDataBuffer = PinotDataBuffer.mapReadOnlyBigEndianFile(indexFile); - CLPForwardIndexReaderV1 clpForwardIndexReaderV1 = new CLPForwardIndexReaderV1(pinotDataBuffer, logLines.size()); - for (int i = 0; i < logLines.size(); i++) { - Assert.assertEquals(clpForwardIndexReaderV1.getString(i, clpForwardIndexReaderV1.createContext()), - logLines.get(i)); + try (PinotDataBuffer pinotDataBuffer = PinotDataBuffer.mapReadOnlyBigEndianFile(indexFile)) { + CLPForwardIndexReaderV1 clpForwardIndexReaderV1 = new CLPForwardIndexReaderV1(pinotDataBuffer, logLines.size()); + for (int i = 0; i < logLines.size(); i++) { + Assert.assertEquals(clpForwardIndexReaderV1.getString(i, clpForwardIndexReaderV1.createContext()), + logLines.get(i)); + } } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/CLPForwardIndexCreatorV2Test.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/CLPForwardIndexCreatorV2Test.java index 65152152e455..52b8292b4971 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/CLPForwardIndexCreatorV2Test.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/CLPForwardIndexCreatorV2Test.java @@ -28,6 +28,7 @@ import java.util.List; import org.apache.commons.compress.compressors.gzip.GzipCompressorInputStream; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterClassCheckRule; import org.apache.pinot.segment.local.io.writer.impl.DirectMemoryManager; import org.apache.pinot.segment.local.realtime.impl.forward.CLPMutableForwardIndexV2; import org.apache.pinot.segment.local.segment.creator.impl.fwd.CLPForwardIndexCreatorV2; @@ -46,7 +47,7 @@ import org.testng.annotations.Test; -public class CLPForwardIndexCreatorV2Test { +public class CLPForwardIndexCreatorV2Test implements PinotBuffersAfterClassCheckRule { private static final String COLUMN_NAME = "column1"; private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(), CLPForwardIndexCreatorV2Test.class.getSimpleName()); @@ -76,6 +77,7 @@ public void setUp() @AfterClass public void tearDown() throws IOException { + _memoryManager.close(); TestUtils.ensureDirectoriesExistAndEmpty(TEMP_DIR); } @@ -83,49 +85,51 @@ public void tearDown() public void testCLPWriter() throws IOException { // Create and ingest into a clp mutable forward indexes - CLPMutableForwardIndexV2 clpMutableForwardIndexV2 = new CLPMutableForwardIndexV2(COLUMN_NAME, _memoryManager); - int rawSizeBytes = 0; - int maxLength = 0; - for (int i = 0; i < _logMessages.size(); i++) { - String logMessage = _logMessages.get(i); - clpMutableForwardIndexV2.setString(i, logMessage); - rawSizeBytes += logMessage.length(); - maxLength = Math.max(maxLength, logMessage.length()); - } + try ( + CLPMutableForwardIndexV2 clpMutableForwardIndexV2 = new CLPMutableForwardIndexV2(COLUMN_NAME, _memoryManager)) { + int rawSizeBytes = 0; + int maxLength = 0; + for (int i = 0; i < _logMessages.size(); i++) { + String logMessage = _logMessages.get(i); + clpMutableForwardIndexV2.setString(i, logMessage); + rawSizeBytes += logMessage.length(); + maxLength = Math.max(maxLength, logMessage.length()); + } - // LZ4 compression type - long rawStringFwdIndexSizeLZ4 = createStringRawForwardIndex(ChunkCompressionType.LZ4, maxLength); - long clpFwdIndexSizeLZ4 = - createAndValidateClpImmutableForwardIndex(clpMutableForwardIndexV2, ChunkCompressionType.LZ4); - // For LZ4 compression: - // 1. CLP raw forward index should achieve at least 40x compression - // 2. at least 25% smaller file size compared to standard raw forward index with LZ4 compression - Assert.assertTrue((float) rawSizeBytes / clpFwdIndexSizeLZ4 >= 40); - Assert.assertTrue((float) rawStringFwdIndexSizeLZ4 / clpFwdIndexSizeLZ4 >= 0.25); - - // ZSTD compression type - long rawStringFwdIndexSizeZSTD = createStringRawForwardIndex(ChunkCompressionType.ZSTANDARD, maxLength); - long clpFwdIndexSizeZSTD = - createAndValidateClpImmutableForwardIndex(clpMutableForwardIndexV2, ChunkCompressionType.ZSTANDARD); - // For ZSTD compression - // 1. CLP raw forward index should achieve at least 66x compression - // 2. at least 19% smaller file size compared to standard raw forward index with ZSTD compression - Assert.assertTrue((float) rawSizeBytes / clpFwdIndexSizeZSTD >= 66); - Assert.assertTrue((float) rawStringFwdIndexSizeZSTD / clpFwdIndexSizeZSTD >= 0.19); + // LZ4 compression type + long rawStringFwdIndexSizeLZ4 = createStringRawForwardIndex(ChunkCompressionType.LZ4, maxLength); + long clpFwdIndexSizeLZ4 = + createAndValidateClpImmutableForwardIndex(clpMutableForwardIndexV2, ChunkCompressionType.LZ4); + // For LZ4 compression: + // 1. CLP raw forward index should achieve at least 40x compression + // 2. at least 25% smaller file size compared to standard raw forward index with LZ4 compression + Assert.assertTrue((float) rawSizeBytes / clpFwdIndexSizeLZ4 >= 40); + Assert.assertTrue((float) rawStringFwdIndexSizeLZ4 / clpFwdIndexSizeLZ4 >= 0.25); + + // ZSTD compression type + long rawStringFwdIndexSizeZSTD = createStringRawForwardIndex(ChunkCompressionType.ZSTANDARD, maxLength); + long clpFwdIndexSizeZSTD = + createAndValidateClpImmutableForwardIndex(clpMutableForwardIndexV2, ChunkCompressionType.ZSTANDARD); + // For ZSTD compression + // 1. CLP raw forward index should achieve at least 66x compression + // 2. at least 19% smaller file size compared to standard raw forward index with ZSTD compression + Assert.assertTrue((float) rawSizeBytes / clpFwdIndexSizeZSTD >= 66); + Assert.assertTrue((float) rawStringFwdIndexSizeZSTD / clpFwdIndexSizeZSTD >= 0.19); + } } private long createStringRawForwardIndex(ChunkCompressionType chunkCompressionType, int maxLength) throws IOException { // Create a raw string immutable forward index TestUtils.ensureDirectoriesExistAndEmpty(TEMP_DIR); - SingleValueVarByteRawIndexCreator index = + try (SingleValueVarByteRawIndexCreator index = new SingleValueVarByteRawIndexCreator(TEMP_DIR, chunkCompressionType, COLUMN_NAME, _logMessages.size(), - FieldSpec.DataType.STRING, maxLength); - for (String logMessage : _logMessages) { - index.putString(logMessage); + FieldSpec.DataType.STRING, maxLength)) { + for (String logMessage : _logMessages) { + index.putString(logMessage); + } + index.seal(); } - index.seal(); - index.close(); File indexFile = new File(TEMP_DIR, COLUMN_NAME + V1Constants.Indexes.RAW_SV_FORWARD_INDEX_FILE_EXTENSION); return indexFile.length(); @@ -138,11 +142,16 @@ private long createAndValidateClpImmutableForwardIndex(CLPMutableForwardIndexV2 // Read from immutable forward index and validate the content File indexFile = new File(TEMP_DIR, COLUMN_NAME + V1Constants.Indexes.RAW_SV_FORWARD_INDEX_FILE_EXTENSION); - PinotDataBuffer pinotDataBuffer = PinotDataBuffer.mapReadOnlyBigEndianFile(indexFile); - CLPForwardIndexReaderV2 clpForwardIndexReaderV2 = new CLPForwardIndexReaderV2(pinotDataBuffer, _logMessages.size()); - CLPForwardIndexReaderV2.CLPReaderContext clpForwardIndexReaderV2Context = clpForwardIndexReaderV2.createContext(); - for (int i = 0; i < _logMessages.size(); i++) { - Assert.assertEquals(clpForwardIndexReaderV2.getString(i, clpForwardIndexReaderV2Context), _logMessages.get(i)); + try (PinotDataBuffer pinotDataBuffer = PinotDataBuffer.mapReadOnlyBigEndianFile(indexFile)) { + CLPForwardIndexReaderV2 clpForwardIndexReaderV2 = + new CLPForwardIndexReaderV2(pinotDataBuffer, _logMessages.size()); + try (CLPForwardIndexReaderV2.CLPReaderContext clpForwardIndexReaderV2Context = + clpForwardIndexReaderV2.createContext()) { + for (int i = 0; i < _logMessages.size(); i++) { + Assert.assertEquals(clpForwardIndexReaderV2.getString(i, clpForwardIndexReaderV2Context), + _logMessages.get(i)); + } + } } return indexSize; diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/LuceneFSTIndexCreatorTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/LuceneFSTIndexCreatorTest.java index 11ed2197fa28..0e88cc73e9f3 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/LuceneFSTIndexCreatorTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/LuceneFSTIndexCreatorTest.java @@ -22,6 +22,7 @@ import java.io.IOException; import java.nio.ByteOrder; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.segment.creator.impl.inv.text.LuceneFSTIndexCreator; import org.apache.pinot.segment.local.segment.index.readers.LuceneFSTIndexReader; import org.apache.pinot.segment.spi.memory.PinotDataBuffer; @@ -35,7 +36,7 @@ import static org.apache.pinot.segment.spi.V1Constants.Indexes.LUCENE_V912_FST_INDEX_FILE_EXTENSION; -public class LuceneFSTIndexCreatorTest { +public class LuceneFSTIndexCreatorTest implements PinotBuffersAfterMethodCheckRule { private static final File INDEX_DIR = new File(FileUtils.getTempDirectory(), "LuceneFSTIndex"); @BeforeClass @@ -63,15 +64,17 @@ public void testIndexWriterReader() INDEX_DIR, "testFSTColumn", uniqueValues); creator.seal(); File fstFile = new File(INDEX_DIR, "testFSTColumn" + LUCENE_V912_FST_INDEX_FILE_EXTENSION); - PinotDataBuffer pinotDataBuffer = + try (PinotDataBuffer pinotDataBuffer = PinotDataBuffer.mapFile(fstFile, true, 0, fstFile.length(), ByteOrder.BIG_ENDIAN, "fstIndexFile"); - LuceneFSTIndexReader reader = new LuceneFSTIndexReader(pinotDataBuffer); - int[] matchedDictIds = reader.getDictIds("hello.*").toArray(); - Assert.assertEquals(2, matchedDictIds.length); - Assert.assertEquals(0, matchedDictIds[0]); - Assert.assertEquals(1, matchedDictIds[1]); + LuceneFSTIndexReader reader = new LuceneFSTIndexReader(pinotDataBuffer)) { - matchedDictIds = reader.getDictIds(".*llo").toArray(); - Assert.assertEquals(0, matchedDictIds.length); + int[] matchedDictIds = reader.getDictIds("hello.*").toArray(); + Assert.assertEquals(2, matchedDictIds.length); + Assert.assertEquals(0, matchedDictIds[0]); + Assert.assertEquals(1, matchedDictIds[1]); + + matchedDictIds = reader.getDictIds(".*llo").toArray(); + Assert.assertEquals(0, matchedDictIds.length); + } } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/MultiValueFixedByteRawIndexCreatorTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/MultiValueFixedByteRawIndexCreatorTest.java index 9a2105726aa7..6828964892cc 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/MultiValueFixedByteRawIndexCreatorTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/MultiValueFixedByteRawIndexCreatorTest.java @@ -30,6 +30,7 @@ import java.util.stream.Collectors; import java.util.stream.IntStream; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.io.writer.impl.VarByteChunkForwardIndexWriterV4; import org.apache.pinot.segment.local.segment.creator.impl.fwd.MultiValueFixedByteRawIndexCreator; import org.apache.pinot.segment.local.segment.index.readers.forward.FixedByteChunkMVForwardIndexReader; @@ -47,7 +48,7 @@ import org.testng.annotations.Test; -public class MultiValueFixedByteRawIndexCreatorTest { +public class MultiValueFixedByteRawIndexCreatorTest implements PinotBuffersAfterMethodCheckRule { protected static String _outputDir; @@ -173,25 +174,27 @@ public void testMV(DataType dataType, List inputs, ToIntFunction sizeo creator.close(); //read - final PinotDataBuffer buffer = PinotDataBuffer.mapFile(file, true, 0, file.length(), ByteOrder.BIG_ENDIAN, ""); - ForwardIndexReader reader = getForwardIndexReader(buffer, dataType, writerVersion); - - final ForwardIndexReaderContext context = reader.createContext(); - T valueBuffer = constructor.apply(maxElements); - for (int i = 0; i < numDocs; i++) { - Assert.assertEquals(inputs.get(i), extractor.extract(reader, context, i, valueBuffer)); - } + try (final PinotDataBuffer buffer = PinotDataBuffer.mapFile(file, true, 0, file.length(), ByteOrder.BIG_ENDIAN, + "")) { + ForwardIndexReader reader = getForwardIndexReader(buffer, dataType, writerVersion); + + final ForwardIndexReaderContext context = reader.createContext(); + T valueBuffer = constructor.apply(maxElements); + for (int i = 0; i < numDocs; i++) { + Assert.assertEquals(inputs.get(i), extractor.extract(reader, context, i, valueBuffer)); + } - // Value byte range test - Assert.assertTrue(reader.isBufferByteRangeInfoSupported()); - Assert.assertFalse(reader.isFixedOffsetMappingType()); - final ForwardIndexReaderContext valueRangeContext = reader.createContext(); - List ranges = new ArrayList<>(); - for (int i = 0; i < numDocs; i++) { - try { - reader.recordDocIdByteRanges(i, valueRangeContext, ranges); - } catch (Exception e) { - Assert.fail("Failed to record byte ranges for docId: " + i, e); + // Value byte range test + Assert.assertTrue(reader.isBufferByteRangeInfoSupported()); + Assert.assertFalse(reader.isFixedOffsetMappingType()); + final ForwardIndexReaderContext valueRangeContext = reader.createContext(); + List ranges = new ArrayList<>(); + for (int i = 0; i < numDocs; i++) { + try { + reader.recordDocIdByteRanges(i, valueRangeContext, ranges); + } catch (Exception e) { + Assert.fail("Failed to record byte ranges for docId: " + i, e); + } } } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/MultiValueVarByteRawIndexCreatorTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/MultiValueVarByteRawIndexCreatorTest.java index cb49f5b37891..8dde0da31f02 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/MultiValueVarByteRawIndexCreatorTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/MultiValueVarByteRawIndexCreatorTest.java @@ -29,6 +29,7 @@ import java.util.stream.IntStream; import java.util.stream.Stream; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.segment.creator.impl.fwd.MultiValueVarByteRawIndexCreator; import org.apache.pinot.segment.local.segment.index.forward.ForwardIndexReaderFactory; import org.apache.pinot.segment.spi.V1Constants.Indexes; @@ -44,7 +45,7 @@ import org.testng.annotations.Test; -public class MultiValueVarByteRawIndexCreatorTest { +public class MultiValueVarByteRawIndexCreatorTest implements PinotBuffersAfterMethodCheckRule { private static final File OUTPUT_DIR = new File(FileUtils.getTempDirectory(), MultiValueVarByteRawIndexCreatorTest.class.getSimpleName()); diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/NativeFSTIndexCreatorTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/NativeFSTIndexCreatorTest.java index 12a0363d44b7..bec54c1c28ba 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/NativeFSTIndexCreatorTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/NativeFSTIndexCreatorTest.java @@ -21,6 +21,7 @@ import java.io.File; import java.io.IOException; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.utils.nativefst.NativeFSTIndexCreator; import org.apache.pinot.segment.local.utils.nativefst.NativeFSTIndexReader; import org.apache.pinot.segment.spi.memory.PinotDataBuffer; @@ -32,7 +33,7 @@ import static org.apache.pinot.segment.spi.V1Constants.Indexes.LUCENE_V912_FST_INDEX_FILE_EXTENSION; -public class NativeFSTIndexCreatorTest { +public class NativeFSTIndexCreatorTest implements PinotBuffersAfterMethodCheckRule { private static final File INDEX_DIR = new File(FileUtils.getTempDirectory(), "NativeFSTIndexCreatorTest"); @BeforeClass @@ -62,14 +63,17 @@ public void testIndexWriterReader() File fstFile = new File(INDEX_DIR, "testFSTColumn" + LUCENE_V912_FST_INDEX_FILE_EXTENSION); try (PinotDataBuffer dataBuffer = PinotDataBuffer.mapReadOnlyBigEndianFile(fstFile); NativeFSTIndexReader reader = new NativeFSTIndexReader(dataBuffer)) { + try { + int[] matchedDictIds = reader.getDictIds("hello.*").toArray(); + Assert.assertEquals(2, matchedDictIds.length); + Assert.assertEquals(0, matchedDictIds[0]); + Assert.assertEquals(1, matchedDictIds[1]); - int[] matchedDictIds = reader.getDictIds("hello.*").toArray(); - Assert.assertEquals(2, matchedDictIds.length); - Assert.assertEquals(0, matchedDictIds[0]); - Assert.assertEquals(1, matchedDictIds[1]); - - matchedDictIds = reader.getDictIds(".*llo").toArray(); - Assert.assertEquals(0, matchedDictIds.length); + matchedDictIds = reader.getDictIds(".*llo").toArray(); + Assert.assertEquals(0, matchedDictIds.length); + } finally { + reader.closeInTest(); + } } } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/NativeTextIndexCreatorTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/NativeTextIndexCreatorTest.java index e7489ca1774b..bf346d3c58d0 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/NativeTextIndexCreatorTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/NativeTextIndexCreatorTest.java @@ -21,6 +21,7 @@ import java.io.File; import java.io.IOException; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.segment.creator.impl.text.NativeTextIndexCreator; import org.apache.pinot.segment.local.segment.index.readers.text.NativeTextIndexReader; import org.apache.pinot.spi.data.DimensionFieldSpec; @@ -33,7 +34,7 @@ import static org.testng.AssertJUnit.assertEquals; -public class NativeTextIndexCreatorTest { +public class NativeTextIndexCreatorTest implements PinotBuffersAfterMethodCheckRule { private static final File INDEX_DIR = new File(FileUtils.getTempDirectory(), "NativeTextIndexCreatorTest"); @BeforeClass @@ -68,25 +69,28 @@ public void testIndexWriterReader() File fstFile = new File(INDEX_DIR, "testFSTColumn" + NATIVE_TEXT_INDEX_FILE_EXTENSION); try (NativeTextIndexReader reader = new NativeTextIndexReader("testFSTColumn", fstFile.getParentFile())) { + try { + int[] matchedDocIds = reader.getDocIds("hello.*").toArray(); + assertEquals(2, matchedDocIds.length); + assertEquals(0, matchedDocIds[0]); + assertEquals(1, matchedDocIds[1]); - int[] matchedDocIds = reader.getDocIds("hello.*").toArray(); - assertEquals(2, matchedDocIds.length); - assertEquals(0, matchedDocIds[0]); - assertEquals(1, matchedDocIds[1]); + matchedDocIds = reader.getDocIds(".*llo").toArray(); + assertEquals(2, matchedDocIds.length); + assertEquals(0, matchedDocIds[0]); + assertEquals(1, matchedDocIds[1]); - matchedDocIds = reader.getDocIds(".*llo").toArray(); - assertEquals(2, matchedDocIds.length); - assertEquals(0, matchedDocIds[0]); - assertEquals(1, matchedDocIds[1]); + matchedDocIds = reader.getDocIds("wor.*").toArray(); + assertEquals(2, matchedDocIds.length); + assertEquals(0, matchedDocIds[0]); + assertEquals(1, matchedDocIds[1]); - matchedDocIds = reader.getDocIds("wor.*").toArray(); - assertEquals(2, matchedDocIds.length); - assertEquals(0, matchedDocIds[0]); - assertEquals(1, matchedDocIds[1]); - - matchedDocIds = reader.getDocIds("zoo.*").toArray(); - assertEquals(1, matchedDocIds.length); - assertEquals(3, matchedDocIds[0]); + matchedDocIds = reader.getDocIds("zoo.*").toArray(); + assertEquals(1, matchedDocIds.length); + assertEquals(3, matchedDocIds[0]); + } finally { + reader.closeInTest(); + } } } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/RangeIndexCreatorTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/RangeIndexCreatorTest.java index 5bee89764a22..2f29bbb3de62 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/RangeIndexCreatorTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/RangeIndexCreatorTest.java @@ -23,6 +23,7 @@ import java.util.Arrays; import java.util.Random; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.segment.creator.impl.inv.RangeIndexCreator; import org.apache.pinot.segment.local.segment.index.readers.RangeIndexReaderImpl; import org.apache.pinot.segment.spi.index.reader.RangeIndexReader; @@ -42,7 +43,7 @@ import static org.testng.Assert.assertNull; -public class RangeIndexCreatorTest { +public class RangeIndexCreatorTest implements PinotBuffersAfterMethodCheckRule { private static final File INDEX_DIR = new File(FileUtils.getTempDirectory(), "RangeIndexCreatorTest"); private static final Random RANDOM = new Random(42); private static final String COLUMN_NAME = "testColumn"; diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/RawIndexCreatorTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/RawIndexCreatorTest.java index 0a20736e1730..69746ae72599 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/RawIndexCreatorTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/RawIndexCreatorTest.java @@ -26,6 +26,7 @@ import java.util.Random; import org.apache.commons.io.FileUtils; import org.apache.commons.lang3.RandomStringUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterClassCheckRule; import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; import org.apache.pinot.segment.local.segment.index.forward.ForwardIndexReaderFactory; import org.apache.pinot.segment.local.segment.index.readers.forward.ChunkReaderContext; @@ -58,7 +59,7 @@ /** * Class for testing Raw index creators. */ -public class RawIndexCreatorTest { +public class RawIndexCreatorTest implements PinotBuffersAfterClassCheckRule { private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(), RawIndexCreatorTest.class.getSimpleName()); @@ -239,7 +240,8 @@ public void testBytesMVRawIndexCreator() throws Exception { PinotDataBuffer indexBuffer = getIndexBufferForColumn(BYTES_MV_COLUMN); try (VarByteChunkMVForwardIndexReader rawIndexReader = new VarByteChunkMVForwardIndexReader(indexBuffer, - DataType.BYTES); ChunkReaderContext readerContext = rawIndexReader.createContext()) { + DataType.BYTES); + ChunkReaderContext readerContext = rawIndexReader.createContext()) { _recordReader.rewind(); int maxNumberOfMultiValues = _segmentDirectory.getSegmentMetadata().getColumnMetadataFor(BYTES_MV_COLUMN).getMaxNumberOfMultiValues(); diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/SegmentGenerationWithBytesTypeTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/SegmentGenerationWithBytesTypeTest.java index 90bff63662bb..7fee205f6341 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/SegmentGenerationWithBytesTypeTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/SegmentGenerationWithBytesTypeTest.java @@ -26,6 +26,7 @@ import java.util.List; import java.util.Random; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterClassCheckRule; import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; import org.apache.pinot.segment.local.segment.index.readers.BaseImmutableDictionary; @@ -54,7 +55,7 @@ /** * Class for testing segment generation with byte[] data type. */ -public class SegmentGenerationWithBytesTypeTest { +public class SegmentGenerationWithBytesTypeTest implements PinotBuffersAfterClassCheckRule { private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(), SegmentGenerationWithBytesTypeTest.class.getSimpleName()); @@ -120,32 +121,35 @@ public void testMetadata() { } @Test - public void testDictionary() { - BaseImmutableDictionary dictionary = (BaseImmutableDictionary) _segment.getDictionary(FIXED_BYTE_SORTED_COLUMN); - assertEquals(dictionary.length(), NUM_SORTED_VALUES); + public void testDictionary() + throws IOException { + try (BaseImmutableDictionary dictionary = (BaseImmutableDictionary) _segment.getDictionary( + FIXED_BYTE_SORTED_COLUMN)) { + assertEquals(dictionary.length(), NUM_SORTED_VALUES); - // Test dictionary indexing. - for (int i = 0; i < NUM_ROWS; i++) { - int value = (i * NUM_SORTED_VALUES) / NUM_ROWS; - // For sorted columns, values are written as 0, 0, 0.., 1, 1, 1...n, n, n - assertEquals(dictionary.indexOf(BytesUtils.toHexString(Ints.toByteArray(value))), value % NUM_SORTED_VALUES); - } + // Test dictionary indexing. + for (int i = 0; i < NUM_ROWS; i++) { + int value = (i * NUM_SORTED_VALUES) / NUM_ROWS; + // For sorted columns, values are written as 0, 0, 0.., 1, 1, 1...n, n, n + assertEquals(dictionary.indexOf(BytesUtils.toHexString(Ints.toByteArray(value))), value % NUM_SORTED_VALUES); + } - // Test value not in dictionary. - assertEquals(dictionary.indexOf(BytesUtils.toHexString(Ints.toByteArray(NUM_SORTED_VALUES + 1))), -1); - assertEquals(dictionary.insertionIndexOf(BytesUtils.toHexString(Ints.toByteArray(NUM_SORTED_VALUES + 1))), - -(NUM_SORTED_VALUES + 1)); + // Test value not in dictionary. + assertEquals(dictionary.indexOf(BytesUtils.toHexString(Ints.toByteArray(NUM_SORTED_VALUES + 1))), -1); + assertEquals(dictionary.insertionIndexOf(BytesUtils.toHexString(Ints.toByteArray(NUM_SORTED_VALUES + 1))), + -(NUM_SORTED_VALUES + 1)); - int[] dictIds = new int[NUM_SORTED_VALUES]; - for (int i = 0; i < NUM_SORTED_VALUES; i++) { - dictIds[i] = i; - } + int[] dictIds = new int[NUM_SORTED_VALUES]; + for (int i = 0; i < NUM_SORTED_VALUES; i++) { + dictIds[i] = i; + } - byte[][] values = new byte[NUM_SORTED_VALUES][]; - dictionary.readBytesValues(dictIds, NUM_SORTED_VALUES, values); - for (int expected = 0; expected < NUM_SORTED_VALUES; expected++) { - int actual = ByteBuffer.wrap(values[expected]).asIntBuffer().get(); - assertEquals(actual, expected); + byte[][] values = new byte[NUM_SORTED_VALUES][]; + dictionary.readBytesValues(dictIds, NUM_SORTED_VALUES, values); + for (int expected = 0; expected < NUM_SORTED_VALUES; expected++) { + int actual = ByteBuffer.wrap(values[expected]).asIntBuffer().get(); + assertEquals(actual, expected); + } } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/SegmentGenerationWithFilterRecordsTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/SegmentGenerationWithFilterRecordsTest.java index b6c495d820e2..06f9938c7a85 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/SegmentGenerationWithFilterRecordsTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/SegmentGenerationWithFilterRecordsTest.java @@ -23,6 +23,7 @@ import java.util.ArrayList; import java.util.List; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader; import org.apache.pinot.segment.local.segment.readers.PinotSegmentRecordReader; @@ -45,7 +46,7 @@ /** * Tests filtering of records during segment generation */ -public class SegmentGenerationWithFilterRecordsTest { +public class SegmentGenerationWithFilterRecordsTest implements PinotBuffersAfterMethodCheckRule { private static final String STRING_COLUMN = "col1"; private static final String[] STRING_VALUES = {"A", "B", "C", "D", "E"}; private static final String LONG_COLUMN = "col2"; @@ -87,11 +88,12 @@ public void testNumDocs() File segmentDir = buildSegment(_tableConfig, _schema); SegmentMetadataImpl metadata = new SegmentMetadataImpl(segmentDir); Assert.assertEquals(metadata.getTotalDocs(), 2); - PinotSegmentRecordReader segmentRecordReader = new PinotSegmentRecordReader(segmentDir); - GenericRow next = segmentRecordReader.next(); - Assert.assertEquals(next.getValue(STRING_COLUMN), "C"); - next = segmentRecordReader.next(); - Assert.assertEquals(next.getValue(STRING_COLUMN), "E"); + try (PinotSegmentRecordReader segmentRecordReader = new PinotSegmentRecordReader(segmentDir)) { + GenericRow next = segmentRecordReader.next(); + Assert.assertEquals(next.getValue(STRING_COLUMN), "C"); + next = segmentRecordReader.next(); + Assert.assertEquals(next.getValue(STRING_COLUMN), "E"); + } } private File buildSegment(final TableConfig tableConfig, final Schema schema) diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/SegmentGenerationWithMinMaxTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/SegmentGenerationWithMinMaxTest.java index 1d6900a0c029..82d4c627a6ac 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/SegmentGenerationWithMinMaxTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/SegmentGenerationWithMinMaxTest.java @@ -22,6 +22,7 @@ import java.util.ArrayList; import java.util.List; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader; import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; @@ -41,7 +42,7 @@ /** * Tests filtering of records during segment generation */ -public class SegmentGenerationWithMinMaxTest { +public class SegmentGenerationWithMinMaxTest implements PinotBuffersAfterMethodCheckRule { private static final String STRING_COLUMN = "col1"; private static final String[] STRING_VALUES_WITH_COMMA_CHARACTER = {"A,,", ",B,", "C,Z,", "D,", "E,"}; diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/SegmentGenerationWithNoRecordsTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/SegmentGenerationWithNoRecordsTest.java index 7394378e8b4c..28c1f7ef35fd 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/SegmentGenerationWithNoRecordsTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/SegmentGenerationWithNoRecordsTest.java @@ -23,6 +23,7 @@ import java.util.Collections; import java.util.concurrent.TimeUnit; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader; import org.apache.pinot.segment.local.segment.readers.PinotSegmentRecordReader; @@ -42,7 +43,7 @@ /** * Tests segment generation for empty files */ -public class SegmentGenerationWithNoRecordsTest { +public class SegmentGenerationWithNoRecordsTest implements PinotBuffersAfterMethodCheckRule { private static final String STRING_COLUMN1 = "string_col1"; private static final String STRING_COLUMN2 = "string_col2"; private static final String STRING_COLUMN3 = "string_col3"; @@ -96,8 +97,9 @@ public void testNumDocs() Assert.assertEquals(metadata.getTimeUnit(), TimeUnit.MILLISECONDS); Assert.assertEquals(metadata.getStartTime(), metadata.getEndTime()); Assert.assertTrue(metadata.getAllColumns().containsAll(_schema.getColumnNames())); - PinotSegmentRecordReader segmentRecordReader = new PinotSegmentRecordReader(segmentDir); - Assert.assertFalse(segmentRecordReader.hasNext()); + try (PinotSegmentRecordReader segmentRecordReader = new PinotSegmentRecordReader(segmentDir)) { + Assert.assertFalse(segmentRecordReader.hasNext()); + } } private File buildSegment(final TableConfig tableConfig, final Schema schema) diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/SegmentGenerationWithTimeColumnTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/SegmentGenerationWithTimeColumnTest.java index 78d84d9a1156..779900d5b585 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/SegmentGenerationWithTimeColumnTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/SegmentGenerationWithTimeColumnTest.java @@ -26,6 +26,7 @@ import java.util.Random; import java.util.concurrent.TimeUnit; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader; import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; @@ -49,7 +50,7 @@ import org.testng.annotations.Test; -public class SegmentGenerationWithTimeColumnTest { +public class SegmentGenerationWithTimeColumnTest implements PinotBuffersAfterMethodCheckRule { private static final String STRING_COL_NAME = "someString"; private static final String TIME_COL_NAME = "date"; private static final String TIME_COL_FORMAT_NO_ZONE = "yyyyMMdd"; diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/SegmentPartitionTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/SegmentPartitionTest.java index 140904b4cb96..189f7a57e5e8 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/SegmentPartitionTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/SegmentPartitionTest.java @@ -28,6 +28,7 @@ import java.util.Random; import java.util.Set; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterClassCheckRule; import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader; @@ -60,7 +61,7 @@ *

  • Test to cover segment pruning during query execution.
    • * */ -public class SegmentPartitionTest { +public class SegmentPartitionTest implements PinotBuffersAfterClassCheckRule { private static final String SEGMENT_DIR_NAME = System.getProperty("java.io.tmpdir") + File.separator + "partitionTest"; private static final String TABLE_NAME = "partitionTable"; @@ -89,6 +90,7 @@ public void init() */ @AfterClass public void cleanup() { + _segment.destroy(); FileUtils.deleteQuietly(new File(SEGMENT_DIR_NAME)); } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/VarByteChunkV4Test.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/VarByteChunkV4Test.java index 8387c93dc838..d329d7d730fb 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/VarByteChunkV4Test.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/VarByteChunkV4Test.java @@ -32,6 +32,7 @@ import java.util.stream.IntStream; import java.util.stream.Stream; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterClassCheckRule; import org.apache.pinot.segment.local.io.writer.impl.VarByteChunkForwardIndexWriterV4; import org.apache.pinot.segment.local.segment.index.readers.forward.VarByteChunkForwardIndexReaderV4; import org.apache.pinot.segment.spi.compression.ChunkCompressionType; @@ -45,7 +46,7 @@ import static org.testng.Assert.assertEquals; -public class VarByteChunkV4Test { +public class VarByteChunkV4Test implements PinotBuffersAfterClassCheckRule { private static File[] _dirs; @@ -162,9 +163,9 @@ private void testWriteRead(File file, ChunkCompressionType compressionType, write.accept(writer, value); } } - try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile(file)) { - try (VarByteChunkForwardIndexReaderV4 reader = new VarByteChunkForwardIndexReaderV4(buffer, dataType, - true); VarByteChunkForwardIndexReaderV4.ReaderContext context = reader.createContext()) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile(file); + VarByteChunkForwardIndexReaderV4 reader = new VarByteChunkForwardIndexReaderV4(buffer, dataType, true); + VarByteChunkForwardIndexReaderV4.ReaderContext context = reader.createContext()) { for (int i = 0; i < values.size(); i++) { assertEquals(read.read(reader, context, i), values.get(i)); } @@ -190,7 +191,6 @@ private void testWriteRead(File file, ChunkCompressionType compressionType, assertEquals(read.read(reader, context, i), values.get(i)); } } - } } protected Stream randomStrings(int count, int lengthOfLongestEntry) { diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/VarByteChunkV5Test.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/VarByteChunkV5Test.java index c5bbc75c2760..29497ea1dd42 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/VarByteChunkV5Test.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/VarByteChunkV5Test.java @@ -29,6 +29,7 @@ import java.util.function.Function; import java.util.stream.Collectors; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterClassCheckRule; import org.apache.pinot.segment.local.io.writer.impl.VarByteChunkForwardIndexWriterV5; import org.apache.pinot.segment.local.segment.creator.impl.fwd.MultiValueFixedByteRawIndexCreator; import org.apache.pinot.segment.local.segment.index.readers.forward.VarByteChunkForwardIndexReaderV5; @@ -44,7 +45,7 @@ import static org.testng.Assert.assertEquals; -public class VarByteChunkV5Test extends VarByteChunkV4Test { +public class VarByteChunkV5Test extends VarByteChunkV4Test implements PinotBuffersAfterClassCheckRule { private static final Random RANDOM = new Random(); private static File[] _dirs; diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/inv/BitmapInvertedIndexWriterTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/inv/BitmapInvertedIndexWriterTest.java index 63842e54fda2..e81a143af1fb 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/inv/BitmapInvertedIndexWriterTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/inv/BitmapInvertedIndexWriterTest.java @@ -27,6 +27,7 @@ import java.util.stream.Collectors; import java.util.stream.IntStream; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.segment.creator.impl.inv.BitmapInvertedIndexWriter; import org.apache.pinot.segment.local.segment.index.readers.BitmapInvertedIndexReader; import org.apache.pinot.segment.spi.memory.PinotDataBuffer; @@ -43,7 +44,7 @@ import static org.testng.Assert.assertEquals; -public class BitmapInvertedIndexWriterTest { +public class BitmapInvertedIndexWriterTest implements PinotBuffersAfterMethodCheckRule { private static final File INDEX_DIR = new File(FileUtils.getTempDirectory(), "BitmapInvertedIndexWriterTest"); diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/FixedBitMVEntryDictForwardIndexTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/FixedBitMVEntryDictForwardIndexTest.java index 12d3b2c352e5..7019ce64f043 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/FixedBitMVEntryDictForwardIndexTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/FixedBitMVEntryDictForwardIndexTest.java @@ -22,6 +22,7 @@ import java.io.IOException; import java.util.Random; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.io.writer.impl.FixedBitMVEntryDictForwardIndexWriter; import org.apache.pinot.segment.local.segment.index.readers.forward.FixedBitMVEntryDictForwardIndexReader; import org.apache.pinot.segment.spi.V1Constants; @@ -33,7 +34,7 @@ import static org.testng.Assert.assertEquals; -public class FixedBitMVEntryDictForwardIndexTest { +public class FixedBitMVEntryDictForwardIndexTest implements PinotBuffersAfterMethodCheckRule { private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(), "FixedBitMVEntryDictForwardIndexTest"); private static final File INDEX_FILE = new File(TEMP_DIR, "testColumn" + V1Constants.Indexes.UNSORTED_MV_FORWARD_INDEX_FILE_EXTENSION); diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/FixedBitMVForwardIndexTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/FixedBitMVForwardIndexTest.java index 5493cb6c3dd6..eb18a2349c6e 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/FixedBitMVForwardIndexTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/FixedBitMVForwardIndexTest.java @@ -24,6 +24,7 @@ import java.util.List; import java.util.Random; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.io.writer.impl.FixedBitMVForwardIndexWriter; import org.apache.pinot.segment.local.segment.index.readers.forward.FixedBitMVForwardIndexReader; import org.apache.pinot.segment.spi.V1Constants; @@ -37,7 +38,7 @@ import static org.testng.Assert.assertEquals; -public class FixedBitMVForwardIndexTest { +public class FixedBitMVForwardIndexTest implements PinotBuffersAfterMethodCheckRule { private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(), "FixedBitMVForwardIndexTest"); private static final File INDEX_FILE = new File(TEMP_DIR, "testColumn" + V1Constants.Indexes.UNSORTED_MV_FORWARD_INDEX_FILE_EXTENSION); diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/FixedBitSVForwardIndexReaderTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/FixedBitSVForwardIndexReaderTest.java index b392b41677d4..26d76661d0c6 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/FixedBitSVForwardIndexReaderTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/FixedBitSVForwardIndexReaderTest.java @@ -23,6 +23,7 @@ import java.util.ArrayList; import java.util.Random; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.io.writer.impl.FixedBitSVForwardIndexWriter; import org.apache.pinot.segment.local.segment.index.readers.forward.FixedBitSVForwardIndexReader; import org.apache.pinot.segment.spi.V1Constants; @@ -34,7 +35,8 @@ import static org.testng.Assert.assertEquals; -public class FixedBitSVForwardIndexReaderTest { + +public class FixedBitSVForwardIndexReaderTest implements PinotBuffersAfterMethodCheckRule { private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(), "FixedBitMVForwardIndexTest"); private static final File INDEX_FILE = new File(TEMP_DIR, "testColumn" + V1Constants.Indexes.UNSORTED_SV_FORWARD_INDEX_FILE_EXTENSION); diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/FixedByteChunkSVForwardIndexTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/FixedByteChunkSVForwardIndexTest.java index 7cfe87039871..3587345435e7 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/FixedByteChunkSVForwardIndexTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/FixedByteChunkSVForwardIndexTest.java @@ -24,6 +24,7 @@ import java.util.Random; import java.util.stream.IntStream; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.io.writer.impl.FixedByteChunkForwardIndexWriter; import org.apache.pinot.segment.local.segment.index.readers.forward.ChunkReaderContext; import org.apache.pinot.segment.local.segment.index.readers.forward.FixedByteChunkSVForwardIndexReader; @@ -46,7 +47,7 @@ * * Number of docs and docs per chunk are chosen to generate complete as well partial chunks. */ -public class FixedByteChunkSVForwardIndexTest { +public class FixedByteChunkSVForwardIndexTest implements PinotBuffersAfterMethodCheckRule { private static final int NUM_VALUES = 10009; private static final int NUM_DOCS_PER_CHUNK = 5003; private static final String TEST_FILE = System.getProperty("java.io.tmpdir") + File.separator + "FixedByteSVRTest"; @@ -84,20 +85,17 @@ public void testInt(ChunkCompressionType compressionType, int version) } } - try (ForwardIndexReader fourByteOffsetReader = version >= 4 - ? new FixedBytePower2ChunkSVForwardIndexReader( - PinotDataBuffer.mapReadOnlyBigEndianFile(outFileFourByte), DataType.INT) - : new FixedByteChunkSVForwardIndexReader( - PinotDataBuffer.mapReadOnlyBigEndianFile(outFileFourByte), DataType.INT); - ChunkReaderContext fourByteOffsetReaderContext = fourByteOffsetReader - .createContext(); + try (PinotDataBuffer buffer1 = PinotDataBuffer.mapReadOnlyBigEndianFile(outFileFourByte); + ForwardIndexReader fourByteOffsetReader = version >= 4 + ? new FixedBytePower2ChunkSVForwardIndexReader(buffer1, DataType.INT) + : new FixedByteChunkSVForwardIndexReader(buffer1, DataType.INT); + ChunkReaderContext fourByteOffsetReaderContext = fourByteOffsetReader.createContext(); + PinotDataBuffer buffer2 = PinotDataBuffer.mapReadOnlyBigEndianFile(outFileEightByte); ForwardIndexReader eightByteOffsetReader = version >= 4 - ? new FixedBytePower2ChunkSVForwardIndexReader( - PinotDataBuffer.mapReadOnlyBigEndianFile(outFileEightByte), DataType.INT) - : new FixedByteChunkSVForwardIndexReader( - PinotDataBuffer.mapReadOnlyBigEndianFile(outFileEightByte), DataType.INT); - ChunkReaderContext eightByteOffsetReaderContext = eightByteOffsetReader - .createContext()) { + ? new FixedBytePower2ChunkSVForwardIndexReader(buffer2, DataType.INT) + : new FixedByteChunkSVForwardIndexReader(buffer2, DataType.INT); + ChunkReaderContext eightByteOffsetReaderContext = eightByteOffsetReader.createContext()) { + for (int i = 0; i < NUM_VALUES; i++) { Assert.assertEquals(fourByteOffsetReader.getInt(i, fourByteOffsetReaderContext), expected[i]); Assert.assertEquals(eightByteOffsetReader.getInt(i, eightByteOffsetReaderContext), expected[i]); @@ -150,20 +148,17 @@ public void testLong(ChunkCompressionType compressionType, int version) } } - try (ForwardIndexReader fourByteOffsetReader = version >= 4 - ? new FixedBytePower2ChunkSVForwardIndexReader( - PinotDataBuffer.mapReadOnlyBigEndianFile(outFileFourByte), DataType.LONG) - : new FixedByteChunkSVForwardIndexReader( - PinotDataBuffer.mapReadOnlyBigEndianFile(outFileFourByte), DataType.LONG); - ChunkReaderContext fourByteOffsetReaderContext = fourByteOffsetReader - .createContext(); + try (PinotDataBuffer buffer1 = PinotDataBuffer.mapReadOnlyBigEndianFile(outFileFourByte); + ForwardIndexReader fourByteOffsetReader = version >= 4 + ? new FixedBytePower2ChunkSVForwardIndexReader(buffer1, DataType.LONG) + : new FixedByteChunkSVForwardIndexReader(buffer1, DataType.LONG); + ChunkReaderContext fourByteOffsetReaderContext = fourByteOffsetReader.createContext(); + PinotDataBuffer buffer2 = PinotDataBuffer.mapReadOnlyBigEndianFile(outFileEightByte); ForwardIndexReader eightByteOffsetReader = version >= 4 - ? new FixedBytePower2ChunkSVForwardIndexReader( - PinotDataBuffer.mapReadOnlyBigEndianFile(outFileEightByte), DataType.LONG) - : new FixedByteChunkSVForwardIndexReader( - PinotDataBuffer.mapReadOnlyBigEndianFile(outFileEightByte), DataType.LONG); - ChunkReaderContext eightByteOffsetReaderContext = eightByteOffsetReader - .createContext()) { + ? new FixedBytePower2ChunkSVForwardIndexReader(buffer2, DataType.LONG) + : new FixedByteChunkSVForwardIndexReader(buffer2, DataType.LONG); + ChunkReaderContext eightByteOffsetReaderContext = eightByteOffsetReader.createContext()) { + for (int i = 0; i < NUM_VALUES; i++) { Assert.assertEquals(fourByteOffsetReader.getLong(i, fourByteOffsetReaderContext), expected[i]); Assert.assertEquals(eightByteOffsetReader.getLong(i, eightByteOffsetReaderContext), expected[i]); @@ -215,20 +210,16 @@ public void testFloat(ChunkCompressionType compressionType, int version) } } - try (ForwardIndexReader fourByteOffsetReader = version >= 4 - ? new FixedBytePower2ChunkSVForwardIndexReader( - PinotDataBuffer.mapReadOnlyBigEndianFile(outFileFourByte), DataType.FLOAT) - : new FixedByteChunkSVForwardIndexReader( - PinotDataBuffer.mapReadOnlyBigEndianFile(outFileFourByte), DataType.FLOAT); - ChunkReaderContext fourByteOffsetReaderContext = fourByteOffsetReader - .createContext(); + try (PinotDataBuffer buffer1 = PinotDataBuffer.mapReadOnlyBigEndianFile(outFileFourByte); + ForwardIndexReader fourByteOffsetReader = version >= 4 + ? new FixedBytePower2ChunkSVForwardIndexReader(buffer1, DataType.FLOAT) + : new FixedByteChunkSVForwardIndexReader(buffer1, DataType.FLOAT); + ChunkReaderContext fourByteOffsetReaderContext = fourByteOffsetReader.createContext(); + PinotDataBuffer buffer2 = PinotDataBuffer.mapReadOnlyBigEndianFile(outFileEightByte); ForwardIndexReader eightByteOffsetReader = version >= 4 - ? new FixedBytePower2ChunkSVForwardIndexReader(PinotDataBuffer.mapReadOnlyBigEndianFile(outFileEightByte), - DataType.FLOAT) - : new FixedByteChunkSVForwardIndexReader( - PinotDataBuffer.mapReadOnlyBigEndianFile(outFileEightByte), DataType.FLOAT); - ChunkReaderContext eightByteOffsetReaderContext = eightByteOffsetReader - .createContext()) { + ? new FixedBytePower2ChunkSVForwardIndexReader(buffer2, DataType.FLOAT) + : new FixedByteChunkSVForwardIndexReader(buffer2, DataType.FLOAT); + ChunkReaderContext eightByteOffsetReaderContext = eightByteOffsetReader.createContext()) { for (int i = 0; i < NUM_VALUES; i++) { Assert.assertEquals(fourByteOffsetReader.getFloat(i, fourByteOffsetReaderContext), expected[i]); Assert.assertEquals(eightByteOffsetReader.getFloat(i, eightByteOffsetReaderContext), expected[i]); @@ -280,18 +271,15 @@ public void testDouble(ChunkCompressionType compressionType, int version) } } - try (ForwardIndexReader fourByteOffsetReader = version >= 4 - ? new FixedBytePower2ChunkSVForwardIndexReader( - PinotDataBuffer.mapReadOnlyBigEndianFile(outFileFourByte), DataType.DOUBLE) - : new FixedByteChunkSVForwardIndexReader( - PinotDataBuffer.mapReadOnlyBigEndianFile(outFileFourByte), DataType.DOUBLE); - ChunkReaderContext fourByteOffsetReaderContext = fourByteOffsetReader - .createContext(); + try (PinotDataBuffer buffer1 = PinotDataBuffer.mapReadOnlyBigEndianFile(outFileFourByte); + ForwardIndexReader fourByteOffsetReader = version >= 4 + ? new FixedBytePower2ChunkSVForwardIndexReader(buffer1, DataType.DOUBLE) + : new FixedByteChunkSVForwardIndexReader(buffer1, DataType.DOUBLE); + ChunkReaderContext fourByteOffsetReaderContext = fourByteOffsetReader.createContext(); + PinotDataBuffer buffer2 = PinotDataBuffer.mapReadOnlyBigEndianFile(outFileEightByte); ForwardIndexReader eightByteOffsetReader = version >= 4 - ? new FixedBytePower2ChunkSVForwardIndexReader( - PinotDataBuffer.mapReadOnlyBigEndianFile(outFileEightByte), DataType.DOUBLE) - : new FixedByteChunkSVForwardIndexReader( - PinotDataBuffer.mapReadOnlyBigEndianFile(outFileEightByte), DataType.DOUBLE); + ? new FixedBytePower2ChunkSVForwardIndexReader(buffer2, DataType.DOUBLE) + : new FixedByteChunkSVForwardIndexReader(buffer2, DataType.DOUBLE); ChunkReaderContext eightByteOffsetReaderContext = eightByteOffsetReader .createContext()) { for (int i = 0; i < NUM_VALUES; i++) { @@ -364,8 +352,8 @@ private void testBackwardCompatibilityHelper(String fileName, int numDocs, doubl throw new RuntimeException("Input file not found: " + fileName); } File file = new File(resource.getFile()); - try (FixedByteChunkSVForwardIndexReader reader = new FixedByteChunkSVForwardIndexReader( - PinotDataBuffer.mapReadOnlyBigEndianFile(file), DataType.DOUBLE); + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile(file); + FixedByteChunkSVForwardIndexReader reader = new FixedByteChunkSVForwardIndexReader(buffer, DataType.DOUBLE); ChunkReaderContext readerContext = reader.createContext()) { for (int i = 0; i < numDocs; i++) { double actual = reader.getDouble(i, readerContext); diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/SortedForwardIndexReaderTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/SortedForwardIndexReaderTest.java index a9eb2899c41e..28dabe9f2f80 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/SortedForwardIndexReaderTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/SortedForwardIndexReaderTest.java @@ -20,6 +20,7 @@ import java.io.File; import java.util.Random; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.io.writer.impl.FixedByteSingleValueMultiColWriter; import org.apache.pinot.segment.local.segment.index.readers.sorted.SortedIndexReaderImpl; import org.apache.pinot.segment.spi.memory.PinotDataBuffer; @@ -28,8 +29,9 @@ @Test -public class SortedForwardIndexReaderTest { +public class SortedForwardIndexReaderTest implements PinotBuffersAfterMethodCheckRule { + @Test public void testSimple() throws Exception { @@ -38,27 +40,30 @@ public void testSimple() File file = new File("test_sortef_fwd_index.dat"); file.delete(); int[] columnSizes = new int[]{4, 4}; - FixedByteSingleValueMultiColWriter writer = - new FixedByteSingleValueMultiColWriter(file, cardinality, columnSizes.length, columnSizes); - Random random = new Random(); + int[] startDocIdArray = new int[cardinality]; int[] endDocIdArray = new int[cardinality]; - int prevEnd = -1; - int totalDocs = 0; - for (int i = 0; i < cardinality; i++) { - int length = random.nextInt(maxLength); - int start = prevEnd + 1; - int end = start + length; - startDocIdArray[i] = start; - endDocIdArray[i] = end; - writer.setInt(i, 0, start); - writer.setInt(i, 1, end); - prevEnd = end; - totalDocs += length; + + try (FixedByteSingleValueMultiColWriter writer = + new FixedByteSingleValueMultiColWriter(file, cardinality, columnSizes.length, columnSizes)) { + Random random = new Random(); + int prevEnd = -1; + int totalDocs = 0; + for (int i = 0; i < cardinality; i++) { + int length = random.nextInt(maxLength); + int start = prevEnd + 1; + int end = start + length; + startDocIdArray[i] = start; + endDocIdArray[i] = end; + writer.setInt(i, 0, start); + writer.setInt(i, 1, end); + prevEnd = end; + totalDocs += length; + } } - writer.close(); - try (SortedIndexReaderImpl reader = new SortedIndexReaderImpl(PinotDataBuffer.loadBigEndianFile(file), cardinality); + try (PinotDataBuffer dataBuffer = PinotDataBuffer.loadBigEndianFile(file); + SortedIndexReaderImpl reader = new SortedIndexReaderImpl(dataBuffer, cardinality); SortedIndexReaderImpl.Context readerContext = reader.createContext()) { for (int i = 0; i < cardinality; i++) { for (int docId = startDocIdArray[i]; docId <= endDocIdArray[i]; docId++) { diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/VarByteChunkSVForwardIndexTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/VarByteChunkSVForwardIndexTest.java index 896cbb5b6cfe..d000c0b62c6a 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/VarByteChunkSVForwardIndexTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/VarByteChunkSVForwardIndexTest.java @@ -27,6 +27,7 @@ import org.apache.commons.io.FileUtils; import org.apache.commons.lang3.RandomStringUtils; import org.apache.commons.lang3.StringUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.io.writer.impl.VarByteChunkForwardIndexWriter; import org.apache.pinot.segment.local.segment.creator.impl.fwd.SingleValueVarByteRawIndexCreator; import org.apache.pinot.segment.local.segment.index.readers.forward.ChunkReaderContext; @@ -45,7 +46,7 @@ /** * Unit test for {@link VarByteChunkSVForwardIndexReader} and {@link VarByteChunkForwardIndexWriter} classes. */ -public class VarByteChunkSVForwardIndexTest { +public class VarByteChunkSVForwardIndexTest implements PinotBuffersAfterMethodCheckRule { private static final int NUM_ENTRIES = 5003; private static final int NUM_DOCS_PER_CHUNK = 1009; private static final int MAX_STRING_LENGTH = 101; @@ -121,11 +122,13 @@ public void test(ChunkCompressionType compressionType) } } - try (VarByteChunkSVForwardIndexReader fourByteOffsetReader = new VarByteChunkSVForwardIndexReader( - PinotDataBuffer.mapReadOnlyBigEndianFile(outFileFourByte), DataType.STRING); + try (PinotDataBuffer buffer1 = PinotDataBuffer.mapReadOnlyBigEndianFile(outFileFourByte); + VarByteChunkSVForwardIndexReader fourByteOffsetReader = new VarByteChunkSVForwardIndexReader( + buffer1, DataType.STRING); ChunkReaderContext fourByteOffsetReaderContext = fourByteOffsetReader.createContext(); + PinotDataBuffer buffer2 = PinotDataBuffer.mapReadOnlyBigEndianFile(outFileEightByte); VarByteChunkSVForwardIndexReader eightByteOffsetReader = new VarByteChunkSVForwardIndexReader( - PinotDataBuffer.mapReadOnlyBigEndianFile(outFileEightByte), DataType.STRING); + buffer2, DataType.STRING); ChunkReaderContext eightByteOffsetReaderContext = eightByteOffsetReader.createContext()) { for (int i = 0; i < NUM_ENTRIES; i++) { Assert.assertEquals(fourByteOffsetReader.getString(i, fourByteOffsetReaderContext), expected[i]); @@ -166,8 +169,8 @@ private void testBackwardCompatibilityHelper(String fileName, String[] data, int throw new RuntimeException("Input file not found: " + fileName); } File file = new File(resource.getFile()); - try (VarByteChunkSVForwardIndexReader reader = new VarByteChunkSVForwardIndexReader( - PinotDataBuffer.mapReadOnlyBigEndianFile(file), DataType.STRING); + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile(file); + VarByteChunkSVForwardIndexReader reader = new VarByteChunkSVForwardIndexReader(buffer, DataType.STRING); ChunkReaderContext readerContext = reader.createContext()) { for (int i = 0; i < numDocs; i++) { String actual = reader.getString(i, readerContext); @@ -246,8 +249,8 @@ private void testLargeVarcharHelper(ChunkCompressionType compressionType, int nu } } - PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile(outFile); - try (VarByteChunkSVForwardIndexReader reader = new VarByteChunkSVForwardIndexReader(buffer, DataType.STRING); + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile(outFile); + VarByteChunkSVForwardIndexReader reader = new VarByteChunkSVForwardIndexReader(buffer, DataType.STRING); ChunkReaderContext readerContext = reader.createContext()) { for (int i = 0; i < numDocs; i++) { Assert.assertEquals(reader.getString(i, readerContext), expected[i]); @@ -261,18 +264,20 @@ private void testLargeVarcharHelper(ChunkCompressionType compressionType, int nu // (75000 characters in each row and 10000 rows will hit this scenario). // So we specifically test for mapping the index file using the default factory // trying to exercise the buffer used in larger cases - buffer = PinotDataBuffer.createDefaultFactory(false) - .mapFile(outFile, outFile.canRead(), 0, outFile.length(), ByteOrder.BIG_ENDIAN); - assert !(buffer instanceof PinotByteBuffer) : "This test tries to exercise the long buffer algorithm"; + try (PinotDataBuffer buffer = PinotDataBuffer.createDefaultFactory(false) + .mapFile(outFile, outFile.canRead(), 0, outFile.length(), ByteOrder.BIG_ENDIAN)) { - try (VarByteChunkSVForwardIndexReader reader = new VarByteChunkSVForwardIndexReader(buffer, DataType.STRING); - ChunkReaderContext readerContext = reader.createContext()) { - for (int i = 0; i < numDocs; i++) { - Assert.assertEquals(reader.getString(i, readerContext), expected[i]); + assert !(buffer instanceof PinotByteBuffer) : "This test tries to exercise the long buffer algorithm"; + + try (VarByteChunkSVForwardIndexReader reader = new VarByteChunkSVForwardIndexReader(buffer, DataType.STRING); + ChunkReaderContext readerContext = reader.createContext()) { + for (int i = 0; i < numDocs; i++) { + Assert.assertEquals(reader.getString(i, readerContext), expected[i]); + } } - } - FileUtils.deleteQuietly(outFile); + FileUtils.deleteQuietly(outFile); + } } @Test(expectedExceptions = IllegalStateException.class) diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/mutable/CLPMutableForwardIndexTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/mutable/CLPMutableForwardIndexTest.java index fde7d677ab84..790be1b8ed4a 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/mutable/CLPMutableForwardIndexTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/mutable/CLPMutableForwardIndexTest.java @@ -21,6 +21,7 @@ import java.io.IOException; import java.util.ArrayList; import java.util.List; +import org.apache.pinot.segment.local.PinotBuffersAfterClassCheckRule; import org.apache.pinot.segment.local.io.writer.impl.DirectMemoryManager; import org.apache.pinot.segment.local.realtime.impl.forward.CLPMutableForwardIndex; import org.apache.pinot.segment.local.realtime.impl.forward.CLPMutableForwardIndexV2; @@ -35,7 +36,7 @@ import org.testng.annotations.Test; -public class CLPMutableForwardIndexTest { +public class CLPMutableForwardIndexTest implements PinotBuffersAfterClassCheckRule { private final List _logLines = new ArrayList<>() {{ add("2023/10/26 00:03:10.168 INFO [PropertyCache] [HelixController-pipeline-default-pinot-(4a02a32c_DEFAULT)] " + "Event pinot::DEFAULT::4a02a32c_DEFAULT : Refreshed 35 property LiveInstance took 5 ms. Selective:" diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/mutable/CLPMutableForwardIndexV2Test.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/mutable/CLPMutableForwardIndexV2Test.java index 6179a4b8bf68..4ab2aca48eeb 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/mutable/CLPMutableForwardIndexV2Test.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/mutable/CLPMutableForwardIndexV2Test.java @@ -27,6 +27,7 @@ import java.util.List; import java.util.concurrent.ThreadLocalRandom; import org.apache.commons.compress.compressors.gzip.GzipCompressorInputStream; +import org.apache.pinot.segment.local.PinotBuffersAfterClassCheckRule; import org.apache.pinot.segment.local.io.writer.impl.DirectMemoryManager; import org.apache.pinot.segment.local.realtime.impl.forward.CLPMutableForwardIndexV2; import org.apache.pinot.segment.spi.memory.PinotDataBufferMemoryManager; @@ -36,7 +37,7 @@ import org.testng.annotations.Test; -public class CLPMutableForwardIndexV2Test { +public class CLPMutableForwardIndexV2Test implements PinotBuffersAfterClassCheckRule { private PinotDataBufferMemoryManager _memoryManager; private List _logMessages = new ArrayList<>(); diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/mutable/FixedByteMVMutableForwardIndexTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/mutable/FixedByteMVMutableForwardIndexTest.java index 77fd706f3c64..993c9ee7c360 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/mutable/FixedByteMVMutableForwardIndexTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/mutable/FixedByteMVMutableForwardIndexTest.java @@ -22,6 +22,7 @@ import java.lang.reflect.Field; import java.util.Arrays; import java.util.Random; +import org.apache.pinot.segment.local.PinotBuffersAfterClassCheckRule; import org.apache.pinot.segment.local.io.writer.impl.DirectMemoryManager; import org.apache.pinot.segment.local.realtime.impl.forward.FixedByteMVMutableForwardIndex; import org.apache.pinot.segment.spi.memory.PinotDataBufferMemoryManager; @@ -32,7 +33,7 @@ import org.testng.annotations.Test; -public class FixedByteMVMutableForwardIndexTest { +public class FixedByteMVMutableForwardIndexTest implements PinotBuffersAfterClassCheckRule { private PinotDataBufferMemoryManager _memoryManager; @BeforeClass diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/mutable/FixedByteSVMutableForwardIndexTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/mutable/FixedByteSVMutableForwardIndexTest.java index cadea9a1fca7..358d6390bce4 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/mutable/FixedByteSVMutableForwardIndexTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/mutable/FixedByteSVMutableForwardIndexTest.java @@ -23,6 +23,7 @@ import java.io.IOException; import java.util.Arrays; import java.util.Random; +import org.apache.pinot.segment.local.PinotBuffersAfterClassCheckRule; import org.apache.pinot.segment.local.io.writer.impl.DirectMemoryManager; import org.apache.pinot.segment.local.realtime.impl.forward.FixedByteSVMutableForwardIndex; import org.apache.pinot.segment.spi.memory.PinotDataBufferMemoryManager; @@ -33,7 +34,7 @@ import org.testng.annotations.Test; -public class FixedByteSVMutableForwardIndexTest { +public class FixedByteSVMutableForwardIndexTest implements PinotBuffersAfterClassCheckRule { private PinotDataBufferMemoryManager _memoryManager; @BeforeClass diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/mutable/VarByteSVMutableForwardIndexTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/mutable/VarByteSVMutableForwardIndexTest.java index 918b52b2486a..46d8741d0a46 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/mutable/VarByteSVMutableForwardIndexTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/mutable/VarByteSVMutableForwardIndexTest.java @@ -21,6 +21,7 @@ import java.io.IOException; import java.util.Random; import org.apache.commons.lang3.RandomStringUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterClassCheckRule; import org.apache.pinot.segment.local.io.writer.impl.DirectMemoryManager; import org.apache.pinot.segment.local.realtime.impl.forward.VarByteSVMutableForwardIndex; import org.apache.pinot.segment.spi.memory.PinotDataBufferMemoryManager; @@ -33,7 +34,7 @@ import static java.nio.charset.StandardCharsets.UTF_8; -public class VarByteSVMutableForwardIndexTest { +public class VarByteSVMutableForwardIndexTest implements PinotBuffersAfterClassCheckRule { private PinotDataBufferMemoryManager _memoryManager; @BeforeClass diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/SegmentPreProcessorTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/SegmentPreProcessorTest.java index 45a1afbf7480..012c2d920d20 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/SegmentPreProcessorTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/SegmentPreProcessorTest.java @@ -35,6 +35,7 @@ import org.apache.commons.configuration2.ex.ConfigurationException; import org.apache.commons.io.FileUtils; import org.apache.commons.lang3.RandomStringUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterClassCheckRule; import org.apache.pinot.segment.local.segment.creator.SegmentTestUtils; import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; import org.apache.pinot.segment.local.segment.index.converter.SegmentV1V2ToV3FormatConverter; @@ -80,7 +81,7 @@ import static org.testng.Assert.*; -public class SegmentPreProcessorTest { +public class SegmentPreProcessorTest implements PinotBuffersAfterClassCheckRule { private static final String RAW_TABLE_NAME = "testTable"; private static final String SEGMENT_NAME = "testSegment"; private static final File TEMP_DIR = diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/reader/FixedByteWidthRowColDataFileReaderTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/reader/FixedByteWidthRowColDataFileReaderTest.java index c06961969792..5c63c7855e5a 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/reader/FixedByteWidthRowColDataFileReaderTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/reader/FixedByteWidthRowColDataFileReaderTest.java @@ -23,6 +23,7 @@ import java.io.FileOutputStream; import java.io.RandomAccessFile; import java.util.Random; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.io.reader.impl.FixedByteSingleValueMultiColReader; import org.apache.pinot.segment.spi.memory.PinotDataBuffer; import org.testng.Assert; @@ -30,7 +31,7 @@ @Test -public class FixedByteWidthRowColDataFileReaderTest { +public class FixedByteWidthRowColDataFileReaderTest implements PinotBuffersAfterMethodCheckRule { @Test void testSingleCol() @@ -52,8 +53,9 @@ void testSingleCol() RandomAccessFile raf = new RandomAccessFile(f, "rw"); raf.close(); - try (FixedByteSingleValueMultiColReader heapReader = new FixedByteSingleValueMultiColReader( - PinotDataBuffer.loadBigEndianFile(f), data.length, new int[]{4})) { + try (PinotDataBuffer buffer = PinotDataBuffer.loadBigEndianFile(f); + FixedByteSingleValueMultiColReader heapReader = new FixedByteSingleValueMultiColReader(buffer, data.length, + new int[]{4})) { heapReader.open(); for (int i = 0; i < data.length; i++) { Assert.assertEquals(heapReader.getInt(i, 0), data[i]); @@ -88,8 +90,9 @@ void testMultiCol() RandomAccessFile raf = new RandomAccessFile(f, "rw"); raf.close(); - try (FixedByteSingleValueMultiColReader heapReader = new FixedByteSingleValueMultiColReader( - PinotDataBuffer.loadBigEndianFile(f), numRows, new int[]{4, 4})) { + try (PinotDataBuffer buffer = PinotDataBuffer.loadBigEndianFile(f); + FixedByteSingleValueMultiColReader heapReader = new FixedByteSingleValueMultiColReader(buffer, numRows, + new int[]{4, 4})) { heapReader.open(); for (int i = 0; i < numRows; i++) { for (int j = 0; j < numCols; j++) { diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/ImmutableDictionaryTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/ImmutableDictionaryTest.java index 9b10688cd688..220028c36bc3 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/ImmutableDictionaryTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/ImmutableDictionaryTest.java @@ -31,6 +31,7 @@ import java.util.TreeSet; import org.apache.commons.io.FileUtils; import org.apache.commons.lang3.RandomStringUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.segment.creator.impl.SegmentDictionaryCreator; import org.apache.pinot.segment.spi.V1Constants; import org.apache.pinot.segment.spi.memory.PinotDataBuffer; @@ -49,7 +50,7 @@ import static org.testng.Assert.assertEquals; -public class ImmutableDictionaryTest { +public class ImmutableDictionaryTest implements PinotBuffersAfterMethodCheckRule { private static final FALFInterner STRING_INTERNER = new FALFInterner<>(500); private static final FALFInterner BYTE_INTERNER = new FALFInterner<>(500, Arrays::hashCode); private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(), "ImmutableDictionaryTest"); @@ -179,9 +180,9 @@ public void setUp() @Test public void testIntDictionary() throws Exception { - try (IntDictionary intDictionary = new IntDictionary( - PinotDataBuffer.mapReadOnlyBigEndianFile(new File(TEMP_DIR, INT_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), - NUM_VALUES)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, INT_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + IntDictionary intDictionary = new IntDictionary(buffer, NUM_VALUES)) { testIntDictionary(intDictionary); } } @@ -189,9 +190,9 @@ public void testIntDictionary() @Test public void testOnHeapIntDictionary() throws Exception { - try (OnHeapIntDictionary onHeapIntDictionary = new OnHeapIntDictionary( - PinotDataBuffer.mapReadOnlyBigEndianFile(new File(TEMP_DIR, INT_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), - NUM_VALUES)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, INT_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + OnHeapIntDictionary onHeapIntDictionary = new OnHeapIntDictionary(buffer, NUM_VALUES)) { testIntDictionary(onHeapIntDictionary); } } @@ -216,8 +217,9 @@ private void testIntDictionary(BaseImmutableDictionary intDictionary) { @Test public void testLongDictionary() throws Exception { - try (LongDictionary longDictionary = new LongDictionary(PinotDataBuffer.mapReadOnlyBigEndianFile( - new File(TEMP_DIR, LONG_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), NUM_VALUES)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, LONG_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + LongDictionary longDictionary = new LongDictionary(buffer, NUM_VALUES)) { testLongDictionary(longDictionary); } } @@ -225,8 +227,9 @@ public void testLongDictionary() @Test public void testOnHeapLongDictionary() throws Exception { - try (OnHeapLongDictionary onHeapLongDictionary = new OnHeapLongDictionary(PinotDataBuffer.mapReadOnlyBigEndianFile( - new File(TEMP_DIR, LONG_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), NUM_VALUES)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, LONG_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + OnHeapLongDictionary onHeapLongDictionary = new OnHeapLongDictionary(buffer, NUM_VALUES)) { testLongDictionary(onHeapLongDictionary); } } @@ -251,8 +254,9 @@ private void testLongDictionary(BaseImmutableDictionary longDictionary) { @Test public void testFloatDictionary() throws Exception { - try (FloatDictionary floatDictionary = new FloatDictionary(PinotDataBuffer.mapReadOnlyBigEndianFile( - new File(TEMP_DIR, FLOAT_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), NUM_VALUES)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, FLOAT_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + FloatDictionary floatDictionary = new FloatDictionary(buffer, NUM_VALUES)) { testFloatDictionary(floatDictionary); } } @@ -260,9 +264,9 @@ public void testFloatDictionary() @Test public void testOnHeapFloatDictionary() throws Exception { - try (OnHeapFloatDictionary onHeapFloatDictionary = new OnHeapFloatDictionary( - PinotDataBuffer.mapReadOnlyBigEndianFile( - new File(TEMP_DIR, FLOAT_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), NUM_VALUES)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, FLOAT_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + OnHeapFloatDictionary onHeapFloatDictionary = new OnHeapFloatDictionary(buffer, NUM_VALUES)) { testFloatDictionary(onHeapFloatDictionary); } } @@ -287,8 +291,9 @@ private void testFloatDictionary(BaseImmutableDictionary floatDictionary) { @Test public void testDoubleDictionary() throws Exception { - try (DoubleDictionary doubleDictionary = new DoubleDictionary(PinotDataBuffer.mapReadOnlyBigEndianFile( - new File(TEMP_DIR, DOUBLE_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), NUM_VALUES)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, DOUBLE_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + DoubleDictionary doubleDictionary = new DoubleDictionary(buffer, NUM_VALUES)) { testDoubleDictionary(doubleDictionary); } } @@ -296,9 +301,9 @@ public void testDoubleDictionary() @Test public void testOnHeapDoubleDictionary() throws Exception { - try (OnHeapDoubleDictionary onHeapDoubleDictionary = new OnHeapDoubleDictionary( - PinotDataBuffer.mapReadOnlyBigEndianFile( - new File(TEMP_DIR, DOUBLE_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), NUM_VALUES)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, DOUBLE_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + OnHeapDoubleDictionary onHeapDoubleDictionary = new OnHeapDoubleDictionary(buffer, NUM_VALUES)) { testDoubleDictionary(onHeapDoubleDictionary); } } @@ -323,8 +328,9 @@ private void testDoubleDictionary(BaseImmutableDictionary doubleDictionary) { @Test public void testBigDecimalDictionary() throws Exception { - try (BigDecimalDictionary bigDecimalDictionary = new BigDecimalDictionary(PinotDataBuffer.mapReadOnlyBigEndianFile( - new File(TEMP_DIR, BIG_DECIMAL_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), NUM_VALUES, + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, BIG_DECIMAL_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + BigDecimalDictionary bigDecimalDictionary = new BigDecimalDictionary(buffer, NUM_VALUES, _bigDecimalByteLength)) { testBigDecimalDictionary(bigDecimalDictionary); } @@ -333,10 +339,10 @@ public void testBigDecimalDictionary() @Test public void testOnHeapBigDecimalDictionary() throws Exception { - try (OnHeapBigDecimalDictionary onHeapBigDecimalDictionary = new OnHeapBigDecimalDictionary( - PinotDataBuffer.mapReadOnlyBigEndianFile( - new File(TEMP_DIR, BIG_DECIMAL_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), NUM_VALUES, - _bigDecimalByteLength)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, BIG_DECIMAL_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + OnHeapBigDecimalDictionary onHeapBigDecimalDictionary = new OnHeapBigDecimalDictionary(buffer, NUM_VALUES, + _bigDecimalByteLength)) { testBigDecimalDictionary(onHeapBigDecimalDictionary); } } @@ -362,9 +368,9 @@ private void testBigDecimalDictionary(BaseImmutableDictionary bigDecimalDictiona @Test public void testStringDictionary() throws Exception { - try (StringDictionary stringDictionary = new StringDictionary(PinotDataBuffer.mapReadOnlyBigEndianFile( - new File(TEMP_DIR, STRING_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), NUM_VALUES, - _numBytesPerStringValue)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, STRING_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + StringDictionary stringDictionary = new StringDictionary(buffer, NUM_VALUES, _numBytesPerStringValue)) { testStringDictionary(stringDictionary); } } @@ -372,10 +378,10 @@ public void testStringDictionary() @Test public void testOnHeapStringDictionary() throws Exception { - try (OnHeapStringDictionary onHeapStringDictionary = new OnHeapStringDictionary( - PinotDataBuffer.mapReadOnlyBigEndianFile( - new File(TEMP_DIR, STRING_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), NUM_VALUES, - _numBytesPerStringValue, null, null)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, STRING_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + OnHeapStringDictionary onHeapStringDictionary = new OnHeapStringDictionary(buffer, NUM_VALUES, + _numBytesPerStringValue, null, null)) { testStringDictionary(onHeapStringDictionary); } } @@ -383,10 +389,10 @@ public void testOnHeapStringDictionary() @Test public void testOnHeapStringDictionaryWithInterning() throws Exception { - try (OnHeapStringDictionary onHeapStringDictionary = new OnHeapStringDictionary( - PinotDataBuffer.mapReadOnlyBigEndianFile( - new File(TEMP_DIR, STRING_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), NUM_VALUES, - _numBytesPerStringValue, STRING_INTERNER, BYTE_INTERNER)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, STRING_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + OnHeapStringDictionary onHeapStringDictionary = new OnHeapStringDictionary(buffer, NUM_VALUES, + _numBytesPerStringValue, STRING_INTERNER, BYTE_INTERNER)) { testStringDictionary(onHeapStringDictionary); } } @@ -407,8 +413,9 @@ private void testStringDictionary(BaseImmutableDictionary stringDictionary) { @Test public void testBytesDictionary() throws Exception { - try (BytesDictionary bytesDictionary = new BytesDictionary(PinotDataBuffer.mapReadOnlyBigEndianFile( - new File(TEMP_DIR, BYTES_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), NUM_VALUES, BYTES_LENGTH)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, BYTES_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + BytesDictionary bytesDictionary = new BytesDictionary(buffer, NUM_VALUES, BYTES_LENGTH)) { testBytesDictionary(bytesDictionary); } } @@ -416,9 +423,10 @@ public void testBytesDictionary() @Test public void testOnHeapBytesDictionary() throws Exception { - try (OnHeapBytesDictionary onHeapBytesDictionary = new OnHeapBytesDictionary( - PinotDataBuffer.mapReadOnlyBigEndianFile( - new File(TEMP_DIR, BYTES_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), NUM_VALUES, BYTES_LENGTH, null)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, BYTES_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + OnHeapBytesDictionary onHeapBytesDictionary = new OnHeapBytesDictionary(buffer, NUM_VALUES, BYTES_LENGTH, + null)) { testBytesDictionary(onHeapBytesDictionary); } } @@ -427,10 +435,10 @@ public void testOnHeapBytesDictionary() public void testOnHeapBytesDictionaryWithInterning() throws Exception { - try (OnHeapBytesDictionary onHeapBytesDictionary = new OnHeapBytesDictionary( - PinotDataBuffer.mapReadOnlyBigEndianFile( - new File(TEMP_DIR, BYTES_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), NUM_VALUES, BYTES_LENGTH, - BYTE_INTERNER)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, BYTES_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + OnHeapBytesDictionary onHeapBytesDictionary = new OnHeapBytesDictionary(buffer, NUM_VALUES, BYTES_LENGTH, + BYTE_INTERNER)) { testBytesDictionary(onHeapBytesDictionary); } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/ImmutableDictionaryTypeConversionTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/ImmutableDictionaryTypeConversionTest.java index e47c415648f5..67509b12d991 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/ImmutableDictionaryTypeConversionTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/ImmutableDictionaryTypeConversionTest.java @@ -24,6 +24,7 @@ import java.util.Arrays; import java.util.Random; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.segment.creator.impl.SegmentDictionaryCreator; import org.apache.pinot.segment.spi.V1Constants; import org.apache.pinot.segment.spi.memory.PinotDataBuffer; @@ -44,7 +45,7 @@ import static org.testng.Assert.assertEquals; -public class ImmutableDictionaryTypeConversionTest { +public class ImmutableDictionaryTypeConversionTest implements PinotBuffersAfterMethodCheckRule { private static final FALFInterner STRING_INTERNER = new FALFInterner<>(128); private static final FALFInterner BYTE_INTERNER = new FALFInterner<>(128, Arrays::hashCode); @@ -174,9 +175,9 @@ public void setUp() @Test public void testIntDictionary() throws Exception { - try (IntDictionary intDictionary = new IntDictionary( - PinotDataBuffer.mapReadOnlyBigEndianFile(new File(TEMP_DIR, INT_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), - NUM_VALUES)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, INT_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + IntDictionary intDictionary = new IntDictionary(buffer, NUM_VALUES)) { testNumericDictionary(intDictionary); } } @@ -184,9 +185,9 @@ public void testIntDictionary() @Test public void testOnHeapIntDictionary() throws Exception { - try (OnHeapIntDictionary onHeapIntDictionary = new OnHeapIntDictionary( - PinotDataBuffer.mapReadOnlyBigEndianFile(new File(TEMP_DIR, INT_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), - NUM_VALUES)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, INT_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + OnHeapIntDictionary onHeapIntDictionary = new OnHeapIntDictionary(buffer, NUM_VALUES)) { testNumericDictionary(onHeapIntDictionary); } } @@ -194,8 +195,9 @@ public void testOnHeapIntDictionary() @Test public void testLongDictionary() throws Exception { - try (LongDictionary longDictionary = new LongDictionary(PinotDataBuffer.mapReadOnlyBigEndianFile( - new File(TEMP_DIR, LONG_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), NUM_VALUES)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, LONG_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + LongDictionary longDictionary = new LongDictionary(buffer, NUM_VALUES)) { testNumericDictionary(longDictionary); } } @@ -203,8 +205,9 @@ public void testLongDictionary() @Test public void testOnHeapLongDictionary() throws Exception { - try (OnHeapLongDictionary onHeapLongDictionary = new OnHeapLongDictionary(PinotDataBuffer.mapReadOnlyBigEndianFile( - new File(TEMP_DIR, LONG_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), NUM_VALUES)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, LONG_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + OnHeapLongDictionary onHeapLongDictionary = new OnHeapLongDictionary(buffer, NUM_VALUES)) { testNumericDictionary(onHeapLongDictionary); } } @@ -212,8 +215,9 @@ public void testOnHeapLongDictionary() @Test public void testFloatDictionary() throws Exception { - try (FloatDictionary floatDictionary = new FloatDictionary(PinotDataBuffer.mapReadOnlyBigEndianFile( - new File(TEMP_DIR, FLOAT_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), NUM_VALUES)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, FLOAT_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + FloatDictionary floatDictionary = new FloatDictionary(buffer, NUM_VALUES)) { testNumericDictionary(floatDictionary); } } @@ -221,9 +225,9 @@ public void testFloatDictionary() @Test public void testOnHeapFloatDictionary() throws Exception { - try (OnHeapFloatDictionary onHeapFloatDictionary = new OnHeapFloatDictionary( - PinotDataBuffer.mapReadOnlyBigEndianFile( - new File(TEMP_DIR, FLOAT_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), NUM_VALUES)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, FLOAT_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + OnHeapFloatDictionary onHeapFloatDictionary = new OnHeapFloatDictionary(buffer, NUM_VALUES)) { testNumericDictionary(onHeapFloatDictionary); } } @@ -231,8 +235,9 @@ public void testOnHeapFloatDictionary() @Test public void testDoubleDictionary() throws Exception { - try (DoubleDictionary doubleDictionary = new DoubleDictionary(PinotDataBuffer.mapReadOnlyBigEndianFile( - new File(TEMP_DIR, DOUBLE_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), NUM_VALUES)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, DOUBLE_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + DoubleDictionary doubleDictionary = new DoubleDictionary(buffer, NUM_VALUES)) { testNumericDictionary(doubleDictionary); } } @@ -240,9 +245,9 @@ public void testDoubleDictionary() @Test public void testOnHeapDoubleDictionary() throws Exception { - try (OnHeapDoubleDictionary onHeapDoubleDictionary = new OnHeapDoubleDictionary( - PinotDataBuffer.mapReadOnlyBigEndianFile( - new File(TEMP_DIR, DOUBLE_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), NUM_VALUES)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, DOUBLE_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + OnHeapDoubleDictionary onHeapDoubleDictionary = new OnHeapDoubleDictionary(buffer, NUM_VALUES)) { testNumericDictionary(onHeapDoubleDictionary); } } @@ -250,9 +255,10 @@ public void testOnHeapDoubleDictionary() @Test public void testBigDecimalDictionary() throws Exception { - try (BigDecimalDictionary bigDecimalDictionary = new BigDecimalDictionary(PinotDataBuffer.mapReadOnlyBigEndianFile( - new File(TEMP_DIR, BIG_DECIMAL_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), NUM_VALUES, - _bigDecimalByteLength)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, BIG_DECIMAL_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + BigDecimalDictionary bigDecimalDictionary = new BigDecimalDictionary(buffer, NUM_VALUES, + _bigDecimalByteLength)) { testNumericDictionary(bigDecimalDictionary); } } @@ -260,10 +266,10 @@ public void testBigDecimalDictionary() @Test public void testOnHeapBigDecimalDictionary() throws Exception { - try (OnHeapBigDecimalDictionary onHeapBigDecimalDictionary = new OnHeapBigDecimalDictionary( - PinotDataBuffer.mapReadOnlyBigEndianFile( - new File(TEMP_DIR, BIG_DECIMAL_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), NUM_VALUES, - _bigDecimalByteLength)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, BIG_DECIMAL_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + OnHeapBigDecimalDictionary onHeapBigDecimalDictionary = new OnHeapBigDecimalDictionary(buffer, NUM_VALUES, + _bigDecimalByteLength)) { testNumericDictionary(onHeapBigDecimalDictionary); } } @@ -315,8 +321,9 @@ private void testNumericDictionary(BaseImmutableDictionary dictionary) { @Test public void testStringDictionary() throws Exception { - try (StringDictionary stringDictionary = new StringDictionary(PinotDataBuffer.mapReadOnlyBigEndianFile( - new File(TEMP_DIR, STRING_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), NUM_VALUES, STRING_LENGTH)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, STRING_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + StringDictionary stringDictionary = new StringDictionary(buffer, NUM_VALUES, STRING_LENGTH)) { testStringDictionary(stringDictionary); } } @@ -324,9 +331,9 @@ public void testStringDictionary() @Test public void testOnHeapStringDictionary() throws Exception { - try (OnHeapStringDictionary onHeapStringDictionary = new OnHeapStringDictionary( - PinotDataBuffer.mapReadOnlyBigEndianFile( - new File(TEMP_DIR, STRING_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), NUM_VALUES, STRING_LENGTH, + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, STRING_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + OnHeapStringDictionary onHeapStringDictionary = new OnHeapStringDictionary(buffer, NUM_VALUES, STRING_LENGTH, null, null)) { testStringDictionary(onHeapStringDictionary); } @@ -335,16 +342,16 @@ public void testOnHeapStringDictionary() @Test public void testOnHeapStringDictionaryWithInterner() throws Exception { - try (OnHeapStringDictionary onHeapStringDictionary = new OnHeapStringDictionary( - PinotDataBuffer.mapReadOnlyBigEndianFile( - new File(TEMP_DIR, STRING_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), NUM_VALUES, STRING_LENGTH, + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, STRING_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + OnHeapStringDictionary onHeapStringDictionary = new OnHeapStringDictionary(buffer, NUM_VALUES, STRING_LENGTH, STRING_INTERNER, BYTE_INTERNER)) { testStringDictionary(onHeapStringDictionary); } - try (OnHeapStringDictionary onHeapStringDictionary = new OnHeapStringDictionary( - PinotDataBuffer.mapReadOnlyBigEndianFile( - new File(TEMP_DIR, STRING_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), NUM_VALUES, STRING_LENGTH, + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, STRING_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + OnHeapStringDictionary onHeapStringDictionary = new OnHeapStringDictionary(buffer, NUM_VALUES, STRING_LENGTH, STRING_INTERNER, BYTE_INTERNER)) { testStringDictionary(onHeapStringDictionary); } @@ -379,8 +386,9 @@ private void testStringDictionary(BaseImmutableDictionary dictionary) { @Test public void testBytesDictionary() throws Exception { - try (BytesDictionary bytesDictionary = new BytesDictionary(PinotDataBuffer.mapReadOnlyBigEndianFile( - new File(TEMP_DIR, BYTES_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), NUM_VALUES, BYTES_LENGTH)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, BYTES_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + BytesDictionary bytesDictionary = new BytesDictionary(buffer, NUM_VALUES, BYTES_LENGTH)) { testBytesDictionary(bytesDictionary); } } @@ -388,9 +396,10 @@ public void testBytesDictionary() @Test public void testOnHeapBytesDictionary() throws Exception { - try (OnHeapBytesDictionary onHeapBytesDictionary = new OnHeapBytesDictionary( - PinotDataBuffer.mapReadOnlyBigEndianFile( - new File(TEMP_DIR, BYTES_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), NUM_VALUES, BYTES_LENGTH, null)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, BYTES_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + OnHeapBytesDictionary onHeapBytesDictionary = new OnHeapBytesDictionary(buffer, NUM_VALUES, BYTES_LENGTH, + null)) { testBytesDictionary(onHeapBytesDictionary); } } @@ -398,10 +407,10 @@ public void testOnHeapBytesDictionary() @Test public void testOnHeapBytesDictionaryWithInterning() throws Exception { - try (OnHeapBytesDictionary onHeapBytesDictionary = new OnHeapBytesDictionary( - PinotDataBuffer.mapReadOnlyBigEndianFile( - new File(TEMP_DIR, BYTES_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)), NUM_VALUES, BYTES_LENGTH, - BYTE_INTERNER)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile( + new File(TEMP_DIR, BYTES_COLUMN_NAME + V1Constants.Dict.FILE_EXTENSION)); + OnHeapBytesDictionary onHeapBytesDictionary = new OnHeapBytesDictionary(buffer, NUM_VALUES, BYTES_LENGTH, + BYTE_INTERNER)) { testBytesDictionary(onHeapBytesDictionary); } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/NullValueVectorReaderImplTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/NullValueVectorReaderImplTest.java index 3bc7db28ede6..62908fba494c 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/NullValueVectorReaderImplTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/NullValueVectorReaderImplTest.java @@ -21,6 +21,7 @@ import java.io.File; import java.io.IOException; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.segment.creator.impl.nullvalue.NullValueVectorCreator; import org.apache.pinot.segment.spi.index.reader.NullValueVectorReader; import org.apache.pinot.segment.spi.memory.PinotDataBuffer; @@ -30,7 +31,7 @@ import org.testng.annotations.Test; -public class NullValueVectorReaderImplTest { +public class NullValueVectorReaderImplTest implements PinotBuffersAfterMethodCheckRule { private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(), "NullValueVectorReaderImplTest"); private static final String COLUMN_NAME = "test"; @@ -53,10 +54,11 @@ public void testNullValueVectorReader() throws IOException { Assert.assertEquals(TEMP_DIR.list().length, 1); File nullValueFile = new File(TEMP_DIR, TEMP_DIR.list()[0]); - PinotDataBuffer buffer = PinotDataBuffer.loadBigEndianFile(nullValueFile); - NullValueVectorReader reader = new NullValueVectorReaderImpl(buffer); - for (int i = 0; i < 100; i++) { - Assert.assertTrue(reader.isNull(i)); + try (PinotDataBuffer buffer = PinotDataBuffer.loadBigEndianFile(nullValueFile)) { + NullValueVectorReader reader = new NullValueVectorReaderImpl(buffer); + for (int i = 0; i < 100; i++) { + Assert.assertTrue(reader.isNull(i)); + } } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/forward/FixedBitSVForwardIndexReaderV2Test.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/forward/FixedBitSVForwardIndexReaderV2Test.java index 5deb60d90a95..37b71514af8b 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/forward/FixedBitSVForwardIndexReaderV2Test.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/forward/FixedBitSVForwardIndexReaderV2Test.java @@ -23,6 +23,7 @@ import java.util.ArrayList; import java.util.Random; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.io.util.PinotDataBitSetV2; import org.apache.pinot.segment.local.io.writer.impl.FixedBitSVForwardIndexWriter; import org.apache.pinot.segment.spi.memory.PinotDataBuffer; @@ -34,7 +35,7 @@ import static org.testng.Assert.assertEquals; -public class FixedBitSVForwardIndexReaderV2Test { +public class FixedBitSVForwardIndexReaderV2Test implements PinotBuffersAfterMethodCheckRule { private static final File INDEX_DIR = new File(FileUtils.getTempDirectory(), "FixedBitIntReaderTest"); private static final int NUM_VALUES = 99_999; private static final int NUM_DOC_IDS = PinotDataBitSetV2.MAX_DOC_PER_CALL; diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/text/LuceneTextIndexCompatibleTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/text/LuceneTextIndexCompatibleTest.java index 5ac2c6a17a55..8dd642b749f9 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/text/LuceneTextIndexCompatibleTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/text/LuceneTextIndexCompatibleTest.java @@ -20,6 +20,7 @@ import com.google.common.collect.ImmutableMap; import java.io.File; +import java.io.IOException; import org.testng.Assert; import org.testng.annotations.Test; @@ -27,10 +28,12 @@ public class LuceneTextIndexCompatibleTest { @Test - public void testLucene80IndexReader() { + public void testLucene80IndexReader() + throws IOException { File indexPath = new File(LuceneTextIndexCompatibleTest.class.getClassLoader().getResource("data/lucene_80_index").getPath()); - LuceneTextIndexReader lucene80Index = new LuceneTextIndexReader("Text", indexPath, 1000, ImmutableMap.of()); - Assert.assertNotNull(lucene80Index); + try (LuceneTextIndexReader lucene80Index = new LuceneTextIndexReader("Text", indexPath, 1000, ImmutableMap.of())) { + Assert.assertNotNull(lucene80Index); + } } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readerwriter/FixedByteSingleValueMultiColumnReaderWriterTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readerwriter/FixedByteSingleValueMultiColumnReaderWriterTest.java index 8e8cfb55ed30..37f4efe798de 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readerwriter/FixedByteSingleValueMultiColumnReaderWriterTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readerwriter/FixedByteSingleValueMultiColumnReaderWriterTest.java @@ -21,6 +21,7 @@ import java.io.IOException; import java.util.Random; import org.apache.commons.lang3.RandomStringUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterClassCheckRule; import org.apache.pinot.segment.local.io.readerwriter.impl.FixedByteSingleValueMultiColumnReaderWriter; import org.apache.pinot.segment.local.io.writer.impl.DirectMemoryManager; import org.apache.pinot.segment.spi.memory.PinotDataBufferMemoryManager; @@ -33,7 +34,7 @@ /** * Unit test for {@link FixedByteSingleValueMultiColumnReaderWriter} */ -public class FixedByteSingleValueMultiColumnReaderWriterTest { +public class FixedByteSingleValueMultiColumnReaderWriterTest implements PinotBuffersAfterClassCheckRule { private static final int NUM_ROWS = 1001; private static final int NUM_ROWS_PER_CHUNK = 23; diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readerwriter/FixedByteValueReaderWriterTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readerwriter/FixedByteValueReaderWriterTest.java index 5e6a5c673dbc..65b09df5644f 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readerwriter/FixedByteValueReaderWriterTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readerwriter/FixedByteValueReaderWriterTest.java @@ -25,6 +25,7 @@ import java.util.Arrays; import java.util.List; import java.util.concurrent.ThreadLocalRandom; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.io.util.FixedByteValueReaderWriter; import org.apache.pinot.segment.spi.memory.PinotDataBuffer; import org.testng.annotations.DataProvider; @@ -33,7 +34,7 @@ import static org.testng.Assert.assertEquals; -public class FixedByteValueReaderWriterTest { +public class FixedByteValueReaderWriterTest implements PinotBuffersAfterMethodCheckRule { @DataProvider public static Object[][] params() { diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readerwriter/ValueReaderComparisonTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readerwriter/ValueReaderComparisonTest.java index f788bdb37821..69855469222f 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readerwriter/ValueReaderComparisonTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readerwriter/ValueReaderComparisonTest.java @@ -32,6 +32,7 @@ import java.util.stream.IntStream; import java.util.stream.Stream; import org.apache.commons.lang3.tuple.Pair; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.io.util.FixedByteValueReaderWriter; import org.apache.pinot.segment.local.io.util.ValueReader; import org.apache.pinot.segment.local.io.util.VarLengthValueReader; @@ -45,7 +46,7 @@ import static org.testng.Assert.assertTrue; -public class ValueReaderComparisonTest { +public class ValueReaderComparisonTest implements PinotBuffersAfterMethodCheckRule { // Number of rounds to run the test for, change this number to test locally for catching the corner cases. private static final int NUM_ROUNDS = 1; @@ -368,11 +369,11 @@ private static Pair prepare(boolean fixed, int num } else { assert byteOrder == ByteOrder.BIG_ENDIAN : "little endian unsupported by VarLengthValueWriter"; Path file = Files.createTempFile(ValueReaderComparisonTest.class.getName() + "-" + UUID.randomUUID(), ".tmp"); - VarLengthValueWriter writer = new VarLengthValueWriter(file.toFile(), storedValues.length); - for (byte[] storedValue : storedValues) { - writer.add(storedValue); + try (VarLengthValueWriter writer = new VarLengthValueWriter(file.toFile(), storedValues.length)) { + for (byte[] storedValue : storedValues) { + writer.add(storedValue); + } } - writer.close(); PinotDataBuffer buffer = PinotDataBuffer.mapFile(file.toFile(), true, 0, Files.size(file), byteOrder, "ValueReaderComparisonTest"); return Pair.of(new VarLengthValueReader(buffer), buffer); diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/writer/FixedByteWidthRowColForwardIndexWriterTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/writer/FixedByteWidthRowColForwardIndexWriterTest.java index 1db5bad250c0..49e0792c4582 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/writer/FixedByteWidthRowColForwardIndexWriterTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/writer/FixedByteWidthRowColForwardIndexWriterTest.java @@ -23,6 +23,7 @@ import java.io.FileInputStream; import java.util.Random; import org.apache.commons.lang3.StringUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.io.reader.impl.FixedByteSingleValueMultiColReader; import org.apache.pinot.segment.local.io.writer.impl.FixedByteSingleValueMultiColWriter; import org.apache.pinot.segment.spi.V1Constants; @@ -32,7 +33,7 @@ @Test -public class FixedByteWidthRowColForwardIndexWriterTest { +public class FixedByteWidthRowColForwardIndexWriterTest implements PinotBuffersAfterMethodCheckRule { @Test public void testSingleColInt() throws Exception { @@ -42,18 +43,20 @@ public void testSingleColInt() int rows = 100; int cols = 1; int[] columnSizes = new int[]{4}; - FixedByteSingleValueMultiColWriter writer = new FixedByteSingleValueMultiColWriter(file, rows, cols, columnSizes); int[] data = new int[rows]; - Random r = new Random(); - for (int i = 0; i < rows; i++) { - data[i] = r.nextInt(); - writer.setInt(i, 0, data[i]); + + try (FixedByteSingleValueMultiColWriter writer = new FixedByteSingleValueMultiColWriter(file, rows, cols, + columnSizes)) { + Random r = new Random(); + for (int i = 0; i < rows; i++) { + data[i] = r.nextInt(); + writer.setInt(i, 0, data[i]); + } } - writer.close(); File rfile = new File("test_single_col_writer.dat"); - try (FixedByteSingleValueMultiColReader reader = new FixedByteSingleValueMultiColReader( - PinotDataBuffer.mapReadOnlyBigEndianFile(rfile), rows, columnSizes)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile(rfile); + FixedByteSingleValueMultiColReader reader = new FixedByteSingleValueMultiColReader(buffer, rows, columnSizes)) { for (int i = 0; i < rows; i++) { Assert.assertEquals(reader.getInt(i, 0), data[i]); } @@ -80,8 +83,8 @@ public void testSingleColFloat() writer.close(); File rfile = new File("test_single_col_writer.dat"); - try (FixedByteSingleValueMultiColReader reader = new FixedByteSingleValueMultiColReader( - PinotDataBuffer.mapReadOnlyBigEndianFile(rfile), rows, columnSizes)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile(rfile); + FixedByteSingleValueMultiColReader reader = new FixedByteSingleValueMultiColReader(buffer, rows, columnSizes)) { for (int i = 0; i < rows; i++) { Assert.assertEquals(reader.getFloat(i, 0), data[i]); } @@ -98,18 +101,20 @@ public void testSingleColDouble() final int rows = 100; final int cols = 1; final int[] columnSizes = new int[]{8}; - FixedByteSingleValueMultiColWriter writer = new FixedByteSingleValueMultiColWriter(wfile, rows, cols, columnSizes); final double[] data = new double[rows]; - Random r = new Random(); - for (int i = 0; i < rows; i++) { - data[i] = r.nextDouble(); - writer.setDouble(i, 0, data[i]); + + try (FixedByteSingleValueMultiColWriter writer = new FixedByteSingleValueMultiColWriter(wfile, rows, cols, + columnSizes)) { + Random r = new Random(); + for (int i = 0; i < rows; i++) { + data[i] = r.nextDouble(); + writer.setDouble(i, 0, data[i]); + } } - writer.close(); File rfile = new File("test_single_col_writer.dat"); - try (FixedByteSingleValueMultiColReader reader = new FixedByteSingleValueMultiColReader( - PinotDataBuffer.mapReadOnlyBigEndianFile(rfile), rows, columnSizes)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile(rfile); + FixedByteSingleValueMultiColReader reader = new FixedByteSingleValueMultiColReader(buffer, rows, columnSizes)) { for (int i = 0; i < rows; i++) { Assert.assertEquals(reader.getDouble(i, 0), data[i]); } @@ -136,8 +141,8 @@ public void testSingleColLong() writer.close(); File rfile = new File("test_single_col_writer.dat"); - try (FixedByteSingleValueMultiColReader reader = new FixedByteSingleValueMultiColReader( - PinotDataBuffer.mapReadOnlyBigEndianFile(rfile), rows, columnSizes)) { + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile(rfile); + FixedByteSingleValueMultiColReader reader = new FixedByteSingleValueMultiColReader(buffer, rows, columnSizes)) { for (int i = 0; i < rows; i++) { Assert.assertEquals(reader.getLong(i, 0), data[i]); } @@ -189,23 +194,27 @@ public void testSpecialCharsForStringReaderWriter() String testString2 = new String(bytes2); int stringColumnMaxLength = Math.max(testString1.getBytes().length, testString2.getBytes().length); int[] columnSizes = new int[]{stringColumnMaxLength}; - FixedByteSingleValueMultiColWriter writer = new FixedByteSingleValueMultiColWriter(file, rows, cols, columnSizes); String[] data = new String[rows]; - for (int i = 0; i < rows; i++) { - String toPut = (i % 2 == 0) ? testString1 : testString2; - final int padding = stringColumnMaxLength - toPut.getBytes().length; - final StringBuilder bld = new StringBuilder(); - bld.append(toPut); - for (int j = 0; j < padding; j++) { - bld.append(V1Constants.Str.DEFAULT_STRING_PAD_CHAR); + try (FixedByteSingleValueMultiColWriter writer = new FixedByteSingleValueMultiColWriter(file, rows, cols, + columnSizes)) { + for (int i = 0; i < rows; i++) { + String toPut = (i % 2 == 0) ? testString1 : testString2; + final int padding = stringColumnMaxLength - toPut.getBytes().length; + + final StringBuilder bld = new StringBuilder(); + bld.append(toPut); + for (int j = 0; j < padding; j++) { + bld.append(V1Constants.Str.DEFAULT_STRING_PAD_CHAR); + } + data[i] = bld.toString(); + writer.setString(i, 0, data[i]); } - data[i] = bld.toString(); - writer.setString(i, 0, data[i]); } - writer.close(); - try (FixedByteSingleValueMultiColReader dataFileReader = new FixedByteSingleValueMultiColReader( - PinotDataBuffer.mapReadOnlyBigEndianFile(file), rows, new int[]{stringColumnMaxLength})) { + + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile(file); + FixedByteSingleValueMultiColReader dataFileReader = new FixedByteSingleValueMultiColReader(buffer, rows, + new int[]{stringColumnMaxLength})) { for (int i = 0; i < rows; i++) { String stringInFile = dataFileReader.getString(i, 0); Assert.assertEquals(stringInFile, data[i]); @@ -234,23 +243,27 @@ public void testSpecialPaddingCharsForStringReaderWriter() String testString2 = new String(bytes2); int stringColumnMaxLength = Math.max(testString1.getBytes().length, testString2.getBytes().length); int[] columnSizes = new int[]{stringColumnMaxLength}; - FixedByteSingleValueMultiColWriter writer = new FixedByteSingleValueMultiColWriter(file, rows, cols, columnSizes); String[] data = new String[rows]; - for (int i = 0; i < rows; i++) { - String toPut = (i % 2 == 0) ? testString1 : testString2; - final int padding = stringColumnMaxLength - toPut.getBytes().length; - final StringBuilder bld = new StringBuilder(); - bld.append(toPut); - for (int j = 0; j < padding; j++) { - bld.append(paddingChar); + try (FixedByteSingleValueMultiColWriter writer = new FixedByteSingleValueMultiColWriter(file, rows, cols, + columnSizes)) { + for (int i = 0; i < rows; i++) { + String toPut = (i % 2 == 0) ? testString1 : testString2; + final int padding = stringColumnMaxLength - toPut.getBytes().length; + + final StringBuilder bld = new StringBuilder(); + bld.append(toPut); + for (int j = 0; j < padding; j++) { + bld.append(paddingChar); + } + data[i] = bld.toString(); + writer.setString(i, 0, data[i]); } - data[i] = bld.toString(); - writer.setString(i, 0, data[i]); } - writer.close(); - try (FixedByteSingleValueMultiColReader dataFileReader = new FixedByteSingleValueMultiColReader( - PinotDataBuffer.mapReadOnlyBigEndianFile(file), rows, new int[]{stringColumnMaxLength})) { + + try (PinotDataBuffer buffer = PinotDataBuffer.mapReadOnlyBigEndianFile(file); + FixedByteSingleValueMultiColReader dataFileReader = new FixedByteSingleValueMultiColReader(buffer, rows, + new int[]{stringColumnMaxLength})) { for (int i = 0; i < rows; i++) { String stringInFile = dataFileReader.getString(i, 0); Assert.assertEquals(stringInFile, data[i]); diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/store/ColumnIndexDirectoryTestHelper.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/store/ColumnIndexDirectoryTestHelper.java index b6718ee460b3..412cadeda10b 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/store/ColumnIndexDirectoryTestHelper.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/store/ColumnIndexDirectoryTestHelper.java @@ -19,6 +19,7 @@ package org.apache.pinot.segment.local.segment.store; import java.io.IOException; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.spi.ColumnMetadata; import org.apache.pinot.segment.spi.creator.SegmentVersion; import org.apache.pinot.segment.spi.index.IndexType; @@ -33,7 +34,7 @@ import static org.mockito.Mockito.when; -public class ColumnIndexDirectoryTestHelper { +public class ColumnIndexDirectoryTestHelper implements PinotBuffersAfterMethodCheckRule { private ColumnIndexDirectoryTestHelper() { } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/store/FilePerIndexDirectoryTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/store/FilePerIndexDirectoryTest.java index 8085c9526242..9e3dcdeab3c3 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/store/FilePerIndexDirectoryTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/store/FilePerIndexDirectoryTest.java @@ -26,6 +26,7 @@ import java.util.HashSet; import java.util.TreeSet; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.segment.creator.impl.text.LuceneTextIndexCreator; import org.apache.pinot.segment.local.segment.creator.impl.text.NativeTextIndexCreator; import org.apache.pinot.segment.local.segment.index.readers.text.LuceneTextIndexReader; @@ -49,7 +50,7 @@ import static org.testng.Assert.assertTrue; -public class FilePerIndexDirectoryTest { +public class FilePerIndexDirectoryTest implements PinotBuffersAfterMethodCheckRule { private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(), FilePerIndexDirectoryTest.class.toString()); @@ -161,8 +162,9 @@ public void testHasIndex() @Test public void testRemoveIndex() throws IOException { - try (FilePerIndexDirectory fpi = new FilePerIndexDirectory(TEMP_DIR, _segmentMetadata, ReadMode.mmap)) { - fpi.newBuffer("col1", StandardIndexes.forward(), 1024); + try (FilePerIndexDirectory fpi = new FilePerIndexDirectory(TEMP_DIR, _segmentMetadata, ReadMode.mmap); + //buff needs closing because removeIndex() doesn't do it. + PinotDataBuffer buff = fpi.newBuffer("col1", StandardIndexes.forward(), 1024)) { fpi.newBuffer("col2", StandardIndexes.dictionary(), 100); assertTrue(fpi.getFileFor("col1", StandardIndexes.forward()).exists()); assertTrue(fpi.getFileFor("col2", StandardIndexes.dictionary()).exists()); @@ -226,10 +228,12 @@ public void testRemoveTextIndices() try (FilePerIndexDirectory fpi = new FilePerIndexDirectory(TEMP_DIR, _segmentMetadata, ReadMode.mmap)) { assertTrue(fpi.hasIndexFor("foo", StandardIndexes.text())); // Use TextIndex once to trigger the creation of mapping files. - LuceneTextIndexReader fooReader = new LuceneTextIndexReader("foo", TEMP_DIR, 1, new HashMap<>()); - fooReader.getDocIds("clean"); - LuceneTextIndexReader barReader = new LuceneTextIndexReader("bar", TEMP_DIR, 3, new HashMap<>()); - barReader.getDocIds("retain hold"); + try (LuceneTextIndexReader fooReader = new LuceneTextIndexReader("foo", TEMP_DIR, 1, new HashMap<>())) { + fooReader.getDocIds("clean"); + } + try (LuceneTextIndexReader barReader = new LuceneTextIndexReader("bar", TEMP_DIR, 3, new HashMap<>())) { + barReader.getDocIds("retain hold"); + } // Both files for TextIndex should be removed. fpi.removeIndex("foo", StandardIndexes.text()); @@ -237,6 +241,7 @@ public void testRemoveTextIndices() assertFalse( new File(TEMP_DIR, "foo" + V1Constants.Indexes.LUCENE_TEXT_INDEX_DOCID_MAPPING_FILE_EXTENSION).exists()); } + assertTrue(new File(TEMP_DIR, "bar" + V1Constants.Indexes.LUCENE_V912_TEXT_INDEX_FILE_EXTENSION).exists()); assertTrue(new File(TEMP_DIR, "bar" + V1Constants.Indexes.LUCENE_TEXT_INDEX_DOCID_MAPPING_FILE_EXTENSION).exists()); @@ -255,10 +260,11 @@ public void testRemoveTextIndices() assertTrue(fpi.hasIndexFor("bar", StandardIndexes.text())); // Check if the text index still work. - LuceneTextIndexReader barReader = new LuceneTextIndexReader("bar", TEMP_DIR, 3, new HashMap<>()); - MutableRoaringBitmap ids = barReader.getDocIds("retain hold"); - assertTrue(ids.contains(0)); - assertTrue(ids.contains(2)); + try (LuceneTextIndexReader barReader = new LuceneTextIndexReader("bar", TEMP_DIR, 3, new HashMap<>())) { + MutableRoaringBitmap ids = barReader.getDocIds("retain hold"); + assertTrue(ids.contains(0)); + assertTrue(ids.contains(2)); + } } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/store/SegmentLocalFSDirectoryTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/store/SegmentLocalFSDirectoryTest.java index f2aed2fdef5d..87e9920b5029 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/store/SegmentLocalFSDirectoryTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/store/SegmentLocalFSDirectoryTest.java @@ -20,6 +20,7 @@ import java.io.File; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterClassCheckRule; import org.apache.pinot.segment.spi.creator.SegmentVersion; import org.apache.pinot.segment.spi.index.StandardIndexes; import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl; @@ -33,7 +34,7 @@ import org.testng.annotations.Test; -public class SegmentLocalFSDirectoryTest { +public class SegmentLocalFSDirectoryTest implements PinotBuffersAfterClassCheckRule { private static final File TEST_DIRECTORY = new File(SingleFileIndexDirectoryTest.class.toString()); private SegmentDirectory _segmentDirectory; private SegmentMetadataImpl _metadata; diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/store/SingleFileIndexDirectoryTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/store/SingleFileIndexDirectoryTest.java index edc8884dc33b..eddf5f99c27f 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/store/SingleFileIndexDirectoryTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/store/SingleFileIndexDirectoryTest.java @@ -34,6 +34,7 @@ import java.util.stream.Collectors; import org.apache.commons.configuration2.ex.ConfigurationException; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.segment.creator.impl.text.LuceneTextIndexCreator; import org.apache.pinot.segment.local.segment.index.readers.text.LuceneTextIndexReader; import org.apache.pinot.segment.spi.V1Constants; @@ -57,7 +58,7 @@ import static org.testng.Assert.assertTrue; -public class SingleFileIndexDirectoryTest { +public class SingleFileIndexDirectoryTest implements PinotBuffersAfterMethodCheckRule { private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(), SingleFileIndexDirectoryTest.class.toString()); @@ -92,19 +93,22 @@ public void testWithEmptyDir() throws Exception { // segmentDir does not have anything to begin with assertEquals(TEMP_DIR.list().length, 0); - SingleFileIndexDirectory columnDirectory = new SingleFileIndexDirectory(TEMP_DIR, _segmentMetadata, ReadMode.mmap); - PinotDataBuffer writtenBuffer = columnDirectory.newBuffer("foo", StandardIndexes.dictionary(), 1024); - String data = "This is a test string"; - final byte[] dataBytes = data.getBytes(); - int pos = 0; - for (byte b : dataBytes) { - writtenBuffer.putByte(pos++, b); + final byte[] dataBytes; + try (SingleFileIndexDirectory columnDirectory = new SingleFileIndexDirectory(TEMP_DIR, _segmentMetadata, + ReadMode.mmap); + PinotDataBuffer writtenBuffer = columnDirectory.newBuffer("foo", StandardIndexes.dictionary(), 1024)) { + String data = "This is a test string"; + dataBytes = data.getBytes(); + int pos = 0; + for (byte b : dataBytes) { + writtenBuffer.putByte(pos++, b); + } } - writtenBuffer.close(); Mockito.when(_segmentMetadata.getAllColumns()).thenReturn(new TreeSet<>(Arrays.asList("foo"))); try (SingleFileIndexDirectory directoryReader = new SingleFileIndexDirectory(TEMP_DIR, _segmentMetadata, - ReadMode.mmap); PinotDataBuffer readBuffer = directoryReader.getBuffer("foo", StandardIndexes.dictionary())) { + ReadMode.mmap); + PinotDataBuffer readBuffer = directoryReader.getBuffer("foo", StandardIndexes.dictionary())) { assertEquals(1024, readBuffer.size()); int length = dataBytes.length; for (int i = 0; i < length; i++) { @@ -259,10 +263,12 @@ public void testRemoveTextIndices() try (SingleFileIndexDirectory sfd = new SingleFileIndexDirectory(TEMP_DIR, _segmentMetadata, ReadMode.mmap)) { assertTrue(sfd.hasIndexFor("foo", StandardIndexes.text())); // Use TextIndex once to trigger the creation of mapping files. - LuceneTextIndexReader fooReader = new LuceneTextIndexReader("foo", TEMP_DIR, 1, new HashMap<>()); - fooReader.getDocIds("clean"); - LuceneTextIndexReader barReader = new LuceneTextIndexReader("bar", TEMP_DIR, 3, new HashMap<>()); - barReader.getDocIds("retain hold"); + try (LuceneTextIndexReader fooReader = new LuceneTextIndexReader("foo", TEMP_DIR, 1, new HashMap<>())) { + fooReader.getDocIds("clean"); + } + try (LuceneTextIndexReader barReader = new LuceneTextIndexReader("bar", TEMP_DIR, 3, new HashMap<>())) { + barReader.getDocIds("retain hold"); + } // Both files for TextIndex should be removed. sfd.removeIndex("foo", StandardIndexes.text()); @@ -289,10 +295,11 @@ public void testRemoveTextIndices() assertTrue(sfd.hasIndexFor("bar", StandardIndexes.text())); // Check if the text index still work. - LuceneTextIndexReader barReader = new LuceneTextIndexReader("bar", TEMP_DIR, 3, new HashMap<>()); - MutableRoaringBitmap ids = barReader.getDocIds("retain hold"); - assertTrue(ids.contains(0)); - assertTrue(ids.contains(2)); + try (LuceneTextIndexReader barReader = new LuceneTextIndexReader("bar", TEMP_DIR, 3, new HashMap<>())) { + MutableRoaringBitmap ids = barReader.getDocIds("retain hold"); + assertTrue(ids.contains(0)); + assertTrue(ids.contains(2)); + } } } @@ -383,8 +390,7 @@ public void testGetColumnIndices() new HashSet<>(Collections.singletonList("col2"))); assertEquals(sfd.getColumnsWithIndex(StandardIndexes.inverted()), new HashSet<>(Collections.singletonList("col4"))); - assertEquals(sfd.getColumnsWithIndex(StandardIndexes.h3()), - new HashSet<>(Collections.singletonList("col5"))); + assertEquals(sfd.getColumnsWithIndex(StandardIndexes.h3()), new HashSet<>(Collections.singletonList("col5"))); assertEquals(sfd.getColumnsWithIndex(StandardIndexes.text()), new HashSet<>(Arrays.asList("foo", "bar"))); sfd.removeIndex("col1", StandardIndexes.forward()); diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/store/StarTreeIndexReaderTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/store/StarTreeIndexReaderTest.java index 02278354e27c..4c9a5dbf5a3d 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/store/StarTreeIndexReaderTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/store/StarTreeIndexReaderTest.java @@ -28,6 +28,7 @@ import org.apache.commons.configuration2.ex.ConfigurationException; import org.apache.commons.io.FileUtils; import org.apache.commons.lang3.tuple.Pair; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.startree.v2.store.StarTreeIndexMapUtils; import org.apache.pinot.segment.spi.AggregationFunctionType; import org.apache.pinot.segment.spi.creator.SegmentVersion; @@ -49,7 +50,7 @@ import static org.testng.Assert.assertTrue; -public class StarTreeIndexReaderTest { +public class StarTreeIndexReaderTest implements PinotBuffersAfterMethodCheckRule { private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(), StarTreeIndexReaderTest.class.toString()); private SegmentMetadataImpl _segmentMetadata; diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/fst/FSTBuilderTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/fst/FSTBuilderTest.java index edee3ebef21e..0b4b99588cf2 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/fst/FSTBuilderTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/fst/FSTBuilderTest.java @@ -30,6 +30,7 @@ import org.apache.lucene.util.fst.FST; import org.apache.lucene.util.fst.Outputs; import org.apache.lucene.util.fst.PositiveIntOutputs; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.spi.memory.PinotDataBuffer; import org.testng.Assert; import org.testng.annotations.AfterClass; @@ -37,7 +38,7 @@ import org.testng.annotations.Test; -public class FSTBuilderTest { +public class FSTBuilderTest implements PinotBuffersAfterMethodCheckRule { private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(), "FST"); @BeforeClass @@ -74,17 +75,18 @@ public void testFSTBuilder() Outputs outputs = PositiveIntOutputs.getSingleton(); File fstFile = new File(outputFile.getAbsolutePath()); - PinotDataBuffer pinotDataBuffer = - PinotDataBuffer.mapFile(fstFile, true, 0, fstFile.length(), ByteOrder.BIG_ENDIAN, ""); - PinotBufferIndexInput indexInput = new PinotBufferIndexInput(pinotDataBuffer, 0L, fstFile.length()); + try (PinotDataBuffer pinotDataBuffer = + PinotDataBuffer.mapFile(fstFile, true, 0, fstFile.length(), ByteOrder.BIG_ENDIAN, "")) { + PinotBufferIndexInput indexInput = new PinotBufferIndexInput(pinotDataBuffer, 0L, fstFile.length()); - List results = RegexpMatcher.regexMatch("hello.*123", fst); - Assert.assertEquals(results.size(), 1); - Assert.assertEquals(results.get(0).longValue(), 21L); + List results = RegexpMatcher.regexMatch("hello.*123", fst); + Assert.assertEquals(results.size(), 1); + Assert.assertEquals(results.get(0).longValue(), 21L); - results = RegexpMatcher.regexMatch(".*world", fst); - Assert.assertEquals(results.size(), 1); - Assert.assertEquals(results.get(0).longValue(), 12L); + results = RegexpMatcher.regexMatch(".*world", fst); + Assert.assertEquals(results.size(), 1); + Assert.assertEquals(results.get(0).longValue(), 12L); + } } @AfterClass diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/nativefst/FSTRegexpWithWeirdTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/nativefst/FSTRegexpWithWeirdTest.java index c27c8e1050df..89f091627f14 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/nativefst/FSTRegexpWithWeirdTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/nativefst/FSTRegexpWithWeirdTest.java @@ -23,11 +23,13 @@ import java.io.IOException; import java.nio.charset.Charset; import java.util.Arrays; +import org.apache.pinot.segment.local.PinotBuffersAfterClassCheckRule; import org.apache.pinot.segment.local.utils.nativefst.builder.FSTBuilder; import org.apache.pinot.segment.local.utils.nativefst.builder.FSTSerializerImpl; import org.apache.pinot.segment.local.utils.nativefst.utils.RegexpMatcher; import org.roaringbitmap.RoaringBitmapWriter; import org.roaringbitmap.buffer.MutableRoaringBitmap; +import org.testng.annotations.AfterClass; import org.testng.annotations.BeforeClass; import org.testng.annotations.Test; @@ -37,7 +39,7 @@ /** * Tests with weird input */ -public class FSTRegexpWithWeirdTest { +public class FSTRegexpWithWeirdTest implements PinotBuffersAfterClassCheckRule { private FST _fst; private static byte[][] convertToBytes(String[] strings) { @@ -69,6 +71,14 @@ public void setUp() _fst = FST.read(new ByteArrayInputStream(fstData), ImmutableFST.class, true); } + @AfterClass + public void tearDown() + throws IOException { + if (_fst instanceof ImmutableFST) { + ((ImmutableFST) _fst)._mutableBytesStore.close(); + } + } + @Test public void testRegex1() { assertEquals(regexQueryNrHits(".*196169"), 1); diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/nativefst/FSTSanityTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/nativefst/FSTSanityTest.java index 96e74043200b..3ab381e532b6 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/nativefst/FSTSanityTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/nativefst/FSTSanityTest.java @@ -27,9 +27,11 @@ import java.util.Objects; import java.util.SortedMap; import java.util.TreeMap; +import org.apache.pinot.segment.local.PinotBuffersAfterClassCheckRule; import org.apache.pinot.segment.local.utils.fst.RegexpMatcher; import org.apache.pinot.segment.local.utils.nativefst.builder.FSTBuilder; import org.apache.pinot.segment.local.utils.nativefst.builder.FSTSerializerImpl; +import org.testng.annotations.AfterClass; import org.testng.annotations.BeforeClass; import org.testng.annotations.Test; @@ -40,7 +42,7 @@ * Load a 0.5 million unique words data set and do the same set of queries on Lucene FST and * native FST and compare results */ -public class FSTSanityTest { +public class FSTSanityTest implements PinotBuffersAfterClassCheckRule { private FST _nativeFST; private org.apache.lucene.util.fst.FST _fst; @@ -64,6 +66,14 @@ public void setUp() _fst = org.apache.pinot.segment.local.utils.fst.FSTBuilder.buildFST(input); } + @AfterClass + public void tearDown() + throws IOException { + if (_nativeFST instanceof ImmutableFST) { + ((ImmutableFST) _nativeFST)._mutableBytesStore.close(); + } + } + @Test public void testRegex() throws IOException { diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/nativefst/FSTTraversalTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/nativefst/FSTTraversalTest.java index f22c7ab0aaea..6715b49c5443 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/nativefst/FSTTraversalTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/nativefst/FSTTraversalTest.java @@ -29,12 +29,14 @@ import java.util.Arrays; import java.util.SortedMap; import java.util.TreeMap; +import org.apache.pinot.segment.local.PinotBuffersAfterClassCheckRule; import org.apache.pinot.segment.local.io.writer.impl.DirectMemoryManager; import org.apache.pinot.segment.local.utils.nativefst.builder.FSTBuilder; import org.apache.pinot.segment.local.utils.nativefst.builder.FSTSerializerImpl; import org.apache.pinot.segment.local.utils.nativefst.utils.RegexpMatcher; import org.roaringbitmap.RoaringBitmapWriter; import org.roaringbitmap.buffer.MutableRoaringBitmap; +import org.testng.annotations.AfterClass; import org.testng.annotations.BeforeClass; import org.testng.annotations.Test; @@ -54,7 +56,7 @@ * * This class also holds tests for {@link RegexpMatcher} since they both perform FST traversals */ -public class FSTTraversalTest { +public class FSTTraversalTest implements PinotBuffersAfterClassCheckRule { private FST _fst; private FST _regexFST; @@ -79,6 +81,19 @@ public void setUp() _regexFST = fstBuilder.complete(); } + @AfterClass + public void tearDown() + throws IOException { + close(_fst); + } + + private void close(FST fst) + throws IOException { + if (fst instanceof ImmutableFST) { + ((ImmutableFST) fst)._mutableBytesStore.close(); + } + } + @Test public void testAutomatonHasPrefixBug() throws Exception { @@ -87,29 +102,33 @@ public void testAutomatonHasPrefixBug() "bcd".getBytes(UTF_8), "bce".getBytes(UTF_8)), new int[]{10, 11, 12, 13, 14, 15}); byte[] fstData = new FSTSerializerImpl().withNumbers().serialize(fst, new ByteArrayOutputStream()).toByteArray(); fst = FST.read(new ByteArrayInputStream(fstData), ImmutableFST.class, true); + try { - FSTTraversal fstTraversal = new FSTTraversal(fst); - assertEquals(fstTraversal.match("a".getBytes(UTF_8))._kind, EXACT_MATCH); - assertEquals(fstTraversal.match("ab".getBytes(UTF_8))._kind, EXACT_MATCH); - assertEquals(fstTraversal.match("abc".getBytes(UTF_8))._kind, EXACT_MATCH); - assertEquals(fstTraversal.match("ad".getBytes(UTF_8))._kind, EXACT_MATCH); + FSTTraversal fstTraversal = new FSTTraversal(fst); + assertEquals(fstTraversal.match("a".getBytes(UTF_8))._kind, EXACT_MATCH); + assertEquals(fstTraversal.match("ab".getBytes(UTF_8))._kind, EXACT_MATCH); + assertEquals(fstTraversal.match("abc".getBytes(UTF_8))._kind, EXACT_MATCH); + assertEquals(fstTraversal.match("ad".getBytes(UTF_8))._kind, EXACT_MATCH); - assertEquals(fstTraversal.match("b".getBytes(UTF_8))._kind, SEQUENCE_IS_A_PREFIX); - assertEquals(fstTraversal.match("bc".getBytes(UTF_8))._kind, SEQUENCE_IS_A_PREFIX); + assertEquals(fstTraversal.match("b".getBytes(UTF_8))._kind, SEQUENCE_IS_A_PREFIX); + assertEquals(fstTraversal.match("bc".getBytes(UTF_8))._kind, SEQUENCE_IS_A_PREFIX); - MatchResult m = fstTraversal.match("abcd".getBytes(UTF_8)); - assertEquals(m._kind, AUTOMATON_HAS_PREFIX); - assertEquals(m._index, 3); + MatchResult m = fstTraversal.match("abcd".getBytes(UTF_8)); + assertEquals(m._kind, AUTOMATON_HAS_PREFIX); + assertEquals(m._index, 3); - m = fstTraversal.match("ade".getBytes(UTF_8)); - assertEquals(m._kind, AUTOMATON_HAS_PREFIX); - assertEquals(m._index, 2); + m = fstTraversal.match("ade".getBytes(UTF_8)); + assertEquals(m._kind, AUTOMATON_HAS_PREFIX); + assertEquals(m._index, 2); - m = fstTraversal.match("ax".getBytes(UTF_8)); - assertEquals(m._kind, AUTOMATON_HAS_PREFIX); - assertEquals(m._index, 1); + m = fstTraversal.match("ax".getBytes(UTF_8)); + assertEquals(m._kind, AUTOMATON_HAS_PREFIX); + assertEquals(m._index, 1); - assertEquals(fstTraversal.match("d".getBytes(UTF_8))._kind, NO_MATCH); + assertEquals(fstTraversal.match("d".getBytes(UTF_8))._kind, NO_MATCH); + } finally { + close(fst); + } } @Test @@ -154,22 +173,26 @@ public void testMatch() throws IOException { File file = new File("./src/test/resources/data/abc.native.fst"); FST fst = FST.read(new FileInputStream(file), false, new DirectMemoryManager(FSTTraversalTest.class.getName())); - FSTTraversal traversalHelper = new FSTTraversal(fst); + try { + FSTTraversal traversalHelper = new FSTTraversal(fst); - MatchResult m = traversalHelper.match("ax".getBytes()); - assertEquals(m._kind, AUTOMATON_HAS_PREFIX); - assertEquals(m._index, 1); - assertEquals(suffixes(fst, m._node), Sets.newHashSet("ba", "c")); + MatchResult m = traversalHelper.match("ax".getBytes()); + assertEquals(m._kind, AUTOMATON_HAS_PREFIX); + assertEquals(m._index, 1); + assertEquals(suffixes(fst, m._node), Sets.newHashSet("ba", "c")); - assertEquals(traversalHelper.match("aba".getBytes())._kind, EXACT_MATCH); + assertEquals(traversalHelper.match("aba".getBytes())._kind, EXACT_MATCH); - m = traversalHelper.match("abalonger".getBytes()); - assertEquals(m._kind, AUTOMATON_HAS_PREFIX); - assertEquals("abalonger".substring(m._index), "longer"); + m = traversalHelper.match("abalonger".getBytes()); + assertEquals(m._kind, AUTOMATON_HAS_PREFIX); + assertEquals("abalonger".substring(m._index), "longer"); - m = traversalHelper.match("ab".getBytes()); - assertEquals(m._kind, SEQUENCE_IS_A_PREFIX); - assertEquals(suffixes(fst, m._node), Sets.newHashSet("a")); + m = traversalHelper.match("ab".getBytes()); + assertEquals(m._kind, SEQUENCE_IS_A_PREFIX); + assertEquals(suffixes(fst, m._node), Sets.newHashSet("a")); + } finally { + close(fst); + } } @Test @@ -185,11 +208,15 @@ public void testRegexMatcherPrefix() FST fst = builder.complete(); byte[] fstData = new FSTSerializerImpl().withNumbers().serialize(fst, new ByteArrayOutputStream()).toByteArray(); fst = FST.read(new ByteArrayInputStream(fstData), ImmutableFST.class, true); + try { - RoaringBitmapWriter writer = RoaringBitmapWriter.bufferWriter().get(); - RegexpMatcher.regexMatch("h.*", fst, writer::add); + RoaringBitmapWriter writer = RoaringBitmapWriter.bufferWriter().get(); + RegexpMatcher.regexMatch("h.*", fst, writer::add); - assertEquals(writer.get().getCardinality(), 2); + assertEquals(writer.get().getCardinality(), 2); + } finally { + close(fst); + } } @Test @@ -205,11 +232,15 @@ public void testRegexMatcherSuffix() FST fst = builder.complete(); byte[] fstData = new FSTSerializerImpl().withNumbers().serialize(fst, new ByteArrayOutputStream()).toByteArray(); fst = FST.read(new ByteArrayInputStream(fstData), ImmutableFST.class, true); + try { - RoaringBitmapWriter writer = RoaringBitmapWriter.bufferWriter().get(); - RegexpMatcher.regexMatch(".*h", fst, writer::add); + RoaringBitmapWriter writer = RoaringBitmapWriter.bufferWriter().get(); + RegexpMatcher.regexMatch(".*h", fst, writer::add); - assertEquals(writer.get().getCardinality(), 2); + assertEquals(writer.get().getCardinality(), 2); + } finally { + close(fst); + } } @Test @@ -223,17 +254,21 @@ public void testRegexMatcherSuffix2() FST fst = FSTBuilder.buildFST(input); byte[] fstData = new FSTSerializerImpl().withNumbers().serialize(fst, new ByteArrayOutputStream()).toByteArray(); fst = FST.read(new ByteArrayInputStream(fstData), ImmutableFST.class, true); + try { - RoaringBitmapWriter writer = RoaringBitmapWriter.bufferWriter().get(); - RegexpMatcher.regexMatch(".*123", fst, writer::add); + RoaringBitmapWriter writer = RoaringBitmapWriter.bufferWriter().get(); + RegexpMatcher.regexMatch(".*123", fst, writer::add); - assertEquals(writer.get().getCardinality(), 1); + assertEquals(writer.get().getCardinality(), 1); - writer.reset(); + writer.reset(); - RegexpMatcher.regexMatch(".till", fst, writer::add); + RegexpMatcher.regexMatch(".till", fst, writer::add); - assertEquals(writer.get().getCardinality(), 1); + assertEquals(writer.get().getCardinality(), 1); + } finally { + close(fst); + } } @Test @@ -247,16 +282,20 @@ public void testRegexMatcherMatchAny() FST fst = FSTBuilder.buildFST(input); byte[] fstData = new FSTSerializerImpl().withNumbers().serialize(fst, new ByteArrayOutputStream()).toByteArray(); fst = FST.read(new ByteArrayInputStream(fstData), ImmutableFST.class, true); + try { - RoaringBitmapWriter writer = RoaringBitmapWriter.bufferWriter().get(); - RegexpMatcher.regexMatch("hello.*123", fst, writer::add); + RoaringBitmapWriter writer = RoaringBitmapWriter.bufferWriter().get(); + RegexpMatcher.regexMatch("hello.*123", fst, writer::add); - assertEquals(writer.get().getCardinality(), 1); + assertEquals(writer.get().getCardinality(), 1); - writer.reset(); - RegexpMatcher.regexMatch("hello.*", fst, writer::add); + writer.reset(); + RegexpMatcher.regexMatch("hello.*", fst, writer::add); - assertEquals(writer.get().getCardinality(), 2); + assertEquals(writer.get().getCardinality(), 2); + } finally { + close(fst); + } } @Test @@ -271,8 +310,11 @@ public void testFSTToString() FST fst = FSTBuilder.buildFST(input); byte[] fstData = new FSTSerializerImpl().withNumbers().serialize(fst, new ByteArrayOutputStream()).toByteArray(); fst = FST.read(new ByteArrayInputStream(fstData), ImmutableFST.class, true); - - ImmutableFST.printToString(fst); + try { + ImmutableFST.printToString(fst); + } finally { + close(fst); + } } @Test @@ -285,11 +327,14 @@ public void testRegexMatcherMatchQuestionMark() FST fst = FSTBuilder.buildFST(input); byte[] fstData = new FSTSerializerImpl().withNumbers().serialize(fst, new ByteArrayOutputStream()).toByteArray(); fst = FST.read(new ByteArrayInputStream(fstData), ImmutableFST.class, true); + try { + RoaringBitmapWriter writer = RoaringBitmapWriter.bufferWriter().get(); + RegexpMatcher.regexMatch("cars?", fst, writer::add); - RoaringBitmapWriter writer = RoaringBitmapWriter.bufferWriter().get(); - RegexpMatcher.regexMatch("cars?", fst, writer::add); - - assertEquals(writer.get().getCardinality(), 2); + assertEquals(writer.get().getCardinality(), 2); + } finally { + close(fst); + } } @Test diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/nativefst/ImmutableFSTDeserializedTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/nativefst/ImmutableFSTDeserializedTest.java index a5df1f860e38..47ab50d1f8e9 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/nativefst/ImmutableFSTDeserializedTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/nativefst/ImmutableFSTDeserializedTest.java @@ -18,8 +18,11 @@ */ package org.apache.pinot.segment.local.utils.nativefst; +import java.io.IOException; import java.io.InputStream; +import org.apache.pinot.segment.local.PinotBuffersAfterClassCheckRule; import org.apache.pinot.segment.local.io.writer.impl.DirectMemoryManager; +import org.testng.annotations.AfterClass; import org.testng.annotations.BeforeClass; import org.testng.annotations.Test; @@ -30,14 +33,24 @@ /** * Deserialize a FST and ensure results are right */ -public class ImmutableFSTDeserializedTest { +public class ImmutableFSTDeserializedTest implements PinotBuffersAfterClassCheckRule { private FST _fst; + private DirectMemoryManager _memManager; @BeforeClass public void setUp() throws Exception { try (InputStream inputStream = getClass().getClassLoader().getResourceAsStream("data/serfst.txt")) { - _fst = FST.read(inputStream, true, new DirectMemoryManager(ImmutableFSTDeserializedTest.class.getName())); + _memManager = new DirectMemoryManager(ImmutableFSTDeserializedTest.class.getName()); + _fst = FST.read(inputStream, true, _memManager); + } + } + + @AfterClass + public void tearDown() + throws IOException { + if (_memManager != null) { + _memManager.close(); } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/nativefst/ImmutableFSTTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/nativefst/ImmutableFSTTest.java index 78d4d161e5f3..92651c51114d 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/nativefst/ImmutableFSTTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/nativefst/ImmutableFSTTest.java @@ -28,6 +28,7 @@ import java.util.Arrays; import java.util.List; import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.utils.nativefst.builder.FSTBuilder; import org.apache.pinot.segment.local.utils.nativefst.builder.FSTInfo; import org.testng.annotations.Test; @@ -41,7 +42,7 @@ /** * Additional tests for {@link ImmutableFST}. */ -public final class ImmutableFSTTest { +public final class ImmutableFSTTest implements PinotBuffersAfterMethodCheckRule { public List _expected = Arrays.asList("a", "aba", "ac", "b", "ba", "c"); public static void walkNode(byte[] buffer, int depth, FST fst, int node, int cnt, List result) { @@ -77,20 +78,26 @@ private static void verifyContent(FST fst, List expected) { @Test public void testVersion5() throws IOException { + FST fst = null; try (InputStream inputStream = getClass().getClassLoader().getResourceAsStream("data/abc.native.fst")) { - FST fst = FST.read(inputStream); + fst = FST.read(inputStream); assertFalse(fst.getFlags().contains(FSTFlags.NUMBERS)); verifyContent(fst, _expected); + } finally { + close(fst); } } @Test public void testVersion5WithNumbers() throws IOException { + FST fst = null; try (InputStream inputStream = getClass().getClassLoader().getResourceAsStream("data/abc-numbers.native.fst")) { - FST fst = FST.read(inputStream); + fst = FST.read(inputStream); assertTrue(fst.getFlags().contains(FSTFlags.NUMBERS)); verifyContent(fst, _expected); + } finally { + close(fst); } } @@ -98,11 +105,14 @@ public void testVersion5WithNumbers() public void testArcsAndNodes() throws IOException { for (String resourceName : new String[]{"data/abc.native.fst", "data/abc-numbers.native.fst"}) { + FST fst = null; try (InputStream inputStream = getClass().getClassLoader().getResourceAsStream(resourceName)) { - FST fst = FST.read(inputStream); + fst = FST.read(inputStream); FSTInfo fstInfo = new FSTInfo(fst); assertEquals(fstInfo._nodeCount, 4); assertEquals(fstInfo._arcsCount, 7); + } finally { + close(fst); } } } @@ -110,8 +120,9 @@ public void testArcsAndNodes() @Test public void testNumbers() throws IOException { + FST fst = null; try (InputStream inputStream = getClass().getClassLoader().getResourceAsStream("data/abc-numbers.native.fst")) { - FST fst = FST.read(inputStream); + fst = FST.read(inputStream); assertTrue(fst.getFlags().contains(FSTFlags.NEXTBIT)); // Get all numbers for nodes. @@ -121,6 +132,8 @@ public void testNumbers() result.sort(null); assertEquals(result, Arrays.asList("0 c", "1 b", "2 ba", "3 a", "4 ac", "5 aba")); + } finally { + close(fst); } } @@ -133,15 +146,28 @@ public void testSave() for (String input : inputList) { builder.add(input.getBytes(UTF_8), 0, input.length(), 127); } - FST fst = builder.complete(); - - File fstFile = new File(FileUtils.getTempDirectory(), "test.native.fst"); - fst.save(new FileOutputStream(fstFile)); - - try (FileInputStream inputStream = new FileInputStream(fstFile)) { - verifyContent(FST.read(inputStream, ImmutableFST.class, true), inputList); + FST fst1 = builder.complete(); + ImmutableFST fst2 = null; + File fstFile = null; + try { + fstFile = new File(FileUtils.getTempDirectory(), "test.native.fst"); + fst1.save(new FileOutputStream(fstFile)); + + try (FileInputStream inputStream = new FileInputStream(fstFile)) { + fst2 = FST.read(inputStream, ImmutableFST.class, true); + verifyContent(fst2, inputList); + } + } finally { + close(fst1); + close(fst2); + FileUtils.deleteQuietly(fstFile); } + } - FileUtils.deleteQuietly(fstFile); + private static void close(FST fst) + throws IOException { + if (fst instanceof ImmutableFST) { + ((ImmutableFST) fst)._mutableBytesStore.close(); + } } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/nativefst/SerializerTestBase.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/nativefst/SerializerTestBase.java index 82af070da8aa..383e7c2f78e0 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/nativefst/SerializerTestBase.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/nativefst/SerializerTestBase.java @@ -26,6 +26,7 @@ import java.util.Arrays; import java.util.HashSet; import java.util.List; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; import org.apache.pinot.segment.local.io.writer.impl.DirectMemoryManager; import org.apache.pinot.segment.local.utils.nativefst.builder.FSTBuilder; import org.apache.pinot.segment.local.utils.nativefst.builder.FSTSerializer; @@ -40,7 +41,7 @@ /** * Base class for serializer tests */ -public abstract class SerializerTestBase { +public abstract class SerializerTestBase implements PinotBuffersAfterMethodCheckRule { /* * Drain bytes from a byte buffer to a string. */ @@ -137,7 +138,13 @@ private void checkSerialization0(FSTSerializer serializer, final byte[][] in, FS byte[] fstData = serializer.serialize(root, new ByteArrayOutputStream()).toByteArray(); FST fst = FST.read(new ByteArrayInputStream(fstData), hasOutputSymbols, new DirectMemoryManager(SerializerTestBase.class.getName())); - checkCorrect(in, fst); + try { + checkCorrect(in, fst); + } finally { + if (fst instanceof ImmutableFST) { + ((ImmutableFST) fst)._mutableBytesStore.close(); + } + } } /* @@ -178,17 +185,23 @@ public void testAutomatonWithNodeNumbers() byte[] fstData = createSerializer().withNumbers().serialize(fst, new ByteArrayOutputStream()).toByteArray(); fst = FST.read(new ByteArrayInputStream(fstData), true, new DirectMemoryManager(SerializerTestBase.class.getName())); + try { - // Ensure we have the NUMBERS flag set. - assertTrue(fst.getFlags().contains(NUMBERS)); + // Ensure we have the NUMBERS flag set. + assertTrue(fst.getFlags().contains(NUMBERS)); - // Get all numbers from nodes. - byte[] buffer = new byte[128]; - List result = new ArrayList<>(); - ImmutableFSTTest.walkNode(buffer, 0, fst, fst.getRootNode(), 0, result); + // Get all numbers from nodes. + byte[] buffer = new byte[128]; + List result = new ArrayList<>(); + ImmutableFSTTest.walkNode(buffer, 0, fst, fst.getRootNode(), 0, result); - result.sort(null); - assertEquals(result, Arrays.asList("0 a", "1 aba", "2 ac", "3 b", "4 ba", "5 c")); + result.sort(null); + assertEquals(result, Arrays.asList("0 a", "1 aba", "2 ac", "3 b", "4 ba", "5 c")); + } finally { + if (fst instanceof ImmutableFST) { + ((ImmutableFST) fst)._mutableBytesStore.close(); + } + } } protected abstract FSTSerializer createSerializer(); diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/memory/PinotDataBuffer.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/memory/PinotDataBuffer.java index c2fb5084a754..e055ca9dfbaf 100644 --- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/memory/PinotDataBuffer.java +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/memory/PinotDataBuffer.java @@ -27,6 +27,7 @@ import java.nio.ByteOrder; import java.nio.channels.FileChannel; import java.util.ArrayList; +import java.util.Collections; import java.util.List; import java.util.Map; import java.util.concurrent.atomic.AtomicLong; @@ -328,6 +329,10 @@ protected static void cleanStats() { } public static List getBufferInfo() { + if (BUFFER_CONTEXT_MAP.isEmpty()) { + return Collections.emptyList(); + } + List bufferInfo = new ArrayList<>(BUFFER_CONTEXT_MAP.size()); for (BufferContext bufferContext : BUFFER_CONTEXT_MAP.values()) { bufferInfo.add(bufferContext.toString()); @@ -335,6 +340,18 @@ public static List getBufferInfo() { return bufferInfo; } + @VisibleForTesting + public static void closeOpenBuffers() { + for (PinotDataBuffer buffer : BUFFER_CONTEXT_MAP.keySet()) { + try { + buffer.close(); + } catch (IOException e) { + e.printStackTrace(); + } + } + BUFFER_CONTEXT_MAP.clear(); + } + private static String getBufferStats() { return String.format("Direct buffer count: %s, size: %s; Mmap buffer count: %s, size: %s", DIRECT_BUFFER_COUNT.get(), DIRECT_BUFFER_USAGE.get(), MMAP_BUFFER_COUNT.get(), MMAP_BUFFER_USAGE.get());