Improve OCF encoder/decoder handling of dynamic types

ocf/ocf.go

@@ -6,6 +6,7 @@
 import (
 	"bytes"
 	"crypto/rand"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
@@ -20,10 +21,11 @@
 	codecKey  = "avro.codec"
 )
 
-var magicBytes = [4]byte{'O', 'b', 'j', 1}
+var (
+	magicBytes = [4]byte{'O', 'b', 'j', 1}
 
-// HeaderSchema is the Avro schema of a container file header.
-var HeaderSchema = avro.MustParse(`{
+	// HeaderSchema is the Avro schema of a container file header.
+	HeaderSchema = avro.MustParse(`{
 	"type": "record",
 	"name": "org.apache.avro.file.Header",
 	"fields": [
@@ -33,6 +35,15 @@
 	]
 }`)
 
+	// DefaultSchemaMarshaler calls the schema's String() method, to produce
+	// a "canonical" schema.
+	DefaultSchemaMarshaler = defaultMarshalSchema
+	// FullSchemaMarshaler calls the schema's MarshalJSON() method, to produce
+	// a schema with all details preserved. The "canonical" schema returned by
+	// the default marshaler does not preserve a type's extra properties.
+	FullSchemaMarshaler = fullMarshalSchema
+)
+
 // Header represents an Avro container file header.
 type Header struct {
 	Magic [4]byte           `avro:"magic"`
@@ -42,6 +53,7 @@
 
 type decoderConfig struct {
 	DecoderConfig avro.API
+	SchemaCache   *avro.SchemaCache
 }
 
 // DecoderFunc represents a configuration function for Decoder.
@@ -54,13 +66,22 @@
 	}
 }
 
+// WithDecoderSchemaCache sets the schema cache for the decoder.
+// If not specified, defaults to avro.DefaultSchemaCache.
+func WithDecoderSchemaCache(cache *avro.SchemaCache) DecoderFunc {
+	return func(cfg *decoderConfig) {
+		cfg.SchemaCache = cache
+	}
+}
+
 // Decoder reads and decodes Avro values from a container file.
 type Decoder struct {
 	reader      *avro.Reader
 	resetReader *bytesx.ResetReader
 	decoder     *avro.Decoder
 	meta        map[string][]byte
 	sync        [16]byte
+	schema      avro.Schema
 
 	codec Codec
 
@@ -71,14 +92,15 @@
 func NewDecoder(r io.Reader, opts ...DecoderFunc) (*Decoder, error) {
 	cfg := decoderConfig{
 		DecoderConfig: avro.DefaultConfig,
+		SchemaCache:   avro.DefaultSchemaCache,
 	}
 	for _, opt := range opts {
 		opt(&cfg)
 	}
 
 	reader := avro.NewReader(r, 1024)
 
-	h, err := readHeader(reader)
+	h, err := readHeader(reader, cfg.SchemaCache)
 	if err != nil {
 		return nil, fmt.Errorf("decoder: %w", err)
 	}
@@ -92,6 +114,7 @@
 		meta:        h.Meta,
 		sync:        h.Sync,
 		codec:       h.Codec,
+		schema:      h.Schema,
 	}, nil
 }
 
@@ -100,6 +123,10 @@
 	return d.meta
 }
 
+func (d *Decoder) Schema() avro.Schema {
+	return d.schema
+}
+
 // HasNext determines if there is another value to read.
 func (d *Decoder) HasNext() bool {
 	if d.count <= 0 {
@@ -174,6 +201,8 @@
 	Metadata         map[string][]byte
 	Sync             [16]byte
 	EncodingConfig   avro.API
+	SchemaCache      *avro.SchemaCache
+	SchemaMarshaler  func(avro.Schema) ([]byte, error)
 }
 
 // EncoderFunc represents a configuration function for Encoder.
@@ -209,6 +238,22 @@
 	}
 }
 
+// WithEncoderSchemaCache sets the schema cache for the encoder.
+// If not specified, defaults to avro.DefaultSchemaCache.
+func WithEncoderSchemaCache(cache *avro.SchemaCache) EncoderFunc {
+	return func(cfg *encoderConfig) {
+		cfg.SchemaCache = cache
+	}
+}
+
+// WithSchemaMarshaler sets the schema marshaler for the encoder.
+// If not specified, defaults to DefaultSchemaMarshaler.
+func WithSchemaMarshaler(m func(avro.Schema) ([]byte, error)) EncoderFunc {
+	return func(cfg *encoderConfig) {
+		cfg.SchemaMarshaler = m
+	}
+}
+
 // WithSyncBlock sets the sync block.
 func WithSyncBlock(sync [16]byte) EncoderFunc {
 	return func(cfg *encoderConfig) {
@@ -241,12 +286,36 @@
 // If the writer is an existing ocf file, it will append data using the
 // existing schema.
 func NewEncoder(s string, w io.Writer, opts ...EncoderFunc) (*Encoder, error) {
+	return newEncoder(
+		func(cache *avro.SchemaCache) (avro.Schema, error) {
+			return avro.ParseWithCache(s, "", cache)
+		},
+		w,
+		opts...)
+}
+
+// NewEncoderWithSchema returns a new encoder that writes to w using schema s.
+//
+// If the writer is an existing ocf file, it will append data using the
+// existing schema.
+func NewEncoderWithSchema(schema avro.Schema, w io.Writer, opts ...EncoderFunc) (*Encoder, error) {
+	return newEncoder(
+		func(_ *avro.SchemaCache) (avro.Schema, error) {
+			return schema, nil
+		},
+		w,
+		opts...)
+}
+
+func newEncoder(getSchema func(*avro.SchemaCache) (avro.Schema, error), w io.Writer, opts ...EncoderFunc) (*Encoder, error) {
 	cfg := encoderConfig{
 		BlockLength:      100,
 		CodecName:        Null,
 		CodecCompression: -1,
 		Metadata:         map[string][]byte{},
 		EncodingConfig:   avro.DefaultConfig,
+		SchemaCache:      avro.DefaultSchemaCache,
+		SchemaMarshaler:  DefaultSchemaMarshaler,
 	}
 	for _, opt := range opts {
 		opt(&cfg)
@@ -263,7 +332,7 @@
 
 		if info.Size() > 0 {
 			reader := avro.NewReader(file, 1024)
-			h, err := readHeader(reader)
+			h, err := readHeader(reader, cfg.SchemaCache)
 			if err != nil {
 				return nil, err
 			}
@@ -285,12 +354,16 @@
 		}
 	}
 
-	schema, err := avro.Parse(s)
+	schema, err := getSchema(cfg.SchemaCache)
+	if err != nil {
+		return nil, err
+	}
+	schemaJSON, err := cfg.SchemaMarshaler(schema)
 	if err != nil {
 		return nil, err
 	}
 
-	cfg.Metadata[schemaKey] = []byte(schema.String())
+	cfg.Metadata[schemaKey] = schemaJSON
 	cfg.Metadata[codecKey] = []byte(cfg.CodecName)
 	header := Header{
 		Magic: magicBytes,
@@ -400,7 +473,7 @@
 	Sync   [16]byte
 }
 
-func readHeader(reader *avro.Reader) (*ocfHeader, error) {
+func readHeader(reader *avro.Reader, schemaCache *avro.SchemaCache) (*ocfHeader, error) {
 	var h Header
 	reader.ReadVal(HeaderSchema, &h)
 	if reader.Error != nil {
@@ -410,7 +483,7 @@
 	if h.Magic != magicBytes {
 		return nil, errors.New("invalid avro file")
 	}
-	schema, err := avro.Parse(string(h.Meta[schemaKey]))
+	schema, err := avro.ParseBytesWithCache(h.Meta[schemaKey], "", schemaCache)
 	if err != nil {
 		return nil, err
 	}
@@ -447,3 +520,11 @@
 		}
 	}
 }
+
+func defaultMarshalSchema(schema avro.Schema) ([]byte, error) {
+	return []byte(schema.String()), nil
+}
+
+func fullMarshalSchema(schema avro.Schema) ([]byte, error) {
+	return json.Marshal(schema)
+}