zeroshade commented on code in PR #36796: URL: https://github.com/apache/arrow/pull/36796#discussion_r1316384892
########## 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: could there end up being multiple *different* schemas with the same name? Would that be a problem for this cache? ########## 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} Review Comment: should this be type `arrow.Null` ? ########## 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 Review Comment: this should probably not be global to the package, otherwise multiple calls to read schemas might get cached schemas that aren't correct....? ########## 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()) + if isLogicalSchemaType(f.Type()) { + c := n.newChild(f.Name(), f.Type()) + avroLogicalToArrowField(c) + } else { + c := n.newChild(f.Name(), f.Type()) + arrowSchemafromAvro(c) + } + + } + } + var fields []arrow.Field + for _, child := range n.Children() { + fields = append(fields, child.arrowField) + } + + namedSchema, ok := isNamedSchema(n.schema) + if !ok || namedSchema == n.name+"_data" || !n.union { + n.arrowField = arrow.Field{Name: n.name, Type: arrow.StructOf(fields...), Nullable: true} + } else { + n.arrowField = arrow.Field{Name: n.name, Type: arrow.StructOf(fields...), Nullable: true, Metadata: arrow.MetadataFrom(map[string]string{"typeName": namedSchema})} + } + +} + +func isLogicalSchemaType(s avro.Schema) bool { + lts, ok := s.(avro.LogicalTypeSchema) + if !ok { + return false + } + if lt := lts.Logical(); lt != nil { + return true + } + return false +} + +func isNamedSchema(s avro.Schema) (string, bool) { + if ns, ok := s.(avro.NamedSchema); ok { + return ns.FullName(), ok + } + return "", false +} + +// Avro primitive type. +// +// NOTE: Arrow Binary type is used as a catchall to avoid potential data loss. Review Comment: instead of returning a "NotImplemented" error? ########## 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()) + if isLogicalSchemaType(f.Type()) { + c := n.newChild(f.Name(), f.Type()) + avroLogicalToArrowField(c) + } else { + c := n.newChild(f.Name(), f.Type()) + arrowSchemafromAvro(c) + } + + } + } + var fields []arrow.Field + for _, child := range n.Children() { + fields = append(fields, child.arrowField) + } + + namedSchema, ok := isNamedSchema(n.schema) + if !ok || namedSchema == n.name+"_data" || !n.union { + n.arrowField = arrow.Field{Name: n.name, Type: arrow.StructOf(fields...), Nullable: true} + } else { + n.arrowField = arrow.Field{Name: n.name, Type: arrow.StructOf(fields...), Nullable: true, Metadata: arrow.MetadataFrom(map[string]string{"typeName": namedSchema})} + } + +} + +func isLogicalSchemaType(s avro.Schema) bool { + lts, ok := s.(avro.LogicalTypeSchema) + if !ok { + return false + } + if lt := lts.Logical(); lt != nil { + return true + } + return false +} + +func isNamedSchema(s avro.Schema) (string, bool) { + if ns, ok := s.(avro.NamedSchema); ok { + return ns.FullName(), ok + } + return "", false +} + +// Avro primitive type. +// +// NOTE: Arrow Binary type is used as a catchall to avoid potential data loss. +func avroPrimitiveToArrowType(avroFieldType string) arrow.DataType { + switch avroFieldType { + // int: 32-bit signed integer + case "int": + return arrow.PrimitiveTypes.Int32 + // long: 64-bit signed integer + case "long": + return arrow.PrimitiveTypes.Int64 + // float: single precision (32-bit) IEEE 754 floating-point number + case "float": + return arrow.PrimitiveTypes.Float32 + // double: double precision (64-bit) IEEE 754 floating-point number + case "double": + return arrow.PrimitiveTypes.Float64 + // bytes: sequence of 8-bit unsigned bytes + case "bytes": + return arrow.BinaryTypes.Binary + // boolean: a binary value + case "boolean": + return arrow.FixedWidthTypes.Boolean + // string: unicode character sequence + case "string": + return arrow.BinaryTypes.String + // fallback to binary type for any unsupported type + // (shouldn't happen as all types specified in avro 1.11.1 are supported) + default: + return arrow.BinaryTypes.Binary + } +} + +func avroLogicalToArrowField(n *schemaNode) arrow.Field { + var dt arrow.DataType + // Avro logical types + switch lt := n.schema.(avro.LogicalTypeSchema).Logical(); lt.Type() { + // The decimal logical type represents an arbitrary-precision signed decimal number of the form unscaled × 10-scale. + // A decimal logical type annotates Avro bytes or fixed types. The byte array must contain the two’s-complement + // representation of the unscaled integer value in big-endian byte order. The scale is fixed, and is specified + // using an attribute. + // + // The following attributes are supported: + // scale, a JSON integer representing the scale (optional). If not specified the scale is 0. + // precision, a JSON integer representing the (maximum) precision of decimals stored in this type (required). + case "decimal": + id := arrow.DECIMAL128 + if lt.(*avro.DecimalLogicalSchema).Precision() > 38 { Review Comment: use the constant, `decimal128.MaxPrecision` -- 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]
