Add concurrent result sets in db.Query.Query(...)

Author: Flak34

internal/query/client.go

@@ -401,11 +401,12 @@ func clientQuery(ctx context.Context, pool sessionPool, q string, opts ...option
 		if err != nil {
 			return xerrors.WithStackTrace(err)
 		}
+
 		defer func() {
 			_ = streamResult.Close(ctx)
 		}()
 
-		r, err = resultToMaterializedResult(ctx, streamResult)
+		r, err = concurrentResultToMaterializedResult(ctx, streamResult)
 		if err != nil {
 			return xerrors.WithStackTrace(err)
 		}

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

@@ -93,7 +93,7 @@ func executeQueryRequest(sessionID, q string, cfg executeSettings) (
 		},
 		Parameters:             params,
 		StatsMode:              Ydb_Query.StatsMode(cfg.StatsMode()),
-		ConcurrentResultSets:   cfg.ConcurrentResultSets(),
+		ConcurrentResultSets:   false,
 		PoolId:                 cfg.ResourcePool(),
 		ResponsePartLimitBytes: cfg.ResponsePartLimitSizeBytes(),
 	}

internal/query/result.go

@@ -13,6 +13,7 @@ import (
 	"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"
@@ -341,6 +342,33 @@ func (r *streamResult) nextPartFunc(
 	}
 }
 
+func (r *streamResult) NextPart(ctx context.Context) (_ result.Part, err error) {
+	if r.lastPart == nil {
+		return nil, xerrors.WithStackTrace(io.EOF)
+	}
+
+	select {
+	case <-r.closer.Done():
+		return nil, xerrors.WithStackTrace(r.closer.Err())
+	case <-ctx.Done():
+		return nil, xerrors.WithStackTrace(ctx.Err())
+	default:
+		part, err := r.nextPart(ctx)
+		if err != nil && !xerrors.Is(err, io.EOF) {
+			return nil, xerrors.WithStackTrace(err)
+		}
+		if part.GetExecStats() != nil && r.statsCallback != nil {
+			r.statsCallback(stats.FromQueryStats(part.GetExecStats()))
+		}
+		defer func() {
+			r.lastPart = part
+			r.resultSetIndex = part.GetResultSetIndex()
+		}()
+
+		return newResultPart(r.lastPart), nil
+	}
+}
+
 func (r *streamResult) NextResultSet(ctx context.Context) (_ result.Set, err error) {
 	if r.trace != nil {
 		onDone := trace.QueryOnResultNextResultSet(r.trace, &ctx,
@@ -459,3 +487,58 @@ func resultToMaterializedResult(ctx context.Context, r result.Result) (result.Re
 		resultSets: resultSets,
 	}, nil
 }
+
+func concurrentResultToMaterializedResult(ctx context.Context, r result.ConcurrentResult) (result.Result, error) {
+	type resultSet struct {
+		rows        []query.Row
+		columnNames []string
+		columnTypes []types.Type
+	}
+	resultSetByIndex := make(map[int64]resultSet)
+
+	for {
+		if ctx.Err() != nil {
+			return nil, xerrors.WithStackTrace(ctx.Err())
+		}
+
+		part, err := r.NextPart(ctx)
+		if err != nil {
+			if xerrors.Is(err, io.EOF) {
+				break
+			}
+
+			return nil, xerrors.WithStackTrace(err)
+		}
+
+		rs := resultSetByIndex[part.ResultSetIndex()]
+		if len(rs.columnNames) == 0 {
+			rs.columnTypes = part.ColumnTypes()
+			rs.columnNames = part.ColumnNames()
+		}
+
+		rows := make([]query.Row, 0)
+		for {
+			row, err := part.NextRow(ctx)
+			if err != nil {
+				if xerrors.Is(err, io.EOF) {
+					break
+				}
+
+				return nil, xerrors.WithStackTrace(err)
+			}
+			rows = append(rows, row)
+		}
+		rs.rows = append(rs.rows, rows...)
+
+		resultSetByIndex[part.ResultSetIndex()] = rs
+	}
+
+	resultSets := make([]result.Set, len(resultSetByIndex))
+	for rsIndex, rs := range resultSetByIndex {
+		resultSets[rsIndex] = MaterializedResultSet(int(rsIndex), rs.columnNames, rs.columnTypes, rs.rows)
+	}
+
+	return &materializedResult{
+		resultSets: resultSets,
+	}, nil
+}

internal/query/result/result.go

@@ -21,6 +21,11 @@ type (
 		// with Go version 1.23+
 		ResultSets(ctx context.Context) xiter.Seq2[Set, error]
 	}
+	ConcurrentResult interface {
+		closer.Closer
+
+		NextPart(ctx context.Context) (Part, error)
+	}
 	Set interface {
 		Index() int
 		Columns() []string
@@ -34,6 +39,13 @@ type (
 		Set
 		closer.Closer
 	}
+	Part interface {
+		ResultSetIndex() int64
+		ColumnNames() []string
+		ColumnTypes() []types.Type
+
+		NextRow(ctx context.Context) (Row, error)
+	}
 	Row interface {
 		Values() []value.Value
 

internal/query/result_part.go

@@ -0,0 +1,76 @@
+package query
+
+import (
+	"context"
+	"io"
+
+	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
+	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Query"
+	"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/query"
+)
+
+var (
+	_ query.Part = (*resultPart)(nil)
+)
+
+type (
+	resultPart struct {
+		resultSetIndex int64
+		columns        []*Ydb.Column
+		rows           []*Ydb.Value
+		columnNames    []string
+		columnTypes    []types.Type
+		rowIndex       int
+	}
+)
+
+func (p *resultPart) ResultSetIndex() int64 {
+	return p.resultSetIndex
+}
+
+func (p *resultPart) ColumnNames() []string {
+	if len(p.columnNames) != 0 {
+		return p.columnNames
+	}
+	names := make([]string, len(p.columns))
+	for i, col := range p.columns {
+		names[i] = col.GetName()
+	}
+	p.columnNames = names
+	return names
+}
+
+func (p *resultPart) ColumnTypes() []types.Type {
+	if len(p.columnTypes) != 0 {
+		return p.columnTypes
+	}
+	colTypes := make([]types.Type, len(p.columns))
+	for i, col := range p.columns {
+		colTypes[i] = types.TypeFromYDB(col.GetType())
+	}
+	p.columnTypes = colTypes
+	return colTypes
+}
+
+func (p *resultPart) NextRow(ctx context.Context) (query.Row, error) {
+	if p.rowIndex == len(p.rows) {
+		return nil, xerrors.WithStackTrace(io.EOF)
+	}
+
+	defer func() {
+		p.rowIndex++
+	}()
+
+	return NewRow(p.columns, p.rows[p.rowIndex]), nil
+}
+
+func newResultPart(part *Ydb_Query.ExecuteQueryResponsePart) *resultPart {
+	return &resultPart{
+		resultSetIndex: part.GetResultSetIndex(),
+		columns:        part.GetResultSet().GetColumns(),
+		rows:           part.GetResultSet().GetRows(),
+		rowIndex:       0,
+	}
+}

query/execute_options.go

@@ -63,9 +63,9 @@ func WithResponsePartLimitSizeBytes(size int64) ExecuteOption {
 	return options.WithResponsePartLimitSizeBytes(size)
 }
 
-//func WithConcurrentResultSets(isEnabled bool) ExecuteOption {
-//	return options.WithConcurrentResultSets(isEnabled)
-//}
+func WithConcurrentResultSets(isEnabled bool) ExecuteOption {
+	return options.WithConcurrentResultSets(isEnabled)
+}
 
 func WithCallOptions(opts ...grpc.CallOption) ExecuteOption {
 	return options.WithCallOptions(opts...)

query/result.go

@@ -8,8 +8,10 @@ import (
 
 type (
 	Result            = result.Result
+	ConcurrentResult  = result.ConcurrentResult
 	ResultSet         = result.Set
 	ClosableResultSet = result.ClosableResultSet
+	Part              = result.Part
 	Row               = result.Row
 	Type              = types.Type
 	NamedDestination  = scanner.NamedDestination