From 304b848e6ea232001dc376ed7867b68944f82c5f Mon Sep 17 00:00:00 2001 From: Clint Wylie Date: Wed, 12 Jan 2022 22:13:29 -0800 Subject: [PATCH 1/2] fix array type strategy write size tracking --- .../segment/column/NullableTypeStrategy.java | 23 ++++--- .../druid/segment/column/TypeStrategies.java | 2 +- .../segment/column/TypeStrategiesTest.java | 62 +++++++++++++------ 3 files changed, 59 insertions(+), 28 deletions(-) diff --git a/core/src/main/java/org/apache/druid/segment/column/NullableTypeStrategy.java b/core/src/main/java/org/apache/druid/segment/column/NullableTypeStrategy.java index 1e0ea8db0390..998e6a92dc4b 100644 --- a/core/src/main/java/org/apache/druid/segment/column/NullableTypeStrategy.java +++ b/core/src/main/java/org/apache/druid/segment/column/NullableTypeStrategy.java @@ -92,20 +92,27 @@ public int write(ByteBuffer buffer, @Nullable T value, int maxSizeBytes) public T read(ByteBuffer buffer, int offset) { final int oldPosition = buffer.position(); - buffer.position(offset); - T value = read(buffer); - buffer.position(oldPosition); - return value; + try { + buffer.position(offset); + T value = read(buffer); + return value; + } + finally { + buffer.position(oldPosition); + } } public int write(ByteBuffer buffer, int offset, @Nullable T value, int maxSizeBytes) { final int oldPosition = buffer.position(); - buffer.position(offset); - final int size = write(buffer, value, maxSizeBytes); - buffer.position(oldPosition); - return size; + try { + buffer.position(offset); + return write(buffer, value, maxSizeBytes); + } + finally { + buffer.position(oldPosition); + } } @Override diff --git a/core/src/main/java/org/apache/druid/segment/column/TypeStrategies.java b/core/src/main/java/org/apache/druid/segment/column/TypeStrategies.java index de5604d46b0f..c7b212fd00e9 100644 --- a/core/src/main/java/org/apache/druid/segment/column/TypeStrategies.java +++ b/core/src/main/java/org/apache/druid/segment/column/TypeStrategies.java @@ -466,7 +466,7 @@ public int write(ByteBuffer buffer, Object[] value, int maxSizeBytes) remaining = 0; } else { sizeBytes += written; - remaining -= sizeBytes; + remaining -= written; } } return extraNeeded < 0 ? extraNeeded : sizeBytes; diff --git a/core/src/test/java/org/apache/druid/segment/column/TypeStrategiesTest.java b/core/src/test/java/org/apache/druid/segment/column/TypeStrategiesTest.java index 1037bc9e30ba..1579741ecc78 100644 --- a/core/src/test/java/org/apache/druid/segment/column/TypeStrategiesTest.java +++ b/core/src/test/java/org/apache/druid/segment/column/TypeStrategiesTest.java @@ -32,6 +32,7 @@ import javax.annotation.Nullable; import java.nio.ByteBuffer; +import java.util.Arrays; public class TypeStrategiesTest { @@ -469,7 +470,7 @@ public void testArrayTypeStrategy() // string arrays strategy = new TypeStrategies.ArrayTypeStrategy(ColumnType.STRING_ARRAY); final String[] someStringArray = new String[]{"hello", "hi", null, "hey"}; - final Object[] someObjectStringArray = new String[]{"hello", "hi", null, "hey"}; + final Object[] someObjectStringArray = new Object[]{"hello", "hi", null, "hey"}; assertArrayStrategy(strategy, empty); assertArrayStrategy(strategy, someStringArray); @@ -493,6 +494,20 @@ public void testArrayTypeStrategy() assertArrayStrategy(strategy, nester); } + @Test + public void testArrayTypeStrategyCloseToTheLimit() + { + TypeStrategy strategy = new TypeStrategies.ArrayTypeStrategy(ColumnType.STRING_ARRAY); + String filler = "AAAAAAAAAA"; + // test runs at offset 10, and 5 bytes for array null byte and size int, so + int size = (int) Math.floor( + (double) (buffer.capacity() - 5) / (double) ColumnType.STRING.getNullableStrategy().estimateSizeBytes(filler) + ); + Object[] filler_array = new Object[size]; + Arrays.fill(filler_array, filler); + assertArrayStrategy(strategy, filler_array, buffer.capacity(), 0); + } + private void assertStrategy(TypeStrategy strategy, @Nullable T value) { final int maxSize = 2048; @@ -549,8 +564,34 @@ private void assertArrayStrategy(TypeStrategy strategy, @Nullable Object[] value final int expectedLength = strategy.estimateSizeBytes(value); Assert.assertNotEquals(0, expectedLength); + // basic tests at some position and offset + assertArrayStrategy(strategy,value, maxSize, 10); + + buffer.position(0); + + // test buffer offset when with different position + NullableTypeStrategy nullableTypeStrategy = new NullableTypeStrategy(strategy); + Assert.assertEquals(expectedLength, strategy.write(buffer, 1024, value, maxSize)); + Assert.assertArrayEquals(value, (Object[]) strategy.read(buffer, 1024)); + Assert.assertEquals(0, buffer.position()); + + // test buffer offset nullable write read value + Assert.assertEquals(1 + expectedLength, nullableTypeStrategy.write(buffer, 1024, value, maxSize)); + Assert.assertArrayEquals(value, (Object[]) nullableTypeStrategy.read(buffer, 1024)); + Assert.assertEquals(0, buffer.position()); + + // test buffer offset nullable write read null + Assert.assertEquals(1, nullableTypeStrategy.write(buffer, 1024, null, maxSize)); + Assert.assertNull(nullableTypeStrategy.read(buffer, 1024)); + Assert.assertEquals(0, buffer.position()); + } + + private void assertArrayStrategy(TypeStrategy strategy, @Nullable Object[] value, int maxSize, int offset) + { + final int expectedLength = strategy.estimateSizeBytes(value); + Assert.assertNotEquals(0, expectedLength); + // test buffer - int offset = 10; buffer.position(offset); Assert.assertEquals(expectedLength, strategy.write(buffer, value, maxSize)); Assert.assertEquals(expectedLength, buffer.position() - offset); @@ -574,23 +615,6 @@ private void assertArrayStrategy(TypeStrategy strategy, @Nullable Object[] value buffer.position(offset); Assert.assertNull(nullableTypeStrategy.read(buffer)); Assert.assertEquals(1, buffer.position() - offset); - - buffer.position(0); - - // test buffer offset - Assert.assertEquals(expectedLength, strategy.write(buffer, 1024, value, maxSize)); - Assert.assertArrayEquals(value, (Object[]) strategy.read(buffer, 1024)); - Assert.assertEquals(0, buffer.position()); - - // test buffer offset nullable write read value - Assert.assertEquals(1 + expectedLength, nullableTypeStrategy.write(buffer, 1024, value, maxSize)); - Assert.assertArrayEquals(value, (Object[]) nullableTypeStrategy.read(buffer, 1024)); - Assert.assertEquals(0, buffer.position()); - - // test buffer offset nullable write read null - Assert.assertEquals(1, nullableTypeStrategy.write(buffer, 1024, null, maxSize)); - Assert.assertNull(nullableTypeStrategy.read(buffer, 1024)); - Assert.assertEquals(0, buffer.position()); } public static class NullableLongPair extends Pair implements Comparable From ae8ea333cb0c9bb9b8faf91a0ed80323e45ee751 Mon Sep 17 00:00:00 2001 From: Clint Wylie Date: Thu, 13 Jan 2022 00:18:38 -0800 Subject: [PATCH 2/2] fix checkstyle --- .../org/apache/druid/segment/column/TypeStrategiesTest.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/core/src/test/java/org/apache/druid/segment/column/TypeStrategiesTest.java b/core/src/test/java/org/apache/druid/segment/column/TypeStrategiesTest.java index 1579741ecc78..861562b5de7a 100644 --- a/core/src/test/java/org/apache/druid/segment/column/TypeStrategiesTest.java +++ b/core/src/test/java/org/apache/druid/segment/column/TypeStrategiesTest.java @@ -501,7 +501,7 @@ public void testArrayTypeStrategyCloseToTheLimit() String filler = "AAAAAAAAAA"; // test runs at offset 10, and 5 bytes for array null byte and size int, so int size = (int) Math.floor( - (double) (buffer.capacity() - 5) / (double) ColumnType.STRING.getNullableStrategy().estimateSizeBytes(filler) + (double) (buffer.capacity() - 5) / (double) ColumnType.STRING.getNullableStrategy().estimateSizeBytes(filler) ); Object[] filler_array = new Object[size]; Arrays.fill(filler_array, filler); @@ -565,7 +565,7 @@ private void assertArrayStrategy(TypeStrategy strategy, @Nullable Object[] value Assert.assertNotEquals(0, expectedLength); // basic tests at some position and offset - assertArrayStrategy(strategy,value, maxSize, 10); + assertArrayStrategy(strategy, value, maxSize, 10); buffer.position(0);