user ops: Allow from statements (#4698)

Closes #4645

Author: mattnibs

compiler/job.go

@@ -40,17 +40,11 @@ func NewJob(octx *op.Context, in ast.Seq, src *data.Source, head *lakeparse.Comm
 	if len(seq) == 0 {
 		return nil, errors.New("internal error: AST seq cannot be empty")
 	}
-	from, reader, err := buildFrom(seq[0], head)
-	if err != nil {
-		return nil, err
-	}
-	if from != nil {
-		seq.Prepend(from)
-	}
 	entry, err := semantic.Analyze(octx.Context, seq, src, head)
 	if err != nil {
 		return nil, err
 	}
+	reader, _ := entry[0].(*kernel.Reader)
 	return &Job{
 		octx:      octx,
 		builder:   kernel.NewBuilder(octx, src),
@@ -60,71 +54,6 @@ func NewJob(octx *op.Context, in ast.Seq, src *data.Source, head *lakeparse.Comm
 	}, nil
 }
 
-func buildFrom(op ast.Op, head *lakeparse.Commitish) (*ast.From, *kernel.Reader, error) {
-	var from *ast.From
-	switch op := op.(type) {
-	case *ast.From:
-		// Already have an entry point with From.  Do nothing.
-		return nil, nil, nil
-	case *ast.Scope:
-		if len(op.Body) == 0 {
-			return nil, nil, errors.New("internal error: scope op has empty body")
-		}
-		return buildFrom(op.Body[0], head)
-	default:
-		var readers *kernel.Reader
-		trunk := ast.Trunk{Kind: "Trunk"}
-		if head != nil {
-			// For the lakes, if there is no from operator, then
-			// we default to scanning HEAD (without any of the
-			// from options).
-			trunk.Source = &ast.Pool{
-				Kind: "Pool",
-				Spec: ast.PoolSpec{
-					Pool: &ast.String{
-						Kind: "String",
-						Text: "HEAD",
-					},
-				},
-			}
-		} else {
-			readers = &kernel.Reader{}
-			trunk.Source = readers
-		}
-		from = &ast.From{
-			Kind:   "From",
-			Trunks: []ast.Trunk{trunk},
-		}
-		// XXX why not move this above? or
-		if isParallelWithLeadingFroms(op) {
-			from = nil
-			readers = nil
-		}
-		return from, readers, nil
-	}
-}
-
-func isParallelWithLeadingFroms(o ast.Op) bool {
-	par, ok := o.(*ast.Parallel)
-	if !ok {
-		return false
-	}
-	for _, seq := range par.Paths {
-		if !hasLeadingFrom(seq) {
-			return false
-		}
-	}
-	return true
-}
-
-func hasLeadingFrom(seq ast.Seq) bool {
-	if len(seq) == 0 {
-		return false
-	}
-	_, ok := seq[0].(*ast.From)
-	return ok
-}
-
 func (j *Job) Entry() dag.Seq {
 	//XXX need to prepend consts depending on context
 	return j.entry

compiler/kernel/op.go

@@ -704,6 +704,8 @@ func isEntry(seq dag.Seq) bool {
 		return true
 	case *dag.Scope:
 		return isEntry(op.Body)
+	case *dag.UserOpCall:
+		return isEntry(op.Body)
 	case *dag.Fork:
 		if len(op.Paths) == 0 {
 			return false

compiler/semantic/analyzer.go

@@ -7,14 +7,27 @@ import (
 	"github.com/brimdata/zed/compiler/ast"
 	"github.com/brimdata/zed/compiler/ast/dag"
 	"github.com/brimdata/zed/compiler/data"
+	"github.com/brimdata/zed/compiler/kernel"
 	"github.com/brimdata/zed/lakeparse"
 )
 
 // Analyze performs a semantic analysis of the AST, translating it from AST
 // to DAG form, resolving syntax ambiguities, and performing constant propagation.
 // After semantic analysis, the DAG is ready for either optimization or compilation.
 func Analyze(ctx context.Context, seq ast.Seq, source *data.Source, head *lakeparse.Commitish) (dag.Seq, error) {
-	return newAnalyzer(ctx, source, head).semSeq(seq)
+	a := newAnalyzer(ctx, source, head)
+	s, err := a.semSeq(seq)
+	if err != nil {
+		return nil, err
+	}
+	op, err := a.buildFrom(s[0])
+	if err != nil {
+		return nil, err
+	}
+	if op != nil {
+		s.Prepend(op)
+	}
+	return s, nil
 }
 
 type analyzer struct {
@@ -49,6 +62,37 @@ func (a *analyzer) exitScope() {
 	a.scope = a.scope.parent
 }
 
+func (a *analyzer) buildFrom(op dag.Op) (dag.Op, error) {
+	switch op := op.(type) {
+	case *dag.FileScan, *dag.HTTPScan, *dag.PoolScan, *dag.LakeMetaScan, *dag.PoolMetaScan, *dag.CommitMetaScan, *dag.DeleteScan:
+		return nil, nil
+	case *dag.Fork:
+		return a.buildFrom(op.Paths[0][0])
+	case *dag.Scope:
+		return a.buildFrom(op.Body[0])
+	case *dag.UserOpCall:
+		return a.buildFrom(op.Body[0])
+	}
+	// No from so add a source.
+	if a.head == nil {
+		return &kernel.Reader{}, nil
+	}
+	pool := &ast.Pool{
+		Kind: "Pool",
+		Spec: ast.PoolSpec{
+			Pool: &ast.String{
+				Kind: "String",
+				Text: "HEAD",
+			},
+		},
+	}
+	ops, err := a.semPool(pool)
+	if err != nil {
+		return nil, err
+	}
+	return ops[0], nil
+}
+
 type opDecl struct {
 	ast   *ast.OpDecl
 	deps  []*dag.UserOp

runtime/op/ztests/user-from.yaml

@@ -0,0 +1,26 @@
+script: |
+  zq -z -I test.zed
+
+inputs:
+  - name: test.zed
+    data: |
+      op fileA(): (
+        file A.zson | sort a 
+      )
+      fileA()
+  - name: A.zson
+    data: |
+      {a:5}
+      {a:1}
+      {a:7}
+      {a:3}
+      {a:8}
+  
+outputs:
+  - name: stdout
+    data: |
+      {a:1}
+      {a:3}
+      {a:5}
+      {a:7}
+      {a:8}