Skip to content

Commit

Permalink
add arrow.ListLikeType interface
Browse files Browse the repository at this point in the history
  • Loading branch information
@candiduslynx
candiduslynx committed Jun 2, 2023
1 parent 018e7d3 commit 5071eaf
Show file tree
Hide file tree
Showing 2 changed files with 56 additions and 26 deletions.
10 changes: 10 additions & 0 deletions go/arrow/array/list.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -611,4 +611,14 @@ var (
_ arrow.Array = (*LargeList)(nil)
_ Builder = (*ListBuilder)(nil)
_ Builder = (*LargeListBuilder)(nil)

_ ListLike = (*List)(nil)
_ ListLike = (*LargeList)(nil)
_ ListLike = (*FixedSizeList)(nil)
_ ListLike = (*Map)(nil)

_ ListLikeBuilder = (*ListBuilder)(nil)
_ ListLikeBuilder = (*LargeListBuilder)(nil)
_ ListLikeBuilder = (*FixedSizeListBuilder)(nil)
_ ListLikeBuilder = (*MapBuilder)(nil)
)
72 changes: 46 additions & 26 deletions go/arrow/datatype_nested.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,13 +25,20 @@ import (
"github.com/apache/arrow/go/v13/arrow/internal/debug"
)

type NestedType interface {
DataType
type (
NestedType interface {
DataType

// Fields method provides a copy of NestedType fields
// (so it can be safely mutated and will not result in updating the NestedType).
Fields() []Field
}
// Fields method provides a copy of NestedType fields
// (so it can be safely mutated and will not result in updating the NestedType).
Fields() []Field
}

ListLikeType interface {
DataType
Elem() DataType
}
)

// ListType describes a nested type in which each array slot contains
// a variable-size sequence of values, all having the same relative type.
Expand Down Expand Up @@ -97,11 +104,11 @@ func (t *ListType) ElemField() Field {

func (t *ListType) Fields() []Field { return []Field{t.ElemField()} }

func (ListType) Layout() DataTypeLayout {
func (*ListType) Layout() DataTypeLayout {
return DataTypeLayout{Buffers: []BufferSpec{SpecBitmap(), SpecFixedWidth(Int32SizeBytes)}}
}

func (ListType) OffsetTypeTraits() OffsetTraits { return Int32Traits }
func (*ListType) OffsetTypeTraits() OffsetTraits { return Int32Traits }

type LargeListType struct {
ListType
Expand All @@ -121,11 +128,11 @@ func (t *LargeListType) Fingerprint() string {
return ""
}

func (LargeListType) Layout() DataTypeLayout {
func (*LargeListType) Layout() DataTypeLayout {
return DataTypeLayout{Buffers: []BufferSpec{SpecBitmap(), SpecFixedWidth(Int64SizeBytes)}}
}

func (LargeListType) OffsetTypeTraits() OffsetTraits { return Int64Traits }
func (*LargeListType) OffsetTypeTraits() OffsetTraits { return Int64Traits }

func LargeListOfField(f Field) *LargeListType {
if f.Type == nil {
Expand All @@ -134,18 +141,18 @@ func LargeListOfField(f Field) *LargeListType {
return &LargeListType{ListType{elem: f}}
}

// ListOf returns the list type with element type t.
// For example, if t represents int32, ListOf(t) represents []int32.
// LargeListOf returns the list type with element type t.
// For example, if t represents int32, LargeListOf(t) represents []int32.
//
// ListOf panics if t is nil or invalid. NullableElem defaults to true
// LargeListOf panics if t is nil or invalid. NullableElem defaults to true
func LargeListOf(t DataType) *LargeListType {
if t == nil {
panic("arrow: nil DataType")
}
return &LargeListType{ListType{elem: Field{Name: "item", Type: t, Nullable: true}}}
}

// ListOfNonNullable is like ListOf but NullableElem defaults to false, indicating
// LargeListOfNonNullable is like ListOf but NullableElem defaults to false, indicating
// that the child type should be marked as non-nullable.
func LargeListOfNonNullable(t DataType) *LargeListType {
if t == nil {
Expand Down Expand Up @@ -230,7 +237,7 @@ func (t *FixedSizeListType) Fingerprint() string {

func (t *FixedSizeListType) Fields() []Field { return []Field{t.ElemField()} }

func (FixedSizeListType) Layout() DataTypeLayout {
func (*FixedSizeListType) Layout() DataTypeLayout {
return DataTypeLayout{Buffers: []BufferSpec{SpecBitmap()}}
}

Expand Down Expand Up @@ -330,7 +337,7 @@ func (t *StructType) Fingerprint() string {
return b.String()
}

func (StructType) Layout() DataTypeLayout {
func (*StructType) Layout() DataTypeLayout {
return DataTypeLayout{Buffers: []BufferSpec{SpecBitmap()}}
}

Expand Down Expand Up @@ -389,12 +396,10 @@ func (t *MapType) KeyType() DataType { return t.KeyField().Type }
func (t *MapType) ItemField() Field { return t.value.Elem().(*StructType).Field(1) }
func (t *MapType) ItemType() DataType { return t.ItemField().Type }
func (t *MapType) ValueType() *StructType { return t.value.Elem().(*StructType) }
func (t *MapType) ValueField() Field {
return Field{
Name: "entries",
Type: t.ValueType(),
}
}
func (t *MapType) ValueField() Field { return Field{Name: "entries", Type: t.ValueType()} }

// Elem returns the MapType's element type (if treating MapType as ListLikeType)
func (t *MapType) Elem() DataType { return t.ValueType() }

func (t *MapType) SetItemNullable(nullable bool) {
t.value.Elem().(*StructType).fields[1].Nullable = nullable
Expand All @@ -420,7 +425,7 @@ func (t *MapType) Layout() DataTypeLayout {
return t.value.Layout()
}

func (MapType) OffsetTypeTraits() OffsetTraits { return Int32Traits }
func (*MapType) OffsetTypeTraits() OffsetTraits { return Int32Traits }

type (
// UnionTypeCode is an alias to int8 which is the type of the ids
Expand Down Expand Up @@ -502,14 +507,14 @@ func (t *unionType) init(fields []Field, typeCodes []UnionTypeCode) {

// Fields method provides a copy of union type fields
// (so it can be safely mutated and will not result in updating the union type).
func (t unionType) Fields() []Field {
func (t *unionType) Fields() []Field {
fields := make([]Field, len(t.children))
copy(fields, t.children)
return fields
}

func (t unionType) TypeCodes() []UnionTypeCode { return t.typeCodes }
func (t unionType) ChildIDs() []int { return t.childIDs[:] }
func (t *unionType) TypeCodes() []UnionTypeCode { return t.typeCodes }
func (t *unionType) ChildIDs() []int { return t.childIDs[:] }

func (t *unionType) validate(fields []Field, typeCodes []UnionTypeCode, _ UnionMode) error {
if len(fields) != len(typeCodes) {
Expand Down Expand Up @@ -767,7 +772,22 @@ func (f Field) String() string {

var (
_ DataType = (*ListType)(nil)
_ DataType = (*LargeListType)(nil)
_ DataType = (*FixedSizeListType)(nil)
_ DataType = (*StructType)(nil)
_ DataType = (*MapType)(nil)
_ DataType = (*DenseUnionType)(nil)
_ DataType = (*SparseUnionType)(nil)

_ NestedType = (*ListType)(nil)
_ NestedType = (*LargeListType)(nil)
_ NestedType = (*FixedSizeListType)(nil)
_ NestedType = (*MapType)(nil)
_ NestedType = (*DenseUnionType)(nil)
_ NestedType = (*SparseUnionType)(nil)

_ ListLikeType = (*ListType)(nil)
_ ListLikeType = (*LargeListType)(nil)
_ ListLikeType = (*FixedSizeListType)(nil)
_ ListLikeType = (*MapType)(nil)
)

0 comments on commit 5071eaf

Please sign in to comment.