reduce CGO calls when scanning rows

charlievieth · charlievieth · commit e1a8f32f7e0e · 2024-12-17T12:32:44.000-05:00
This commit improves the performance of queries by at least 20% by only
collecting the information needed to convert sqlite3 values to Go values
and by batching CGO calls when scanning rows (the performance
improvement scales with the number of columns being scanned).

This commit adds a new coltype field to the SQLiteRows struct which
stores the declared column type (either data/time or boolean) and the
sqlite3 datatype. Previously, this library would fetch the string
representation of each column, which is inefficient and rarely needed
since the non-standard SQLiteRows.DeclTypes method is rarely called.

It also changes the benchmark suite to use an in-memory database since
we do not want the file system interfering with benchmark results.

goos: darwin
goarch: arm64
pkg: github.com/mattn/go-sqlite3
cpu: Apple M1 Max
                             │ base.10.txt  │             new.10.txt              │
                             │    sec/op    │   sec/op     vs base                │
CustomFunctions-10              3.318µ ± 2%   3.115µ ± 2%   -6.10% (p=0.000 n=10)
Suite/BenchmarkExec-10          1.236µ ± 1%   1.240µ ± 2%        ~ (p=0.617 n=10)
Suite/BenchmarkQuery-10         4.004µ ± 7%   3.363µ ± 2%  -16.02% (p=0.000 n=10)
Suite/BenchmarkParams-10        4.241µ ± 1%   3.758µ ± 2%  -11.40% (p=0.000 n=10)
Suite/BenchmarkStmt-10          2.830µ ± 0%   2.378µ ± 2%  -15.97% (p=0.000 n=10)
Suite/BenchmarkRows-10          126.3µ ± 1%   101.3µ ± 1%  -19.79% (p=0.000 n=10)
Suite/BenchmarkStmtRows-10      124.9µ ± 1%   100.5µ ± 2%  -19.56% (p=0.000 n=10)
Suite/BenchmarkStmt10Cols-10   10.130µ ± 0%   7.042µ ± 1%  -30.48% (p=0.000 n=10)
geomean                         8.655µ        7.328µ       -15.33%

                             │ base.10.txt  │               new.10.txt               │
                             │     B/op     │     B/op      vs base                  │
CustomFunctions-10               568.0 ± 0%     576.0 ± 0%   +1.41% (p=0.000 n=10)
Suite/BenchmarkExec-10           128.0 ± 0%     128.0 ± 0%        ~ (p=1.000 n=10) ¹
Suite/BenchmarkQuery-10          688.0 ± 0%     648.0 ± 0%   -5.81% (p=0.000 n=10)
Suite/BenchmarkParams-10       1.078Ki ± 0%   1.031Ki ± 0%   -4.35% (p=0.000 n=10)
Suite/BenchmarkStmt-10           920.0 ± 0%     872.0 ± 0%   -5.22% (p=0.000 n=10)
Suite/BenchmarkRows-10         9.305Ki ± 0%   9.188Ki ± 0%   -1.26% (p=0.000 n=10)
Suite/BenchmarkStmtRows-10     9.289Ki ± 0%   9.164Ki ± 0%   -1.35% (p=0.000 n=10)
Suite/BenchmarkStmt10Cols-10     992.0 ± 0%     696.0 ± 0%  -29.84% (p=0.000 n=10)
geomean                        1.181Ki        1.106Ki        -6.35%
¹ all samples are equal

                             │ base.10.txt │              new.10.txt              │
                             │  allocs/op  │ allocs/op   vs base                  │
CustomFunctions-10              18.00 ± 0%   18.00 ± 0%        ~ (p=1.000 n=10) ¹
Suite/BenchmarkExec-10          7.000 ± 0%   7.000 ± 0%        ~ (p=1.000 n=10) ¹
Suite/BenchmarkQuery-10         23.00 ± 0%   23.00 ± 0%        ~ (p=1.000 n=10) ¹
Suite/BenchmarkParams-10        27.00 ± 0%   27.00 ± 0%        ~ (p=1.000 n=10) ¹
Suite/BenchmarkStmt-10          25.00 ± 0%   25.00 ± 0%        ~ (p=1.000 n=10) ¹
Suite/BenchmarkRows-10          525.0 ± 0%   519.0 ± 0%   -1.14% (p=0.000 n=10)
Suite/BenchmarkStmtRows-10      524.0 ± 0%   518.0 ± 0%   -1.15% (p=0.000 n=10)
Suite/BenchmarkStmt10Cols-10    39.00 ± 0%   19.00 ± 0%  -51.28% (p=0.000 n=10)
geomean                         46.26        42.17        -8.86%
¹ all samples are equal
diff --git a/sqlite3.go b/sqlite3.go
@@ -201,6 +201,53 @@ static int sqlite3_system_errno(sqlite3 *db) {
   return 0;
 }
 #endif
+
+#define GO_SQLITE3_DECL_DATE (1 << 7)
+#define GO_SQLITE3_DECL_BOOL (1 << 6)
+#define GO_SQLITE3_DECL_MASK (GO_SQLITE3_DECL_DATE | GO_SQLITE3_DECL_BOOL)
+#define GO_SQLITE3_TYPE_MASK (GO_SQLITE3_DECL_BOOL - 1)
+
+// _sqlite3_column_decltypes stores the declared column type in the typs array.
+// This function must always be called before _sqlite3_column_types since it
+// overwrites the datatype.
+static void _sqlite3_column_decltypes(sqlite3_stmt* stmt, uint8_t *typs, int ntyps) {
+	for (int i = 0; i < ntyps; i++) {
+		const char *typ = sqlite3_column_decltype(stmt, i);
+		if (typ == NULL) {
+			typs[i] = 0;
+			continue;
+		}
+		switch (typ[0]) {
+		case 'b':
+		case 'B':
+			if (!sqlite3_stricmp(typ, "boolean")) {
+				typs[i] = GO_SQLITE3_DECL_BOOL;
+			}
+			break;
+		case 'd':
+		case 'D':
+			if (!sqlite3_stricmp(typ, "date") || !sqlite3_stricmp(typ, "datetime")) {
+				typs[i] = GO_SQLITE3_DECL_DATE;
+			}
+			break;
+		case 't':
+		case 'T':
+			if (!sqlite3_stricmp(typ, "timestamp")) {
+				typs[i] = GO_SQLITE3_DECL_DATE;
+			}
+			break;
+		default:
+			typs[i] = 0;
+		}
+	}
+}
+
+static void _sqlite3_column_types(sqlite3_stmt *stmt, uint8_t *typs, int ntyps) {
+	for (int i = 0; i < ntyps; i++) {
+		typs[i] &= GO_SQLITE3_DECL_MASK; // clear lower bits
+		typs[i] |= (uint8_t)sqlite3_column_type(stmt, i);
+	}
+}
 */
 import "C"
 import (
@@ -239,12 +286,6 @@ var SQLiteTimestampFormats = []string{
 	"2006-01-02",
 }
 
-const (
-	columnDate      string = "date"
-	columnDatetime  string = "datetime"
-	columnTimestamp string = "timestamp"
-)
-
 // This variable can be replaced with -ldflags like below:
 // go build -ldflags="-X 'github.com/mattn/go-sqlite3.driverName=my-sqlite3'"
 var driverName = "sqlite3"
@@ -390,13 +431,32 @@ type SQLiteResult struct {
 	changes int64
 }
 
+// A columnType is a compact representation of sqlite3 columns datatype and
+// declared type. The first two bits store the declared type and the remaining
+// six bits store the sqlite3 datatype.
+type columnType uint8
+
+// declType returns the declared type, which is currently GO_SQLITE3_DECL_DATE
+// or GO_SQLITE3_DECL_BOOL, since those are the only two types that we need for
+// converting values.
+func (c columnType) declType() int {
+	return int(c) & C.GO_SQLITE3_DECL_MASK
+}
+
+// dataType returns the sqlite3 datatype code of the column, which is the
+// result of sqlite3_column_type.
+func (c columnType) dataType() int {
+	return int(c) & C.GO_SQLITE3_TYPE_MASK
+}
+
 // SQLiteRows implements driver.Rows.
 type SQLiteRows struct {
 	s        *SQLiteStmt
 	nc       int32 // Number of columns
 	cls      bool  // True if we need to close the parent statement in Close
 	cols     []string
 	decltype []string
+	coltype  []columnType
 	ctx      context.Context // no better alternative to pass context into Next() method
 	closemu  sync.Mutex
 }
@@ -2148,7 +2208,10 @@ func (rc *SQLiteRows) Columns() []string {
 	return rc.cols
 }
 
-func (rc *SQLiteRows) declTypes() []string {
+// DeclTypes return column types.
+func (rc *SQLiteRows) DeclTypes() []string {
+	rc.s.mu.Lock()
+	defer rc.s.mu.Unlock()
 	if rc.s.s != nil && rc.decltype == nil {
 		rc.decltype = make([]string, rc.nc)
 		for i := 0; i < int(rc.nc); i++ {
@@ -2158,13 +2221,6 @@ func (rc *SQLiteRows) declTypes() []string {
 	return rc.decltype
 }
 
-// DeclTypes return column types.
-func (rc *SQLiteRows) DeclTypes() []string {
-	rc.s.mu.Lock()
-	defer rc.s.mu.Unlock()
-	return rc.declTypes()
-}
-
 // Next move cursor to next. Attempts to honor context timeout from QueryContext call.
 func (rc *SQLiteRows) Next(dest []driver.Value) error {
 	rc.s.mu.Lock()
@@ -2197,6 +2253,13 @@ func (rc *SQLiteRows) Next(dest []driver.Value) error {
 	}
 }
 
+func (rc *SQLiteRows) colTypePtr() *C.uint8_t {
+	if len(rc.coltype) == 0 {
+		return nil
+	}
+	return (*C.uint8_t)(unsafe.Pointer(&rc.coltype[0]))
+}
+
 // nextSyncLocked moves cursor to next; must be called with locked mutex.
 func (rc *SQLiteRows) nextSyncLocked(dest []driver.Value) error {
 	rv := C._sqlite3_step_internal(rc.s.s)
@@ -2210,15 +2273,24 @@ func (rc *SQLiteRows) nextSyncLocked(dest []driver.Value) error {
 		}
 		return nil
 	}
+	if len(dest) == 0 {
+		return nil
+	}
 
-	rc.declTypes()
+	if rc.coltype == nil {
+		rc.coltype = make([]columnType, rc.nc)
+		C._sqlite3_column_decltypes(rc.s.s, rc.colTypePtr(), C.int(rc.nc))
+	}
+	// Must call this each time since sqlite3 is loosely
+	// typed and the column types can vary between rows.
+	C._sqlite3_column_types(rc.s.s, rc.colTypePtr(), C.int(rc.nc))
 
 	for i := range dest {
-		switch C.sqlite3_column_type(rc.s.s, C.int(i)) {
+		switch rc.coltype[i].dataType() {
 		case C.SQLITE_INTEGER:
 			val := int64(C.sqlite3_column_int64(rc.s.s, C.int(i)))
-			switch rc.decltype[i] {
-			case columnTimestamp, columnDatetime, columnDate:
+			switch rc.coltype[i].declType() {
+			case C.GO_SQLITE3_DECL_DATE:
 				var t time.Time
 				// Assume a millisecond unix timestamp if it's 13 digits -- too
 				// large to be a reasonable timestamp in seconds.
@@ -2233,7 +2305,7 @@ func (rc *SQLiteRows) nextSyncLocked(dest []driver.Value) error {
 					t = t.In(rc.s.c.loc)
 				}
 				dest[i] = t
-			case "boolean":
+			case C.GO_SQLITE3_DECL_BOOL:
 				dest[i] = val > 0
 			default:
 				dest[i] = val
@@ -2257,8 +2329,7 @@ func (rc *SQLiteRows) nextSyncLocked(dest []driver.Value) error {
 			n := int(C.sqlite3_column_bytes(rc.s.s, C.int(i)))
 			s := C.GoStringN((*C.char)(unsafe.Pointer(C.sqlite3_column_text(rc.s.s, C.int(i)))), C.int(n))
 
-			switch rc.decltype[i] {
-			case columnTimestamp, columnDatetime, columnDate:
+			if rc.coltype[i].declType() == C.GO_SQLITE3_DECL_DATE {
 				var t time.Time
 				s = strings.TrimSuffix(s, "Z")
 				for _, format := range SQLiteTimestampFormats {
@@ -2275,7 +2346,7 @@ func (rc *SQLiteRows) nextSyncLocked(dest []driver.Value) error {
 					t = t.In(rc.s.c.loc)
 				}
 				dest[i] = t
-			default:
+			} else {
 				dest[i] = s
 			}
 		}
diff --git a/sqlite3_test.go b/sqlite3_test.go
@@ -2030,7 +2030,7 @@ func BenchmarkCustomFunctions(b *testing.B) {
 }
 
 func TestSuite(t *testing.T) {
-	initializeTestDB(t)
+	initializeTestDB(t, false)
 	defer freeTestDB()
 
 	for _, test := range tests {
@@ -2039,7 +2039,7 @@ func TestSuite(t *testing.T) {
 }
 
 func BenchmarkSuite(b *testing.B) {
-	initializeTestDB(b)
+	initializeTestDB(b, true)
 	defer freeTestDB()
 
 	for _, benchmark := range benchmarks {
@@ -2068,8 +2068,13 @@ type TestDB struct {
 
 var db *TestDB
 
-func initializeTestDB(t testing.TB) {
-	tempFilename := TempFilename(t)
+func initializeTestDB(t testing.TB, memory bool) {
+	var tempFilename string
+	if memory {
+		tempFilename = ":memory:"
+	} else {
+		tempFilename = TempFilename(t)
+	}
 	d, err := sql.Open("sqlite3", tempFilename+"?_busy_timeout=99999")
 	if err != nil {
 		os.Remove(tempFilename)
@@ -2084,9 +2089,11 @@ func freeTestDB() {
 	if err != nil {
 		panic(err)
 	}
-	err = os.Remove(db.tempFilename)
-	if err != nil {
-		panic(err)
+	if db.tempFilename != "" && db.tempFilename != ":memory:" {
+		err := os.Remove(db.tempFilename)
+		if err != nil {
+			panic(err)
+		}
 	}
 }
 
@@ -2112,6 +2119,7 @@ var benchmarks = []testing.InternalBenchmark{
 	{Name: "BenchmarkRows", F: benchmarkRows},
 	{Name: "BenchmarkStmtRows", F: benchmarkStmtRows},
 	{Name: "BenchmarkQueryParallel", F: benchmarkQueryParallel},
+	{Name: "BenchmarkStmt10Cols", F: benchmarkStmt10Cols},
 }
 
 func (db *TestDB) mustExec(sql string, args ...any) sql.Result {
@@ -2586,3 +2594,60 @@ func benchmarkQueryParallel(b *testing.B) {
 		}
 	})
 }
+
+func benchmarkStmt10Cols(b *testing.B) {
+	db.once.Do(makeBench)
+
+	const createTableStmt = `
+	DROP TABLE IF EXISTS bench_cols;
+	VACUUM;
+	CREATE TABLE bench_cols (
+		r0 INTEGER NOT NULL,
+		r1 INTEGER NOT NULL,
+		r2 INTEGER NOT NULL,
+		r3 INTEGER NOT NULL,
+		r4 INTEGER NOT NULL,
+		r5 INTEGER NOT NULL,
+		r6 INTEGER NOT NULL,
+		r7 INTEGER NOT NULL,
+		r8 INTEGER NOT NULL,
+		r9 INTEGER NOT NULL
+	);`
+	if _, err := db.Exec(createTableStmt); err != nil {
+		b.Fatal(err)
+	}
+	for i := int64(0); i < 4; i++ {
+		_, err := db.Exec("INSERT INTO bench_cols VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?);",
+			i, i, i, i, i, i, i, i, i, i)
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+
+	stmt, err := db.Prepare("SELECT * FROM bench_cols;")
+	if err != nil {
+		b.Fatal(err)
+	}
+	defer stmt.Close()
+
+	b.ResetTimer()
+	var (
+		v0, v1, v2, v3, v4 int64
+		v5, v6, v7, v8, v9 int64
+	)
+	for i := 0; i < b.N; i++ {
+		rows, err := stmt.Query()
+		if err != nil {
+			b.Fatal(err)
+		}
+		for rows.Next() {
+			err := rows.Scan(&v0, &v1, &v2, &v3, &v4, &v5, &v6, &v7, &v8, &v9)
+			if err != nil {
+				b.Fatal(err)
+			}
+		}
+		if err := rows.Err(); err != nil {
+			b.Fatal(err)
+		}
+	}
+}
