custom row exec (dolthub#2593)

* support custom row execution operators * test panics for missing builders * fix strict key bug * more tests * more lookup join coverage * more tests
wk8 · Jul 23, 2024 · f685de1 · f685de1
1 parent 9dc0861
commit f685de1
Show file tree

Hide file tree

Showing 6 changed files with 235 additions and 3 deletions.
diff --git a/enginetest/join_op_tests.go b/enginetest/join_op_tests.go
@@ -133,6 +133,194 @@ var DefaultJoinOpTests = []joinOpTest{
 			},
 		},
 	},
+	{
+		name: "keyless lookup join indexes",
+		setup: [][]string{
+			setup.MydbData[0],
+			{
+				"CREATE table xy (x int, y int, z int, index y_idx(y));",
+				"CREATE table ab (a int primary key, b int, c int);",
+				"insert into xy values (1,0,3), (1,0,3), (0,2,1),(0,2,1);",
+				"insert into ab values (0,1,1), (1,2,2), (2,3,3), (3,2,2);",
+			},
+		},
+		tests: []JoinOpTests{
+			// covering tablescan
+			{
+				Query:    "select /*+ JOIN_ORDER(ab,xy) */ y,a from xy join ab on y = a",
+				Expected: []sql.Row{{0, 0}, {0, 0}, {2, 2}, {2, 2}},
+			},
+			{
+				Query:    "select /*+ JOIN_ORDER(xy,ab) */ y,a from xy join ab on y = a",
+				Expected: []sql.Row{{0, 0}, {0, 0}, {2, 2}, {2, 2}},
+			},
+			// covering indexed source
+			{
+				Query:    "select /*+ JOIN_ORDER(ab,xy) */ y,a from xy join ab on y = a where y = 2",
+				Expected: []sql.Row{{2, 2}, {2, 2}},
+			},
+			{
+				Query:    "select /*+ JOIN_ORDER(xy,ab) */ y,a from xy join ab on y = a where y = 2",
+				Expected: []sql.Row{{2, 2}, {2, 2}},
+			},
+			// non-covering tablescan
+			{
+				Query:    "select /*+ JOIN_ORDER(ab,xy) */ y,a,z from xy join ab on y = a",
+				Expected: []sql.Row{{0, 0, 3}, {0, 0, 3}, {2, 2, 1}, {2, 2, 1}},
+			},
+			{
+				Query:    "select /*+ JOIN_ORDER(xy,ab) */ y,a,z from xy join ab on y = a",
+				Expected: []sql.Row{{0, 0, 3}, {0, 0, 3}, {2, 2, 1}, {2, 2, 1}},
+			},
+			// non-covering indexed source
+			{
+				Query:    "select /*+ JOIN_ORDER(ab,xy) */ y,a,z from xy join ab on y = a where y = 2",
+				Expected: []sql.Row{{2, 2, 1}, {2, 2, 1}},
+			},
+			{
+				Query:    "select /*+ JOIN_ORDER(xy,ab) */ y,a,z from xy join ab on y = a where y = 2",
+				Expected: []sql.Row{{2, 2, 1}, {2, 2, 1}},
+			},
+		},
+	},
+	{
+		name: "keyed null lookup join indexes",
+		setup: [][]string{
+			setup.MydbData[0],
+			{
+				"CREATE table xy (x int, y int, z int primary key, index y_idx(y));",
+				"CREATE table ab (a int, b int primary key, c int);",
+				"insert into xy values (1,0,0), (1,null,1), (0,2,2),(0,2,3);",
+				"insert into ab values (0,1,0), (1,2,1), (2,3,2), (null,4,3);",
+			},
+		},
+		tests: []JoinOpTests{
+			// non-covering tablescan
+			{
+				Query:    "select /*+ JOIN_ORDER(ab,xy) */ y,a,x from xy join ab on y = a",
+				Expected: []sql.Row{{0, 0, 1}, {2, 2, 0}, {2, 2, 0}},
+			},
+			// covering
+			{
+				Query:    "select /*+ JOIN_ORDER(ab,xy) */ y,a,z from xy join ab on y = a",
+				Expected: []sql.Row{{0, 0, 0}, {2, 2, 2}, {2, 2, 3}},
+			},
+		},
+	},
+	{
+		name: "partial key null lookup join indexes",
+		setup: [][]string{
+			setup.MydbData[0],
+			{
+				"CREATE table xy (x int, y int, z int primary key, index y_idx(y,x));",
+				"CREATE table ab (a int, b int primary key, c int);",
+				"insert into xy values (1,0,0), (1,null,1), (0,2,2),(0,2,3);",
+				"insert into ab values (0,1,0), (1,2,1), (2,3,2), (null,4,3);",
+			},
+		},
+		tests: []JoinOpTests{
+			// non-covering tablescan
+			{
+				Query:    "select /*+ JOIN_ORDER(ab,xy) */ y,a,x from xy join ab on y = a",
+				Expected: []sql.Row{{0, 0, 1}, {2, 2, 0}, {2, 2, 0}},
+			},
+			// covering
+			{
+				Query:    "select /*+ JOIN_ORDER(ab,xy) */ y,a,z from xy join ab on y = a",
+				Expected: []sql.Row{{0, 0, 0}, {2, 2, 2}, {2, 2, 3}},
+			},
+		},
+	},
+	{
+		name: "keyed lookup join indexes",
+		setup: [][]string{
+			setup.MydbData[0],
+			{
+				"CREATE table xy (x int, y int, z int primary key, index y_idx(y));",
+				"CREATE table ab (a int, b int primary key, c int);",
+				"insert into xy values (1,0,0), (1,0,1), (0,2,2),(0,2,3);",
+				"insert into ab values (0,1,0), (1,2,1), (2,3,2), (3,4,3);",
+			},
+		},
+		tests: []JoinOpTests{
+			// covering tablescan
+			{
+				Query:    "select /*+ JOIN_ORDER(ab,xy) */ y,a from xy join ab on y = a",
+				Expected: []sql.Row{{0, 0}, {0, 0}, {2, 2}, {2, 2}},
+			},
+			{
+				Query:    "select /*+ JOIN_ORDER(xy,ab) */ y,a from xy join ab on y = a",
+				Expected: []sql.Row{{0, 0}, {0, 0}, {2, 2}, {2, 2}},
+			},
+			// covering indexed source
+			{
+				Query:    "select /*+ JOIN_ORDER(ab,xy) */ y,a from xy join ab on y = a where y = 2",
+				Expected: []sql.Row{{2, 2}, {2, 2}},
+			},
+			{
+				Query:    "select /*+ JOIN_ORDER(xy,ab) */ y,a from xy join ab on y = a where y = 2",
+				Expected: []sql.Row{{2, 2}, {2, 2}},
+			},
+			// non-covering tablescan
+			{
+				Query:    "select /*+ JOIN_ORDER(ab,xy) */ y,a,x from xy join ab on y = a",
+				Expected: []sql.Row{{0, 0, 1}, {0, 0, 1}, {2, 2, 0}, {2, 2, 0}},
+			},
+			{
+				Query:    "select /*+ JOIN_ORDER(xy,ab) */ y,a,x from xy join ab on y = a",
+				Expected: []sql.Row{{0, 0, 1}, {0, 0, 1}, {2, 2, 0}, {2, 2, 0}},
+			},
+			// non-covering indexed source
+			{
+				Query:    "select /*+ JOIN_ORDER(ab,xy) */ y,a,x from xy join ab on y = a where y = 2",
+				Expected: []sql.Row{{2, 2, 0}, {2, 2, 0}},
+			},
+			{
+				Query:    "select /*+ JOIN_ORDER(xy,ab) */ y,a,x from xy join ab on y = a where y = 2",
+				Expected: []sql.Row{{2, 2, 0}, {2, 2, 0}},
+			},
+		},
+	},
+	{
+		name: "multi pk lookup join indexes",
+		setup: [][]string{
+			setup.MydbData[0],
+			{
+				"CREATE table wxyz (w int, x int, y int, z int, primary key (x,w), index yw_idx(y,w));",
+				"CREATE table abcd (a int, b int, c int, d int, primary key (a,b), index ca_idx(c,a));",
+				"insert into wxyz values (1,0,0,0), (1,1,1,1), (0,2,2,1),(0,1,3,1);",
+				"insert into abcd values (0,0,0,0), (0,1,1,1), (0,2,2,1),(2,1,3,1);",
+			},
+		},
+		tests: []JoinOpTests{
+			{
+				Query:    "select /*+ JOIN_ORDER(abcd,wxyz) */ y,a,z from wxyz join abcd on y = a",
+				Expected: []sql.Row{{0, 0, 0}, {0, 0, 0}, {0, 0, 0}, {2, 2, 1}},
+			},
+			{
+				Query:    "select /*+ JOIN_ORDER(abcd,wxyz) */ y,a,w from wxyz join abcd on y = a",
+				Expected: []sql.Row{{0, 0, 1}, {0, 0, 1}, {0, 0, 1}, {2, 2, 0}},
+			},
+		},
+	},
+	{
+		name: "redundant keyless index",
+		setup: [][]string{
+			setup.MydbData[0],
+			{
+				"CREATE table xy (x int, y int, z int, index y_idx(x,y,z));",
+				"CREATE table ab (a int, b int primary key, c int);",
+				"insert into xy values (1,0,0), (1,0,1), (0,2,2),(0,2,3);",
+				"insert into ab values (0,1,0), (1,2,1), (2,3,2), (3,4,3);",
+			},
+		},
+		tests: []JoinOpTests{
+			{
+				Query:    "select /*+ JOIN_ORDER(ab,xy) */ y,a,z from xy join ab on y = a",
+				Expected: []sql.Row{{0, 0, 1}, {0, 0, 0}, {2, 2, 3}, {2, 2, 2}},
+			},
+		},
+	},
 	{
 		name: "issue 5633, nil comparison in merge join",
 		setup: [][]string{

diff --git a/sql/plan/indexed_table_access.go b/sql/plan/indexed_table_access.go
@@ -276,6 +276,31 @@ func (i *IndexedTableAccess) CanBuildIndex(ctx *sql.Context) (bool, error) {
 	return err == nil && !lookup.IsEmpty(), nil
 }
 
+func (i *IndexedTableAccess) IsStrictLookup() bool {
+	if !i.lb.index.IsUnique() {
+		return false
+	}
+	for _, m := range i.lb.matchesNullMask {
+		if m {
+			return false
+		}
+	}
+	if len(i.lb.keyExprs) != len(i.lb.index.Expressions()) {
+		// only partial key
+		return false
+	}
+	if strings.EqualFold(i.lb.index.ID(), "primary") {
+		return true
+	}
+	for _, e := range i.lb.keyExprs {
+		if e.IsNullable() {
+			// nullable key may not be
+			return false
+		}
+	}
+	return true
+}
+
 func (i *IndexedTableAccess) GetLookup(ctx *sql.Context, row sql.Row) (sql.IndexLookup, error) {
 	// if the lookup was provided at analysis time (static evaluation), use it.
 	if !i.lookup.IsEmpty() {

diff --git a/sql/rowexec/builder.go b/sql/rowexec/builder.go
@@ -29,9 +29,16 @@ type ExecBuilderFunc func(ctx *sql.Context, n sql.Node, r sql.Row) (sql.RowIter,
 // BaseBuilder converts a plan tree into a RowIter tree. All relational nodes
 // have a build statement. Custom source nodes that provide rows that implement
 // sql.ExecSourceRel are also built into the tree.
-type BaseBuilder struct{}
+type BaseBuilder struct {
+	// if override is provided, we try to build executor with this first
+	override sql.NodeExecBuilder
+}
 
 func (b *BaseBuilder) Build(ctx *sql.Context, n sql.Node, r sql.Row) (sql.RowIter, error) {
 	defer trace.StartRegion(ctx, "ExecBuilder.Build").End()
 	return b.buildNodeExec(ctx, n, r)
 }
+
+func NewOverrideBuilder(override sql.NodeExecBuilder) sql.NodeExecBuilder {
+	return &BaseBuilder{override: override}
+}
diff --git a/sql/rowexec/dml.go b/sql/rowexec/dml.go
@@ -293,6 +293,7 @@ func (b *BaseBuilder) buildTriggerBeginEndBlock(ctx *sql.Context, n *plan.Trigge
 		statements: n.Children(),
 		row:        row,
 		once:       &sync.Once{},
+		b:          b,
 	}, nil
 }
 
@@ -308,6 +309,7 @@ func (b *BaseBuilder) buildTriggerExecutor(ctx *sql.Context, n *plan.TriggerExec
 		triggerEvent:   n.TriggerEvent,
 		executionLogic: n.Right(),
 		ctx:            ctx,
+		b:              b,
 	}, nil
 }
 

diff --git a/sql/rowexec/node_builder.gen.go b/sql/rowexec/node_builder.gen.go
@@ -24,7 +24,17 @@ import (
 )
 
 func (b *BaseBuilder) buildNodeExec(ctx *sql.Context, n sql.Node, row sql.Row) (sql.RowIter, error) {
-	iter, err := b.buildNodeExecNoAnalyze(ctx, n, row)
+	var iter sql.RowIter
+	var err error
+	if b.override != nil {
+		iter, err = b.override.Build(ctx, n, row)
+	}
+	if err != nil {
+		return nil, err
+	}
+	if iter == nil {
+		iter, err = b.buildNodeExecNoAnalyze(ctx, n, row)
+	}
 	if err != nil {
 		return nil, err
 	}

diff --git a/sql/rowexec/proc.go b/sql/rowexec/proc.go
@@ -345,7 +345,7 @@ func (b *BaseBuilder) buildElseCaseError(ctx *sql.Context, n plan.ElseCaseError,
 }
 
 func (b *BaseBuilder) buildOpen(ctx *sql.Context, n *plan.Open, row sql.Row) (sql.RowIter, error) {
-	return &openIter{pRef: n.Pref, name: n.Name, row: row}, nil
+	return &openIter{pRef: n.Pref, name: n.Name, row: row, b: b}, nil
 }
 
 func (b *BaseBuilder) buildClose(ctx *sql.Context, n *plan.Close, row sql.Row) (sql.RowIter, error) {