From 76a243deef97d4d907afe8510d50cdccf4833ca7 Mon Sep 17 00:00:00 2001
From: Matthew Nibecker <hello@mattnibecker.com>
Date: Wed, 26 Aug 2026 11:40:22 -0700
Subject: [PATCH] optimizer: make optimizeSourcePaths work with sub queries

Fixes #6206
---
 compiler/dag/op.go              | 29 +++++++++++++++++++++++++++++
 compiler/optimizer/optimizer.go |  3 ++-
 compiler/ztests/pruner.yaml     | 14 ++++++++++++++
 3 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/compiler/dag/op.go b/compiler/dag/op.go
index f10915684d..bcdc1a45c4 100644
--- a/compiler/dag/op.go
+++ b/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
+}
diff --git a/compiler/optimizer/optimizer.go b/compiler/optimizer/optimizer.go
index 056bb19c92..a7c78c7258 100644
--- a/compiler/optimizer/optimizer.go
+++ b/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) {
diff --git a/compiler/ztests/pruner.yaml b/compiler/ztests/pruner.yaml
index fdc9ea5e77..31608640ec 100644
--- a/compiler/ztests/pruner.yaml
+++ b/compiler/ztests/pruner.yaml
@@ -1,11 +1,15 @@
 # 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
@@ -13,3 +17,13 @@ outputs:
         )
       | 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