inspect all aliases

Signed-off-by: kuba-- <[email protected]>

Author: kuba--

engine_test.go

@@ -1638,22 +1638,32 @@ func TestIndexes(t *testing.T) {
 
 	_, _, err = e.Query(
 		newCtx(),
-		"CREATE INDEX myidx ON mytable USING pilosa (i) WITH (async = false)",
+		"CREATE INDEX idx_i ON mytable USING pilosa (i) WITH (async = false)",
 	)
 	require.NoError(t, err)
 
 	_, _, err = e.Query(
 		newCtx(),
-		"CREATE INDEX myidx_multi ON mytable USING pilosa (i, s) WITH (async = false)",
+		"CREATE INDEX idx_s ON mytable USING pilosa (s) WITH (async = false)",
+	)
+	require.NoError(t, err)
+
+	_, _, err = e.Query(
+		newCtx(),
+		"CREATE INDEX idx_is ON mytable USING pilosa (i, s) WITH (async = false)",
 	)
 	require.NoError(t, err)
 
 	defer func() {
-		done, err := e.Catalog.DeleteIndex("mydb", "myidx", true)
+		done, err := e.Catalog.DeleteIndex("mydb", "idx_i", true)
 		require.NoError(t, err)
 		<-done
 
-		done, err = e.Catalog.DeleteIndex("foo", "myidx_multi", true)
+		done, err = e.Catalog.DeleteIndex("mydb", "idx_s", true)
+		require.NoError(t, err)
+		<-done
+
+		done, err = e.Catalog.DeleteIndex("foo", "idx_is", true)
 		require.NoError(t, err)
 		<-done
 	}()
@@ -1743,6 +1753,25 @@ func TestIndexes(t *testing.T) {
 				{int64(2), "second row"},
 			},
 		},
+		{
+			"SELECT s, SUBSTRING(s, 1, 1) AS sub_s FROM mytable WHERE sub_s = 's'",
+			[]sql.Row{
+				{"second row", "s"},
+			},
+		},
+		{
+			"SELECT count(i) AS mytable_i, SUBSTR(s, -3) AS mytable_s FROM mytable WHERE i > 0 AND mytable_s='row' GROUP BY mytable_s",
+			[]sql.Row{
+				{int32(3), "row"},
+			},
+		},
+		{
+			"SELECT mytable_i FROM (SELECT i AS mytable_i FROM mytable) as t WHERE mytable_i > 1",
+			[]sql.Row{
+				{int64(2)},
+				{int64(3)},
+			},
+		},
 	}
 
 	for _, tt := range testCases {

sql/analyzer/assign_indexes.go

@@ -35,22 +35,40 @@ func assignIndexes(a *Analyzer, node sql.Node) (map[string]*indexLookup, error)
 		}
 	}()
 
-	var err error
-	plan.Inspect(node, func(node sql.Node) bool {
-		filter, ok := node.(*plan.Filter)
-		if !ok {
+	aliases := make(map[string]sql.Expression)
+	var (
+		err error
+		fn  func(node sql.Node) bool
+	)
+	fn = func(n sql.Node) bool {
+		if n == nil {
 			return true
 		}
 
-		aliases := make(map[string]sql.Expression)
-		if prj, ok := filter.Child.(*plan.Project); ok {
+		if prj, ok := n.(*plan.Project); ok {
 			for _, ex := range prj.Expressions() {
 				if alias, ok := ex.(*expression.Alias); ok {
-					aliases[alias.Name()] = alias.Child
+					if _, ok := aliases[alias.Name()]; !ok {
+						aliases[alias.Name()] = alias.Child
+					}
 				}
 			}
+		} else {
+			for _, ch := range n.Children() {
+				plan.Inspect(ch, fn)
+			}
 		}
 
+		return true
+	}
+
+	plan.Inspect(node, func(node sql.Node) bool {
+		filter, ok := node.(*plan.Filter)
+		if !ok {
+			return true
+		}
+		fn(filter.Child)
+
 		var result map[string]*indexLookup
 		result, err = getIndexes(filter.Expression, aliases, a)
 		if err != nil {
@@ -276,8 +294,13 @@ func unifyExpressions(aliases map[string]sql.Expression, expr ...sql.Expression)
 
 	for i, e := range expr {
 		uex := e
-		if aliases != nil {
-			if alias, ok := aliases[e.String()]; ok {
+		name := e.String()
+		if n, ok := e.(sql.Nameable); ok {
+			name = n.Name()
+		}
+
+		if aliases != nil && len(aliases) > 0 {
+			if alias, ok := aliases[name]; ok {
 				uex = alias
 			}
 		}