perf: wire putVectorBuf into connection write path

Return pooled vector buffers to vectorBufPool after the framer copies
marshalled bytes in executeQuery and executeBatch. This completes the
zero-alloc steady-state cycle for vector marshal operations.

In executeQuery, a defer after the marshal loop returns buffers for
columns identified as pooled vector types (float32, float64, int32,
int64). In executeBatch, vector buffers are collected across all batch
statements and returned via a single defer.

The vectorBufPoolSubtype helper centralizes the type check to keep
the two call sites consistent with the marshal fast paths.

Includes unit tests covering vectorBufPoolSubtype classification,
single-query and batch pool return simulation, and non-pooled type
safety.

Author: mykaul

conn.go

@@ -1574,6 +1574,29 @@ func (c *Conn) executeQuery(ctx context.Context, qry *Query) (iter *Iter) {
 			}
 		}
 
+		// Return pooled vector buffers after the framer has copied the
+		// marshalled bytes (which happens inside c.exec → buildFrame).
+		// Only install the defer when at least one column uses a poolable
+		// vector buffer to avoid ~50ns defer overhead on non-vector queries.
+		cols := info.request.columns
+		vals := params.values
+		hasPooledVec := false
+		for _, col := range cols {
+			if vt, ok := col.TypeInfo.(VectorType); ok && vectorBufPoolSubtype(vt) {
+				hasPooledVec = true
+				break
+			}
+		}
+		if hasPooledVec {
+			defer func() {
+				for i, col := range cols {
+					if vt, ok := col.TypeInfo.(VectorType); ok && vectorBufPoolSubtype(vt) {
+						putVectorBuf(vals[i].value)
+					}
+				}
+			}()
+		}
+
 		// if the metadata was not present in the response then we should not skip it
 		params.skipMeta = !(c.session.cfg.DisableSkipMetadata || qry.disableSkipMetadata) && len(info.response.columns) != 0
 
@@ -1761,6 +1784,10 @@ func (c *Conn) executeBatch(ctx context.Context, batch *Batch) (iter *Iter) {
 
 	hasLwtEntries := false
 
+	// vectorBufs collects marshalled byte slices from vector fast paths
+	// so they can be returned to vectorBufPool after the framer copies them.
+	var vectorBufs [][]byte
+
 	for i := 0; i < n; i++ {
 		entry := &batch.Entries[i]
 		b := &req.statements[i]
@@ -1802,6 +1829,9 @@ func (c *Conn) executeBatch(ctx context.Context, batch *Batch) (iter *Iter) {
 				if err := marshalQueryValue(typ, value, v); err != nil {
 					return &Iter{err: err}
 				}
+				if vt, ok := typ.(VectorType); ok && vectorBufPoolSubtype(vt) {
+					vectorBufs = append(vectorBufs, v.value)
+				}
 			}
 
 			if !hasLwtEntries && info.request.lwt {
@@ -1812,6 +1842,16 @@ func (c *Conn) executeBatch(ctx context.Context, batch *Batch) (iter *Iter) {
 		}
 	}
 
+	// Return pooled vector buffers after the framer has copied the
+	// marshalled bytes (which happens inside c.exec → buildFrame).
+	if len(vectorBufs) > 0 {
+		defer func() {
+			for _, buf := range vectorBufs {
+				putVectorBuf(buf)
+			}
+		}()
+	}
+
 	// The batch is considered to be conditional if even one of the
 	// statements is conditional.
 	batch.routingInfo.mu.Lock()

marshal.go

@@ -1013,9 +1013,10 @@ func unmarshalVector(info VectorType, data []byte, value interface{}) error {
 // vectorBufPool pools []byte buffers used by vector marshal fast paths.
 // Buffers are returned to the pool by putVectorBuf after the framer copies them.
 //
-// NOTE: putVectorBuf is currently exercised only by benchmarks/tests.
-// Wiring it into the connection write path (so production callers return
-// buffers after the framer copies them) is planned for a follow-up change.
+// vectorBufPool is used by the marshal fast paths (marshalVectorFloat32 etc.)
+// to allocate temporary byte buffers. After the framer copies the marshalled
+// bytes via writeBytes (append into framer.buf), the caller returns the buffer
+// to the pool via putVectorBuf.
 var vectorBufPool = sync.Pool{}
 
 func getVectorBuf(size int) []byte {
@@ -1044,6 +1045,17 @@ func putVectorBuf(buf []byte) {
 	vectorBufPool.Put(buf) //nolint:staticcheck // SA6002: []byte is a value type; boxing cost is acceptable for pool reuse
 }
 
+// vectorBufPoolSubtype returns true if the given VectorType's SubType uses a
+// marshal fast path that allocates from vectorBufPool via getVectorBuf.
+func vectorBufPoolSubtype(vt VectorType) bool {
+	switch vt.SubType.Type() {
+	case TypeFloat, TypeDouble, TypeInt, TypeBigInt:
+		return true
+	default:
+		return false
+	}
+}
+
 // vectorByteSize computes dim * elemBytes and returns an error if the result
 // would overflow int. This guards against corrupt or adversarial schema metadata
 // on 32-bit platforms where int is 32 bits.
@@ -1072,10 +1084,12 @@ func marshalVectorFloat32(dim int, vec []float32) ([]byte, error) {
 		return nil, marshalErrorf("%v", err)
 	}
 	buf := getVectorBuf(size)
-	off := 0
-	for _, v := range vec {
-		binary.BigEndian.PutUint32(buf[off:off+4], math.Float32bits(v))
-		off += 4
+	if dim == 0 {
+		return buf, nil
+	}
+	_ = buf[dim*4-1] // BCE hint
+	for i, v := range vec {
+		binary.BigEndian.PutUint32(buf[i*4:i*4+4], math.Float32bits(v))
 	}
 	return buf, nil
 }
@@ -1097,10 +1111,12 @@ func marshalVectorFloat64(dim int, vec []float64) ([]byte, error) {
 		return nil, marshalErrorf("%v", err)
 	}
 	buf := getVectorBuf(size)
-	off := 0
-	for _, v := range vec {
-		binary.BigEndian.PutUint64(buf[off:off+8], math.Float64bits(v))
-		off += 8
+	if dim == 0 {
+		return buf, nil
+	}
+	_ = buf[dim*8-1] // BCE hint
+	for i, v := range vec {
+		binary.BigEndian.PutUint64(buf[i*8:i*8+8], math.Float64bits(v))
 	}
 	return buf, nil
 }
@@ -1136,9 +1152,9 @@ func unmarshalVectorFloat32(dim int, data []byte, dst *[]float32) error {
 	} else {
 		vec = make([]float32, dim)
 	}
+	_ = data[dim*4-1] // BCE hint: compiler can prove data[i*4:i*4+4] is in-bounds
 	for i := range vec {
-		vec[i] = math.Float32frombits(binary.BigEndian.Uint32(data[:4]))
-		data = data[4:]
+		vec[i] = math.Float32frombits(binary.BigEndian.Uint32(data[i*4 : i*4+4]))
 	}
 	*dst = vec
 	return nil
@@ -1175,9 +1191,9 @@ func unmarshalVectorFloat64(dim int, data []byte, dst *[]float64) error {
 	} else {
 		vec = make([]float64, dim)
 	}
+	_ = data[dim*8-1] // BCE hint: compiler can prove data[i*8:i*8+8] is in-bounds
 	for i := range vec {
-		vec[i] = math.Float64frombits(binary.BigEndian.Uint64(data[:8]))
-		data = data[8:]
+		vec[i] = math.Float64frombits(binary.BigEndian.Uint64(data[i*8 : i*8+8]))
 	}
 	*dst = vec
 	return nil
@@ -1200,10 +1216,12 @@ func marshalVectorInt32(dim int, vec []int32) ([]byte, error) {
 		return nil, marshalErrorf("%v", err)
 	}
 	buf := getVectorBuf(size)
-	off := 0
-	for _, v := range vec {
-		binary.BigEndian.PutUint32(buf[off:off+4], uint32(v))
-		off += 4
+	if dim == 0 {
+		return buf, nil
+	}
+	_ = buf[dim*4-1] // BCE hint
+	for i, v := range vec {
+		binary.BigEndian.PutUint32(buf[i*4:i*4+4], uint32(v))
 	}
 	return buf, nil
 }
@@ -1225,10 +1243,12 @@ func marshalVectorInt64(dim int, vec []int64) ([]byte, error) {
 		return nil, marshalErrorf("%v", err)
 	}
 	buf := getVectorBuf(size)
-	off := 0
-	for _, v := range vec {
-		binary.BigEndian.PutUint64(buf[off:off+8], uint64(v))
-		off += 8
+	if dim == 0 {
+		return buf, nil
+	}
+	_ = buf[dim*8-1] // BCE hint
+	for i, v := range vec {
+		binary.BigEndian.PutUint64(buf[i*8:i*8+8], uint64(v))
 	}
 	return buf, nil
 }
@@ -1264,9 +1284,9 @@ func unmarshalVectorInt32(dim int, data []byte, dst *[]int32) error {
 	} else {
 		vec = make([]int32, dim)
 	}
+	_ = data[dim*4-1] // BCE hint: compiler can prove data[i*4:i*4+4] is in-bounds
 	for i := range vec {
-		vec[i] = int32(binary.BigEndian.Uint32(data[:4]))
-		data = data[4:]
+		vec[i] = int32(binary.BigEndian.Uint32(data[i*4 : i*4+4]))
 	}
 	*dst = vec
 	return nil
@@ -1303,9 +1323,9 @@ func unmarshalVectorInt64(dim int, data []byte, dst *[]int64) error {
 	} else {
 		vec = make([]int64, dim)
 	}
+	_ = data[dim*8-1] // BCE hint: compiler can prove data[i*8:i*8+8] is in-bounds
 	for i := range vec {
-		vec[i] = int64(binary.BigEndian.Uint64(data[:8]))
-		data = data[8:]
+		vec[i] = int64(binary.BigEndian.Uint64(data[i*8 : i*8+8]))
 	}
 	*dst = vec
 	return nil

marshal_vector_test.go

@@ -1491,3 +1491,153 @@ func TestGetVectorBuf_NonPositiveSize(t *testing.T) {
 		}
 	}
 }
+
+// TestVectorBufPoolSubtype verifies that vectorBufPoolSubtype correctly
+// identifies which vector subtypes use the pooled fast path.
+func TestVectorBufPoolSubtype(t *testing.T) {
+	tests := []struct {
+		name   string
+		vt     VectorType
+		expect bool
+	}{
+		{"float32", makeFloat32VectorType(3), true},
+		{"float64", makeFloat64VectorType(3), true},
+		{"int32", makeInt32VectorType(3), true},
+		{"int64", makeInt64VectorType(3), true},
+		{"uuid", makeUUIDVectorType(3), false},
+		{"varchar", VectorType{
+			SubType:    NativeType{proto: protoVersion4, typ: TypeVarchar},
+			Dimensions: 3,
+		}, false},
+		{"boolean", VectorType{
+			SubType:    NativeType{proto: protoVersion4, typ: TypeBoolean},
+			Dimensions: 3,
+		}, false},
+		{"timestamp", VectorType{
+			SubType:    NativeType{proto: protoVersion4, typ: TypeTimestamp},
+			Dimensions: 3,
+		}, false},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := vectorBufPoolSubtype(tt.vt)
+			if got != tt.expect {
+				t.Errorf("vectorBufPoolSubtype(%s) = %v, want %v", tt.name, got, tt.expect)
+			}
+		})
+	}
+}
+
+// TestVectorBufPoolReturnSimulation simulates the buffer lifecycle used in
+// executeQuery and executeBatch: marshal a vector value via marshalQueryValue,
+// then return the buffer to the pool via putVectorBuf. Verifies that the
+// returned buffer is reused by a subsequent getVectorBuf call.
+func TestVectorBufPoolReturnSimulation(t *testing.T) {
+	const dim = 128
+	vt := makeFloat32VectorType(dim)
+
+	// Create test data.
+	vec := make([]float32, dim)
+	for i := range vec {
+		vec[i] = float32(i)
+	}
+
+	// Marshal like marshalQueryValue does.
+	data, err := Marshal(vt, vec)
+	if err != nil {
+		t.Fatalf("Marshal: %v", err)
+	}
+	expectedSize := dim * 4
+	if len(data) != expectedSize {
+		t.Fatalf("marshalled size = %d, want %d", len(data), expectedSize)
+	}
+
+	// Remember the backing array pointer before returning to pool.
+	origPtr := &data[:1][0]
+
+	// Simulate what the defer in executeQuery/executeBatch does.
+	if vectorBufPoolSubtype(vt) {
+		putVectorBuf(data)
+	}
+
+	// Get a buffer of the same size — should reuse the pooled one.
+	reused := getVectorBuf(expectedSize)
+	if len(reused) != expectedSize {
+		t.Fatalf("getVectorBuf(%d) returned len %d", expectedSize, len(reused))
+	}
+	reusedPtr := &reused[:1][0]
+	if origPtr != reusedPtr {
+		t.Error("expected pooled buffer to be reused, but got a different allocation")
+	}
+
+	// Clean up.
+	putVectorBuf(reused)
+}
+
+// TestVectorBufPoolReturnNonPooledType verifies that putVectorBuf is a no-op
+// for vector types that don't use the fast path (e.g. UUID vectors), and
+// that vectorBufPoolSubtype correctly excludes them.
+func TestVectorBufPoolReturnNonPooledType(t *testing.T) {
+	const dim = 4
+	vt := makeUUIDVectorType(dim)
+
+	if vectorBufPoolSubtype(vt) {
+		t.Fatal("vectorBufPoolSubtype should be false for UUID vectors")
+	}
+
+	// The UUID marshal path does not use getVectorBuf, so putting its
+	// result back should be skipped by the type check. Verify no panic.
+	data, err := Marshal(vt, []UUID{
+		{0x01}, {0x02}, {0x03}, {0x04},
+	})
+	if err != nil {
+		t.Fatalf("Marshal UUID vector: %v", err)
+	}
+	// Even if we mistakenly call putVectorBuf, it should not panic.
+	putVectorBuf(data)
+}
+
+// TestVectorBufPoolBatchSimulation simulates the batch buffer lifecycle:
+// multiple statements with vector columns get their buffers collected and
+// returned after the framer copies them.
+func TestVectorBufPoolBatchSimulation(t *testing.T) {
+	const dim = 64
+
+	types := []struct {
+		vt  VectorType
+		val interface{}
+	}{
+		{makeFloat32VectorType(dim), make([]float32, dim)},
+		{makeFloat64VectorType(dim), make([]float64, dim)},
+		{makeInt32VectorType(dim), make([]int32, dim)},
+		{makeInt64VectorType(dim), make([]int64, dim)},
+	}
+
+	// Simulate batch: marshal all, collect buffers.
+	var vectorBufs [][]byte
+	for _, tt := range types {
+		data, err := Marshal(tt.vt, tt.val)
+		if err != nil {
+			t.Fatalf("Marshal %v: %v", tt.vt.SubType.Type(), err)
+		}
+		if vectorBufPoolSubtype(tt.vt) {
+			vectorBufs = append(vectorBufs, data)
+		}
+	}
+
+	if len(vectorBufs) != 4 {
+		t.Fatalf("expected 4 pooled buffers, got %d", len(vectorBufs))
+	}
+
+	// Return all to pool (like the defer in executeBatch).
+	for _, buf := range vectorBufs {
+		putVectorBuf(buf)
+	}
+
+	// Verify at least one can be reused.
+	reused := getVectorBuf(dim * 4) // float32 size
+	if reused == nil {
+		t.Fatal("expected to get a buffer from pool after returning batch buffers")
+	}
+	putVectorBuf(reused)
+}