diff --git a/arrow/array/struct.go b/arrow/array/struct.go index 07505de1..47e30071 100644 --- a/arrow/array/struct.go +++ b/arrow/array/struct.go @@ -61,6 +61,9 @@ func NewStructArrayWithFieldsAndNulls(cols []arrow.Array, fields []arrow.Field, } length := cols[0].Len() + if err := validateStructArrayOffset(offset, length); err != nil { + return nil, err + } children := make([]arrow.ArrayData, len(cols)) for i, c := range cols { if length != c.Len() { @@ -97,6 +100,9 @@ func NewStructArrayWithNulls(cols []arrow.Array, names []string, nullBitmap *mem return nil, fmt.Errorf("%w: can't infer struct array length with 0 child arrays", arrow.ErrInvalid) } length := cols[0].Len() + if err := validateStructArrayOffset(offset, length); err != nil { + return nil, err + } children := make([]arrow.ArrayData, len(cols)) fields := make([]arrow.Field, len(cols)) for i, c := range cols { @@ -108,11 +114,24 @@ func NewStructArrayWithNulls(cols []arrow.Array, names []string, nullBitmap *mem fields[i].Type = c.DataType() fields[i].Nullable = true } - data := NewData(arrow.StructOf(fields...), length, []*memory.Buffer{nullBitmap}, children, nullCount, offset) + if nullBitmap == nil { + if nullCount > 0 { + return nil, fmt.Errorf("%w: null count is greater than 0 but null bitmap is nil", arrow.ErrInvalid) + } + nullCount = 0 + } + data := NewData(arrow.StructOf(fields...), length-offset, []*memory.Buffer{nullBitmap}, children, nullCount, offset) defer data.Release() return NewStructData(data), nil } +func validateStructArrayOffset(offset, length int) error { + if offset < 0 || offset > length { + return fmt.Errorf("%w: invalid offset %d for length %d", arrow.ErrInvalid, offset, length) + } + return nil +} + // NewStructData returns a new Struct array value from data. func NewStructData(data arrow.ArrayData) *Struct { a := &Struct{} @@ -149,7 +168,7 @@ func (a *Struct) String() string { if arrow.IsUnion(v.DataType().ID()) { fmt.Fprintf(o, "%v", v) continue - } else if !bytes.Equal(structBitmap, v.NullBitmapBytes()) { + } else if a.Offset() != v.Data().Offset() || !bytes.Equal(structBitmap, v.NullBitmapBytes()) { masked := a.newStructFieldWithParentValidityMask(i) fmt.Fprintf(o, "%v", masked) masked.Release() @@ -169,17 +188,19 @@ func (a *Struct) String() string { func (a *Struct) newStructFieldWithParentValidityMask(fieldIndex int) arrow.Array { field := a.Field(fieldIndex) nullBitmapBytes := field.NullBitmapBytes() + fieldOffset := field.Data().Offset() var maskedNullBitmapBytes []byte if len(nullBitmapBytes) == 0 { - maskedNullBitmapBytes = make([]byte, int(bitutil.BytesForBits(int64(field.Len())))) - bitutil.SetBitsTo(maskedNullBitmapBytes, 0, int64(field.Len()), true) + fieldEnd := int64(fieldOffset + field.Len()) + maskedNullBitmapBytes = make([]byte, int(bitutil.BytesForBits(fieldEnd))) + bitutil.SetBitsTo(maskedNullBitmapBytes, 0, fieldEnd, true) } else { maskedNullBitmapBytes = make([]byte, len(nullBitmapBytes)) copy(maskedNullBitmapBytes, nullBitmapBytes) } for i := 0; i < field.Len(); i++ { if a.IsNull(i) { - bitutil.ClearBit(maskedNullBitmapBytes, i) + bitutil.ClearBit(maskedNullBitmapBytes, fieldOffset+i) } } data := NewSliceData(field.Data(), 0, int64(field.Len())).(*Data) @@ -191,6 +212,7 @@ func (a *Struct) newStructFieldWithParentValidityMask(fieldIndex int) arrow.Arra } bufs[0] = memory.NewBufferBytes(maskedNullBitmapBytes) data.buffers = bufs + data.SetNullN(UnknownNullCount) maskedField := MakeFromData(data) return maskedField } diff --git a/arrow/array/struct_test.go b/arrow/array/struct_test.go index 216b353a..a6d4f1c7 100644 --- a/arrow/array/struct_test.go +++ b/arrow/array/struct_test.go @@ -475,6 +475,122 @@ func TestStructArrayStringerMasksRequiredChildWithoutNullBitmap(t *testing.T) { }) } +func TestNewStructArrayWithNullsValidatesOffset(t *testing.T) { + builder := array.NewInt32Builder(memory.DefaultAllocator) + defer builder.Release() + builder.AppendValues([]int32{1, 2, 3}, nil) + child := builder.NewArray() + defer child.Release() + + for _, test := range []struct { + name string + offset int + length int + }{ + {name: "zero", offset: 0, length: 3}, + {name: "sliced", offset: 1, length: 2}, + {name: "empty", offset: 3, length: 0}, + } { + t.Run(test.name, func(t *testing.T) { + arr, err := array.NewStructArrayWithNulls( + []arrow.Array{child}, []string{"value"}, nil, 0, test.offset) + require.NoError(t, err) + defer arr.Release() + assert.Equal(t, test.length, arr.Len()) + assert.Equal(t, test.length, arr.Field(0).Len()) + }) + } + + for _, offset := range []int{-1, 4} { + arr, err := array.NewStructArrayWithNulls( + []arrow.Array{child}, []string{"value"}, nil, 0, offset) + require.ErrorIs(t, err, arrow.ErrInvalid) + require.Nil(t, arr) + } + + _, err := array.NewStructArrayWithNulls( + []arrow.Array{child}, []string{"value"}, nil, 1, 0) + require.ErrorIs(t, err, arrow.ErrInvalid) + + fields := []arrow.Field{{Name: "value", Type: arrow.PrimitiveTypes.Int32}} + for _, offset := range []int{-1, 4} { + arr, err := array.NewStructArrayWithFieldsAndNulls( + []arrow.Array{child}, fields, nil, 0, offset) + require.ErrorIs(t, err, arrow.ErrInvalid) + require.Nil(t, arr) + } +} + +func TestStructStringMasksParentValidityAtPhysicalOffset(t *testing.T) { + fields := []arrow.Field{{Name: "value", Type: arrow.PrimitiveTypes.Int32}} + parentNulls := memory.NewBufferBytes([]byte{0x05}) + + t.Run("child validity bitmap", func(t *testing.T) { + values := make([]int32, 11) + for i := range values { + values[i] = int32(i) + } + validity := memory.NewBufferBytes([]byte{0xff, 0x07}) + valueBuffer := memory.NewBufferBytes(arrow.Int32Traits.CastToBytes(values)) + childData := array.NewData(arrow.PrimitiveTypes.Int32, len(values), []*memory.Buffer{ + validity, valueBuffer, + }, nil, 0, 0) + defer childData.Release() + child := array.MakeFromData(childData) + defer child.Release() + childSlice := array.NewSlice(child, 8, 11) + defer childSlice.Release() + + arr, err := array.NewStructArrayWithFieldsAndNulls( + []arrow.Array{childSlice}, fields, parentNulls, 1, 1) + require.NoError(t, err) + defer arr.Release() + assert.Equal(t, "{[(null) 10]}", arr.String()) + }) + + t.Run("child without validity bitmap", func(t *testing.T) { + builder := array.NewInt32Builder(memory.DefaultAllocator) + defer builder.Release() + values := make([]int32, 11) + for i := range values { + values[i] = int32(i) + } + builder.AppendValues(values, nil) + child := builder.NewArray() + defer child.Release() + childSlice := array.NewSlice(child, 8, 11) + defer childSlice.Release() + + arr, err := array.NewStructArrayWithFieldsAndNulls( + []arrow.Array{childSlice}, fields, parentNulls, 1, 1) + require.NoError(t, err) + defer arr.Release() + assert.Equal(t, "{[(null) 10]}", arr.String()) + }) + + t.Run("identical bitmap with different physical offsets", func(t *testing.T) { + values := []int32{10, 11, 12, 13} + validity := memory.NewBufferBytes([]byte{0x0d}) + valueBuffer := memory.NewBufferBytes(arrow.Int32Traits.CastToBytes(values)) + childData := array.NewData(arrow.PrimitiveTypes.Int32, len(values), []*memory.Buffer{ + validity, valueBuffer, + }, nil, 0, 0) + defer childData.Release() + child := array.MakeFromData(childData) + defer child.Release() + childSlice := array.NewSlice(child, 1, 4) + defer childSlice.Release() + parentNulls := memory.NewBufferBytes([]byte{0x0d}) + defer parentNulls.Release() + + arr, err := array.NewStructArrayWithFieldsAndNulls( + []arrow.Array{childSlice}, fields, parentNulls, 1, 1) + require.NoError(t, err) + defer arr.Release() + assert.Equal(t, "{[(null) 13]}", arr.String()) + }) +} + func TestStructArrayUnmarshalJSONMissingFields(t *testing.T) { pool := memory.NewGoAllocator()