This is an automated email from the ASF dual-hosted git repository.
zeroshade pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-go.git
The following commit(s) were added to refs/heads/main by this push:
new 227f6e3e fix(arrow/array): validate union JSON type codes (#1157)
227f6e3e is described below
commit 227f6e3ea469622d9b3f46924841e2dc4eed943d
Author: Minh Vu <[email protected]>
AuthorDate: Wed Aug 26 23:00:16 2026 +0200
fix(arrow/array): validate union JSON type codes (#1157)
## What
Validate union type IDs before narrowing them to int8 in both sparse and
dense JSON decoders. Out-of-range, negative, and nonnumeric type IDs are
rejected.
## Test
- go test ./arrow/array -run ^TestUnion -count=1
---
arrow/array/record.go | 21 +++--
arrow/array/record_test.go | 44 +++++++++++
arrow/array/union.go | 96 ++++++++++++-----------
arrow/array/union_test.go | 187 +++++++++++++++++++++++++++++++++++++++++++++
4 files changed, 298 insertions(+), 50 deletions(-)
diff --git a/arrow/array/record.go b/arrow/array/record.go
index 443d3492..704759b1 100644
--- a/arrow/array/record.go
+++ b/arrow/array/record.go
@@ -469,6 +469,12 @@ type builderCheckpoint struct {
lastStr *string
}
+func (checkpoint *builderCheckpoint) syncChildren(builders []Builder) {
+ for i := len(checkpoint.children); i < len(builders); i++ {
+ checkpoint.children = append(checkpoint.children,
newBuilderCheckpoint(builders[i]))
+ }
+}
+
func newBuilderCheckpoint(builder Builder) *builderCheckpoint {
checkpoint := &builderCheckpoint{
builder: builder,
@@ -499,13 +505,9 @@ func newBuilderCheckpoint(builder Builder)
*builderCheckpoint {
checkpoint.children = append(checkpoint.children,
newBuilderCheckpoint(field))
}
case *SparseUnionBuilder:
- for _, child := range builder.children {
- checkpoint.children = append(checkpoint.children,
newBuilderCheckpoint(child))
- }
+ checkpoint.syncChildren(builder.children)
case *DenseUnionBuilder:
- for _, child := range builder.children {
- checkpoint.children = append(checkpoint.children,
newBuilderCheckpoint(child))
- }
+ checkpoint.syncChildren(builder.children)
case storageBuilder:
checkpoint.children = append(checkpoint.children,
newBuilderCheckpoint(builder.StorageBuilder()))
case *RunEndEncodedBuilder:
@@ -520,6 +522,13 @@ func newBuilderCheckpoint(builder Builder)
*builderCheckpoint {
}
func (checkpoint *builderCheckpoint) capture() {
+ switch builder := checkpoint.builder.(type) {
+ case *SparseUnionBuilder:
+ checkpoint.syncChildren(builder.children)
+ case *DenseUnionBuilder:
+ checkpoint.syncChildren(builder.children)
+ }
+
checkpoint.length = checkpoint.builder.Len()
if checkpoint.state != nil {
checkpoint.state.capture()
diff --git a/arrow/array/record_test.go b/arrow/array/record_test.go
index 6ad3c7e3..ed7c005d 100644
--- a/arrow/array/record_test.go
+++ b/arrow/array/record_test.go
@@ -560,6 +560,50 @@ func TestRecordBuilderRollsBackRowsAfterDecodeError(t
*testing.T) {
}
+func TestRecordBuilderRollsBackDynamicallyAddedUnionChild(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+
+ for _, tc := range []struct {
+ name string
+ mode arrow.UnionMode
+ }{{name: "dense", mode: arrow.DenseMode}, {name: "sparse", mode:
arrow.SparseMode}} {
+ t.Run(tc.name, func(t *testing.T) {
+ unionType := arrow.UnionOf(tc.mode,
+ []arrow.Field{{Name: "first", Type:
arrow.PrimitiveTypes.Int32}},
+ []arrow.UnionTypeCode{0})
+ schema := arrow.NewSchema([]arrow.Field{
+ {Name: "value", Type: unionType},
+ {Name: "other", Type:
arrow.PrimitiveTypes.Int32},
+ }, nil)
+ builder := array.NewRecordBuilder(mem, schema)
+ defer builder.Release()
+
+ unionBuilder := builder.Field(0).(array.UnionBuilder)
+ assert.NoError(t,
builder.UnmarshalJSON([]byte(`{"value":[0,1],"other":1}`)))
+
+ secondChild := array.NewInt32Builder(mem)
+ defer secondChild.Release()
+ if tc.mode == arrow.SparseMode {
+ secondChild.AppendNull()
+ }
+ assert.EqualValues(t, 1,
unionBuilder.AppendChild(secondChild, "second"))
+
+ assert.Error(t,
builder.UnmarshalJSON([]byte(`{"value":[1,2],"other":"invalid"}`)))
+ wantChildLen := 0
+ if tc.mode == arrow.SparseMode {
+ wantChildLen = 1
+ }
+ assert.Equal(t, wantChildLen,
unionBuilder.Child(1).Len())
+
+ assert.NoError(t,
builder.UnmarshalJSON([]byte(`{"value":[1,3],"other":2}`)))
+ union := unionBuilder.NewArray().(array.Union)
+ defer union.Release()
+ assert.Equal(t, `[1,3]`, union.ValueStr(1))
+ })
+ }
+}
+
func TestRecordBuilderRollsBackVariableWidthState(t *testing.T) {
mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
defer mem.AssertSize(t, 0)
diff --git a/arrow/array/union.go b/arrow/array/union.go
index afa4e84c..e6137eae 100644
--- a/arrow/array/union.go
+++ b/arrow/array/union.go
@@ -772,10 +772,32 @@ type unionBuilder struct {
typeIDtoBuilder []Builder
typeIDtoChildID []int
// for all typeID < denseTypeID, typeIDtoBuilder[typeID] != nil
- denseTypeID arrow.UnionTypeCode
+ denseTypeID int
typesBuilder *int8BufferBuilder
}
+func unionTypeCodeFromJSON(dec *json.Decoder, typeID json.RawMessage, typ
arrow.DataType) (arrow.UnionTypeCode, error) {
+ id, err := json.Number(string(typeID)).Int64()
+ if err != nil {
+ return 0, &json.UnmarshalTypeError{
+ Offset: dec.InputOffset(),
+ Type: reflect.TypeOf(int8(0)),
+ Struct: fmt.Sprint(typ),
+ Value: "integer",
+ }
+ }
+
+ if id < 0 || id > int64(arrow.MaxUnionTypeCode) {
+ return 0, &json.UnmarshalTypeError{
+ Offset: dec.InputOffset(),
+ Type: reflect.TypeOf(int8(0)),
+ Struct: fmt.Sprint(typ),
+ Value: "integer",
+ }
+ }
+ return arrow.UnionTypeCode(id), nil
+}
+
func newUnionBuilder(mem memory.Allocator, children []Builder, typ
arrow.UnionType) *unionBuilder {
if children == nil {
children = make([]Builder, 0)
@@ -864,9 +886,9 @@ func (b *unionBuilder) Type() arrow.DataType {
}
func (b *unionBuilder) AppendChild(newChild Builder, fieldName string)
arrow.UnionTypeCode {
+ newType := b.nextTypeID()
newChild.Retain()
b.children = append(b.children, newChild)
- newType := b.nextTypeID()
b.typeIDtoChildID[newType] = len(b.children) - 1
b.typeIDtoBuilder[newType] = newChild
@@ -880,21 +902,23 @@ func (b *unionBuilder) nextTypeID() arrow.UnionTypeCode {
// find typeID such that typeIDtoBuilder[typeID] == nil
// use that for the new child. Start searching at denseTypeID
// since typeIDtoBuilder is densely packed up at least to denseTypeID
- for ; int(b.denseTypeID) < len(b.typeIDtoBuilder); b.denseTypeID++ {
+ for ; b.denseTypeID < len(b.typeIDtoBuilder); b.denseTypeID++ {
if b.typeIDtoBuilder[b.denseTypeID] == nil {
id := b.denseTypeID
b.denseTypeID++
- return id
+ return arrow.UnionTypeCode(id)
}
}
- debug.Assert(len(b.typeIDtoBuilder) < int(arrow.MaxUnionTypeCode), "too
many children typeids")
+ if b.denseTypeID > int(arrow.MaxUnionTypeCode) {
+ panic("arrow/array: too many children typeids")
+ }
// typeIDtoBuilder is already densely packed, so just append the new
child
b.typeIDtoBuilder = append(b.typeIDtoBuilder, nil)
b.typeIDtoChildID = append(b.typeIDtoChildID, arrow.InvalidUnionChildID)
id := b.denseTypeID
b.denseTypeID++
- return id
+ return arrow.UnionTypeCode(id)
}
func (b *unionBuilder) newData() *Data {
@@ -1070,6 +1094,7 @@ func (b *SparseUnionBuilder) AppendValueFromString(s
string) error {
return nil
}
dec := json.NewDecoder(strings.NewReader(s))
+ dec.UseNumber()
return b.UnmarshalOne(dec)
}
@@ -1082,30 +1107,21 @@ func (b *SparseUnionBuilder) UnmarshalOne(dec
*json.Decoder) error {
switch t {
case json.Delim('['):
// should be [type_id, Value]
- typeID, err := dec.Token()
- if err != nil {
+ var typeID json.RawMessage
+ if err := dec.Decode(&typeID); err != nil {
return err
}
- var typeCode int8
+ typeCode, err := unionTypeCodeFromJSON(dec, typeID, b.Type())
+ if err != nil {
+ return err
+ }
- switch tid := typeID.(type) {
- case json.Number:
- id, err := tid.Int64()
- if err != nil {
- return err
- }
- typeCode = int8(id)
- case float64:
- if tid != float64(int64(tid)) {
- return &json.UnmarshalTypeError{
- Offset: dec.InputOffset(),
- Type: reflect.TypeOf(int8(0)),
- Struct: fmt.Sprint(b.Type()),
- Value: "float",
- }
+ if int(typeCode) >= len(b.typeIDtoChildID) {
+ return &json.UnmarshalTypeError{
+ Offset: dec.InputOffset(),
+ Value: "invalid type code",
}
- typeCode = int8(tid)
}
childNum := b.typeIDtoChildID[typeCode]
@@ -1343,6 +1359,7 @@ func (d *DenseUnionBuilder) AppendValueFromString(s
string) error {
return nil
}
dec := json.NewDecoder(strings.NewReader(s))
+ dec.UseNumber()
return d.UnmarshalOne(dec)
}
@@ -1355,30 +1372,21 @@ func (b *DenseUnionBuilder) UnmarshalOne(dec
*json.Decoder) error {
switch t {
case json.Delim('['):
// should be [type_id, Value]
- typeID, err := dec.Token()
- if err != nil {
+ var typeID json.RawMessage
+ if err := dec.Decode(&typeID); err != nil {
return err
}
- var typeCode int8
+ typeCode, err := unionTypeCodeFromJSON(dec, typeID, b.Type())
+ if err != nil {
+ return err
+ }
- switch tid := typeID.(type) {
- case json.Number:
- id, err := tid.Int64()
- if err != nil {
- return err
- }
- typeCode = int8(id)
- case float64:
- if tid != float64(int64(tid)) {
- return &json.UnmarshalTypeError{
- Offset: dec.InputOffset(),
- Type: reflect.TypeOf(int8(0)),
- Struct: fmt.Sprint(b.Type()),
- Value: "float",
- }
+ if int(typeCode) >= len(b.typeIDtoChildID) {
+ return &json.UnmarshalTypeError{
+ Offset: dec.InputOffset(),
+ Value: "invalid type code",
}
- typeCode = int8(tid)
}
childNum := b.typeIDtoChildID[typeCode]
diff --git a/arrow/array/union_test.go b/arrow/array/union_test.go
index 7402f82b..391348ed 100644
--- a/arrow/array/union_test.go
+++ b/arrow/array/union_test.go
@@ -24,6 +24,7 @@ import (
"github.com/apache/arrow-go/v18/arrow"
"github.com/apache/arrow-go/v18/arrow/array"
"github.com/apache/arrow-go/v18/arrow/memory"
+ internaljson "github.com/apache/arrow-go/v18/internal/json"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/stretchr/testify/suite"
@@ -78,6 +79,192 @@ func TestUnionBuilderChildBounds(t *testing.T) {
}
}
+type unionJSONDecoder struct {
+ name string
+ apply func(array.UnionBuilder, string) error
+}
+
+func unionJSONDecoders() []unionJSONDecoder {
+ return []unionJSONDecoder{
+ {
+ name: "UnmarshalJSON",
+ apply: func(builder array.UnionBuilder, typeCode
string) error {
+ return builder.UnmarshalJSON([]byte("[[" +
typeCode + ", 1]]"))
+ },
+ },
+ {
+ name: "AppendValueFromString",
+ apply: func(builder array.UnionBuilder, typeCode
string) error {
+ return builder.AppendValueFromString("[" +
typeCode + ", 1]")
+ },
+ },
+ {
+ name: "UnmarshalOne",
+ apply: func(builder array.UnionBuilder, typeCode
string) error {
+ dec :=
internaljson.NewDecoder(strings.NewReader("[" + typeCode + ", 1]"))
+ return builder.UnmarshalOne(dec)
+ },
+ },
+ }
+}
+
+func TestUnionBuilderRejectsInvalidJSONTypeCodes(t *testing.T) {
+ fields := []arrow.Field{{Name: "value", Type:
arrow.PrimitiveTypes.Int32}}
+ cases := []struct {
+ name string
+ new func() array.UnionBuilder
+ }{
+ {
+ name: "dense",
+ new: func() array.UnionBuilder {
+ return
array.NewDenseUnionBuilder(memory.DefaultAllocator, arrow.DenseUnionOf(fields,
[]arrow.UnionTypeCode{0}))
+ },
+ },
+ {
+ name: "sparse",
+ new: func() array.UnionBuilder {
+ return
array.NewSparseUnionBuilder(memory.DefaultAllocator,
arrow.SparseUnionOf(fields, []arrow.UnionTypeCode{0}))
+ },
+ },
+ }
+
+ for _, tc := range cases {
+ t.Run(tc.name, func(t *testing.T) {
+ for _, decoder := range unionJSONDecoders() {
+ t.Run(decoder.name, func(t *testing.T) {
+ for _, typeCode := range
[]string{"256", "-1", "127", "1.5", "null", "-1e-400", "127.00000000000000001"}
{
+ t.Run(typeCode, func(t
*testing.T) {
+ builder := tc.new()
+ defer builder.Release()
+ assert.Error(t,
decoder.apply(builder, typeCode))
+ })
+ }
+ })
+ }
+ })
+ }
+}
+
+func TestUnionBuilderRejectsRoundedJSONTypeCodes(t *testing.T) {
+ fields := []arrow.Field{{Name: "value", Type:
arrow.PrimitiveTypes.Int32}}
+ cases := []struct {
+ name string
+ new func() array.UnionBuilder
+ }{
+ {
+ name: "dense",
+ new: func() array.UnionBuilder {
+ return
array.NewDenseUnionBuilder(memory.DefaultAllocator, arrow.DenseUnionOf(fields,
[]arrow.UnionTypeCode{127}))
+ },
+ },
+ {
+ name: "sparse",
+ new: func() array.UnionBuilder {
+ return
array.NewSparseUnionBuilder(memory.DefaultAllocator,
arrow.SparseUnionOf(fields, []arrow.UnionTypeCode{127}))
+ },
+ },
+ }
+
+ for _, tc := range cases {
+ t.Run(tc.name, func(t *testing.T) {
+ for _, decoder := range unionJSONDecoders() {
+ t.Run(decoder.name, func(t *testing.T) {
+ builder := tc.new()
+ defer builder.Release()
+ assert.Error(t, decoder.apply(builder,
"127.00000000000000001"))
+ })
+ }
+ })
+ }
+}
+
+func TestUnionBuilderUnmarshalOnePreservesDecoderConfiguration(t *testing.T) {
+ unionFields := []arrow.Field{{Name: "value", Type:
arrow.PrimitiveTypes.Int32}}
+ cases := []struct {
+ name string
+ typ arrow.DataType
+ }{
+ {
+ name: "dense",
+ typ: arrow.DenseUnionOf(unionFields,
[]arrow.UnionTypeCode{0}),
+ },
+ {
+ name: "sparse",
+ typ: arrow.SparseUnionOf(unionFields,
[]arrow.UnionTypeCode{0}),
+ },
+ }
+
+ for _, tc := range cases {
+ t.Run(tc.name, func(t *testing.T) {
+ dtype := arrow.StructOf(
+ arrow.Field{Name: "u", Type: tc.typ},
+ arrow.Field{Name: "i", Type:
arrow.PrimitiveTypes.Int32},
+ )
+ for _, input := range []struct {
+ name string
+ json string
+ }{
+ {name: "integer first", json:
`{"i":1.5,"u":[0,1]}`},
+ {name: "union first", json:
`{"u":[0,1],"i":1.5}`},
+ } {
+ t.Run(input.name, func(t *testing.T) {
+ builder :=
array.NewStructBuilder(memory.DefaultAllocator, dtype)
+ defer builder.Release()
+
+ dec :=
internaljson.NewDecoder(strings.NewReader(input.json))
+ require.NoError(t,
builder.UnmarshalOne(dec))
+ })
+ }
+ })
+ }
+}
+
+func TestUnionBuilderCanAppendMaxTypeCodeChild(t *testing.T) {
+ cases := []struct {
+ name string
+ new func(memory.Allocator) array.UnionBuilder
+ }{
+ {
+ name: "dense",
+ new: func(mem memory.Allocator) array.UnionBuilder {
+ return array.NewEmptyDenseUnionBuilder(mem)
+ },
+ },
+ {
+ name: "sparse",
+ new: func(mem memory.Allocator) array.UnionBuilder {
+ return array.NewEmptySparseUnionBuilder(mem)
+ },
+ },
+ }
+
+ for _, tc := range cases {
+ t.Run(tc.name, func(t *testing.T) {
+ mem :=
memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+
+ builder := tc.new(mem)
+ defer builder.Release()
+ appendChild := func(name string) arrow.UnionTypeCode {
+ child := array.NewInt32Builder(mem)
+ defer child.Release()
+ return builder.AppendChild(child, name)
+ }
+
+ for i := 0; i <= int(arrow.MaxUnionTypeCode); i++ {
+ code := appendChild(fmt.Sprintf("child-%d", i))
+ assert.EqualValues(t, i, code)
+ }
+ assert.Equal(t, int(arrow.MaxUnionTypeCode)+1,
builder.Type().(arrow.UnionType).NumFields())
+
+ assert.PanicsWithValue(t, "arrow/array: too many
children typeids", func() {
+ appendChild("overflow")
+ })
+ assert.Equal(t, int(arrow.MaxUnionTypeCode)+1,
builder.Type().(arrow.UnionType).NumFields())
+ })
+ }
+}
+
func TestUnionSliceEquals(t *testing.T) {
unionFields := []arrow.Field{
{Name: "u0", Type: arrow.PrimitiveTypes.Int32, Nullable: true},
