diff --git a/parquet/metadata/statistics_test.go b/parquet/metadata/statistics_test.go index a0d107fa..6809eeef 100644 --- a/parquet/metadata/statistics_test.go +++ b/parquet/metadata/statistics_test.go @@ -21,6 +21,8 @@ import ( "reflect" "testing" + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" "github.com/apache/arrow-go/v18/arrow/bitutil" "github.com/apache/arrow-go/v18/arrow/float16" "github.com/apache/arrow-go/v18/arrow/memory" @@ -640,3 +642,183 @@ func TestNewStatisticsDistinctCountUnset(t *testing.T) { }) } } + +// TestByteArrayStatisticsDoesNotAliasInput is a regression test for +// adbc-drivers/bigquery#229: ByteArrayStatistics stored the min/max as slices +// aliasing the caller's value buffer (e.g. an Arrow array value buffer) instead +// of copying the bytes. When a streaming writer released a batch before the +// statistics were serialized during Close, the lazily-accessed min/max pointed +// at freed memory and bytes.Compare in less() segfaulted. SetMinMax must copy +// the min/max into statistics-owned memory so they remain valid after the +// source buffer is released. +func TestByteArrayStatisticsDoesNotAliasInput(t *testing.T) { + descr := schema.NewColumn(schema.NewByteArrayNode("ba", parquet.Repetitions.Required, -1), 0, 0) + + t.Run("Update copies min/max", func(t *testing.T) { + stats := metadata.NewStatistics(descr, memory.DefaultAllocator).(*metadata.ByteArrayStatistics) + + buf := []byte("mmmmm") + stats.Update([]parquet.ByteArray{parquet.ByteArray(buf)}, 0) + require.True(t, stats.HasMinMax()) + require.Equal(t, "mmmmm", string(stats.Min())) + require.Equal(t, "mmmmm", string(stats.Max())) + + // Simulate the source buffer being freed/reused after the batch is written. + copy(buf, "zzzzz") + assert.Equal(t, "mmmmm", string(stats.Min()), "min must not alias the input buffer") + assert.Equal(t, "mmmmm", string(stats.Max()), "max must not alias the input buffer") + }) + + t.Run("stored min/max survive freeing a prior batch", func(t *testing.T) { + // Mirrors the crash flow directly: batch N establishes min/max, its buffer + // is released, then batch N+1's SetMinMax compares new values against the + // stored (previously dangling) min/max via less(). + stats := metadata.NewStatistics(descr, memory.DefaultAllocator).(*metadata.ByteArrayStatistics) + + batch1 := []byte("aaazzz") + stats.UpdateSpaced([]parquet.ByteArray{batch1[0:3], batch1[3:6]}, []byte{0x3}, 0, 0) + require.Equal(t, "aaa", string(stats.Min())) + require.Equal(t, "zzz", string(stats.Max())) + + // Batch 1's buffer is released and its memory reused for other data. + copy(batch1, "######") + + batch2 := []byte("mmm") + stats.UpdateSpaced([]parquet.ByteArray{parquet.ByteArray(batch2)}, []byte{0x1}, 0, 0) + + assert.Equal(t, "aaa", string(stats.Min()), "min from a released batch must survive") + assert.Equal(t, "zzz", string(stats.Max()), "max from a released batch must survive") + }) + + t.Run("Merge copies min/max", func(t *testing.T) { + src := metadata.NewStatistics(descr, memory.DefaultAllocator).(*metadata.ByteArrayStatistics) + src.Update([]parquet.ByteArray{parquet.ByteArray("kkkkk")}, 0) + + dst := metadata.NewStatistics(descr, memory.DefaultAllocator).(*metadata.ByteArrayStatistics) + dst.Merge(src) + require.Equal(t, "kkkkk", string(dst.Min())) + require.Equal(t, "kkkkk", string(dst.Max())) + + // Updating src overwrites its statistics-owned min/max buffers in place. If + // Merge had aliased src's buffers instead of copying, this would corrupt dst. + src.Update([]parquet.ByteArray{parquet.ByteArray("aaaaa"), parquet.ByteArray("zzzzz")}, 0) + require.Equal(t, "aaaaa", string(src.Min())) + require.Equal(t, "zzzzz", string(src.Max())) + assert.Equal(t, "kkkkk", string(dst.Min()), "merged min must be independent of src") + assert.Equal(t, "kkkkk", string(dst.Max()), "merged max must be independent of src") + }) + + t.Run("encoded min/max survive freeing the source", func(t *testing.T) { + stats := metadata.NewStatistics(descr, memory.DefaultAllocator).(*metadata.ByteArrayStatistics) + buf := []byte("hello") + stats.Update([]parquet.ByteArray{parquet.ByteArray(buf)}, 0) + + copy(buf, "world") + assert.Equal(t, []byte("hello"), stats.EncodeMin()) + assert.Equal(t, []byte("hello"), stats.EncodeMax()) + }) +} + +// TestFixedLenByteArrayStatisticsDoesNotAliasInput is the FIXED_LEN_BYTE_ARRAY +// counterpart to TestByteArrayStatisticsDoesNotAliasInput (adbc-drivers/bigquery#229). +func TestFixedLenByteArrayStatisticsDoesNotAliasInput(t *testing.T) { + n, err := schema.NewPrimitiveNode("flba", parquet.Repetitions.Required, parquet.Types.FixedLenByteArray, -1, 3) + require.NoError(t, err) + descr := schema.NewColumn(n, 0, 0) + + stats := metadata.NewStatistics(descr, memory.DefaultAllocator).(*metadata.FixedLenByteArrayStatistics) + + buf := []byte("mmm") + stats.Update([]parquet.FixedLenByteArray{parquet.FixedLenByteArray(buf)}, 0) + require.True(t, stats.HasMinMax()) + require.Equal(t, "mmm", string(stats.Min())) + require.Equal(t, "mmm", string(stats.Max())) + + // Simulate the source buffer being freed/reused after the batch is written. + copy(buf, "zzz") + assert.Equal(t, "mmm", string(stats.Min()), "min must not alias the input buffer") + assert.Equal(t, "mmm", string(stats.Max()), "max must not alias the input buffer") +} + +// TestFloat16StatisticsDoesNotAliasInput covers the FLOAT16 logical type, which +// is backed by parquet.FixedLenByteArray and shares the copy-on-store fix +// (adbc-drivers/bigquery#229). +func TestFloat16StatisticsDoesNotAliasInput(t *testing.T) { + descr := schema.NewColumn(newFloat16Node("f16", parquet.Repetitions.Required, -1), 0, 0) + stats := metadata.NewStatistics(descr, memory.DefaultAllocator).(*metadata.Float16Statistics) + + want := float16.New(1.5).ToLEBytes() + mutable := float16.New(1.5).ToLEBytes() + stats.Update([]parquet.FixedLenByteArray{parquet.FixedLenByteArray(mutable)}, 0) + require.True(t, stats.HasMinMax()) + require.Equal(t, []byte(want), []byte(stats.Min())) + + // Overwrite the source buffer to model it being freed/reused after the batch. + copy(mutable, float16.New(9.5).ToLEBytes()) + assert.Equal(t, []byte(want), []byte(stats.Min()), "min must not alias the input buffer") + assert.Equal(t, []byte(want), []byte(stats.Max()), "max must not alias the input buffer") +} + +// TestFixedLenByteArrayStatisticsUpdateFromArrowMax is a regression test for a +// separate pre-existing bug in FixedLenByteArray.UpdateFromArrow: it computed +// the running max against the running min instead of the running max, producing +// an incorrect maximum for FIXED_LEN_BYTE_ARRAY columns updated from Arrow. +func TestFixedLenByteArrayStatisticsUpdateFromArrowMax(t *testing.T) { + mem := memory.DefaultAllocator + n, err := schema.NewPrimitiveNode("flba", parquet.Repetitions.Required, parquet.Types.FixedLenByteArray, -1, 3) + require.NoError(t, err) + descr := schema.NewColumn(n, 0, 0) + stats := metadata.NewStatistics(descr, mem).(*metadata.FixedLenByteArrayStatistics) + + bldr := array.NewFixedSizeBinaryBuilder(mem, &arrow.FixedSizeBinaryType{ByteWidth: 3}) + defer bldr.Release() + for _, v := range []string{"bbb", "aaa", "ccc", "abc"} { + bldr.Append([]byte(v)) + } + arr := bldr.NewArray() + defer arr.Release() + + require.NoError(t, stats.UpdateFromArrow(arr, true)) + require.True(t, stats.HasMinMax()) + assert.Equal(t, "aaa", string(stats.Min())) + assert.Equal(t, "ccc", string(stats.Max()), "max must be computed against the running max, not the running min") +} + +type encodedStatProvider struct { + min, max []byte +} + +func (e *encodedStatProvider) GetMin() []byte { return e.min } +func (e *encodedStatProvider) GetMax() []byte { return e.max } +func (e *encodedStatProvider) GetNullCount() int64 { return 0 } +func (e *encodedStatProvider) GetDistinctCount() int64 { return 0 } +func (e *encodedStatProvider) IsSetMax() bool { return e.max != nil } +func (e *encodedStatProvider) IsSetMin() bool { return e.min != nil } +func (e *encodedStatProvider) IsSetNullCount() bool { return false } +func (e *encodedStatProvider) IsSetDistinctCount() bool { return false } + +// TestByteArrayStatisticsFromEncodedOwnsMinMax is a regression test for a review +// finding on the use-after-free fix: statistics built from encoded metadata via +// New*StatisticsFromEncoded stored min/max slices aliasing the caller's metadata +// buffer. Because SetMinMax now reuses those buffers on update, a later update +// could write through and corrupt the caller's metadata, so the decoded min/max +// must be copied into statistics-owned buffers. +func TestByteArrayStatisticsFromEncodedOwnsMinMax(t *testing.T) { + descr := schema.NewColumn(schema.NewByteArrayNode("ba", parquet.Repetitions.Required, -1), 0, 0) + + encMin := []byte("bbb") + encMax := []byte("yyy") + stats := metadata.NewStatisticsFromEncoded(descr, memory.DefaultAllocator, 4, + &encodedStatProvider{min: encMin, max: encMax}).(*metadata.ByteArrayStatistics) + require.True(t, stats.HasMinMax()) + require.Equal(t, "bbb", string(stats.Min())) + require.Equal(t, "yyy", string(stats.Max())) + + // Updating with new extremes must not write through into the encoded source + // buffers, which would happen if the stored min/max still aliased them. + stats.Update([]parquet.ByteArray{parquet.ByteArray("aaa"), parquet.ByteArray("zzz")}, 0) + assert.Equal(t, "bbb", string(encMin), "encoded min source buffer must not be overwritten") + assert.Equal(t, "yyy", string(encMax), "encoded max source buffer must not be overwritten") + assert.Equal(t, "aaa", string(stats.Min())) + assert.Equal(t, "zzz", string(stats.Max())) +} diff --git a/parquet/metadata/statistics_types.gen.go b/parquet/metadata/statistics_types.gen.go index a9a19ea1..e450fb7c 100644 --- a/parquet/metadata/statistics_types.gen.go +++ b/parquet/metadata/statistics_types.gen.go @@ -1899,11 +1899,13 @@ func NewByteArrayStatisticsFromEncoded(descr *schema.Column, mem memory.Allocato encodedMin := encoded.GetMin() if encodedMin != nil && len(encodedMin) > 0 { - ret.min = ret.plainDecode(encodedMin) + // Copy into a statistics-owned buffer so the stored min does not alias the + // encoded metadata buffer; SetMinMax reuses this buffer on later updates. + ret.min = append(ret.min[:0], ret.plainDecode(encodedMin)...) } encodedMax := encoded.GetMax() if encodedMax != nil && len(encodedMax) > 0 { - ret.max = ret.plainDecode(encodedMax) + ret.max = append(ret.max[:0], ret.plainDecode(encodedMax)...) } ret.hasMinMax = encoded.IsSetMax() || encoded.IsSetMin() return ret @@ -2007,7 +2009,18 @@ func (s *ByteArrayStatistics) getMinMaxSpaced(values []parquet.ByteArray, validB return } +// Min returns the current minimum value. +// +// The returned slice references a statistics-owned buffer that is reused on +// subsequent updates, so it is only valid until the next update that replaces +// the minimum. Copy it if you need to retain it. func (s *ByteArrayStatistics) Min() parquet.ByteArray { return s.min } + +// Max returns the current maximum value. +// +// The returned slice references a statistics-owned buffer that is reused on +// subsequent updates, so it is only valid until the next update that replaces +// the maximum. Copy it if you need to retain it. func (s *ByteArrayStatistics) Max() parquet.ByteArray { return s.max } // Merge merges the stats from other into this stat object, updating @@ -2093,6 +2106,11 @@ func (s *ByteArrayStatistics) UpdateFromArrow(values arrow.Array, updateCounts b // SetMinMax updates the min and max values only if they are not currently set // or if argMin is less than the current min / argMax is greater than the current max +// +// The min and max are copied into statistics-owned buffers rather than retaining +// the provided slices. Those slices typically alias the source column data (e.g. +// an Arrow value buffer) which may be released before the statistics are +// serialized during Close, so aliasing them would leave dangling references. func (s *ByteArrayStatistics) SetMinMax(argMin, argMax parquet.ByteArray) { maybeMinMax := s.cleanStat([2]parquet.ByteArray{argMin, argMax}) if maybeMinMax == nil { @@ -2104,14 +2122,14 @@ func (s *ByteArrayStatistics) SetMinMax(argMin, argMax parquet.ByteArray) { if !s.hasMinMax { s.hasMinMax = true - s.min = min - s.max = max + s.min = append(s.min[:0], min...) + s.max = append(s.max[:0], max...) } else { if !s.less(s.min, min) { - s.min = min + s.min = append(s.min[:0], min...) } if s.less(s.max, max) { - s.max = max + s.max = append(s.max[:0], max...) } } } @@ -2203,11 +2221,13 @@ func NewFixedLenByteArrayStatisticsFromEncoded(descr *schema.Column, mem memory. encodedMin := encoded.GetMin() if encodedMin != nil && len(encodedMin) > 0 { - ret.min = ret.plainDecode(encodedMin) + // Copy into a statistics-owned buffer so the stored min does not alias the + // encoded metadata buffer; SetMinMax reuses this buffer on later updates. + ret.min = append(ret.min[:0], ret.plainDecode(encodedMin)...) } encodedMax := encoded.GetMax() if encodedMax != nil && len(encodedMax) > 0 { - ret.max = ret.plainDecode(encodedMax) + ret.max = append(ret.max[:0], ret.plainDecode(encodedMax)...) } ret.hasMinMax = encoded.IsSetMax() || encoded.IsSetMin() return ret @@ -2323,7 +2343,18 @@ func (s *FixedLenByteArrayStatistics) getMinMaxSpaced(values []parquet.FixedLenB return } +// Min returns the current minimum value. +// +// The returned slice references a statistics-owned buffer that is reused on +// subsequent updates, so it is only valid until the next update that replaces +// the minimum. Copy it if you need to retain it. func (s *FixedLenByteArrayStatistics) Min() parquet.FixedLenByteArray { return s.min } + +// Max returns the current maximum value. +// +// The returned slice references a statistics-owned buffer that is reused on +// subsequent updates, so it is only valid until the next update that replaces +// the maximum. Copy it if you need to retain it. func (s *FixedLenByteArrayStatistics) Max() parquet.FixedLenByteArray { return s.max } // Merge merges the stats from other into this stat object, updating @@ -2394,7 +2425,7 @@ func (s *FixedLenByteArrayStatistics) UpdateFromArrow(values arrow.Array, update for i := 0; i < values.Len(); i++ { v := data[i*width : (i+1)*width] min = s.minval(min, v) - max = s.maxval(min, v) + max = s.maxval(max, v) } s.SetMinMax(min, max) @@ -2403,6 +2434,11 @@ func (s *FixedLenByteArrayStatistics) UpdateFromArrow(values arrow.Array, update // SetMinMax updates the min and max values only if they are not currently set // or if argMin is less than the current min / argMax is greater than the current max +// +// The min and max are copied into statistics-owned buffers rather than retaining +// the provided slices. Those slices typically alias the source column data (e.g. +// an Arrow value buffer) which may be released before the statistics are +// serialized during Close, so aliasing them would leave dangling references. func (s *FixedLenByteArrayStatistics) SetMinMax(argMin, argMax parquet.FixedLenByteArray) { maybeMinMax := s.cleanStat([2]parquet.FixedLenByteArray{argMin, argMax}) if maybeMinMax == nil { @@ -2414,14 +2450,14 @@ func (s *FixedLenByteArrayStatistics) SetMinMax(argMin, argMax parquet.FixedLenB if !s.hasMinMax { s.hasMinMax = true - s.min = min - s.max = max + s.min = append(s.min[:0], min...) + s.max = append(s.max[:0], max...) } else { if !s.less(s.min, min) { - s.min = min + s.min = append(s.min[:0], min...) } if s.less(s.max, max) { - s.max = max + s.max = append(s.max[:0], max...) } } } @@ -2517,11 +2553,13 @@ func NewFloat16StatisticsFromEncoded(descr *schema.Column, mem memory.Allocator, encodedMin := encoded.GetMin() if encodedMin != nil && len(encodedMin) > 0 { - ret.min = ret.plainDecode(encodedMin) + // Copy into a statistics-owned buffer so the stored min does not alias the + // encoded metadata buffer; SetMinMax reuses this buffer on later updates. + ret.min = append(ret.min[:0], ret.plainDecode(encodedMin)...) } encodedMax := encoded.GetMax() if encodedMax != nil && len(encodedMax) > 0 { - ret.max = ret.plainDecode(encodedMax) + ret.max = append(ret.max[:0], ret.plainDecode(encodedMax)...) } ret.hasMinMax = encoded.IsSetMax() || encoded.IsSetMin() return ret @@ -2644,7 +2682,18 @@ func (s *Float16Statistics) getMinMaxSpaced(values []parquet.FixedLenByteArray, return } +// Min returns the current minimum value. +// +// The returned slice references a statistics-owned buffer that is reused on +// subsequent updates, so it is only valid until the next update that replaces +// the minimum. Copy it if you need to retain it. func (s *Float16Statistics) Min() parquet.FixedLenByteArray { return s.min } + +// Max returns the current maximum value. +// +// The returned slice references a statistics-owned buffer that is reused on +// subsequent updates, so it is only valid until the next update that replaces +// the maximum. Copy it if you need to retain it. func (s *Float16Statistics) Max() parquet.FixedLenByteArray { return s.max } // Merge merges the stats from other into this stat object, updating @@ -2704,6 +2753,11 @@ func (s *Float16Statistics) UpdateFromArrow(values arrow.Array, updateCounts boo // SetMinMax updates the min and max values only if they are not currently set // or if argMin is less than the current min / argMax is greater than the current max +// +// The min and max are copied into statistics-owned buffers rather than retaining +// the provided slices. Those slices typically alias the source column data (e.g. +// an Arrow value buffer) which may be released before the statistics are +// serialized during Close, so aliasing them would leave dangling references. func (s *Float16Statistics) SetMinMax(argMin, argMax parquet.FixedLenByteArray) { maybeMinMax := s.cleanStat([2]parquet.FixedLenByteArray{argMin, argMax}) if maybeMinMax == nil { @@ -2715,14 +2769,14 @@ func (s *Float16Statistics) SetMinMax(argMin, argMax parquet.FixedLenByteArray) if !s.hasMinMax { s.hasMinMax = true - s.min = min - s.max = max + s.min = append(s.min[:0], min...) + s.max = append(s.max[:0], max...) } else { if !s.less(s.min, min) { - s.min = min + s.min = append(s.min[:0], min...) } if s.less(s.max, max) { - s.max = max + s.max = append(s.max[:0], max...) } } } diff --git a/parquet/metadata/statistics_types.gen.go.tmpl b/parquet/metadata/statistics_types.gen.go.tmpl index 76657943..0f0bc377 100644 --- a/parquet/metadata/statistics_types.gen.go.tmpl +++ b/parquet/metadata/statistics_types.gen.go.tmpl @@ -87,11 +87,21 @@ func New{{.Name}}StatisticsFromEncoded(descr *schema.Column, mem memory.Allocato encodedMin := encoded.GetMin() if encodedMin != nil && len(encodedMin) > 0 { +{{- if or (eq .name "parquet.ByteArray") (eq .name "parquet.FixedLenByteArray")}} + // Copy into a statistics-owned buffer so the stored min does not alias the + // encoded metadata buffer; SetMinMax reuses this buffer on later updates. + ret.min = append(ret.min[:0], ret.plainDecode(encodedMin)...) +{{- else}} ret.min = ret.plainDecode(encodedMin) +{{- end}} } encodedMax := encoded.GetMax() if encodedMax != nil && len(encodedMax) > 0 { +{{- if or (eq .name "parquet.ByteArray") (eq .name "parquet.FixedLenByteArray")}} + ret.max = append(ret.max[:0], ret.plainDecode(encodedMax)...) +{{- else}} ret.max = ret.plainDecode(encodedMax) +{{- end}} } ret.hasMinMax = encoded.IsSetMax() || encoded.IsSetMin() return ret @@ -299,8 +309,24 @@ func (s *{{.Name}}Statistics) getMinMaxSpaced(values []{{.name}}, validBits []by return } +{{if or (eq .name "parquet.ByteArray") (eq .name "parquet.FixedLenByteArray")}} +// Min returns the current minimum value. +// +// The returned slice references a statistics-owned buffer that is reused on +// subsequent updates, so it is only valid until the next update that replaces +// the minimum. Copy it if you need to retain it. +func (s *{{.Name}}Statistics) Min() {{.name}} { return s.min } + +// Max returns the current maximum value. +// +// The returned slice references a statistics-owned buffer that is reused on +// subsequent updates, so it is only valid until the next update that replaces +// the maximum. Copy it if you need to retain it. +func (s *{{.Name}}Statistics) Max() {{.name}} { return s.max } +{{else}} func (s *{{.Name}}Statistics) Min() {{.name}} { return s.min } func (s *{{.Name}}Statistics) Max() {{.name}} { return s.max } +{{end}} // Merge merges the stats from other into this stat object, updating // the null count, distinct count, number of values and the min/max if @@ -370,7 +396,7 @@ func (s *{{.Name}}Statistics) UpdateFromArrow(values arrow.Array, updateCounts b for i := 0; i < values.Len(); i++ { v := data[i * width : (i+1) * width] min = s.minval(min, v) - max = s.maxval(min, v) + max = s.maxval(max, v) } s.SetMinMax(min, max) @@ -501,6 +527,13 @@ func (s *{{.Name}}Statistics) getMinMaxFromBitmap(bitmap []byte, bitmapOffset in // SetMinMax updates the min and max values only if they are not currently set // or if argMin is less than the current min / argMax is greater than the current max +{{- if or (eq .name "parquet.ByteArray") (eq .name "parquet.FixedLenByteArray")}} +// +// The min and max are copied into statistics-owned buffers rather than retaining +// the provided slices. Those slices typically alias the source column data (e.g. +// an Arrow value buffer) which may be released before the statistics are +// serialized during Close, so aliasing them would leave dangling references. +{{- end}} func (s *{{.Name}}Statistics) SetMinMax(argMin, argMax {{.name}}) { maybeMinMax := s.cleanStat([2]{{.name}}{argMin, argMax}) if maybeMinMax == nil { @@ -512,14 +545,27 @@ func (s *{{.Name}}Statistics) SetMinMax(argMin, argMax {{.name}}) { if !s.hasMinMax { s.hasMinMax = true +{{- if or (eq .name "parquet.ByteArray") (eq .name "parquet.FixedLenByteArray")}} + s.min = append(s.min[:0], min...) + s.max = append(s.max[:0], max...) +{{- else}} s.min = min s.max = max +{{- end}} } else { if !s.less(s.min, min) { +{{- if or (eq .name "parquet.ByteArray") (eq .name "parquet.FixedLenByteArray")}} + s.min = append(s.min[:0], min...) +{{- else}} s.min = min +{{- end}} } if s.less(s.max, max) { +{{- if or (eq .name "parquet.ByteArray") (eq .name "parquet.FixedLenByteArray")}} + s.max = append(s.max[:0], max...) +{{- else}} s.max = max +{{- end}} } } } diff --git a/parquet/pqarrow/encode_arrow_test.go b/parquet/pqarrow/encode_arrow_test.go index 95a5b3e5..197f6dcb 100644 --- a/parquet/pqarrow/encode_arrow_test.go +++ b/parquet/pqarrow/encode_arrow_test.go @@ -2780,3 +2780,70 @@ func TestReadWriteShreddedVariant(t *testing.T) { assert.Truef(t, array.Equal(arr, tbl.Column(0).Data().Chunk(0)), "expected: %s\ngot: %s", arr, tbl.Column(0).Data().Chunk(0)) } + +// poisonOnFreeAllocator overwrites every buffer with 0xFF as it is freed so that +// any lingering alias into released memory becomes observable. This gives a +// deterministic reproduction of the adbc-drivers/bigquery#229 use-after-free +// instead of depending on the Go runtime to reuse the freed pages. +type poisonOnFreeAllocator struct { + memory.Allocator +} + +func (p poisonOnFreeAllocator) Free(b []byte) { + for i := range b { + b[i] = 0xff + } + p.Allocator.Free(b) +} + +// TestByteArrayStatisticsStreamingReleaseBetweenBatches reproduces +// adbc-drivers/bigquery#229 end to end. With statistics enabled (the default) +// and dictionary encoding disabled (so the dense Arrow write path runs), a +// streaming writer releases each batch before writing the next. The ByteArray +// column statistics must copy their min/max, so the round-tripped values stay +// correct even though batch 1's buffer was poisoned when it was released before +// batch 2's write compared against the stored min/max. +func TestByteArrayStatisticsStreamingReleaseBetweenBatches(t *testing.T) { + checked := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer checked.AssertSize(t, 0) + mem := poisonOnFreeAllocator{checked} + + sc := arrow.NewSchema([]arrow.Field{{Name: "s", Type: arrow.BinaryTypes.String}}, nil) + + newBatch := func(vals ...string) arrow.RecordBatch { + bldr := array.NewRecordBuilder(mem, sc) + defer bldr.Release() + bldr.Field(0).(*array.StringBuilder).AppendValues(vals, nil) + return bldr.NewRecordBatch() + } + + var buf bytes.Buffer + w, err := pqarrow.NewFileWriter(sc, &buf, + parquet.NewWriterProperties(parquet.WithDictionaryDefault(false)), + pqarrow.NewArrowWriterProperties(pqarrow.WithAllocator(mem))) + require.NoError(t, err) + + batch1 := newBatch("mmm", "zzz") + require.NoError(t, w.WriteBuffered(batch1)) + batch1.Release() + + batch2 := newBatch("aaa", "nnn") + require.NoError(t, w.WriteBuffered(batch2)) + batch2.Release() + + require.NoError(t, w.Close()) + + rdr, err := file.NewParquetReader(bytes.NewReader(buf.Bytes())) + require.NoError(t, err) + defer rdr.Close() + + md := rdr.MetaData() + require.Equal(t, 1, len(md.RowGroups)) + col, err := md.RowGroup(0).ColumnChunk(0) + require.NoError(t, err) + stats, err := col.Statistics() + require.NoError(t, err) + require.True(t, stats.HasMinMax()) + assert.Equal(t, "aaa", string(stats.EncodeMin()), "min corrupted by released batch buffer") + assert.Equal(t, "zzz", string(stats.EncodeMax()), "max corrupted by released batch buffer") +}