feat(arrow/array): convert RecordReader and iterators (#314)

### Rationale for this change
With the advent of Go iterators via the [iter](https://pkg.go.dev/iter)
module, we should provide some easy compatibility and canonicity for
handling streams of record batches.

### What changes are included in this PR?
Two new functions: `ReaderFromIter` and `IterFromReader` for converting
between the `RecordReader` interface and an iterator of Records and
errors. This should make it easy to integrate in various packages
without forcing refactors or boilerplate code.

### Are these changes tested?
Yes, unit tests are added for them.

### Are there any user-facing changes?
No.

Author: zeroshade

arrow/array/record.go

@@ -19,6 +19,7 @@ package array
 import (
 	"bytes"
 	"fmt"
+	"iter"
 	"strings"
 	"sync/atomic"
 
@@ -405,6 +406,84 @@ func (b *RecordBuilder) UnmarshalJSON(data []byte) error {
 	return nil
 }
 
+type iterReader struct {
+	refCount atomic.Int64
+
+	schema *arrow.Schema
+	cur    arrow.Record
+
+	next func() (arrow.Record, error, bool)
+	stop func()
+
+	err error
+}
+
+func (ir *iterReader) Schema() *arrow.Schema { return ir.schema }
+
+func (ir *iterReader) Retain() { ir.refCount.Add(1) }
+func (ir *iterReader) Release() {
+	debug.Assert(ir.refCount.Load() > 0, "too many releases")
+
+	if ir.refCount.Add(-1) == 0 {
+		ir.stop()
+		ir.schema, ir.next = nil, nil
+		if ir.cur != nil {
+			ir.cur.Release()
+		}
+	}
+}
+
+func (ir *iterReader) Record() arrow.Record { return ir.cur }
+func (ir *iterReader) Err() error           { return ir.err }
+
+func (ir *iterReader) Next() bool {
+	if ir.cur != nil {
+		ir.cur.Release()
+	}
+
+	var ok bool
+	ir.cur, ir.err, ok = ir.next()
+	if ir.err != nil {
+		ir.stop()
+		return false
+	}
+
+	return ok
+}
+
+// ReaderFromIter wraps a go iterator for arrow.Record + error into a RecordReader
+// interface object for ease of use.
+func ReaderFromIter(schema *arrow.Schema, itr iter.Seq2[arrow.Record, error]) RecordReader {
+	next, stop := iter.Pull2(itr)
+	rdr := &iterReader{
+		schema: schema,
+		next:   next,
+		stop:   stop,
+	}
+	rdr.refCount.Add(1)
+	return rdr
+}
+
+// IterFromReader converts a RecordReader interface into an iterator that
+// you can use range on. The semantics are still important, if a record
+// that is returned is desired to be utilized beyond the scope of an iteration
+// then Retain must be called on it.
+func IterFromReader(rdr RecordReader) iter.Seq2[arrow.Record, error] {
+	rdr.Retain()
+	return func(yield func(arrow.Record, error) bool) {
+		defer rdr.Release()
+		for rdr.Next() {
+			if !yield(rdr.Record(), nil) {
+				return
+			}
+		}
+
+		if rdr.Err() != nil {
+			yield(nil, rdr.Err())
+		}
+	}
+}
+
 var (
 	_ arrow.Record = (*simpleRecord)(nil)
 	_ RecordReader = (*simpleRecords)(nil)

arrow/array/record_test.go

@@ -301,33 +301,97 @@ func TestRecordReader(t *testing.T) {
 	defer rec2.Release()
 
 	recs := []arrow.Record{rec1, rec2}
-	itr, err := array.NewRecordReader(schema, recs)
-	if err != nil {
-		t.Fatal(err)
-	}
-	defer itr.Release()
+	t.Run("simple reader", func(t *testing.T) {
+		itr, err := array.NewRecordReader(schema, recs)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer itr.Release()
 
-	itr.Retain()
-	itr.Release()
+		itr.Retain()
+		itr.Release()
 
-	if got, want := itr.Schema(), schema; !got.Equal(want) {
-		t.Fatalf("invalid schema. got=%#v, want=%#v", got, want)
-	}
+		if got, want := itr.Schema(), schema; !got.Equal(want) {
+			t.Fatalf("invalid schema. got=%#v, want=%#v", got, want)
+		}
 
-	n := 0
-	for itr.Next() {
-		n++
-		if got, want := itr.Record(), recs[n-1]; !reflect.DeepEqual(got, want) {
-			t.Fatalf("itr[%d], invalid record. got=%#v, want=%#v", n-1, got, want)
+		n := 0
+		for itr.Next() {
+			n++
+			if got, want := itr.Record(), recs[n-1]; !reflect.DeepEqual(got, want) {
+				t.Fatalf("itr[%d], invalid record. got=%#v, want=%#v", n-1, got, want)
+			}
+		}
+		if err := itr.Err(); err != nil {
+			t.Fatalf("itr error: %#v", err)
 		}
-	}
-	if err := itr.Err(); err != nil {
-		t.Fatalf("itr error: %#v", err)
-	}
 
-	if n != len(recs) {
-		t.Fatalf("invalid number of iterations. got=%d, want=%d", n, len(recs))
-	}
+		if n != len(recs) {
+			t.Fatalf("invalid number of iterations. got=%d, want=%d", n, len(recs))
+		}
+	})
+
+	t.Run("iter to reader", func(t *testing.T) {
+		itr := func(yield func(arrow.Record, error) bool) {
+			for _, r := range recs {
+				if !yield(r, nil) {
+					return
+				}
+			}
+		}
+
+		rdr := array.ReaderFromIter(schema, itr)
+		defer rdr.Release()
+
+		rdr.Retain()
+		rdr.Release()
+
+		if got, want := rdr.Schema(), schema; !got.Equal(want) {
+			t.Fatalf("invalid schema. got=%#v, want=%#v", got, want)
+		}
+
+		n := 0
+		for rdr.Next() {
+			n++
+			// facet of using the simple record reader with a slice
+			// by default it will release records when the reader is released
+			// leading to too many releases on the original record
+			// so we retain it to keep it from going away while the test runs
+			rdr.Record().Retain()
+			if got, want := rdr.Record(), recs[n-1]; !reflect.DeepEqual(got, want) {
+				t.Fatalf("itr[%d], invalid record. got=%#v, want=%#v", n-1, got, want)
+			}
+		}
+		if err := rdr.Err(); err != nil {
+			t.Fatalf("itr error: %#v", err)
+		}
+
+		if n != len(recs) {
+			t.Fatalf("invalid number of iterations. got=%d, want=%d", n, len(recs))
+		}
+	})
+
+	t.Run("reader to iter", func(t *testing.T) {
+		rdr, err := array.NewRecordReader(schema, recs)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		itr := array.IterFromReader(rdr)
+		rdr.Release()
+
+		n := 0
+		for rec, err := range itr {
+			if err != nil {
+				t.Fatalf("itr error: %#v", err)
+			}
+
+			n++
+			if got, want := rec, recs[n-1]; !reflect.DeepEqual(got, want) {
+				t.Fatalf("itr[%d], invalid record. got=%#v, want=%#v", n-1, got, want)
+			}
+		}
+	})
 
 	for _, tc := range []struct {
 		name   string