diff --git a/vector/src/main/java/org/apache/arrow/vector/dictionary/DictionaryEncoder.java b/vector/src/main/java/org/apache/arrow/vector/dictionary/DictionaryEncoder.java index 4af1a8693f..284f02032c 100644 --- a/vector/src/main/java/org/apache/arrow/vector/dictionary/DictionaryEncoder.java +++ b/vector/src/main/java/org/apache/arrow/vector/dictionary/DictionaryEncoder.java @@ -165,7 +165,7 @@ static void retrieveIndexVector( for (int i = start; i < end; i++) { if (!indices.isNull(i)) { int indexAsInt = (int) indices.getValueAsLong(i); - if (indexAsInt > dictionaryCount) { + if (indexAsInt < 0 || indexAsInt >= dictionaryCount) { throw new IllegalArgumentException( "Provided dictionary does not contain value for index " + indexAsInt); } diff --git a/vector/src/test/java/org/apache/arrow/vector/TestDictionaryVector.java b/vector/src/test/java/org/apache/arrow/vector/TestDictionaryVector.java index 0945919b91..d97ce9e39c 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestDictionaryVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestDictionaryVector.java @@ -21,6 +21,7 @@ import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.fail; @@ -942,6 +943,35 @@ public void testNoMemoryLeak() { assertEquals(0, allocator.getAllocatedMemory(), "decode memory leak"); } + @Test + public void testDecodeRejectsDictionaryIndicesOutsideBounds() { + try (final IntVector indices = newVector(IntVector.class, "", Types.MinorType.INT, allocator); + final VarCharVector dictionaryVector = newVarCharVector("dict", allocator)) { + setVector(dictionaryVector, zero, one); + Dictionary dictionary = + new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null)); + + setVector(indices, dictionaryVector.getValueCount()); + IllegalArgumentException upperBoundException = + assertThrows( + IllegalArgumentException.class, + () -> DictionaryEncoder.decode(indices, dictionary, allocator)); + assertEquals( + "Provided dictionary does not contain value for index 2", + upperBoundException.getMessage()); + + setVector(indices, -1); + IllegalArgumentException negativeException = + assertThrows( + IllegalArgumentException.class, + () -> DictionaryEncoder.decode(indices, dictionary, allocator)); + assertEquals( + "Provided dictionary does not contain value for index -1", + negativeException.getMessage()); + } + assertEquals(0, allocator.getAllocatedMemory(), "decode memory leak"); + } + @Test public void testListNoMemoryLeak() { // Create a new value vector @@ -1053,7 +1083,7 @@ public void testStructNoMemoryLeak() { NullableStructWriter writer = indices.getWriter(); writer.allocate(); writer.start(); - writer.integer("f0").writeInt(1); + writer.integer("f0").writeInt(0); writer.integer("f1").writeInt(3); writer.end(); writer.setValueCount(1);