Skip to content

⚡ Bolt: Optimize row scanning and query execution allocations #95

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
cungminh2710 merged 2 commits into from
May 31, 2026
Merged

⚡ Bolt: Optimize row scanning and query execution allocations #95

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5fb480c
perf(rain): optimize row scanning and query execution allocations
google-labs-jules[bot] May 31, 2026
8cd0d39
perf(rain): optimize row scanning and query execution allocations
google-labs-jules[bot] May 31, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 19 additions & 18 deletions pkg/rain/model.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -334,20 +334,23 @@ func scanRowsAgainstTableDirect(rows *sql.Rows, dest any, table *schema.TableDef
return err
}
} else {
// Re-derive a reflect.Value for the element to reset it and handle non-offset columns.
item := reflect.NewAt(elemType, ptr).Elem()

// OPTIMIZATION: Skip zeroing existing elements if the plan is a full scan.
// This allows us to reuse existing pointer allocations in the struct fields.
if !plan.isFullScan {
// Reset existing element to its zero state before reuse to avoid data carry-over.
item.Set(zeroElem)
}

var targetVal reflect.Value
if plan.needsTargetValue {
targetVal = item
// Re-derive a reflect.Value for the element to reset it and handle non-offset columns.
targetVal = reflect.NewAt(elemType, ptr).Elem()

// OPTIMIZATION: Skip zeroing existing elements if the plan is a full scan.
// This allows us to reuse existing pointer allocations in the struct fields.
if !plan.isFullScan {
// Reset existing element to its zero state before reuse to avoid data carry-over.
targetVal.Set(zeroElem)
}
} else if !plan.isFullScan {
// Even if we don't need a reflect.Value for assignments, we still need it to zero the struct
// if this isn't a full scan to avoid data carry-over from previous rows.
reflect.NewAt(elemType, ptr).Elem().Set(zeroElem)
}

if err := scanDirectRowAddr(ptr, targetVal, plan, scratch); err != nil {
return err
}
Expand Down Expand Up @@ -1026,19 +1029,15 @@ func scanCachedRowsAgainstTable(result *cachedSelectRows, dest any, table *schem
return err
}

items := target
items := reflect.MakeSlice(target.Type(), 0, len(result.Rows))
for _, row := range result.Rows {
var item reflect.Value
if pointerElems {
item = reflect.New(structType)
} else {
item = reflect.New(structType).Elem()
}

var scanTarget reflect.Value
if pointerElems {
item = reflect.New(structType)
scanTarget = item.Elem()
} else {
item = reflect.New(structType).Elem()
scanTarget = item
}

Expand Down Expand Up @@ -1226,9 +1225,11 @@ func newRowScanPlanForColumns(cols []string, modelType reflect.Type, table *sche
plan.timeValueCols = append(plan.timeValueCols, colPlan)
}
} else {
plan.needsTargetValue = true
plan.otherCols = append(plan.otherCols, colPlan)
}
default:
plan.needsTargetValue = true
plan.otherCols = append(plan.otherCols, colPlan)
}
} else {
Expand Down
59 changes: 34 additions & 25 deletions pkg/rain/query_compile.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,45 +27,41 @@ type compiledQuery struct {
sql string
argPlan []compiledArg
hasNames bool
args []any
}

func (q compiledQuery) literalArgs() ([]any, error) {
if q.hasNames {
return nil, ErrPreparedArgsRequired
}

args := make([]any, 0, len(q.argPlan))
for _, arg := range q.argPlan {
args = append(args, arg.value)
}
return args, nil
// OPTIMIZATION: Return a fresh copy of the pre-calculated arguments to avoid
// shared state footguns if the caller modifies the slice, while still
// avoiding the overhead of rebuilding the slice from argPlan.
return append([]any(nil), q.args...), nil
}

func (q compiledQuery) bind(args PreparedArgs) ([]any, error) {
if !q.hasNames {
bound := make([]any, 0, len(q.argPlan))
for _, arg := range q.argPlan {
bound = append(bound, arg.value)
}
if len(args) > 0 {
return nil, fmt.Errorf("rain: unexpected prepared args for query without placeholders")
}
return bound, nil
// OPTIMIZATION: Return a fresh copy of the pre-calculated arguments to avoid
// shared state footguns if the caller modifies the slice.
return append([]any(nil), q.args...), nil
}
Comment on lines 44 to 51
Copy link
Copy Markdown
Contributor

@greptile-apps greptile-apps Bot May 31, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Shared backing array returned to callers

literalArgs() and the !q.hasNames branch of bind() now return q.args directly instead of a fresh copy. For PreparedSelectQuery, the compiledQuery is stored permanently and reused across many calls — every invocation of Scan, Count, or Exists receives the same backing array. Today all call sites only read the slice (cache-key hashing, variadic spread to QueryContext), so concurrent reads are safe. However, a future change that modifies the returned slice (e.g., appending a tenant ID or an extra bound value) before passing it to the driver would silently corrupt the prepared query's args for every subsequent call, because the original never creates a defensive copy anymore. The previous code always returned an independent slice, which kept this footgun impossible.

Prompt To Fix With AI 
This is a comment left during a code review.
Path: pkg/rain/query_compile.go
Line: 41-46

Comment:
**Shared backing array returned to callers**

`literalArgs()` and the `!q.hasNames` branch of `bind()` now return `q.args` directly instead of a fresh copy. For `PreparedSelectQuery`, the `compiledQuery` is stored permanently and reused across many calls — every invocation of `Scan`, `Count`, or `Exists` receives the same backing array. Today all call sites only read the slice (cache-key hashing, variadic spread to `QueryContext`), so concurrent reads are safe. However, a future change that modifies the returned slice (e.g., appending a tenant ID or an extra bound value) before passing it to the driver would silently corrupt the prepared query's `args` for every subsequent call, because the original never creates a defensive copy anymore. The previous code always returned an independent slice, which kept this footgun impossible.

How can I resolve this? If you propose a fix, please make it concise.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

Fix in Codex


seen := make(map[string]struct{}, len(args))
bound := make([]any, 0, len(q.argPlan))
for _, arg := range q.argPlan {
bound := append([]any(nil), q.args...)
for i, arg := range q.argPlan {
if arg.kind == compiledArgLiteral {
bound = append(bound, arg.value)
continue
}
value, ok := args[arg.name]
if !ok {
return nil, fmt.Errorf("rain: missing prepared arg %q", arg.name)
}
seen[arg.name] = struct{}{}
bound = append(bound, value)
bound[i] = value
}
for name := range args {
if _, ok := seen[name]; ok {
Expand Down Expand Up @@ -120,19 +116,32 @@ func (c *compileContext) String() string {
}

func (c *compileContext) compiledQuery() compiledQuery {
// OPTIMIZATION: Explicitly copy the argPlan slice so that the compileContext
// and its underlying array can be safely returned to the sync.Pool without
// causing data corruption for the caller of compiledQuery.
argPlan := append([]compiledArg(nil), c.argPlan...)
var hasNames bool
for _, arg := range c.argPlan {
if arg.kind == compiledArgNamedPlaceholder {
hasNames = true
break
}
}

// OPTIMIZATION: Only copy the argPlan if the query has named placeholders.
// Queries with only literals can use the pre-calculated args slice.
var argPlan []compiledArg
if hasNames {
argPlan = append([]compiledArg(nil), c.argPlan...)
}

compiled := compiledQuery{
sql: c.String(),
argPlan: argPlan,
sql: c.String(),
argPlan: argPlan,
hasNames: hasNames,
}
for _, arg := range compiled.argPlan {
if arg.kind == compiledArgNamedPlaceholder {
compiled.hasNames = true
break
if len(c.argPlan) > 0 {
compiled.args = make([]any, len(c.argPlan))
for i, arg := range c.argPlan {
if arg.kind == compiledArgLiteral {
compiled.args[i] = arg.value
}
}
}
return compiled
Expand Down
13 changes: 6 additions & 7 deletions pkg/rain/query_select.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -975,11 +975,10 @@ func (q *SelectQuery) compileExists() (compiledQuery, error) {
}

func wrapExistsCompiled(compiled compiledQuery) (compiledQuery, error) {
existsQuery := compiledQuery{
sql: "SELECT EXISTS(" + compiled.sql + ")",
argPlan: make([]compiledArg, len(compiled.argPlan)),
hasNames: compiled.hasNames,
}
copy(existsQuery.argPlan, compiled.argPlan)
return existsQuery, nil
// NOTE: This shallow copies the input compiledQuery and wraps the SQL.
// The argPlan and args slices are shared with the original. This is safe
// because compileExists (the only caller) does not use the original after
// this call.
compiled.sql = "SELECT EXISTS(" + compiled.sql + ")"
return compiled, nil
Comment on lines 977 to +983
Copy link
Copy Markdown
Contributor

@greptile-apps greptile-apps Bot May 31, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Silent slice aliasing after removing explicit argPlan copy

The previous wrapExistsCompiled made an explicit deep copy of argPlan so that the returned compiledQuery had no shared mutable state with its source. The new code mutates and returns a value-type copy, meaning argPlan and args are still backed by the same arrays as the compiled passed in. In the one call site (compileExists), compiled is a local that immediately goes out of scope, so there is no live aliasing today. But the function signature accepts any compiledQuery by value, and a future caller that retains the original compiled alongside the returned exists-variant would silently share mutable slice state between the two. A comment documenting the ownership assumption (caller must not use the original after this call) would prevent a subtle future regression.

Prompt To Fix With AI 
This is a comment left during a code review.
Path: pkg/rain/query_select.go
Line: 977-979

Comment:
**Silent slice aliasing after removing explicit argPlan copy**

The previous `wrapExistsCompiled` made an explicit deep copy of `argPlan` so that the returned `compiledQuery` had no shared mutable state with its source. The new code mutates and returns a value-type copy, meaning `argPlan` and `args` are still backed by the same arrays as the `compiled` passed in. In the one call site (`compileExists`), `compiled` is a local that immediately goes out of scope, so there is no live aliasing today. But the function signature accepts any `compiledQuery` by value, and a future caller that retains the original `compiled` alongside the returned exists-variant would silently share mutable slice state between the two. A comment documenting the ownership assumption (caller must not use the original after this call) would prevent a subtle future regression.

How can I resolve this? If you propose a fix, please make it concise.

Fix in Codex

}