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..73809a3136 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestDictionaryVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestDictionaryVector.java @@ -942,6 +942,33 @@ public void testNoMemoryLeak() { assertEquals(0, allocator.getAllocatedMemory(), "decode memory leak"); } + @Test + public void testDecodeIndexOutOfBounds() { + // valid indices are 0..dictionaryCount-1; index == dictionaryCount and negative indices + // must be rejected before dereferencing the dictionary vector. + 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, 2); + try (final ValueVector decoded = DictionaryEncoder.decode(indices, dictionary, allocator)) { + fail("There should be an exception when decoding an index equal to the dictionary size"); + } catch (IllegalArgumentException e) { + assertEquals("Provided dictionary does not contain value for index 2", e.getMessage()); + } + + setVector(indices, -1); + try (final ValueVector decoded = DictionaryEncoder.decode(indices, dictionary, allocator)) { + fail("There should be an exception when decoding a negative index"); + } catch (IllegalArgumentException e) { + assertEquals("Provided dictionary does not contain value for index -1", e.getMessage()); + } + } + assertEquals(0, allocator.getAllocatedMemory(), "decode memory leak"); + } + @Test public void testListNoMemoryLeak() { // Create a new value vector @@ -1053,7 +1080,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);