lightning: fix corrupted DECIMAL values in dictionary-encoded parquet files (#70436)
Project: pingcap/tidb | Language: Go | Concepts: memory-management, error-handling, testing-strategy | Difficulty: mid
Link: https://github.com/pingcap/tidb/pull/70461
The problem
When importing Parquet files with dictionary-encoded DECIMAL columns (stored as BYTE_ARRAY or FIXED_LEN_BYTE_ARRAY), repeated negative values could be silently corrupted. Only the first occurrence would decode correctly; subsequent rows referencing the same dictionary entry would get incorrect results (e.g., negative values turning positive). This happened because the decoding logic mutated a shared buffer, affecting all rows that referenced the same dictionary entry.
Why it's non-trivial
A naive fix might simply decode values as usual, not realizing that dictionary decoding in the Parquet reader reuses the same backing buffer for all rows referring to a given dictionary entry. If the decoding function mutates this buffer (as it does for negative values), all subsequent reads of that dictionary entry are corrupted. The fix must ensure that no mutation happens to shared buffers, without introducing unnecessary allocations for every row.
The approach
1. Introduce a safe buffer copy in the decimal setter
The core fix is to ensure that before converting a byte array to a decimal string, the code copies the value into a private buffer, so any mutation doesn't affect shared data:
+// getDecimalByteSetter returns a setter for byte array DECIMAL columns.
+// Dictionary decoded values alias the shared dictionary buffer, and
+// binaryToDecimalStr negates its input in place, so copy it into our own
+// buffer first. The buffer is reused across rows, which is safe because
+// binaryToDecimalStr never retains it.
+func getDecimalByteSetter[T parquet.ByteArray | parquet.FixedLenByteArray](scale int) setter[T] {
+ var buf []byte
+ return func(val T, d *types.Datum) {
+ buf = append(buf[:0], val...)
+ d.SetString(binaryToDecimalStr(buf, scale), "utf8mb4_bin")
+ }
+}
This function creates a reusable buffer per column. Each time a value is set, it copies the bytes into this buffer before passing it to the conversion function, ensuring no shared dictionary buffer is mutated.
2. Use the safe setter for both BYTE_ARRAY and FIXED_LEN_BYTE_ARRAY decimal columns
- return func(val parquet.ByteArray, d *types.Datum) {
- str := binaryToDecimalStr(val, int(converted.decimalMeta.Scale))
- d.SetString(str, "utf8mb4_bin")
- }
+ return getDecimalByteSetter[parquet.ByteArray](int(converted.decimalMeta.Scale))
and
- return func(val parquet.FixedLenByteArray, d *types.Datum) {
- str := binaryToDecimalStr(val, int(converted.decimalMeta.Scale))
- d.SetString(str, "utf8mb4_bin")
- }
+ return getDecimalByteSetter[parquet.FixedLenByteArray](int(converted.decimalMeta.Scale))
Both conversion functions now use the safe buffer-copying setter, preventing buffer corruption for both variable and fixed-length byte arrays.
3. Add regression tests for repeated dictionary-encoded decimals
A new test case is added to ensure that repeated negative decimal values are correctly decoded even when dictionary encoding is used:
+ // Enough repeated rows that the writer keeps the dictionary page instead of
+ // falling back to PLAIN, so every row decodes to the same shared buffer.
+ const dictRows = 256
+ repeatedByteArray := func(val []byte) func(int) (any, []int16) {
+ return func(numRows int) (any, []int16) {
+ values := make([]parquet.ByteArray, numRows)
+ levels := make([]int16, numRows)
+ for i := range numRows {
+ values[i] = val
+ levels[i] = 1
+ }
+ return values, levels
+ }
+ }
+ repeatedFixedLenByteArray := func(val []byte) func(int) (any, []int16) {
+ return func(numRows int) (any, []int16) {
+ values := make([]parquet.FixedLenByteArray, numRows)
+ levels := make([]int16, numRows)
+ for i := range numRows {
+ values[i] = val
+ levels[i] = 1
+ }
+ return values, levels
+ }
+ }
+
+ pc = []ParquetColumn{
+ {
+ Name: "positive",
+ Type: parquet.Types.ByteArray,
+ Converted: schema.ConvertedTypes.Decimal,
+ Precision: 9,
+ Scale: 2,
+ Gen: repeatedByteArray([]byte{0x30, 0x39}), // 12345
+ },
+ {
+ // binaryToDecimalStr only rewrites its input for negative values,
+ // so this is the column that catches the shared-buffer aliasing.
+ Name: "negative",
+ Type: parquet.Types.ByteArray,
+ Converted: schema.ConvertedTypes.Decimal,
+ Precision: 9,
+ Scale: 2,
+ Gen: repeatedByteArray([]byte{0xcf, 0xc7}), // -12345
+ },
+ {
+ Name: "fixed_len_negative",
+ Type: parquet.Types.FixedLenByteArray,
+ Converted: schema.ConvertedTypes.Decimal,
+ TypeLen: 4,
+ Precision: 9,
+ Scale: 2,
+ Gen: repeatedFixedLenByteArray([]byte{0xff, 0xff, 0xcf, 0xc7}), // -12345
+ },
+ }
This test creates columns with repeated positive and negative decimal values, ensuring that dictionary encoding is used and that negative values are not corrupted.
4. Add support for writing fixed-length byte array columns in test utilities
+ case *file.FixedLenByteArrayColumnChunkWriter:
+ buf, _ := vals.([]parquet.FixedLenByteArray)
+ _, err = w.WriteBatch(buf, defLevel, nil)
This enables the test utility to write fixed-length byte array columns, which is needed for the regression test.
From the review thread
No notable review discussion on this PR.
Takeaways
-
Never mutate shared buffers returned by external libraries unless you own them.
When working with data structures like dictionary-encoded arrays, always assume the underlying storage may be shared and immutable from your perspective. -
Use per-column reusable buffers to avoid per-row allocations in high-throughput code.
Allocating a buffer once per column and reusing it for each row is both safe and efficient, compared to allocating a new buffer for every value. -
Regression tests should reproduce the exact conditions that triggered the bug, including encoding modes and data patterns.
Tests that specifically force dictionary encoding and repeated values help ensure the fix is robust and prevent silent regressions.