add session variable to activate in-memory inner joins

Signed-off-by: Miguel Molina <[email protected]>

Author: erizocosmico

sql/plan/innerjoin.go

@@ -10,6 +10,7 @@ import (
 )
 
 const experimentalInMemoryJoinKey = "EXPERIMENTAL_IN_MEMORY_JOIN"
+const inMemoryJoinSessionVar = "inmemory_joins"
 
 var useInMemoryJoins = os.Getenv(experimentalInMemoryJoinKey) != ""
 
@@ -66,8 +67,14 @@ func (j *InnerJoin) RowIter(ctx *sql.Context) (sql.RowIter, error) {
 		return nil, err
 	}
 
+	var inMemorySession bool
+	_, val := ctx.Get(inMemoryJoinSessionVar)
+	if val != nil {
+		inMemorySession = true
+	}
+
 	var iter sql.RowIter
-	if useInMemoryJoins {
+	if useInMemoryJoins || inMemorySession {
 		r, err := j.Right.RowIter(ctx)
 		if err != nil {
 			span.Finish()

sql/plan/innerjoin_test.go

@@ -39,11 +39,6 @@ func TestInnerJoin(t *testing.T) {
 }
 
 func TestInMemoryInnerJoin(t *testing.T) {
-	useInMemoryJoins = true
-	defer func() {
-		useInMemoryJoins = false
-	}()
-
 	require := require.New(t)
 	finalSchema := append(lSchema, rSchema...)
 
@@ -62,7 +57,14 @@ func TestInMemoryInnerJoin(t *testing.T) {
 
 	require.Equal(finalSchema, j.Schema())
 
-	rows := collectRows(t, j)
+	ctx := sql.NewEmptyContext()
+	ctx.Set(inMemoryJoinSessionVar, sql.Text, "true")
+
+	iter, err := j.RowIter(ctx)
+	require.NoError(err)
+
+	rows, err := sql.RowIterToRows(iter)
+	require.NoError(err)
 	require.Len(rows, 2)
 
 	require.Equal([]sql.Row{