create better interface

Author: yevgenypats

go/arrow/array/builder.go

@@ -318,11 +318,11 @@ func NewBuilder(mem memory.Allocator, dtype arrow.DataType) Builder {
 		return NewMapBuilderWithType(mem, typ)
 	case arrow.EXTENSION:
 		typ := dtype.(arrow.ExtensionType)
-		bldr := typ.NewBuilder(mem, typ)
-		if bldr != nil {
-			return bldr.(Builder)
+		bldr := NewExtensionBuilder(mem, typ)
+		if custom, ok := typ.(ExtensionBuilderWrapper); ok {
+			return custom.NewBuilder(bldr)
 		}
-		return NewExtensionBuilder(mem, typ)
+		return bldr
 	case arrow.FIXED_SIZE_LIST:
 		typ := dtype.(*arrow.FixedSizeListType)
 		return NewFixedSizeListBuilder(mem, typ.Len(), typ.Elem())

go/arrow/array/extension_builder.go

@@ -0,0 +1,5 @@
+package array
+
+type ExtensionBuilderWrapper interface {
+	NewBuilder(bldr *ExtensionBuilder) Builder
+}

go/arrow/datatype_extension.go

@@ -20,8 +20,6 @@ import (
 	"fmt"
 	"reflect"
 	"sync"
-
-	"github.com/apache/arrow/go/v12/arrow/memory"
 )
 
 var (
@@ -124,20 +122,16 @@ type ExtensionType interface {
 	// If the storage type is incorrect or something else is invalid with the data this should
 	// return nil and an appropriate error.
 	Deserialize(storageType DataType, data string) (ExtensionType, error)
-	// this should return array.Builder interface but we cannot import due to cycle import, so we use 
-	// interface{} instead. At least for 
-	NewBuilder(mem memory.Allocator, dt ExtensionType) interface{}
 	mustEmbedExtensionBase()
 }
 
 // ExtensionBase is the base struct for user-defined Extension Types which must be
 // embedded in any user-defined types like so:
 //
-//     type UserDefinedType struct {
-//         arrow.ExtensionBase
-//         // any other data
-//     }
-//
+//	type UserDefinedType struct {
+//	    arrow.ExtensionBase
+//	    // any other data
+//	}
 type ExtensionBase struct {
 	// Storage is the underlying storage type
 	Storage DataType
@@ -166,13 +160,10 @@ func (e *ExtensionBase) Fields() []Field {
 	return nil
 }
 
-func (e *ExtensionBase) NewBuilder(mem memory.Allocator, dt ExtensionType) interface{} {
-	return nil
-}
-
 func (e *ExtensionBase) Layout() DataTypeLayout { return e.Storage.Layout() }
 
 // this no-op exists to ensure that this type must be embedded in any user-defined extension type.
+//
 //lint:ignore U1000 this function is intentionally unused as it only exists to ensure embedding happens
 func (ExtensionBase) mustEmbedExtensionBase() {}
 

go/arrow/internal/testing/types/extension_test.go

@@ -33,15 +33,16 @@ import (
 var testUUID = uuid.New()
 
 func TestExtensionBuilder(t *testing.T) {
-	builder := types.NewUUIDBuilder(memory.DefaultAllocator, types.NewUUIDType())
+	extBuilder := array.NewExtensionBuilder(memory.DefaultAllocator, types.NewUUIDType())
+	builder := types.NewUUIDBuilder(extBuilder)
 	builder.Append(testUUID)
 	arr := builder.NewArray()
 	arrStr := arr.String()
 	assert.Equal(t, "[\""+testUUID.String()+"\"]", arrStr)
 	jsonStr, err := json.Marshal(arr)
 	assert.NoError(t, err)
 
-	arr1, _ , err := array.FromJSON(memory.DefaultAllocator, types.NewUUIDType(), bytes.NewReader(jsonStr))	
+	arr1, _, err := array.FromJSON(memory.DefaultAllocator, types.NewUUIDType(), bytes.NewReader(jsonStr))
 	assert.NoError(t, err)
 	assert.Equal(t, arr, arr1)
 }

go/arrow/internal/testing/types/extension_types.go

@@ -28,7 +28,6 @@ import (
 
 	"github.com/apache/arrow/go/v12/arrow"
 	"github.com/apache/arrow/go/v12/arrow/array"
-	"github.com/apache/arrow/go/v12/arrow/memory"
 	"github.com/google/uuid"
 	"golang.org/x/xerrors"
 )
@@ -38,10 +37,9 @@ type UUIDBuilder struct {
 	dtype *UUIDType
 }
 
-func NewUUIDBuilder(mem memory.Allocator, dtype arrow.ExtensionType) *UUIDBuilder {
+func NewUUIDBuilder(bldr *array.ExtensionBuilder) *UUIDBuilder {
 	b := &UUIDBuilder{
-		ExtensionBuilder: array.NewExtensionBuilder(mem, dtype),
-		dtype:            dtype.(*UUIDType),
+		ExtensionBuilder: bldr,
 	}
 	return b
 }
@@ -192,7 +190,6 @@ type UUIDType struct {
 	arrow.ExtensionBase
 }
 
-
 // NewUUIDType is a convenience function to create an instance of UuidType
 // with the correct storage type
 func NewUUIDType() *UUIDType {
@@ -226,11 +223,10 @@ func (u UUIDType) ExtensionEquals(other arrow.ExtensionType) bool {
 	return u.ExtensionName() == other.ExtensionName()
 }
 
-func (u UUIDType) NewBuilder(mem memory.Allocator, dt arrow.ExtensionType) interface{} {
-	return NewUUIDBuilder(mem, dt)
+func (u UUIDType) NewBuilder(bldr *array.ExtensionBuilder) array.Builder {
+	return NewUUIDBuilder(bldr)
 }
 
-
 // Parametric1Array is a simple int32 array for use with the Parametric1Type
 // in testing a parameterized user-defined extension type.
 type Parametric1Array struct {