Unnesting a slice of structures into query params: best practices?

[wchargin]

Hello! I am quite happy with the ergonomics of the collectors like pgx.RowToAddrOfStructByName for mapping columns from the result set of a query onto columns in a struct. Is there a nice facility for the dual operation, writing a slice of structures to query parameters?

In particular, suppose that I have a query like:

INSERT INTO users(id, name, email)
SELECT id, name, email
FROM unnest($1::int8[], $2::text[], $3::text[]) q(id, name, email)
-- or VALUES(unnest($1::int8[]), unnest($2[]::text[]), unnest($3::text[]))
-- or many other examples that use this idea

…and I have a slice users []*User, with the following structure definition:

type User struct {
	ID    int64  `db:"user"`
	Name  string `db:"name"`
	Email string `db:"email"`
}

The only way that I know to get that data into that query is to unroll it into separate slices, to match the query parameters:

ids := make([]int64, len(users))
names := make([]string, len(users))
emails := make([]string, len(users))
for i, user := range users {
	ids[i] = user.ID
	names[i] = user.Name
	emails[i] = user.Email
}
_, err := conn.Exec(ctx, sqlInsert, ids, names, emails)

But this isn't super ergonomic. The collectors do a nice SOA-to-AOS transform on the read side. Does pgx endorse any pattern with similar ergonomics to do the AOS-to-SOA transform on the write side?

(Thank you for your work on pgx! I've so far been delighted with its architecture and layering. It's had all the extension points that I needed and has been very reliable.)

[wchargin]

Here is a kind of weird little API that I just sketched out that might help with this kind of thing…

package soa

// Call Collect with a slice of items, and provide a callback that calls Emit
// once for each column of interest. Every item should make the same number of
// calls to Emit, emitting fields in the same order each time.
//
// The return value parameter is a slice of slices, each of length len(items).
// If each invocation of the callback made N calls to Emit, then the result
// will have N slices. If the ith lexical call to Emit had type parameter V,
// then result[i] will have runtime type []V.
//
// For example, if the callback starts by calling `Emit(wr, item.SomeString)`,
// then the first value of the result will be a []string that has the values of
// the SomeString field for each item.
func Collect[T any](items []T, process func(wr *Writer, item T)) []any

// Emit the next column for the current item.
func Emit[V any](wr *Writer, value V)

type Writer struct { /* unexported fields */ }

Sample usage (main.go)

package main

import (
	"fmt"

	"example.com/m/soa"
)

type User struct {
	ID    int
	Name  string
	Email string
}

func main() {
	users := []*User{
		{ID: 123, Name: "Alice", Email: "[email protected]"},
		{ID: 456, Name: "Bob", Email: "[email protected]"},
	}
	args := soa.Collect(users, func(wr *soa.Writer, user *User) {
		soa.Emit(wr, user.ID)
		soa.Emit(wr, user.Name)
		soa.Emit(wr, user.Email)
	})
	for i, slice := range args {
		fmt.Printf("args[%d] is a %T: %v\n", i, slice, slice)
	}
	// then you could call conn.Query(ctx, sql, args...)
}

// Prints:
// args[0] is a []int: [123 456]
// args[1] is a []string: [Alice Bob]
// args[2] is a []string: [[email protected] [email protected]]

Implementation (soa/soa.go)

package soa

type Writer struct {
	firstPass bool
	sinks     []any
	sinkIdx   int
}

func Collect[T any](items []T, process func(wr *Writer, item T)) []any {
	if len(items) == 0 {
		return nil
	}

	wr := &Writer{firstPass: true}
	process(wr, items[0])
	expected := wr.sinkIdx

	wr.firstPass = false
	for _, item := range items[1:] {
		wr.sinkIdx = 0
		process(wr, item)
		if wr.sinkIdx != expected {
			panic("calls to Emit differed across items")
		}
	}

	return wr.sinks
}

func Emit[V any](wr *Writer, value V) {
	if wr.firstPass {
		slice := ([]V)(nil)
		wr.sinks = append(wr.sinks, slice)
	}
	sinkPtr := &wr.sinks[wr.sinkIdx]
	sink, ok := (*sinkPtr).([]V)
	if !ok {
		panic("heap pollution: mismatched SOA types")
	}
	wr.sinkIdx++
	*sinkPtr = append(sink, value)
}

It's a bit weird because of the Golang restriction that methods can't take additional generic parameters. And it still does require manually defining a callback that calls Emit once per field. But it does save roughly half the typing: no need to define the output slices or manage indices, and you get an []any that you can spread like args... into the query parameters. It doesn't use reflection, and it's flexible enough to use with probably any query no matter how exactly it uses its arguments.

I haven't tried to use this in real code yet. If I do, I'll try to remember to report back how it feels!

[wchargin]

It's a bit weird because of the Golang restriction that methods can't take additional generic parameters. And it still does require manually defining a callback that calls Emit once per field.

Ooh… and you can solve both of these by binding the type parameter to a helper struct that runs a CPS transform, then hiding it behind an interface to erase the parameter:

func main() {
	users := []*User{
		{ID: 123, Name: "Alice", Email: "[email protected]"},
		{ID: 456, Name: "Bob", Email: "[email protected]"},
	}
	columns, ok := soa.Collect(users, func(u *User) soa.Fields {
		return soa.Fields{soa.V(u.ID), soa.V(u.Name), soa.V(u.Email)}
	})
	fmt.Printf("ok: %v\n", ok)
	for i, slice := range columns {
		fmt.Printf("columns[%d] is a %T: %v\n", i, slice, slice)
	}
	// then you could call conn.Query(ctx, sql, columns...)
}

You still have to wrap each field in soa.V, and this adds some extra interface indirection cost. But it makes the client call a one-liner and hides more of the implementation details.

Implementation v2 (soa/soa.go)

package soa

type Fields []Field

// Collect converts a slice of items into a slice of columns. Provide a
// callback that returns a list of fields of interest for a given item.
// Collect will return a slice of columns, each of length len(items), where
// columns[i][j] is the ith field of items[j]. Fields can be constructed by
// calling `V`.
//
// Across invocations of the callback, it should always return the same number
// of fields, and the field at any given index should always have the same
// static type parameter. Collect will panic if this is not the case.
//
// If the field at index i had type parameter E, then result[i] will have
// runtime type []E. For example, if `project(item)` returns a slice whose
// first element is `V(item.SomeString)`, then result[0] will be a []string
// that has the values of the SomeString field for each item.
//
// If items is empty, then columns will be nil and ok will be false, because
// there is no way invoke the callback to learn what the columns should be.
func Collect[T any](items []T, project func(item T) Fields) (columns []any, ok bool) {
	if len(items) == 0 {
		return nil, false
	}

	wr := &writer{firstPass: true, itemsLen: len(items)}
	fields := project(items[0])
	wr.sinks = make([]any, len(fields))
	for _, f := range fields {
		f.emit(wr)
	}
	wr.firstPass = false
	wr.itemIndex++

	for _, item := range items[1:] {
		wr.sinkIndex = 0
		fields = project(item)
		if len(fields) != len(wr.sinks) {
			panic(reasonFieldCount)
		}
		for _, f := range project(item) {
			f.emit(wr)
		}
		wr.itemIndex++
	}

	return wr.sinks, true
}

// A Field is returned by V. See Collect for more details.
type Field interface {
	emit(wr *writer)
}

type field[E any] struct {
	value E
}

// V bundles a field with its static type parameter. V should be called by the
// callback passed to Collect.
func V[E any](value E) Field {
	return field[E]{value: value}
}

func (f field[E]) emit(wr *writer) {
	var sink []E
	if wr.firstPass {
		sink = make([]E, wr.itemsLen)
		wr.sinks[wr.sinkIndex] = sink
	} else {
		var ok bool
		sink, ok = wr.sinks[wr.sinkIndex].([]E)
		if !ok {
			panic(reasonFieldTypes)
		}
	}
	sink[wr.itemIndex] = f.value
	wr.sinkIndex++
}

type writer struct {
	sinks     []any
	sinkIndex int
	firstPass bool
	itemsLen  int
	itemIndex int
}

type panicReason string

const (
	reasonFieldCount panicReason = "soa: inconsistent number of fields"
	reasonFieldTypes             = "soa: inconsistent type parameter for field"
)

[jackc]

pgx does not have anything builtin to do the type of conversion you are trying. But I would probably try a few other paths for multiple insert before the approach above.

1. CopyFrom would give the highest performance for multiple insert.
2. Use a batch and insert one row per query. This has the advantage of making it easier to handle errors because you know exactly which row caused the problem.
3. I would look into using an array of composites directly as the parameter type.