airframe-sql: Resolve CTE before resolving group indexes (#2414)

wvlet · Sep 15, 2022 · 346ac97 · 346ac97
1 parent 687a6f7
commit 346ac97
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 3 deletions.
diff --git a/airframe-sql/src/main/scala/wvlet/airframe/sql/analyzer/TypeResolver.scala b/airframe-sql/src/main/scala/wvlet/airframe/sql/analyzer/TypeResolver.scala
@@ -24,17 +24,19 @@ import wvlet.log.LogSupport
   */
 object TypeResolver extends LogSupport {
 
-  def typerRules: List[Rule] =
+  def typerRules: List[Rule] = {
     // First resolve all input table types
-    TypeResolver.resolveAggregationIndexes _ ::
+    // CTE Table Refs must be resolved before resolving aggregation indexes
+    TypeResolver.resolveCTETableRef _ ::
+      TypeResolver.resolveAggregationIndexes _ ::
       TypeResolver.resolveAggregationKeys _ ::
-      TypeResolver.resolveCTETableRef _ ::
       TypeResolver.resolveTableRef _ ::
       TypeResolver.resolveJoinUsing _ ::
       TypeResolver.resolveRegularRelation _ ::
       TypeResolver.resolveColumns _ ::
       TypeResolver.resolveUnion _ ::
       Nil
+  }
 
   def resolve(analyzerContext: AnalyzerContext, plan: LogicalPlan): LogicalPlan = {
     val resolvedPlan = TypeResolver.typerRules

diff --git a/airframe-sql/src/test/scala/wvlet/airframe/sql/analyzer/TypeResolverTest.scala b/airframe-sql/src/test/scala/wvlet/airframe/sql/analyzer/TypeResolverTest.scala
@@ -226,6 +226,14 @@ class TypeResolverTest extends AirSpec {
         ResolvedAttribute("p2", DataType.StringType, Some(tableA), Some(a2))
       )
     }
+
+    test("parse WITH statement referenced in grouping indexes") {
+      val p = analyze("""WITH X AS (
+        |  SELECT id FROM A
+        |)
+        |SELECT max(id), COUNT(*)
+        |FROM X GROUP BY 1""".stripMargin)
+    }
   }
 
   test("resolve join attributes") {