gen4: Add hash join primitive and planning

go/test/endtoend/vtgate/gen4/column_name_test.go

@@ -20,6 +20,8 @@ import (
 	"context"
 	"testing"
 
+	"vitess.io/vitess/go/test/endtoend/vtgate/utils"
+
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
@@ -35,12 +37,10 @@ func TestColumnNames(t *testing.T) {
 	require.NoError(t, err)
 	defer conn.Close()
 
-	_, err = exec(t, conn, "create table uks.t2(id bigint,phone bigint,msg varchar(100),primary key(id)) Engine=InnoDB")
-	require.NoError(t, err)
-	defer exec(t, conn, "drop table uks.t2")
+	utils.Exec(t, conn, "create table uks.t2(id bigint,phone bigint,msg varchar(100),primary key(id)) Engine=InnoDB")
+	defer utils.Exec(t, conn, "drop table uks.t2")
 
-	qr, err := exec(t, conn, "SELECT t1.id as t1id, t2.id as t2id, t2.phone as t2phn FROM ks.t1 cross join uks.t2 where t1.id = t2.id ORDER BY t2.phone")
-	require.NoError(t, err)
+	qr := utils.Exec(t, conn, "SELECT t1.id as t1id, t2.id as t2id, t2.phone as t2phn FROM ks.t1 cross join uks.t2 where t1.id = t2.id ORDER BY t2.phone")
 
 	assert.Equal(t, 3, len(qr.Fields))
 	assert.Equal(t, "t1id", qr.Fields[0].Name)

go/test/endtoend/vtgate/gen4/system_schema_test.go

@@ -21,6 +21,8 @@ import (
 	"fmt"
 	"testing"
 
+	"vitess.io/vitess/go/test/endtoend/vtgate/utils"
+
 	"github.com/stretchr/testify/assert"
 
 	"github.com/stretchr/testify/require"
@@ -84,7 +86,7 @@ func TestInformationSchemaWithSubquery(t *testing.T) {
 	require.NoError(t, err)
 	defer conn.Close()
 
-	result := checkedExec(t, conn, "SELECT column_name FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_SCHEMA = (SELECT SCHEMA()) AND TABLE_NAME = 'not_exists'")
+	result := utils.Exec(t, conn, "SELECT column_name FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_SCHEMA = (SELECT SCHEMA()) AND TABLE_NAME = 'not_exists'")
 	assert.Empty(t, result.Rows)
 }
 
@@ -95,8 +97,8 @@ func TestInformationSchemaQueryGetsRoutedToTheRightTableAndKeyspace(t *testing.T
 	require.NoError(t, err)
 	defer conn.Close()
 
-	_ = checkedExec(t, conn, "SELECT id FROM ks.t1000") // test that the routed table is available to us
-	result := checkedExec(t, conn, "SELECT * FROM information_schema.tables WHERE table_schema = database() and table_name='ks.t1000'")
+	_ = utils.Exec(t, conn, "SELECT id FROM ks.t1000") // test that the routed table is available to us
+	result := utils.Exec(t, conn, "SELECT * FROM information_schema.tables WHERE table_schema = database() and table_name='ks.t1000'")
 	assert.NotEmpty(t, result.Rows)
 }
 
@@ -107,12 +109,12 @@ func TestFKConstraintUsingInformationSchema(t *testing.T) {
 	require.NoError(t, err)
 	defer conn.Close()
 
-	checkedExec(t, conn, "create table ks.t7_xxhash(uid varchar(50),phone bigint,msg varchar(100),primary key(uid)) Engine=InnoDB")
-	checkedExec(t, conn, "create table ks.t7_fk(id bigint,t7_uid varchar(50),primary key(id),CONSTRAINT t7_fk_ibfk_1 foreign key (t7_uid) references t7_xxhash(uid)	on delete set null on update cascade) Engine=InnoDB;")
-	defer checkedExec(t, conn, "drop table ks.t7_fk, ks.t7_xxhash")
+	utils.Exec(t, conn, "create table ks.t7_xxhash(uid varchar(50),phone bigint,msg varchar(100),primary key(uid)) Engine=InnoDB")
+	utils.Exec(t, conn, "create table ks.t7_fk(id bigint,t7_uid varchar(50),primary key(id),CONSTRAINT t7_fk_ibfk_1 foreign key (t7_uid) references t7_xxhash(uid)	on delete set null on update cascade) Engine=InnoDB;")
+	defer utils.Exec(t, conn, "drop table ks.t7_fk, ks.t7_xxhash")
 
 	query := "select fk.referenced_table_name as to_table, fk.referenced_column_name as primary_key, fk.column_name as `column`, fk.constraint_name as name, rc.update_rule as on_update, rc.delete_rule as on_delete from information_schema.referential_constraints as rc join information_schema.key_column_usage as fk using (constraint_schema, constraint_name) where fk.referenced_column_name is not null and fk.table_schema = database() and fk.table_name = 't7_fk' and rc.constraint_schema = database() and rc.table_name = 't7_fk'"
-	assertMatches(t, conn, query, `[[VARCHAR("t7_xxhash") VARCHAR("uid") VARCHAR("t7_uid") VARCHAR("t7_fk_ibfk_1") VARCHAR("CASCADE") VARCHAR("SET NULL")]]`)
+	utils.AssertMatches(t, conn, query, `[[VARCHAR("t7_xxhash") VARCHAR("uid") VARCHAR("t7_uid") VARCHAR("t7_fk_ibfk_1") VARCHAR("CASCADE") VARCHAR("SET NULL")]]`)
 }
 
 func TestConnectWithSystemSchema(t *testing.T) {
@@ -123,7 +125,7 @@ func TestConnectWithSystemSchema(t *testing.T) {
 		connParams.DbName = dbname
 		conn, err := mysql.Connect(ctx, &connParams)
 		require.NoError(t, err)
-		checkedExec(t, conn, `select @@max_allowed_packet from dual`)
+		utils.Exec(t, conn, `select @@max_allowed_packet from dual`)
 		conn.Close()
 	}
 }
@@ -135,8 +137,8 @@ func TestUseSystemSchema(t *testing.T) {
 	require.NoError(t, err)
 	defer conn.Close()
 	for _, dbname := range []string{"information_schema", "mysql", "performance_schema", "sys"} {
-		checkedExec(t, conn, fmt.Sprintf("use %s", dbname))
-		checkedExec(t, conn, `select @@max_allowed_packet from dual`)
+		utils.Exec(t, conn, fmt.Sprintf("use %s", dbname))
+		utils.Exec(t, conn, `select @@max_allowed_packet from dual`)
 	}
 }
 
@@ -153,16 +155,16 @@ func TestSystemSchemaQueryWithoutQualifier(t *testing.T) {
 		"on c.table_schema = t.table_schema and c.table_name = t.table_name "+
 		"where t.table_schema = '%s' and c.table_schema = '%s' "+
 		"order by t.table_schema,t.table_name,c.column_name", shardedKs, shardedKs)
-	qr1 := checkedExec(t, conn, queryWithQualifier)
+	qr1 := utils.Exec(t, conn, queryWithQualifier)
 
-	checkedExec(t, conn, "use information_schema")
+	utils.Exec(t, conn, "use information_schema")
 	queryWithoutQualifier := fmt.Sprintf("select t.table_schema,t.table_name,c.column_name,c.column_type "+
 		"from tables t "+
 		"join columns c "+
 		"on c.table_schema = t.table_schema and c.table_name = t.table_name "+
 		"where t.table_schema = '%s' and c.table_schema = '%s' "+
 		"order by t.table_schema,t.table_name,c.column_name", shardedKs, shardedKs)
-	qr2 := checkedExec(t, conn, queryWithoutQualifier)
+	qr2 := utils.Exec(t, conn, queryWithoutQualifier)
 	require.Equal(t, qr1, qr2)
 
 	connParams := vtParams
@@ -171,7 +173,7 @@ func TestSystemSchemaQueryWithoutQualifier(t *testing.T) {
 	require.NoError(t, err)
 	defer conn2.Close()
 
-	qr3 := checkedExec(t, conn2, queryWithoutQualifier)
+	qr3 := utils.Exec(t, conn2, queryWithoutQualifier)
 	require.Equal(t, qr2, qr3)
 }
 
@@ -187,7 +189,7 @@ func TestMultipleSchemaPredicates(t *testing.T) {
 		"join information_schema.columns c "+
 		"on c.table_schema = t.table_schema and c.table_name = t.table_name "+
 		"where t.table_schema = '%s' and c.table_schema = '%s' and c.table_schema = '%s' and c.table_schema = '%s'", shardedKs, shardedKs, shardedKs, shardedKs)
-	qr1 := checkedExec(t, conn, query)
+	qr1 := utils.Exec(t, conn, query)
 	require.EqualValues(t, 4, len(qr1.Fields))
 
 	// test a query with two keyspace names

go/vt/vtgate/engine/hash_join.go

@@ -0,0 +1,236 @@
+/*
+Copyright 2021 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package engine
+
+import (
+	"fmt"
+	"strings"
+
+	"vitess.io/vitess/go/mysql/collations"
+	"vitess.io/vitess/go/sqltypes"
+	querypb "vitess.io/vitess/go/vt/proto/query"
+	"vitess.io/vitess/go/vt/sqlparser"
+	"vitess.io/vitess/go/vt/vtgate/evalengine"
+)
+
+var _ Primitive = (*HashJoin)(nil)
+
+// HashJoin specifies the parameters for a join primitive.
+type HashJoin struct {
+	Opcode JoinOpcode
+
+	// Left and Right are the LHS and RHS primitives
+	// of the Join. They can be any primitive.
+	Left, Right Primitive `json:",omitempty"`
+
+	// Cols defines which columns from the left
+	// or right results should be used to build the
+	// return result. For results coming from the
+	// left query, the index values go as -1, -2, etc.
+	// For the right query, they're 1, 2, etc.
+	// If Cols is {-1, -2, 1, 2}, it means that
+	// the returned result will be {Left0, Left1, Right0, Right1}.
+	Cols []int `json:",omitempty"`
+
+	// The keys correspond to the column offset in the inputs where
+	// the join columns can be found
+	LHSKey, RHSKey int
+
+	ASTPred sqlparser.Expr
+}
+
+type hashKey = int64
+
+// TryExecute implements the Primitive interface
+func (hj *HashJoin) TryExecute(vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool) (*sqltypes.Result, error) {
+	lresult, err := vcursor.ExecutePrimitive(hj.Left, bindVars, wantfields)
+	if err != nil {
+		return nil, err
+	}
+
+	// build the probe table from the LHS result
+	probeTable := map[hashKey][]row{}
+	for _, current := range lresult.Rows {
+		joinVal := current[hj.LHSKey]
+		if joinVal.IsNull() {
+			continue
+		}
+		hashcode, err := evalengine.NullsafeHashcode(joinVal)
+		if err != nil {
+			return nil, err
+		}
+		probeTable[hashcode] = append(probeTable[hashcode], current)
+	}
+
+	rresult, err := vcursor.ExecutePrimitive(hj.Right, bindVars, wantfields)
+	if err != nil {
+		return nil, err
+	}
+
+	result := &sqltypes.Result{
+		Fields: joinFields(lresult.Fields, rresult.Fields, hj.Cols),
+	}
+
+	for _, currentRHSRow := range rresult.Rows {
+		joinVal := currentRHSRow[hj.RHSKey]
+		if joinVal.IsNull() {
+			continue
+		}
+		hashcode, err := evalengine.NullsafeHashcode(joinVal)
+		if err != nil {
+			return nil, err
+		}
+		lftRows := probeTable[hashcode]
+		for _, currentLHSRow := range lftRows {
+			lhsVal := currentLHSRow[hj.LHSKey]
+			// hash codes can give false positives, so we need to check with a real comparison as well
+			cmp, err := evalengine.NullsafeCompare(joinVal, lhsVal, collations.Unknown)
+			if err != nil {
+				return nil, err
+			}
+
+			if cmp == 0 {
+				// we have a match!
+				result.Rows = append(result.Rows, joinRows(currentLHSRow, currentRHSRow, hj.Cols))
+			}
+		}
+	}
+
+	return result, nil
+}
+
+// TryStreamExecute implements the Primitive interface
+func (hj *HashJoin) TryStreamExecute(vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool, callback func(*sqltypes.Result) error) error {
+	// build the probe table from the LHS result
+	probeTable := map[hashKey][]row{}
+	var lfields []*querypb.Field
+	err := vcursor.StreamExecutePrimitive(hj.Left, bindVars, wantfields, func(result *sqltypes.Result) error {
+		if len(lfields) == 0 && len(result.Fields) != 0 {
+			lfields = result.Fields
+		}
+		for _, current := range result.Rows {
+			joinVal := current[hj.LHSKey]
+			if joinVal.IsNull() {
+				continue
+			}
+			hashcode, err := evalengine.NullsafeHashcode(joinVal)
+			if err != nil {
+				return err
+			}
+			probeTable[hashcode] = append(probeTable[hashcode], current)
+		}
+		return nil
+	})
+	if err != nil {
+		return err
+	}
+
+	return vcursor.StreamExecutePrimitive(hj.Right, bindVars, wantfields, func(result *sqltypes.Result) error {
+		res := &sqltypes.Result{}
+		if len(result.Fields) != 0 {
+			res = &sqltypes.Result{
+				Fields: joinFields(lfields, result.Fields, hj.Cols),
+			}
+		}
+		for _, currentRHSRow := range result.Rows {
+			joinVal := currentRHSRow[hj.RHSKey]
+			if joinVal.IsNull() {
+				continue
+			}
+			hashcode, err := evalengine.NullsafeHashcode(joinVal)
+			if err != nil {
+				return err
+			}
+			lftRows := probeTable[hashcode]
+			for _, currentLHSRow := range lftRows {
+				lhsVal := currentLHSRow[hj.LHSKey]
+				// hash codes can give false positives, so we need to check with a real comparison as well
+				cmp, err := evalengine.NullsafeCompare(joinVal, lhsVal, collations.Unknown)
+				if err != nil {
+					return err
+				}
+
+				if cmp == 0 {
+					// we have a match!
+					res.Rows = append(res.Rows, joinRows(currentLHSRow, currentRHSRow, hj.Cols))
+				}
+			}
+		}
+		if len(res.Rows) != 0 || len(res.Fields) != 0 {
+			return callback(res)
+		}
+		return nil
+	})
+}
+
+// RouteType implements the Primitive interface
+func (hj *HashJoin) RouteType() string {
+	return "HashJoin"
+}
+
+// GetKeyspaceName implements the Primitive interface
+func (hj *HashJoin) GetKeyspaceName() string {
+	if hj.Left.GetKeyspaceName() == hj.Right.GetKeyspaceName() {
+		return hj.Left.GetKeyspaceName()
+	}
+	return hj.Left.GetKeyspaceName() + "_" + hj.Right.GetKeyspaceName()
+}
+
+// GetTableName implements the Primitive interface
+func (hj *HashJoin) GetTableName() string {
+	return hj.Left.GetTableName() + "_" + hj.Right.GetTableName()
+}
+
+// GetFields implements the Primitive interface
+func (hj *HashJoin) GetFields(vcursor VCursor, bindVars map[string]*querypb.BindVariable) (*sqltypes.Result, error) {
+	joinVars := make(map[string]*querypb.BindVariable)
+	lresult, err := hj.Left.GetFields(vcursor, bindVars)
+	if err != nil {
+		return nil, err
+	}
+	result := &sqltypes.Result{}
+	rresult, err := hj.Right.GetFields(vcursor, combineVars(bindVars, joinVars))
+	if err != nil {
+		return nil, err
+	}
+	result.Fields = joinFields(lresult.Fields, rresult.Fields, hj.Cols)
+	return result, nil
+}
+
+// NeedsTransaction implements the Primitive interface
+func (hj *HashJoin) NeedsTransaction() bool {
+	return hj.Right.NeedsTransaction() || hj.Left.NeedsTransaction()
+}
+
+// Inputs implements the Primitive interface
+func (hj *HashJoin) Inputs() []Primitive {
+	return []Primitive{hj.Left, hj.Right}
+}
+
+// description implements the Primitive interface
+func (hj *HashJoin) description() PrimitiveDescription {
+	other := map[string]interface{}{
+		"TableName":         hj.GetTableName(),
+		"JoinColumnIndexes": strings.Trim(strings.Join(strings.Fields(fmt.Sprint(hj.Cols)), ","), "[]"),
+		"Predicate":         sqlparser.String(hj.ASTPred),
+	}
+	return PrimitiveDescription{
+		OperatorType: "Join",
+		Variant:      "Hash" + hj.Opcode.String(),
+		Other:        other,
+	}
+}