loicalleyne commented on code in PR #36796: URL: https://github.com/apache/arrow/pull/36796#discussion_r1316460234
########## go/arrow/avro/schema.go: ########## @@ -0,0 +1,407 @@ +// 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 reads Avro OCF files and presents the extracted data as records + +package avro + +import ( + "fmt" + "math" + "strconv" + + "github.com/apache/arrow/go/v14/arrow" + "github.com/apache/arrow/go/v14/internal/types" + avro "github.com/hamba/avro/v2" +) + +type schemaNode struct { + name string + parent *schemaNode + schema avro.Schema + union bool + children []*schemaNode + arrowField arrow.Field + path []string + index, depth int32 +} + +var schemaCache avro.SchemaCache + +func newSchemaNode() *schemaNode { return &schemaNode{name: "", index: -1} } + +func (node *schemaNode) newChild(n string, s avro.Schema) *schemaNode { + child := &schemaNode{ + name: n, + parent: node, + schema: s, + index: int32(len(node.children)), + depth: node.depth + 1} + node.children = append(node.children, child) + return child +} +func (node *schemaNode) Children() []*schemaNode { return node.children } + +func (node *schemaNode) Name() string { return node.name } + +// NamePath returns a slice of keys making up the path to the field +func (node *schemaNode) NamePath() []string { + if len(node.path) == 0 { + var path []string + cur := node + for i := node.depth - 1; i >= 0; i-- { + path = append([]string{cur.name}, path...) + cur = cur.parent + } + return path + } + return node.path +} + +// ArrowSchemaFromAvro returns a new Arrow schema from an Avro schema +func ArrowSchemaFromAvro(schema avro.Schema) (s *arrow.Schema, err error) { + defer func() { + if r := recover(); r != nil { + s = nil + err = fmt.Errorf("invalid avro schema: could not convert schema") + } + }() + n := newSchemaNode() + n.schema = schema + c := n.newChild(n.schema.(avro.NamedSchema).Name(), n.schema) + arrowSchemafromAvro(c) + var fields []arrow.Field + for _, g := range c.Children() { + fields = append(fields, g.arrowField) + } + s = arrow.NewSchema(fields, nil) + return s, nil +} + +func arrowSchemafromAvro(n *schemaNode) { + if ns, ok := n.schema.(avro.NamedSchema); ok { + schemaCache.Add(ns.Name(), ns) + } + switch st := n.schema.Type(); st { + case "record": + iterateFields(n) + case "enum": + symbols := make(map[string]string) + for index, symbol := range n.schema.(avro.PropertySchema).(*avro.EnumSchema).Symbols() { + k := strconv.FormatInt(int64(index), 10) + symbols[k] = symbol + } + var dt arrow.DictionaryType = arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Uint64, ValueType: arrow.BinaryTypes.String, Ordered: false} + sl := len(symbols) + switch { + case sl <= math.MaxUint8: + dt.IndexType = arrow.PrimitiveTypes.Uint8 + case sl > math.MaxUint8 && sl <= math.MaxUint16: + dt.IndexType = arrow.PrimitiveTypes.Uint16 + case sl > math.MaxUint16 && sl <= math.MaxUint32: + dt.IndexType = arrow.PrimitiveTypes.Uint32 + } + n.arrowField = arrow.Field{Name: n.Name(), Type: &dt, Nullable: true, Metadata: arrow.MetadataFrom(symbols)} + case "array": + // logical items type + c := n.newChild(n.name, n.schema.(*avro.ArraySchema).Items()) + if isLogicalSchemaType(n.schema.(*avro.ArraySchema).Items()) { + avroLogicalToArrowField(c) + } else { + arrowSchemafromAvro(c) + } + n.arrowField = arrow.Field{Name: n.name, Type: arrow.ListOf(c.arrowField.Type)} + case "map": + c := n.newChild(n.name, n.schema.(*avro.MapSchema).Values()) + arrowSchemafromAvro(c) + n.arrowField = arrow.Field{Name: n.name, Type: arrow.MapOf(arrow.BinaryTypes.String, c.arrowField.Type), Metadata: c.arrowField.Metadata, Nullable: true} + case "union": + if n.schema.(*avro.UnionSchema).Nullable() { + if len(n.schema.(*avro.UnionSchema).Types()) > 1 { + n.schema = n.schema.(*avro.UnionSchema).Types()[1] + n.union = true + arrowSchemafromAvro(n) + } + } + // Avro "fixed" field type = Arrow FixedSize Primitive BinaryType + case "fixed": + if isLogicalSchemaType(n.schema) { + avroLogicalToArrowField(n) + } else { + n.arrowField = arrow.Field{Name: n.Name(), Type: &arrow.FixedSizeBinaryType{ByteWidth: n.schema.(*avro.FixedSchema).Size()}, Nullable: true} + } + case "string": + if isLogicalSchemaType(n.schema) { + avroLogicalToArrowField(n) + } else { + n.arrowField = arrow.Field{Name: n.name, Type: avroPrimitiveToArrowType(string(st)), Nullable: true} + } + case "bytes": + if isLogicalSchemaType(n.schema) { + avroLogicalToArrowField(n) + } else { + n.arrowField = arrow.Field{Name: n.name, Type: avroPrimitiveToArrowType(string(st)), Nullable: true} + } + case "int": + if isLogicalSchemaType(n.schema) { + avroLogicalToArrowField(n) + } else { + n.arrowField = arrow.Field{Name: n.name, Type: avroPrimitiveToArrowType(string(st)), Nullable: true} + } + case "long": + if isLogicalSchemaType(n.schema) { + avroLogicalToArrowField(n) + } else { + n.arrowField = arrow.Field{Name: n.name, Type: avroPrimitiveToArrowType(string(st)), Nullable: true} + } + case "float": + n.arrowField = arrow.Field{Name: n.name, Type: avroPrimitiveToArrowType(string(st)), Nullable: true} + case "double": + n.arrowField = arrow.Field{Name: n.name, Type: avroPrimitiveToArrowType(string(st)), Nullable: true} + case "boolean": + n.arrowField = arrow.Field{Name: n.name, Type: avroPrimitiveToArrowType(string(st)), Nullable: true} + case "<ref>": + n.schema = schemaCache.Get(n.schema.(*avro.RefSchema).Schema().Name()) + arrowSchemafromAvro(n) + case "null": + n.arrowField = arrow.Field{Name: n.name, Type: arrow.BinaryTypes.Binary, Nullable: true} + } +} + +// iterate record Fields() +func iterateFields(n *schemaNode) { + for _, f := range n.schema.(avro.PropertySchema).(*avro.RecordSchema).Fields() { + switch ft := f.Type().(type) { + // Avro "array" field type + case *avro.ArraySchema: + schemaCache.Add(f.Name(), ft.Items()) + // logical items type + c := n.newChild(f.Name(), ft.Items()) + if isLogicalSchemaType(ft.Items()) { + avroLogicalToArrowField(c) + } else { + arrowSchemafromAvro(c) + } + c.arrowField = arrow.Field{Name: c.name, Type: arrow.ListOf(c.arrowField.Type), Metadata: c.arrowField.Metadata, Nullable: true} + // Avro "enum" field type = Arrow dictionary type + case *avro.EnumSchema: + schemaCache.Add(f.Name(), f.Type()) + c := n.newChild(f.Name(), f.Type()) + symbols := make(map[string]string) + for index, symbol := range ft.Symbols() { + k := strconv.FormatInt(int64(index), 10) + symbols[k] = symbol + } + var dt arrow.DictionaryType = arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Uint64, ValueType: arrow.BinaryTypes.String, Ordered: false} + sl := len(symbols) + switch { + case sl <= math.MaxUint8: + dt.IndexType = arrow.PrimitiveTypes.Uint8 + case sl > math.MaxUint8 && sl <= math.MaxUint16: + dt.IndexType = arrow.PrimitiveTypes.Uint16 + case sl > math.MaxUint16 && sl <= math.MaxUint32: + dt.IndexType = arrow.PrimitiveTypes.Uint32 + } + c.arrowField = arrow.Field{Name: f.Name(), Type: &dt, Nullable: true, Metadata: arrow.MetadataFrom(symbols)} + // Avro "fixed" field type = Arrow FixedSize Primitive BinaryType + case *avro.FixedSchema: + schemaCache.Add(f.Name(), f.Type()) + if isLogicalSchemaType(f.Type()) { + c := n.newChild(f.Name(), f.Type()) + avroLogicalToArrowField(c) + } else { + c := n.newChild(f.Name(), f.Type()) + arrowSchemafromAvro(c) + } + case *avro.RecordSchema: + schemaCache.Add(f.Name(), f.Type()) + c := n.newChild(f.Name(), f.Type()) + iterateFields(c) + // Avro "map" field type - KVP with value of one type - keys are strings + case *avro.MapSchema: + schemaCache.Add(f.Name(), ft.Values()) + c := n.newChild(f.Name(), ft.Values()) + arrowSchemafromAvro(c) + if ft.Values().Type() == "union" { + c.arrowField = arrow.Field{Name: f.Name(), Type: arrow.MapOf(arrow.BinaryTypes.String, arrow.BinaryTypes.String), Metadata: c.arrowField.Metadata, Nullable: true} + } else { + c.arrowField = arrow.Field{Name: f.Name(), Type: arrow.MapOf(arrow.BinaryTypes.String, c.arrowField.Type), Metadata: c.arrowField.Metadata, Nullable: true} + } + case *avro.UnionSchema: + if ft.Nullable() { + if len(ft.Types()) > 1 { + schemaCache.Add(f.Name(), ft.Types()[1]) + c := n.newChild(f.Name(), ft.Types()[1]) + c.union = true + arrowSchemafromAvro(c) + } + } + default: + schemaCache.Add(f.Name(), f.Type()) Review Comment: As far as I can determine from the spec and samples I've seen, in Avro a referenced field has to be defined first in the depth-first, left-to-right traversal of the JSON parse tree so if the name is reused it would get overwritten in the schema cache before being referenced again. -- 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]
