Skip to content

Commit 5219de3

Browse files
yevgenypatsyevgenypats
authored
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 <[email protected]> Signed-off-by: Matt Topol <[email protected]>
1 parent 71f3c56 commit 5219de3

35 files changed

+472
-234
lines changed

‎go/arrow/array/array.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ import (
2727
type arraymarshal interface {
2828
arrow.Array
2929

30-
getOneForMarshal(i int) interface{}
30+
GetOneForMarshal(i int) interface{}
3131
}
3232

3333
const (

‎go/arrow/array/binary.go

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -141,7 +141,7 @@ func (a *Binary) setData(data *Data) {
141141
}
142142
}
143143

144-
func (a *Binary) getOneForMarshal(i int) interface{} {
144+
func (a *Binary) GetOneForMarshal(i int) interface{} {
145145
if a.IsNull(i) {
146146
return nil
147147
}
@@ -151,7 +151,7 @@ func (a *Binary) getOneForMarshal(i int) interface{} {
151151
func (a *Binary) MarshalJSON() ([]byte, error) {
152152
vals := make([]interface{}, a.Len())
153153
for i := 0; i < a.Len(); i++ {
154-
vals[i] = a.getOneForMarshal(i)
154+
vals[i] = a.GetOneForMarshal(i)
155155
}
156156
// golang marshal standard says that []byte will be marshalled
157157
// as a base64-encoded string
@@ -274,7 +274,7 @@ func (a *LargeBinary) setData(data *Data) {
274274
}
275275
}
276276

277-
func (a *LargeBinary) getOneForMarshal(i int) interface{} {
277+
func (a *LargeBinary) GetOneForMarshal(i int) interface{} {
278278
if a.IsNull(i) {
279279
return nil
280280
}
@@ -284,7 +284,7 @@ func (a *LargeBinary) getOneForMarshal(i int) interface{} {
284284
func (a *LargeBinary) MarshalJSON() ([]byte, error) {
285285
vals := make([]interface{}, a.Len())
286286
for i := 0; i < a.Len(); i++ {
287-
vals[i] = a.getOneForMarshal(i)
287+
vals[i] = a.GetOneForMarshal(i)
288288
}
289289
// golang marshal standard says that []byte will be marshalled
290290
// as a base64-encoded string

‎go/arrow/array/binarybuilder.go

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -289,7 +289,7 @@ func (b *BinaryBuilder) appendNextOffset() {
289289
b.appendOffsetVal(numBytes)
290290
}
291291

292-
func (b *BinaryBuilder) unmarshalOne(dec *json.Decoder) error {
292+
func (b *BinaryBuilder) UnmarshalOne(dec *json.Decoder) error {
293293
t, err := dec.Token()
294294
if err != nil {
295295
return err
@@ -316,9 +316,9 @@ func (b *BinaryBuilder) unmarshalOne(dec *json.Decoder) error {
316316
return nil
317317
}
318318

319-
func (b *BinaryBuilder) unmarshal(dec *json.Decoder) error {
319+
func (b *BinaryBuilder) Unmarshal(dec *json.Decoder) error {
320320
for dec.More() {
321-
if err := b.unmarshalOne(dec); err != nil {
321+
if err := b.UnmarshalOne(dec); err != nil {
322322
return err
323323
}
324324
}
@@ -336,7 +336,7 @@ func (b *BinaryBuilder) UnmarshalJSON(data []byte) error {
336336
return fmt.Errorf("binary builder must unpack from json array, found %s", delim)
337337
}
338338

339-
return b.unmarshal(dec)
339+
return b.Unmarshal(dec)
340340
}
341341

342342
var (

‎go/arrow/array/boolean.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,7 @@ func (a *Boolean) setData(data *Data) {
8181
}
8282
}
8383

84-
func (a *Boolean) getOneForMarshal(i int) interface{} {
84+
func (a *Boolean) GetOneForMarshal(i int) interface{} {
8585
if a.IsValid(i) {
8686
return a.Value(i)
8787
}

‎go/arrow/array/booleanbuilder.go

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -172,7 +172,7 @@ func (b *BooleanBuilder) newData() *Data {
172172
return res
173173
}
174174

175-
func (b *BooleanBuilder) unmarshalOne(dec *json.Decoder) error {
175+
func (b *BooleanBuilder) UnmarshalOne(dec *json.Decoder) error {
176176
t, err := dec.Token()
177177
if err != nil {
178178
return err
@@ -205,9 +205,9 @@ func (b *BooleanBuilder) unmarshalOne(dec *json.Decoder) error {
205205
return nil
206206
}
207207

208-
func (b *BooleanBuilder) unmarshal(dec *json.Decoder) error {
208+
func (b *BooleanBuilder) Unmarshal(dec *json.Decoder) error {
209209
for dec.More() {
210-
if err := b.unmarshalOne(dec); err != nil {
210+
if err := b.UnmarshalOne(dec); err != nil {
211211
return err
212212
}
213213
}
@@ -226,7 +226,7 @@ func (b *BooleanBuilder) UnmarshalJSON(data []byte) error {
226226
return fmt.Errorf("boolean builder must unpack from json array, found %s", delim)
227227
}
228228

229-
return b.unmarshal(dec)
229+
return b.Unmarshal(dec)
230230
}
231231

232232
var (

‎go/arrow/array/builder.go

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -79,8 +79,8 @@ type Builder interface {
7979
init(capacity int)
8080
resize(newBits int, init func(int))
8181

82-
unmarshalOne(*json.Decoder) error
83-
unmarshal(*json.Decoder) error
82+
UnmarshalOne(*json.Decoder) error
83+
Unmarshal(*json.Decoder) error
8484

8585
newData() *Data
8686
}
@@ -318,7 +318,11 @@ func NewBuilder(mem memory.Allocator, dtype arrow.DataType) Builder {
318318
return NewMapBuilderWithType(mem, typ)
319319
case arrow.EXTENSION:
320320
typ := dtype.(arrow.ExtensionType)
321-
return NewExtensionBuilder(mem, typ)
321+
bldr := NewExtensionBuilder(mem, typ)
322+
if custom, ok := typ.(ExtensionBuilderWrapper); ok {
323+
return custom.NewBuilder(bldr)
324+
}
325+
return bldr
322326
case arrow.FIXED_SIZE_LIST:
323327
typ := dtype.(*arrow.FixedSizeListType)
324328
return NewFixedSizeListBuilder(mem, typ.Len(), typ.Elem())

‎go/arrow/array/decimal128.go

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,7 @@ func (a *Decimal128) setData(data *Data) {
8080
}
8181
}
8282

83-
func (a *Decimal128) getOneForMarshal(i int) interface{} {
83+
func (a *Decimal128) GetOneForMarshal(i int) interface{} {
8484
if a.IsNull(i) {
8585
return nil
8686
}
@@ -94,7 +94,7 @@ func (a *Decimal128) getOneForMarshal(i int) interface{} {
9494
func (a *Decimal128) MarshalJSON() ([]byte, error) {
9595
vals := make([]interface{}, a.Len())
9696
for i := 0; i < a.Len(); i++ {
97-
vals[i] = a.getOneForMarshal(i)
97+
vals[i] = a.GetOneForMarshal(i)
9898
}
9999
return json.Marshal(vals)
100100
}
@@ -259,7 +259,7 @@ func (b *Decimal128Builder) newData() (data *Data) {
259259
return
260260
}
261261

262-
func (b *Decimal128Builder) unmarshalOne(dec *json.Decoder) error {
262+
func (b *Decimal128Builder) UnmarshalOne(dec *json.Decoder) error {
263263
t, err := dec.Token()
264264
if err != nil {
265265
return err
@@ -298,9 +298,9 @@ func (b *Decimal128Builder) unmarshalOne(dec *json.Decoder) error {
298298
return nil
299299
}
300300

301-
func (b *Decimal128Builder) unmarshal(dec *json.Decoder) error {
301+
func (b *Decimal128Builder) Unmarshal(dec *json.Decoder) error {
302302
for dec.More() {
303-
if err := b.unmarshalOne(dec); err != nil {
303+
if err := b.UnmarshalOne(dec); err != nil {
304304
return err
305305
}
306306
}
@@ -322,7 +322,7 @@ func (b *Decimal128Builder) UnmarshalJSON(data []byte) error {
322322
return fmt.Errorf("decimal128 builder must unpack from json array, found %s", delim)
323323
}
324324

325-
return b.unmarshal(dec)
325+
return b.Unmarshal(dec)
326326
}
327327

328328
var (

‎go/arrow/array/decimal256.go

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,7 @@ func (a *Decimal256) setData(data *Data) {
8080
}
8181
}
8282

83-
func (a *Decimal256) getOneForMarshal(i int) interface{} {
83+
func (a *Decimal256) GetOneForMarshal(i int) interface{} {
8484
if a.IsNull(i) {
8585
return nil
8686
}
@@ -94,7 +94,7 @@ func (a *Decimal256) getOneForMarshal(i int) interface{} {
9494
func (a *Decimal256) MarshalJSON() ([]byte, error) {
9595
vals := make([]interface{}, a.Len())
9696
for i := 0; i < a.Len(); i++ {
97-
vals[i] = a.getOneForMarshal(i)
97+
vals[i] = a.GetOneForMarshal(i)
9898
}
9999
return json.Marshal(vals)
100100
}
@@ -259,7 +259,7 @@ func (b *Decimal256Builder) newData() (data *Data) {
259259
return
260260
}
261261

262-
func (b *Decimal256Builder) unmarshalOne(dec *json.Decoder) error {
262+
func (b *Decimal256Builder) UnmarshalOne(dec *json.Decoder) error {
263263
t, err := dec.Token()
264264
if err != nil {
265265
return err
@@ -298,9 +298,9 @@ func (b *Decimal256Builder) unmarshalOne(dec *json.Decoder) error {
298298
return nil
299299
}
300300

301-
func (b *Decimal256Builder) unmarshal(dec *json.Decoder) error {
301+
func (b *Decimal256Builder) Unmarshal(dec *json.Decoder) error {
302302
for dec.More() {
303-
if err := b.unmarshalOne(dec); err != nil {
303+
if err := b.UnmarshalOne(dec); err != nil {
304304
return err
305305
}
306306
}
@@ -322,7 +322,7 @@ func (b *Decimal256Builder) UnmarshalJSON(data []byte) error {
322322
return fmt.Errorf("arrow/array: decimal256 builder must unpack from json array, found %s", delim)
323323
}
324324

325-
return b.unmarshal(dec)
325+
return b.Unmarshal(dec)
326326
}
327327

328328
var (

‎go/arrow/array/dictionary.go

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -281,18 +281,18 @@ func (d *Dictionary) GetValueIndex(i int) int {
281281
return -1
282282
}
283283

284-
func (d *Dictionary) getOneForMarshal(i int) interface{} {
284+
func (d *Dictionary) GetOneForMarshal(i int) interface{} {
285285
if d.IsNull(i) {
286286
return nil
287287
}
288288
vidx := d.GetValueIndex(i)
289-
return d.Dictionary().(arraymarshal).getOneForMarshal(vidx)
289+
return d.Dictionary().(arraymarshal).GetOneForMarshal(vidx)
290290
}
291291

292292
func (d *Dictionary) MarshalJSON() ([]byte, error) {
293293
vals := make([]interface{}, d.Len())
294294
for i := 0; i < d.Len(); i++ {
295-
vals[i] = d.getOneForMarshal(i)
295+
vals[i] = d.GetOneForMarshal(i)
296296
}
297297
return json.Marshal(vals)
298298
}
@@ -721,14 +721,14 @@ func (b *dictionaryBuilder) UnmarshalJSON(data []byte) error {
721721
return fmt.Errorf("dictionary builder must upack from json array, found %s", delim)
722722
}
723723

724-
return b.unmarshal(dec)
724+
return b.Unmarshal(dec)
725725
}
726726

727-
func (b *dictionaryBuilder) unmarshal(dec *json.Decoder) error {
727+
func (b *dictionaryBuilder) Unmarshal(dec *json.Decoder) error {
728728
bldr := NewBuilder(b.mem, b.dt.ValueType)
729729
defer bldr.Release()
730730

731-
if err := bldr.unmarshal(dec); err != nil {
731+
if err := bldr.Unmarshal(dec); err != nil {
732732
return err
733733
}
734734

@@ -737,7 +737,7 @@ func (b *dictionaryBuilder) unmarshal(dec *json.Decoder) error {
737737
return b.AppendArray(arr)
738738
}
739739

740-
func (b *dictionaryBuilder) unmarshalOne(dec *json.Decoder) error {
740+
func (b *dictionaryBuilder) UnmarshalOne(dec *json.Decoder) error {
741741
return errors.New("unmarshal json to dictionary not yet implemented")
742742
}
743743

‎go/arrow/array/encoded.go

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -200,22 +200,22 @@ func (r *RunEndEncoded) String() string {
200200
buf.WriteByte(',')
201201
}
202202

203-
value := r.values.(arraymarshal).getOneForMarshal(i)
203+
value := r.values.(arraymarshal).GetOneForMarshal(i)
204204
if byts, ok := value.(json.RawMessage); ok {
205205
value = string(byts)
206206
}
207207
fmt.Fprintf(&buf, "{%d -> %v}",
208-
r.ends.(arraymarshal).getOneForMarshal(i),
208+
r.ends.(arraymarshal).GetOneForMarshal(i),
209209
value)
210210
}
211211

212212
buf.WriteByte(']')
213213
return buf.String()
214214
}
215215

216-
func (r *RunEndEncoded) getOneForMarshal(i int) interface{} {
216+
func (r *RunEndEncoded) GetOneForMarshal(i int) interface{} {
217217
physIndex := encoded.FindPhysicalIndex(r.data, i+r.data.offset)
218-
return r.values.(arraymarshal).getOneForMarshal(physIndex)
218+
return r.values.(arraymarshal).GetOneForMarshal(physIndex)
219219
}
220220

221221
func (r *RunEndEncoded) MarshalJSON() ([]byte, error) {
@@ -226,7 +226,7 @@ func (r *RunEndEncoded) MarshalJSON() ([]byte, error) {
226226
if i != 0 {
227227
buf.WriteByte(',')
228228
}
229-
if err := enc.Encode(r.getOneForMarshal(i)); err != nil {
229+
if err := enc.Encode(r.GetOneForMarshal(i)); err != nil {
230230
return nil, err
231231
}
232232
}
@@ -397,7 +397,7 @@ func (b *RunEndEncodedBuilder) newData() (data *Data) {
397397
return
398398
}
399399

400-
func (b *RunEndEncodedBuilder) unmarshalOne(dec *json.Decoder) error {
400+
func (b *RunEndEncodedBuilder) UnmarshalOne(dec *json.Decoder) error {
401401
var value interface{}
402402
if err := dec.Decode(&value); err != nil {
403403
return err
@@ -422,13 +422,13 @@ func (b *RunEndEncodedBuilder) unmarshalOne(dec *json.Decoder) error {
422422

423423
b.Append(1)
424424
b.lastUnmarshalled = value
425-
return b.ValueBuilder().unmarshalOne(json.NewDecoder(bytes.NewReader(data)))
425+
return b.ValueBuilder().UnmarshalOne(json.NewDecoder(bytes.NewReader(data)))
426426
}
427427

428-
func (b *RunEndEncodedBuilder) unmarshal(dec *json.Decoder) error {
428+
func (b *RunEndEncodedBuilder) Unmarshal(dec *json.Decoder) error {
429429
b.finishRun()
430430
for dec.More() {
431-
if err := b.unmarshalOne(dec); err != nil {
431+
if err := b.UnmarshalOne(dec); err != nil {
432432
return err
433433
}
434434
}
@@ -446,7 +446,7 @@ func (b *RunEndEncodedBuilder) UnmarshalJSON(data []byte) error {
446446
return fmt.Errorf("list builder must unpack from json array, found %s", delim)
447447
}
448448

449-
return b.unmarshal(dec)
449+
return b.Unmarshal(dec)
450450
}
451451

452452
var (

‎go/arrow/array/extension.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -134,8 +134,8 @@ func (e *ExtensionArrayBase) String() string {
134134
return fmt.Sprintf("(%s)%s", e.data.dtype, e.storage)
135135
}
136136

137-
func (e *ExtensionArrayBase) getOneForMarshal(i int) interface{} {
138-
return e.storage.getOneForMarshal(i)
137+
func (e *ExtensionArrayBase) GetOneForMarshal(i int) interface{} {
138+
return e.storage.GetOneForMarshal(i)
139139
}
140140

141141
func (e *ExtensionArrayBase) MarshalJSON() ([]byte, error) {

‎go/arrow/array/extension_builder.go

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

0 commit comments

Comments
 (0)