Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

ARROW-17390: [Go] Add union scalar types #13860

Merged
merged 1 commit into from
Aug 11, 2022
Merged

ARROW-17390: [Go] Add union scalar types #13860

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
192 changes: 192 additions & 0 deletions go/arrow/scalar/nested.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -520,3 +520,195 @@ func (s *Dictionary) GetEncodedValue() (Scalar, error) {
func (s *Dictionary) value() interface{} {
return s.Value.Index.value()
}

type Union interface {
Scalar
ChildValue() Scalar
Release()
}

type SparseUnion struct {
scalar

TypeCode arrow.UnionTypeCode
Value []Scalar
ChildID int
}

func (s *SparseUnion) equals(rhs Scalar) bool {
right := rhs.(*SparseUnion)
return Equals(s.ChildValue(), right.ChildValue())
}

func (s *SparseUnion) value() interface{} { return s.ChildValue() }

func (s *SparseUnion) String() string {
dt := s.Type.(*arrow.SparseUnionType)
val := s.ChildValue()
return "union{" + dt.Fields()[dt.ChildIDs()[s.TypeCode]].String() + " = " + val.String() + "}"
}

func (s *SparseUnion) Release() {
for _, v := range s.Value {
if v, ok := v.(Releasable); ok {
v.Release()
}
}
}

func (s *SparseUnion) Validate() (err error) {
dt := s.Type.(*arrow.SparseUnionType)
if len(dt.Fields()) != len(s.Value) {
return fmt.Errorf("sparse union scalar value had %d fields but type has %d fields", len(dt.Fields()), len(s.Value))
}

if s.TypeCode < 0 || int(s.TypeCode) >= len(dt.ChildIDs()) || dt.ChildIDs()[s.TypeCode] == arrow.InvalidUnionChildID {
return fmt.Errorf("%s scalar has invalid type code %d", dt, s.TypeCode)
}

for i, f := range dt.Fields() {
v := s.Value[i]
if !arrow.TypeEqual(f.Type, v.DataType()) {
return fmt.Errorf("%s value for field %s had incorrect type of %s", dt, f, v.DataType())
}
if err = v.Validate(); err != nil {
return err
}
}
return
}

func (s *SparseUnion) ValidateFull() (err error) {
dt := s.Type.(*arrow.SparseUnionType)
if len(dt.Fields()) != len(s.Value) {
return fmt.Errorf("sparse union scalar value had %d fields but type has %d fields", len(dt.Fields()), len(s.Value))
}

if s.TypeCode < 0 || int(s.TypeCode) >= len(dt.ChildIDs()) || dt.ChildIDs()[s.TypeCode] == arrow.InvalidUnionChildID {
return fmt.Errorf("%s scalar has invalid type code %d", dt, s.TypeCode)
}

for i, f := range dt.Fields() {
v := s.Value[i]
if !arrow.TypeEqual(f.Type, v.DataType()) {
return fmt.Errorf("%s value for field %s had incorrect type of %s", dt, f, v.DataType())
}
if err = v.ValidateFull(); err != nil {
return err
}
}
return
}

func (s *SparseUnion) CastTo(to arrow.DataType) (Scalar, error) {
if !s.Valid {
return MakeNullScalar(to), nil
}

switch to.ID() {
case arrow.STRING:
return NewStringScalar(s.String()), nil
case arrow.LARGE_STRING:
return NewLargeStringScalar(s.String()), nil
}

return nil, fmt.Errorf("cannot cast non-nil union to type other than string")
}

func (s *SparseUnion) ChildValue() Scalar { return s.Value[s.ChildID] }

func NewSparseUnionScalar(val []Scalar, code arrow.UnionTypeCode, dt *arrow.SparseUnionType) *SparseUnion {
ret := &SparseUnion{
scalar: scalar{dt, true},
TypeCode: code,
Value: val,
ChildID: dt.ChildIDs()[code],
}
ret.Valid = ret.Value[ret.ChildID].IsValid()
return ret
}

func NewSparseUnionScalarFromValue(val Scalar, idx int, dt *arrow.SparseUnionType) *SparseUnion {
code := dt.TypeCodes()[idx]
values := make([]Scalar, len(dt.Fields()))
for i, f := range dt.Fields() {
if i == idx {
values[i] = val
} else {
values[i] = MakeNullScalar(f.Type)
}
}
return NewSparseUnionScalar(values, code, dt)
}

type DenseUnion struct {
scalar

TypeCode arrow.UnionTypeCode
Value Scalar
}

func (s *DenseUnion) equals(rhs Scalar) bool {
right := rhs.(*DenseUnion)
return Equals(s.Value, right.Value)
}

func (s *DenseUnion) value() interface{} { return s.ChildValue() }

func (s *DenseUnion) String() string {
dt := s.Type.(*arrow.DenseUnionType)
return "union{" + dt.Fields()[dt.ChildIDs()[s.TypeCode]].String() + " = " + s.Value.String() + "}"
}

func (s *DenseUnion) Release() {
if v, ok := s.Value.(Releasable); ok {
v.Release()
}
}

func (s *DenseUnion) Validate() (err error) {
dt := s.Type.(*arrow.DenseUnionType)
if s.TypeCode < 0 || int(s.TypeCode) >= len(dt.ChildIDs()) || dt.ChildIDs()[s.TypeCode] == arrow.InvalidUnionChildID {
return fmt.Errorf("%s scalar has invalid type code %d", dt, s.TypeCode)
}
fieldType := dt.Fields()[dt.ChildIDs()[s.TypeCode]].Type
if !arrow.TypeEqual(fieldType, s.Value.DataType()) {
return fmt.Errorf("%s scalar with type code %d should have an underlying value of type %s, got %s",
s.Type, s.TypeCode, fieldType, s.Value.DataType())
}
return s.Value.Validate()
}

func (s *DenseUnion) ValidateFull() error {
dt := s.Type.(*arrow.DenseUnionType)
if s.TypeCode < 0 || int(s.TypeCode) >= len(dt.ChildIDs()) || dt.ChildIDs()[s.TypeCode] == arrow.InvalidUnionChildID {
return fmt.Errorf("%s scalar has invalid type code %d", dt, s.TypeCode)
}
fieldType := dt.Fields()[dt.ChildIDs()[s.TypeCode]].Type
if !arrow.TypeEqual(fieldType, s.Value.DataType()) {
return fmt.Errorf("%s scalar with type code %d should have an underlying value of type %s, got %s",
s.Type, s.TypeCode, fieldType, s.Value.DataType())
}
return s.Value.ValidateFull()
}

func (s *DenseUnion) CastTo(to arrow.DataType) (Scalar, error) {
if !s.Valid {
return MakeNullScalar(to), nil
}

switch to.ID() {
case arrow.STRING:
return NewStringScalar(s.String()), nil
case arrow.LARGE_STRING:
return NewLargeStringScalar(s.String()), nil
}

return nil, fmt.Errorf("cannot cast non-nil union to type other than string")
}

func (s *DenseUnion) ChildValue() Scalar { return s.Value }

func NewDenseUnionScalar(v Scalar, code arrow.UnionTypeCode, dt *arrow.DenseUnionType) *DenseUnion {
return &DenseUnion{scalar: scalar{dt, v.IsValid()}, TypeCode: code, Value: v}
}
81 changes: 66 additions & 15 deletions go/arrow/scalar/scalar.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -466,10 +466,6 @@ func MakeNullScalar(dt arrow.DataType) Scalar {
return makeNullFn[byte(dt.ID()&0x3f)](dt)
}

func unsupportedScalarType(dt arrow.DataType) Scalar {
panic("unsupported scalar data type: " + dt.ID().String())
}

func invalidScalarType(dt arrow.DataType) Scalar {
panic("invalid scalar type: " + dt.ID().String())
}
Expand Down Expand Up @@ -516,17 +512,33 @@ func init() {
arrow.DECIMAL128: func(dt arrow.DataType) Scalar { return &Decimal128{scalar: scalar{dt, false}} },
arrow.LIST: func(dt arrow.DataType) Scalar { return &List{scalar: scalar{dt, false}} },
arrow.STRUCT: func(dt arrow.DataType) Scalar { return &Struct{scalar: scalar{dt, false}} },
arrow.SPARSE_UNION: unsupportedScalarType,
arrow.DENSE_UNION: unsupportedScalarType,
arrow.DICTIONARY: func(dt arrow.DataType) Scalar { return NewNullDictScalar(dt) },
arrow.LARGE_STRING: func(dt arrow.DataType) Scalar { return &LargeString{&String{&Binary{scalar: scalar{dt, false}}}} },
arrow.LARGE_BINARY: func(dt arrow.DataType) Scalar { return &LargeBinary{&Binary{scalar: scalar{dt, false}}} },
arrow.LARGE_LIST: func(dt arrow.DataType) Scalar { return &LargeList{&List{scalar: scalar{dt, false}}} },
arrow.DECIMAL256: func(dt arrow.DataType) Scalar { return &Decimal256{scalar: scalar{dt, false}} },
arrow.MAP: func(dt arrow.DataType) Scalar { return &Map{&List{scalar: scalar{dt, false}}} },
arrow.EXTENSION: func(dt arrow.DataType) Scalar { return &Extension{scalar: scalar{dt, false}} },
arrow.FIXED_SIZE_LIST: func(dt arrow.DataType) Scalar { return &FixedSizeList{&List{scalar: scalar{dt, false}}} },
arrow.DURATION: func(dt arrow.DataType) Scalar { return &Duration{scalar: scalar{dt, false}} },
arrow.SPARSE_UNION: func(dt arrow.DataType) Scalar {
typ := dt.(*arrow.SparseUnionType)
if len(typ.Fields()) == 0 {
panic("cannot make scalar of empty union type")
}
values := make([]Scalar, len(typ.Fields()))
for i, f := range typ.Fields() {
values[i] = MakeNullScalar(f.Type)
}
return NewSparseUnionScalar(values, typ.TypeCodes()[0], typ)
},
arrow.DENSE_UNION: func(dt arrow.DataType) Scalar {
typ := dt.(*arrow.DenseUnionType)
if len(typ.Fields()) == 0 {
panic("cannot make scalar of empty union type")
}
return NewDenseUnionScalar(MakeNullScalar(typ.Fields()[0].Type), typ.TypeCodes()[0], typ)
},
arrow.DICTIONARY: func(dt arrow.DataType) Scalar { return NewNullDictScalar(dt) },
arrow.LARGE_STRING: func(dt arrow.DataType) Scalar { return &LargeString{&String{&Binary{scalar: scalar{dt, false}}}} },
arrow.LARGE_BINARY: func(dt arrow.DataType) Scalar { return &LargeBinary{&Binary{scalar: scalar{dt, false}}} },
arrow.LARGE_LIST: func(dt arrow.DataType) Scalar { return &LargeList{&List{scalar: scalar{dt, false}}} },
arrow.DECIMAL256: func(dt arrow.DataType) Scalar { return &Decimal256{scalar: scalar{dt, false}} },
arrow.MAP: func(dt arrow.DataType) Scalar { return &Map{&List{scalar: scalar{dt, false}}} },
arrow.EXTENSION: func(dt arrow.DataType) Scalar { return &Extension{scalar: scalar{dt, false}} },
arrow.FIXED_SIZE_LIST: func(dt arrow.DataType) Scalar { return &FixedSizeList{&List{scalar: scalar{dt, false}}} },
arrow.DURATION: func(dt arrow.DataType) Scalar { return &Duration{scalar: scalar{dt, false}} },
// invalid data types to fill out array size 2^6 - 1
63: invalidScalarType,
}
Expand Down Expand Up @@ -646,6 +658,39 @@ func GetScalar(arr arrow.Array, idx int) (Scalar, error) {
scalar.Value.Dict = arr.Dictionary()
scalar.Value.Dict.Retain()
return scalar, nil
case *array.SparseUnion:
var err error
typeCode := arr.TypeCode(idx)
children := make([]Scalar, arr.NumFields())
defer func() {
if err != nil {
for _, c := range children {
if c == nil {
break
}

if v, ok := c.(Releasable); ok {
v.Release()
}
}
}
}()

for i := range arr.UnionType().Fields() {
if children[i], err = GetScalar(arr.Field(i), idx); err != nil {
return nil, err
}
}
return NewSparseUnionScalar(children, typeCode, arr.UnionType().(*arrow.SparseUnionType)), nil
case *array.DenseUnion:
typeCode := arr.TypeCode(idx)
child := arr.Field(arr.ChildID(idx))
offset := arr.ValueOffset(idx)
value, err := GetScalar(child, int(offset))
if err != nil {
return nil, err
}
return NewDenseUnionScalar(value, typeCode, arr.UnionType().(*arrow.DenseUnionType)), nil
}

return nil, fmt.Errorf("cannot create scalar from array of type %s", arr.DataType())
Expand Down Expand Up @@ -902,6 +947,12 @@ func Hash(seed maphash.Seed, s Scalar) uint64 {
return valueHash(s.Value.Days) & valueHash(s.Value.Milliseconds)
case *MonthDayNanoInterval:
return valueHash(s.Value.Months) & valueHash(s.Value.Days) & valueHash(s.Value.Nanoseconds)
case *SparseUnion:
// typecode is ignored when comparing for equality, so don't hash it either
out ^= Hash(seed, s.Value[s.ChildID])
case *DenseUnion:
// typecode is ignored when comparing equality, so don't hash it either
out ^= Hash(seed, s.Value)
case PrimitiveScalar:
h.Write(s.Data())
hash()
Expand Down
Loading