feat: initial release

Assisted-by: GLM 5.3 Flash

io/bench_test.go

@@ -0,0 +1,353 @@
+// 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))
+		}
+	}
+}