feat: stream local scans to Sirius sidecar

pkg/container/vector/vector.go

@@ -3243,6 +3243,11 @@ func (v *Vector) IsConst() bool {
 	return v.class == CONSTANT
 }
 
+// IsFlat reports whether the vector stores one physical value per logical row.
+func (v *Vector) IsFlat() bool {
+	return v.class == FLAT
+}
+
 func (v *Vector) IsGrouping() bool {
 	return v.length > 0 &&
 		v.length == v.gsp.Count() &&

pkg/frontend/computation_wrapper.go

@@ -2191,6 +2191,9 @@ func compileStatementContexts(
 ) (requestCtx, compileCtx context.Context) {
 	requestCtx = perfcounter.AttachCompilePlanMarkKey(ctx, crs)
 	if siriusStatementSelected(sql, stmt) {
+		if strings.HasPrefix(strings.ToUpper(strings.TrimSpace(sql)), sidecarStreamHintPrefix) {
+			return requestCtx, compile.WithSiriusStreamOffload(requestCtx)
+		}
 		return requestCtx, compile.WithSiriusOffload(requestCtx)
 	}
 	return requestCtx, requestCtx

pkg/frontend/computation_wrapper_test.go

@@ -2331,6 +2331,7 @@ func TestCompileStatementContextsPreserveCounterWithoutLeakingSelection(t *testi
 		separateChild bool
 	}{
 		{name: "selected", sql: "/*+ SIDECAR */ select 1", separateChild: true},
+		{name: "stream selected", sql: "/*+ SIDECAR STREAM */ select 1", separateChild: true},
 		{name: "unselected", sql: "select 1"},
 	} {
 		t.Run(test.name, func(t *testing.T) {

pkg/frontend/sidecar_offload.go

@@ -36,9 +36,10 @@ import (
 )
 
 const (
-	sidecarHintPrefix      = "/*+ SIDECAR */"
-	sidecarGPUHintPrefix   = "/*+ SIDECAR GPU */"
-	sidecarMaxResponseSize = 512 << 20 // 512 MB
+	sidecarHintPrefix       = "/*+ SIDECAR */"
+	sidecarGPUHintPrefix    = "/*+ SIDECAR GPU */"
+	sidecarStreamHintPrefix = "/*+ SIDECAR STREAM */"
+	sidecarMaxResponseSize  = 512 << 20 // 512 MB
 )
 
 // errSidecarNotConfigured is a sentinel indicating sidecar offload should
@@ -94,6 +95,9 @@ func getManifestBaseURL() string {
 func isSidecarQuery(sql string) (bool, bool) {
 	trimmed := strings.TrimSpace(sql)
 	upper := strings.ToUpper(trimmed)
+	if strings.HasPrefix(upper, sidecarStreamHintPrefix) {
+		return true, true
+	}
 	if strings.HasPrefix(upper, sidecarGPUHintPrefix) {
 		return true, true
 	}
@@ -108,6 +112,9 @@ func isSidecarQuery(sql string) (bool, bool) {
 func stripSidecarHint(sql string) string {
 	trimmed := strings.TrimSpace(sql)
 	upper := strings.ToUpper(trimmed)
+	if strings.HasPrefix(upper, sidecarStreamHintPrefix) {
+		return strings.TrimSpace(trimmed[len(sidecarStreamHintPrefix):])
+	}
 	if strings.HasPrefix(upper, sidecarGPUHintPrefix) {
 		return strings.TrimSpace(trimmed[len(sidecarGPUHintPrefix):])
 	}

pkg/frontend/sidecar_offload_test.go

@@ -57,6 +57,10 @@ func TestIsSidecarQuery(t *testing.T) {
 	assert.True(t, isSidecar)
 	assert.True(t, useGPU)
 
+	isSidecar, useGPU = isSidecarQuery("/*+ SIDECAR STREAM */ SELECT * FROM t")
+	assert.True(t, isSidecar)
+	assert.True(t, useGPU)
+
 	// Non-sidecar queries
 	isSidecar, _ = isSidecarQuery("SELECT * FROM t")
 	assert.False(t, isSidecar)
@@ -73,6 +77,7 @@ func TestStripSidecarHint(t *testing.T) {
 	assert.Equal(t, "SELECT * FROM t", stripSidecarHint("  /*+ SIDECAR */ SELECT * FROM t"))
 	assert.Equal(t, "SELECT * FROM t", stripSidecarHint("/*+ SIDECAR GPU */ SELECT * FROM t"))
 	assert.Equal(t, "SELECT * FROM t", stripSidecarHint("  /*+ sidecar gpu */ SELECT * FROM t"))
+	assert.Equal(t, "SELECT * FROM t", stripSidecarHint("/*+ SIDECAR STREAM */ SELECT * FROM t"))
 	assert.Equal(t, "SELECT * FROM t", stripSidecarHint("SELECT * FROM t"))
 }
 

pkg/sql/colexec/output/output.go

@@ -89,6 +89,9 @@ func (output *Output) Call(proc *process.Process) (vm.CallResult, error) {
 			result.Status = vm.ExecStop
 			return result, err
 		}
+		if output.stop != nil && output.stop() {
+			result.Status = vm.ExecStop
+		}
 
 		// TODO: analyzer.Output(result.Batch)
 		return result, nil
@@ -140,6 +143,9 @@ func (output *Output) Call(proc *process.Process) (vm.CallResult, error) {
 					result.Status = vm.ExecStop
 					return result, err
 				}
+				if output.stop != nil && output.stop() {
+					result.Status = vm.ExecStop
+				}
 
 				result.Batch = bat
 				// same as nonBlock

pkg/sql/colexec/output/output_test.go

@@ -156,6 +156,46 @@ func TestOutputCallbackCPUIsNotOutputWait(t *testing.T) {
 	require.Zero(t, proc.Mp().CurrNB())
 }
 
+func TestOutputCanStopOneProducerAfterSuccessfulCallback(t *testing.T) {
+	proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero())
+	calls := 0
+	arg := NewArgument().WithFunc(func(_ *batch.Batch, _ *perfcounter.CounterSet) error {
+		calls++
+		return nil
+	}).WithShouldStop(func() bool { return calls == 1 })
+	require.NoError(t, arg.Prepare(proc))
+	first := newBatch([]types.Type{types.T_int8.ToType()}, proc, 1)
+	second := newBatch([]types.Type{types.T_int8.ToType()}, proc, 1)
+	resetChildren(arg, []*batch.Batch{first, second})
+	result, err := vm.Exec(arg, proc)
+	require.NoError(t, err)
+	require.Equal(t, vm.ExecStop, result.Status)
+	require.Equal(t, 1, calls)
+	arg.GetChildren(0).Free(proc, false, nil)
+	arg.Free(proc, false, nil)
+	proc.Free()
+}
+
+func TestBlockingOutputCanStopAfterFirstReleasedBatch(t *testing.T) {
+	proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero())
+	calls := 0
+	arg := NewArgument().WithBlock(true).WithFunc(func(_ *batch.Batch, _ *perfcounter.CounterSet) error {
+		calls++
+		return nil
+	}).WithShouldStop(func() bool { return calls == 1 })
+	require.NoError(t, arg.Prepare(proc))
+	first := newBatch([]types.Type{types.T_int8.ToType()}, proc, 1)
+	second := newBatch([]types.Type{types.T_int8.ToType()}, proc, 1)
+	resetChildren(arg, []*batch.Batch{first, second})
+	result, err := vm.Exec(arg, proc)
+	require.NoError(t, err)
+	require.Equal(t, vm.ExecStop, result.Status)
+	require.Equal(t, 1, calls)
+	arg.GetChildren(0).Free(proc, false, nil)
+	arg.Free(proc, false, nil)
+	proc.Free()
+}
+
 // create a new block based on the type information
 func newBatch(ts []types.Type, proc *process.Process, rows int64) *batch.Batch {
 	return testutil.NewBatch(ts, false, int(rows), proc.Mp())

pkg/sql/colexec/output/types.go

@@ -46,6 +46,7 @@ type Output struct {
 
 	Data interface{}
 	Func func(*batch.Batch, *perfcounter.CounterSet) error
+	stop func() bool
 
 	// IsAdaptive enables the adaptive vector search fallback mechanism.
 	// When set to true and the query completes with zero results (rowCount == 0),
@@ -99,6 +100,14 @@ func (output *Output) WithFunc(Func func(*batch.Batch, *perfcounter.CounterSet)
 	return output
 }
 
+// WithShouldStop lets a sink end only its own producer pipeline after a
+// successful callback. It is used when a remote consumer declares that this
+// input relation is no longer needed; nil preserves the ordinary output path.
+func (output *Output) WithShouldStop(stop func() bool) *Output {
+	output.stop = stop
+	return output
+}
+
 // WithBlocck set the output is blocked. If true output will block the current pipeline, and cache
 // all input batches. And wait for all the input's batch to be locked before outputting the cached batch
 // to the downstream operator.

pkg/sql/compile/compile2.go

@@ -340,7 +340,9 @@ func (c *Compile) Run(_ uint64) (queryResult *util2.RunResult, err error) {
 	queryResult = &util2.RunResult{}
 	v2.TxnStatementTotalCounter.Inc()
 	if c.siriusRead != nil {
-		err = c.runSiriusRead(execTopContext)
+		err = c.runSiriusRead(execTopContext, func(snapshot mpool.AllocationAccountTerminalSnapshot) {
+			resourceRecorder.recordAllocationAccountTerminal(snapshot)
+		})
 		return queryResult, err
 	}
 	attemptStart := time.Now()