sql/analyzer: add rule to fix field's indexes in inner joins

Signed-off-by: Manuel Carmona <[email protected]>

Author: mcarmonaa

engine_test.go

@@ -206,6 +206,32 @@ func TestQueries(t *testing.T) {
 	}
 }
 
+func TestJoinsFields(t *testing.T) {
+	const (
+		query        = `SELECT * FROM tabletest, mytable mt INNER JOIN othertable ot ON mt.i = ot.i2`
+		expectedRows = 9
+	)
+
+	expectedSchema := sql.Schema{
+		&sql.Column{Name: "text", Type: sql.Text, Default: interface{}(nil), Nullable: false, Source: "tabletest"},
+		&sql.Column{Name: "number", Type: sql.Int32, Default: interface{}(nil), Nullable: false, Source: "tabletest"},
+		&sql.Column{Name: "i", Type: sql.Int64, Default: interface{}(nil), Nullable: false, Source: "mytable"},
+		&sql.Column{Name: "s", Type: sql.Text, Default: interface{}(nil), Nullable: false, Source: "mytable"},
+		&sql.Column{Name: "s2", Type: sql.Text, Default: interface{}(nil), Nullable: false, Source: "othertable"},
+		&sql.Column{Name: "i2", Type: sql.Int64, Default: interface{}(nil), Nullable: false, Source: "othertable"},
+	}
+
+	e := newEngine(t)
+	session := sql.NewEmptyContext()
+	schema, rowIter, err := e.Query(session, query)
+	require.NoError(t, err)
+	require.Exactly(t, expectedSchema, schema)
+
+	rows, err := sql.RowIterToRows(rowIter)
+	require.NoError(t, err)
+	require.Len(t, rows, expectedRows)
+}
+
 func TestOrderByColumns(t *testing.T) {
 	require := require.New(t)
 	e := newEngine(t)

sql/analyzer/analyzer_test.go

@@ -208,9 +208,9 @@ func TestAddRule(t *testing.T) {
 	require := require.New(t)
 
 	a := New(nil)
-	require.Len(a.Rules, 12)
+	require.Len(a.Rules, len(DefaultRules))
 	a.AddRule("foo", pushdown)
-	require.Len(a.Rules, 13)
+	require.Len(a.Rules, len(DefaultRules)+1)
 }
 
 func TestAddValidationRule(t *testing.T) {

sql/analyzer/rules.go

@@ -24,6 +24,7 @@ var DefaultRules = []Rule{
 	{"pushdown", pushdown},
 	{"optimize_distinct", optimizeDistinct},
 	{"erase_projection", eraseProjection},
+	{"fix_join_fields", fixJoinFields},
 }
 
 var (
@@ -809,3 +810,48 @@ func fixFieldIndexes(schema sql.Schema, exp sql.Expression) (sql.Expression, err
 		return e, nil
 	})
 }
+
+func fixJoinFields(ctx *sql.Context, a *Analyzer, node sql.Node) (sql.Node, error) {
+	span, ctx := ctx.Span("fix join fields")
+	defer span.Finish()
+
+	a.Log("fix join fields, node of type: %T", node)
+	return node.TransformUp(func(n sql.Node) (sql.Node, error) {
+		a.Log("transforming node of type: %T", n)
+		if !n.Resolved() {
+			return n, nil
+		}
+
+		_, ok := n.(*plan.InnerJoin)
+		if !ok {
+			return n, nil
+		}
+
+		return n.TransformExpressionsUp(func(e sql.Expression) (sql.Expression, error) {
+			a.Log("transforming expression of type: %T", e)
+			if !e.Resolved() {
+				return e, nil
+			}
+
+			field, ok := e.(*expression.GetField)
+			if !ok {
+				return e, nil
+			}
+
+			idx, ok := n.Schema().Contains(field.Name())
+			if !ok {
+				return nil, ErrColumnNotFound.New(field.Name())
+			}
+
+			fixedField := expression.NewGetField(
+				idx,
+				field.Type(),
+				field.Name(),
+				field.IsNullable(),
+			)
+
+			a.Log("fixed field for %T", n)
+			return fixedField, nil
+		})
+	})
+}

sql/type.go

@@ -63,14 +63,15 @@ func (s Schema) CheckRow(row Row) error {
 	return nil
 }
 
-// Contains returns whether the schema contains a column with the given name.
-func (s Schema) Contains(column string) bool {
-	for _, col := range s {
+// Contains returns whether the schema contains a column with the given name
+// and the index of that column in the schema.
+func (s Schema) Contains(column string) (int, bool) {
+	for i, col := range s {
 		if col.Name == column {
-			return true
+			return i, true
 		}
 	}
-	return false
+	return 0, false
 }
 
 // Equals checks whether the given schema is equal to this one.