Batch of small optimizations (apache#556)

This primarily consists of reducing allocations.

e.g. 03be6d3 adds a specific case to TypeEqual for ListView types,
which avoids reflect.DeepEqual. l.elem.Equal does not allocate;
reflect.DeepEqual does, so this reduces allocations when comparing a
ListViewType.

8c7729f changes from referencing flatbug internal types by reference to
by-value, which avoids letting them escape to the heap.

e340658 switches from reflect.DeepEqual to slices.Equal when comparing
the child IDs of a Union type, which is functionally equivalent for
slices and, again, does not allocate.

9eba6d2 does similar for an arrowflight cookie test, replacing
reflect.DeepEquals with maps.Equal; I did a grep for reflect.DeepEqual
and noticed it.

7557d15 is probably the most objectionable; it replaces a *float64 with
float64 directly for minSpaceSavings, since the current nil value has
identical behavior to zero anyways. It's a tiny cleanup IMO, but doesn't
really have any practical value.

2d3593f is a bit ugly; it inlines a copy of NewReaderFromMessageReader
and a copy of NewMessageReader into NewReader - by implementing the two
directly, it's able to avoid allocating two copies of the `config`
structure, and to avoid looping over and invoking config functions
twice. It's kinda ugly though, and the `// types: make(dictTypeMap)`
which is copied over suggests that there's further room for improvement.

I didn't drop the two APIs, as they are exposed to users, but IMO the
performance win is probably worth it, and it can be cleaned up later if
needed.

0273ffe is a small optimization to the allocator: the Go allocator
banches on each allocation, but both branches are functionally
equivalent, so I just merged the two. I doubt this would have a _major_
impact, since the cost of a branch is negligible compared to the cost of
a heap allocation, but performance is often won with a lot of small
impacts IMO :)

Lastly, for now at least, f1a6a33 statically allocates 4 bytes per IPC
messageReader (I did _not_ touch the ArrowFlight implementation), and
uses it on every message read, rather than heap-allocating 4 bytes each
time Message is called.


There's more to be done, but this should be a much more practically
sized chunk to review, and I can add more later :)

---------

Co-authored-by: Noam Preil <[email protected]>

Author: pixelherodev

arrow/compare.go

@@ -18,6 +18,7 @@ package arrow
 
 import (
 	"reflect"
+	"slices"
 )
 
 type typeEqualsConfig struct {
@@ -122,7 +123,7 @@ func TypeEqual(left, right DataType, opts ...TypeEqualOption) bool {
 			return false
 		}
 
-		if !reflect.DeepEqual(l.ChildIDs(), r.ChildIDs()) {
+		if !slices.Equal(l.ChildIDs(), r.ChildIDs()) {
 			return false
 		}
 
@@ -149,6 +150,8 @@ func TypeEqual(left, right DataType, opts ...TypeEqualOption) bool {
 		r := right.(*RunEndEncodedType)
 		return TypeEqual(l.Encoded(), r.Encoded(), opts...) &&
 			TypeEqual(l.runEnds, r.runEnds, opts...)
+	case *ListViewType:
+		return l.elem.Equal(right.(*ListViewType).elem)
 	default:
 		return reflect.DeepEqual(left, right)
 	}

arrow/flight/basic_auth_flight_test.go

@@ -73,11 +73,17 @@ func (*validator) IsValid(bearerToken string) (interface{}, error) {
 func TestErrorAuths(t *testing.T) {
 	unary, stream := flight.CreateServerBearerTokenAuthInterceptors(&validator{})
 	s := flight.NewFlightServer(grpc.UnaryInterceptor(unary), grpc.StreamInterceptor(stream))
-	s.Init("localhost:0")
+	if err := s.Init("localhost:0"); err != nil {
+		panic(err)
+	}
 	f := &HeaderAuthTestFlight{}
 	s.RegisterFlightService(f)
 
-	go s.Serve()
+	go func() {
+		if err := s.Serve(); err != nil {
+			panic(err)
+		}
+	}()
 	defer s.Shutdown()
 
 	client, err := flight.NewFlightClient(s.Addr().String(), nil, grpc.WithTransportCredentials(insecure.NewCredentials()))

arrow/flight/cookie_middleware_test.go

@@ -21,9 +21,9 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"maps"
 	"net/http"
 	"net/textproto"
-	"reflect"
 	"strings"
 	"testing"
 	"time"
@@ -85,7 +85,7 @@ func (s *serverAddCookieMiddleware) StartCall(ctx context.Context) context.Conte
 		}
 	}
 
-	if !reflect.DeepEqual(s.expectedCookies, got) {
+	if !maps.Equal(s.expectedCookies, got) {
 		panic(fmt.Sprintf("did not get expected cookies, expected %+v, got %+v", s.expectedCookies, got))
 	}
 

arrow/ipc/file_reader.go

@@ -531,12 +531,12 @@ func (src *ipcSource) buffer(i int) *memory.Buffer {
 	return raw
 }
 
-func (src *ipcSource) fieldMetadata(i int) *flatbuf.FieldNode {
+func (src *ipcSource) fieldMetadata(i int) flatbuf.FieldNode {
 	var node flatbuf.FieldNode
 	if !src.meta.Nodes(&node, i) {
 		panic("arrow/ipc: field metadata out of bound")
 	}
-	return &node
+	return node
 }
 
 func (src *ipcSource) variadicCount(i int) int64 {
@@ -553,7 +553,7 @@ type arrayLoaderContext struct {
 	version   MetadataVersion
 }
 
-func (ctx *arrayLoaderContext) field() *flatbuf.FieldNode {
+func (ctx *arrayLoaderContext) field() flatbuf.FieldNode {
 	field := ctx.src.fieldMetadata(ctx.ifield)
 	ctx.ifield++
 	return field
@@ -647,7 +647,7 @@ func (ctx *arrayLoaderContext) loadArray(dt arrow.DataType) arrow.ArrayData {
 	}
 }
 
-func (ctx *arrayLoaderContext) loadCommon(typ arrow.Type, nbufs int) (*flatbuf.FieldNode, []*memory.Buffer) {
+func (ctx *arrayLoaderContext) loadCommon(typ arrow.Type, nbufs int) (flatbuf.FieldNode, []*memory.Buffer) {
 	buffers := make([]*memory.Buffer, 0, nbufs)
 	field := ctx.field()
 

arrow/ipc/file_writer.go

@@ -247,7 +247,7 @@ type FileWriter struct {
 	codec           flatbuf.CompressionType
 	compressNP      int
 	compressors     []compressor
-	minSpaceSavings *float64
+	minSpaceSavings float64
 
 	// map of the last written dictionaries by id
 	// so we can avoid writing the same dictionary over and over

arrow/ipc/ipc.go

@@ -71,7 +71,7 @@ type config struct {
 	ensureNativeEndian bool
 	noAutoSchema       bool
 	emitDictDeltas     bool
-	minSpaceSavings    *float64
+	minSpaceSavings    float64
 }
 
 func newConfig(opts ...Option) *config {
@@ -189,7 +189,7 @@ func WithDictionaryDeltas(v bool) Option {
 // Go and C++ versions prior to 12.0.0.
 func WithMinSpaceSavings(savings float64) Option {
 	return func(cfg *config) {
-		cfg.minSpaceSavings = &savings
+		cfg.minSpaceSavings = savings
 	}
 }
 

arrow/ipc/message.go

@@ -149,7 +149,8 @@ type messageReader struct {
 	refCount atomic.Int64
 	msg      *Message
 
-	mem memory.Allocator
+	mem    memory.Allocator
+	header [4]byte
 }
 
 // NewMessageReader returns a reader that reads messages from an input stream.
@@ -188,7 +189,7 @@ func (r *messageReader) Release() {
 // underlying stream.
 // It is valid until the next call to Message.
 func (r *messageReader) Message() (*Message, error) {
-	buf := make([]byte, 4)
+	buf := r.header[:]
 	_, err := io.ReadFull(r.r, buf)
 	if err != nil {
 		return nil, fmt.Errorf("arrow/ipc: could not read continuation indicator: %w", err)

arrow/ipc/reader.go

@@ -89,8 +89,32 @@ func NewReaderFromMessageReader(r MessageReader, opts ...Option) (reader *Reader
 }
 
 // NewReader returns a reader that reads records from an input stream.
-func NewReader(r io.Reader, opts ...Option) (*Reader, error) {
-	return NewReaderFromMessageReader(NewMessageReader(r, opts...), opts...)
+func NewReader(r io.Reader, opts ...Option) (rr *Reader, err error) {
+	defer func() {
+		if pErr := recover(); pErr != nil {
+			err = utils.FormatRecoveredError("arrow/ipc: unknown error while reading", pErr)
+		}
+	}()
+	cfg := newConfig(opts...)
+	mr := &messageReader{r: r, mem: cfg.alloc}
+	mr.refCount.Add(1)
+	rr = &Reader{
+		r:        mr,
+		refCount: atomic.Int64{},
+		// types:    make(dictTypeMap),
+		memo:               dictutils.NewMemo(),
+		mem:                cfg.alloc,
+		ensureNativeEndian: cfg.ensureNativeEndian,
+		expectedSchema:     cfg.schema,
+	}
+	rr.refCount.Add(1)
+
+	if !cfg.noAutoSchema {
+		if err := rr.readSchema(cfg.schema); err != nil {
+			return nil, err
+		}
+	}
+	return rr, nil
 }
 
 // Err returns the last error encountered during the iteration over the

arrow/ipc/writer.go

@@ -87,7 +87,7 @@ type Writer struct {
 	codec           flatbuf.CompressionType
 	compressNP      int
 	compressors     []compressor
-	minSpaceSavings *float64
+	minSpaceSavings float64
 
 	// map of the last written dictionaries by id
 	// so we can avoid writing the same dictionary over and over
@@ -328,7 +328,7 @@ type recordEncoder struct {
 	codec           flatbuf.CompressionType
 	compressNP      int
 	compressors     []compressor
-	minSpaceSavings *float64
+	minSpaceSavings float64
 }
 
 func newRecordEncoder(
@@ -338,7 +338,7 @@ func newRecordEncoder(
 	allow64b bool,
 	codec flatbuf.CompressionType,
 	compressNP int,
-	minSpaceSavings *float64,
+	minSpaceSavings float64,
 	compressors []compressor,
 ) *recordEncoder {
 	return &recordEncoder{
@@ -355,12 +355,12 @@ func newRecordEncoder(
 
 func (w *recordEncoder) shouldCompress(uncompressed, compressed int) bool {
 	debug.Assert(uncompressed > 0, "uncompressed size is 0")
-	if w.minSpaceSavings == nil {
+	if w.minSpaceSavings == 0 {
 		return true
 	}
 
 	savings := 1.0 - float64(compressed)/float64(uncompressed)
-	return savings >= *w.minSpaceSavings
+	return savings >= w.minSpaceSavings
 }
 
 func (w *recordEncoder) reset() {
@@ -477,13 +477,10 @@ func (w *recordEncoder) encode(p *Payload, rec arrow.RecordBatch) error {
 	}
 
 	if w.codec != -1 {
-		if w.minSpaceSavings != nil {
-			pct := *w.minSpaceSavings
-			if pct < 0 || pct > 1 {
-				p.Release()
-				return fmt.Errorf("%w: minSpaceSavings not in range [0,1]. Provided %.05f",
-					arrow.ErrInvalid, pct)
-			}
+		if w.minSpaceSavings < 0 || w.minSpaceSavings > 1 {
+			p.Release()
+			return fmt.Errorf("%w: minSpaceSavings not in range [0,1]. Provided %.05f",
+				arrow.ErrInvalid, w.minSpaceSavings)
 		}
 		w.compressBodyBuffers(p)
 	}

arrow/ipc/writer_test.go

@@ -195,7 +195,7 @@ func TestWriteWithCompressionAndMinSavings(t *testing.T) {
 
 	for _, codec := range []flatbuf.CompressionType{flatbuf.CompressionTypeLZ4_FRAME, flatbuf.CompressionTypeZSTD} {
 		compressors := []compressor{getCompressor(codec)}
-		enc := newRecordEncoder(mem, 0, 5, true, codec, 1, nil, compressors)
+		enc := newRecordEncoder(mem, 0, 5, true, codec, 1, 0, compressors)
 		var payload Payload
 		require.NoError(t, enc.encode(&payload, batch))
 		assert.Len(t, payload.body, 2)
@@ -207,7 +207,7 @@ func TestWriteWithCompressionAndMinSavings(t *testing.T) {
 		assert.Greater(t, compressedSize, int64(0))
 		expectedSavings := 1.0 - float64(compressedSize)/float64(uncompressedSize)
 
-		compressEncoder := newRecordEncoder(mem, 0, 5, true, codec, 1, &expectedSavings, compressors)
+		compressEncoder := newRecordEncoder(mem, 0, 5, true, codec, 1, expectedSavings, compressors)
 		payload.Release()
 		payload.body = payload.body[:0]
 		require.NoError(t, compressEncoder.encode(&payload, batch))
@@ -220,7 +220,7 @@ func TestWriteWithCompressionAndMinSavings(t *testing.T) {
 		// slightly bump the threshold. the body buffer should now be prefixed
 		// with -1 and its content left uncompressed
 		minSavings := math.Nextafter(expectedSavings, 1.0)
-		compressEncoder.minSpaceSavings = &minSavings
+		compressEncoder.minSpaceSavings = minSavings
 		require.NoError(t, compressEncoder.encode(&payload, batch))
 		assert.Len(t, payload.body, 2)
 		assert.EqualValues(t, -1, prefixedSize(payload.body[1]))
@@ -229,7 +229,7 @@ func TestWriteWithCompressionAndMinSavings(t *testing.T) {
 		payload.body = payload.body[:0]
 
 		for _, outOfRange := range []float64{math.Nextafter(1.0, 2.0), math.Nextafter(0, -1)} {
-			compressEncoder.minSpaceSavings = &outOfRange
+			compressEncoder.minSpaceSavings = outOfRange
 			err := compressEncoder.encode(&payload, batch)
 			assert.ErrorIs(t, err, arrow.ErrInvalid)
 			assert.ErrorContains(t, err, "minSpaceSavings not in range [0,1]")