loicalleyne commented on code in PR #37115: URL: https://github.com/apache/arrow/pull/37115#discussion_r1349378306
########## go/arrow/avro/reader_types.go: ########## @@ -0,0 +1,892 @@ +// 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 avro + +import ( + "bytes" + "encoding/binary" + "fmt" + "math/big" + + "github.com/apache/arrow/go/v14/internal/types" + + "github.com/apache/arrow/go/v14/arrow" + "github.com/apache/arrow/go/v14/arrow/array" + "github.com/apache/arrow/go/v14/arrow/decimal128" + "github.com/apache/arrow/go/v14/arrow/decimal256" + "github.com/apache/arrow/go/v14/arrow/memory" +) + +type dataLoader struct { + idx, depth int32 + list *fieldPos + item *fieldPos + mapField *fieldPos + mapKey *fieldPos + mapValue *fieldPos + fields []*fieldPos + children []*dataLoader +} + +func newDataLoader() *dataLoader { return &dataLoader{idx: 0, depth: 0} } + +// drawTree takes the tree of field builders produced by mapFieldBuilders() +// and produces another tree structure and aggregates fields whose values can +// be retrieved from a `map[string]any` into a slice of builders, and creates a hierarchy to +// deal with nested types (lists and maps). +func (d *dataLoader) drawTree(field *fieldPos) { + for _, f := range field.children() { + if f.isList || f.isMap { + if f.isList { + c := d.newListChild(f) + if !f.childrens[0].isList { + c.item = f.childrens[0] + c.drawTree(f.childrens[0]) + } else { + c.drawTree(f.childrens[0].childrens[0]) + } + } + if f.isMap { + c := d.newMapChild(f) + if !arrow.IsNested(f.childrens[1].builder.Type().ID()) { + c.mapKey = f.childrens[0] + c.mapValue = f.childrens[1] + } else { + c.mapKey = f.childrens[0] + m := c.newChild() + m.mapValue = f.childrens[1] + m.drawTree(f.childrens[1]) + } + } + } else { + d.fields = append(d.fields, f) + if len(f.children()) > 0 { + d.drawTree(f) + } + } + } +} + +func (d *dataLoader) loadDatum(data any) error { + if d.list == nil && d.mapField == nil { + if d.mapValue != nil { + d.mapValue.appendFunc(data) + } + for _, f := range d.fields { + if d.mapValue == nil { + err := f.appendFunc(f.getValue(data.(map[string]any))) + if err != nil { + return err + } + } else { + switch dt := data.(type) { + case nil: + err := f.appendFunc(dt) + if err != nil { + return err + } + case []any: + if len(d.children) < 1 { + for _, e := range dt { + err := f.appendFunc(e) + if err != nil { + return err + } + } + } else { + for _, e := range dt { + d.children[0].loadDatum(e) + } + } + case map[string]any: + err := f.appendFunc(f.getValue(dt)) + if err != nil { + return err + } + } + } + } + for _, c := range d.children { + if c.list != nil { + c.loadDatum(c.list.getValue(data)) + } + if c.mapField != nil { + switch dt := data.(type) { + case nil: + c.loadDatum(dt) + case map[string]any: + c.loadDatum(c.mapField.getValue(dt)) + default: + c.loadDatum(c.mapField.getValue(data).(map[string]any)) + } + } + } + } else { + if d.list != nil { + switch dt := data.(type) { + case nil: + d.list.appendFunc(dt) + case []any: + d.list.appendFunc(dt) + for _, e := range dt { + if d.item != nil { + d.item.appendFunc(e) + } + for _, f := range d.fields { + err := f.appendFunc(f.getValue(e)) + if err != nil { + return err + } + } + for _, c := range d.children { + if c.list != nil { + c.loadDatum(c.list.getValue(e)) + } + if c.mapField != nil { + c.loadDatum(c.mapField.getValue(e).(map[string]any)) + } + } + } + case map[string]any: + d.list.appendFunc(dt["array"]) + for _, e := range dt["array"].([]any) { + if d.item != nil { + d.item.appendFunc(e) + } + for _, f := range d.fields { + err := f.appendFunc(f.getValue(e)) + if err != nil { + return err + } + } + for _, c := range d.children { + c.loadDatum(c.list.getValue(e)) + } + } + default: + d.list.appendFunc(data) + d.item.appendFunc(dt) + } + } + if d.mapField != nil { + switch dt := data.(type) { + case nil: + d.mapField.appendFunc(dt) + case map[string]any: + + d.mapField.appendFunc(dt) + for k, v := range dt { + d.mapKey.appendFunc(k) + if d.mapValue != nil { + d.mapValue.appendFunc(v) + } else { + d.children[0].loadDatum(v) + } + } + + } + } + } + return nil +} + +func (d *dataLoader) newChild() *dataLoader { + var child *dataLoader = &dataLoader{ + depth: d.depth + 1, + } + d.children = append(d.children, child) + return child +} + +func (d *dataLoader) newListChild(list *fieldPos) *dataLoader { + var child *dataLoader = &dataLoader{ + list: list, + item: list.childrens[0], + depth: d.depth + 1, + } + d.children = append(d.children, child) + return child +} + +func (d *dataLoader) newMapChild(mapField *fieldPos) *dataLoader { + var child *dataLoader = &dataLoader{ + mapField: mapField, + depth: d.depth + 1, + } + d.children = append(d.children, child) + return child +} + +type fieldPos struct { + parent *fieldPos + fieldName string + builder array.Builder + path []string + isList bool + isItem bool + isStruct bool + isMap bool + typeName string + appendFunc func(val interface{}) error + metadatas arrow.Metadata + childrens []*fieldPos + index, depth int32 +} + +func newFieldPos() *fieldPos { return &fieldPos{index: -1} } + +func (f *fieldPos) name() string { return f.fieldName } + +func (f *fieldPos) child(index int) (*fieldPos, error) { + if index < len(f.children()) { + return f.childrens[index], nil + } + return nil, fmt.Errorf("%v child index %d not found", f.namePath(), index) +} + +func (f *fieldPos) children() []*fieldPos { return f.childrens } + +func (f *fieldPos) metadata() arrow.Metadata { return f.metadatas } + +func (f *fieldPos) newChild(childName string, childBuilder array.Builder, meta arrow.Metadata) *fieldPos { + var child fieldPos = fieldPos{ + parent: f, + fieldName: childName, + builder: childBuilder, + metadatas: meta, + index: int32(len(f.childrens)), + depth: f.depth + 1, + } + if f.isList { + child.isItem = true + } + child.path = child.namePath() + f.childrens = append(f.childrens, &child) + return &child +} + +// NamePath returns a slice of keys making up the path to the field +func (f *fieldPos) namePath() []string { + if len(f.path) == 0 { + var path []string + var listPath []string + cur := f + for i := f.depth - 1; i >= 0; i-- { + if cur.typeName == "" { + path = append([]string{cur.fieldName}, path...) + } else { + path = append([]string{cur.fieldName, cur.typeName}, path...) + } + if !cur.parent.isMap { + cur = cur.parent + } + } + if f.parent.parent != nil && f.parent.parent.isList { + for i := len(path) - 1; i >= 0; i-- { + if path[i] != "item" { + listPath = append([]string{path[i]}, listPath...) + } else { + return listPath + } + } + } + if f.parent != nil && f.parent.fieldName == "value" { + for i := len(path) - 1; i >= 0; i-- { + if path[i] != "value" { + listPath = append([]string{path[i]}, listPath...) + } else { + return listPath + } + } + } + return path + } + return f.path +} + +// GetValue retrieves the value from the map[string]any +// by following the field's key path +func (f *fieldPos) getValue(m any) any { + var value any = m + // The value of the type assertion is not assigned to a variable as `value` is reassigned + // and its type reasserted as the data is traversed + switch value.(type) { + case map[string]any: + for _, key := range f.namePath() { + valueMap, ok := value.(map[string]any) + if !ok { + if key == "item" { + return value + } + return nil + } + value, ok = valueMap[key] + if !ok { + return nil + } + } + default: + return value + } + return value +} + +// Avro data is loaded to Arrow arrays using the following type mapping: +// +// Avro Go Arrow +// null nil Null +// boolean bool Boolean +// bytes []byte Binary +// float float32 Float32 +// double float64 Float64 +// long int64 Int64 +// int int32 Int32 +// string string String +// array []interface{} List +// enum string Dictionary +// fixed []byte FixedSizeBinary +// map and record map[string]any Struct +// +// mapFieldBuilders builds a tree of field builders matching the Arrow schema +func mapFieldBuilders(b array.Builder, field arrow.Field, parent *fieldPos) { + switch bt := b.(type) { + case *array.BinaryBuilder: + f := parent.newChild(field.Name, bt, field.Metadata) + f.appendFunc = func(data interface{}) error { + appendBinaryData(bt, data) + return nil + } + case *array.BinaryDictionaryBuilder: + // has metadata for Avro enum symbols + f := parent.newChild(field.Name, b, field.Metadata) + f.appendFunc = func(data interface{}) error { + appendBinaryDictData(bt, data) + return nil + } + // add Avro enum symbols to builder + sb := array.NewStringBuilder(memory.DefaultAllocator) + for _, v := range field.Metadata.Values() { + sb.Append(v) + } + sa := sb.NewStringArray() + bt.InsertStringDictValues(sa) + case *array.BooleanBuilder: + f := parent.newChild(field.Name, b, field.Metadata) + f.appendFunc = func(data interface{}) error { + appendBoolData(bt, data) + return nil + } + case *array.Date32Builder: + f := parent.newChild(field.Name, b, field.Metadata) + f.appendFunc = func(data interface{}) error { + appendDate32Data(bt, data) + return nil + } + case *array.Decimal128Builder: + f := parent.newChild(field.Name, b, field.Metadata) + f.appendFunc = func(data interface{}) error { + err := appendDecimal128Data(bt, data) + if err != nil { + return err + } + return nil + } + case *array.Decimal256Builder: + f := parent.newChild(field.Name, b, field.Metadata) + f.appendFunc = func(data interface{}) error { + err := appendDecimal256Data(bt, data) + if err != nil { + return err + } + return nil + } + case *types.UUIDBuilder: + f := parent.newChild(field.Name, b, field.Metadata) + f.appendFunc = func(data interface{}) error { + switch dt := data.(type) { + case nil: + bt.AppendNull() + case string: + err := bt.AppendValueFromString(dt) + if err != nil { + return err + } + case []byte: + err := bt.AppendValueFromString(string(dt)) + if err != nil { + return err + } + } + + return nil + } + case *array.FixedSizeBinaryBuilder: + f := parent.newChild(field.Name, b, field.Metadata) + f.appendFunc = func(data interface{}) error { + appendFixedSizeBinaryData(bt, data) + return nil + } + case *array.Float32Builder: + f := parent.newChild(field.Name, b, field.Metadata) + f.appendFunc = func(data interface{}) error { + appendFloat32Data(bt, data) + return nil + } + case *array.Float64Builder: + f := parent.newChild(field.Name, b, field.Metadata) + f.appendFunc = func(data interface{}) error { + appendFloat64Data(bt, data) + return nil + } + case *array.Int32Builder: + f := parent.newChild(field.Name, b, field.Metadata) + f.appendFunc = func(data interface{}) error { + appendInt32Data(bt, data) + return nil + } + case *array.Int64Builder: + f := parent.newChild(field.Name, b, field.Metadata) + f.appendFunc = func(data interface{}) error { + appendInt64Data(bt, data) + return nil + } + case *array.LargeListBuilder: + lf := parent.newChild(field.Name, b, field.Metadata) + vb := bt.ValueBuilder() + lf.isList = true + mapFieldBuilders(vb, field.Type.(*arrow.LargeListType).ElemField(), lf) + lf.appendFunc = func(data interface{}) error { + switch dt := data.(type) { + case nil: + bt.AppendNull() + case []interface{}: + if len(dt) == 0 { + bt.AppendEmptyValue() + } else { + bt.Append(true) + } + default: + bt.Append(true) + } + return nil + } + case *array.ListBuilder: + lf := parent.newChild(field.Name, b, field.Metadata) + vb := bt.ValueBuilder() + lf.isList = true + mapFieldBuilders(vb, field.Type.(*arrow.ListType).ElemField(), lf) + lf.appendFunc = func(data interface{}) error { + switch dt := data.(type) { + case nil: + bt.AppendNull() + case []interface{}: + if len(dt) == 0 { + bt.AppendEmptyValue() + } else { + bt.Append(true) + } + default: + bt.Append(true) + } + return nil + } + case *array.MapBuilder: + // has metadata for objects in values + mf := parent.newChild(field.Name, b, field.Metadata) + mf.isMap = true + kb := bt.KeyBuilder() + ib := bt.ItemBuilder() + mapFieldBuilders(kb, field.Type.(*arrow.MapType).KeyField(), mf) + mapFieldBuilders(ib, field.Type.(*arrow.MapType).ItemField(), mf) + mf.appendFunc = func(data interface{}) error { + switch data.(type) { + case nil: + bt.AppendNull() + default: + bt.Append(true) + } + return nil + } + case *array.MonthDayNanoIntervalBuilder: + f := parent.newChild(field.Name, b, field.Metadata) + f.appendFunc = func(data interface{}) error { + appendDurationData(bt, data) + return nil + } + case *array.StringBuilder: + f := parent.newChild(field.Name, b, field.Metadata) + f.appendFunc = func(data interface{}) error { + appendStringData(bt, data) + return nil + } + case *array.StructBuilder: + // has metadata for Avro Union named types + sf := parent.newChild(field.Name, b, field.Metadata) + sf.typeName, _ = field.Metadata.GetValue("typeName") + sf.isStruct = true + // create children + for i, f := range field.Type.(*arrow.StructType).Fields() { + mapFieldBuilders(bt.FieldBuilder(i), f, sf) + } + sf.appendFunc = func(data interface{}) error { + switch data.(type) { + case nil: + bt.AppendNull() + default: + appendStructData(bt, data) + } + return nil + } + case *array.Time32Builder: + f := parent.newChild(field.Name, b, field.Metadata) + f.appendFunc = func(data interface{}) error { + appendTime32Data(bt, data) + return nil + } + case *array.Time64Builder: + f := parent.newChild(field.Name, b, field.Metadata) + f.appendFunc = func(data interface{}) error { + appendTime64Data(bt, data) + return nil + } + case *array.TimestampBuilder: + f := parent.newChild(field.Name, b, field.Metadata) + f.appendFunc = func(data interface{}) error { + appendTimestampData(bt, data) + return nil + } + } +} + +func appendBinaryData(b *array.BinaryBuilder, data interface{}) { + switch dt := data.(type) { + case nil: + b.AppendNull() + case map[string]any: + switch ct := dt["bytes"].(type) { + case nil: + b.AppendNull() + default: + b.Append(ct.([]byte)) + } + default: + b.Append([]byte(fmt.Sprint(data))) + } +} + +func appendBinaryDictData(b *array.BinaryDictionaryBuilder, data interface{}) { + switch dt := data.(type) { + case nil: + b.AppendNull() + case string: + b.AppendString(dt) + case map[string]any: + switch dt["string"].(type) { + case nil: + b.AppendNull() + case string: + b.AppendString(dt["string"].(string)) + } + } +} + +func appendBoolData(b *array.BooleanBuilder, data interface{}) { + switch dt := data.(type) { + case nil: + b.AppendNull() + case bool: + b.Append(dt) + case map[string]any: + switch dt["boolean"].(type) { + case nil: + b.AppendNull() + case bool: + b.Append(dt["boolean"].(bool)) + } + } +} + +func appendDate32Data(b *array.Date32Builder, data interface{}) { + switch dt := data.(type) { + case nil: + b.AppendNull() + case int32: + b.Append(arrow.Date32(dt)) + case map[string]any: + switch dt["int"].(type) { + case nil: + b.AppendNull() + case int32: + b.Append(arrow.Date32(dt["int"].(int32))) + } + } +} + +func appendDecimal128Data(b *array.Decimal128Builder, data interface{}) error { + switch dt := data.(type) { + case nil: + b.AppendNull() + case []byte: + if len(dt) <= 38 { + var intData int64 + buf := bytes.NewBuffer(dt) + err := binary.Read(buf, binary.BigEndian, &intData) + if err != nil { + return err + } + b.Append(decimal128.FromI64(intData)) + } else { + var bigIntData big.Int + buf := bytes.NewBuffer(dt) + err := binary.Read(buf, binary.BigEndian, &bigIntData) Review Comment: TBH I didn't have any test files that contain that type so I don't think that code path was ever run; you're right, it should be ```err := binary.Read(buf, binary.BigEndian, &bigIntData.Bytes())``` or something like that. I'll do a bit of research. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
