diff --git a/arrow/array.go b/arrow/array.go index d42ca6d05..09fafe42f 100644 --- a/arrow/array.go +++ b/arrow/array.go @@ -128,6 +128,21 @@ type Array interface { Release() } +// NullableMarshaler is an optional interface an Array may implement to control +// whether a value is rendered as JSON null based on the field-local nullability +// from the schema, rather than solely on the array's own validity bitmap. +// +// It exists so containers (struct/union/dictionary/run-end-encoded/extension) +// and RecordToJSON can propagate each field's Nullable flag down to its values +// without a breaking change to the Array interface. Arrays that do not implement +// it fall back to GetOneForMarshal, preserving the previous validity-based behavior. +type NullableMarshaler interface { + // GetOneForMarshalNullable returns the value at i for json.Marshal. When + // nullable is false, a value is returned even if the array's validity bitmap + // marks it null (a non-nullable field must not serialize as null). + GetOneForMarshalNullable(i int, nullable bool) interface{} +} + // ValueType is a generic constraint for valid Arrow primitive types type ValueType interface { bool | FixedWidthType | string | []byte diff --git a/arrow/array/binary.go b/arrow/array/binary.go index a8e77ae9c..1801d9d64 100644 --- a/arrow/array/binary.go +++ b/arrow/array/binary.go @@ -152,13 +152,17 @@ func (a *Binary) setData(data *Data) { } } -func (a *Binary) GetOneForMarshal(i int) interface{} { - if a.IsNull(i) { +func (a *Binary) GetOneForMarshalNullable(i int, nullable bool) interface{} { + if nullable && a.IsNull(i) { return nil } return a.Value(i) } +func (a *Binary) GetOneForMarshal(i int) interface{} { + return a.GetOneForMarshalNullable(i, true) +} + func (a *Binary) MarshalJSON() ([]byte, error) { vals := make([]interface{}, a.Len()) for i := 0; i < a.Len(); i++ { @@ -223,9 +227,9 @@ func (a *Binary) ValidateFull() error { return nil } -func arrayEqualBinary(left, right *Binary) bool { +func arrayEqualBinary(left, right *Binary, opt equalOption) bool { for i := 0; i < left.Len(); i++ { - if left.IsNull(i) { + if opt.nullable && left.IsNull(i) { continue } if !bytes.Equal(left.Value(i), right.Value(i)) { @@ -346,13 +350,17 @@ func (a *LargeBinary) setData(data *Data) { } } -func (a *LargeBinary) GetOneForMarshal(i int) interface{} { - if a.IsNull(i) { +func (a *LargeBinary) GetOneForMarshalNullable(i int, nullable bool) interface{} { + if nullable && a.IsNull(i) { return nil } return a.Value(i) } +func (a *LargeBinary) GetOneForMarshal(i int) interface{} { + return a.GetOneForMarshalNullable(i, true) +} + func (a *LargeBinary) MarshalJSON() ([]byte, error) { vals := make([]interface{}, a.Len()) for i := 0; i < a.Len(); i++ { @@ -417,9 +425,9 @@ func (a *LargeBinary) ValidateFull() error { return nil } -func arrayEqualLargeBinary(left, right *LargeBinary) bool { +func arrayEqualLargeBinary(left, right *LargeBinary, opt equalOption) bool { for i := 0; i < left.Len(); i++ { - if left.IsNull(i) { + if opt.nullable && left.IsNull(i) { continue } if !bytes.Equal(left.Value(i), right.Value(i)) { @@ -522,13 +530,17 @@ func (a *BinaryView) ValueStr(i int) string { return base64.StdEncoding.EncodeToString(a.Value(i)) } -func (a *BinaryView) GetOneForMarshal(i int) interface{} { - if a.IsNull(i) { +func (a *BinaryView) GetOneForMarshalNullable(i int, nullable bool) interface{} { + if nullable && a.IsNull(i) { return nil } return a.Value(i) } +func (a *BinaryView) GetOneForMarshal(i int) interface{} { + return a.GetOneForMarshalNullable(i, true) +} + func (a *BinaryView) MarshalJSON() ([]byte, error) { vals := make([]interface{}, a.Len()) for i := 0; i < a.Len(); i++ { @@ -539,10 +551,10 @@ func (a *BinaryView) MarshalJSON() ([]byte, error) { return json.Marshal(vals) } -func arrayEqualBinaryView(left, right *BinaryView) bool { +func arrayEqualBinaryView(left, right *BinaryView, opt equalOption) bool { leftBufs, rightBufs := left.dataBuffers, right.dataBuffers for i := 0; i < left.Len(); i++ { - if left.IsNull(i) { + if opt.nullable && left.IsNull(i) { continue } if !left.ValueHeader(i).Equals(leftBufs, right.ValueHeader(i), rightBufs) { diff --git a/arrow/array/boolean.go b/arrow/array/boolean.go index d579fa0c8..183d4955a 100644 --- a/arrow/array/boolean.go +++ b/arrow/array/boolean.go @@ -90,13 +90,17 @@ func (a *Boolean) setData(data *Data) { } } -func (a *Boolean) GetOneForMarshal(i int) interface{} { - if a.IsValid(i) { +func (a *Boolean) GetOneForMarshalNullable(i int, nullable bool) interface{} { + if !nullable || a.IsValid(i) { return a.Value(i) } return nil } +func (a *Boolean) GetOneForMarshal(i int) interface{} { + return a.GetOneForMarshalNullable(i, true) +} + func (a *Boolean) MarshalJSON() ([]byte, error) { vals := make([]interface{}, a.Len()) for i := 0; i < a.Len(); i++ { @@ -109,9 +113,9 @@ func (a *Boolean) MarshalJSON() ([]byte, error) { return json.Marshal(vals) } -func arrayEqualBoolean(left, right *Boolean) bool { +func arrayEqualBoolean(left, right *Boolean, opt equalOption) bool { for i := 0; i < left.Len(); i++ { - if left.IsNull(i) { + if opt.nullable && left.IsNull(i) { continue } if left.Value(i) != right.Value(i) { diff --git a/arrow/array/compare.go b/arrow/array/compare.go index d8a9552e9..119a77596 100644 --- a/arrow/array/compare.go +++ b/arrow/array/compare.go @@ -25,8 +25,23 @@ import ( "github.com/apache/arrow-go/v18/internal/bitutils" ) +// schemaFieldEqual reports whether two fields agree on the schema-level +// attributes that record/table equality must match between the two sides: name +// and nullability. The field type and values are validated per column by the +// array comparison (baseArrayEqual), so they are not re-checked here, and field +// metadata is intentionally not compared so records that differ only in metadata +// (for example a PARQUET:field_id attached during a round-trip) still compare as +// equal, matching the pre-existing behavior of these comparison APIs. +func schemaFieldEqual(l, r arrow.Field) bool { + return l.Name == r.Name && l.Nullable == r.Nullable +} + // RecordEqual reports whether the two provided records are equal. -func RecordEqual(left, right arrow.RecordBatch) bool { +func RecordEqual(left, right arrow.RecordBatch, opts ...EqualOption) bool { + return recordEqual(left, right, newEqualOption(opts...)) +} + +func recordEqual(left, right arrow.RecordBatch, opt equalOption) bool { switch { case left.NumCols() != right.NumCols(): return false @@ -35,9 +50,17 @@ func RecordEqual(left, right arrow.RecordBatch) bool { } for i := range left.Columns() { + lf := left.Schema().Field(i) + rf := right.Schema().Field(i) + if !schemaFieldEqual(lf, rf) { + return false + } + + opt.nullable = lf.Nullable + lc := left.Column(i) rc := right.Column(i) - if !Equal(lc, rc) { + if !equal(lc, rc, opt) { return false } } @@ -47,6 +70,10 @@ func RecordEqual(left, right arrow.RecordBatch) bool { // RecordApproxEqual reports whether the two provided records are approximately equal. // For non-floating point columns, it is equivalent to RecordEqual. func RecordApproxEqual(left, right arrow.RecordBatch, opts ...EqualOption) bool { + return recordApproxEqual(left, right, newEqualOption(opts...)) +} + +func recordApproxEqual(left, right arrow.RecordBatch, opt equalOption) bool { switch { case left.NumCols() != right.NumCols(): return false @@ -54,9 +81,15 @@ func RecordApproxEqual(left, right arrow.RecordBatch, opts ...EqualOption) bool return false } - opt := newEqualOption(opts...) - for i := range left.Columns() { + lf := left.Schema().Field(i) + rf := right.Schema().Field(i) + if !schemaFieldEqual(lf, rf) { + return false + } + + opt.nullable = lf.Nullable + lc := left.Column(i) rc := right.Column(i) if !arrayApproxEqual(lc, rc, opt) { @@ -106,13 +139,17 @@ func chunkedBinaryApply(left, right *arrow.Chunked, fn func(left arrow.Array, lb } // ChunkedEqual reports whether two chunked arrays are equal regardless of their chunkings -func ChunkedEqual(left, right *arrow.Chunked) bool { +func ChunkedEqual(left, right *arrow.Chunked, opts ...EqualOption) bool { + return chunkedEqual(left, right, newEqualOption(opts...)) +} + +func chunkedEqual(left, right *arrow.Chunked, opt equalOption) bool { switch { case left == right: return true case left.Len() != right.Len(): return false - case left.NullN() != right.NullN(): + case opt.nullable && hasTopLevelValidityBitmap(left.DataType().ID()) && left.NullN() != right.NullN(): return false case !arrow.TypeEqual(left.DataType(), right.DataType()): return false @@ -120,7 +157,7 @@ func ChunkedEqual(left, right *arrow.Chunked) bool { var isequal = true chunkedBinaryApply(left, right, func(left arrow.Array, lbeg, lend int64, right arrow.Array, rbeg, rend int64) bool { - isequal = SliceEqual(left, lbeg, lend, right, rbeg, rend) + isequal = sliceEqual(left, lbeg, lend, right, rbeg, rend, opt) return isequal }) @@ -130,12 +167,16 @@ func ChunkedEqual(left, right *arrow.Chunked) bool { // ChunkedApproxEqual reports whether two chunked arrays are approximately equal regardless of their chunkings // for non-floating point arrays, this is equivalent to ChunkedEqual func ChunkedApproxEqual(left, right *arrow.Chunked, opts ...EqualOption) bool { + return chunkedApproxEqual(left, right, newEqualOption(opts...)) +} + +func chunkedApproxEqual(left, right *arrow.Chunked, opt equalOption) bool { switch { case left == right: return true case left.Len() != right.Len(): return false - case left.NullN() != right.NullN(): + case opt.nullable && hasTopLevelValidityBitmap(left.DataType().ID()) && left.NullN() != right.NullN(): return false case !arrow.TypeEqual(left.DataType(), right.DataType()): return false @@ -143,7 +184,7 @@ func ChunkedApproxEqual(left, right *arrow.Chunked, opts ...EqualOption) bool { var isequal = true chunkedBinaryApply(left, right, func(left arrow.Array, lbeg, lend int64, right arrow.Array, rbeg, rend int64) bool { - isequal = SliceApproxEqual(left, lbeg, lend, right, rbeg, rend, opts...) + isequal = sliceApproxEqual(left, lbeg, lend, right, rbeg, rend, opt) return isequal }) @@ -151,7 +192,11 @@ func ChunkedApproxEqual(left, right *arrow.Chunked, opts ...EqualOption) bool { } // TableEqual returns if the two tables have the same data in the same schema -func TableEqual(left, right arrow.Table) bool { +func TableEqual(left, right arrow.Table, opts ...EqualOption) bool { + return tableEqual(left, right, newEqualOption(opts...)) +} + +func tableEqual(left, right arrow.Table, opt equalOption) bool { switch { case left.NumCols() != right.NumCols(): return false @@ -160,21 +205,29 @@ func TableEqual(left, right arrow.Table) bool { } for i := 0; int64(i) < left.NumCols(); i++ { - lc := left.Column(i) - rc := right.Column(i) - if !lc.Field().Equal(rc.Field()) { + lf := left.Schema().Field(i) + rf := right.Schema().Field(i) + if !schemaFieldEqual(lf, rf) { return false } - if !ChunkedEqual(lc.Data(), rc.Data()) { + opt.nullable = lf.Nullable + + lc := left.Column(i) + rc := right.Column(i) + if !chunkedEqual(lc.Data(), rc.Data(), opt) { return false } } return true } -// TableEqual returns if the two tables have the approximately equal data in the same schema +// TableApproxEqual returns if the two tables have the approximately equal data in the same schema func TableApproxEqual(left, right arrow.Table, opts ...EqualOption) bool { + return tableApproxEqual(left, right, newEqualOption(opts...)) +} + +func tableApproxEqual(left, right arrow.Table, opt equalOption) bool { switch { case left.NumCols() != right.NumCols(): return false @@ -183,13 +236,17 @@ func TableApproxEqual(left, right arrow.Table, opts ...EqualOption) bool { } for i := 0; int64(i) < left.NumCols(); i++ { - lc := left.Column(i) - rc := right.Column(i) - if !lc.Field().Equal(rc.Field()) { + lf := left.Schema().Field(i) + rf := right.Schema().Field(i) + if !schemaFieldEqual(lf, rf) { return false } - if !ChunkedApproxEqual(lc.Data(), rc.Data(), opts...) { + opt.nullable = lf.Nullable + + lc := left.Column(i) + rc := right.Column(i) + if !chunkedApproxEqual(lc.Data(), rc.Data(), opt) { return false } } @@ -197,13 +254,17 @@ func TableApproxEqual(left, right arrow.Table, opts ...EqualOption) bool { } // Equal reports whether the two provided arrays are equal. -func Equal(left, right arrow.Array) bool { +func Equal(left, right arrow.Array, opts ...EqualOption) bool { + return equal(left, right, newEqualOption(opts...)) +} + +func equal(left, right arrow.Array, opt equalOption) bool { switch { - case !baseArrayEqual(left, right): + case !baseArrayEqual(left, right, opt): return false case left.Len() == 0: return true - case left.NullN() == left.Len(): + case opt.nullable && hasTopLevelValidityBitmap(left.DataType().ID()) && left.NullN() == left.Len(): return true } @@ -216,155 +277,158 @@ func Equal(left, right arrow.Array) bool { return true case *Boolean: r := right.(*Boolean) - return arrayEqualBoolean(l, r) + return arrayEqualBoolean(l, r, opt) case *FixedSizeBinary: r := right.(*FixedSizeBinary) - return arrayEqualFixedSizeBinary(l, r) + return arrayEqualFixedSizeBinary(l, r, opt) case *Binary: r := right.(*Binary) - return arrayEqualBinary(l, r) + return arrayEqualBinary(l, r, opt) case *String: r := right.(*String) - return arrayEqualString(l, r) + return arrayEqualString(l, r, opt) case *LargeBinary: r := right.(*LargeBinary) - return arrayEqualLargeBinary(l, r) + return arrayEqualLargeBinary(l, r, opt) case *LargeString: r := right.(*LargeString) - return arrayEqualLargeString(l, r) + return arrayEqualLargeString(l, r, opt) case *BinaryView: r := right.(*BinaryView) - return arrayEqualBinaryView(l, r) + return arrayEqualBinaryView(l, r, opt) case *StringView: r := right.(*StringView) - return arrayEqualStringView(l, r) + return arrayEqualStringView(l, r, opt) case *Int8: r := right.(*Int8) - return arrayEqualFixedWidth(l, r) + return arrayEqualFixedWidth(l, r, opt) case *Int16: r := right.(*Int16) - return arrayEqualFixedWidth(l, r) + return arrayEqualFixedWidth(l, r, opt) case *Int32: r := right.(*Int32) - return arrayEqualFixedWidth(l, r) + return arrayEqualFixedWidth(l, r, opt) case *Int64: r := right.(*Int64) - return arrayEqualFixedWidth(l, r) + return arrayEqualFixedWidth(l, r, opt) case *Uint8: r := right.(*Uint8) - return arrayEqualFixedWidth(l, r) + return arrayEqualFixedWidth(l, r, opt) case *Uint16: r := right.(*Uint16) - return arrayEqualFixedWidth(l, r) + return arrayEqualFixedWidth(l, r, opt) case *Uint32: r := right.(*Uint32) - return arrayEqualFixedWidth(l, r) + return arrayEqualFixedWidth(l, r, opt) case *Uint64: r := right.(*Uint64) - return arrayEqualFixedWidth(l, r) + return arrayEqualFixedWidth(l, r, opt) case *Float16: r := right.(*Float16) - return arrayEqualFixedWidth(l, r) + return arrayEqualFixedWidth(l, r, opt) case *Float32: r := right.(*Float32) - return arrayEqualFixedWidth(l, r) + return arrayEqualFixedWidth(l, r, opt) case *Float64: r := right.(*Float64) - return arrayEqualFixedWidth(l, r) + return arrayEqualFixedWidth(l, r, opt) case *Decimal32: r := right.(*Decimal32) - return arrayEqualDecimal(l, r) + return arrayEqualDecimal(l, r, opt) case *Decimal64: r := right.(*Decimal64) - return arrayEqualDecimal(l, r) + return arrayEqualDecimal(l, r, opt) case *Decimal128: r := right.(*Decimal128) - return arrayEqualDecimal(l, r) + return arrayEqualDecimal(l, r, opt) case *Decimal256: r := right.(*Decimal256) - return arrayEqualDecimal(l, r) + return arrayEqualDecimal(l, r, opt) case *Date32: r := right.(*Date32) - return arrayEqualFixedWidth(l, r) + return arrayEqualFixedWidth(l, r, opt) case *Date64: r := right.(*Date64) - return arrayEqualFixedWidth(l, r) + return arrayEqualFixedWidth(l, r, opt) case *Time32: r := right.(*Time32) - return arrayEqualFixedWidth(l, r) + return arrayEqualFixedWidth(l, r, opt) case *Time64: r := right.(*Time64) - return arrayEqualFixedWidth(l, r) + return arrayEqualFixedWidth(l, r, opt) case *Timestamp: r := right.(*Timestamp) - return arrayEqualTimestamp(l, r) + return arrayEqualTimestamp(l, r, opt) case *List: r := right.(*List) - return arrayEqualList(l, r) + return arrayEqualList(l, r, opt) case *LargeList: r := right.(*LargeList) - return arrayEqualLargeList(l, r) + return arrayEqualLargeList(l, r, opt) case *ListView: r := right.(*ListView) - return arrayEqualListView(l, r) + return arrayEqualListView(l, r, opt) case *LargeListView: r := right.(*LargeListView) - return arrayEqualLargeListView(l, r) + return arrayEqualLargeListView(l, r, opt) case *FixedSizeList: r := right.(*FixedSizeList) - return arrayEqualFixedSizeList(l, r) + return arrayEqualFixedSizeList(l, r, opt) case *Struct: r := right.(*Struct) - return arrayEqualStruct(l, r) + return arrayEqualStruct(l, r, opt) case *MonthInterval: r := right.(*MonthInterval) - return arrayEqualMonthInterval(l, r) + return arrayEqualMonthInterval(l, r, opt) case *DayTimeInterval: r := right.(*DayTimeInterval) - return arrayEqualDayTimeInterval(l, r) + return arrayEqualDayTimeInterval(l, r, opt) case *MonthDayNanoInterval: r := right.(*MonthDayNanoInterval) - return arrayEqualMonthDayNanoInterval(l, r) + return arrayEqualMonthDayNanoInterval(l, r, opt) case *Duration: r := right.(*Duration) - return arrayEqualFixedWidth(l, r) + return arrayEqualFixedWidth(l, r, opt) case *Map: r := right.(*Map) - return arrayEqualMap(l, r) + return arrayEqualMap(l, r, opt) case ExtensionArray: r := right.(ExtensionArray) - return arrayEqualExtension(l, r) + return arrayEqualExtension(l, r, opt) case *Dictionary: r := right.(*Dictionary) - return arrayEqualDict(l, r) + return arrayEqualDict(l, r, opt) case *SparseUnion: r := right.(*SparseUnion) - return arraySparseUnionEqual(l, r) + return arraySparseUnionEqual(l, r, opt) case *DenseUnion: r := right.(*DenseUnion) - return arrayDenseUnionEqual(l, r) + return arrayDenseUnionEqual(l, r, opt) case *RunEndEncoded: r := right.(*RunEndEncoded) - return arrayRunEndEncodedEqual(l, r) + return arrayRunEndEncodedEqual(l, r, opt) default: panic(fmt.Errorf("arrow/array: unknown array type %T", l)) } } // SliceEqual reports whether slices left[lbeg:lend] and right[rbeg:rend] are equal. -func SliceEqual(left arrow.Array, lbeg, lend int64, right arrow.Array, rbeg, rend int64) bool { +func SliceEqual(left arrow.Array, lbeg, lend int64, right arrow.Array, rbeg, rend int64, opts ...EqualOption) bool { + return sliceEqual(left, lbeg, lend, right, rbeg, rend, newEqualOption(opts...)) +} + +func sliceEqual(left arrow.Array, lbeg, lend int64, right arrow.Array, rbeg, rend int64, opt equalOption) bool { l := NewSlice(left, lbeg, lend) defer l.Release() r := NewSlice(right, rbeg, rend) defer r.Release() - return Equal(l, r) + return equal(l, r, opt) } // SliceApproxEqual reports whether slices left[lbeg:lend] and right[rbeg:rend] are approximately equal. func SliceApproxEqual(left arrow.Array, lbeg, lend int64, right arrow.Array, rbeg, rend int64, opts ...EqualOption) bool { - opt := newEqualOption(opts...) - return sliceApproxEqual(left, lbeg, lend, right, rbeg, rend, opt) + return sliceApproxEqual(left, lbeg, lend, right, rbeg, rend, newEqualOption(opts...)) } func sliceApproxEqual(left arrow.Array, lbeg, lend int64, right arrow.Array, rbeg, rend int64, opt equalOption) bool { @@ -382,6 +446,7 @@ type equalOption struct { atol float64 // absolute tolerance nansEq bool // whether NaNs are considered equal. unorderedMapKeys bool // whether maps are allowed to have different entries order + nullable bool // whether the fields being compared are considered nullable } func (eq equalOption) f16(f1, f2 float16.Num) bool { @@ -417,8 +482,9 @@ func (eq equalOption) f64(v1, v2 float64) bool { func newEqualOption(opts ...EqualOption) equalOption { eq := equalOption{ - atol: defaultAbsoluteTolerance, - nansEq: false, + atol: defaultAbsoluteTolerance, + nansEq: false, + nullable: true, } for _, opt := range opts { opt(&eq) @@ -452,20 +518,27 @@ func WithUnorderedMapKeys(v bool) EqualOption { } } +// WithNullable sets whether the comparison function will consider both fields as nullable. If they're non-nullable, their +// valids buffer will be ignored for comparison and the underlying values will be used instead +func WithNullable(v bool) EqualOption { + return func(o *equalOption) { + o.nullable = v + } +} + // ApproxEqual reports whether the two provided arrays are approximately equal. // For non-floating point arrays, it is equivalent to Equal. func ApproxEqual(left, right arrow.Array, opts ...EqualOption) bool { - opt := newEqualOption(opts...) - return arrayApproxEqual(left, right, opt) + return arrayApproxEqual(left, right, newEqualOption(opts...)) } func arrayApproxEqual(left, right arrow.Array, opt equalOption) bool { switch { - case !baseArrayEqual(left, right): + case !baseArrayEqual(left, right, opt): return false case left.Len() == 0: return true - case left.NullN() == left.Len(): + case opt.nullable && hasTopLevelValidityBitmap(left.DataType().ID()) && left.NullN() == left.Len(): return true } @@ -478,52 +551,52 @@ func arrayApproxEqual(left, right arrow.Array, opt equalOption) bool { return true case *Boolean: r := right.(*Boolean) - return arrayEqualBoolean(l, r) + return arrayEqualBoolean(l, r, opt) case *FixedSizeBinary: r := right.(*FixedSizeBinary) - return arrayEqualFixedSizeBinary(l, r) + return arrayEqualFixedSizeBinary(l, r, opt) case *Binary: r := right.(*Binary) - return arrayEqualBinary(l, r) + return arrayEqualBinary(l, r, opt) case *String: r := right.(*String) - return arrayApproxEqualString(l, r) + return arrayApproxEqualString(l, r, opt) case *LargeBinary: r := right.(*LargeBinary) - return arrayEqualLargeBinary(l, r) + return arrayEqualLargeBinary(l, r, opt) case *LargeString: r := right.(*LargeString) - return arrayApproxEqualLargeString(l, r) + return arrayApproxEqualLargeString(l, r, opt) case *BinaryView: r := right.(*BinaryView) - return arrayEqualBinaryView(l, r) + return arrayEqualBinaryView(l, r, opt) case *StringView: r := right.(*StringView) - return arrayApproxEqualStringView(l, r) + return arrayApproxEqualStringView(l, r, opt) case *Int8: r := right.(*Int8) - return arrayEqualFixedWidth(l, r) + return arrayEqualFixedWidth(l, r, opt) case *Int16: r := right.(*Int16) - return arrayEqualFixedWidth(l, r) + return arrayEqualFixedWidth(l, r, opt) case *Int32: r := right.(*Int32) - return arrayEqualFixedWidth(l, r) + return arrayEqualFixedWidth(l, r, opt) case *Int64: r := right.(*Int64) - return arrayEqualFixedWidth(l, r) + return arrayEqualFixedWidth(l, r, opt) case *Uint8: r := right.(*Uint8) - return arrayEqualFixedWidth(l, r) + return arrayEqualFixedWidth(l, r, opt) case *Uint16: r := right.(*Uint16) - return arrayEqualFixedWidth(l, r) + return arrayEqualFixedWidth(l, r, opt) case *Uint32: r := right.(*Uint32) - return arrayEqualFixedWidth(l, r) + return arrayEqualFixedWidth(l, r, opt) case *Uint64: r := right.(*Uint64) - return arrayEqualFixedWidth(l, r) + return arrayEqualFixedWidth(l, r, opt) case *Float16: r := right.(*Float16) return arrayApproxEqualFloat16(l, r, opt) @@ -535,31 +608,31 @@ func arrayApproxEqual(left, right arrow.Array, opt equalOption) bool { return arrayApproxEqualFloat64(l, r, opt) case *Decimal32: r := right.(*Decimal32) - return arrayEqualDecimal(l, r) + return arrayEqualDecimal(l, r, opt) case *Decimal64: r := right.(*Decimal64) - return arrayEqualDecimal(l, r) + return arrayEqualDecimal(l, r, opt) case *Decimal128: r := right.(*Decimal128) - return arrayEqualDecimal(l, r) + return arrayEqualDecimal(l, r, opt) case *Decimal256: r := right.(*Decimal256) - return arrayEqualDecimal(l, r) + return arrayEqualDecimal(l, r, opt) case *Date32: r := right.(*Date32) - return arrayEqualFixedWidth(l, r) + return arrayEqualFixedWidth(l, r, opt) case *Date64: r := right.(*Date64) - return arrayEqualFixedWidth(l, r) + return arrayEqualFixedWidth(l, r, opt) case *Time32: r := right.(*Time32) - return arrayEqualFixedWidth(l, r) + return arrayEqualFixedWidth(l, r, opt) case *Time64: r := right.(*Time64) - return arrayEqualFixedWidth(l, r) + return arrayEqualFixedWidth(l, r, opt) case *Timestamp: r := right.(*Timestamp) - return arrayEqualTimestamp(l, r) + return arrayEqualTimestamp(l, r, opt) case *List: r := right.(*List) return arrayApproxEqualList(l, r, opt) @@ -580,16 +653,16 @@ func arrayApproxEqual(left, right arrow.Array, opt equalOption) bool { return arrayApproxEqualStruct(l, r, opt) case *MonthInterval: r := right.(*MonthInterval) - return arrayEqualMonthInterval(l, r) + return arrayEqualMonthInterval(l, r, opt) case *DayTimeInterval: r := right.(*DayTimeInterval) - return arrayEqualDayTimeInterval(l, r) + return arrayEqualDayTimeInterval(l, r, opt) case *MonthDayNanoInterval: r := right.(*MonthDayNanoInterval) - return arrayEqualMonthDayNanoInterval(l, r) + return arrayEqualMonthDayNanoInterval(l, r, opt) case *Duration: r := right.(*Duration) - return arrayEqualFixedWidth(l, r) + return arrayEqualFixedWidth(l, r, opt) case *Map: r := right.(*Map) if opt.unorderedMapKeys { @@ -616,15 +689,35 @@ func arrayApproxEqual(left, right arrow.Array, opt equalOption) bool { } } -func baseArrayEqual(left, right arrow.Array) bool { +func withNullable(opt equalOption, nullable bool) equalOption { + opt.nullable = nullable + return opt +} + +// hasTopLevelValidityBitmap reports whether arrays of the given type id carry +// their own top-level validity bitmap. Union and run-end-encoded arrays do not: +// their nullness is encoded entirely in their children, so top-level null-count +// and validity comparisons are not meaningful for them and can legitimately +// differ between logically-equal arrays (e.g. built from JSON vs. read from IPC). +func hasTopLevelValidityBitmap(id arrow.Type) bool { + switch id { + case arrow.SPARSE_UNION, arrow.DENSE_UNION, arrow.RUN_END_ENCODED: + return false + } + return true +} + +func baseArrayEqual(left, right arrow.Array, opt equalOption) bool { switch { case left.Len() != right.Len(): return false - case left.NullN() != right.NullN(): - return false case !arrow.TypeEqual(left.DataType(), right.DataType()): // We do not check for metadata as in the C++ implementation. return false - case !validityBitmapEqual(left, right): + case !hasTopLevelValidityBitmap(left.DataType().ID()): + return true + case opt.nullable && left.NullN() != right.NullN(): + return false + case opt.nullable && !validityBitmapEqual(left, right): return false } return true @@ -644,9 +737,9 @@ func validityBitmapEqual(left, right arrow.Array) bool { return true } -func arrayApproxEqualString(left, right *String) bool { +func arrayApproxEqualString(left, right *String, opt equalOption) bool { for i := 0; i < left.Len(); i++ { - if left.IsNull(i) { + if opt.nullable && left.IsNull(i) { continue } if stripNulls(left.Value(i)) != stripNulls(right.Value(i)) { @@ -656,9 +749,9 @@ func arrayApproxEqualString(left, right *String) bool { return true } -func arrayApproxEqualLargeString(left, right *LargeString) bool { +func arrayApproxEqualLargeString(left, right *LargeString, opt equalOption) bool { for i := 0; i < left.Len(); i++ { - if left.IsNull(i) { + if opt.nullable && left.IsNull(i) { continue } if stripNulls(left.Value(i)) != stripNulls(right.Value(i)) { @@ -668,9 +761,9 @@ func arrayApproxEqualLargeString(left, right *LargeString) bool { return true } -func arrayApproxEqualStringView(left, right *StringView) bool { +func arrayApproxEqualStringView(left, right *StringView, opt equalOption) bool { for i := 0; i < left.Len(); i++ { - if left.IsNull(i) { + if opt.nullable && left.IsNull(i) { continue } if stripNulls(left.Value(i)) != stripNulls(right.Value(i)) { @@ -682,7 +775,7 @@ func arrayApproxEqualStringView(left, right *StringView) bool { func arrayApproxEqualFloat16(left, right *Float16, opt equalOption) bool { for i := 0; i < left.Len(); i++ { - if left.IsNull(i) { + if opt.nullable && left.IsNull(i) { continue } if !opt.f16(left.Value(i), right.Value(i)) { @@ -694,7 +787,7 @@ func arrayApproxEqualFloat16(left, right *Float16, opt equalOption) bool { func arrayApproxEqualFloat32(left, right *Float32, opt equalOption) bool { for i := 0; i < left.Len(); i++ { - if left.IsNull(i) { + if opt.nullable && left.IsNull(i) { continue } if !opt.f32(left.Value(i), right.Value(i)) { @@ -706,7 +799,7 @@ func arrayApproxEqualFloat32(left, right *Float32, opt equalOption) bool { func arrayApproxEqualFloat64(left, right *Float64, opt equalOption) bool { for i := 0; i < left.Len(); i++ { - if left.IsNull(i) { + if opt.nullable && left.IsNull(i) { continue } if !opt.f64(left.Value(i), right.Value(i)) { @@ -717,8 +810,9 @@ func arrayApproxEqualFloat64(left, right *Float64, opt equalOption) bool { } func arrayApproxEqualList(left, right *List, opt equalOption) bool { + childOpt := withNullable(opt, left.DataType().(arrow.ListLikeType).ElemField().Nullable) for i := 0; i < left.Len(); i++ { - if left.IsNull(i) { + if opt.nullable && left.IsNull(i) { continue } o := func() bool { @@ -726,7 +820,7 @@ func arrayApproxEqualList(left, right *List, opt equalOption) bool { defer l.Release() r := right.newListValue(i) defer r.Release() - return arrayApproxEqual(l, r, opt) + return arrayApproxEqual(l, r, childOpt) }() if !o { return false @@ -736,8 +830,9 @@ func arrayApproxEqualList(left, right *List, opt equalOption) bool { } func arrayApproxEqualLargeList(left, right *LargeList, opt equalOption) bool { + childOpt := withNullable(opt, left.DataType().(arrow.ListLikeType).ElemField().Nullable) for i := 0; i < left.Len(); i++ { - if left.IsNull(i) { + if opt.nullable && left.IsNull(i) { continue } o := func() bool { @@ -745,7 +840,7 @@ func arrayApproxEqualLargeList(left, right *LargeList, opt equalOption) bool { defer l.Release() r := right.newListValue(i) defer r.Release() - return arrayApproxEqual(l, r, opt) + return arrayApproxEqual(l, r, childOpt) }() if !o { return false @@ -755,8 +850,9 @@ func arrayApproxEqualLargeList(left, right *LargeList, opt equalOption) bool { } func arrayApproxEqualListView(left, right *ListView, opt equalOption) bool { + childOpt := withNullable(opt, left.DataType().(arrow.ListLikeType).ElemField().Nullable) for i := 0; i < left.Len(); i++ { - if left.IsNull(i) { + if opt.nullable && left.IsNull(i) { continue } o := func() bool { @@ -764,7 +860,7 @@ func arrayApproxEqualListView(left, right *ListView, opt equalOption) bool { defer l.Release() r := right.newListValue(i) defer r.Release() - return arrayApproxEqual(l, r, opt) + return arrayApproxEqual(l, r, childOpt) }() if !o { return false @@ -774,8 +870,9 @@ func arrayApproxEqualListView(left, right *ListView, opt equalOption) bool { } func arrayApproxEqualLargeListView(left, right *LargeListView, opt equalOption) bool { + childOpt := withNullable(opt, left.DataType().(arrow.ListLikeType).ElemField().Nullable) for i := 0; i < left.Len(); i++ { - if left.IsNull(i) { + if opt.nullable && left.IsNull(i) { continue } o := func() bool { @@ -783,7 +880,7 @@ func arrayApproxEqualLargeListView(left, right *LargeListView, opt equalOption) defer l.Release() r := right.newListValue(i) defer r.Release() - return arrayApproxEqual(l, r, opt) + return arrayApproxEqual(l, r, childOpt) }() if !o { return false @@ -793,8 +890,9 @@ func arrayApproxEqualLargeListView(left, right *LargeListView, opt equalOption) } func arrayApproxEqualFixedSizeList(left, right *FixedSizeList, opt equalOption) bool { + childOpt := withNullable(opt, left.DataType().(arrow.ListLikeType).ElemField().Nullable) for i := 0; i < left.Len(); i++ { - if left.IsNull(i) { + if opt.nullable && left.IsNull(i) { continue } o := func() bool { @@ -802,7 +900,7 @@ func arrayApproxEqualFixedSizeList(left, right *FixedSizeList, opt equalOption) defer l.Release() r := right.newListValue(i) defer r.Release() - return arrayApproxEqual(l, r, opt) + return arrayApproxEqual(l, r, childOpt) }() if !o { return false @@ -812,17 +910,23 @@ func arrayApproxEqualFixedSizeList(left, right *FixedSizeList, opt equalOption) } func arrayApproxEqualStruct(left, right *Struct, opt equalOption) bool { - return bitutils.VisitSetBitRuns( - left.NullBitmapBytes(), - int64(left.Offset()), int64(left.Len()), - approxEqualStructRun(left, right, opt), - ) == nil + visitFn := approxEqualStructRun(left, right, opt) + if opt.nullable { + return bitutils.VisitSetBitRuns( + left.NullBitmapBytes(), + int64(left.Offset()), int64(left.Len()), + visitFn, + ) == nil + } + return visitFn(0, int64(left.Len())) == nil } func approxEqualStructRun(left, right *Struct, opt equalOption) bitutils.VisitFn { + st := left.DataType().(*arrow.StructType) return func(pos int64, length int64) error { for i := range left.fields { - if !sliceApproxEqual(left.fields[i], pos, pos+length, right.fields[i], pos, pos+length, opt) { + childOpt := withNullable(opt, st.Field(i).Nullable) + if !sliceApproxEqual(left.fields[i], pos, pos+length, right.fields[i], pos, pos+length, childOpt) { return arrow.ErrInvalid } } @@ -833,7 +937,7 @@ func approxEqualStructRun(left, right *Struct, opt equalOption) bitutils.VisitFn // arrayApproxEqualMap doesn't care about the order of keys (in Go map traversal order is undefined) func arrayApproxEqualMap(left, right *Map, opt equalOption) bool { for i := 0; i < left.Len(); i++ { - if left.IsNull(i) { + if opt.nullable && left.IsNull(i) { continue } if !arrayApproxEqualSingleMapEntry(left.newListValue(i).(*Struct), right.newListValue(i).(*Struct), opt) { @@ -854,7 +958,7 @@ func arrayApproxEqualSingleMapEntry(left, right *Struct, opt equalOption) bool { switch { case left.Len() != right.Len(): return false - case left.NullN() != right.NullN(): + case opt.nullable && left.NullN() != right.NullN(): return false case !arrow.TypeEqual(left.DataType(), right.DataType()): // We do not check for metadata as in the C++ implementation. return false @@ -862,9 +966,13 @@ func arrayApproxEqualSingleMapEntry(left, right *Struct, opt equalOption) bool { return true } + st := left.DataType().(*arrow.StructType) + keyOpt := withNullable(opt, st.Field(0).Nullable) + valOpt := withNullable(opt, st.Field(1).Nullable) + used := make(map[int]bool, right.Len()) for i := 0; i < left.Len(); i++ { - if left.IsNull(i) { + if opt.nullable && left.IsNull(i) { continue } @@ -874,7 +982,7 @@ func arrayApproxEqualSingleMapEntry(left, right *Struct, opt equalOption) bool { if used[j] { continue } - if right.IsNull(j) { + if opt.nullable && right.IsNull(j) { used[j] = true continue } @@ -882,12 +990,12 @@ func arrayApproxEqualSingleMapEntry(left, right *Struct, opt equalOption) bool { rBeg, rEnd := int64(j), int64(j+1) // check keys (field 0) - if !sliceApproxEqual(left.Field(0), lBeg, lEnd, right.Field(0), rBeg, rEnd, opt) { + if !sliceApproxEqual(left.Field(0), lBeg, lEnd, right.Field(0), rBeg, rEnd, keyOpt) { continue } // only now check the values - if sliceApproxEqual(left.Field(1), lBeg, lEnd, right.Field(1), rBeg, rEnd, opt) { + if sliceApproxEqual(left.Field(1), lBeg, lEnd, right.Field(1), rBeg, rEnd, valOpt) { found = true used[j] = true break diff --git a/arrow/array/compare_test.go b/arrow/array/compare_test.go index 0671def62..9de3445f6 100644 --- a/arrow/array/compare_test.go +++ b/arrow/array/compare_test.go @@ -391,6 +391,49 @@ func TestArrayApproxEqualFloats(t *testing.T) { } } +func TestArrayEqualNonNullable(t *testing.T) { + for name, recs := range arrdata.Records { + t.Run(name, func(t *testing.T) { + rec := recs[0] + + // Clone the schema and make everything non-nullable + fields := rec.Schema().Fields() + meta := rec.Schema().Metadata() + for i := range fields { + fields[i].Nullable = false + } + schema := arrow.NewSchema(fields, &meta) + + for i, rawCol := range rec.Columns() { + // make a clone of the column with NullN=0 + col := array.MakeFromData(array.NewData( + rawCol.DataType(), + rawCol.Len(), + rawCol.Data().Buffers(), + rawCol.Data().Children(), + 0, + 0, + )) + t.Run(schema.Field(i).Name, func(t *testing.T) { + arr := col + if !array.Equal(arr, arr, array.WithNullable(false)) { + t.Fatalf("identical arrays should compare equal:\narray=%v", arr) + } + sub1 := array.NewSlice(arr, 1, int64(arr.Len())) + defer sub1.Release() + + sub2 := array.NewSlice(arr, 0, int64(arr.Len()-1)) + defer sub2.Release() + + if array.Equal(sub1, sub2) && name != "nulls" { + t.Fatalf("non-identical arrays should not compare equal:\nsub1=%v\nsub2=%v\narrf=%v\n", sub1, sub2, arr) + } + }) + } + }) + } +} + func testStringMap(mem memory.Allocator, m map[string]string, keys []string) *array.Map { dt := arrow.MapOf(arrow.BinaryTypes.String, arrow.BinaryTypes.String) builder := array.NewMapBuilderWithType(mem, dt) diff --git a/arrow/array/decimal.go b/arrow/array/decimal.go index 704b1d932..269ca2904 100644 --- a/arrow/array/decimal.go +++ b/arrow/array/decimal.go @@ -89,8 +89,8 @@ func (a *baseDecimal[T]) setData(data *Data) { } } -func (a *baseDecimal[T]) GetOneForMarshal(i int) any { - if a.IsNull(i) { +func (a *baseDecimal[T]) GetOneForMarshalNullable(i int, nullable bool) any { + if nullable && a.IsNull(i) { return nil } @@ -99,6 +99,10 @@ func (a *baseDecimal[T]) GetOneForMarshal(i int) any { return n.ToBigFloat(scale).Text('g', int(typ.GetPrecision())) } +func (a *baseDecimal[T]) GetOneForMarshal(i int) any { + return a.GetOneForMarshalNullable(i, true) +} + func (a *baseDecimal[T]) MarshalJSON() ([]byte, error) { vals := make([]any, a.Len()) for i := 0; i < a.Len(); i++ { @@ -110,9 +114,9 @@ func (a *baseDecimal[T]) MarshalJSON() ([]byte, error) { func arrayEqualDecimal[T interface { decimal.DecimalTypes decimal.Num[T] -}](left, right *baseDecimal[T]) bool { +}](left, right *baseDecimal[T], opt equalOption) bool { for i := 0; i < left.Len(); i++ { - if left.IsNull(i) { + if opt.nullable && left.IsNull(i) { continue } diff --git a/arrow/array/dictionary.go b/arrow/array/dictionary.go index 38a43f77c..982aa2e3e 100644 --- a/arrow/array/dictionary.go +++ b/arrow/array/dictionary.go @@ -286,13 +286,16 @@ func (d *Dictionary) GetValueIndex(i int) int { debug.Assert(false, "unreachable dictionary index") return -1 } - -func (d *Dictionary) GetOneForMarshal(i int) interface{} { - if d.IsNull(i) { +func (d *Dictionary) GetOneForMarshalNullable(i int, nullable bool) interface{} { + if nullable && d.IsNull(i) { return nil } vidx := d.GetValueIndex(i) - return d.Dictionary().GetOneForMarshal(vidx) + return getOneForMarshalNullable(d.Dictionary(), vidx, nullable) +} + +func (d *Dictionary) GetOneForMarshal(i int) interface{} { + return d.GetOneForMarshalNullable(i, true) } func (d *Dictionary) MarshalJSON() ([]byte, error) { @@ -303,8 +306,8 @@ func (d *Dictionary) MarshalJSON() ([]byte, error) { return json.Marshal(vals) } -func arrayEqualDict(l, r *Dictionary) bool { - return Equal(l.Dictionary(), r.Dictionary()) && Equal(l.indices, r.indices) +func arrayEqualDict(l, r *Dictionary, opt equalOption) bool { + return equal(l.Dictionary(), r.Dictionary(), opt) && equal(l.indices, r.indices, opt) } func arrayApproxEqualDict(l, r *Dictionary, opt equalOption) bool { diff --git a/arrow/array/encoded.go b/arrow/array/encoded.go index 8d628ffc2..2724c1c0f 100644 --- a/arrow/array/encoded.go +++ b/arrow/array/encoded.go @@ -239,9 +239,16 @@ func (r *RunEndEncoded) String() string { buf.WriteByte(']') return buf.String() } +func (r *RunEndEncoded) GetOneForMarshalNullable(i int, nullable bool) interface{} { + // The values child may serialize as JSON null only when both the outer REE + // field and the REE values child (ValueNullable) are nullable; a non-nullable + // outer field must never emit null even if the values child is nullable. + nullable = nullable && r.data.dtype.(*arrow.RunEndEncodedType).ValueNullable + return getOneForMarshalNullable(r.values, r.GetPhysicalIndex(i), nullable) +} func (r *RunEndEncoded) GetOneForMarshal(i int) interface{} { - return r.values.GetOneForMarshal(r.GetPhysicalIndex(i)) + return r.GetOneForMarshalNullable(i, true) } func (r *RunEndEncoded) MarshalJSON() ([]byte, error) { @@ -260,14 +267,17 @@ func (r *RunEndEncoded) MarshalJSON() ([]byte, error) { return buf.Bytes(), nil } -func arrayRunEndEncodedEqual(l, r *RunEndEncoded) bool { +func arrayRunEndEncodedEqual(l, r *RunEndEncoded, opt equalOption) bool { // types were already checked before getting here, so we know // the encoded types are equal + childOpt := withNullable(opt, opt.nullable && + l.DataType().(*arrow.RunEndEncodedType).ValueNullable && + r.DataType().(*arrow.RunEndEncodedType).ValueNullable) mr := encoded.NewMergedRuns([2]arrow.Array{l, r}) for mr.Next() { lIndex := mr.IndexIntoArray(0) rIndex := mr.IndexIntoArray(1) - if !SliceEqual(l.values, lIndex, lIndex+1, r.values, rIndex, rIndex+1) { + if !sliceEqual(l.values, lIndex, lIndex+1, r.values, rIndex, rIndex+1, childOpt) { return false } } @@ -277,11 +287,14 @@ func arrayRunEndEncodedEqual(l, r *RunEndEncoded) bool { func arrayRunEndEncodedApproxEqual(l, r *RunEndEncoded, opt equalOption) bool { // types were already checked before getting here, so we know // the encoded types are equal + childOpt := withNullable(opt, opt.nullable && + l.DataType().(*arrow.RunEndEncodedType).ValueNullable && + r.DataType().(*arrow.RunEndEncodedType).ValueNullable) mr := encoded.NewMergedRuns([2]arrow.Array{l, r}) for mr.Next() { lIndex := mr.IndexIntoArray(0) rIndex := mr.IndexIntoArray(1) - if !sliceApproxEqual(l.values, lIndex, lIndex+1, r.values, rIndex, rIndex+1, opt) { + if !sliceApproxEqual(l.values, lIndex, lIndex+1, r.values, rIndex, rIndex+1, childOpt) { return false } } diff --git a/arrow/array/extension.go b/arrow/array/extension.go index e509b5e0f..a74f6dee2 100644 --- a/arrow/array/extension.go +++ b/arrow/array/extension.go @@ -46,12 +46,12 @@ type ExtensionArray interface { // two extension arrays are equal if their data types are equal and // their underlying storage arrays are equal. -func arrayEqualExtension(l, r ExtensionArray) bool { +func arrayEqualExtension(l, r ExtensionArray, opt equalOption) bool { if !arrow.TypeEqual(l.DataType(), r.DataType()) { return false } - return Equal(l.Storage(), r.Storage()) + return equal(l.Storage(), r.Storage(), opt) } // two extension arrays are approximately equal if their data types are @@ -116,6 +116,11 @@ func (e *ExtensionArrayBase) String() string { return fmt.Sprintf("(%s)%s", e.data.dtype, e.storage) } +// GetOneForMarshal returns the value at i from the underlying storage array. +// ExtensionArrayBase deliberately does not implement arrow.NullableMarshaler: +// doing so would promote a nullable-aware method onto every embedding extension +// array and bypass a concrete type's own GetOneForMarshal override. Field-local +// nullability for plain extension arrays is handled by the marshaling helper. func (e *ExtensionArrayBase) GetOneForMarshal(i int) interface{} { return e.storage.GetOneForMarshal(i) } diff --git a/arrow/array/fixed_size_list.go b/arrow/array/fixed_size_list.go index d382ebe93..65d0654cb 100644 --- a/arrow/array/fixed_size_list.go +++ b/arrow/array/fixed_size_list.go @@ -84,9 +84,10 @@ func (a *FixedSizeList) setData(data *Data) { a.values = MakeFromData(data.childData[0]) } -func arrayEqualFixedSizeList(left, right *FixedSizeList) bool { +func arrayEqualFixedSizeList(left, right *FixedSizeList, opt equalOption) bool { + childOpt := withNullable(opt, left.DataType().(arrow.ListLikeType).ElemField().Nullable) for i := 0; i < left.Len(); i++ { - if left.IsNull(i) { + if opt.nullable && left.IsNull(i) { continue } o := func() bool { @@ -94,7 +95,7 @@ func arrayEqualFixedSizeList(left, right *FixedSizeList) bool { defer l.Release() r := right.newListValue(i) defer r.Release() - return Equal(l, r) + return equal(l, r, childOpt) }() if !o { return false @@ -123,18 +124,17 @@ func (a *FixedSizeList) Release() { a.values.Release() } -func (a *FixedSizeList) GetOneForMarshal(i int) interface{} { - if a.IsNull(i) { +func (a *FixedSizeList) GetOneForMarshalNullable(i int, nullable bool) interface{} { + if nullable && a.IsNull(i) { return nil } slice := a.newListValue(i) defer slice.Release() - v, err := json.Marshal(slice) - if err != nil { - panic(err) - } + return marshalListElemsNullable(slice, a.DataType().(arrow.ListLikeType).ElemField().Nullable) +} - return json.RawMessage(v) +func (a *FixedSizeList) GetOneForMarshal(i int) interface{} { + return a.GetOneForMarshalNullable(i, true) } func (a *FixedSizeList) MarshalJSON() ([]byte, error) { diff --git a/arrow/array/fixedsize_binary.go b/arrow/array/fixedsize_binary.go index 31d507c5b..11dc76154 100644 --- a/arrow/array/fixedsize_binary.go +++ b/arrow/array/fixedsize_binary.go @@ -86,14 +86,18 @@ func (a *FixedSizeBinary) setData(data *Data) { } } -func (a *FixedSizeBinary) GetOneForMarshal(i int) interface{} { - if a.IsNull(i) { +func (a *FixedSizeBinary) GetOneForMarshalNullable(i int, nullable bool) interface{} { + if nullable && a.IsNull(i) { return nil } return a.Value(i) } +func (a *FixedSizeBinary) GetOneForMarshal(i int) interface{} { + return a.GetOneForMarshalNullable(i, true) +} + func (a *FixedSizeBinary) MarshalJSON() ([]byte, error) { vals := make([]interface{}, a.Len()) for i := 0; i < a.Len(); i++ { @@ -106,9 +110,9 @@ func (a *FixedSizeBinary) MarshalJSON() ([]byte, error) { return json.Marshal(vals) } -func arrayEqualFixedSizeBinary(left, right *FixedSizeBinary) bool { +func arrayEqualFixedSizeBinary(left, right *FixedSizeBinary, opt equalOption) bool { for i := 0; i < left.Len(); i++ { - if left.IsNull(i) { + if opt.nullable && left.IsNull(i) { continue } if !bytes.Equal(left.Value(i), right.Value(i)) { diff --git a/arrow/array/float16.go b/arrow/array/float16.go index 41276803b..8923eb438 100644 --- a/arrow/array/float16.go +++ b/arrow/array/float16.go @@ -77,13 +77,17 @@ func (a *Float16) setData(data *Data) { } } -func (a *Float16) GetOneForMarshal(i int) interface{} { - if a.IsValid(i) { +func (a *Float16) GetOneForMarshalNullable(i int, nullable bool) interface{} { + if !nullable || a.IsValid(i) { return a.values[i].Float32() } return nil } +func (a *Float16) GetOneForMarshal(i int) interface{} { + return a.GetOneForMarshalNullable(i, true) +} + func (a *Float16) MarshalJSON() ([]byte, error) { vals := make([]interface{}, a.Len()) for i, v := range a.values { diff --git a/arrow/array/interval.go b/arrow/array/interval.go index 2c029c252..076a7a9d4 100644 --- a/arrow/array/interval.go +++ b/arrow/array/interval.go @@ -94,13 +94,17 @@ func (a *MonthInterval) setData(data *Data) { } } -func (a *MonthInterval) GetOneForMarshal(i int) interface{} { - if a.IsValid(i) { +func (a *MonthInterval) GetOneForMarshalNullable(i int, nullable bool) interface{} { + if !nullable || a.IsValid(i) { return a.values[i] } return nil } +func (a *MonthInterval) GetOneForMarshal(i int) interface{} { + return a.GetOneForMarshalNullable(i, true) +} + // MarshalJSON will create a json array out of a MonthInterval array, // each value will be an object of the form {"months": #} where // # is the numeric value of that index @@ -120,9 +124,9 @@ func (a *MonthInterval) MarshalJSON() ([]byte, error) { return json.Marshal(vals) } -func arrayEqualMonthInterval(left, right *MonthInterval) bool { +func arrayEqualMonthInterval(left, right *MonthInterval, opt equalOption) bool { for i := 0; i < left.Len(); i++ { - if left.IsNull(i) { + if opt.nullable && left.IsNull(i) { continue } if left.Value(i) != right.Value(i) { @@ -399,13 +403,17 @@ func (a *DayTimeInterval) setData(data *Data) { } } -func (a *DayTimeInterval) GetOneForMarshal(i int) interface{} { - if a.IsValid(i) { +func (a *DayTimeInterval) GetOneForMarshalNullable(i int, nullable bool) interface{} { + if !nullable || a.IsValid(i) { return a.values[i] } return nil } +func (a *DayTimeInterval) GetOneForMarshal(i int) interface{} { + return a.GetOneForMarshalNullable(i, true) +} + // MarshalJSON will marshal this array to JSON as an array of objects, // consisting of the form {"days": #, "milliseconds": #} for each element. func (a *DayTimeInterval) MarshalJSON() ([]byte, error) { @@ -423,9 +431,9 @@ func (a *DayTimeInterval) MarshalJSON() ([]byte, error) { return json.Marshal(vals) } -func arrayEqualDayTimeInterval(left, right *DayTimeInterval) bool { +func arrayEqualDayTimeInterval(left, right *DayTimeInterval, opt equalOption) bool { for i := 0; i < left.Len(); i++ { - if left.IsNull(i) { + if opt.nullable && left.IsNull(i) { continue } if left.Value(i) != right.Value(i) { @@ -703,13 +711,17 @@ func (a *MonthDayNanoInterval) setData(data *Data) { } } -func (a *MonthDayNanoInterval) GetOneForMarshal(i int) interface{} { - if a.IsValid(i) { +func (a *MonthDayNanoInterval) GetOneForMarshalNullable(i int, nullable bool) interface{} { + if !nullable || a.IsValid(i) { return a.values[i] } return nil } +func (a *MonthDayNanoInterval) GetOneForMarshal(i int) interface{} { + return a.GetOneForMarshalNullable(i, true) +} + // MarshalJSON will marshal this array to a JSON array with elements // marshalled to the form {"months": #, "days": #, "nanoseconds": #} func (a *MonthDayNanoInterval) MarshalJSON() ([]byte, error) { @@ -727,9 +739,9 @@ func (a *MonthDayNanoInterval) MarshalJSON() ([]byte, error) { return json.Marshal(vals) } -func arrayEqualMonthDayNanoInterval(left, right *MonthDayNanoInterval) bool { +func arrayEqualMonthDayNanoInterval(left, right *MonthDayNanoInterval, opt equalOption) bool { for i := 0; i < left.Len(); i++ { - if left.IsNull(i) { + if opt.nullable && left.IsNull(i) { continue } if left.Value(i) != right.Value(i) { diff --git a/arrow/array/list.go b/arrow/array/list.go index d72887bb8..5ce57a341 100644 --- a/arrow/array/list.go +++ b/arrow/array/list.go @@ -98,18 +98,18 @@ func (a *List) setData(data *Data) { a.values = MakeFromData(data.childData[0]) } -func (a *List) GetOneForMarshal(i int) interface{} { - if a.IsNull(i) { +func (a *List) GetOneForMarshalNullable(i int, nullable bool) interface{} { + if nullable && a.IsNull(i) { return nil } slice := a.newListValue(i) defer slice.Release() - v, err := json.Marshal(slice) - if err != nil { - panic(err) - } - return json.RawMessage(v) + return marshalListElemsNullable(slice, a.DataType().(arrow.ListLikeType).ElemField().Nullable) +} + +func (a *List) GetOneForMarshal(i int) interface{} { + return a.GetOneForMarshalNullable(i, true) } func (a *List) MarshalJSON() ([]byte, error) { @@ -129,9 +129,10 @@ func (a *List) MarshalJSON() ([]byte, error) { return buf.Bytes(), nil } -func arrayEqualList(left, right *List) bool { +func arrayEqualList(left, right *List, opt equalOption) bool { + childOpt := withNullable(opt, left.DataType().(arrow.ListLikeType).ElemField().Nullable) for i := 0; i < left.Len(); i++ { - if left.IsNull(i) { + if opt.nullable && left.IsNull(i) { continue } o := func() bool { @@ -139,7 +140,7 @@ func arrayEqualList(left, right *List) bool { defer l.Release() r := right.newListValue(i) defer r.Release() - return Equal(l, r) + return equal(l, r, childOpt) }() if !o { return false @@ -230,18 +231,18 @@ func (a *LargeList) setData(data *Data) { a.values = MakeFromData(data.childData[0]) } -func (a *LargeList) GetOneForMarshal(i int) interface{} { - if a.IsNull(i) { +func (a *LargeList) GetOneForMarshalNullable(i int, nullable bool) interface{} { + if nullable && a.IsNull(i) { return nil } slice := a.newListValue(i) defer slice.Release() - v, err := json.Marshal(slice) - if err != nil { - panic(err) - } - return json.RawMessage(v) + return marshalListElemsNullable(slice, a.DataType().(arrow.ListLikeType).ElemField().Nullable) +} + +func (a *LargeList) GetOneForMarshal(i int) interface{} { + return a.GetOneForMarshalNullable(i, true) } func (a *LargeList) MarshalJSON() ([]byte, error) { @@ -261,9 +262,10 @@ func (a *LargeList) MarshalJSON() ([]byte, error) { return buf.Bytes(), nil } -func arrayEqualLargeList(left, right *LargeList) bool { +func arrayEqualLargeList(left, right *LargeList, opt equalOption) bool { + childOpt := withNullable(opt, left.DataType().(arrow.ListLikeType).ElemField().Nullable) for i := 0; i < left.Len(); i++ { - if left.IsNull(i) { + if opt.nullable && left.IsNull(i) { continue } o := func() bool { @@ -271,7 +273,7 @@ func arrayEqualLargeList(left, right *LargeList) bool { defer l.Release() r := right.newListValue(i) defer r.Release() - return Equal(l, r) + return equal(l, r, childOpt) }() if !o { return false @@ -705,18 +707,18 @@ func (a *ListView) setData(data *Data) { a.values = MakeFromData(data.childData[0]) } -func (a *ListView) GetOneForMarshal(i int) interface{} { - if a.IsNull(i) { +func (a *ListView) GetOneForMarshalNullable(i int, nullable bool) interface{} { + if nullable && a.IsNull(i) { return nil } slice := a.newListValue(i) defer slice.Release() - v, err := json.Marshal(slice) - if err != nil { - panic(err) - } - return json.RawMessage(v) + return marshalListElemsNullable(slice, a.DataType().(arrow.ListLikeType).ElemField().Nullable) +} + +func (a *ListView) GetOneForMarshal(i int) interface{} { + return a.GetOneForMarshalNullable(i, true) } func (a *ListView) MarshalJSON() ([]byte, error) { @@ -736,9 +738,10 @@ func (a *ListView) MarshalJSON() ([]byte, error) { return buf.Bytes(), nil } -func arrayEqualListView(left, right *ListView) bool { +func arrayEqualListView(left, right *ListView, opt equalOption) bool { + childOpt := withNullable(opt, left.DataType().(arrow.ListLikeType).ElemField().Nullable) for i := 0; i < left.Len(); i++ { - if left.IsNull(i) { + if opt.nullable && left.IsNull(i) { continue } o := func() bool { @@ -746,7 +749,7 @@ func arrayEqualListView(left, right *ListView) bool { defer l.Release() r := right.newListValue(i) defer r.Release() - return Equal(l, r) + return equal(l, r, childOpt) }() if !o { return false @@ -852,18 +855,18 @@ func (a *LargeListView) setData(data *Data) { a.values = MakeFromData(data.childData[0]) } -func (a *LargeListView) GetOneForMarshal(i int) interface{} { - if a.IsNull(i) { +func (a *LargeListView) GetOneForMarshalNullable(i int, nullable bool) interface{} { + if nullable && a.IsNull(i) { return nil } slice := a.newListValue(i) defer slice.Release() - v, err := json.Marshal(slice) - if err != nil { - panic(err) - } - return json.RawMessage(v) + return marshalListElemsNullable(slice, a.DataType().(arrow.ListLikeType).ElemField().Nullable) +} + +func (a *LargeListView) GetOneForMarshal(i int) interface{} { + return a.GetOneForMarshalNullable(i, true) } func (a *LargeListView) MarshalJSON() ([]byte, error) { @@ -883,9 +886,10 @@ func (a *LargeListView) MarshalJSON() ([]byte, error) { return buf.Bytes(), nil } -func arrayEqualLargeListView(left, right *LargeListView) bool { +func arrayEqualLargeListView(left, right *LargeListView, opt equalOption) bool { + childOpt := withNullable(opt, left.DataType().(arrow.ListLikeType).ElemField().Nullable) for i := 0; i < left.Len(); i++ { - if left.IsNull(i) { + if opt.nullable && left.IsNull(i) { continue } o := func() bool { @@ -893,7 +897,7 @@ func arrayEqualLargeListView(left, right *LargeListView) bool { defer l.Release() r := right.newListValue(i) defer r.Release() - return Equal(l, r) + return equal(l, r, childOpt) }() if !o { return false diff --git a/arrow/array/map.go b/arrow/array/map.go index 71d4d1382..af6c28183 100644 --- a/arrow/array/map.go +++ b/arrow/array/map.go @@ -106,9 +106,9 @@ func (a *Map) Release() { a.items.Release() } -func arrayEqualMap(left, right *Map) bool { +func arrayEqualMap(left, right *Map, opt equalOption) bool { // since Map is implemented using a list, we can just use arrayEqualList - return arrayEqualList(left.List, right.List) + return arrayEqualList(left.List, right.List, opt) } type MapBuilder struct { diff --git a/arrow/array/null.go b/arrow/array/null.go index 8f8f58056..bdb4ee22b 100644 --- a/arrow/array/null.go +++ b/arrow/array/null.go @@ -80,10 +80,14 @@ func (a *Null) setData(data *Data) { a.data.nulls = a.data.length } -func (a *Null) GetOneForMarshal(i int) interface{} { +func (a *Null) GetOneForMarshalNullable(i int, nullable bool) interface{} { return nil } +func (a *Null) GetOneForMarshal(i int) interface{} { + return a.GetOneForMarshalNullable(i, true) +} + func (a *Null) MarshalJSON() ([]byte, error) { return json.Marshal(make([]interface{}, a.Len())) } diff --git a/arrow/array/numeric_generic.go b/arrow/array/numeric_generic.go index 1b671fc76..5abade6d2 100644 --- a/arrow/array/numeric_generic.go +++ b/arrow/array/numeric_generic.go @@ -83,14 +83,18 @@ func (a *numericArray[T]) ValueStr(i int) string { return fmt.Sprintf("%v", a.values[i]) } -func (a *numericArray[T]) GetOneForMarshal(i int) any { - if a.IsNull(i) { +func (a *numericArray[T]) GetOneForMarshalNullable(i int, nullable bool) any { + if nullable && a.IsNull(i) { return nil } return a.values[i] } +func (a *numericArray[T]) GetOneForMarshal(i int) any { + return a.GetOneForMarshalNullable(i, true) +} + func (a *numericArray[T]) MarshalJSON() ([]byte, error) { vals := make([]any, a.Len()) for i := range a.Len() { @@ -107,14 +111,18 @@ type oneByteArrs[T int8 | uint8] struct { numericArray[T] } -func (a *oneByteArrs[T]) GetOneForMarshal(i int) any { - if a.IsNull(i) { +func (a *oneByteArrs[T]) GetOneForMarshalNullable(i int, nullable bool) any { + if nullable && a.IsNull(i) { return nil } return float64(a.values[i]) // prevent uint8/int8 from being seen as binary data } +func (a *oneByteArrs[T]) GetOneForMarshal(i int) any { + return a.GetOneForMarshalNullable(i, true) +} + func (a *oneByteArrs[T]) MarshalJSON() ([]byte, error) { vals := make([]any, a.Len()) for i := range a.Len() { @@ -141,8 +149,8 @@ func (a *floatArray[T]) ValueStr(i int) string { return strconv.FormatFloat(float64(a.Value(i)), 'g', -1, bitWidth) } -func (a *floatArray[T]) GetOneForMarshal(i int) any { - if a.IsNull(i) { +func (a *floatArray[T]) GetOneForMarshalNullable(i int, nullable bool) any { + if nullable && a.IsNull(i) { return nil } @@ -157,6 +165,10 @@ func (a *floatArray[T]) GetOneForMarshal(i int) any { } } +func (a *floatArray[T]) GetOneForMarshal(i int) any { + return a.GetOneForMarshalNullable(i, true) +} + func (a *floatArray[T]) MarshalJSON() ([]byte, error) { vals := make([]any, a.Len()) for i := range a.values { @@ -189,14 +201,18 @@ func (d *dateArray[T]) ValueStr(i int) string { return d.values[i].FormattedString() } -func (d *dateArray[T]) GetOneForMarshal(i int) interface{} { - if d.IsNull(i) { +func (d *dateArray[T]) GetOneForMarshalNullable(i int, nullable bool) interface{} { + if nullable && d.IsNull(i) { return nil } return d.values[i].FormattedString() } +func (d *dateArray[T]) GetOneForMarshal(i int) interface{} { + return d.GetOneForMarshalNullable(i, true) +} + type timeType interface { TimeUnit() arrow.TimeUnit } @@ -225,14 +241,18 @@ func (a *timeArray[T]) ValueStr(i int) string { return a.values[i].FormattedString(a.DataType().(timeType).TimeUnit()) } -func (a *timeArray[T]) GetOneForMarshal(i int) interface{} { - if a.IsNull(i) { +func (a *timeArray[T]) GetOneForMarshalNullable(i int, nullable bool) interface{} { + if nullable && a.IsNull(i) { return nil } return a.values[i].ToTime(a.DataType().(timeType).TimeUnit()).Format("15:04:05.999999999") } +func (a *timeArray[T]) GetOneForMarshal(i int) interface{} { + return a.GetOneForMarshalNullable(i, true) +} + type Duration struct { numericArray[arrow.Duration] } @@ -261,13 +281,17 @@ func (a *Duration) ValueStr(i int) string { return fmt.Sprintf("%d%s", a.values[i], a.DataType().(timeType).TimeUnit()) } -func (a *Duration) GetOneForMarshal(i int) any { - if a.IsNull(i) { +func (a *Duration) GetOneForMarshalNullable(i int, nullable bool) any { + if nullable && a.IsNull(i) { return nil } return fmt.Sprintf("%d%s", a.values[i], a.DataType().(timeType).TimeUnit()) } +func (a *Duration) GetOneForMarshal(i int) any { + return a.GetOneForMarshalNullable(i, true) +} + type Int64 struct { numericArray[int64] } @@ -436,9 +460,9 @@ func NewDate64Data(data arrow.ArrayData) *Date64 { func (a *Date64) Date64Values() []arrow.Date64 { return a.Values() } -func arrayEqualFixedWidth[T arrow.FixedWidthType](left, right arrow.TypedArray[T]) bool { +func arrayEqualFixedWidth[T arrow.FixedWidthType](left, right arrow.TypedArray[T], opt equalOption) bool { for i := range left.Len() { - if left.IsNull(i) { + if opt.nullable && left.IsNull(i) { continue } if left.Value(i) != right.Value(i) { diff --git a/arrow/array/record.go b/arrow/array/record.go index b7f84180a..4b7b17bd4 100644 --- a/arrow/array/record.go +++ b/arrow/array/record.go @@ -455,10 +455,22 @@ func (b *RecordBuilder) UnmarshalOne(dec *json.Decoder) error { } continue } + idx := indices[0] - if err := b.fields[indices[0]].UnmarshalOne(dec); err != nil { + var next json.RawMessage + if err := dec.Decode(&next); err != nil { return err } + + if json.IsNullMessage(next) && !b.schema.Field(idx).Nullable { + b.fields[idx].AppendEmptyValue() + } else { + sub := json.NewDecoder(bytes.NewReader(next)) + sub.UseNumber() + if err := b.fields[idx].UnmarshalOne(sub); err != nil { + return err + } + } } // consume the closing '}' diff --git a/arrow/array/record_test.go b/arrow/array/record_test.go index a3924382a..d07325c73 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,6 +26,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" + "github.com/apache/arrow-go/v18/internal/json" "github.com/stretchr/testify/assert" ) @@ -485,9 +487,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 +500,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 +516,12 @@ func TestRecordBuilder(t *testing.T) { } } + err := b.UnmarshalJSON([]byte(`{"f1-i32": 6, "f2-f64-notnull": 6.6, "map": [{"key": "4", "value": "d"}]}`)) + assert.NoError(t, err) + + err = b.UnmarshalJSON([]byte(`{"f1-i32": null, "f2-f64-notnull": null, "map": null}`)) + assert.NoError(t, err) + rec := b.NewRecordBatch() defer rec.Release() @@ -518,7 +529,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(7); 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 +538,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 (null)]`; 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 0]`; 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"]} (null)]`; 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) + assert.NoError(t, err) + + roundtripped, _, err := array.FromJSON(mem, arr.DataType(), bytes.NewReader(jsonStr)) + defer roundtripped.Release() + assert.NoError(t, err) + assert.Truef(t, array.Equal(arr, roundtripped), "JSON round trip returns different array: got=%q, want=%d", arr, roundtripped) } func TestRecordBuilderResize(t *testing.T) { diff --git a/arrow/array/string.go b/arrow/array/string.go index 7c2ab0744..bbeba8879 100644 --- a/arrow/array/string.go +++ b/arrow/array/string.go @@ -154,13 +154,17 @@ func (a *String) setData(data *Data) { } } -func (a *String) GetOneForMarshal(i int) interface{} { - if a.IsValid(i) { +func (a *String) GetOneForMarshalNullable(i int, nullable bool) interface{} { + if !nullable || a.IsValid(i) { return a.Value(i) } return nil } +func (a *String) GetOneForMarshal(i int) interface{} { + return a.GetOneForMarshalNullable(i, true) +} + func (a *String) MarshalJSON() ([]byte, error) { vals := make([]interface{}, a.Len()) for i := 0; i < a.Len(); i++ { @@ -231,9 +235,9 @@ func (a *String) ValidateFull() error { return nil } -func arrayEqualString(left, right *String) bool { +func arrayEqualString(left, right *String, opt equalOption) bool { for i := 0; i < left.Len(); i++ { - if left.IsNull(i) { + if opt.nullable && left.IsNull(i) { continue } if left.Value(i) != right.Value(i) { @@ -362,13 +366,17 @@ func (a *LargeString) setData(data *Data) { } } -func (a *LargeString) GetOneForMarshal(i int) interface{} { - if a.IsValid(i) { +func (a *LargeString) GetOneForMarshalNullable(i int, nullable bool) interface{} { + if !nullable || a.IsValid(i) { return a.Value(i) } return nil } +func (a *LargeString) GetOneForMarshal(i int) interface{} { + return a.GetOneForMarshalNullable(i, true) +} + func (a *LargeString) MarshalJSON() ([]byte, error) { vals := make([]interface{}, a.Len()) for i := 0; i < a.Len(); i++ { @@ -435,9 +443,9 @@ func (a *LargeString) ValidateFull() error { return nil } -func arrayEqualLargeString(left, right *LargeString) bool { +func arrayEqualLargeString(left, right *LargeString, opt equalOption) bool { for i := 0; i < left.Len(); i++ { - if left.IsNull(i) { + if opt.nullable && left.IsNull(i) { continue } if left.Value(i) != right.Value(i) { @@ -526,13 +534,17 @@ func (a *StringView) ValueStr(i int) string { return a.Value(i) } -func (a *StringView) GetOneForMarshal(i int) interface{} { - if a.IsNull(i) { +func (a *StringView) GetOneForMarshalNullable(i int, nullable bool) interface{} { + if nullable && a.IsNull(i) { return nil } return a.Value(i) } +func (a *StringView) GetOneForMarshal(i int) interface{} { + return a.GetOneForMarshalNullable(i, true) +} + func (a *StringView) MarshalJSON() ([]byte, error) { vals := make([]interface{}, a.Len()) for i := 0; i < a.Len(); i++ { @@ -541,10 +553,10 @@ func (a *StringView) MarshalJSON() ([]byte, error) { return json.Marshal(vals) } -func arrayEqualStringView(left, right *StringView) bool { +func arrayEqualStringView(left, right *StringView, opt equalOption) bool { leftBufs, rightBufs := left.dataBuffers, right.dataBuffers for i := 0; i < left.Len(); i++ { - if left.IsNull(i) { + if opt.nullable && left.IsNull(i) { continue } if !left.ValueHeader(i).Equals(leftBufs, right.ValueHeader(i), rightBufs) { diff --git a/arrow/array/struct.go b/arrow/array/struct.go index 07505de1c..5579158d0 100644 --- a/arrow/array/struct.go +++ b/arrow/array/struct.go @@ -209,19 +209,24 @@ func (a *Struct) setData(data *Data) { } } -func (a *Struct) GetOneForMarshal(i int) interface{} { - if a.IsNull(i) { +func (a *Struct) GetOneForMarshalNullable(i int, nullable bool) interface{} { + if nullable && a.IsNull(i) { return nil } 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) + tmp[fieldList[j].Name] = getOneForMarshalNullable(d, i, dtype.Field(j).Nullable) } return tmp } +func (a *Struct) GetOneForMarshal(i int) interface{} { + return a.GetOneForMarshalNullable(i, true) +} + func (a *Struct) MarshalJSON() ([]byte, error) { var buf bytes.Buffer enc := json.NewEncoder(&buf) @@ -239,10 +244,11 @@ func (a *Struct) MarshalJSON() ([]byte, error) { return buf.Bytes(), nil } -func arrayEqualStruct(left, right *Struct) bool { +func arrayEqualStruct(left, right *Struct, opt equalOption) bool { for i, lf := range left.fields { rf := right.fields[i] - if !Equal(lf, rf) { + opt.nullable = left.data.dtype.(*arrow.StructType).Field(i).Nullable + if !equal(lf, rf, opt) { return false } } @@ -487,9 +493,20 @@ func (b *StructBuilder) UnmarshalOne(dec *json.Decoder) error { continue } - if err := b.fields[idx].UnmarshalOne(dec); err != nil { + var next json.RawMessage + if err := dec.Decode(&next); err != nil { return err } + + if json.IsNullMessage(next) && !b.dtype.(*arrow.StructType).Field(idx).Nullable { + b.fields[idx].AppendEmptyValue() + } else { + sub := json.NewDecoder(bytes.NewReader(next)) + sub.UseNumber() + if err := b.fields[idx].UnmarshalOne(sub); err != nil { + return err + } + } } // Append null values to all optional fields that were not presented in the json input diff --git a/arrow/array/struct_test.go b/arrow/array/struct_test.go index 216b353ab..6027d0288 100644 --- a/arrow/array/struct_test.go +++ b/arrow/array/struct_test.go @@ -513,6 +513,12 @@ func TestStructArrayUnmarshalJSONMissingFields(t *testing.T) { panic: false, want: `{[(null)] [3] {[(null)] [(null)] ["test"]}}`, }, + { + name: "explicit null in required field", + jsonInput: `[{"f2": 3, "f3": {"f3_3": null}}]`, + panic: false, + want: `{[(null)] [3] {[(null)] [(null)] [""]}}`, + }, } for _, tc := range tests { diff --git a/arrow/array/timestamp.go b/arrow/array/timestamp.go index 5ac0ee6cc..e2897fac5 100644 --- a/arrow/array/timestamp.go +++ b/arrow/array/timestamp.go @@ -109,13 +109,17 @@ func (a *Timestamp) ValueStr(i int) string { return toTime(a.values[i]).Format(layout) } -func (a *Timestamp) GetOneForMarshal(i int) interface{} { - if val := a.ValueStr(i); val != NullValueStr { +func (a *Timestamp) GetOneForMarshalNullable(i int, nullable bool) interface{} { + if val := a.ValueStr(i); !nullable || val != NullValueStr { return val } return nil } +func (a *Timestamp) GetOneForMarshal(i int) interface{} { + return a.GetOneForMarshalNullable(i, true) +} + func (a *Timestamp) MarshalJSON() ([]byte, error) { vals := make([]interface{}, a.Len()) for i := range a.values { @@ -125,9 +129,9 @@ func (a *Timestamp) MarshalJSON() ([]byte, error) { return json.Marshal(vals) } -func arrayEqualTimestamp(left, right *Timestamp) bool { +func arrayEqualTimestamp(left, right *Timestamp, opt equalOption) bool { for i := 0; i < left.Len(); i++ { - if left.IsNull(i) { + if opt.nullable && left.IsNull(i) { continue } if left.Value(i) != right.Value(i) { diff --git a/arrow/array/union.go b/arrow/array/union.go index 58524a768..f8e00fe74 100644 --- a/arrow/array/union.go +++ b/arrow/array/union.go @@ -320,17 +320,22 @@ func (a *SparseUnion) setData(data *Data) { debug.Assert(a.data.buffers[0] == nil, "arrow/array: validity bitmap for sparse unions should be nil") } -func (a *SparseUnion) GetOneForMarshal(i int) interface{} { +func (a *SparseUnion) GetOneForMarshalNullable(i int, nullable bool) any { typeID := a.RawTypeCodes()[i] childID := a.ChildID(i) data := a.Field(childID) - if data.IsNull(i) { - return nil + childNullable := a.unionType.Fields()[childID].Nullable + if childNullable && data.IsNull(i) { + return []any{typeID, nil} } - return []interface{}{typeID, data.GetOneForMarshal(i)} + return []any{typeID, getOneForMarshalNullable(data, i, childNullable)} +} + +func (a *SparseUnion) GetOneForMarshal(i int) any { + return a.GetOneForMarshalNullable(i, true) } func (a *SparseUnion) MarshalJSON() ([]byte, error) { @@ -440,7 +445,7 @@ func (a *SparseUnion) GetFlattenedField(mem memory.Allocator, index int) (arrow. return MakeFromData(childData), nil } -func arraySparseUnionEqual(l, r *SparseUnion) bool { +func arraySparseUnionEqual(l, r *SparseUnion, opt equalOption) bool { childIDs := l.unionType.ChildIDs() leftCodes, rightCodes := l.RawTypeCodes(), r.RawTypeCodes() @@ -451,8 +456,9 @@ func arraySparseUnionEqual(l, r *SparseUnion) bool { } childNum := childIDs[typeID] - eq := SliceEqual(l.children[childNum], int64(i), int64(i+1), - r.children[childNum], int64(i), int64(i+1)) + childOpt := withNullable(opt, l.unionType.Fields()[childNum].Nullable) + eq := sliceEqual(l.children[childNum], int64(i), int64(i+1), + r.children[childNum], int64(i), int64(i+1), childOpt) if !eq { return false } @@ -471,8 +477,9 @@ func arraySparseUnionApproxEqual(l, r *SparseUnion, opt equalOption) bool { } childNum := childIDs[typeID] + childOpt := withNullable(opt, l.unionType.Fields()[childNum].Nullable) eq := sliceApproxEqual(l.children[childNum], int64(i+l.data.offset), int64(i+l.data.offset+1), - r.children[childNum], int64(i+r.data.offset), int64(i+r.data.offset+1), opt) + r.children[childNum], int64(i+r.data.offset), int64(i+r.data.offset+1), childOpt) if !eq { return false } @@ -613,18 +620,23 @@ func (a *DenseUnion) setData(data *Data) { } } -func (a *DenseUnion) GetOneForMarshal(i int) interface{} { +func (a *DenseUnion) GetOneForMarshalNullable(i int, nullable bool) any { typeID := a.RawTypeCodes()[i] childID := a.ChildID(i) data := a.Field(childID) offset := int(a.RawValueOffsets()[i]) - if data.IsNull(offset) { - return nil + childNullable := a.unionType.Fields()[childID].Nullable + if childNullable && data.IsNull(offset) { + return []any{typeID, nil} } - return []interface{}{typeID, data.GetOneForMarshal(offset)} + return []any{typeID, getOneForMarshalNullable(data, offset, childNullable)} +} + +func (a *DenseUnion) GetOneForMarshal(i int) any { + return a.GetOneForMarshalNullable(i, true) } func (a *DenseUnion) MarshalJSON() ([]byte, error) { @@ -682,7 +694,7 @@ func (a *DenseUnion) String() string { return b.String() } -func arrayDenseUnionEqual(l, r *DenseUnion) bool { +func arrayDenseUnionEqual(l, r *DenseUnion, opt equalOption) bool { childIDs := l.unionType.ChildIDs() leftCodes, rightCodes := l.RawTypeCodes(), r.RawTypeCodes() leftOffsets, rightOffsets := l.RawValueOffsets(), r.RawValueOffsets() @@ -694,8 +706,9 @@ func arrayDenseUnionEqual(l, r *DenseUnion) bool { } childNum := childIDs[typeID] - eq := SliceEqual(l.children[childNum], int64(leftOffsets[i]), int64(leftOffsets[i]+1), - r.children[childNum], int64(rightOffsets[i]), int64(rightOffsets[i]+1)) + childOpt := withNullable(opt, l.unionType.Fields()[childNum].Nullable) + eq := sliceEqual(l.children[childNum], int64(leftOffsets[i]), int64(leftOffsets[i]+1), + r.children[childNum], int64(rightOffsets[i]), int64(rightOffsets[i]+1), childOpt) if !eq { return false } @@ -715,8 +728,9 @@ func arrayDenseUnionApproxEqual(l, r *DenseUnion, opt equalOption) bool { } childNum := childIDs[typeID] + childOpt := withNullable(opt, l.unionType.Fields()[childNum].Nullable) eq := sliceApproxEqual(l.children[childNum], int64(leftOffsets[i]), int64(leftOffsets[i]+1), - r.children[childNum], int64(rightOffsets[i]), int64(rightOffsets[i]+1), opt) + r.children[childNum], int64(rightOffsets[i]), int64(rightOffsets[i]+1), childOpt) if !eq { return false } diff --git a/arrow/array/util.go b/arrow/array/util.go index 136ed3537..83b892bab 100644 --- a/arrow/array/util.go +++ b/arrow/array/util.go @@ -284,7 +284,7 @@ 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) + cols[fields[j].Name] = getOneForMarshalNullable(c, i, rec.Schema().Field(j).Nullable) } if err := enc.Encode(cols); err != nil { return err @@ -293,6 +293,50 @@ func RecordToJSON(rec arrow.RecordBatch, w io.Writer) error { return nil } +// getOneForMarshalNullable dispatches to an Array's field-local-nullability-aware +// marshaler when it implements arrow.NullableMarshaler, otherwise it falls back to +// the plain GetOneForMarshal (which decides null purely from the validity bitmap). +// This lets containers propagate a field's Nullable flag without every Array having +// to implement the optional interface. +func getOneForMarshalNullable(a arrow.Array, i int, nullable bool) interface{} { + if nm, ok := a.(arrow.NullableMarshaler); ok { + return nm.GetOneForMarshalNullable(i, nullable) + } + + // ExtensionArrayBase intentionally does not implement arrow.NullableMarshaler + // (see extension.go), so extension arrays reach here. Preserve the extension's + // own logical JSON by default; the only field-local nullability we can safely + // add is that a null slot in a non-nullable field must not serialize as null. + // Ask the extension first (honoring any custom GetOneForMarshal override); only + // if it would emit null do we fall back to the storage value. Plain wrapper + // extensions (e.g. Parametric*Array) return nil at a null slot and thus round + // trip through storage, without bypassing a custom representation. + if ext, ok := a.(ExtensionArray); ok && !nullable && a.IsNull(i) { + if v := a.GetOneForMarshal(i); v != nil { + return v + } + return getOneForMarshalNullable(ext.Storage(), i, false) + } + + return a.GetOneForMarshal(i) +} + +// marshalListElemsNullable marshals a single list element's child slice +// element-by-element, honoring the element field's nullability so a +// non-nullable element field serializes underlying values instead of JSON +// null at null child slots (matching struct/record field-local behavior). +func marshalListElemsNullable(slice arrow.Array, elemNullable bool) json.RawMessage { + vals := make([]interface{}, slice.Len()) + for k := 0; k < slice.Len(); k++ { + vals[k] = getOneForMarshalNullable(slice, k, elemNullable) + } + v, err := json.Marshal(vals) + if err != nil { + panic(err) + } + return json.RawMessage(v) +} + func TableFromJSON(mem memory.Allocator, sc *arrow.Schema, recJSON []string, opt ...FromJSONOption) (arrow.Table, error) { batches := make([]arrow.RecordBatch, len(recJSON)) for i, batchJSON := range recJSON { diff --git a/arrow/array/util_test.go b/arrow/array/util_test.go index eb3de6a8b..84103d454 100644 --- a/arrow/array/util_test.go +++ b/arrow/array/util_test.go @@ -452,29 +452,50 @@ func TestArrRecordsJSONRoundTrip(t *testing.T) { continue } t.Run(k, func(t *testing.T) { - var buf bytes.Buffer - assert.NotPanics(t, func() { - enc := json.NewEncoder(&buf) - for _, r := range v { - if err := enc.Encode(r); err != nil { - panic(err) - } + for _, nullable := range []bool{true, false} { + var name string + if nullable { + name = "nullable" + } else { + name = "non-nullable" } - }) - - rdr := bytes.NewReader(buf.Bytes()) - var cur int64 - - mem := memory.NewCheckedAllocator(memory.NewGoAllocator()) - defer mem.AssertSize(t, 0) - - for _, r := range v { - rec, off, err := array.RecordFromJSON(mem, r.Schema(), rdr, array.WithStartOffset(cur)) - assert.NoError(t, err) - defer rec.Release() - assert.Truef(t, array.RecordApproxEqual(r, rec), "expected: %s\ngot: %s\n", r, rec) - cur += off + t.Run(name, func(t *testing.T) { + fields := v[0].Schema().Fields() + for i := range fields { + fields[i].Nullable = nullable + } + meta := v[0].Schema().Metadata() + schema := arrow.NewSchema(fields, &meta) + + var buf bytes.Buffer + assert.NotPanics(t, func() { + enc := json.NewEncoder(&buf) + for _, rawBatch := range v { + batch := array.NewRecordBatch(schema, rawBatch.Columns(), rawBatch.NumRows()) + if err := enc.Encode(batch); err != nil { + panic(err) + } + } + }) + + rdr := bytes.NewReader(buf.Bytes()) + var cur int64 + + mem := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer mem.AssertSize(t, 0) + + for _, rawBatch := range v { + batch := array.NewRecordBatch(schema, rawBatch.Columns(), rawBatch.NumRows()) + + rec, off, err := array.RecordFromJSON(mem, schema, rdr, array.WithStartOffset(cur)) + assert.NoError(t, err) + defer rec.Release() + + assert.Truef(t, array.RecordApproxEqual(batch, rec), "expected: %s\ngot: %s\n", batch, rec) + cur += off + } + }) } }) } diff --git a/arrow/compute/vector_sort_test.go b/arrow/compute/vector_sort_test.go index 39bf5e95f..5a15428e7 100644 --- a/arrow/compute/vector_sort_test.go +++ b/arrow/compute/vector_sort_test.go @@ -1349,8 +1349,8 @@ func TestVectorSortIndicesCppRecordBatchParity(t *testing.T) { t.Run("NoNull", 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": 3, "b": 5}, @@ -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}, @@ -1396,8 +1396,8 @@ func TestVectorSortIndicesCppRecordBatchParity(t *testing.T) { t.Run("NaN", 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() @@ -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}, @@ -1486,9 +1486,9 @@ func TestVectorSortIndicesCppRecordBatchParity(t *testing.T) { ts := &arrow.TimestampType{Unit: arrow.Microsecond} fsb3 := &arrow.FixedSizeBinaryType{ByteWidth: 3} schema := arrow.NewSchema([]arrow.Field{ - {Name: "a", Type: ts}, - {Name: "b", Type: arrow.BinaryTypes.LargeString}, - {Name: "c", Type: fsb3}, + {Name: "a", Type: ts, Nullable: true}, + {Name: "b", Type: arrow.BinaryTypes.LargeString, Nullable: true}, + {Name: "c", Type: fsb3, Nullable: true}, }, nil) ba := array.NewTimestampBuilder(mem, ts) defer ba.Release() @@ -1535,8 +1535,8 @@ func TestVectorSortIndicesCppRecordBatchParity(t *testing.T) { d128 := &arrow.Decimal128Type{Precision: 3, Scale: 1} d256 := &arrow.Decimal256Type{Precision: 4, Scale: 2} schema := arrow.NewSchema([]arrow.Field{ - {Name: "a", Type: d128}, - {Name: "b", Type: d256}, + {Name: "a", Type: d128, Nullable: true}, + {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) @@ -1719,8 +1719,8 @@ func TestVectorSortIndicesCppTableParity(t *testing.T) { t.Run("HeterogenousChunking", func(t *testing.T) { s := 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) a0, _, err := array.FromJSON(mem, arrow.PrimitiveTypes.Float32, strings.NewReader("[null, 1]")) require.NoError(t, err) diff --git a/arrow/extensions/bool8.go b/arrow/extensions/bool8.go index 97038a1bf..8f784a043 100644 --- a/arrow/extensions/bool8.go +++ b/arrow/extensions/bool8.go @@ -114,13 +114,17 @@ func (a *Bool8Array) MarshalJSON() ([]byte, error) { return json.Marshal(values) } -func (a *Bool8Array) GetOneForMarshal(i int) interface{} { - if a.IsNull(i) { +func (a *Bool8Array) GetOneForMarshalNullable(i int, nullable bool) interface{} { + if nullable && a.IsNull(i) { return nil } return a.Value(i) } +func (a *Bool8Array) GetOneForMarshal(i int) interface{} { + return a.GetOneForMarshalNullable(i, true) +} + // boolToInt8 performs the simple scalar conversion of bool to the canonical int8 // value for the Bool8Type. func boolToInt8(v bool) int8 { diff --git a/arrow/extensions/json.go b/arrow/extensions/json.go index 3f46b50e3..4b7194f94 100644 --- a/arrow/extensions/json.go +++ b/arrow/extensions/json.go @@ -116,16 +116,20 @@ func (a *JSONArray) ValueBytes(i int) []byte { return b } -// ValueJSON wraps the underlying string value as a json.RawMessage, -// or returns nil if the array value is null. -func (a *JSONArray) ValueJSON(i int) json.RawMessage { +func (a *JSONArray) valueJSON(i int, nullable bool) json.RawMessage { var val json.RawMessage - if a.IsValid(i) { + if !nullable || a.IsValid(i) { val = json.RawMessage(a.Storage().(array.StringLike).Value(i)) } return val } +// ValueJSON wraps the underlying string value as a json.RawMessage, +// or returns nil if the array value is null. +func (a *JSONArray) ValueJSON(i int) json.RawMessage { + return a.valueJSON(i, true) +} + // MarshalJSON implements json.Marshaler. // Marshaling json.RawMessage is a no-op, except that nil values will // be marshaled as a JSON null. @@ -138,8 +142,12 @@ func (a *JSONArray) MarshalJSON() ([]byte, error) { } // GetOneForMarshal implements arrow.Array. +func (a *JSONArray) GetOneForMarshalNullable(i int, nullable bool) interface{} { + return a.valueJSON(i, nullable) +} + func (a *JSONArray) GetOneForMarshal(i int) interface{} { - return a.ValueJSON(i) + return a.GetOneForMarshalNullable(i, true) } var ( diff --git a/arrow/extensions/uuid.go b/arrow/extensions/uuid.go index 9aac02253..cbe4d7c22 100644 --- a/arrow/extensions/uuid.go +++ b/arrow/extensions/uuid.go @@ -195,13 +195,17 @@ func (a *UUIDArray) MarshalJSON() ([]byte, error) { return json.Marshal(vals) } -func (a *UUIDArray) GetOneForMarshal(i int) interface{} { - if a.IsValid(i) { +func (a *UUIDArray) GetOneForMarshalNullable(i int, nullable bool) interface{} { + if !nullable || a.IsValid(i) { return a.Value(i) } return nil } +func (a *UUIDArray) GetOneForMarshal(i int) interface{} { + return a.GetOneForMarshalNullable(i, true) +} + // UUIDType is a simple extension type that represents a FixedSizeBinary(16) // to be used for representing UUIDs type UUIDType struct { 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/extensions/variant.go b/arrow/extensions/variant.go index 379822c4b..e2123485b 100644 --- a/arrow/extensions/variant.go +++ b/arrow/extensions/variant.go @@ -599,8 +599,8 @@ func (v *VariantArray) MarshalJSON() ([]byte, error) { return json.Marshal(values) } -func (v *VariantArray) GetOneForMarshal(i int) any { - if v.IsNull(i) { +func (v *VariantArray) GetOneForMarshalNullable(i int, nullable bool) any { + if nullable && v.IsNull(i) { return nil } @@ -612,6 +612,10 @@ func (v *VariantArray) GetOneForMarshal(i int) any { return val.Value() } +func (v *VariantArray) GetOneForMarshal(i int) any { + return v.GetOneForMarshalNullable(i, true) +} + type variantReader interface { IsNull(i int) bool Value(i int) (variant.Value, error) diff --git a/arrow/internal/arrdata/arrdata.go b/arrow/internal/arrdata/arrdata.go index 095571a8f..d95ee7cc7 100644 --- a/arrow/internal/arrdata/arrdata.go +++ b/arrow/internal/arrdata/arrdata.go @@ -192,59 +192,59 @@ 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) - mask := []bool{true, false, false, true, true, true, false, true} + innerValids := []bool{true, false, false, true, true} chunks := [][]arrow.Array{ { structOf(mem, dtype, [][]arrow.Array{ { - arrayOf(mem, []int32{-1, -2, -3, -4, -5}, mask[:5]), - arrayOf(mem, []string{"111", "222", "333", "444", "555"}, mask[:5]), + arrayOf(mem, []int32{-1, -2, -3, -4, -5}, innerValids), + arrayOf(mem, []string{"111", "222", "333", "444", "555"}, innerValids), }, { - arrayOf(mem, []int32{-11, -12, -13, -14, -15}, mask[:5]), - arrayOf(mem, []string{"1111", "1222", "1333", "1444", "1555"}, mask[:5]), + arrayOf(mem, []int32{-11, -12, -13, -14, -15}, innerValids), + arrayOf(mem, []string{"1111", "1222", "1333", "1444", "1555"}, innerValids), }, { - arrayOf(mem, []int32{-21, -22, -23, -24, -25}, mask[:5]), - arrayOf(mem, []string{"2111", "2222", "2333", "2444", "2555"}, mask[:5]), + arrayOf(mem, []int32{-21, -22, -23, -24, -25}, innerValids), + arrayOf(mem, []string{"2111", "2222", "2333", "2444", "2555"}, innerValids), }, { - arrayOf(mem, []int32{-31, -32, -33, -34, -35}, mask[:5]), - arrayOf(mem, []string{"3111", "3222", "3333", "3444", "3555"}, mask[:5]), + arrayOf(mem, []int32{-31, -32, -33, -34, -35}, innerValids), + arrayOf(mem, []string{"3111", "3222", "3333", "3444", "3555"}, innerValids), }, { - arrayOf(mem, []int32{-41, -42, -43, -44, -45}, mask[:5]), - arrayOf(mem, []string{"4111", "4222", "4333", "4444", "4555"}, mask[:5]), + arrayOf(mem, []int32{-41, -42, -43, -44, -45}, innerValids), + arrayOf(mem, []string{"4111", "4222", "4333", "4444", "4555"}, innerValids), }, }, []bool{true, false, true, true, true}), }, { structOf(mem, dtype, [][]arrow.Array{ { - arrayOf(mem, []int32{1, 2, 3, 4, 5}, mask[:5]), - arrayOf(mem, []string{"-111", "-222", "-333", "-444", "-555"}, mask[:5]), + arrayOf(mem, []int32{1, 2, 3, 4, 5}, innerValids), + arrayOf(mem, []string{"-111", "-222", "-333", "-444", "-555"}, innerValids), }, { - arrayOf(mem, []int32{11, 12, 13, 14, 15}, mask[:5]), - arrayOf(mem, []string{"-1111", "-1222", "-1333", "-1444", "-1555"}, mask[:5]), + arrayOf(mem, []int32{11, 12, 13, 14, 15}, innerValids), + arrayOf(mem, []string{"-1111", "-1222", "-1333", "-1444", "-1555"}, innerValids), }, { - arrayOf(mem, []int32{21, 22, 23, 24, 25}, mask[:5]), - arrayOf(mem, []string{"-2111", "-2222", "-2333", "-2444", "-2555"}, mask[:5]), + arrayOf(mem, []int32{21, 22, 23, 24, 25}, innerValids), + arrayOf(mem, []string{"-2111", "-2222", "-2333", "-2444", "-2555"}, innerValids), }, { - arrayOf(mem, []int32{31, 32, 33, 34, 35}, mask[:5]), - arrayOf(mem, []string{"-3111", "-3222", "-3333", "-3444", "-3555"}, mask[:5]), + arrayOf(mem, []int32{31, 32, 33, 34, 35}, innerValids), + arrayOf(mem, []string{"-3111", "-3222", "-3333", "-3444", "-3555"}, innerValids), }, { - arrayOf(mem, []int32{41, 42, 43, 44, 45}, mask[:5]), - arrayOf(mem, []string{"-4111", "-4222", "-4333", "-4444", "-4555"}, mask[:5]), + arrayOf(mem, []int32{41, 42, 43, 44, 45}, innerValids), + arrayOf(mem, []string{"-4111", "-4222", "-4333", "-4444", "-4555"}, innerValids), }, }, []bool{true, false, false, true, true}), }, diff --git a/arrow/internal/arrjson/arrjson_test.go b/arrow/internal/arrjson/arrjson_test.go index 7e2f386fc..faeecfc01 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": [] } ] diff --git a/arrow/ipc/cmd/arrow-ls/main_test.go b/arrow/ipc/cmd/arrow-ls/main_test.go index f90e4a800..0f3b5377c 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 `, }, @@ -221,7 +221,7 @@ records: 3 name: "structs", want: `schema: fields: 1 - - struct_nullable: type=struct, nullable + - struct_nullable: type=struct, nullable records: 2 `, }, @@ -230,7 +230,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/metadata_test.go b/arrow/ipc/metadata_test.go index ac8820dcc..63e700b8b 100644 --- a/arrow/ipc/metadata_test.go +++ b/arrow/ipc/metadata_test.go @@ -257,12 +257,15 @@ func TestUnrecognizedExtensionType(t *testing.T) { // create a record batch with the same data, but the field should contain the // extension metadata and be of the storage type instead of being the extension type. - extMetadata := arrow.NewMetadata([]string{ExtensionTypeKeyName, ExtensionMetadataKeyName}, []string{"uuid", "uuid-serialized"}) + extMetadata := arrow.NewMetadata([]string{ExtensionTypeKeyName, ExtensionMetadataKeyName}, []string{"arrow.uuid", ""}) batchNoExt := array.NewRecordBatch( arrow.NewSchema([]arrow.Field{ {Name: "f0", Type: storageArr.DataType(), Nullable: true, Metadata: extMetadata}, }, nil), []arrow.Array{storageArr}, 4) defer batchNoExt.Release() + // RecordEqual ignores field metadata, so explicitly verify the unrecognized + // extension metadata is preserved on the read-back field. + assert.Truef(t, rec.Schema().Field(0).Metadata.Equal(extMetadata), "expected metadata %v, got %v", extMetadata, rec.Schema().Field(0).Metadata) assert.Truef(t, array.RecordEqual(rec, batchNoExt), "expected: %s\ngot: %s\n", batchNoExt, rec) } diff --git a/internal/json/json.go b/internal/json/json.go index b4c4c9f6e..7fdd0d863 100644 --- a/internal/json/json.go +++ b/internal/json/json.go @@ -20,6 +20,7 @@ package json import ( + "bytes" "io" "github.com/goccy/go-json" @@ -49,3 +50,7 @@ func NewDecoder(r io.Reader) *Decoder { func NewEncoder(w io.Writer) *Encoder { return json.NewEncoder(w) } + +func IsNullMessage(m RawMessage) bool { + return bytes.Equal(m, []byte("null")) +} diff --git a/internal/json/json_stdlib.go b/internal/json/json_stdlib.go index 3031029d8..06071fada 100644 --- a/internal/json/json_stdlib.go +++ b/internal/json/json_stdlib.go @@ -20,6 +20,7 @@ package json import ( + "bytes" "io" "encoding/json" @@ -49,3 +50,7 @@ func NewDecoder(r io.Reader) *Decoder { func NewEncoder(w io.Writer) *Encoder { return json.NewEncoder(w) } + +func IsNullMessage(m RawMessage) bool { + return bytes.Equal(m, []byte("null")) +}