loicalleyne commented on code in PR #37115: URL: https://github.com/apache/arrow/pull/37115#discussion_r1349492122
########## 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) Review Comment: This type assertion appears counter-intuitive but is necessary because the for loop is traversing the map by the keys in f.namePath by reassigning value to valueMap and looking for the key, if the type assertion fails before the end of the path then the data isn't there (that part of the tree doesn't exist because of a null value). I've added more details to the function comments. -- 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]
