CHORE: Stream job results directly to destination instead of buffering in memory

[prasadlohakpure]

Problem

Sync job responses and result persistence were fully materializing query results in memory before writing them out:

• payloadHandler (handler.go) did json.Marshal(result) into a []byte, then w.Write(...) — buffering the entire response before the first byte reached the client.
• storeResults (job.go) did the same json.Marshal + write pattern for the result file/S3 object.
• Plugins that fetch rows from a live cursor (Postgres, Snowflake, StarRocks) always built the complete [][]any in memory (via FromRows or per-endpoint accumulation) before handing it back, regardless of result size.

For large result sets this meant multiple full copies of the same data alive in memory at once, and no bytes reaching the client/destination until everything was ready.

Changes

Response and persistence encoding

• handler.go: encode directly into the http.ResponseWriter via json.NewEncoder instead of marshal-then-write.
• job.go / internal/pkg/aws/s3.go: storeResults now encodes directly into the destination (local file via a temp-file-then-rename for atomicity, or S3 via a new StreamToS3 helper backed by io.Pipe + the multipart uploader) instead of marshaling to a []byte first.

Optional streaming plugin contract

• pkg/plugin/plugin.go: added an optional ResultStreamer interface (mirrors the existing HealthChecker pattern) — a plugin holding a live row cursor can implement StreamResult(ctx, w, ...) to write rows straight to the destination as they're fetched, instead of accumulating them.
• job.go: dispatch prefers StreamResult over Execute when a plugin implements it, for both sync and async execution.
• pkg/result/stream.go: shared RowWriter + StreamRows helper so any row source (currently database/sql) can emit the {"columns":...,"data":...} shape incrementally, one row at a time.
• Implemented StreamResult for postgres and snowflake (both hold a *sql.Rows cursor) and starrocks (fans out to multiple backend endpoints concurrently, then streams each endpoint's rows to the destination in order, releasing each endpoint's buffer as soon as it's written).
• plugin.Handler.Execute is unchanged — required by the base interface, but now unreachable for any plugin implementing ResultStreamer; each plugin's Execute returns an explicit error rather than carrying a second, unexercised implementation of the same query logic.

Smaller fixes defered for now

• starrocks: fans a query out to multiple backend endpoints concurrently and already buffers each endpoint's full row set in memory (rowsByEndpoint [][][]any) before anything gets written — it has no live cursor to stream from. Wiring it into RowWriter would require its own shared-writer plumbing, increasing the scope suffficiently.

Behavior notes

• For a sync job whose plugin implements ResultStreamer, the result is no longer inlined in the POST /api/v1/job response body — it's written straight to the result destination, and the caller fetches it via the existing GET /api/v1/job/{id}/result endpoint (already streams the file to the client). Non-streaming plugins are unaffected.
• Sync API responses now end with a trailing newline (an artifact of json.Encoder) where they previously did not.

Benchmarks

Local micro-benchmarks (Apple M2, Go testing.B/-benchmem), synthetic 4-column rows (int64, string, float64, bool), comparing the pre-change approach (build the full [][]any, then json.Marshal it) against the post-change approach (stream row-by-row) at 1,000 and 10,000 rows. Re-run after scoping this PR down to postgres/snowflake (starrocks streaming deferred to a follow-up) and collapsing the row writer into a single self-contained StreamRows function.

Peak retained memory (the actual point of this change)

Cumulative -benchmem bytes/op don't capture this — this is live heap held at once for the duration a result sits in memory before being written out.

Rows	Before (fully materialized)	After (streamed)	Reduction
1,000	~421 KB	~3.5 KB	~99.2%
10,000	~1.89 MB	~3.4 KB	~99.8%

Row processing (pkg/result: accumulate-then-marshal vs. stream)

Rows	Metric	Before	After	Δ
1,000	bytes/op	448,883	230,611	~49% less
1,000	allocs/op	9,203	11,202	~22% more (many small per-row marshals vs. one big one)
1,000	time/op	~559 µs	~652 µs	~17% slower
10,000	bytes/op	2,893,777	2,298,763	~21% less
10,000	allocs/op	96,219	116,216	~21% more
10,000	time/op	~5.27 ms	~5.99 ms	~14% slower

HTTP response encode (handler.go: marshal-then-write vs. Encoder.Encode)

Rows	Metric	Before	After	Δ
1,000	bytes/op	56,492	28,837	~49% less
1,000	allocs/op	12	12	unchanged
10,000	bytes/op	732,434	361,241	~51% less
10,000	allocs/op	43	43	unchanged

Testing:

Verified S3 result files as well as API response against live PROD running jobs

1. Snowflake plugin:
   job_id: 00a04fc7-b33e-4aad-bb16-4bf3f49d10d7

2. Starrocks plugin:
   job_id: 554e4a3c-5f1f-48fd-a993-68db76675e0e

3. Trino plugin
   job_id: a71fdd55-1373-4325-bcd0-9b11dd8024b3

Contract change:

Previously under process result files used to display :

{"error":"operation error S3: GetObject, 
https response error StatusCode: 404, 
RequestID: VHD6WMJVAS4ZX6E0, 
HostID: FGmBuPPDxSMPFpgcxDgR/q9xXfRgxQYKJCGtSHyHbwaU5+dt3R29VgyJBPZPBFnpGaRwU3mMCv0=, NoSuchKey: "}

Now: underprocess result files will display

Reason:
With now data streaming enabled, local result/ s3 file state can be incomplete since its not a whole copy but rather data being streamed. So we will defer to show the file unless streaming complete.

Takeaway: memory footprint drops sharply at both sizes (the actual goal), CPU time is roughly a wash or slightly higher for row-by-row streaming due to per-row allocation overhead — a deliberate, expected memory/CPU tradeoff, not a free win.

[wiz-55ccc8b716]

Wiz Scan Summary

Scanner	Findings
Vulnerabilities	-
Sensitive Data	-
Secrets	-
IaC Misconfigurations	-
SAST Findings	1
Software Management Findings	-
Total	1

View scan details in Wiz

To detect these findings earlier in the dev lifecycle, try the Wiz Code extension for VS Code, JetBrains, or Visual Studio.

[Copilot]

Pull request overview

This PR reduces memory pressure and improves time-to-first-byte by streaming job results directly to their output destination (HTTP response, local file, or S3) instead of fully materializing and marshaling results in memory.

Changes:

• Added a streaming result writer (RowWriter / StreamRows) to emit the existing {"columns":[...],"data":[...]} JSON shape incrementally.
• Introduced an optional plugin.ResultStreamer interface and updated job dispatch to prefer streaming when implemented (postgres/snowflake/starrocks updated accordingly).
• Switched result persistence to streaming writes (atomic temp-file-then-rename locally; io.Pipe + multipart uploader for S3), and updated sync HTTP responses to use json.Encoder.

Reviewed changes

Copilot reviewed 10 out of 10 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
pkg/result/stream.go	Adds incremental JSON result streaming helpers (RowWriter, StreamRows).
pkg/result/result.go	Refactors shared row/column extraction into helpers used by both buffered and streaming paths.
pkg/plugin/plugin.go	Adds optional ResultStreamer plugin contract for streaming results.
internal/pkg/object/command/starrocks/starrocks.go	Implements StreamResult entrypoint and makes Execute explicitly unreachable for streaming plugins.
internal/pkg/object/command/starrocks/job_context.go	Streams StarRocks results to an io.Writer and reworks endpoint fetch to support streaming emission.
internal/pkg/object/command/snowflake/snowflake.go	Implements streaming result output via result.StreamRows and refactors query setup.
internal/pkg/object/command/postgres/postgres.go	Implements streaming result output via result.StreamRows, with shared preparation/query helpers.
internal/pkg/heimdall/job.go	Prefers ResultStreamer in dispatch and adds writeResultToDestination for streaming persistence (local + S3).
internal/pkg/heimdall/handler.go	Encodes HTTP responses via json.NewEncoder instead of marshal-then-write.
internal/pkg/aws/s3.go	Adds StreamToS3 helper using io.Pipe and the AWS SDK uploader for streaming writes.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.