- Fixed corner case in Parquet page v1 that didn't account for padded trailing values in bit-packed runs.

Author: avalerio-tkd

src/processing/parquet_utils.cpp

@@ -76,10 +76,22 @@ uint32_t ReadV1RunHeaderUleb128(tcb::span<const uint8_t> bytes, size_t& offset)
 // - present_count = number of decoded definition levels equal to max_def_level.
 //
 size_t CountPresentFromDefinitionLevelsV1(tcb::span<const uint8_t> def_payload, int32_t num_values, int32_t max_def_level) {
+    if (num_values < 0) {
+        throw InvalidInputException("Invalid V1 definition levels: num_values must be non-negative, got " + 
+            std::to_string(num_values));
+    }
+    if (max_def_level <= 0) {
+        throw InvalidInputException("Invalid V1 definition levels: max_def_level must be positive, got " + 
+            std::to_string(max_def_level));
+    }
+
     // Definition level bit width is ceil(log2(max_def_level + 1)).
+    // Computes the minimum number of bits needed to represent definition levels from 0..max_def_level.
     int bit_width = 0;
-    while ((1u << bit_width) <= static_cast<uint32_t>(max_def_level)) {
+    uint32_t def_level_domain = static_cast<uint32_t>(max_def_level);
+    while (def_level_domain > 0) {
         ++bit_width;
+        def_level_domain >>= 1;
     }
     if (bit_width <= 0) {
         throw InvalidInputException("Invalid V1 definition levels: computed bit_width must be positive");
@@ -124,7 +136,7 @@ size_t CountPresentFromDefinitionLevelsV1(tcb::span<const uint8_t> def_payload,
             const size_t num_groups = static_cast<size_t>(header >> 1);
             const size_t run_len = num_groups * 8;
             const size_t remaining = static_cast<size_t>(num_values) - decoded_values;
-            if (num_groups == 0 || run_len > remaining) {
+            if (num_groups == 0) {
                 throw InvalidInputException("Invalid DATA_PAGE_V1 definition levels: invalid bit-packed run length");
             }
 
@@ -148,7 +160,10 @@ size_t CountPresentFromDefinitionLevelsV1(tcb::span<const uint8_t> def_payload,
                 return v;
             };
 
-            for (size_t i = 0; i < run_len; ++i) {
+            // A final bit-packed run may include padded trailing values to complete
+            // 8-value groups. Decode only the logical values still remaining.
+            const size_t values_to_decode = std::min(run_len, remaining);
+            for (size_t i = 0; i < values_to_decode; ++i) {
                 uint32_t level = ReadPacked(i * static_cast<size_t>(bit_width));
                 if (level > static_cast<uint32_t>(max_def_level)) {
                     throw InvalidInputException("Invalid DATA_PAGE_V1 definition levels: decoded level exceeds max_def_level");
@@ -157,9 +172,12 @@ size_t CountPresentFromDefinitionLevelsV1(tcb::span<const uint8_t> def_payload,
                     ++present_count;
                 }
             }
-            decoded_values += run_len;
+            decoded_values += values_to_decode;
         }
     }
+    if (def_offset != def_payload.size()) {
+        throw InvalidInputException("Invalid DATA_PAGE_V1 definition levels: trailing bytes after decoding");
+    }
     return present_count;
 }
 

src/processing/parquet_utils_test.cpp

@@ -279,10 +279,21 @@ TEST(ParquetUtils, CountPresentFromDefinitionLevelsV1_BitPackedCanonical_0to7_88
     // Canonical example from Parquet Encodings.md:
     // values 0..7 with bit_width=3 are packed as bytes 0x88, 0xC6, 0xFA.
     // Bit-packed header for one group of 8 values is varint((1 << 1) | 1) = 0x03.
+    //
+    // Important detail: decoder bit_width is derived from max_def_level at runtime.
+    // So this payload is valid only when max_def_level implies bit_width=3 (e.g. 7).
+    // Reusing these same bytes with max_def_level=3 (bit_width=2) is a different
+    // encoding contract and can leave unread trailing bytes by design.
     std::vector<uint8_t> def_payload = {0x03, 0x88, 0xC6, 0xFA};
 
     EXPECT_EQ(1u, CountPresentFromDefinitionLevelsV1(def_payload, 8, 7));  // only value 7
-    EXPECT_EQ(1u, CountPresentFromDefinitionLevelsV1(def_payload, 8, 3));  // only value 3
+
+    // Separate bit_width=2 payload for max_def_level=3.
+    // This keeps payload bit-width consistent with decoder configuration and
+    // avoids mixing a 3-bit packed stream with a 2-bit decode expectation.
+    std::vector<uint32_t> levels_bw2 = {0, 1, 2, 3, 0, 1, 2, 0};  // one value at level 3
+    auto def_payload_bw2 = MakeBitPackedDefPayload(levels_bw2, 2);
+    EXPECT_EQ(1u, CountPresentFromDefinitionLevelsV1(def_payload_bw2, 8, 3));
 }
 
 TEST(ParquetUtils, CountPresentFromDefinitionLevelsV1_ManualBytes_RleRunLen4_Level1) {
@@ -338,6 +349,31 @@ TEST(ParquetUtils, CountPresentFromDefinitionLevelsV1_RejectsZeroBitPackedGroups
     EXPECT_THROW(CountPresentFromDefinitionLevelsV1(def_payload, 8, 1), InvalidInputException);
 }
 
+TEST(ParquetUtils, CountPresentFromDefinitionLevelsV1_BitPackedFinalRunAllowsPadding) {
+    // Corner case: a final bit-packed run is encoded as a full 8-value group,
+    // while logical num_values ends mid-group. Decode only the logical values
+    // and ignore padded trailing values in the last group.
+    // Payload: header=0x03 (1 bit-packed group => 8 values), packed=0x07 (bits 1,1,1,0,0,0,0,0).
+    std::vector<uint8_t> def_payload = {0x03, 0x07};
+    EXPECT_EQ(3u, CountPresentFromDefinitionLevelsV1(def_payload, 3, 1));
+}
+
+TEST(ParquetUtils, CountPresentFromDefinitionLevelsV1_RejectsTrailingBytesAfterDecoding) {
+    // One full bit-packed group (8 values) plus extra trailing byte that must be rejected.
+    std::vector<uint8_t> def_payload = {0x03, 0xAA, 0xFF};
+    EXPECT_THROW(CountPresentFromDefinitionLevelsV1(def_payload, 8, 1), InvalidInputException);
+}
+
+TEST(ParquetUtils, CountPresentFromDefinitionLevelsV1_RejectsNonPositiveMaxDefLevel) {
+    auto def_payload = MakeRleDefPayload(1, 0, 1);
+    EXPECT_THROW(CountPresentFromDefinitionLevelsV1(def_payload, 1, 0), InvalidInputException);
+}
+
+TEST(ParquetUtils, CountPresentFromDefinitionLevelsV1_RejectsNegativeNumValues) {
+    auto def_payload = MakeRleDefPayload(1, 1, 1);
+    EXPECT_THROW(CountPresentFromDefinitionLevelsV1(def_payload, -1, 1), InvalidInputException);
+}
+
 // -----------------------------------------------------------------------------
 // Tests for DecompressAndSplit function.
 // -----------------------------------------------------------------------------