apacheGH-38918: [Go] Avoid schema.Fields allocations in some places (a…

…pache#38919)

### Rationale for this change

Unnecessary allocations.

### What changes are included in this PR?

This PR is split into several commits. The first addresses allocations in the `dictutils` package, the second adds `NumFields` to `NestedType` so that the third commit, which is a purely mechanical change from `len(type.Fields())` to `type.NumFields` to avoid allocations in these specific cases can pass tests with no further changes.

The last commit removes some Fields allocations that specifically hurt our project. Note that this is not an all-encompassing change (therefore this PR should probably not close the linked issue).

### Are these changes tested?

These changes are implicitly tested by the existing test-suite. No functionality has been changed and they should be invisible to the user.

### Are there any user-facing changes?

No.

* Addresses: apache#38918
* Closes: apache#38918

Authored-by: Alfonso Subiotto Marques <[email protected]>
Signed-off-by: Matt Topol <[email protected]>

go/arrow/array/concat.go

@@ -666,7 +666,7 @@ func concat(data []arrow.ArrayData, mem memory.Allocator) (arr arrow.ArrayData,
 		}
 		out.childData = []arrow.ArrayData{children}
 	case *arrow.StructType:
-		out.childData = make([]arrow.ArrayData, len(dt.Fields()))
+		out.childData = make([]arrow.ArrayData, dt.NumFields())
 		for i := range dt.Fields() {
 			children := gatherChildren(data, i)
 			for _, c := range children {

go/arrow/array/record.go

@@ -185,7 +185,7 @@ func (rec *simpleRecord) validate() error {
 		return nil
 	}
 
-	if len(rec.arrs) != len(rec.schema.Fields()) {
+	if len(rec.arrs) != rec.schema.NumFields() {
 		return fmt.Errorf("arrow/array: number of columns/fields mismatch")
 	}
 
@@ -285,11 +285,11 @@ func NewRecordBuilder(mem memory.Allocator, schema *arrow.Schema) *RecordBuilder
 		refCount: 1,
 		mem:      mem,
 		schema:   schema,
-		fields:   make([]Builder, len(schema.Fields())),
+		fields:   make([]Builder, schema.NumFields()),
 	}
 
-	for i, f := range schema.Fields() {
-		b.fields[i] = NewBuilder(b.mem, f.Type)
+	for i := 0; i < schema.NumFields(); i++ {
+		b.fields[i] = NewBuilder(b.mem, schema.Field(i).Type)
 	}
 
 	return b
@@ -397,8 +397,8 @@ func (b *RecordBuilder) UnmarshalJSON(data []byte) error {
 		}
 	}
 
-	for i, f := range b.schema.Fields() {
-		if !keylist[f.Name] {
+	for i := 0; i < b.schema.NumFields(); i++ {
+		if !keylist[b.schema.Field(i).Name] {
 			b.fields[i].AppendNull()
 		}
 	}

go/arrow/array/struct.go

@@ -224,7 +224,7 @@ func NewStructBuilder(mem memory.Allocator, dtype *arrow.StructType) *StructBuil
 	b := &StructBuilder{
 		builder: builder{refCount: 1, mem: mem},
 		dtype:   dtype,
-		fields:  make([]Builder, len(dtype.Fields())),
+		fields:  make([]Builder, dtype.NumFields()),
 	}
 	for i, f := range dtype.Fields() {
 		b.fields[i] = NewBuilder(b.mem, f.Type)

go/arrow/array/table.go

@@ -137,11 +137,11 @@ func NewTable(schema *arrow.Schema, cols []arrow.Column, rows int64) *simpleTabl
 //   - the total length of each column's array slice (ie: number of rows
 //     in the column) aren't the same for all columns.
 func NewTableFromSlice(schema *arrow.Schema, data [][]arrow.Array) *simpleTable {
-	if len(data) != len(schema.Fields()) {
+	if len(data) != schema.NumFields() {
 		panic("array/table: mismatch in number of columns and data for creating a table")
 	}
 
-	cols := make([]arrow.Column, len(schema.Fields()))
+	cols := make([]arrow.Column, schema.NumFields())
 	for i, arrs := range data {
 		field := schema.Field(i)
 		chunked := arrow.NewChunked(field.Type, arrs)
@@ -177,7 +177,7 @@ func NewTableFromSlice(schema *arrow.Schema, data [][]arrow.Array) *simpleTable
 // NewTableFromRecords panics if the records and schema are inconsistent.
 func NewTableFromRecords(schema *arrow.Schema, recs []arrow.Record) *simpleTable {
 	arrs := make([]arrow.Array, len(recs))
-	cols := make([]arrow.Column, len(schema.Fields()))
+	cols := make([]arrow.Column, schema.NumFields())
 
 	defer func(cols []arrow.Column) {
 		for i := range cols {
@@ -224,7 +224,7 @@ func (tbl *simpleTable) NumCols() int64             { return int64(len(tbl.cols)
 func (tbl *simpleTable) Column(i int) *arrow.Column { return &tbl.cols[i] }
 
 func (tbl *simpleTable) validate() {
-	if len(tbl.cols) != len(tbl.schema.Fields()) {
+	if len(tbl.cols) != tbl.schema.NumFields() {
 		panic(errors.New("arrow/array: table schema mismatch"))
 	}
 	for i, col := range tbl.cols {

go/arrow/array/union.go

@@ -896,7 +896,7 @@ func NewEmptySparseUnionBuilder(mem memory.Allocator) *SparseUnionBuilder {
 // children and type codes. Builders will be constructed for each child
 // using the fields in typ
 func NewSparseUnionBuilder(mem memory.Allocator, typ *arrow.SparseUnionType) *SparseUnionBuilder {
-	children := make([]Builder, len(typ.Fields()))
+	children := make([]Builder, typ.NumFields())
 	for i, f := range typ.Fields() {
 		children[i] = NewBuilder(mem, f.Type)
 		defer children[i].Release()
@@ -1129,7 +1129,7 @@ func NewEmptyDenseUnionBuilder(mem memory.Allocator) *DenseUnionBuilder {
 // children and type codes. Builders will be constructed for each child
 // using the fields in typ
 func NewDenseUnionBuilder(mem memory.Allocator, typ *arrow.DenseUnionType) *DenseUnionBuilder {
-	children := make([]Builder, 0, len(typ.Fields()))
+	children := make([]Builder, 0, typ.NumFields())
 	defer func() {
 		for _, child := range children {
 			child.Release()

go/arrow/array/util.go

@@ -428,7 +428,7 @@ func (n *nullArrayFactory) create() *Data {
 	}
 
 	if nf, ok := dt.(arrow.NestedType); ok {
-		childData = make([]arrow.ArrayData, len(nf.Fields()))
+		childData = make([]arrow.ArrayData, nf.NumFields())
 	}
 
 	switch dt := dt.(type) {

go/arrow/cdata/cdata_exports.go

@@ -274,7 +274,7 @@ func (exp *schemaExporter) export(field arrow.Field) {
 		exp.dict = new(schemaExporter)
 		exp.dict.export(arrow.Field{Type: dt.ValueType})
 	case arrow.NestedType:
-		exp.children = make([]schemaExporter, len(dt.Fields()))
+		exp.children = make([]schemaExporter, dt.NumFields())
 		for i, f := range dt.Fields() {
 			exp.children[i].export(f)
 		}

go/arrow/compute/cast.go

@@ -266,8 +266,8 @@ func CastStruct(ctx *exec.KernelCtx, batch *exec.ExecSpan, out *exec.ExecResult)
 		opts          = ctx.State.(kernels.CastState)
 		inType        = batch.Values[0].Array.Type.(*arrow.StructType)
 		outType       = out.Type.(*arrow.StructType)
-		inFieldCount  = len(inType.Fields())
-		outFieldCount = len(outType.Fields())
+		inFieldCount  = inType.NumFields()
+		outFieldCount = outType.NumFields()
 	)
 
 	fieldsToSelect := make([]int, outFieldCount)

go/arrow/compute/exec/span.go

@@ -633,7 +633,7 @@ func FillZeroLength(dt arrow.DataType, span *ArraySpan) {
 		return
 	}
 
-	span.resizeChildren(len(nt.Fields()))
+	span.resizeChildren(nt.NumFields())
 	for i, f := range nt.Fields() {
 		FillZeroLength(f.Type, &span.Children[i])
 	}

go/arrow/compute/exprs/builders.go

@@ -162,7 +162,7 @@ func NewFieldRefFromDotPath(dotpath string, rootSchema *arrow.Schema) (expr.Refe
 			idx, _ := strconv.Atoi(dotpath[:subend])
 			switch ct := curType.(type) {
 			case *arrow.StructType:
-				if idx > len(ct.Fields()) {
+				if idx > ct.NumFields() {
 					return nil, fmt.Errorf("%w: field out of bounds in dotpath", arrow.ErrIndex)
 				}
 				curType = ct.Field(idx).Type

go/arrow/compute/exprs/exec.go

@@ -53,7 +53,7 @@ func makeExecBatch(ctx context.Context, schema *arrow.Schema, partial compute.Da
 		partialBatch := partial.(*compute.RecordDatum).Value
 		batchSchema := partialBatch.Schema()
 
-		out.Values = make([]compute.Datum, len(schema.Fields()))
+		out.Values = make([]compute.Datum, schema.NumFields())
 		out.Len = partialBatch.NumRows()
 
 		for i, field := range schema.Fields() {
@@ -99,7 +99,7 @@ func makeExecBatch(ctx context.Context, schema *arrow.Schema, partial compute.Da
 			return makeExecBatch(ctx, schema, compute.NewDatumWithoutOwning(batch))
 		case *compute.ScalarDatum:
 			out.Len = 1
-			out.Values = make([]compute.Datum, len(schema.Fields()))
+			out.Values = make([]compute.Datum, schema.NumFields())
 
 			s := part.Value.(*scalar.Struct)
 			dt := s.Type.(*arrow.StructType)

go/arrow/compute/exprs/types.go

@@ -594,7 +594,7 @@ func ToSubstraitType(dt arrow.DataType, nullable bool, ext ExtensionIDSet) (type
 			Precision: dt.GetPrecision(), Scale: dt.GetScale()}, nil
 	case arrow.STRUCT:
 		dt := dt.(*arrow.StructType)
-		fields := make([]types.Type, len(dt.Fields()))
+		fields := make([]types.Type, dt.NumFields())
 		var err error
 		for i, f := range dt.Fields() {
 			fields[i], err = ToSubstraitType(f.Type, f.Nullable, ext)

go/arrow/compute/fieldref_test.go

@@ -44,7 +44,7 @@ func TestFieldPathBasics(t *testing.T) {
 	assert.Nil(t, f)
 	assert.ErrorIs(t, err, compute.ErrEmpty)
 
-	f, err = compute.FieldPath{len(s.Fields()) * 2}.Get(s)
+	f, err = compute.FieldPath{s.NumFields() * 2}.Get(s)
 	assert.Nil(t, f)
 	assert.ErrorIs(t, err, compute.ErrIndexRange)
 }
@@ -63,7 +63,7 @@ func TestFieldRefBasics(t *testing.T) {
 	}
 
 	// out of range index results in failure to match
-	assert.Empty(t, compute.FieldRefIndex(len(s.Fields())*2).FindAll(s.Fields()))
+	assert.Empty(t, compute.FieldRefIndex(s.NumFields()*2).FindAll(s.Fields()))
 
 	// lookup by name returns the indices of both matching fields
 	assert.Equal(t, []compute.FieldPath{{0}, {2}}, compute.FieldRefName("alpha").FindAll(s.Fields()))

go/arrow/datatype_encoded.go

@@ -58,6 +58,8 @@ func (t *RunEndEncodedType) Fields() []Field {
 	}
 }
 
+func (t *RunEndEncodedType) NumFields() int { return 2 }
+
 func (*RunEndEncodedType) ValidRunEndsType(dt DataType) bool {
 	switch dt.ID() {
 	case INT16, INT32, INT64:

go/arrow/datatype_extension.go

@@ -161,6 +161,13 @@ func (e *ExtensionBase) Fields() []Field {
 	return nil
 }
 
+func (e *ExtensionBase) NumFields() int {
+	if nested, ok := e.Storage.(NestedType); ok {
+		return nested.NumFields()
+	}
+	return 0
+}
+
 func (e *ExtensionBase) Layout() DataTypeLayout { return e.Storage.Layout() }
 
 // this no-op exists to ensure that this type must be embedded in any user-defined extension type.

go/arrow/datatype_nested.go

@@ -32,6 +32,8 @@ type (
 		// Fields method provides a copy of NestedType fields
 		// (so it can be safely mutated and will not result in updating the NestedType).
 		Fields() []Field
+		// NumFields provides the number of fields without allocating.
+		NumFields() int
 	}
 
 	ListLikeType interface {
@@ -109,6 +111,8 @@ func (t *ListType) ElemField() Field {
 
 func (t *ListType) Fields() []Field { return []Field{t.ElemField()} }
 
+func (t *ListType) NumFields() int { return 1 }
+
 func (*ListType) Layout() DataTypeLayout {
 	return DataTypeLayout{Buffers: []BufferSpec{SpecBitmap(), SpecFixedWidth(Int32SizeBytes)}}
 }
@@ -242,6 +246,8 @@ func (t *FixedSizeListType) Fingerprint() string {
 
 func (t *FixedSizeListType) Fields() []Field { return []Field{t.ElemField()} }
 
+func (t *FixedSizeListType) NumFields() int { return 1 }
+
 func (*FixedSizeListType) Layout() DataTypeLayout {
 	return DataTypeLayout{Buffers: []BufferSpec{SpecBitmap()}}
 }
@@ -308,6 +314,8 @@ func (t *ListViewType) ElemField() Field {
 
 func (t *ListViewType) Fields() []Field { return []Field{t.ElemField()} }
 
+func (t *ListViewType) NumFields() int { return 1 }
+
 func (*ListViewType) Layout() DataTypeLayout {
 	return DataTypeLayout{Buffers: []BufferSpec{SpecBitmap(), SpecFixedWidth(Int32SizeBytes), SpecFixedWidth(Int32SizeBytes)}}
 }
@@ -376,6 +384,8 @@ func (t *LargeListViewType) ElemField() Field {
 
 func (t *LargeListViewType) Fields() []Field { return []Field{t.ElemField()} }
 
+func (t *LargeListViewType) NumFields() int { return 1 }
+
 func (*LargeListViewType) Layout() DataTypeLayout {
 	return DataTypeLayout{Buffers: []BufferSpec{SpecBitmap(), SpecFixedWidth(Int64SizeBytes), SpecFixedWidth(Int64SizeBytes)}}
 }
@@ -447,6 +457,8 @@ func (t *StructType) Fields() []Field {
 	return fields
 }
 
+func (t *StructType) NumFields() int { return len(t.fields) }
+
 func (t *StructType) Field(i int) Field { return t.fields[i] }
 
 // FieldByName gets the field with the given name.
@@ -598,6 +610,8 @@ func (t *MapType) Fingerprint() string {
 
 func (t *MapType) Fields() []Field { return []Field{t.ElemField()} }
 
+func (t *MapType) NumFields() int { return 1 }
+
 func (t *MapType) Layout() DataTypeLayout {
 	return t.value.Layout()
 }
@@ -690,6 +704,8 @@ func (t *unionType) Fields() []Field {
 	return fields
 }
 
+func (t *unionType) NumFields() int { return len(t.children) }
+
 func (t *unionType) TypeCodes() []UnionTypeCode { return t.typeCodes }
 func (t *unionType) ChildIDs() []int            { return t.childIDs[:] }
 

go/arrow/datatype_nested_test.go

@@ -205,7 +205,7 @@ func TestStructOf(t *testing.T) {
 				t.Fatalf("invalid name. got=%q, want=%q", got, want)
 			}
 
-			if got, want := len(got.Fields()), len(tc.fields); got != want {
+			if got, want := got.NumFields(), len(tc.fields); got != want {
 				t.Fatalf("invalid number of fields. got=%d, want=%d", got, want)
 			}
 

go/arrow/flight/flightsql/driver/driver.go

@@ -159,7 +159,7 @@ func (s *Stmt) NumInput() int {
 	// If NumInput returns >= 0, the sql package will sanity check argument
 	// counts from callers and return errors to the caller before the
 	// statement's Exec or Query methods are called.
-	return len(schema.Fields())
+	return schema.NumFields()
 }
 
 // Exec executes a query that doesn't return rows, such

go/arrow/flight/flightsql/example/sql_batch_reader.go

@@ -111,7 +111,7 @@ type SqlBatchReader struct {
 }
 
 func NewSqlBatchReaderWithSchema(mem memory.Allocator, schema *arrow.Schema, rows *sql.Rows) (*SqlBatchReader, error) {
-	rowdest := make([]interface{}, len(schema.Fields()))
+	rowdest := make([]interface{}, schema.NumFields())
 	for i, f := range schema.Fields() {
 		switch f.Type.ID() {
 		case arrow.DENSE_UNION, arrow.SPARSE_UNION:

go/arrow/internal/arrjson/arrjson.go

@@ -1160,7 +1160,7 @@ func arrayFromJSON(mem memory.Allocator, dt arrow.DataType, arr Array) arrow.Arr
 
 		nulls := arr.Count - bitutil.CountSetBits(bitmap.Bytes(), 0, arr.Count)
 
-		fields := make([]arrow.ArrayData, len(dt.Fields()))
+		fields := make([]arrow.ArrayData, dt.NumFields())
 		for i := range fields {
 			child := arrayFromJSON(mem, dt.Field(i).Type, arr.Children[i])
 			defer child.Release()
@@ -1307,7 +1307,7 @@ func arrayFromJSON(mem memory.Allocator, dt arrow.DataType, arr Array) arrow.Arr
 		return array.NewData(dt, arr.Count, []*memory.Buffer{nil}, []arrow.ArrayData{runEnds, values}, 0, 0)
 
 	case arrow.UnionType:
-		fields := make([]arrow.ArrayData, len(dt.Fields()))
+		fields := make([]arrow.ArrayData, dt.NumFields())
 		for i, f := range dt.Fields() {
 			child := arrayFromJSON(mem, f.Type, arr.Children[i])
 			defer child.Release()
@@ -1581,7 +1581,7 @@ func arrayToJSON(field arrow.Field, arr arrow.Array) Array {
 			Name:     field.Name,
 			Count:    arr.Len(),
 			Valids:   validsToJSON(arr),
-			Children: make([]Array, len(dt.Fields())),
+			Children: make([]Array, dt.NumFields()),
 		}
 		for i := range o.Children {
 			o.Children[i] = arrayToJSON(dt.Field(i), arr.Field(i))
@@ -1702,7 +1702,7 @@ func arrayToJSON(field arrow.Field, arr arrow.Array) Array {
 			Count:    arr.Len(),
 			Valids:   validsToJSON(arr),
 			TypeID:   arr.RawTypeCodes(),
-			Children: make([]Array, len(dt.Fields())),
+			Children: make([]Array, dt.NumFields()),
 		}
 		if dt.Mode() == arrow.DenseMode {
 			o.Offset = arr.(*array.DenseUnion).RawValueOffsets()

go/arrow/internal/dictutils/dict.go

@@ -104,7 +104,7 @@ func (d *Mapper) InsertPath(pos FieldPos) {
 	d.hasher.Reset()
 }
 
-func (d *Mapper) ImportField(pos FieldPos, field *arrow.Field) {
+func (d *Mapper) ImportField(pos FieldPos, field arrow.Field) {
 	dt := field.Type
 	if dt.ID() == arrow.EXTENSION {
 		dt = dt.(arrow.ExtensionType).StorageType()
@@ -126,13 +126,18 @@ func (d *Mapper) ImportField(pos FieldPos, field *arrow.Field) {
 
 func (d *Mapper) ImportFields(pos FieldPos, fields []arrow.Field) {
 	for i := range fields {
-		d.ImportField(pos.Child(int32(i)), &fields[i])
+		d.ImportField(pos.Child(int32(i)), fields[i])
 	}
 }
 
 func (d *Mapper) ImportSchema(schema *arrow.Schema) {
 	d.pathToID = make(map[uint64]int64)
-	d.ImportFields(NewFieldPos(), schema.Fields())
+	// This code path intentionally avoids calling ImportFields with
+	// schema.Fields to avoid allocations.
+	pos := NewFieldPos()
+	for i := 0; i < schema.NumFields(); i++ {
+		d.ImportField(pos.Child(int32(i)), schema.Field(i))
+	}
 }
 
 func hasUnresolvedNestedDict(data arrow.ArrayData) bool {