-
Notifications
You must be signed in to change notification settings - Fork 1.6k
PARQUET-373: Fix flaky MemoryManager tests. #269
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -18,14 +18,14 @@ | |
| */ | ||
| package org.apache.parquet.hadoop; | ||
|
|
||
| import org.apache.commons.io.FileUtils; | ||
| import org.apache.hadoop.conf.Configuration; | ||
| import org.apache.hadoop.fs.Path; | ||
| import org.apache.hadoop.mapreduce.RecordWriter; | ||
| import org.junit.After; | ||
| import org.junit.Assert; | ||
| import org.junit.Before; | ||
| import org.junit.Rule; | ||
| import org.junit.Test; | ||
| import org.junit.rules.TemporaryFolder; | ||
| import org.apache.parquet.hadoop.example.GroupWriteSupport; | ||
| import org.apache.parquet.hadoop.metadata.CompressionCodecName; | ||
| import org.apache.parquet.schema.MessageTypeParser; | ||
|
|
@@ -44,89 +44,150 @@ public class TestMemoryManager { | |
| "required int32 line;\n" + | ||
| "required binary content;\n" + | ||
| "}"; | ||
| long expectPoolSize; | ||
| int rowGroupSize; | ||
| long expectedPoolSize; | ||
| ParquetOutputFormat parquetOutputFormat; | ||
| CompressionCodecName codec; | ||
| int counter = 0; | ||
| boolean firstRegister = true; | ||
|
|
||
| @Before | ||
| public void setUp() { | ||
| GroupWriteSupport.setSchema(MessageTypeParser.parseMessageType(writeSchema),conf); | ||
| expectPoolSize = Math.round((double) ManagementFactory.getMemoryMXBean().getHeapMemoryUsage().getMax | ||
| () * MemoryManager.DEFAULT_MEMORY_POOL_RATIO); | ||
| rowGroupSize = (int) Math.floor(expectPoolSize / 2); | ||
| conf.setInt(ParquetOutputFormat.BLOCK_SIZE, rowGroupSize); | ||
| codec = CompressionCodecName.UNCOMPRESSED; | ||
| public void setUp() throws Exception { | ||
| parquetOutputFormat = new ParquetOutputFormat(new GroupWriteSupport()); | ||
|
|
||
| GroupWriteSupport.setSchema(MessageTypeParser.parseMessageType(writeSchema), conf); | ||
| expectedPoolSize = Math.round((double) | ||
| ManagementFactory.getMemoryMXBean().getHeapMemoryUsage().getMax() * | ||
| MemoryManager.DEFAULT_MEMORY_POOL_RATIO); | ||
|
|
||
| long rowGroupSize = expectedPoolSize / 2; | ||
| conf.setLong(ParquetOutputFormat.BLOCK_SIZE, rowGroupSize); | ||
|
|
||
| // the memory manager is not initialized until a writer is created | ||
| createWriter(1).close(null); | ||
| } | ||
|
|
||
| @After | ||
| public void tearDown() throws Exception{ | ||
| FileUtils.deleteDirectory(new File("target/test")); | ||
| @Test | ||
| public void testMemoryManagerUpperLimit() { | ||
| // Verify the memory pool size | ||
| // this value tends to change a little between setup and tests, so this | ||
| // validates that it is within 5% of the expected value | ||
| long poolSize = ParquetOutputFormat.getMemoryManager().getTotalMemoryPool(); | ||
| Assert.assertTrue("Pool size should be within 5% of the expected value", | ||
| Math.abs(expectedPoolSize - poolSize) < (long) (expectedPoolSize * 0.05)); | ||
| } | ||
|
|
||
| @Test | ||
| public void testMemoryManager() throws Exception { | ||
| //Verify the adjusted rowGroupSize of writers | ||
| long poolSize = ParquetOutputFormat.getMemoryManager().getTotalMemoryPool(); | ||
| long rowGroupSize = poolSize / 2; | ||
| conf.setLong(ParquetOutputFormat.BLOCK_SIZE, rowGroupSize); | ||
|
|
||
| Assert.assertTrue("Pool should hold 2 full row groups", | ||
| (2 * rowGroupSize) <= poolSize); | ||
| Assert.assertTrue("Pool should not hold 3 full row groups", | ||
| poolSize < (3 * rowGroupSize)); | ||
|
|
||
| Assert.assertEquals("Allocations should start out at 0", | ||
| 0, getTotalAllocation()); | ||
|
|
||
| RecordWriter writer1 = createWriter(1); | ||
| verifyRowGroupSize(rowGroupSize); | ||
| Assert.assertTrue("Allocations should never exceed pool size", | ||
| getTotalAllocation() <= poolSize); | ||
| Assert.assertEquals("First writer should be limited by row group size", | ||
| rowGroupSize, getTotalAllocation()); | ||
|
|
||
| RecordWriter writer2 = createWriter(2); | ||
| verifyRowGroupSize(rowGroupSize); | ||
| Assert.assertTrue("Allocations should never exceed pool size", | ||
| getTotalAllocation() <= poolSize); | ||
| Assert.assertEquals("Second writer should be limited by row group size", | ||
| 2 * rowGroupSize, getTotalAllocation()); | ||
|
|
||
| RecordWriter writer3 = createWriter(3); | ||
| verifyRowGroupSize((int) Math.floor(expectPoolSize / 3)); | ||
| Assert.assertTrue("Allocations should never exceed pool size", | ||
| getTotalAllocation() <= poolSize); | ||
|
|
||
| writer1.close(null); | ||
| verifyRowGroupSize(rowGroupSize); | ||
| Assert.assertTrue("Allocations should never exceed pool size", | ||
| getTotalAllocation() <= poolSize); | ||
| Assert.assertEquals("Allocations should be increased to the row group size", | ||
| 2 * rowGroupSize, getTotalAllocation()); | ||
|
|
||
| writer2.close(null); | ||
| verifyRowGroupSize(rowGroupSize); | ||
| Assert.assertTrue("Allocations should never exceed pool size", | ||
| getTotalAllocation() <= poolSize); | ||
| Assert.assertEquals("Allocations should be increased to the row group size", | ||
| rowGroupSize, getTotalAllocation()); | ||
|
|
||
| writer3.close(null); | ||
| Assert.assertEquals("Allocations should be increased to the row group size", | ||
| 0, getTotalAllocation()); | ||
| } | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Should we add a validation that checks that getTotalAllocation() is 0 after closing all the writers?
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Done. |
||
|
|
||
| //Verify the memory pool | ||
| Assert.assertEquals("memory pool size is incorrect.", expectPoolSize, | ||
| parquetOutputFormat.getMemoryManager().getTotalMemoryPool()); | ||
| @Test | ||
| public void testReallocationCallback() throws Exception { | ||
| // validate assumptions | ||
| long poolSize = ParquetOutputFormat.getMemoryManager().getTotalMemoryPool(); | ||
| long rowGroupSize = poolSize / 2; | ||
| conf.setLong(ParquetOutputFormat.BLOCK_SIZE, rowGroupSize); | ||
|
|
||
| Assert.assertTrue("Pool should hold 2 full row groups", | ||
| (2 * rowGroupSize) <= poolSize); | ||
| Assert.assertTrue("Pool should not hold 3 full row groups", | ||
| poolSize < (3 * rowGroupSize)); | ||
|
|
||
| Runnable callback = new Runnable() { | ||
| @Override | ||
| public void run() { | ||
| counter++; | ||
| } | ||
| }; | ||
|
|
||
| //Verify Callback mechanism | ||
| Assert.assertEquals("counter calculated by callback is incorrect.", 1, counter); | ||
| Assert.assertEquals("CallBack is duplicated.", 1, parquetOutputFormat.getMemoryManager() | ||
| .getScaleCallBacks().size()); | ||
| } | ||
| // first-time registration should succeed | ||
| ParquetOutputFormat.getMemoryManager() | ||
| .registerScaleCallBack("increment-test-counter", callback); | ||
|
|
||
| private RecordWriter createWriter(int index) throws Exception{ | ||
| Path file = new Path("target/test/", "parquet" + index); | ||
| parquetOutputFormat = new ParquetOutputFormat(new GroupWriteSupport()); | ||
| RecordWriter writer = parquetOutputFormat.getRecordWriter(conf, file, codec); | ||
| try { | ||
| parquetOutputFormat.getMemoryManager().registerScaleCallBack("increment-test-counter", | ||
| new Runnable() { | ||
| @Override | ||
| public void run() { | ||
| counter++; | ||
| } | ||
| }); | ||
| if (!firstRegister) { | ||
| Assert.fail("Duplicated registering callback should throw duplicates exception."); | ||
| } | ||
| firstRegister = false; | ||
| ParquetOutputFormat.getMemoryManager() | ||
| .registerScaleCallBack("increment-test-counter", callback); | ||
| Assert.fail("Duplicated registering callback should throw duplicates exception."); | ||
| } catch (IllegalArgumentException e) { | ||
| if (firstRegister) { | ||
| Assert.fail("Registering the same callback first time should succeed."); | ||
| } | ||
| // expected | ||
| } | ||
|
|
||
| // hit the limit once and clean up | ||
| RecordWriter writer1 = createWriter(1); | ||
| RecordWriter writer2 = createWriter(2); | ||
| RecordWriter writer3 = createWriter(3); | ||
| writer1.close(null); | ||
| writer2.close(null); | ||
| writer3.close(null); | ||
|
|
||
| //Verify Callback mechanism | ||
| Assert.assertEquals("Allocations should be adjusted once", 1, counter); | ||
| Assert.assertEquals("Should not allow duplicate callbacks", | ||
| 1, ParquetOutputFormat.getMemoryManager().getScaleCallBacks().size()); | ||
| } | ||
|
|
||
| @Rule | ||
| public TemporaryFolder temp = new TemporaryFolder(); | ||
|
|
||
| private RecordWriter createWriter(int index) throws Exception { | ||
| File file = temp.newFile(String.valueOf(index) + ".parquet"); | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Are we deleting this file after it's being used? Mabye a file.deleteOnExit() here would help to avoid leaving temporary files after tests.
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Removing this file is done by the TemporaryFolder. Because we use the Rule annotation, it is created and cleaned up by the framework. |
||
| if (!file.delete()) { | ||
| throw new RuntimeException("Could not delete file: " + file); | ||
| } | ||
| RecordWriter writer = parquetOutputFormat.getRecordWriter( | ||
| conf, new Path(file.toString()), | ||
| CompressionCodecName.UNCOMPRESSED); | ||
|
|
||
| return writer; | ||
| } | ||
|
|
||
| private void verifyRowGroupSize(int expectRowGroupSize) { | ||
| Set<InternalParquetRecordWriter> writers = parquetOutputFormat.getMemoryManager() | ||
| .getWriterList().keySet(); | ||
| private long getTotalAllocation() { | ||
| Set<InternalParquetRecordWriter> writers = ParquetOutputFormat | ||
| .getMemoryManager().getWriterList().keySet(); | ||
| long total = 0; | ||
| for (InternalParquetRecordWriter writer : writers) { | ||
| Assert.assertEquals("wrong rowGroupSize", expectRowGroupSize, | ||
| writer.getRowGroupSizeThreshold(), 1); | ||
| total += writer.getRowGroupSizeThreshold(); | ||
| } | ||
| return total; | ||
| } | ||
| } | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I'm not sure I understand these 2 tests. Are we testing that 'poolsize / 2' works correctly?
I see this test (2 * rowGroupsize) <= poolSize, but before that we do rowGroupsize = poolSize / 2. Isn't this always true? Same for poolSize < (3 * rowGroupSize).
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
These tests validate assumptions about the current settings and environment. The idea is to catch environment problems early rather than complaining about a code failure.