Dynamic runtime selection logic

This commit changes the logic for selection of either vam or sam
runtime. If a given query contains a vectorized file then the vector
runtime is selected. This commit also introduces the -vam -sam runtime
flags that if enabled will force the selection of either specified
runtime.

Author: mattnibs

cli/runtimeflags/flags.go

@@ -3,8 +3,11 @@ package runtimeflags
 import (
 	"errors"
 	"flag"
+	"fmt"
+	"os"
 
 	"github.com/brimdata/super/cli/auto"
+	"github.com/brimdata/super/runtime/exec"
 	"github.com/brimdata/super/runtime/sam/expr/agg"
 	"github.com/brimdata/super/runtime/sam/op/fuse"
 	"github.com/brimdata/super/runtime/sam/op/sort"
@@ -28,35 +31,80 @@ func defaultMemMaxBytes() uint64 {
 	}
 }
 
+type EngineFlags struct {
+	sam     bool
+	vam     bool
+	Runtime exec.Runtime
+}
+
 type Flags struct {
 	// these memory limits should be based on a shared resource model
 	aggMemMax  auto.Bytes
 	sortMemMax auto.Bytes
 	fuseMemMax auto.Bytes
+	EngineFlags
+}
+
+func (e *EngineFlags) SetFlags(fs *flag.FlagSet) {
+	fs.BoolVar(&e.sam, "sam", false, "execute query in sequential runtime")
+	fs.BoolVar(&e.vam, "vam", false, "execute query in vector runtime")
 }
 
-func (f *Flags) SetFlags(fs *flag.FlagSet) {
-	f.aggMemMax = auto.NewBytes(uint64(agg.MaxValueSize))
-	fs.Var(&f.aggMemMax, "aggmem", "maximum memory used per aggregate function value in MiB, MB, etc")
+func (e *Flags) SetFlags(fs *flag.FlagSet) {
+	e.aggMemMax = auto.NewBytes(uint64(agg.MaxValueSize))
+	fs.Var(&e.aggMemMax, "aggmem", "maximum memory used per aggregate function value in MiB, MB, etc")
 	def := defaultMemMaxBytes()
-	f.sortMemMax = auto.NewBytes(def)
-	fs.Var(&f.sortMemMax, "sortmem", "maximum memory used by sort in MiB, MB, etc")
-	f.fuseMemMax = auto.NewBytes(def)
-	fs.Var(&f.fuseMemMax, "fusemem", "maximum memory used by fuse in MiB, MB, etc")
+	e.sortMemMax = auto.NewBytes(def)
+	fs.Var(&e.sortMemMax, "sortmem", "maximum memory used by sort in MiB, MB, etc")
+	e.fuseMemMax = auto.NewBytes(def)
+	fs.Var(&e.fuseMemMax, "fusemem", "maximum memory used by fuse in MiB, MB, etc")
+	e.EngineFlags.SetFlags(fs)
 }
 
-func (f *Flags) Init() error {
-	if f.aggMemMax.Bytes <= 0 {
+func (e *EngineFlags) Init() error {
+	var err error
+	e.Runtime, err = e.getRuntime()
+	return err
+}
+
+func (e *Flags) Init() error {
+	if e.aggMemMax.Bytes <= 0 {
 		return errors.New("aggmem value must be greater than zero")
 	}
-	agg.MaxValueSize = int(f.aggMemMax.Bytes)
-	if f.sortMemMax.Bytes <= 0 {
+	agg.MaxValueSize = int(e.aggMemMax.Bytes)
+	if e.sortMemMax.Bytes <= 0 {
 		return errors.New("sortmem value must be greater than zero")
 	}
-	sort.MemMaxBytes = int(f.sortMemMax.Bytes)
-	if f.fuseMemMax.Bytes <= 0 {
+	sort.MemMaxBytes = int(e.sortMemMax.Bytes)
+	if e.fuseMemMax.Bytes <= 0 {
 		return errors.New("fusemem value must be greater than zero")
 	}
-	fuse.MemMaxBytes = int(f.fuseMemMax.Bytes)
-	return nil
+	fuse.MemMaxBytes = int(e.fuseMemMax.Bytes)
+	if e.sam && e.vam {
+		return errors.New("sam and vam flags cannot both be enabled")
+	}
+	return e.EngineFlags.Init()
+}
+
+func (e *EngineFlags) getRuntime() (exec.Runtime, error) {
+	// Flags take precedence.
+	if e.sam {
+		return exec.RuntimeSAM, nil
+	}
+	if e.vam {
+		return exec.RuntimeVAM, nil
+	}
+	// Then environment variable.
+	if rt := os.Getenv("SUPER_RUNTIME"); rt != "" {
+		switch rt {
+		case "sam":
+			return exec.RuntimeSAM, nil
+		case "vam":
+			return exec.RuntimeVAM, nil
+		default:
+			return exec.RuntimeAuto, fmt.Errorf("invalid SUPER_RUNTIME value: %q (must be \"vam\" or \"sam\")", rt)
+		}
+	}
+	return exec.RuntimeAuto, nil
+
 }

cmd/super/compile/shared.go

@@ -5,11 +5,13 @@ import (
 	"errors"
 	"flag"
 	"fmt"
+	"os"
 
 	"github.com/brimdata/super"
 	"github.com/brimdata/super/cli/dbflags"
 	"github.com/brimdata/super/cli/outputflags"
 	"github.com/brimdata/super/cli/queryflags"
+	"github.com/brimdata/super/cli/runtimeflags"
 	"github.com/brimdata/super/compiler"
 	"github.com/brimdata/super/compiler/describe"
 	"github.com/brimdata/super/compiler/parser"
@@ -25,13 +27,15 @@ import (
 )
 
 type Shared struct {
-	dag         bool
-	dynamic     bool
-	optimize    bool
-	parallel    int
-	query       bool
-	queryFlags  queryflags.QueryTextFlags
-	OutputFlags outputflags.Flags
+	dag          bool
+	dynamic      bool
+	optimize     bool
+	parallel     int
+	query        bool
+	runtime      bool
+	queryFlags   queryflags.QueryTextFlags
+	runtimeFlags runtimeflags.EngineFlags
+	OutputFlags  outputflags.Flags
 }
 
 func (s *Shared) SetFlags(fs *flag.FlagSet) {
@@ -40,11 +44,16 @@ func (s *Shared) SetFlags(fs *flag.FlagSet) {
 	fs.BoolVar(&s.optimize, "O", false, "display optimized DAG")
 	fs.IntVar(&s.parallel, "P", 0, "display parallelized DAG")
 	fs.BoolVar(&s.query, "C", false, "display DAG or AST as query text")
+	fs.BoolVar(&s.runtime, "runtime", false, "print selected runtime to stderr")
 	s.OutputFlags.SetFlags(fs)
 	s.queryFlags.SetFlags(fs)
+	s.runtimeFlags.SetFlags(fs)
 }
 
 func (s *Shared) Run(ctx context.Context, args []string, dbFlags *dbflags.Flags, desc bool) error {
+	if err := s.runtimeFlags.Init(); err != nil {
+		return err
+	}
 	if len(s.queryFlags.Query) == 0 && len(args) == 0 {
 		return errors.New("no query specified")
 	}
@@ -72,6 +81,9 @@ func (s *Shared) Run(ctx context.Context, args []string, dbFlags *dbflags.Flags,
 		s.dag = true
 	}
 	if !s.dag {
+		if s.runtime {
+			defer printRuntime(s.runtimeFlags.Runtime)
+		}
 		if s.query {
 			fmt.Println(sfmt.AST(ast.Parsed()))
 			return nil
@@ -83,11 +95,15 @@ func (s *Shared) Run(ctx context.Context, args []string, dbFlags *dbflags.Flags,
 	}
 	rctx := runtime.DefaultContext()
 	env := exec.NewEnvironment(storage.NewLocalEngine(), root)
+	env.Runtime = s.runtimeFlags.Runtime
 	env.Dynamic = s.dynamic
 	dag, err := compiler.Analyze(rctx, ast, env, false)
 	if err != nil {
 		return err
 	}
+	if s.runtime {
+		defer printRuntime(env.Runtime)
+	}
 	if desc {
 		description, err := describe.AnalyzeDAG(ctx, dag, env)
 		if err != nil {
@@ -122,3 +138,7 @@ func (s *Shared) writeValue(ctx context.Context, v any) error {
 	}
 	return err
 }
+
+func printRuntime(r exec.Runtime) {
+	fmt.Fprintf(os.Stderr, "runtime: %s\n", r)
+}

cmd/super/root/command.go

@@ -148,6 +148,7 @@ func (c *Command) Run(args []string) error {
 	env.Dynamic = c.inputFlags.Dynamic
 	env.IgnoreOpenErrors = !c.stopErr
 	env.ReaderOpts = c.inputFlags.ReaderOpts
+	env.Runtime = c.runtimeFlags.Runtime
 	comp := compiler.NewCompilerWithEnv(env)
 	query, err := runtime.CompileQuery(ctx, super.NewContext(), comp, ast, nil)
 	if err != nil {

compiler/semantic/analyzer.go

@@ -43,6 +43,9 @@ func Analyze(ctx context.Context, p *parser.AST, env *exec.Environment, extInput
 		}
 	}
 	main := newDagen(t.reporter).assemble(seq, t.resolver.funcs)
+	if env.Runtime == exec.RuntimeAuto && t.hasVectorizedInput {
+		env.Runtime = exec.RuntimeVAM
+	}
 	return main, t.Error()
 }
 
@@ -52,15 +55,16 @@ func Analyze(ctx context.Context, p *parser.AST, env *exec.Environment, extInput
 // to dataflow.
 type translator struct {
 	reporter
-	ctx      context.Context
-	resolver *resolver
-	checker  *checker
-	opCnt    map[*ast.OpDecl]int
-	opStack  []string
-	cteStack []*ast.SQLCTE
-	env      *exec.Environment
-	scope    *Scope
-	sctx     *super.Context
+	ctx                context.Context
+	resolver           *resolver
+	checker            *checker
+	hasVectorizedInput bool
+	opCnt              map[*ast.OpDecl]int
+	opStack            []string
+	cteStack           []*ast.SQLCTE
+	env                *exec.Environment
+	scope              *Scope
+	sctx               *super.Context
 }
 
 func newTranslator(ctx context.Context, r reporter, env *exec.Environment) *translator {

compiler/semantic/op.go

@@ -284,6 +284,9 @@ func (t *translator) file(n ast.Node, name string, args []ast.OpArg) sem.Op {
 	if format == "" {
 		format = sio.FormatFromPath(name)
 	}
+	if format == "csup" || format == "parquet" {
+		t.hasVectorizedInput = true
+	}
 	typ, err := t.fileType(name, format)
 	if err != nil {
 		t.error(n, err)

compiler/ztests/runtime-select.yaml

@@ -0,0 +1,40 @@
+script: |
+  # Check the environment variable works.
+  SUPER_RUNTIME=vam super compile -dag -C -runtime 'values "foo"'
+  >&2 echo "// ==="
+  # Check the flag overrides environment variable.
+  SUPER_RUNTIME=vam super compile -sam -dag -C -runtime 'values "foo"'
+  >&2 echo "// ==="
+  # A vectorized file runs in vam.
+  super -f parquet -o test.parquet -c 'values {x:1}'
+  super compile -dag -C -runtime 'from test.parquet'
+  >&2 echo "// ==="
+  # Normal file runs in sam.
+  super -f sup -o test.sup -c 'values {y:1}'
+  super compile -dag -C -runtime 'from test.sup'
+  >&2 echo "// ==="
+  # A vectorized and normal file run in vam.
+  super compile -dag -C -runtime 'fork ( from test.parquet ) ( from test.sup )'
+  >&2 echo "// ==="
+  # Check -sam with a vector input still runs in sam.
+  super compile -dag -sam -C -runtime 'from test.parquet'
+  >&2 echo "// ==="
+  # Error if both -sam and -vam are enabled.
+  ! super -vam -sam
+
+outputs:
+  - name: stderr
+    data: |
+      runtime: vam
+      // ===
+      runtime: sam
+      // ===
+      runtime: vam
+      // ===
+      runtime: auto
+      // ===
+      runtime: vam
+      // ===
+      runtime: sam
+      // ===
+      sam and vam flags cannot both be enabled

runtime/exec/environment.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"io"
 	"net/http"
-	"os"
 
 	"github.com/brimdata/super"
 	"github.com/brimdata/super/compiler/dag"
@@ -28,21 +27,39 @@ type VectorConcurrentPuller interface {
 	ConcurrentPull(done bool, id int) (vector.Any, error)
 }
 
+type Runtime int
+
+const (
+	RuntimeAuto Runtime = iota
+	RuntimeSAM
+	RuntimeVAM
+)
+
+func (r Runtime) String() string {
+	switch r {
+	case RuntimeSAM:
+		return "sam"
+	case RuntimeVAM:
+		return "vam"
+	default:
+		return "auto"
+	}
+}
+
 type Environment struct {
 	engine storage.Engine
 	db     *db.Root
-	useVAM bool
 
 	Dynamic          bool
 	IgnoreOpenErrors bool
 	ReaderOpts       anyio.ReaderOpts
+	Runtime          Runtime
 }
 
 func NewEnvironment(engine storage.Engine, d *db.Root) *Environment {
 	return &Environment{
 		engine: engine,
 		db:     d,
-		useVAM: os.Getenv("SUPER_VAM") != "",
 	}
 }
 
@@ -51,11 +68,7 @@ func (e *Environment) Engine() storage.Engine {
 }
 
 func (e *Environment) UseVAM() bool {
-	return e.useVAM
-}
-
-func (e *Environment) SetUseVAM() {
-	e.useVAM = true
+	return e.Runtime == RuntimeVAM
 }
 
 func (e *Environment) IsAttached() bool {

ztest/ztest.go

@@ -508,7 +508,7 @@ func runInternal(ctx context.Context, query string, input *string, outputFlags,
 	}
 	env := exec.NewEnvironment(nil, nil)
 	if vector {
-		env.SetUseVAM()
+		env.Runtime = exec.RuntimeVAM
 	}
 	q, err := runtime.CompileQuery(ctx, sctx, compiler.NewCompilerWithEnv(env), ast, readers)
 	if err != nil {