1package stress23import (4 "database/sql"5 "fmt"6 "math/rand/v2"7 "os"8 "path/filepath"9 "time"1011 "github.com/parquet-go/parquet-go"12 _ "modernc.org/sqlite" // registers the "sqlite" driver13)1415// Sources for the benchmarks of linked sources: a Parquet file and a16// SQLite table of sales, rows rows each (ten million for make stress),17// written once into $STRESS_DIR/generated, or a directory the caller18// gives, and kept between runs, as generating them takes a while.1920// SaleRow is a row of the generated sources: an id, one of eight21// categories, an amount and a date. Its columns are A to D of the22// sources' tabs.23type SaleRow struct {24 ID int64 `parquet:"id"`25 Cat string `parquet:"cat,dict"`26 Amount float64 `parquet:"amount"`27 Day int32 `parquet:"day,date"`28}2930// SaleCats are the categories, in the order rows draw them from.31var SaleCats = []string{"north", "south", "east", "west", "central", "coast", "hills", "islands"}3233// Sale is row i of the sources, the same whenever it's asked for.34func Sale(i int) SaleRow {35 r := rand.New(rand.NewPCG(uint64(i), 12))36 return SaleRow{ID: int64(i), Cat: SaleCats[r.IntN(len(SaleCats))], Amount: float64(r.IntN(1_000_000)) / 100,37 Day: int32(time.Date(2020, 1, 1, 0, 0, 0, 0, time.UTC).Unix()/86400) + int32(r.IntN(2000))}38}3940// GeneratedDir is where generated sources are kept: $STRESS_DIR's41// generated folder, or fallback.42func GeneratedDir(fallback string) string {43 if d := os.Getenv("STRESS_DIR"); d != "" {44 dir := filepath.Join(d, "generated")45 if err := os.MkdirAll(dir, 0o755); err == nil {46 return dir47 }48 }49 return fallback50}5152// SalesParquet is the Parquet source of rows rows in dir, written unless53// it's there.54func SalesParquet(dir string, rows int) (string, error) {55 name := filepath.Join(dir, fmt.Sprintf("sales-%d.parquet", rows))56 if _, err := os.Stat(name); err == nil {57 return name, nil58 }59 tmp := name + ".part"60 f, err := os.Create(tmp)61 if err != nil {62 return "", err63 }64 defer f.Close()65 w := parquet.NewGenericWriter[SaleRow](f, parquet.MaxRowsPerRowGroup(1<<20))66 batch := make([]SaleRow, 1<<16)67 for i := 0; i < rows; i += len(batch) {68 n := min(len(batch), rows-i)69 for j := range n {70 batch[j] = Sale(i + j)71 }72 if _, err := w.Write(batch[:n]); err != nil {73 return "", err74 }75 }76 if err := w.Close(); err != nil {77 return "", err78 }79 if err := f.Close(); err != nil {80 return "", err81 }82 return name, os.Rename(tmp, name)83}8485// SalesSQLite is the SQLite source of rows rows in dir, a table named86// sales, written unless it's there.87func SalesSQLite(dir string, rows int) (string, error) {88 name := filepath.Join(dir, fmt.Sprintf("sales-%d.sqlite", rows))89 if _, err := os.Stat(name); err == nil {90 return name, nil91 }92 tmp := name + ".part"93 os.Remove(tmp)94 db, err := sql.Open("sqlite", tmp+"?_pragma=journal_mode(OFF)&_pragma=synchronous(OFF)")95 if err != nil {96 return "", err97 }98 defer db.Close()99 tx, err := db.Begin()100 if err != nil {101 return "", err102 }103 defer tx.Rollback()104 if _, err := tx.Exec("CREATE TABLE sales (id INTEGER, cat TEXT, amount REAL, day DATE)"); err != nil {105 return "", err106 }107 ins, err := tx.Prepare("INSERT INTO sales VALUES (?, ?, ?, ?)")108 if err != nil {109 return "", err110 }111 for i := range rows {112 r := Sale(i)113 day := time.Unix(int64(r.Day)*86400, 0).UTC().Format(time.DateOnly)114 if _, err := ins.Exec(r.ID, r.Cat, r.Amount, day); err != nil {115 return "", err116 }117 }118 if err := tx.Commit(); err != nil {119 return "", err120 }121 if err := db.Close(); err != nil {122 return "", err123 }124 return name, os.Rename(tmp, name)125}