diff --git a/parquet-column/src/main/java/org/apache/parquet/column/values/dictionary/DictionaryValuesReader.java b/parquet-column/src/main/java/org/apache/parquet/column/values/dictionary/DictionaryValuesReader.java index 53fafc55dc..16fa214100 100644 --- a/parquet-column/src/main/java/org/apache/parquet/column/values/dictionary/DictionaryValuesReader.java +++ b/parquet-column/src/main/java/org/apache/parquet/column/values/dictionary/DictionaryValuesReader.java @@ -39,10 +39,13 @@ public class DictionaryValuesReader extends ValuesReader { private Dictionary dictionary; + private final int maxDictionaryId; + private RunLengthBitPackingHybridDecoder decoder; public DictionaryValuesReader(Dictionary dictionary) { this.dictionary = dictionary; + this.maxDictionaryId = dictionary.getMaxId(); } @Override @@ -65,62 +68,51 @@ public int readInt() throws IOException { @Override public int readValueDictionaryId() { - try { - return decoder.readInt(); - } catch (IOException e) { - throw new ParquetDecodingException(e); - } + return readDictionaryId(); } @Override public Binary readBytes() { - try { - return dictionary.decodeToBinary(decoder.readInt()); - } catch (IOException e) { - throw new ParquetDecodingException(e); - } + return dictionary.decodeToBinary(readDictionaryId()); } @Override public float readFloat() { - try { - return dictionary.decodeToFloat(decoder.readInt()); - } catch (IOException e) { - throw new ParquetDecodingException(e); - } + return dictionary.decodeToFloat(readDictionaryId()); } @Override public double readDouble() { - try { - return dictionary.decodeToDouble(decoder.readInt()); - } catch (IOException e) { - throw new ParquetDecodingException(e); - } + return dictionary.decodeToDouble(readDictionaryId()); } @Override public int readInteger() { - try { - return dictionary.decodeToInt(decoder.readInt()); - } catch (IOException e) { - throw new ParquetDecodingException(e); - } + return dictionary.decodeToInt(readDictionaryId()); } @Override public long readLong() { + return dictionary.decodeToLong(readDictionaryId()); + } + + @Override + public void skip() { try { - return dictionary.decodeToLong(decoder.readInt()); + decoder.readInt(); // Type does not matter as we are just skipping dictionary keys } catch (IOException e) { throw new ParquetDecodingException(e); } } - @Override - public void skip() { + private int readDictionaryId() { try { - decoder.readInt(); // Type does not matter as we are just skipping dictionary keys + int id = decoder.readInt(); + if (id < 0 || id > maxDictionaryId) { + throw new ParquetDecodingException( + "Dictionary id " + id + " is outside the valid range 0 to " + maxDictionaryId); + } + return id; } catch (IOException e) { throw new ParquetDecodingException(e); } diff --git a/parquet-column/src/test/java/org/apache/parquet/column/values/dictionary/TestDictionary.java b/parquet-column/src/test/java/org/apache/parquet/column/values/dictionary/TestDictionary.java index 13033404ce..f2a0ab41e8 100644 --- a/parquet-column/src/test/java/org/apache/parquet/column/values/dictionary/TestDictionary.java +++ b/parquet-column/src/test/java/org/apache/parquet/column/values/dictionary/TestDictionary.java @@ -25,6 +25,7 @@ import static org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.FLOAT; import static org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.INT32; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.assertj.core.data.Offset.offset; import java.io.IOException; @@ -50,6 +51,7 @@ import org.apache.parquet.column.values.plain.BinaryPlainValuesReader; import org.apache.parquet.column.values.plain.PlainValuesReader; import org.apache.parquet.column.values.plain.PlainValuesWriter; +import org.apache.parquet.io.ParquetDecodingException; import org.apache.parquet.io.api.Binary; import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName; import org.junit.jupiter.api.AfterEach; @@ -803,6 +805,22 @@ public void testZeroValues() throws IOException { } } + @Test + public void testInvalidDictionaryId() throws IOException { + BytesInput dictionaryBytes = BytesInput.fromInt(34); + DictionaryPage dictionaryPage = new DictionaryPage(dictionaryBytes, 1, PLAIN); + ColumnDescriptor descriptor = new ColumnDescriptor(new String[] {"foo"}, INT32, 0, 0); + Dictionary dictionary = PLAIN.initDictionary(descriptor, dictionaryPage); + DictionaryValuesReader reader = new DictionaryValuesReader(dictionary); + + // A bit width of 1 can represent dictionary id 1, but this dictionary only contains id 0. + reader.initFromPage(1, BytesInput.from(new byte[] {0x01, 0x02, 0x01}).toInputStream()); + + assertThatThrownBy(reader::readInteger) + .isInstanceOf(ParquetDecodingException.class) + .hasMessage("Dictionary id 1 is outside the valid range 0 to 0"); + } + @Test public void testBooleanDictionary() throws IOException { // Create a dictionary page with boolean values (false, true)