tensor/io/bench_test.go

// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT

package io

// Benchmarks for the read and write paths that carry the per-value
// work: a table decode, a CSV parse, the HDF5 and NetCDF writers and
// readers. Every input is deterministic and written once, before the
// measured loop.

import (
	"encoding/binary"
	"math"
	"os"
	"path/filepath"
	"strconv"
	"testing"

	"sourcedock.dev/petrbalvin/tensor/internal/core"
)

// Sizes: several thousand rows for the tables, and a few hundred
// thousand values for the array formats, which is large enough for the
// per-cell work to dominate the fixed cost of each entry point.
const (
	benchRows = 4000
	benchCols = 24
	benchSide = 512
)

// benchFloats builds a deterministic float64 array.
func benchFloats(shape ...int) *core.Array {
	a := core.New(core.Float, shape...)
	raw := a.RawFloats()
	for i := range raw {
		raw[i] = math.Sin(float64(i)*0.03125)*1000 + float64(i%97)*0.5
	}
	return a
}

// benchInts builds a deterministic int array.
func benchInts(shape ...int) *core.Array {
	a := core.New(core.Int, shape...)
	raw := a.RawInts()
	for i := range raw {
		raw[i] = int64(i)*7 - 3
	}
	return a
}

// benchBinaryTableFile writes a BINTABLE holding one column of every
// numeric form the reader decodes plus a character column, and returns
// its path. The package's own writer emits a subset of those forms
// (B, I and J arrive from other writers), so the table is laid out
// here.
func benchBinaryTableFile(b testing.TB, rows int) string {
	b.Helper()
	forms := []string{"K", "D", "E", "J", "I", "B", "L", "8A"}
	widths := []int{8, 8, 4, 4, 2, 1, 1, 8}
	rowBytes := 0
	for _, w := range widths {
		rowBytes += w
	}
	cards := []string{
		fitsStringCardRaw("XTENSION", "BINTABLE"),
		fitsIntCard("BITPIX", 8),
		fitsIntCard("NAXIS", 2),
		fitsIntCard("NAXIS1", rowBytes),
		fitsIntCard("NAXIS2", rows),
		fitsIntCard("PCOUNT", 0),
		fitsIntCard("GCOUNT", 1),
		fitsIntCard("TFIELDS", len(forms)),
	}
	for i, form := range forms {
		n := strconv.Itoa(i + 1)
		cards = append(cards,
			fitsStringCardRaw("TTYPE"+n, "COL"+n),
			fitsStringCardRaw("TFORM"+n, form))
	}
	cards = append(cards, fitsEndCard())
	out := fitsAppendCards(nil, []string{
		fitsBoolCard("SIMPLE", true),
		fitsIntCard("BITPIX", 8),
		fitsIntCard("NAXIS", 0),
		fitsBoolCard("EXTEND", true),
		fitsEndCard(),
	})
	out = fitsAppendCards(out, cards)
	body := make([]byte, rows*rowBytes)
	star := []byte("star    ")
	for r := range rows {
		p := r * rowBytes
		binary.BigEndian.PutUint64(body[p:], uint64(1000+r))
		p += 8
		binary.BigEndian.PutUint64(body[p:], math.Float64bits(float64(r)*0.25-1))
		p += 8
		binary.BigEndian.PutUint32(body[p:], math.Float32bits(float32(r)*0.5))
		p += 4
		binary.BigEndian.PutUint32(body[p:], uint32(int32(-r)))
		p += 4
		binary.BigEndian.PutUint16(body[p:], uint16(int16(r)))
		p += 2
		body[p] = byte(r)
		p++
		if r%2 == 0 {
			body[p] = 'T'
		} else {
			body[p] = 'F'
		}
		p++
		copy(body[p:], star)
	}
	out = append(out, body...)
	out = fitsAppendZeroPad(out)
	path := filepath.Join(b.TempDir(), "binary.fits")
	if err := os.WriteFile(path, out, 0o644); err != nil {
		b.Fatal(err)
	}
	return path
}

// benchASCIITableFile writes an ASCII table with a character, an
// integer and a float column, and returns its path.
func benchASCIITableFile(b testing.TB, rows int) string {
	b.Helper()
	text := make([]string, rows)
	ints := core.New(core.Int, rows)
	floats := core.New(core.Float, rows)
	for i := range rows {
		text[i] = "star-" + strconv.Itoa(i)
		ints.RawInts()[i] = int64(1000 + i)
		floats.RawFloats()[i] = float64(i)*0.125 - 42
	}
	cols := []FITSTableColumn{
		{Name: "STAR", Form: "12A", Text: text},
		{Name: "ID", Form: "I10", Data: ints},
		{Name: "MAG", Form: "D20.12", Data: floats},
	}
	path := filepath.Join(b.TempDir(), "ascii.fits")
	if err := SaveFITSTable(path, true, cols, nil); err != nil {
		b.Fatal(err)
	}
	return path
}

// benchCSVFile writes a float64 matrix as CSV and returns its path.
func benchCSVFile(b testing.TB, rows, cols int) string {
	b.Helper()
	path := filepath.Join(b.TempDir(), "bench.csv")
	if err := SaveCSV(path, benchFloats(rows, cols)); err != nil {
		b.Fatal(err)
	}
	return path
}

// benchHDF5File writes a classic HDF5 file of one float64 and one
// int64 dataset, optionally through the deflate and shuffle filters,
// and returns its path.
func benchHDF5File(b testing.TB, filtered bool) string {
	b.Helper()
	sets := []HDF5Dataset{
		{Path: "/field", Values: benchFloats(benchSide, benchSide)},
		{Path: "/ids", Values: benchInts(benchSide, benchSide)},
	}
	attrs := map[string]map[string]string{"/": {"origin": "benchmark"}}
	opts := HDF5WriteOptions{}
	if filtered {
		opts = HDF5WriteOptions{Gzip: 6, Shuffle: true}
	}
	path := filepath.Join(b.TempDir(), "bench.h5")
	if err := SaveHDF5(path, sets, attrs, opts); err != nil {
		b.Fatal(err)
	}
	return path
}

// benchNetCDFFile writes a classic NetCDF file of one float64 variable
// in two dimensions and returns its path.
func benchNetCDFFile(b testing.TB) string {
	b.Helper()
	dims := []NetCDFDim{{Name: "row", Length: benchSide}, {Name: "col", Length: benchSide}}
	vars := []NetCDFVar{{Name: "field", Dims: []string{"row", "col"}, Values: benchFloats(benchSide, benchSide)}}
	path := filepath.Join(b.TempDir(), "bench.nc")
	if err := SaveNetCDF(path, dims, vars, map[string]string{"title": "benchmark"}); err != nil {
		b.Fatal(err)
	}
	return path
}

func BenchmarkLoadFITSTableBinary(b *testing.B) {
	path := benchBinaryTableFile(b, benchRows)
	b.ReportAllocs()
	for b.Loop() {
		table, err := LoadFITSTable(path)
		if err != nil {
			b.Fatal(err)
		}
		if table.Rows != benchRows || len(table.Columns) != 8 || table.Text[7] == nil {
			b.Fatalf("table came back as %s with %d columns", table.Kind, len(table.Columns))
		}
	}
}

func BenchmarkLoadFITSTableASCII(b *testing.B) {
	path := benchASCIITableFile(b, benchRows)
	b.ReportAllocs()
	for b.Loop() {
		table, err := LoadFITSTable(path)
		if err != nil {
			b.Fatal(err)
		}
		if table.Rows != benchRows || len(table.Columns) != 3 {
			b.Fatalf("table came back with %d rows and %d columns", table.Rows, len(table.Columns))
		}
	}
}

func BenchmarkLoadCSV(b *testing.B) {
	path := benchCSVFile(b, benchRows, benchCols)
	b.ReportAllocs()
	for b.Loop() {
		a, err := LoadCSV(path, false)
		if err != nil {
			b.Fatal(err)
		}
		if a.Shape()[0] != benchRows || a.Shape()[1] != benchCols {
			b.Fatalf("array came back with shape %v", a.Shape())
		}
	}
}

func BenchmarkSaveHDF5(b *testing.B) {
	sets := []HDF5Dataset{
		{Path: "/field", Values: benchFloats(benchSide, benchSide)},
		{Path: "/ids", Values: benchInts(benchSide, benchSide)},
	}
	attrs := map[string]map[string]string{"/": {"origin": "benchmark"}}
	path := filepath.Join(b.TempDir(), "bench.h5")
	b.ReportAllocs()
	b.SetBytes(int64(2 * benchSide * benchSide * 8))
	for b.Loop() {
		if err := SaveHDF5(path, sets, attrs); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkSaveHDF5Large writes one 8 MiB float64 field, the size
// where moving the payload around shows above the per-value work.
func BenchmarkSaveHDF5Large(b *testing.B) {
	sets := []HDF5Dataset{{Path: "/field", Values: benchFloats(1024, 1024)}}
	path := filepath.Join(b.TempDir(), "bench.h5")
	b.ReportAllocs()
	b.SetBytes(8 << 20)
	for b.Loop() {
		if err := SaveHDF5(path, sets, nil); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkSaveHDF5Filtered writes the two benchmark datasets through
// the shuffle and deflate filters, the chunked path.
func BenchmarkSaveHDF5Filtered(b *testing.B) {
	sets := []HDF5Dataset{
		{Path: "/field", Values: benchFloats(benchSide, benchSide)},
		{Path: "/ids", Values: benchInts(benchSide, benchSide)},
	}
	path := filepath.Join(b.TempDir(), "bench.h5")
	b.ReportAllocs()
	b.SetBytes(int64(2 * benchSide * benchSide * 8))
	for b.Loop() {
		if err := SaveHDF5(path, sets, nil, HDF5WriteOptions{Gzip: 6, Shuffle: true}); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkSaveHDF5Text writes one string dataset, the fixed-width
// text path whose elements are padded to the longest of the file.
func BenchmarkSaveHDF5Text(b *testing.B) {
	const side = 512
	text := make([]string, side*side)
	width := 0
	for i := range text {
		text[i] = "row" + strconv.Itoa(i)
		width = max(width, len(text[i]))
	}
	sets := []HDF5TextDataset{{Path: "/labels", Shape: []int{side, side}, Text: text}}
	path := filepath.Join(b.TempDir(), "bench.h5")
	b.ReportAllocs()
	b.SetBytes(int64(side * side * width))
	for b.Loop() {
		if err := SaveHDF5Text(path, sets); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkLoadHDF5(b *testing.B) {
	path := benchHDF5File(b, false)
	b.ReportAllocs()
	for b.Loop() {
		sets, err := LoadHDF5(path)
		if err != nil {
			b.Fatal(err)
		}
		if len(sets) != 2 {
			b.Fatalf("file came back with %d datasets", len(sets))
		}
	}
}

func BenchmarkLoadHDF5Filtered(b *testing.B) {
	path := benchHDF5File(b, true)
	b.ReportAllocs()
	for b.Loop() {
		sets, err := LoadHDF5(path)
		if err != nil {
			b.Fatal(err)
		}
		if len(sets) != 2 {
			b.Fatalf("file came back with %d datasets", len(sets))
		}
	}
}

func BenchmarkSaveNetCDF(b *testing.B) {
	dims := []NetCDFDim{{Name: "row", Length: benchSide}, {Name: "col", Length: benchSide}}
	vars := []NetCDFVar{{Name: "field", Dims: []string{"row", "col"}, Values: benchFloats(benchSide, benchSide)}}
	attrs := map[string]string{"title": "benchmark"}
	path := filepath.Join(b.TempDir(), "bench.nc")
	b.ReportAllocs()
	for b.Loop() {
		if err := SaveNetCDF(path, dims, vars, attrs); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkLoadNetCDF(b *testing.B) {
	path := benchNetCDFFile(b)
	b.ReportAllocs()
	for b.Loop() {
		_, vars, _, err := LoadNetCDF(path)
		if err != nil {
			b.Fatal(err)
		}
		if len(vars) != 1 {
			b.Fatalf("file came back with %d variables", len(vars))
		}
	}
}