diff --git a/pkg/lightning/mydump/parquet_parser_test.go b/pkg/lightning/mydump/parquet_parser_test.go index 1dad28130ddc5..d66721de02a08 100644 --- a/pkg/lightning/mydump/parquet_parser_test.go +++ b/pkg/lightning/mydump/parquet_parser_test.go @@ -337,6 +337,86 @@ func TestParquetVariousTypes(t *testing.T) { assert.NoError(t, reader.ReadRow()) assert.Equal(t, types.KindUint64, reader.lastRow.Row[0].Kind()) assert.Equal(t, uint64(1), reader.lastRow.Row[0].GetValue()) + + // 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 + }, + } + + fileName = "test.decimal.dictionary.parquet" + require.NoError(t, WriteParquetFile(dir, fileName, pc, dictRows)) + + r, err = store.Open(context.TODO(), fileName, nil) + require.NoError(t, err) + reader, err = NewParquetParser(context.TODO(), store, r, fileName, ParquetFileMeta{}) + require.NoError(t, err) + defer reader.Close() + + for i := range pc { + cc, err := reader.readers[0].MetaData().RowGroup(0).ColumnChunk(i) + require.NoError(t, err) + require.True(t, cc.HasDictionaryPage(), "column %s is not dictionary encoded", pc[i].Name) + } + + expectedValues := []string{"123.45", "-123.45", "-123.45"} + for range dictRows { + require.NoError(t, reader.ReadRow()) + require.Len(t, reader.lastRow.Row, len(expectedValues)) + for i, expectValue := range expectedValues { + require.Equal(t, expectValue, reader.lastRow.Row[i].GetString()) + } + } } func TestParquetAurora(t *testing.T) { diff --git a/pkg/lightning/mydump/parquet_type_converter.go b/pkg/lightning/mydump/parquet_type_converter.go index ea885cb5e4f90..f043ab4630085 100644 --- a/pkg/lightning/mydump/parquet_type_converter.go +++ b/pkg/lightning/mydump/parquet_type_converter.go @@ -224,6 +224,19 @@ func setFloat64Data(val float64, d *types.Datum) { d.SetFloat64(val) } +// 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") + } +} + func getByteArraySetter(converted *convertedType) setter[parquet.ByteArray] { switch converted.converted { case schema.ConvertedTypes.None, schema.ConvertedTypes.BSON, schema.ConvertedTypes.JSON, schema.ConvertedTypes.UTF8, schema.ConvertedTypes.Enum: @@ -232,10 +245,7 @@ func getByteArraySetter(converted *convertedType) setter[parquet.ByteArray] { d.SetBytesAsString(val, "utf8mb4_bin", 0) } case schema.ConvertedTypes.Decimal: - 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)) } return nil @@ -249,10 +259,7 @@ func getFixedLenByteArraySetter(converted *convertedType) setter[parquet.FixedLe d.SetBytesAsString(val, "utf8mb4_bin", 0) } case schema.ConvertedTypes.Decimal: - 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)) } return nil diff --git a/pkg/lightning/mydump/parquet_writer.go b/pkg/lightning/mydump/parquet_writer.go index 0fbf53d42f18e..195297534ff6d 100644 --- a/pkg/lightning/mydump/parquet_writer.go +++ b/pkg/lightning/mydump/parquet_writer.go @@ -135,6 +135,9 @@ func WriteParquetFile(path, fileName string, pcolumns []ParquetColumn, rows int, case *file.ByteArrayColumnChunkWriter: buf, _ := vals.([]parquet.ByteArray) _, err = w.WriteBatch(buf, defLevel, nil) + case *file.FixedLenByteArrayColumnChunkWriter: + buf, _ := vals.([]parquet.FixedLenByteArray) + _, err = w.WriteBatch(buf, defLevel, nil) case *file.Int32ColumnChunkWriter: buf, _ := vals.([]int32) _, err = w.WriteBatch(buf, defLevel, nil)