optimizer: make optimizeSourcePaths work with sub queries

compiler/dag/op.go

@@ -410,3 +410,32 @@ func WalkT[T any](v reflect.Value, post func(T) T) {
 		}
 	}
 }
+
+func WalkTWithError[T any](v reflect.Value, post func(T) (T, error)) error {
+	switch v.Kind() {
+	case reflect.Array, reflect.Slice:
+		for i := range v.Len() {
+			if err := WalkTWithError(v.Index(i), post); err != nil {
+				return err
+			}
+		}
+	case reflect.Interface, reflect.Pointer:
+		return WalkTWithError(v.Elem(), post)
+	case reflect.Struct:
+		for _, field := range v.Fields() {
+			if err := WalkTWithError(field, post); err != nil {
+				return err
+			}
+		}
+	}
+	if v.CanSet() {
+		if t, ok := v.Interface().(T); ok {
+			r, err := post(t)
+			if err != nil {
+				return err
+			}
+			v.Set(reflect.ValueOf(r))
+		}
+	}
+	return nil
+}

compiler/optimizer/optimizer.go

@@ -199,7 +199,7 @@ func (o *Optimizer) OptimizeDeleter(main *dag.Main, replicas int) error {
 }
 
 func (o *Optimizer) optimizeSourcePaths(seq dag.Seq) (dag.Seq, error) {
-	return walkEntries(seq, func(seq dag.Seq) (dag.Seq, error) {
+	err := dag.WalkTWithError(reflect.ValueOf(&seq), func(seq dag.Seq) (dag.Seq, error) {
 		if len(seq) == 0 {
 			return nil, errors.New("internal error: optimizer encountered empty sequential operator")
 		}
@@ -275,6 +275,7 @@ func (o *Optimizer) optimizeSourcePaths(seq dag.Seq) (dag.Seq, error) {
 		}
 		return seq, nil
 	})
+	return seq, err
 }
 
 func (o *Optimizer) SortKeys(seq dag.Seq) ([]order.SortKeys, error) {

compiler/ztests/pruner.yaml

@@ -1,15 +1,29 @@
 # Test pruner flips comparator with literal/path
 script: |
   super -f parquet -o test.parquet -c 'values {v:1}'
+  echo === Test pruner flips comparator with literal/path
   super compile -O -C -c 'from test.parquet | where 1 <= v and 1 >= v'
+  echo === Test pruner optimization works with in subquery expressions
+  super compile -O -C -c 'values {x:(from test.parquet | where 1 <= v and 1 >= v)}'
 
 outputs:
   - name: stdout
     data: |
+      === Test pruner flips comparator with literal/path
       file test.parquet format parquet filter (1<=v and 1>=v)
          pruner (
            expr compare(v.max, 1, true)>=0 and compare(v.min, 1, true)<=0
            fields v.max,v.min
         )
       | where 1<=v and 1>=v
       | output main
+      === Test pruner optimization works with in subquery expressions
+      null
+      | values {x:(
+        file test.parquet format parquet filter (1<=v and 1>=v)
+           pruner (
+             expr compare(v.max, 1, true)>=0 and compare(v.min, 1, true)<=0
+             fields v.max,v.min
+          )
+        | where 1<=v and 1>=v)}
+      | output main