diff --git a/parquet/pqarrow/encode_arrow.go b/parquet/pqarrow/encode_arrow.go index 73b52a34..d866a801 100644 --- a/parquet/pqarrow/encode_arrow.go +++ b/parquet/pqarrow/encode_arrow.go @@ -35,7 +35,6 @@ import ( "github.com/apache/arrow-go/v18/internal/utils" "github.com/apache/arrow-go/v18/parquet" "github.com/apache/arrow-go/v18/parquet/file" - "github.com/apache/arrow-go/v18/parquet/internal/debug" ) // get the count of the number of leaf arrays for the type @@ -350,15 +349,21 @@ func writeDenseArrow(ctx *arrowWriteContext, cw file.ColumnChunkWriter, leafArr data[idx] = int32(val / 86400000) // coerce date64 values } case arrow.DECIMAL128: - for idx, val := range leafArr.(*array.Decimal128).Values() { - debug.Assert(val.HighBits() == 0 || val.HighBits() == -1, "casting Decimal128 greater than the value range; high bits must be 0 or -1") - debug.Assert(int64(val.LowBits()) <= math.MaxUint32, "casting Decimal128 to int32 when value > MaxUint32") + decimalArr := leafArr.(*array.Decimal128) + precision := min(decimalArr.DataType().(arrow.DecimalType).GetPrecision(), int32(9)) + for idx, val := range decimalArr.Values() { + if decimalArr.IsValid(idx) && (precision <= 0 || !val.FitsInPrecision(precision)) { + return fmt.Errorf("%w: decimal value at index %d does not fit precision %d", arrow.ErrInvalid, idx, precision) + } data[idx] = int32(val.LowBits()) } case arrow.DECIMAL256: - for idx, val := range leafArr.(*array.Decimal256).Values() { - debug.Assert(val.Array()[3] == 0 || val.Array()[3] == 0xFFFFFFFF, "casting Decimal128 greater than the value range; high bits must be 0 or -1") - debug.Assert(val.LowBits() <= math.MaxUint32, "casting Decimal128 to int32 when value > MaxUint32") + decimalArr := leafArr.(*array.Decimal256) + precision := min(decimalArr.DataType().(arrow.DecimalType).GetPrecision(), int32(9)) + for idx, val := range decimalArr.Values() { + if decimalArr.IsValid(idx) && (precision <= 0 || !val.FitsInPrecision(precision)) { + return fmt.Errorf("%w: decimal value at index %d does not fit precision %d", arrow.ErrInvalid, idx, precision) + } data[idx] = int32(val.LowBits()) } default: @@ -433,15 +438,23 @@ func writeDenseArrow(ctx *arrowWriteContext, cw file.ColumnChunkWriter, leafArr case arrow.DECIMAL128: ctx.dataBuffer.ResizeNoShrink(arrow.Int64Traits.BytesRequired(leafArr.Len())) data = arrow.Int64Traits.CastFromBytes(ctx.dataBuffer.Bytes()) - for idx, val := range leafArr.(*array.Decimal128).Values() { - debug.Assert(val.HighBits() == 0 || val.HighBits() == -1, "trying to cast Decimal128 to int64 greater than range, high bits must be 0 or -1") + decimalArr := leafArr.(*array.Decimal128) + precision := min(decimalArr.DataType().(arrow.DecimalType).GetPrecision(), int32(18)) + for idx, val := range decimalArr.Values() { + if decimalArr.IsValid(idx) && (precision <= 0 || !val.FitsInPrecision(precision)) { + return fmt.Errorf("%w: decimal value at index %d does not fit precision %d", arrow.ErrInvalid, idx, precision) + } data[idx] = int64(val.LowBits()) } case arrow.DECIMAL256: ctx.dataBuffer.ResizeNoShrink(arrow.Int64Traits.BytesRequired(leafArr.Len())) data = arrow.Int64Traits.CastFromBytes(ctx.dataBuffer.Bytes()) - for idx, val := range leafArr.(*array.Decimal256).Values() { - debug.Assert(val.Array()[3] == 0 || val.Array()[3] == 0xFFFFFFFF, "trying to cast Decimal128 to int64 greater than range, high bits must be 0 or -1") + decimalArr := leafArr.(*array.Decimal256) + precision := min(decimalArr.DataType().(arrow.DecimalType).GetPrecision(), int32(18)) + for idx, val := range decimalArr.Values() { + if decimalArr.IsValid(idx) && (precision <= 0 || !val.FitsInPrecision(precision)) { + return fmt.Errorf("%w: decimal value at index %d does not fit precision %d", arrow.ErrInvalid, idx, precision) + } data[idx] = int64(val.LowBits()) } default: diff --git a/parquet/pqarrow/encode_arrow_test.go b/parquet/pqarrow/encode_arrow_test.go index 6c9e979d..77b7ca91 100644 --- a/parquet/pqarrow/encode_arrow_test.go +++ b/parquet/pqarrow/encode_arrow_test.go @@ -309,6 +309,60 @@ func TestWriteArrowCols(t *testing.T) { } } +func TestWriteDecimalRejectsValuesOutsidePrecision(t *testing.T) { + tests := []struct { + name string + dtype arrow.DecimalType + value string + valuePrec int32 + }{ + {"decimal128_int32", &arrow.Decimal128Type{Precision: 9}, "3000000000", 10}, + {"decimal128_int64", &arrow.Decimal128Type{Precision: 18}, "10000000000000000000", 20}, + {"decimal256_int32", &arrow.Decimal256Type{Precision: 9}, "3000000000", 10}, + {"decimal256_int64", &arrow.Decimal256Type{Precision: 18}, "10000000000000000000", 20}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + var values arrow.Array + switch dtype := tt.dtype.(type) { + case *arrow.Decimal128Type: + value, err := decimal128.FromString(tt.value, tt.valuePrec, 0) + require.NoError(t, err) + builder := array.NewDecimal128Builder(mem, dtype) + builder.Append(value) + values = builder.NewDecimal128Array() + builder.Release() + case *arrow.Decimal256Type: + value, err := decimal256.FromString(tt.value, tt.valuePrec, 0) + require.NoError(t, err) + builder := array.NewDecimal256Builder(mem, dtype) + builder.Append(value) + values = builder.NewDecimal256Array() + builder.Release() + } + defer values.Release() + + sc := arrow.NewSchema([]arrow.Field{{Name: "decimal", Type: tt.dtype}}, nil) + rec := array.NewRecordBatch(sc, []arrow.Array{values}, int64(values.Len())) + defer rec.Release() + + var sink bytes.Buffer + writer, err := pqarrow.NewFileWriter(sc, &sink, + parquet.NewWriterProperties(parquet.WithStoreDecimalAsInteger(true)), + pqarrow.NewArrowWriterProperties(pqarrow.WithAllocator(mem))) + require.NoError(t, err) + err = writer.Write(rec) + require.ErrorIs(t, err, arrow.ErrInvalid) + assert.ErrorContains(t, err, "decimal value at index 0 does not fit precision") + require.NoError(t, writer.Close()) + }) + } +} + func TestWriteArrowInt96(t *testing.T) { mem := memory.NewCheckedAllocator(memory.DefaultAllocator) defer mem.AssertSize(t, 0)