caldempsey opened a new pull request, #152:
URL: https://github.com/apache/spark-connect-go/pull/152
### What changes were proposed in this pull request?
This PR adds streaming support to the Spark Connect Go client by introducing
`ToLocalIterator(ctx)` on the DataFrame API, matching the functionality
available in Scala's implementation. The implementation includes a new
`types.RowIterator` interface that streams Arrow records batch-by-batch, and a
`ToRecordBatches()` method on `ExecutePlanClient` that returns channels for
streaming Arrow records and errors. Context cancellation is handled throughout
the streaming pipeline to ensure proper resource cleanup.
### Why are the changes needed?
The Go client currently only supports `Collect()`, which loads all rows into
memory at once. This approach doesn't work for large OLAP queries that return
gigabytes of data, and prevents lightweight services like API gateways or
serverless functions from using Spark Connect effectively.
This PR brings the same streaming capability that Scala and Java clients
have through `toLocalIterator()`. Rows are deserialised batch-by-batch and
released immediately after consumption. Since `ToTable` was pulling entire
tables of Arrow batches, we needed a new method to stream batches individually.
For now, I've reused `ReadArrowTableToRows` for the batch-to-row conversion,
though this could be optimised in a follow-up PR.
The changes are purely additive to maintain backward compatibility. External
consumers can pull data at their own pace, making it suitable for streaming
over HTTP or gRPC. I believe there are many opportunities to improve
performance by reducing allocs and GC pressure, which can also be addressed in
a follow up.
### Does this PR introduce any user-facing change?
Yes, it adds a new optional API:
```go
it, err := df.ToLocalIterator(ctx)
if err != nil {
return err
}
defer it.Close()
for {
row, err := it.Next()
if err == io.EOF {
break
}
if err != nil {
return err
}
// process row ...
}
```
Existing code using `Collect()` or `ToArrow()` remains unchanged.
### How was this patch tested?
Unit tests verify that streaming returns rows in the correct order, handles
EOF properly, cancels contexts gracefully without goroutine leaks, converts
Arrow batches correctly, and propagates errors through the streaming pipeline.
The implementation was also tested against a Unity Catalog/Databricks Spark
cluster using this streaming HTTP handler:
```go
// Example usage with proper context handling in a REST server
func StreamingHandler(w http.ResponseWriter, r *http.Request) {
// Create a context with timeout for the request
ctx, cancel := context.WithTimeout(r.Context(), 30*time.Minute)
defer cancel()
// Set up proper cleanup on client disconnect
notify := w.(http.CloseNotifier).CloseNotify()
go func() {
select {
case <-notify:
cancel() // Cancel context if client disconnects
case <-ctx.Done():
// Context already done
}
}()
var df sql.DataFrame // Your DataFrame
rowIter, err := df.ToLocalIterator(ctx)
if err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError)
return
}
defer func() {
if closeErr := rowIter.Close(); closeErr != nil {
log.Printf("Error closing iterator: %v", closeErr)
}
}()
w.Header().Set("Content-Type", "application/x-ndjson")
w.WriteHeader(http.StatusOK)
encoder := json.NewEncoder(w)
for {
select {
case <-ctx.Done():
// Context cancelled - exit gracefully
return
default:
}
row, err := rowIter.Next()
if err != nil {
if errors.Is(err, io.EOF) {
return // Normal completion
}
if errors.Is(err, context.Canceled) || errors.Is(err,
context.DeadlineExceeded) {
return // Context cancellation - exit gracefully
}
// Log error but don't crash
log.Printf("Error reading row: %v", err)
return
}
// Stream the row as JSON
if err := encoder.Encode(map[string]interface{}{
"values": row.Values(),
}); err != nil {
// Client disconnected or other write error
return
}
// Flush the response to ensure streaming
if flusher, ok := w.(http.Flusher); ok {
flusher.Flush()
}
}
}
```
All tests pass locally with `go test ./...`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
