GH-34453: [Go] Support Builders for user defined extensions (#34454)

This should serve as discussion as it's a medium change but this should Close #34453 and give users the ability to define custom Builder for their extensions just like they define ExtensionTypes and ExtensionArrays
* Closes: #34453

Authored-by: Yevgeny Pats <16490766+yevgenypats@users.noreply.github.com>
Signed-off-by: Matt Topol <zotthewizard@gmail.com>

go/arrow/array/array.go

@@ -27,7 +27,7 @@ import (
 type arraymarshal interface {
 	arrow.Array
 
-	getOneForMarshal(i int) interface{}
+	GetOneForMarshal(i int) interface{}
 }
 
 const (

go/arrow/array/binary.go

@@ -141,7 +141,7 @@ func (a *Binary) setData(data *Data) {
 	}
 }
 
-func (a *Binary) getOneForMarshal(i int) interface{} {
+func (a *Binary) GetOneForMarshal(i int) interface{} {
 	if a.IsNull(i) {
 		return nil
 	}
@@ -151,7 +151,7 @@ func (a *Binary) getOneForMarshal(i int) interface{} {
 func (a *Binary) MarshalJSON() ([]byte, error) {
 	vals := make([]interface{}, a.Len())
 	for i := 0; i < a.Len(); i++ {
-		vals[i] = a.getOneForMarshal(i)
+		vals[i] = a.GetOneForMarshal(i)
 	}
 	// golang marshal standard says that []byte will be marshalled
 	// as a base64-encoded string
@@ -274,7 +274,7 @@ func (a *LargeBinary) setData(data *Data) {
 	}
 }
 
-func (a *LargeBinary) getOneForMarshal(i int) interface{} {
+func (a *LargeBinary) GetOneForMarshal(i int) interface{} {
 	if a.IsNull(i) {
 		return nil
 	}
@@ -284,7 +284,7 @@ func (a *LargeBinary) getOneForMarshal(i int) interface{} {
 func (a *LargeBinary) MarshalJSON() ([]byte, error) {
 	vals := make([]interface{}, a.Len())
 	for i := 0; i < a.Len(); i++ {
-		vals[i] = a.getOneForMarshal(i)
+		vals[i] = a.GetOneForMarshal(i)
 	}
 	// golang marshal standard says that []byte will be marshalled
 	// as a base64-encoded string

go/arrow/array/binarybuilder.go

@@ -289,7 +289,7 @@ func (b *BinaryBuilder) appendNextOffset() {
 	b.appendOffsetVal(numBytes)
 }
 
-func (b *BinaryBuilder) unmarshalOne(dec *json.Decoder) error {
+func (b *BinaryBuilder) UnmarshalOne(dec *json.Decoder) error {
 	t, err := dec.Token()
 	if err != nil {
 		return err
@@ -316,9 +316,9 @@ func (b *BinaryBuilder) unmarshalOne(dec *json.Decoder) error {
 	return nil
 }
 
-func (b *BinaryBuilder) unmarshal(dec *json.Decoder) error {
+func (b *BinaryBuilder) Unmarshal(dec *json.Decoder) error {
 	for dec.More() {
-		if err := b.unmarshalOne(dec); err != nil {
+		if err := b.UnmarshalOne(dec); err != nil {
 			return err
 		}
 	}
@@ -336,7 +336,7 @@ func (b *BinaryBuilder) UnmarshalJSON(data []byte) error {
 		return fmt.Errorf("binary builder must unpack from json array, found %s", delim)
 	}
 
-	return b.unmarshal(dec)
+	return b.Unmarshal(dec)
 }
 
 var (

go/arrow/array/boolean.go

@@ -81,7 +81,7 @@ func (a *Boolean) setData(data *Data) {
 	}
 }
 
-func (a *Boolean) getOneForMarshal(i int) interface{} {
+func (a *Boolean) GetOneForMarshal(i int) interface{} {
 	if a.IsValid(i) {
 		return a.Value(i)
 	}

go/arrow/array/booleanbuilder.go

@@ -172,7 +172,7 @@ func (b *BooleanBuilder) newData() *Data {
 	return res
 }
 
-func (b *BooleanBuilder) unmarshalOne(dec *json.Decoder) error {
+func (b *BooleanBuilder) UnmarshalOne(dec *json.Decoder) error {
 	t, err := dec.Token()
 	if err != nil {
 		return err
@@ -205,9 +205,9 @@ func (b *BooleanBuilder) unmarshalOne(dec *json.Decoder) error {
 	return nil
 }
 
-func (b *BooleanBuilder) unmarshal(dec *json.Decoder) error {
+func (b *BooleanBuilder) Unmarshal(dec *json.Decoder) error {
 	for dec.More() {
-		if err := b.unmarshalOne(dec); err != nil {
+		if err := b.UnmarshalOne(dec); err != nil {
 			return err
 		}
 	}
@@ -226,7 +226,7 @@ func (b *BooleanBuilder) UnmarshalJSON(data []byte) error {
 		return fmt.Errorf("boolean builder must unpack from json array, found %s", delim)
 	}
 
-	return b.unmarshal(dec)
+	return b.Unmarshal(dec)
 }
 
 var (

go/arrow/array/builder.go

@@ -79,8 +79,8 @@ type Builder interface {
 	init(capacity int)
 	resize(newBits int, init func(int))
 
-	unmarshalOne(*json.Decoder) error
-	unmarshal(*json.Decoder) error
+	UnmarshalOne(*json.Decoder) error
+	Unmarshal(*json.Decoder) error
 
 	newData() *Data
 }
@@ -318,7 +318,11 @@ func NewBuilder(mem memory.Allocator, dtype arrow.DataType) Builder {
 		return NewMapBuilderWithType(mem, typ)
 	case arrow.EXTENSION:
 		typ := dtype.(arrow.ExtensionType)
-		return NewExtensionBuilder(mem, typ)
+		bldr := NewExtensionBuilder(mem, typ)
+		if custom, ok := typ.(ExtensionBuilderWrapper); ok {
+			return custom.NewBuilder(bldr)
+		}
+		return bldr
 	case arrow.FIXED_SIZE_LIST:
 		typ := dtype.(*arrow.FixedSizeListType)
 		return NewFixedSizeListBuilder(mem, typ.Len(), typ.Elem())

go/arrow/array/decimal128.go

@@ -80,7 +80,7 @@ func (a *Decimal128) setData(data *Data) {
 	}
 }
 
-func (a *Decimal128) getOneForMarshal(i int) interface{} {
+func (a *Decimal128) GetOneForMarshal(i int) interface{} {
 	if a.IsNull(i) {
 		return nil
 	}
@@ -94,7 +94,7 @@ func (a *Decimal128) getOneForMarshal(i int) interface{} {
 func (a *Decimal128) MarshalJSON() ([]byte, error) {
 	vals := make([]interface{}, a.Len())
 	for i := 0; i < a.Len(); i++ {
-		vals[i] = a.getOneForMarshal(i)
+		vals[i] = a.GetOneForMarshal(i)
 	}
 	return json.Marshal(vals)
 }
@@ -259,7 +259,7 @@ func (b *Decimal128Builder) newData() (data *Data) {
 	return
 }
 
-func (b *Decimal128Builder) unmarshalOne(dec *json.Decoder) error {
+func (b *Decimal128Builder) UnmarshalOne(dec *json.Decoder) error {
 	t, err := dec.Token()
 	if err != nil {
 		return err
@@ -298,9 +298,9 @@ func (b *Decimal128Builder) unmarshalOne(dec *json.Decoder) error {
 	return nil
 }
 
-func (b *Decimal128Builder) unmarshal(dec *json.Decoder) error {
+func (b *Decimal128Builder) Unmarshal(dec *json.Decoder) error {
 	for dec.More() {
-		if err := b.unmarshalOne(dec); err != nil {
+		if err := b.UnmarshalOne(dec); err != nil {
 			return err
 		}
 	}
@@ -322,7 +322,7 @@ func (b *Decimal128Builder) UnmarshalJSON(data []byte) error {
 		return fmt.Errorf("decimal128 builder must unpack from json array, found %s", delim)
 	}
 
-	return b.unmarshal(dec)
+	return b.Unmarshal(dec)
 }
 
 var (

go/arrow/array/decimal256.go

@@ -80,7 +80,7 @@ func (a *Decimal256) setData(data *Data) {
 	}
 }
 
-func (a *Decimal256) getOneForMarshal(i int) interface{} {
+func (a *Decimal256) GetOneForMarshal(i int) interface{} {
 	if a.IsNull(i) {
 		return nil
 	}
@@ -94,7 +94,7 @@ func (a *Decimal256) getOneForMarshal(i int) interface{} {
 func (a *Decimal256) MarshalJSON() ([]byte, error) {
 	vals := make([]interface{}, a.Len())
 	for i := 0; i < a.Len(); i++ {
-		vals[i] = a.getOneForMarshal(i)
+		vals[i] = a.GetOneForMarshal(i)
 	}
 	return json.Marshal(vals)
 }
@@ -259,7 +259,7 @@ func (b *Decimal256Builder) newData() (data *Data) {
 	return
 }
 
-func (b *Decimal256Builder) unmarshalOne(dec *json.Decoder) error {
+func (b *Decimal256Builder) UnmarshalOne(dec *json.Decoder) error {
 	t, err := dec.Token()
 	if err != nil {
 		return err
@@ -298,9 +298,9 @@ func (b *Decimal256Builder) unmarshalOne(dec *json.Decoder) error {
 	return nil
 }
 
-func (b *Decimal256Builder) unmarshal(dec *json.Decoder) error {
+func (b *Decimal256Builder) Unmarshal(dec *json.Decoder) error {
 	for dec.More() {
-		if err := b.unmarshalOne(dec); err != nil {
+		if err := b.UnmarshalOne(dec); err != nil {
 			return err
 		}
 	}
@@ -322,7 +322,7 @@ func (b *Decimal256Builder) UnmarshalJSON(data []byte) error {
 		return fmt.Errorf("arrow/array: decimal256 builder must unpack from json array, found %s", delim)
 	}
 
-	return b.unmarshal(dec)
+	return b.Unmarshal(dec)
 }
 
 var (

go/arrow/array/dictionary.go

@@ -281,18 +281,18 @@ func (d *Dictionary) GetValueIndex(i int) int {
 	return -1
 }
 
-func (d *Dictionary) getOneForMarshal(i int) interface{} {
+func (d *Dictionary) GetOneForMarshal(i int) interface{} {
 	if d.IsNull(i) {
 		return nil
 	}
 	vidx := d.GetValueIndex(i)
-	return d.Dictionary().(arraymarshal).getOneForMarshal(vidx)
+	return d.Dictionary().(arraymarshal).GetOneForMarshal(vidx)
 }
 
 func (d *Dictionary) MarshalJSON() ([]byte, error) {
 	vals := make([]interface{}, d.Len())
 	for i := 0; i < d.Len(); i++ {
-		vals[i] = d.getOneForMarshal(i)
+		vals[i] = d.GetOneForMarshal(i)
 	}
 	return json.Marshal(vals)
 }
@@ -721,14 +721,14 @@ func (b *dictionaryBuilder) UnmarshalJSON(data []byte) error {
 		return fmt.Errorf("dictionary builder must upack from json array, found %s", delim)
 	}
 
-	return b.unmarshal(dec)
+	return b.Unmarshal(dec)
 }
 
-func (b *dictionaryBuilder) unmarshal(dec *json.Decoder) error {
+func (b *dictionaryBuilder) Unmarshal(dec *json.Decoder) error {
 	bldr := NewBuilder(b.mem, b.dt.ValueType)
 	defer bldr.Release()
 
-	if err := bldr.unmarshal(dec); err != nil {
+	if err := bldr.Unmarshal(dec); err != nil {
 		return err
 	}
 
@@ -737,7 +737,7 @@ func (b *dictionaryBuilder) unmarshal(dec *json.Decoder) error {
 	return b.AppendArray(arr)
 }
 
-func (b *dictionaryBuilder) unmarshalOne(dec *json.Decoder) error {
+func (b *dictionaryBuilder) UnmarshalOne(dec *json.Decoder) error {
 	return errors.New("unmarshal json to dictionary not yet implemented")
 }
 

go/arrow/array/encoded.go

@@ -200,22 +200,22 @@ func (r *RunEndEncoded) String() string {
 			buf.WriteByte(',')
 		}
 
-		value := r.values.(arraymarshal).getOneForMarshal(i)
+		value := r.values.(arraymarshal).GetOneForMarshal(i)
 		if byts, ok := value.(json.RawMessage); ok {
 			value = string(byts)
 		}
 		fmt.Fprintf(&buf, "{%d -> %v}",
-			r.ends.(arraymarshal).getOneForMarshal(i),
+			r.ends.(arraymarshal).GetOneForMarshal(i),
 			value)
 	}
 
 	buf.WriteByte(']')
 	return buf.String()
 }
 
-func (r *RunEndEncoded) getOneForMarshal(i int) interface{} {
+func (r *RunEndEncoded) GetOneForMarshal(i int) interface{} {
 	physIndex := encoded.FindPhysicalIndex(r.data, i+r.data.offset)
-	return r.values.(arraymarshal).getOneForMarshal(physIndex)
+	return r.values.(arraymarshal).GetOneForMarshal(physIndex)
 }
 
 func (r *RunEndEncoded) MarshalJSON() ([]byte, error) {
@@ -226,7 +226,7 @@ func (r *RunEndEncoded) MarshalJSON() ([]byte, error) {
 		if i != 0 {
 			buf.WriteByte(',')
 		}
-		if err := enc.Encode(r.getOneForMarshal(i)); err != nil {
+		if err := enc.Encode(r.GetOneForMarshal(i)); err != nil {
 			return nil, err
 		}
 	}
@@ -397,7 +397,7 @@ func (b *RunEndEncodedBuilder) newData() (data *Data) {
 	return
 }
 
-func (b *RunEndEncodedBuilder) unmarshalOne(dec *json.Decoder) error {
+func (b *RunEndEncodedBuilder) UnmarshalOne(dec *json.Decoder) error {
 	var value interface{}
 	if err := dec.Decode(&value); err != nil {
 		return err
@@ -422,13 +422,13 @@ func (b *RunEndEncodedBuilder) unmarshalOne(dec *json.Decoder) error {
 
 	b.Append(1)
 	b.lastUnmarshalled = value
-	return b.ValueBuilder().unmarshalOne(json.NewDecoder(bytes.NewReader(data)))
+	return b.ValueBuilder().UnmarshalOne(json.NewDecoder(bytes.NewReader(data)))
 }
 
-func (b *RunEndEncodedBuilder) unmarshal(dec *json.Decoder) error {
+func (b *RunEndEncodedBuilder) Unmarshal(dec *json.Decoder) error {
 	b.finishRun()
 	for dec.More() {
-		if err := b.unmarshalOne(dec); err != nil {
+		if err := b.UnmarshalOne(dec); err != nil {
 			return err
 		}
 	}
@@ -446,7 +446,7 @@ func (b *RunEndEncodedBuilder) UnmarshalJSON(data []byte) error {
 		return fmt.Errorf("list builder must unpack from json array, found %s", delim)
 	}
 
-	return b.unmarshal(dec)
+	return b.Unmarshal(dec)
 }
 
 var (

go/arrow/array/extension.go

@@ -134,8 +134,8 @@ func (e *ExtensionArrayBase) String() string {
 	return fmt.Sprintf("(%s)%s", e.data.dtype, e.storage)
 }
 
-func (e *ExtensionArrayBase) getOneForMarshal(i int) interface{} {
-	return e.storage.getOneForMarshal(i)
+func (e *ExtensionArrayBase) GetOneForMarshal(i int) interface{} {
+	return e.storage.GetOneForMarshal(i)
 }
 
 func (e *ExtensionArrayBase) MarshalJSON() ([]byte, error) {

go/arrow/array/extension_builder.go

@@ -0,0 +1,23 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package array
+
+// ExtensionBuilderWrapper is an interface that you need to implement in your custom extension type if you want to provide a customer builder as well.
+// See example in ./arrow/internal/testing/types/extension_types.go
+type ExtensionBuilderWrapper interface {
+	NewBuilder(bldr *ExtensionBuilder) Builder
+}