diff --git a/arrow/array/builder.go b/arrow/array/builder.go index 48113f431..2f3b6b283 100644 --- a/arrow/array/builder.go +++ b/arrow/array/builder.go @@ -17,6 +17,8 @@ package array import ( + "bytes" + "errors" "fmt" "math/bits" "sync/atomic" @@ -388,6 +390,170 @@ func (b *builder) UnsafeAppendBoolToBitmap(isValid bool) { b.length++ } +var jsonNull = []byte("null") + +func unmarshalChild(dec *json.Decoder, child Builder, field arrow.Field) error { + if field.Nullable { + return child.UnmarshalOne(dec) + } + + nulls := child.NullN() + if nulls == UnknownNullCount { + var val json.RawMessage + if err := dec.Decode(&val); err != nil { + return err + } + return unmarshalBufferedChild(val, child, field) + } + + // Every builder appends a null for a JSON null and for nothing else, so + // the child's null count going up by one over the call means the input + // had a null. + length := child.Len() + if err := child.UnmarshalOne(dec); err != nil { + return err + } + if child.Len() == length+1 && child.NullN() == nulls+1 { + return fmt.Errorf("field '%s' is non-nullable but got null", field.Name) + } + return nil +} + +func unmarshalBufferedChild(val json.RawMessage, child Builder, field arrow.Field) error { + if !field.Nullable && bytes.Equal(val, jsonNull) { + return fmt.Errorf("field '%s' is non-nullable but got null", field.Name) + } + + valDec := json.NewDecoder(bytes.NewReader(val)) + valDec.UseNumber() + return child.UnmarshalOne(valDec) +} + +// rowDecoder decodes a value out of a reused copy of that value, so that +// unescaping a string costs O(value) instead of O(remaining document). +// +// NOTE: goccy/go-json unescapes in place and shifts every byte after the escape, +// which is a big performance cost for large JSON documents: +// https://github.com/goccy/go-json/blob/v0.10.6/internal/decoder/string.go#L190 +type rowDecoder struct { + buf json.RawMessage + reader bytes.Reader + scratch json.RawMessage +} + +func (r *rowDecoder) next(dec *json.Decoder) (*json.Decoder, error) { + if err := dec.Decode(&r.buf); err != nil { + return nil, err + } + + r.reader.Reset(r.buf) + rowDec := json.NewDecoder(&r.reader) + rowDec.UseNumber() + return rowDec, nil +} + +func (r *rowDecoder) skip(dec *json.Decoder) error { + return dec.Decode(&r.scratch) +} + +// fieldContainer is what nestedJSONDecoder needs out of *arrow.Schema and +// *arrow.StructType. +type fieldContainer interface { + NumFields() int + Field(i int) arrow.Field +} + +// nestedJSONDecoder decodes JSON objects into one builder per field. +type nestedJSONDecoder struct { + rowDecoder + + fields fieldContainer + seen []bool + fieldIdx map[string]int +} + +func newNestedJSONDecoder(fields fieldContainer) nestedJSONDecoder { + return nestedJSONDecoder{fields: fields} +} + +func (d *nestedJSONDecoder) fieldIndexByName(name string) (int, bool) { + if d.fieldIdx == nil { + d.fieldIdx = make(map[string]int, d.fields.NumFields()) + for i := 0; i < d.fields.NumFields(); i++ { + d.fieldIdx[d.fields.Field(i).Name] = i + } + } + idx, ok := d.fieldIdx[name] + return idx, ok +} + +// unmarshalFields consumes one object from the decoder and appends its fields to +// each nested field builder. It consumes the opening and closing '{' and '}'. +// +// Nullable fields that are not present in the JSON-object are assumed to be null. +func (d *nestedJSONDecoder) unmarshalFields(dec *json.Decoder, builders []Builder) error { + // grow the "seen" buffer if needed + if cap(d.seen) < d.fields.NumFields() { + d.seen = make([]bool, d.fields.NumFields()) + } else { + d.seen = d.seen[:d.fields.NumFields()] + clear(d.seen) + } + + for dec.More() { + keyTok, err := dec.Token() + if err != nil { + return err + } + + key, ok := keyTok.(string) + if !ok { + return errors.New("JSON row must be an object with keys and values, but key is missing") + } + + idx, ok := d.fieldIndexByName(key) + if !ok { + // TODO(serramatutu): this is equivalent to ParseOptions::Ignore in Arrow C++, i.e silently drop extra keys. + // We should eventually support ParseOptions::InferType and ParseOptions::Error + if err := d.skip(dec); err != nil { + return err + } + continue + } + + if d.seen[idx] { + return fmt.Errorf("key '%s' is specified twice", key) + } + d.seen[idx] = true + + if err := unmarshalChild(dec, builders[idx], d.fields.Field(idx)); err != nil { + return err + } + } + + // consume the closing '}' + if _, err := dec.Token(); err != nil { + return err + } + + // check that all non-nullable fields were specified + for i := 0; i < d.fields.NumFields(); i++ { + field := d.fields.Field(i) + if !d.seen[i] && !field.Nullable { + return fmt.Errorf("field '%s' is required but no value was given", field.Name) + } + } + + // missing fields are nullable at this point, so they get a null + for i := 0; i < d.fields.NumFields(); i++ { + if !d.seen[i] { + builders[i].AppendNull() + } + } + + return nil +} + func NewBuilder(mem memory.Allocator, dtype arrow.DataType) Builder { // FIXME(sbinet): use a type switch on dtype instead? switch dtype.ID() { diff --git a/arrow/array/encoded.go b/arrow/array/encoded.go index 8e174f961..72955ee92 100644 --- a/arrow/array/encoded.go +++ b/arrow/array/encoded.go @@ -606,6 +606,12 @@ func (b *RunEndEncodedBuilder) UnmarshalOne(dec *json.Decoder) error { return err } + if value == nil { + if valueField := b.dt.(*arrow.RunEndEncodedType).Fields()[1]; !valueField.Nullable { + return fmt.Errorf("field '%s' is non-nullable but got null", valueField.Name) + } + } + // if we unmarshalled the same value as the previous one, we want to // continue the run. However, there's an edge case. At the start of // unmarshalling, lastUnmarshalled will be nil, but we might get diff --git a/arrow/array/encoded_test.go b/arrow/array/encoded_test.go index c8a5c4e2f..8da6961cf 100644 --- a/arrow/array/encoded_test.go +++ b/arrow/array/encoded_test.go @@ -744,3 +744,46 @@ func TestRunEndEncodedUnmarshalNestedJSON(t *testing.T) { assert.Truef(t, array.Equal(logicalValues, expectedValues), "expected: %s\ngot: %s", expectedValues, logicalValues) } + +func TestRunEndEncodedBuilderUnmarshalNonNullableValue(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + dt := arrow.RunEndEncodedOf(arrow.PrimitiveTypes.Int16, arrow.PrimitiveTypes.Int32) + dt.ValueNullable = false + + bldr := array.NewBuilder(mem, dt).(*array.RunEndEncodedBuilder) + defer bldr.Release() + + dec := json.NewDecoder(strings.NewReader("1 null")) + require.NoError(t, bldr.UnmarshalOne(dec)) + require.ErrorContains(t, bldr.UnmarshalOne(dec), "field 'values' is non-nullable but got null") + + nullableDt := arrow.RunEndEncodedOf(arrow.PrimitiveTypes.Int16, arrow.PrimitiveTypes.Int32) + nullableBldr := array.NewBuilder(mem, nullableDt).(*array.RunEndEncodedBuilder) + defer nullableBldr.Release() + + dec = json.NewDecoder(strings.NewReader("1 null")) + require.NoError(t, nullableBldr.UnmarshalOne(dec)) + require.NoError(t, nullableBldr.UnmarshalOne(dec)) + + arr := nullableBldr.NewRunEndEncodedArray() + defer arr.Release() + + values := arr.Values().(*array.Int32) + assert.False(t, values.IsNull(0)) + assert.True(t, values.IsNull(1)) +} + +func TestRunEndEncodedBuilderKeepsValueNullable(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + dt := arrow.RunEndEncodedOf(arrow.PrimitiveTypes.Int16, arrow.PrimitiveTypes.Int32) + dt.ValueNullable = false + + bldr := array.NewBuilder(mem, dt) + defer bldr.Release() + + require.False(t, bldr.Type().(*arrow.RunEndEncodedType).ValueNullable) +} diff --git a/arrow/array/fixed_size_list.go b/arrow/array/fixed_size_list.go index c3aafb315..78362b338 100644 --- a/arrow/array/fixed_size_list.go +++ b/arrow/array/fixed_size_list.go @@ -365,6 +365,19 @@ func (b *FixedSizeListBuilder) AppendValueFromString(s string) error { } func (b *FixedSizeListBuilder) UnmarshalOne(dec *json.Decoder) error { + if b.checkpoint == nil { + b.checkpoint = newBuilderCheckpoint(b) + } + b.checkpoint.capture() + + if err := b.unmarshalOne(dec); err != nil { + b.checkpoint.restore() + return err + } + return nil +} + +func (b *FixedSizeListBuilder) unmarshalOne(dec *json.Decoder) error { t, err := dec.Token() if err != nil { return err @@ -373,7 +386,7 @@ func (b *FixedSizeListBuilder) UnmarshalOne(dec *json.Decoder) error { switch t { case json.Delim('['): b.Append(true) - if err := b.values.Unmarshal(dec); err != nil { + if err := unmarshalListValues(dec, b.values, b.dt); err != nil { return err } // consume ']' diff --git a/arrow/array/list.go b/arrow/array/list.go index 463da03be..42407a1b0 100644 --- a/arrow/array/list.go +++ b/arrow/array/list.go @@ -297,6 +297,8 @@ type baseListBuilder struct { // actual list type dt arrow.DataType appendOffsetVal func(int) + + checkpoint *builderCheckpoint } type ListLikeBuilder interface { @@ -611,7 +613,40 @@ func (b *baseListBuilder) AppendValueFromString(s string) error { return b.UnmarshalOne(json.NewDecoder(strings.NewReader(s))) } +func unmarshalListValues(dec *json.Decoder, values Builder, dt arrow.DataType) error { + listLike, ok := dt.(arrow.ListLikeType) + if !ok { + return values.Unmarshal(dec) + } + + elem := listLike.ElemField() + if elem.Nullable { + return values.Unmarshal(dec) + } + + for dec.More() { + if err := unmarshalChild(dec, values, elem); err != nil { + return err + } + } + + return nil +} + func (b *baseListBuilder) UnmarshalOne(dec *json.Decoder) error { + if b.checkpoint == nil { + b.checkpoint = newBuilderCheckpoint(b) + } + b.checkpoint.capture() + + if err := b.unmarshalOne(dec); err != nil { + b.checkpoint.restore() + return err + } + return nil +} + +func (b *baseListBuilder) unmarshalOne(dec *json.Decoder) error { t, err := dec.Token() if err != nil { return err @@ -620,7 +655,7 @@ func (b *baseListBuilder) UnmarshalOne(dec *json.Decoder) error { switch t { case json.Delim('['): b.Append(true) - if err := b.values.Unmarshal(dec); err != nil { + if err := unmarshalListValues(dec, b.values, b.dt); err != nil { return err } // consume ']' @@ -1104,6 +1139,8 @@ type baseListViewBuilder struct { dt arrow.DataType appendOffsetVal func(int) appendSizeVal func(int) + + checkpoint *builderCheckpoint } type ListViewBuilder struct { @@ -1421,6 +1458,19 @@ func (b *baseListViewBuilder) AppendValueFromString(s string) error { } func (b *baseListViewBuilder) UnmarshalOne(dec *json.Decoder) error { + if b.checkpoint == nil { + b.checkpoint = newBuilderCheckpoint(b) + } + b.checkpoint.capture() + + if err := b.unmarshalOne(dec); err != nil { + b.checkpoint.restore() + return err + } + return nil +} + +func (b *baseListViewBuilder) unmarshalOne(dec *json.Decoder) error { t, err := dec.Token() if err != nil { return err @@ -1431,7 +1481,7 @@ func (b *baseListViewBuilder) UnmarshalOne(dec *json.Decoder) error { offset := b.values.Len() // 0 is a placeholder size as we don't know the actual size yet b.AppendWithSize(true, 0) - if err := b.values.Unmarshal(dec); err != nil { + if err := unmarshalListValues(dec, b.values, b.dt); err != nil { return err } // consume ']' diff --git a/arrow/array/list_test.go b/arrow/array/list_test.go index 6c513c8f1..2c519c43c 100644 --- a/arrow/array/list_test.go +++ b/arrow/array/list_test.go @@ -18,12 +18,15 @@ package array_test import ( "reflect" + "strings" "testing" "github.com/apache/arrow-go/v18/arrow" "github.com/apache/arrow-go/v18/arrow/array" "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/apache/arrow-go/v18/internal/json" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestListArray(t *testing.T) { @@ -862,3 +865,192 @@ func TestRangeOfValuesUsed(t *testing.T) { }) } } + +func TestListUnmarshalNonNullableElem(t *testing.T) { + pool := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer pool.AssertSize(t, 0) + + tests := []struct { + name string + dt arrow.DataType + jsonInput string + wantErr string + want string + }{ + { + name: "list", + dt: arrow.ListOfNonNullable(arrow.PrimitiveTypes.Int32), + jsonInput: `[[1, null]]`, + wantErr: "field 'item' is non-nullable but got null", + }, + { + name: "large list", + dt: arrow.LargeListOfNonNullable(arrow.PrimitiveTypes.Int32), + jsonInput: `[[1, null]]`, + wantErr: "field 'item' is non-nullable but got null", + }, + { + name: "list view", + dt: arrow.ListViewOfNonNullable(arrow.PrimitiveTypes.Int32), + jsonInput: `[[1, null]]`, + wantErr: "field 'item' is non-nullable but got null", + }, + { + name: "large list view", + dt: arrow.LargeListViewOfNonNullable(arrow.PrimitiveTypes.Int32), + jsonInput: `[[1, null]]`, + wantErr: "field 'item' is non-nullable but got null", + }, + { + name: "fixed size list", + dt: arrow.FixedSizeListOfNonNullable(2, arrow.PrimitiveTypes.Int32), + jsonInput: `[[1, null]]`, + wantErr: "field 'item' is non-nullable but got null", + }, + { + name: "nested non-nullable elem", + dt: arrow.ListOf(arrow.ListOfNonNullable(arrow.PrimitiveTypes.Int32)), + jsonInput: `[[[1, null]]]`, + wantErr: "field 'item' is non-nullable but got null", + }, + { + name: "non-nullable struct elem", + dt: arrow.ListOfNonNullable(arrow.StructOf(arrow.Field{Name: "x", Type: arrow.PrimitiveTypes.Int32, Nullable: true})), + jsonInput: `[[null]]`, + wantErr: "field 'item' is non-nullable but got null", + }, + { + name: "nullable elem accepts null", + dt: arrow.ListOf(arrow.PrimitiveTypes.Int32), + jsonInput: `[[1, null]]`, + want: `[[1 (null)]]`, + }, + { + name: "non-nullable elem accepts values", + dt: arrow.ListOfNonNullable(arrow.PrimitiveTypes.Int64), + jsonInput: `[[9007199254740993]]`, + want: `[[9007199254740993]]`, + }, + { + name: "non-nullable elem accepts null list", + dt: arrow.ListOfNonNullable(arrow.PrimitiveTypes.Int32), + jsonInput: `[null, [1]]`, + want: `[(null) [1]]`, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + arr, _, err := array.FromJSON(pool, tc.dt, strings.NewReader(tc.jsonInput)) + if tc.wantErr != "" { + require.ErrorContains(t, err, tc.wantErr) + return + } + require.NoError(t, err) + defer arr.Release() + + require.NoError(t, array.ValidateFull(arr)) + require.Equal(t, tc.want, arr.String()) + }) + } +} + +func TestListBuilderUnmarshalOneRollback(t *testing.T) { + tests := []struct { + name string + dt arrow.DataType + bad string + good string + wantErr string + want string + wantElems int + }{ + { + name: "list null elem", + dt: arrow.ListOfNonNullable(arrow.PrimitiveTypes.Int32), + bad: `[1, null]`, + good: `[1, 2]`, + wantErr: "field 'item' is non-nullable but got null", + want: `[[1 2]]`, + wantElems: 2, + }, + { + name: "list wrong elem type", + dt: arrow.ListOf(arrow.PrimitiveTypes.Int32), + bad: `[1, "nope"]`, + good: `[1, 2]`, + wantErr: "cannot unmarshal", + want: `[[1 2]]`, + wantElems: 2, + }, + { + name: "large list null elem", + dt: arrow.LargeListOfNonNullable(arrow.PrimitiveTypes.Int32), + bad: `[1, null]`, + good: `[1, 2]`, + wantErr: "field 'item' is non-nullable but got null", + want: `[[1 2]]`, + wantElems: 2, + }, + { + name: "list view null elem", + dt: arrow.ListViewOfNonNullable(arrow.PrimitiveTypes.Int32), + bad: `[1, null]`, + good: `[1, 2]`, + wantErr: "field 'item' is non-nullable but got null", + want: `[[1 2]]`, + wantElems: 2, + }, + { + name: "large list view null elem", + dt: arrow.LargeListViewOfNonNullable(arrow.PrimitiveTypes.Int32), + bad: `[1, null]`, + good: `[1, 2]`, + wantErr: "field 'item' is non-nullable but got null", + want: `[[1 2]]`, + wantElems: 2, + }, + { + name: "fixed size list null elem", + dt: arrow.FixedSizeListOfNonNullable(2, arrow.PrimitiveTypes.Int32), + bad: `[1, null]`, + good: `[1, 2]`, + wantErr: "field 'item' is non-nullable but got null", + want: `[[1 2]]`, + wantElems: 2, + }, + { + name: "nested list null elem", + dt: arrow.ListOf(arrow.ListOfNonNullable(arrow.PrimitiveTypes.Int32)), + bad: `[[1, null]]`, + good: `[[1, 2]]`, + wantErr: "field 'item' is non-nullable but got null", + want: `[[[1 2]]]`, + wantElems: 1, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + pool := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer pool.AssertSize(t, 0) + + bldr := array.NewBuilder(pool, tc.dt).(array.ListLikeBuilder) + defer bldr.Release() + + require.ErrorContains(t, bldr.UnmarshalOne(json.NewDecoder(strings.NewReader(tc.bad))), tc.wantErr) + require.Zero(t, bldr.Len()) + require.Zero(t, bldr.ValueBuilder().Len()) + + require.NoError(t, bldr.UnmarshalOne(json.NewDecoder(strings.NewReader(tc.good)))) + require.Equal(t, 1, bldr.Len()) + require.Equal(t, tc.wantElems, bldr.ValueBuilder().Len()) + + arr := bldr.NewArray() + defer arr.Release() + + require.NoError(t, array.ValidateFull(arr)) + require.Equal(t, tc.want, arr.String()) + }) + } +} diff --git a/arrow/array/map.go b/arrow/array/map.go index 8bc6332d2..0b0b9634b 100644 --- a/arrow/array/map.go +++ b/arrow/array/map.go @@ -150,7 +150,7 @@ type MapBuilder struct { func NewMapBuilder(mem memory.Allocator, keytype, itemtype arrow.DataType, keysSorted bool) *MapBuilder { etype := arrow.MapOf(keytype, itemtype) etype.KeysSorted = keysSorted - listBldr := NewListBuilder(mem, etype.Elem()) + listBldr := NewListBuilderWithField(mem, etype.ElemField()) keyBldr := listBldr.ValueBuilder().(*StructBuilder).FieldBuilder(0) keyBldr.Retain() itemBldr := listBldr.ValueBuilder().(*StructBuilder).FieldBuilder(1) @@ -167,7 +167,7 @@ func NewMapBuilder(mem memory.Allocator, keytype, itemtype arrow.DataType, keysS } func NewMapBuilderWithType(mem memory.Allocator, dt *arrow.MapType) *MapBuilder { - listBldr := NewListBuilder(mem, dt.Elem()) + listBldr := NewListBuilderWithField(mem, dt.ElemField()) keyBldr := listBldr.ValueBuilder().(*StructBuilder).FieldBuilder(0) keyBldr.Retain() itemBldr := listBldr.ValueBuilder().(*StructBuilder).FieldBuilder(1) diff --git a/arrow/array/map_test.go b/arrow/array/map_test.go index 3e9f16678..a5f55aefb 100644 --- a/arrow/array/map_test.go +++ b/arrow/array/map_test.go @@ -19,12 +19,15 @@ package array_test import ( "fmt" "strconv" + "strings" "testing" "github.com/apache/arrow-go/v18/arrow" "github.com/apache/arrow-go/v18/arrow/array" "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/apache/arrow-go/v18/internal/json" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestMapArray(t *testing.T) { @@ -413,3 +416,99 @@ func appendMapBuilderPrefix(builder *array.MapBuilder, n int) { } } } + +func TestMapUnmarshalNonNullableFields(t *testing.T) { + nonNullableItem := func() *arrow.MapType { + dt := arrow.MapOf(arrow.BinaryTypes.String, arrow.PrimitiveTypes.Int32) + dt.SetItemNullable(false) + return dt + } + + tests := []struct { + name string + dt arrow.DataType + jsonInput string + wantErr string + want string + }{ + { + name: "null entry", + dt: arrow.MapOf(arrow.BinaryTypes.String, arrow.PrimitiveTypes.Int32), + jsonInput: `[[null]]`, + wantErr: "field 'entries' is non-nullable but got null", + }, + { + name: "null key", + dt: arrow.MapOf(arrow.BinaryTypes.String, arrow.PrimitiveTypes.Int32), + jsonInput: `[[{"key": null, "value": 1}]]`, + wantErr: "field 'key' is non-nullable but got null", + }, + { + name: "null item with non-nullable item", + dt: nonNullableItem(), + jsonInput: `[[{"key": "a", "value": null}]]`, + wantErr: "field 'value' is non-nullable but got null", + }, + { + name: "null map", + dt: arrow.MapOf(arrow.BinaryTypes.String, arrow.PrimitiveTypes.Int32), + jsonInput: `[null]`, + want: `[(null)]`, + }, + { + name: "null item with nullable item", + dt: arrow.MapOf(arrow.BinaryTypes.String, arrow.PrimitiveTypes.Int32), + jsonInput: `[[{"key": "a", "value": null}]]`, + want: `[{["a"] [(null)]}]`, + }, + { + name: "non-nullable item with values", + dt: nonNullableItem(), + jsonInput: `[[{"key": "a", "value": 1}]]`, + want: `[{["a"] [1]}]`, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + pool := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer pool.AssertSize(t, 0) + + arr, _, err := array.FromJSON(pool, tc.dt, strings.NewReader(tc.jsonInput)) + if tc.wantErr != "" { + require.ErrorContains(t, err, tc.wantErr) + return + } + require.NoError(t, err) + defer arr.Release() + + require.NoError(t, array.ValidateFull(arr)) + require.Equal(t, tc.want, arr.String()) + }) + } +} + +func TestMapBuilderUnmarshalOneRollback(t *testing.T) { + pool := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer pool.AssertSize(t, 0) + + dt := arrow.MapOf(arrow.BinaryTypes.String, arrow.PrimitiveTypes.Int32) + dt.SetItemNullable(false) + + bldr := array.NewMapBuilderWithType(pool, dt) + defer bldr.Release() + + err := bldr.UnmarshalOne(json.NewDecoder(strings.NewReader(`[{"key": "a", "value": null}]`))) + require.ErrorContains(t, err, "field 'value' is non-nullable but got null") + require.Zero(t, bldr.Len()) + require.Zero(t, bldr.KeyBuilder().Len()) + require.Zero(t, bldr.ItemBuilder().Len()) + + require.NoError(t, bldr.UnmarshalOne(json.NewDecoder(strings.NewReader(`[{"key": "a", "value": 1}]`)))) + + arr := bldr.NewArray() + defer arr.Release() + + require.NoError(t, array.ValidateFull(arr)) + require.Equal(t, `[{["a"] [1]}]`, arr.String()) +} diff --git a/arrow/array/record.go b/arrow/array/record.go index 704759b10..31c0b07be 100644 --- a/arrow/array/record.go +++ b/arrow/array/record.go @@ -303,14 +303,16 @@ type RecordBuilder struct { schema *arrow.Schema fields []Builder checkpoints []*builderCheckpoint + jsonDec nestedJSONDecoder } // NewRecordBuilder returns a builder, using the provided memory allocator and a schema. func NewRecordBuilder(mem memory.Allocator, schema *arrow.Schema) *RecordBuilder { b := &RecordBuilder{ - mem: mem, - schema: schema, - fields: make([]Builder, schema.NumFields()), + mem: mem, + schema: schema, + fields: make([]Builder, schema.NumFields()), + jsonDec: newNestedJSONDecoder(schema), } b.refCount.Add(1) @@ -459,8 +461,13 @@ type storageBuilder interface { StorageBuilder() Builder } +type truncatableBuilder interface { + Len() int + truncate(n int) +} + type builderCheckpoint struct { - builder Builder + builder truncatableBuilder length int children []*builderCheckpoint state checkpointState @@ -475,7 +482,7 @@ func (checkpoint *builderCheckpoint) syncChildren(builders []Builder) { } } -func newBuilderCheckpoint(builder Builder) *builderCheckpoint { +func newBuilderCheckpoint(builder truncatableBuilder) *builderCheckpoint { checkpoint := &builderCheckpoint{ builder: builder, } @@ -488,6 +495,10 @@ func newBuilderCheckpoint(builder Builder) *builderCheckpoint { // Keep this switch in sync with builder types that own children. An omitted // nested builder would restore its own state but leave its children changed. switch builder := builder.(type) { + case *baseListBuilder: + checkpoint.children = append(checkpoint.children, newBuilderCheckpoint(builder.values)) + case *baseListViewBuilder: + checkpoint.children = append(checkpoint.children, newBuilderCheckpoint(builder.values)) case *ListBuilder: checkpoint.children = append(checkpoint.children, newBuilderCheckpoint(builder.values)) case *LargeListBuilder: @@ -564,8 +575,10 @@ func (checkpoint *builderCheckpoint) restore() { } // UnmarshalOne reads one row (a JSON object) from the supplied decoder and -// appends a value to each field in the RecordBuilder. Missing fields are -// appended as nulls and unrecognized keys are silently ignored. +// appends a value to each field in the RecordBuilder. +// +// Missing nullable fields get nulls. If an error is found while decoding +// any field, the entire row gets discarded and the error is returned. // // Unlike UnmarshalJSON, this method receives an already-configured // json.Decoder, so options such as UseNumber set by the caller are honored @@ -597,9 +610,13 @@ func (b *RecordBuilder) UnmarshalOne(dec *json.Decoder) error { } func (b *RecordBuilder) unmarshalOne(dec *json.Decoder) (err error) { + rowDec, err := b.jsonDec.next(dec) + if err != nil { + return err + } // should start with a '{' - t, err := dec.Token() + t, err := rowDec.Token() if err != nil { return err } @@ -608,49 +625,12 @@ func (b *RecordBuilder) unmarshalOne(dec *json.Decoder) (err error) { return fmt.Errorf("record should start with '{', not %s", t) } - keylist := make(map[string]bool) - for dec.More() { - keyTok, err := dec.Token() - if err != nil { - return err - } - - key := keyTok.(string) - if keylist[key] { - return fmt.Errorf("key %s shows up twice in row to be decoded", key) - } - keylist[key] = true - - indices := b.schema.FieldIndices(key) - if len(indices) == 0 { - var extra interface{} - if err := dec.Decode(&extra); err != nil { - return err - } - continue - } - - if err := b.fields[indices[0]].UnmarshalOne(dec); err != nil { - return err - } - } - - // consume the closing '}' - if _, err := dec.Token(); err != nil { - return err - } - - for i := 0; i < b.schema.NumFields(); i++ { - if !keylist[b.schema.Field(i).Name] { - b.fields[i].AppendNull() - } - } - return nil + return b.jsonDec.unmarshalFields(rowDec, b.fields) } // Unmarshal reads multiple rows from the decoder, calling UnmarshalOne in a // loop until dec.More() reports there are no more values. Like UnmarshalOne, -// this honors decoder configuration such as UseNumber set by the caller. +// field values are always re-decoded with UseNumber enabled. func (b *RecordBuilder) Unmarshal(dec *json.Decoder) error { for dec.More() { if err := b.UnmarshalOne(dec); err != nil { diff --git a/arrow/array/record_test.go b/arrow/array/record_test.go index ed7c005d5..08ac248eb 100644 --- a/arrow/array/record_test.go +++ b/arrow/array/record_test.go @@ -17,6 +17,7 @@ package array_test import ( + "bytes" "fmt" "reflect" "strings" @@ -25,7 +26,9 @@ import ( "github.com/apache/arrow-go/v18/arrow" "github.com/apache/arrow-go/v18/arrow/array" "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/apache/arrow-go/v18/internal/json" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestRecord(t *testing.T) { @@ -485,9 +488,9 @@ func TestRecordBuilder(t *testing.T) { mapDt.SetItemNullable(false) schema := arrow.NewSchema( []arrow.Field{ - {Name: "f1-i32", Type: arrow.PrimitiveTypes.Int32}, - {Name: "f2-f64", Type: arrow.PrimitiveTypes.Float64}, - {Name: "map", Type: mapDt}, + {Name: "f1-i32", Type: arrow.PrimitiveTypes.Int32, Nullable: true}, + {Name: "f2-f64-notnull", Type: arrow.PrimitiveTypes.Float64, Nullable: false}, + {Name: "map", Type: mapDt, Nullable: true}, }, nil, ) @@ -498,11 +501,14 @@ func TestRecordBuilder(t *testing.T) { b.Retain() b.Release() - b.Field(0).(*array.Int32Builder).AppendValues([]int32{1, 2, 3}, nil) + b.Field(0).(*array.Int32Builder).AppendNull() + b.Field(0).(*array.Int32Builder).AppendValues([]int32{2, 3}, nil) b.Field(0).(*array.Int32Builder).AppendValues([]int32{4, 5}, nil) - b.Field(1).(*array.Float64Builder).AppendValues([]float64{1, 2, 3, 4, 5}, nil) + + b.Field(1).(*array.Float64Builder).AppendValues([]float64{1.1, 2.2, 3.3, 4.4, 5.5}, nil) + mb := b.Field(2).(*array.MapBuilder) - for i := 0; i < 5; i++ { + for i := range 5 { mb.Append(true) if i%3 == 0 { @@ -511,6 +517,15 @@ func TestRecordBuilder(t *testing.T) { } } + err := b.UnmarshalJSON([]byte(`{"f1-i32": null, "f2-f64-notnull": null, "map": null}`)) + assert.Contains(t, err.Error(), "field 'f2-f64-notnull' is non-nullable but got null") + + err = b.UnmarshalJSON([]byte(`{"f1-i32": null, "map": null}`)) + assert.Contains(t, err.Error(), "field 'f2-f64-notnull' is required but no value was given") + + err = b.UnmarshalJSON([]byte(`{"f1-i32": 6, "f2-f64-notnull": 6.6, "map": [{"key": "4", "value": "d"}]}`)) + assert.NoError(t, err) + rec := b.NewRecordBatch() defer rec.Release() @@ -518,7 +533,7 @@ func TestRecordBuilder(t *testing.T) { t.Fatalf("invalid schema: got=%#v, want=%#v", got, want) } - if got, want := rec.NumRows(), int64(5); got != want { + if got, want := rec.NumRows(), int64(6); got != want { t.Fatalf("invalid number of rows: got=%d, want=%d", got, want) } if got, want := rec.NumCols(), int64(3); got != want { @@ -527,9 +542,27 @@ func TestRecordBuilder(t *testing.T) { if got, want := rec.ColumnName(0), schema.Field(0).Name; got != want { t.Fatalf("invalid column name: got=%q, want=%q", got, want) } - if got, want := rec.Column(2).String(), `[{["0" "2" "3"] ["a" "b" "c"]} {[] []} {[] []} {["3" "2" "3"] ["a" "b" "c"]} {[] []}]`; got != want { - t.Fatalf("invalid column name: got=%q, want=%q", got, want) + + if got, want := rec.Column(0).String(), `[(null) 2 3 4 5 6]`; got != want { + t.Fatalf("invalid column values: got=%q, want=%q", got, want) + } + if got, want := rec.Column(1).String(), `[1.1 2.2 3.3 4.4 5.5 6.6]`; got != want { + t.Fatalf("invalid column values: got=%q, want=%q", got, want) + } + if got, want := rec.Column(2).String(), `[{["0" "2" "3"] ["a" "b" "c"]} {[] []} {[] []} {["3" "2" "3"] ["a" "b" "c"]} {[] []} {["4"] ["d"]}]`; got != want { + t.Fatalf("invalid column values: got=%q, want=%q", got, want) } + + // roundtripping from JSON with array.FromJSON should work + arr := array.RecordToStructArray(rec) + defer arr.Release() + jsonStr, err := json.Marshal(arr) + require.NoError(t, err) + + roundtripped, _, err := array.FromJSON(mem, arr.DataType(), bytes.NewReader(jsonStr)) + require.NoError(t, err) + defer roundtripped.Release() + assert.Truef(t, array.Equal(arr, roundtripped), "JSON round trip returns different array: got=%q, want=%q", roundtripped, arr) } func TestRecordBuilderRollsBackRowsAfterDecodeError(t *testing.T) { @@ -698,7 +731,7 @@ func TestRecordBuilderRollsBackBooleanAndNullLengths(t *testing.T) { } { t.Run(tc.name, func(t *testing.T) { schema := arrow.NewSchema([]arrow.Field{ - {Name: "value", Type: tc.typ}, + {Name: "value", Type: tc.typ, Nullable: true}, {Name: "other", Type: arrow.PrimitiveTypes.Int32}, }, nil) builder := array.NewRecordBuilder(mem, schema) @@ -723,7 +756,7 @@ func TestRecordBuilderRollsBackDiscardedValidityBits(t *testing.T) { defer mem.AssertSize(t, 0) schema := arrow.NewSchema([]arrow.Field{ - {Name: "value", Type: arrow.PrimitiveTypes.Int32}, + {Name: "value", Type: arrow.PrimitiveTypes.Int32, Nullable: true}, {Name: "other", Type: arrow.PrimitiveTypes.Int32}, }, nil) builder := array.NewRecordBuilder(mem, schema) diff --git a/arrow/array/struct.go b/arrow/array/struct.go index a9d80835d..b7fb3bb32 100644 --- a/arrow/array/struct.go +++ b/arrow/array/struct.go @@ -18,7 +18,6 @@ package array import ( "bytes" - "errors" "fmt" "math" "strings" @@ -273,9 +272,14 @@ func (a *Struct) GetOneForMarshal(i int) interface{} { } tmp := make(map[string]interface{}) - fieldList := a.data.dtype.(*arrow.StructType).Fields() + dtype := a.data.dtype.(*arrow.StructType) + fieldList := dtype.Fields() for j, d := range a.fields { - tmp[fieldList[j].Name] = d.GetOneForMarshal(i) + if dtype.Field(j).Nullable && d.IsNull(i) { + tmp[fieldList[j].Name] = nil + } else { + tmp[fieldList[j].Name] = d.GetOneForMarshal(i) + } } return tmp } @@ -336,8 +340,10 @@ func (a *Struct) Release() { type StructBuilder struct { builder - dtype arrow.DataType - fields []Builder + dtype arrow.DataType + fields []Builder + checkpoint *builderCheckpoint + jsonDec nestedJSONDecoder } // NewStructBuilder returns a builder, using the provided memory allocator. @@ -346,6 +352,7 @@ func NewStructBuilder(mem memory.Allocator, dtype *arrow.StructType) *StructBuil builder: builder{mem: mem}, dtype: dtype, fields: make([]Builder, dtype.NumFields()), + jsonDec: newNestedJSONDecoder(dtype), } b.refCount.Add(1) @@ -378,6 +385,7 @@ func (b *StructBuilder) Release() { for _, f := range b.fields { f.Release() } + b.checkpoint = nil } } @@ -536,7 +544,26 @@ func (b *StructBuilder) AppendValueFromString(s string) error { } func (b *StructBuilder) UnmarshalOne(dec *json.Decoder) error { - t, err := dec.Token() + if b.checkpoint == nil { + b.checkpoint = newBuilderCheckpoint(b) + } + b.checkpoint.capture() + + if err := b.unmarshalOne(dec); err != nil { + b.checkpoint.restore() + return err + } + return nil +} + +func (b *StructBuilder) unmarshalOne(dec *json.Decoder) error { + offset := dec.InputOffset() + rowDec, err := b.jsonDec.next(dec) + if err != nil { + return err + } + + t, err := rowDec.Token() if err != nil { return err } @@ -544,57 +571,12 @@ func (b *StructBuilder) UnmarshalOne(dec *json.Decoder) error { switch t { case json.Delim('{'): b.Append(true) - keylist := make(map[string]bool) - for dec.More() { - keyTok, err := dec.Token() - if err != nil { - return err - } - - key, ok := keyTok.(string) - if !ok { - return errors.New("missing key") - } - - if keylist[key] { - return fmt.Errorf("key %s is specified twice", key) - } - - keylist[key] = true - - idx, ok := b.dtype.(*arrow.StructType).FieldIdx(key) - if !ok { - var extra interface{} - if err := dec.Decode(&extra); err != nil { - return err - } - continue - } - - if err := b.fields[idx].UnmarshalOne(dec); err != nil { - return err - } - } - - // Append null values to all optional fields that were not presented in the json input - for _, field := range b.dtype.(*arrow.StructType).Fields() { - if !field.Nullable { - continue - } - idx, _ := b.dtype.(*arrow.StructType).FieldIdx(field.Name) - if _, hasKey := keylist[field.Name]; !hasKey { - b.fields[idx].AppendNull() - } - } - - // consume '}' - _, err := dec.Token() - return err + return b.jsonDec.unmarshalFields(rowDec, b.fields) case nil: b.AppendNull() default: return &json.UnmarshalTypeError{ - Offset: dec.InputOffset(), + Offset: offset + rowDec.InputOffset(), Struct: fmt.Sprint(b.dtype), } } diff --git a/arrow/array/struct_test.go b/arrow/array/struct_test.go index 74fda5039..beec80c5c 100644 --- a/arrow/array/struct_test.go +++ b/arrow/array/struct_test.go @@ -656,20 +656,25 @@ func TestStructArrayUnmarshalJSONMissingFields(t *testing.T) { name string jsonInput string want string - invalid bool + wantErr string }{ { name: "missing required field", jsonInput: `[{"f2": 3, "f3": {"f3_1": "test"}}]`, - invalid: true, + wantErr: "field 'f3_3' is required but no value was given", want: "", }, { name: "missing optional fields", jsonInput: `[{"f2": 3, "f3": {"f3_3": "test"}}]`, - invalid: false, want: `{[(null)] [3] {[(null)] [(null)] ["test"]}}`, }, + { + name: "explicit null in required field", + jsonInput: `[{"f2": 3, "f3": {"f3_3": null}}]`, + wantErr: "field 'f3_3' is non-nullable but got null", + want: "", + }, } for _, tc := range tests { @@ -679,16 +684,14 @@ func TestStructArrayUnmarshalJSONMissingFields(t *testing.T) { defer sb.Release() err := sb.UnmarshalJSON([]byte(tc.jsonInput)) - if err != nil { - t.Fatal(err) + if tc.wantErr != "" { + require.ErrorContains(t, err, tc.wantErr) + return } + require.NoError(t, err) arr := sb.NewArray().(*array.Struct) defer arr.Release() - if tc.invalid { - require.ErrorIs(t, array.ValidateFull(arr), arrow.ErrInvalid) - return - } require.NoError(t, array.ValidateFull(arr)) got := arr.String() @@ -700,6 +703,32 @@ func TestStructArrayUnmarshalJSONMissingFields(t *testing.T) { } } +func TestStructBuilderRollsBackRowAfterNestedDecodeError(t *testing.T) { + pool := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer pool.AssertSize(t, 0) + + dtype := arrow.StructOf( + arrow.Field{Name: "a", Type: arrow.PrimitiveTypes.Int32}, + arrow.Field{Name: "b", Type: arrow.ListOf(arrow.PrimitiveTypes.Int32)}, + ) + + sb := array.NewStructBuilder(pool, dtype) + defer sb.Release() + + require.Error(t, sb.UnmarshalJSON([]byte(`[{"a":1,"b":[2,"bad"]}]`))) + assert.Zero(t, sb.Len()) + assert.Zero(t, sb.FieldBuilder(0).Len()) + assert.Zero(t, sb.FieldBuilder(1).Len()) + + require.NoError(t, sb.UnmarshalJSON([]byte(`[{"a":1,"b":[2,3]}]`))) + + arr := sb.NewArray().(*array.Struct) + defer arr.Release() + require.NoError(t, array.ValidateFull(arr)) + assert.Equal(t, 1, arr.Len()) + assert.Equal(t, `{[1] [[2 3]]}`, arr.String()) +} + func TestCreateStructWithNulls(t *testing.T) { pool := memory.NewCheckedAllocator(memory.NewGoAllocator()) defer pool.AssertSize(t, 0) diff --git a/arrow/array/table_test.go b/arrow/array/table_test.go index d42c0d106..e734d68e3 100644 --- a/arrow/array/table_test.go +++ b/arrow/array/table_test.go @@ -789,7 +789,7 @@ func TestTableFromRecords(t *testing.T) { schema := arrow.NewSchema( []arrow.Field{ - {Name: "f1-i32", Type: arrow.PrimitiveTypes.Int32}, + {Name: "f1-i32", Type: arrow.PrimitiveTypes.Int32, Nullable: true}, {Name: "f2-f64", Type: arrow.PrimitiveTypes.Float64}, }, nil, @@ -1059,7 +1059,7 @@ func TestTableToString(t *testing.T) { schema := arrow.NewSchema( []arrow.Field{ - {Name: "f1-i32", Type: arrow.PrimitiveTypes.Int32}, + {Name: "f1-i32", Type: arrow.PrimitiveTypes.Int32, Nullable: true}, {Name: "f2-f64", Type: arrow.PrimitiveTypes.Float64}, }, nil, @@ -1088,7 +1088,7 @@ func TestTableToString(t *testing.T) { expected_str := `schema: fields: 2 - - f1-i32: type=int32 + - f1-i32: type=int32, nullable - f2-f64: type=float64 f1-i32: [[1 2 3 4 5 6 7 8 (null) 10], [111 112 113 114 115 116 117 118 119 120]] f2-f64: [[11 12 13 14 15 16 17 18 19 20], [211 212 213 214 215 216 217 218 219 220]] diff --git a/arrow/array/union.go b/arrow/array/union.go index 17138f8bd..534be6a9d 100644 --- a/arrow/array/union.go +++ b/arrow/array/union.go @@ -805,6 +805,8 @@ type unionBuilder struct { // for all typeID < denseTypeID, typeIDtoBuilder[typeID] != nil denseTypeID int typesBuilder *int8BufferBuilder + + checkpoint *builderCheckpoint } func unionTypeCodeFromJSON(dec *json.Decoder, typeID json.RawMessage, typ arrow.DataType) (arrow.UnionTypeCode, error) { @@ -1164,6 +1166,19 @@ func (b *SparseUnionBuilder) AppendValueFromString(s string) error { } func (b *SparseUnionBuilder) UnmarshalOne(dec *json.Decoder) error { + if b.checkpoint == nil { + b.checkpoint = newBuilderCheckpoint(b) + } + b.checkpoint.capture() + + if err := b.unmarshalOne(dec); err != nil { + b.checkpoint.restore() + return err + } + return nil +} + +func (b *SparseUnionBuilder) unmarshalOne(dec *json.Decoder) error { t, err := dec.Token() if err != nil { return err @@ -1204,7 +1219,7 @@ func (b *SparseUnionBuilder) UnmarshalOne(dec *json.Decoder) error { } b.Append(typeCode) - if err := b.children[childNum].UnmarshalOne(dec); err != nil { + if err := unmarshalChild(dec, b.children[childNum], b.childFields[childNum]); err != nil { return err } @@ -1425,6 +1440,19 @@ func (d *DenseUnionBuilder) AppendValueFromString(s string) error { } func (b *DenseUnionBuilder) UnmarshalOne(dec *json.Decoder) error { + if b.checkpoint == nil { + b.checkpoint = newBuilderCheckpoint(b) + } + b.checkpoint.capture() + + if err := b.unmarshalOne(dec); err != nil { + b.checkpoint.restore() + return err + } + return nil +} + +func (b *DenseUnionBuilder) unmarshalOne(dec *json.Decoder) error { t, err := dec.Token() if err != nil { return err @@ -1459,7 +1487,7 @@ func (b *DenseUnionBuilder) UnmarshalOne(dec *json.Decoder) error { } b.Append(typeCode) - if err := b.children[childNum].UnmarshalOne(dec); err != nil { + if err := unmarshalChild(dec, b.children[childNum], b.childFields[childNum]); err != nil { return err } diff --git a/arrow/array/union_test.go b/arrow/array/union_test.go index 391348ede..c7c33955c 100644 --- a/arrow/array/union_test.go +++ b/arrow/array/union_test.go @@ -204,8 +204,8 @@ func TestUnionBuilderUnmarshalOnePreservesDecoderConfiguration(t *testing.T) { name string json string }{ - {name: "integer first", json: `{"i":1.5,"u":[0,1]}`}, - {name: "union first", json: `{"u":[0,1],"i":1.5}`}, + {name: "integer first", json: `{"i":2,"u":[0,1]}`}, + {name: "union first", json: `{"u":[0,1],"i":2}`}, } { t.Run(input.name, func(t *testing.T) { builder := array.NewStructBuilder(memory.DefaultAllocator, dtype) @@ -1470,3 +1470,140 @@ func TestNestedUnionDictUnion(t *testing.T) { defer arr.Release() assert.Equal(t, 0, arr.Len()) } + +func TestUnionUnmarshalNonNullableChild(t *testing.T) { + unionFields := func(nullable bool) []arrow.Field { + return []arrow.Field{ + {Name: "u0", Type: arrow.PrimitiveTypes.Int32, Nullable: nullable}, + {Name: "u1", Type: arrow.BinaryTypes.String, Nullable: true}, + } + } + codes := []arrow.UnionTypeCode{0, 1} + + tests := []struct { + name string + dt arrow.DataType + jsonInput string + wantErr string + want string + }{ + { + name: "sparse non-nullable child", + dt: arrow.SparseUnionOf(unionFields(false), codes), + jsonInput: `[[0, null]]`, + wantErr: "field 'u0' is non-nullable but got null", + }, + { + name: "dense non-nullable child", + dt: arrow.DenseUnionOf(unionFields(false), codes), + jsonInput: `[[0, null]]`, + wantErr: "field 'u0' is non-nullable but got null", + }, + { + name: "sparse nullable child", + dt: arrow.SparseUnionOf(unionFields(true), codes), + jsonInput: `[[0, null], [1, "x"]]`, + want: `[{u0=} {u1=x}]`, + }, + { + name: "dense nullable child", + dt: arrow.DenseUnionOf(unionFields(true), codes), + jsonInput: `[[0, null], [1, "x"]]`, + want: `[{u0=} {u1=x}]`, + }, + { + name: "sparse non-nullable child with values", + dt: arrow.SparseUnionOf(unionFields(false), codes), + jsonInput: `[[0, 1], [1, null]]`, + want: `[{u0=1} {u1=}]`, + }, + { + name: "dense non-nullable child with values", + dt: arrow.DenseUnionOf(unionFields(false), codes), + jsonInput: `[[0, 1], [1, null]]`, + want: `[{u0=1} {u1=}]`, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer mem.AssertSize(t, 0) + + arr, _, err := array.FromJSON(mem, tc.dt, strings.NewReader(tc.jsonInput)) + if tc.wantErr != "" { + require.ErrorContains(t, err, tc.wantErr) + return + } + require.NoError(t, err) + defer arr.Release() + + require.NoError(t, array.ValidateFull(arr)) + require.Equal(t, tc.want, arr.String()) + }) + } +} + +func TestUnionBuilderUnmarshalOneRollback(t *testing.T) { + unionFields := []arrow.Field{ + {Name: "u0", Type: arrow.PrimitiveTypes.Int32, Nullable: false}, + {Name: "u1", Type: arrow.BinaryTypes.String, Nullable: true}, + } + codes := []arrow.UnionTypeCode{0, 1} + + tests := []struct { + name string + dt arrow.DataType + bad string + want string + }{ + { + name: "sparse null in non-nullable child", + dt: arrow.SparseUnionOf(unionFields, codes), + bad: `[0, null]`, + want: `[{u0=7}]`, + }, + { + name: "dense null in non-nullable child", + dt: arrow.DenseUnionOf(unionFields, codes), + bad: `[0, null]`, + want: `[{u0=7}]`, + }, + { + name: "sparse wrong child type", + dt: arrow.SparseUnionOf(unionFields, codes), + bad: `[0, "nope"]`, + want: `[{u0=7}]`, + }, + { + name: "dense wrong child type", + dt: arrow.DenseUnionOf(unionFields, codes), + bad: `[0, "nope"]`, + want: `[{u0=7}]`, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + pool := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer pool.AssertSize(t, 0) + + bldr := array.NewBuilder(pool, tc.dt).(array.UnionBuilder) + defer bldr.Release() + + require.Error(t, bldr.UnmarshalOne(internaljson.NewDecoder(strings.NewReader(tc.bad)))) + require.Zero(t, bldr.Len()) + for i := range unionFields { + require.Zero(t, bldr.Child(i).Len(), "child %d", i) + } + + require.NoError(t, bldr.UnmarshalOne(internaljson.NewDecoder(strings.NewReader(`[0, 7]`)))) + + arr := bldr.NewArray() + defer arr.Release() + + require.NoError(t, array.ValidateFull(arr)) + require.Equal(t, tc.want, arr.String()) + }) + } +} diff --git a/arrow/array/util.go b/arrow/array/util.go index ca6929119..07da71197 100644 --- a/arrow/array/util.go +++ b/arrow/array/util.go @@ -296,7 +296,11 @@ func RecordToJSON(rec arrow.RecordBatch, w io.Writer) error { cols := make(map[string]interface{}) for i := 0; int64(i) < rec.NumRows(); i++ { for j, c := range rec.Columns() { - cols[fields[j].Name] = c.GetOneForMarshal(i) + if rec.Schema().Field(j).Nullable && c.IsNull(i) { + cols[fields[j].Name] = nil + } else { + cols[fields[j].Name] = c.GetOneForMarshal(i) + } } if err := enc.Encode(cols); err != nil { return err diff --git a/arrow/array/validate_test.go b/arrow/array/validate_test.go index ac6417872..5dda77a92 100644 --- a/arrow/array/validate_test.go +++ b/arrow/array/validate_test.go @@ -440,7 +440,7 @@ func TestTopLevelValidate(t *testing.T) { t.Run("ValidateRecord validates all columns", func(t *testing.T) { validArr := makeStringArrayRaw(t, []int32{0, 3, 6}, "abcdef", 2, 0) - corruptArr := makeStringArrayRaw(t, []int32{0, 5, 3, 5}, "hello", 3, 0) + corruptArr := makeStringArrayRaw(t, []int32{0, 4, 3}, "hello", 2, 0) schema := arrow.NewSchema([]arrow.Field{ {Name: "ok", Type: arrow.BinaryTypes.String}, diff --git a/arrow/compute/vector_sort_test.go b/arrow/compute/vector_sort_test.go index 39bf5e95f..fb87af1b4 100644 --- a/arrow/compute/vector_sort_test.go +++ b/arrow/compute/vector_sort_test.go @@ -1373,8 +1373,8 @@ func TestVectorSortIndicesCppRecordBatchParity(t *testing.T) { t.Run("Null", func(t *testing.T) { schema := arrow.NewSchema([]arrow.Field{ - {Name: "a", Type: arrow.PrimitiveTypes.Uint8}, - {Name: "b", Type: arrow.PrimitiveTypes.Uint32}, + {Name: "a", Type: arrow.PrimitiveTypes.Uint8, Nullable: true}, + {Name: "b", Type: arrow.PrimitiveTypes.Uint32, Nullable: true}, }, nil) jsonRows := `[ {"a": null, "b": 5}, @@ -1426,8 +1426,8 @@ func TestVectorSortIndicesCppRecordBatchParity(t *testing.T) { t.Run("NaNAndNull", func(t *testing.T) { schema := arrow.NewSchema([]arrow.Field{ - {Name: "a", Type: arrow.PrimitiveTypes.Float32}, - {Name: "b", Type: arrow.PrimitiveTypes.Float64}, + {Name: "a", Type: arrow.PrimitiveTypes.Float32, Nullable: true}, + {Name: "b", Type: arrow.PrimitiveTypes.Float64, Nullable: true}, }, nil) ba := array.NewFloat32Builder(mem) defer ba.Release() @@ -1460,8 +1460,8 @@ func TestVectorSortIndicesCppRecordBatchParity(t *testing.T) { t.Run("Boolean", func(t *testing.T) { schema := arrow.NewSchema([]arrow.Field{ - {Name: "a", Type: arrow.FixedWidthTypes.Boolean}, - {Name: "b", Type: arrow.FixedWidthTypes.Boolean}, + {Name: "a", Type: arrow.FixedWidthTypes.Boolean, Nullable: true}, + {Name: "b", Type: arrow.FixedWidthTypes.Boolean, Nullable: true}, }, nil) jsonRows := `[ {"a": true, "b": null}, @@ -1536,7 +1536,7 @@ func TestVectorSortIndicesCppRecordBatchParity(t *testing.T) { d256 := &arrow.Decimal256Type{Precision: 4, Scale: 2} schema := arrow.NewSchema([]arrow.Field{ {Name: "a", Type: d128}, - {Name: "b", Type: d256}, + {Name: "b", Type: d256, Nullable: true}, }, nil) jsonRows := `[ {"a": "12.3", "b": "12.34"}, @@ -1561,8 +1561,8 @@ func TestVectorSortIndicesCppRecordBatchParity(t *testing.T) { t.Run("DuplicateSortKeys", func(t *testing.T) { schema := arrow.NewSchema([]arrow.Field{ - {Name: "a", Type: arrow.PrimitiveTypes.Float32}, - {Name: "b", Type: arrow.PrimitiveTypes.Float64}, + {Name: "a", Type: arrow.PrimitiveTypes.Float32, Nullable: true}, + {Name: "b", Type: arrow.PrimitiveTypes.Float64, Nullable: true}, }, nil) ba := array.NewFloat32Builder(mem) defer ba.Release() @@ -1610,8 +1610,8 @@ func TestVectorSortIndicesCppTableParity(t *testing.T) { ctx := context.Background() schemaAB := arrow.NewSchema([]arrow.Field{ - {Name: "a", Type: arrow.PrimitiveTypes.Uint8}, - {Name: "b", Type: arrow.PrimitiveTypes.Uint32}, + {Name: "a", Type: arrow.PrimitiveTypes.Uint8, Nullable: true}, + {Name: "b", Type: arrow.PrimitiveTypes.Uint32, Nullable: true}, }, nil) t.Run("EmptyTable", func(t *testing.T) { @@ -1667,8 +1667,8 @@ func TestVectorSortIndicesCppTableParity(t *testing.T) { t.Run("BinaryLikeTwoChunks", func(t *testing.T) { fsb3 := &arrow.FixedSizeBinaryType{ByteWidth: 3} s := arrow.NewSchema([]arrow.Field{ - {Name: "a", Type: arrow.BinaryTypes.LargeString}, - {Name: "b", Type: fsb3}, + {Name: "a", Type: arrow.BinaryTypes.LargeString, Nullable: true}, + {Name: "b", Type: fsb3, Nullable: true}, }, nil) buildBatch := func(a []string, b [][]byte, bNulls []bool) arrow.RecordBatch { ab := array.NewLargeStringBuilder(mem) diff --git a/arrow/example_test.go b/arrow/example_test.go index 2b0dc5d4f..c4ca468eb 100644 --- a/arrow/example_test.go +++ b/arrow/example_test.go @@ -261,7 +261,7 @@ func Example_structArray() { pool := memory.NewGoAllocator() dtype := arrow.StructOf([]arrow.Field{ - {Name: "f1", Type: arrow.ListOf(arrow.PrimitiveTypes.Uint8)}, + {Name: "f1", Type: arrow.ListOf(arrow.PrimitiveTypes.Uint8), Nullable: true}, {Name: "f2", Type: arrow.PrimitiveTypes.Int32}, }...) @@ -331,7 +331,7 @@ func Example_structArray() { // Output: // NullN() = 1 // Len() = 4 - // Type() = struct, f2: int32> + // Type() = struct nullable, f2: int32> // Struct[0] = [[j, o, e], 1] // Struct[1] = [[], 2] // Struct[2] = (null) @@ -463,7 +463,7 @@ func Example_record() { schema := arrow.NewSchema( []arrow.Field{ - {Name: "f1-i32", Type: arrow.PrimitiveTypes.Int32}, + {Name: "f1-i32", Type: arrow.PrimitiveTypes.Int32, Nullable: true}, {Name: "f2-f64", Type: arrow.PrimitiveTypes.Float64}, }, nil, @@ -493,7 +493,7 @@ func Example_recordReader() { schema := arrow.NewSchema( []arrow.Field{ - {Name: "f1-i32", Type: arrow.PrimitiveTypes.Int32}, + {Name: "f1-i32", Type: arrow.PrimitiveTypes.Int32, Nullable: true}, {Name: "f2-f64", Type: arrow.PrimitiveTypes.Float64}, }, nil, @@ -542,7 +542,7 @@ func Example_table() { schema := arrow.NewSchema( []arrow.Field{ - {Name: "f1-i32", Type: arrow.PrimitiveTypes.Int32}, + {Name: "f1-i32", Type: arrow.PrimitiveTypes.Int32, Nullable: true}, {Name: "f2-f64", Type: arrow.PrimitiveTypes.Float64}, }, nil, diff --git a/arrow/extensions/uuid_test.go b/arrow/extensions/uuid_test.go index a76b77a91..36bb812b9 100644 --- a/arrow/extensions/uuid_test.go +++ b/arrow/extensions/uuid_test.go @@ -62,7 +62,7 @@ func TestUUIDExtensionBuilder(t *testing.T) { func TestUUIDExtensionRecordBuilder(t *testing.T) { schema := arrow.NewSchema([]arrow.Field{ - {Name: "uuid", Type: extensions.NewUUIDType()}, + {Name: "uuid", Type: extensions.NewUUIDType(), Nullable: true}, }, nil) builder := array.NewRecordBuilder(memory.DefaultAllocator, schema) builder.Field(0).(*extensions.UUIDBuilder).Append(testUUID) diff --git a/arrow/internal/arrdata/arrdata.go b/arrow/internal/arrdata/arrdata.go index b4f0e626b..df50ada17 100644 --- a/arrow/internal/arrdata/arrdata.go +++ b/arrow/internal/arrdata/arrdata.go @@ -192,8 +192,8 @@ func makeStructsRecords() []arrow.RecordBatch { mem := memory.NewGoAllocator() fields := []arrow.Field{ - {Name: "f1", Type: arrow.PrimitiveTypes.Int32}, - {Name: "f2", Type: arrow.BinaryTypes.String}, + {Name: "f1", Type: arrow.PrimitiveTypes.Int32, Nullable: true}, + {Name: "f2", Type: arrow.BinaryTypes.String, Nullable: true}, } dtype := arrow.StructOf(fields...) schema := arrow.NewSchema([]arrow.Field{{Name: "struct_nullable", Type: dtype, Nullable: true}}, nil) @@ -433,8 +433,8 @@ func makeFixedSizeListsRecords() []arrow.RecordBatch { func makeStringsRecords() []arrow.RecordBatch { mem := memory.NewGoAllocator() schema := arrow.NewSchema([]arrow.Field{ - {Name: "strings", Type: arrow.BinaryTypes.String}, - {Name: "bytes", Type: arrow.BinaryTypes.Binary}, + {Name: "strings", Type: arrow.BinaryTypes.String, Nullable: true}, + {Name: "bytes", Type: arrow.BinaryTypes.Binary, Nullable: true}, }, nil) mask := []bool{true, false, false, true, true} @@ -1110,9 +1110,21 @@ func makeRunEndEncodedRecords() []arrow.RecordBatch { ree32Type := arrow.RunEndEncodedOf(arrow.PrimitiveTypes.Int32, arrow.PrimitiveTypes.Int32) ree32Type.ValueNullable = false schema := arrow.NewSchema([]arrow.Field{ - {Name: "ree16", Type: arrow.RunEndEncodedOf(arrow.PrimitiveTypes.Int16, arrow.BinaryTypes.String)}, - {Name: "ree32", Type: ree32Type}, - {Name: "ree64", Type: arrow.RunEndEncodedOf(arrow.PrimitiveTypes.Int64, arrow.BinaryTypes.Binary)}, + { + Name: "ree16", + Type: arrow.RunEndEncodedOf(arrow.PrimitiveTypes.Int16, arrow.BinaryTypes.String), + Nullable: true, + }, + { + Name: "ree32", + Type: ree32Type, + Nullable: true, + }, + { + Name: "ree64", + Type: arrow.RunEndEncodedOf(arrow.PrimitiveTypes.Int64, arrow.BinaryTypes.Binary), + Nullable: true, + }, }, nil) isValid := []bool{true, false, true, false, true} chunks := [][]arrow.Array{ diff --git a/arrow/internal/arrjson/arrjson_test.go b/arrow/internal/arrjson/arrjson_test.go index 232270b9d..5115c3f33 100644 --- a/arrow/internal/arrjson/arrjson_test.go +++ b/arrow/internal/arrjson/arrjson_test.go @@ -948,7 +948,7 @@ func makeStructsWantJSONs() string { "isSigned": true, "bitWidth": 32 }, - "nullable": false, + "nullable": true, "children": [] }, { @@ -956,7 +956,7 @@ func makeStructsWantJSONs() string { "type": { "name": "utf8" }, - "nullable": false, + "nullable": true, "children": [] } ] @@ -1957,7 +1957,7 @@ func makeStringsWantJSONs() string { "type": { "name": "utf8" }, - "nullable": false, + "nullable": true, "children": [] }, { @@ -1965,7 +1965,7 @@ func makeStringsWantJSONs() string { "type": { "name": "binary" }, - "nullable": false, + "nullable": true, "children": [] } ] @@ -5765,7 +5765,7 @@ func makeRunEndEncodedWantJSONs() string { "type": { "name": "runendencoded" }, - "nullable": false, + "nullable": true, "children": [ { "name": "run_ends", @@ -5792,7 +5792,7 @@ func makeRunEndEncodedWantJSONs() string { "type": { "name": "runendencoded" }, - "nullable": false, + "nullable": true, "children": [ { "name": "run_ends", @@ -5821,7 +5821,7 @@ func makeRunEndEncodedWantJSONs() string { "type": { "name": "runendencoded" }, - "nullable": false, + "nullable": true, "children": [ { "name": "run_ends", diff --git a/arrow/ipc/cmd/arrow-ls/main_test.go b/arrow/ipc/cmd/arrow-ls/main_test.go index cb1ab8fd6..98f3cf8b1 100644 --- a/arrow/ipc/cmd/arrow-ls/main_test.go +++ b/arrow/ipc/cmd/arrow-ls/main_test.go @@ -59,7 +59,7 @@ records: 3 name: "structs", want: `schema: fields: 1 - - struct_nullable: type=struct, nullable + - struct_nullable: type=struct, nullable records: 2 `, }, @@ -75,8 +75,8 @@ records: 4 name: "strings", want: `schema: fields: 2 - - strings: type=utf8 - - bytes: type=binary + - strings: type=utf8, nullable + - bytes: type=binary, nullable records: 3 `, }, @@ -249,7 +249,7 @@ records: 3 name: "structs", want: `schema: fields: 1 - - struct_nullable: type=struct, nullable + - struct_nullable: type=struct, nullable records: 2 `, }, @@ -258,7 +258,7 @@ records: 2 want: `version: V5 schema: fields: 1 - - struct_nullable: type=struct, nullable + - struct_nullable: type=struct, nullable records: 2 `, }, diff --git a/arrow/ipc/ipc_test.go b/arrow/ipc/ipc_test.go index f1f310bd0..27ef919f1 100644 --- a/arrow/ipc/ipc_test.go +++ b/arrow/ipc/ipc_test.go @@ -316,7 +316,7 @@ func TestWriteColumnWithOffset(t *testing.T) { func TestIPCTable(t *testing.T) { pool := memory.NewGoAllocator() - schema := arrow.NewSchema([]arrow.Field{{Name: "f1", Type: arrow.PrimitiveTypes.Int32}}, nil) + schema := arrow.NewSchema([]arrow.Field{{Name: "f1", Type: arrow.PrimitiveTypes.Int32, Nullable: true}}, nil) b := array.NewRecordBuilder(pool, schema) defer b.Release() b.Field(0).(*array.Int32Builder).AppendValues([]int32{1, 2, 3, 4}, []bool{true, true, false, true}) diff --git a/arrow/ipc/reader_test.go b/arrow/ipc/reader_test.go index 29db5985a..8883f3ba8 100644 --- a/arrow/ipc/reader_test.go +++ b/arrow/ipc/reader_test.go @@ -100,7 +100,7 @@ func TestReaderCheckedAllocator(t *testing.T) { func TestMappedReader(t *testing.T) { pool := memory.NewCheckedAllocator(memory.NewGoAllocator()) defer pool.AssertSize(t, 0) - schema := arrow.NewSchema([]arrow.Field{{Name: "f1", Type: arrow.PrimitiveTypes.Int32}}, nil) + schema := arrow.NewSchema([]arrow.Field{{Name: "f1", Type: arrow.PrimitiveTypes.Int32, Nullable: true}}, nil) b := array.NewRecordBuilder(pool, schema) defer b.Release() b.Field(0).(*array.Int32Builder).AppendValues([]int32{1, 2, 3, 4}, []bool{true, true, false, true})