Concurrent result sets

CHANGELOG.md

@@ -1,4 +1,5 @@
 * Supported `sql.Null*` from `database/sql` as query params in `toValue` func
+* Added `WithConcurrentResultSets` option for `db.Query().Query()`
 
 ## v3.118.0
 * Added support for nullable `Date32`, `Datetime64`, `Timestamp64`, and `Interval64` types in the `optional` parameter builder

internal/query/client.go

@@ -404,6 +404,7 @@ func clientQuery(ctx context.Context, pool sessionPool, q string, opts ...option
 		if err != nil {
 			return xerrors.WithStackTrace(err)
 		}
+
 		defer func() {
 			_ = streamResult.Close(ctx)
 		}()

internal/query/client_test.go

@@ -922,6 +922,24 @@ func TestClient(t *testing.T) {
 					Status:         Ydb.StatusIds_SUCCESS,
 					ResultSetIndex: 0,
 					ResultSet: &Ydb.ResultSet{
+						Columns: []*Ydb.Column{
+							{
+								Name: "a",
+								Type: &Ydb.Type{
+									Type: &Ydb.Type_TypeId{
+										TypeId: Ydb.Type_UINT64,
+									},
+								},
+							},
+							{
+								Name: "b",
+								Type: &Ydb.Type{
+									Type: &Ydb.Type_TypeId{
+										TypeId: Ydb.Type_UTF8,
+									},
+								},
+							},
+						},
 						Rows: []*Ydb.Value{
 							{
 								Items: []*Ydb.Value{{

internal/query/execute_query.go

@@ -32,6 +32,7 @@ type executeSettings interface {
 	ResourcePool() string
 	ResponsePartLimitSizeBytes() int64
 	Label() string
+	ConcurrentResultSets() bool
 }
 
 type executeScriptConfig interface {
@@ -92,7 +93,7 @@ func executeQueryRequest(sessionID, q string, cfg executeSettings) (
 		},
 		Parameters:             params,
 		StatsMode:              Ydb_Query.StatsMode(cfg.StatsMode()),
-		ConcurrentResultSets:   false,
+		ConcurrentResultSets:   cfg.ConcurrentResultSets(),
 		PoolId:                 cfg.ResourcePool(),
 		ResponsePartLimitBytes: cfg.ResponsePartLimitSizeBytes(),
 	}

internal/query/options/execute.go

@@ -46,6 +46,7 @@ type (
 		issueCallback          func(issues []*Ydb_Issue.IssueMessage)
 		responsePartLimitBytes int64
 		label                  string
+		concurrentResultSets   bool
 	}
 
 	// Execute is an interface for execute method options
@@ -72,9 +73,12 @@ type (
 	}
 	execModeOption         = ExecMode
 	responsePartLimitBytes int64
-	issuesOption           struct {
+
+	issuesOption struct {
 		callback func([]*Ydb_Issue.IssueMessage)
 	}
+
+	concurrentResultSets bool
 )
 
 func (poolID resourcePool) applyExecuteOption(s *executeSettings) {
@@ -132,6 +136,10 @@ func (opts issuesOption) applyExecuteOption(s *executeSettings) {
 	s.issueCallback = opts.callback
 }
 
+func (opt concurrentResultSets) applyExecuteOption(s *executeSettings) {
+	s.concurrentResultSets = bool(opt)
+}
+
 const (
 	ExecModeParse    = ExecMode(Ydb_Query.ExecMode_EXEC_MODE_PARSE)
 	ExecModeValidate = ExecMode(Ydb_Query.ExecMode_EXEC_MODE_VALIDATE)
@@ -205,6 +213,10 @@ func (s *executeSettings) Label() string {
 	return s.label
 }
 
+func (s *executeSettings) ConcurrentResultSets() bool {
+	return s.concurrentResultSets
+}
+
 func WithParameters(params params.Parameters) parametersOption {
 	return parametersOption{
 		params: params,
@@ -237,6 +249,10 @@ func WithResponsePartLimitSizeBytes(size int64) responsePartLimitBytes {
 	return responsePartLimitBytes(size)
 }
 
+func WithConcurrentResultSets(isEnabled bool) concurrentResultSets {
+	return concurrentResultSets(isEnabled)
+}
+
 func (size responsePartLimitBytes) applyExecuteOption(s *executeSettings) {
 	s.responsePartLimitBytes = int64(size)
 }

internal/query/result.go

@@ -8,12 +8,14 @@
 	"time"
 
 	"github.com/ydb-platform/ydb-go-genproto/Ydb_Query_V1"
+	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
 	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Issue"
 	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Query"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/result"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stats"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/types"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xiter"
 	"github.com/ydb-platform/ydb-go-sdk/v3/query"
@@ -433,34 +435,59 @@
 	return MaterializedResultSet(rs.Index(), rs.Columns(), rs.ColumnTypes(), rows), nil
 }
 
-func resultToMaterializedResult(ctx context.Context, r result.Result) (result.Result, error) {
-	var resultSets []result.Set
+func resultToMaterializedResult(ctx context.Context, r *streamResult) (result.Result, error) {
+	type resultSet struct {
+		rows    []query.Row
+		columns []*Ydb.Column
+	}
+	resultSetByIndex := make(map[int64]resultSet)
 
 	for {
-		rs, err := r.NextResultSet(ctx)
+		if ctx.Err() != nil {
+			return nil, xerrors.WithStackTrace(ctx.Err())
+		}
+		if r.closer.Err() != nil {
+			return nil, xerrors.WithStackTrace(r.closer.Err())
+		}
+
+		rs := resultSetByIndex[r.lastPart.GetResultSetIndex()]
+		if len(rs.columns) == 0 {
+			rs.columns = r.lastPart.GetResultSet().GetColumns()
+		}
+
+		rows := make([]query.Row, len(r.lastPart.GetResultSet().GetRows()))
+		for i := range r.lastPart.GetResultSet().GetRows() {
+			rows[i] = NewRow(rs.columns, r.lastPart.GetResultSet().GetRows()[i])
+		}
+		rs.rows = append(rs.rows, rows...)
+
+		resultSetByIndex[r.lastPart.GetResultSetIndex()] = rs
+
+		var err error
+		r.lastPart, err = r.nextPart(ctx)
 		if err != nil {
 			if xerrors.Is(err, io.EOF) {
 				break
 			}
 
 			return nil, xerrors.WithStackTrace(err)
 		}
+		if r.lastPart.GetExecStats() != nil && r.statsCallback != nil {
+			r.statsCallback(stats.FromQueryStats(r.lastPart.GetExecStats()))
+		}
+	}
 
-		var rows []query.Row
-		for {
-			row, err := rs.NextRow(ctx)
-			if err != nil {
-				if xerrors.Is(err, io.EOF) {
-					break
-				}
-
-				return nil, xerrors.WithStackTrace(err)
-			}
+	resultSets := make([]result.Set, len(resultSetByIndex))
+	for rsIndex, rs := range resultSetByIndex {
+		columnNames := make([]string, len(rs.columns))
+		columnTypes := make([]types.Type, len(rs.columns))
 
-			rows = append(rows, row)
+		for i := range rs.columns {
+			columnNames[i] = rs.columns[i].Name
+			columnTypes[i] = types.TypeFromYDB(rs.columns[i].Type)
 		}
 
-		resultSets = append(resultSets, MaterializedResultSet(rs.Index(), rs.Columns(), rs.ColumnTypes(), rows))
+		resultSets[rsIndex] = MaterializedResultSet(int(rsIndex), columnNames, columnTypes, rs.rows)
 	}
 
 	return &materializedResult{

query/execute_options.go

@@ -70,6 +70,10 @@ func WithResponsePartLimitSizeBytes(size int64) ExecuteOption {
 	return options.WithResponsePartLimitSizeBytes(size)
 }
 
+func WithConcurrentResultSets(isEnabled bool) ExecuteOption {
+	return options.WithConcurrentResultSets(isEnabled)
+}
+
 func WithCallOptions(opts ...grpc.CallOption) ExecuteOption {
 	return options.WithCallOptions(opts...)
 }