refactor(arrow/avro): migrate from hamba/avro to twmb/avro

arrow/avro/loader.go

@@ -24,7 +24,7 @@ import (
 
 func (r *OCFReader) decodeOCFToChan() {
 	defer close(r.avroChan)
-	for r.r.HasNext() {
+	for {
 		select {
 		case <-r.readerCtx.Done():
 			r.err = fmt.Errorf("avro decoding cancelled, %d records read", r.avroDatumCount)
@@ -34,7 +34,6 @@ func (r *OCFReader) decodeOCFToChan() {
 			err := r.r.Decode(&datum)
 			if err != nil {
 				if errors.Is(err, io.EOF) {
-					r.err = nil
 					return
 				}
 				r.err = err

arrow/avro/reader.go

@@ -27,10 +27,9 @@ import (
 	"github.com/apache/arrow-go/v18/arrow/array"
 	"github.com/apache/arrow-go/v18/arrow/internal/debug"
 	"github.com/apache/arrow-go/v18/arrow/memory"
-	"github.com/hamba/avro/v2/ocf"
 	"github.com/tidwall/sjson"
-
-	avro "github.com/hamba/avro/v2"
+	"github.com/twmb/avro"
+	"github.com/twmb/avro/ocf"
 )
 
 var ErrMismatchFields = errors.New("arrow/avro: number of records mismatch")
@@ -47,9 +46,9 @@ type schemaEdit struct {
 	value  any
 }
 
-// Reader wraps goavro/OCFReader and creates array.RecordBatches from a schema.
+// OCFReader reads Avro OCF files and exposes them as array.RecordBatches.
 type OCFReader struct {
-	r               *ocf.Decoder
+	r               *ocf.Reader
 	avroSchema      string
 	avroSchemaEdits []schemaEdit
 	schema          *arrow.Schema
@@ -82,7 +81,7 @@ type OCFReader struct {
 // NewReader returns a reader that reads from an Avro OCF file and creates
 // arrow.RecordBatches from the converted avro data.
 func NewOCFReader(r io.Reader, opts ...Option) (*OCFReader, error) {
-	ocfr, err := ocf.NewDecoder(r)
+	ocfr, err := ocf.NewReader(r)
 	if err != nil {
 		return nil, fmt.Errorf("%w: could not create avro ocfreader", arrow.ErrInvalid)
 	}
@@ -108,22 +107,20 @@ func NewOCFReader(r io.Reader, opts ...Option) (*OCFReader, error) {
 	}
 	rr.avroSchema = schema.String()
 	if len(rr.avroSchemaEdits) > 0 {
-		// execute schema edits
 		for _, e := range rr.avroSchemaEdits {
 			err := rr.editAvroSchema(e)
 			if err != nil {
 				return nil, fmt.Errorf("%w: could not edit avro schema", arrow.ErrInvalid)
 			}
 		}
-		// validate edited schema
-		schema, err = avro.Parse(rr.avroSchema)
-		if err != nil {
-			return nil, fmt.Errorf("%w: could not parse modified avro schema", arrow.ErrInvalid)
-		}
 	}
-	rr.schema, err = ArrowSchemaFromAvro(schema)
+	rr.schema, err = ArrowSchemaFromAvroJSON(rr.avroSchema)
 	if err != nil {
-		return nil, fmt.Errorf("%w: could not convert avro schema", arrow.ErrInvalid)
+		msg := "could not convert avro schema"
+		if len(rr.avroSchemaEdits) > 0 {
+			msg = "could not parse modified avro schema"
+		}
+		return nil, fmt.Errorf("%w: %s: %w", arrow.ErrInvalid, msg, err)
 	}
 	if rr.mem == nil {
 		rr.mem = memory.DefaultAllocator
@@ -147,7 +144,7 @@ func NewOCFReader(r io.Reader, opts ...Option) (*OCFReader, error) {
 func (rr *OCFReader) Reuse(r io.Reader, opts ...Option) error {
 	rr.Close()
 	rr.err = nil
-	ocfr, err := ocf.NewDecoder(r)
+	ocfr, err := ocf.NewReader(r)
 	if err != nil {
 		return fmt.Errorf("%w: could not create avro ocfreader", arrow.ErrInvalid)
 	}

arrow/avro/reader_test.go

@@ -19,17 +19,13 @@ package avro
 import (
 	"bytes"
 	"encoding/json"
-	"fmt"
 	"os"
 	"path/filepath"
 	"testing"
 
 	"github.com/apache/arrow-go/v18/arrow"
-	"github.com/apache/arrow-go/v18/arrow/array"
 	"github.com/apache/arrow-go/v18/arrow/avro/testdata"
-	"github.com/apache/arrow-go/v18/arrow/memory"
-	hamba "github.com/hamba/avro/v2"
-	"github.com/hamba/avro/v2/ocf"
+	"github.com/apache/arrow-go/v18/arrow/extensions"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -130,6 +126,10 @@ func TestReader(t *testing.T) {
 					Name: "uuidField",
 					Type: arrow.BinaryTypes.String,
 				},
+				{
+					Name: "fixedUuidField",
+					Type: extensions.NewUUIDType(),
+				},
 				{
 					Name: "timemillis",
 					Type: arrow.FixedWidthTypes.Time32ms,
@@ -178,20 +178,13 @@ func TestReader(t *testing.T) {
 				t.Fatal(err)
 			}
 			r := new(OCFReader)
-			r.avroSchema = schema.String()
+			r.avroSchema = schema
 			r.editAvroSchema(schemaEdit{method: "delete", path: "fields.0"})
-			schema, err = hamba.Parse(r.avroSchema)
+			got, err := ArrowSchemaFromAvroJSON(r.avroSchema)
 			if err != nil {
 				t.Fatalf("%v: could not parse modified avro schema", arrow.ErrInvalid)
 			}
-			got, err := ArrowSchemaFromAvro(schema)
-			if err != nil {
-				t.Fatalf("%v", err)
-			}
 			assert.Equal(t, want.String(), got.String())
-			if fmt.Sprintf("%+v", want.String()) != fmt.Sprintf("%+v", got.String()) {
-				t.Fatalf("got=%v,\n want=%v", got.String(), want.String())
-			}
 		})
 
 		t.Run("ShouldLoadExpectedRecords", func(t *testing.T) {
@@ -211,7 +204,7 @@ func TestReader(t *testing.T) {
 			exists := ar.Next()
 
 			if ar.Err() != nil {
-				t.Error("failed to read next record: %w", ar.Err())
+				t.Errorf("failed to read next record: %v", ar.Err())
 			}
 			if !exists {
 				t.Error("no record exists")
@@ -230,109 +223,3 @@ func TestReader(t *testing.T) {
 		})
 	}
 }
-
-// TestOCFReaderBytesValues exercises avro `bytes` fields, both plain and as a
-// ["null","bytes"] union: hamba hands the decoded value to the appenders as a
-// bare []byte, which previously fell into appendBinaryData's fmt fallback and
-// appended the formatted text (e.g. "[1 2 3]") instead of the payload.
-func TestOCFReaderBytesValues(t *testing.T) {
-	schema := `{
-		"type": "record",
-		"name": "rec",
-		"fields": [
-			{"name": "plain", "type": "bytes"},
-			{"name": "nullable", "type": ["null", "bytes"]}
-		]
-	}`
-	payload := []byte{0x00, 0x01, 0xfe, 0xff}
-
-	var buf bytes.Buffer
-	enc, err := ocf.NewEncoder(schema, &buf)
-	assert.NoError(t, err)
-	assert.NoError(t, enc.Encode(map[string]any{
-		"plain":    payload,
-		"nullable": map[string]any{"bytes": payload},
-	}))
-	assert.NoError(t, enc.Encode(map[string]any{
-		"plain":    []byte{},
-		"nullable": nil,
-	}))
-	assert.NoError(t, enc.Close())
-
-	ar, err := NewOCFReader(bytes.NewReader(buf.Bytes()), WithChunk(-1))
-	assert.NoError(t, err)
-	defer ar.Close()
-
-	assert.True(t, ar.Next())
-	assert.NoError(t, ar.Err())
-	rec := ar.RecordBatch()
-
-	plain := rec.Column(0).(*array.Binary)
-	assert.Equal(t, payload, plain.Value(0))
-	assert.Equal(t, []byte{}, plain.Value(1))
-
-	nullable := rec.Column(1).(*array.Binary)
-	assert.Equal(t, payload, nullable.Value(0))
-	assert.True(t, nullable.IsNull(1))
-}
-
-// Types outside what the hamba decoder produces must error rather than append
-// a fmt-formatted rendering of the value.
-func TestAppendBinaryAndStringDataUnexpectedTypes(t *testing.T) {
-	bb := array.NewBinaryBuilder(memory.DefaultAllocator, arrow.BinaryTypes.Binary)
-	defer bb.Release()
-
-	assert.NoError(t, appendBinaryData(bb, []byte{0x01}))
-	assert.NoError(t, appendBinaryData(bb, nil))
-	assert.NoError(t, appendBinaryData(bb, map[string]any{"bytes": []byte{0x02}}))
-	assert.ErrorContains(t, appendBinaryData(bb, 42), "unexpected type int")
-	assert.ErrorContains(t, appendBinaryData(bb, map[string]any{"bytes": "text"}), "unexpected type string")
-	assert.Equal(t, 3, bb.Len())
-
-	sb := array.NewStringBuilder(memory.DefaultAllocator)
-	defer sb.Release()
-
-	assert.NoError(t, appendStringData(sb, "ok"))
-	assert.NoError(t, appendStringData(sb, []byte("ok")))
-	assert.NoError(t, appendStringData(sb, nil))
-	assert.NoError(t, appendStringData(sb, map[string]any{"string": "ok"}))
-	assert.ErrorContains(t, appendStringData(sb, 42), "unexpected type int")
-	assert.ErrorContains(t, appendStringData(sb, map[string]any{"string": 42}), "unexpected type int")
-	assert.Equal(t, 4, sb.Len())
-}
-
-// loadDatum must surface appender errors from nested paths (map values,
-// list items), not only from top-level and struct fields.
-func TestLoadDatumPropagatesNestedAppendErrors(t *testing.T) {
-	newLoader := func(t *testing.T, avroSchema string) (*dataLoader, *array.RecordBuilder) {
-		t.Helper()
-		schema, err := hamba.Parse(avroSchema)
-		assert.NoError(t, err)
-		arrowSchema, err := ArrowSchemaFromAvro(schema)
-		assert.NoError(t, err)
-		bld := array.NewRecordBuilder(memory.DefaultAllocator, arrowSchema)
-		pos := newFieldPos()
-		ldr := newDataLoader()
-		for idx, fb := range bld.Fields() {
-			mapFieldBuilders(fb, arrowSchema.Field(idx), pos)
-		}
-		ldr.drawTree(pos)
-		return ldr, bld
-	}
-
-	t.Run("map value", func(t *testing.T) {
-		ldr, bld := newLoader(t, `{"type":"record","name":"r","fields":[
-			{"name":"m","type":{"type":"map","values":"bytes"}}]}`)
-		defer bld.Release()
-		assert.NoError(t, ldr.loadDatum(map[string]any{"m": map[string]any{"k": []byte{0x01}}}))
-		assert.ErrorContains(t, ldr.loadDatum(map[string]any{"m": map[string]any{"k": 42}}), "unexpected type int")
-	})
-
-	t.Run("list item", func(t *testing.T) {
-		ldr, bld := newLoader(t, `{"type":"record","name":"r","fields":[
-			{"name":"l","type":{"type":"array","items":"bytes"}}]}`)
-		defer bld.Release()
-		assert.NoError(t, ldr.loadDatum(map[string]any{"l": []any{[]byte{0x01}}}))
-		assert.ErrorContains(t, ldr.loadDatum(map[string]any{"l": []any{42}}), "unexpected type int")
-	})
-}