fix(arrow/array): Parity with Arrow C++ when reading/writing null to/from JSON

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() {

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

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)
+}

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 ']'

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 ']'