Scan of a large matrix performance

[ivaaaan]

Hey. In our application, we scan quite large 2D arrays from the database. We noticed that our application is quite slow, and I made a few benchmarks:

package main

var db *pgxpool.Pool

func TestMain(m *testing.M) {
	// Setup
	pool, err := dockertest.NewPool("")
	if err != nil {
		log.Fatalf("Could not connect to docker: %s", err)
	}

	resource, err := pool.Run("postgres", "16", []string{"POSTGRES_PASSWORD=secret"})
	if err != nil {
		log.Fatalf("Could not start resource: %s", err)
	}

	if err := pool.Retry(func() error {
		var err error
		db, err = pgxpool.New(context.TODO(), fmt.Sprintf("postgres://postgres:secret@localhost:%s/postgres?sslmode=disable", resource.GetPort("5432/tcp")))
		if err != nil {
			return err
		}
		return db.Ping(context.TODO())
	}); err != nil {
		log.Fatalf("Could not connect to docker: %s", err)
	}

	code := m.Run()

	if err := pool.Purge(resource); err != nil {
		log.Fatalf("Could not purge resource: %s", err)
	}

	os.Exit(code)
}

func BenchmarkDecode2D(b *testing.B) {
	for i := 0; i < b.N; i++ {
		rows, _ := db.Query(context.Background(), `SELECT ARRAY(SELECT ARRAY[1::float, 2::float] FROM generate_series(1, 9000))`)

		var matrix [][]float64

		if rows.Next() {
			err := rows.Scan(&matrix)
			if err != nil {
				b.Fatalf("Scan failed: %v", err)
			}
		}
		rows.Close()
	}
}

func BenchmarkDecode1D(b *testing.B) {
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		rows, _ := db.Query(context.Background(), `SELECT ARRAY(SELECT n::float FROM generate_series(1, 1000) n)`)

		var array []float64

		if rows.Next() {
			err := rows.Scan(&array)
			if err != nil {
				b.Fatalf("Scan failed: %v", err)
			}
		}
		rows.Close()
	}
}

Results:

BenchmarkDecode2D-12                 469           2522202 ns/op          592071 B/op       9014 allocs/op
BenchmarkDecode1D-12                6483            183489 ns/op            8766 B/op          9 allocs/op

I'm wondering if there are native ways to improve the scanning of 2D arrays, or if it's better to implement a custom codec for this?

[ivaaaan]

Another idea I have is to unnest array - SELECT array_agg(x::float) FROM unnest(ARRAY(SELECT ARRAY[1::float, 2::float] FROM generate_series(1, 9000))) and return it as 1D to the client. They can figure out dimensions then by themselves

[jackc]

Multi-dimensional array scanning into a slice of slices uses a reflection path. As you found, that is much slower. Try using pgtype.Array[T] instead.

[ivaaaan]

Thanks. Do you have an example of how to use it properly? I'm getting cannot scan _float8 (OID 1022) in binary format into *pgtype.Array[[]float64] with pgtype.Array[[]float64]{}. Or did you mean pg.Array[T] with unnest?

[jackc]

pgtype.Array[float64] is what you would want to use. You would need look at its Elements and Dims fields to determine its exact dimensions.

[ivaaaan]

Thanks a lot, this works!