From 5071eaf8f420a004bb119e39ca4015f8fa4530d3 Mon Sep 17 00:00:00 2001 From: candiduslynx Date: Fri, 2 Jun 2023 10:15:40 +0300 Subject: [PATCH] add `arrow.ListLikeType` interface --- go/arrow/array/list.go | 10 ++++++ go/arrow/datatype_nested.go | 72 +++++++++++++++++++++++-------------- 2 files changed, 56 insertions(+), 26 deletions(-) diff --git a/go/arrow/array/list.go b/go/arrow/array/list.go index d61d980aac0e3..ddec6c5307e90 100644 --- a/go/arrow/array/list.go +++ b/go/arrow/array/list.go @@ -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) ) diff --git a/go/arrow/datatype_nested.go b/go/arrow/datatype_nested.go index 368527b4b2de9..314172a21b47f 100644 --- a/go/arrow/datatype_nested.go +++ b/go/arrow/datatype_nested.go @@ -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. @@ -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 @@ -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 { @@ -134,10 +141,10 @@ 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") @@ -145,7 +152,7 @@ func LargeListOf(t DataType) *LargeListType { 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 { @@ -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()}} } @@ -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()}} } @@ -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 @@ -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 @@ -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) { @@ -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) )