Add QueryChunksContext

Reading a result through database/sql copies and boxes every value, which is
significant overhead when the caller only forwards the rows. QueryChunksContext
hands the native chunks to a callback instead, so values can be read straight
out of the vectors.

Author: OzeroDev

chunk_query.go

@@ -0,0 +1,148 @@
+package duckdb
+
+import (
+	"context"
+	"database/sql"
+	"database/sql/driver"
+	"errors"
+
+	"github.com/duckdb/duckdb-go/v2/mapping"
+)
+
+// QueryChunksContext executes a query and passes each native result chunk to
+// consume. It exists to read results without the per-row allocations of
+// database/sql: values can be read straight out of DuckDB's vectors, e.g., via
+// DataChunk.GetVarcharView.
+//
+// The chunks arrive in result order, and consume observes one chunk at a time.
+// Returning an error from consume stops the iteration and returns that error
+// unchanged, so a sentinel error stops the iteration early.
+// consume is not called at all for a result without rows.
+//
+// The connection is occupied for the whole call. consume must therefore not
+// issue queries on sqlConn, which would deadlock. Note that DuckDB materializes
+// the result before the first chunk is fetched, so this API bounds the memory
+// that the caller allocates, not the memory that DuckDB does.
+//
+// The same *DataChunk is reused for every chunk, and its C-allocated memory is
+// released once consume returns. Neither the chunk nor anything aliasing it may
+// be retained beyond that. args bypass the argument conversion of database/sql
+// and are bound as-is. If query holds multiple statements, all but the last are
+// executed, and only the last one yields chunks.
+func QueryChunksContext(
+	ctx context.Context,
+	sqlConn *sql.Conn,
+	query string,
+	consume func(*DataChunk) error,
+	args ...any,
+) error {
+	if sqlConn == nil {
+		return getError(errAPI, errNilConn)
+	}
+	if consume == nil {
+		return getError(errAPI, errNilChunkConsumer)
+	}
+
+	values := make([]driver.Value, len(args))
+	for i, arg := range args {
+		values[i] = arg
+	}
+
+	return sqlConn.Raw(func(driverConn any) error {
+		conn, ok := driverConn.(*Conn)
+		if !ok {
+			return getError(errAPI, invalidConnError(driverConn))
+		}
+		return conn.queryChunksContext(ctx, query, argsToNamedArgs(values), consume)
+	})
+}
+
+// queryChunksContext prepares query on conn, streams its chunks to consume, and
+// closes the statement before returning.
+func (conn *Conn) queryChunksContext(
+	ctx context.Context,
+	query string,
+	args []driver.NamedValue,
+	consume func(*DataChunk) error,
+) error {
+	if conn.closed {
+		return errClosedCon
+	}
+
+	cleanupCtx := conn.setContext(ctx)
+	defer cleanupCtx()
+
+	return runWithCtxInterrupt(ctx, conn.conn, func(wctx context.Context) (err error) {
+		prepared, err := conn.prepareStmts(wctx, query)
+		if err != nil {
+			return err
+		}
+
+		defer func() {
+			closeErr := prepared.Close()
+			switch {
+			case err != nil && closeErr != nil:
+				err = errors.Join(err, closeErr)
+			case closeErr != nil:
+				err = closeErr
+			}
+		}()
+
+		return consumePreparedChunks(wctx, prepared, args, consume)
+	})
+}
+
+// consumePreparedChunks executes prepared and passes every result chunk to
+// consume, one at a time. It owns all C-allocated memory involved: the result is
+// destroyed on return, and each fetched chunk is freed after consume observes it.
+func consumePreparedChunks(
+	ctx context.Context,
+	prepared *Stmt,
+	args []driver.NamedValue,
+	consume func(*DataChunk) error,
+) error {
+	result, err := prepared.execute(ctx, args)
+	if err != nil {
+		return err
+	}
+	defer mapping.DestroyResult(result)
+
+	columnCount := mapping.ColumnCount(result)
+	columnNames := make([]string, columnCount)
+	for i := range columnNames {
+		columnNames[i] = mapping.ColumnName(result, mapping.IdxT(i))
+	}
+
+	var chunk DataChunk
+	chunk.columnNames = columnNames
+	for {
+		if err := ctx.Err(); err != nil {
+			return err
+		}
+		nativeChunk := mapping.FetchChunk(*result)
+		if nativeChunk.Ptr == nil {
+			if errMsg := mapping.ResultError(result); errMsg != "" {
+				if err := ctx.Err(); err != nil {
+					return err
+				}
+				return getError(errAPI, getDuckDBError(errMsg))
+			}
+			return ctx.Err()
+		}
+		if err := chunk.initFromDuckDataChunk(nativeChunk, false); err != nil {
+			mapping.DestroyDataChunk(&nativeChunk)
+			return getError(errAPI, err)
+		}
+
+		err := consumeResultChunk(&chunk, consume)
+		if err != nil {
+			return err
+		}
+	}
+}
+
+// consumeResultChunk passes chunk to consume and frees its C-allocated memory.
+func consumeResultChunk(chunk *DataChunk, consume func(*DataChunk) error) error {
+	defer chunk.close()
+	return consume(chunk)
+}

chunk_query_bench_test.go

@@ -0,0 +1,88 @@
+package duckdb
+
+import (
+	"context"
+	"database/sql"
+	"fmt"
+	"io"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+const benchChunkRowCount = 100_000
+
+// benchChunkQuery serializes every row to a single VARCHAR column.
+var benchChunkQuery = fmt.Sprintf(
+	`SELECT to_json(r)::VARCHAR FROM (
+		SELECT i AS id, i::VARCHAR AS name, i * 1.5 AS value FROM range(%d) AS values(i)
+	) AS r`, benchChunkRowCount)
+
+func setupChunkBench(b *testing.B) (*sql.Conn, []byte) {
+	b.Helper()
+	db := openDbWrapper(b, "")
+	conn := openConnWrapper(b, db, context.Background())
+	b.Cleanup(func() {
+		closeConnWrapper(b, conn)
+		closeDbWrapper(b, db)
+	})
+	return conn, make([]byte, 0, 128*1024)
+}
+
+// BenchmarkQueryChunks reads the result through the native chunk API.
+func BenchmarkQueryChunks(b *testing.B) {
+	conn, buf := setupChunkBench(b)
+
+	b.ReportAllocs()
+	for b.Loop() {
+		buf = buf[:0]
+		err := QueryChunksContext(context.Background(), conn, benchChunkQuery,
+			func(chunk *DataChunk) error {
+				for rowIdx := range chunk.GetSize() {
+					value, valid, err := chunk.GetVarcharView(0, rowIdx)
+					if err != nil {
+						return err
+					}
+					if !valid {
+						continue
+					}
+					if len(value)+1 > cap(buf)-len(buf) {
+						_, _ = io.Discard.Write(buf)
+						buf = buf[:0]
+					}
+					buf = append(buf, value...)
+					buf = append(buf, '\n')
+				}
+				return nil
+			},
+		)
+		require.NoError(b, err)
+	}
+}
+
+// BenchmarkQueryRows reads the result through database/sql.
+func BenchmarkQueryRows(b *testing.B) {
+	conn, buf := setupChunkBench(b)
+
+	b.ReportAllocs()
+	for b.Loop() {
+		buf = buf[:0]
+		rows, err := conn.QueryContext(context.Background(), benchChunkQuery)
+		require.NoError(b, err)
+
+		var value string
+		for rows.Next() {
+			if err := rows.Scan(&value); err != nil {
+				b.Fatal(err)
+			}
+			if len(value)+1 > cap(buf)-len(buf) {
+				_, _ = io.Discard.Write(buf)
+				buf = buf[:0]
+			}
+			buf = append(buf, value...)
+			buf = append(buf, '\n')
+		}
+		require.NoError(b, rows.Err())
+		require.NoError(b, rows.Close())
+	}
+}