[kvexec] fix more lookup bugs related to schema/projection inconsiste…

…ncies (#8311) * [kvexec[ fix more lookup bugs related to schema/projection inconsistencies * cleanup * [ga-format-pr] Run go/utils/repofmt/format_repo.sh and go/Godeps/update.sh * bump * bump --------- Co-authored-by: max-hoffman <[email protected]>
dolthub · Aug 30, 2024 · d2ac2de · d2ac2de
1 parent b4475a8
commit d2ac2de
Show file tree

Hide file tree

Showing 5 changed files with 38 additions and 11 deletions.
diff --git a/go/go.mod b/go/go.mod
@@ -57,7 +57,7 @@ require (
 	github.com/cespare/xxhash/v2 v2.2.0
 	github.com/creasty/defaults v1.6.0
 	github.com/dolthub/flatbuffers/v23 v23.3.3-dh.2
-	github.com/dolthub/go-mysql-server v0.18.2-0.20240827225330-65d79910212e
+	github.com/dolthub/go-mysql-server v0.18.2-0.20240829233246-ed8de8d3a4e6
 	github.com/dolthub/gozstd v0.0.0-20240423170813-23a2903bca63
 	github.com/dolthub/swiss v0.1.0
 	github.com/goccy/go-json v0.10.2

diff --git a/go/go.sum b/go/go.sum
@@ -183,8 +183,8 @@ github.com/dolthub/fslock v0.0.3 h1:iLMpUIvJKMKm92+N1fmHVdxJP5NdyDK5bK7z7Ba2s2U=
 github.com/dolthub/fslock v0.0.3/go.mod h1:QWql+P17oAAMLnL4HGB5tiovtDuAjdDTPbuqx7bYfa0=
 github.com/dolthub/go-icu-regex v0.0.0-20230524105445-af7e7991c97e h1:kPsT4a47cw1+y/N5SSCkma7FhAPw7KeGmD6c9PBZW9Y=
 github.com/dolthub/go-icu-regex v0.0.0-20230524105445-af7e7991c97e/go.mod h1:KPUcpx070QOfJK1gNe0zx4pA5sicIK1GMikIGLKC168=
-github.com/dolthub/go-mysql-server v0.18.2-0.20240827225330-65d79910212e h1:Fl9cRq6WCRusdbXDNl+iO/1NxVuVHRF7BPh0p0OAhLc=
-github.com/dolthub/go-mysql-server v0.18.2-0.20240827225330-65d79910212e/go.mod h1:nbdOzd0ceWONE80vbfwoRBjut7z3CIj69ZgDF/cKuaA=
+github.com/dolthub/go-mysql-server v0.18.2-0.20240829233246-ed8de8d3a4e6 h1:BcCI8gy6XyhgUZwxMALsHlhyRJR/NuSJ2XGf/CIaAjo=
+github.com/dolthub/go-mysql-server v0.18.2-0.20240829233246-ed8de8d3a4e6/go.mod h1:nbdOzd0ceWONE80vbfwoRBjut7z3CIj69ZgDF/cKuaA=
 github.com/dolthub/gozstd v0.0.0-20240423170813-23a2903bca63 h1:OAsXLAPL4du6tfbBgK0xXHZkOlos63RdKYS3Sgw/dfI=
 github.com/dolthub/gozstd v0.0.0-20240423170813-23a2903bca63/go.mod h1:lV7lUeuDhH5thVGDCKXbatwKy2KW80L4rMT46n+Y2/Q=
 github.com/dolthub/ishell v0.0.0-20240701202509-2b217167d718 h1:lT7hE5k+0nkBdj/1UOSFwjWpNxf+LCApbRHgnCA17XE=

diff --git a/go/libraries/doltcore/sqle/index/index_reader.go b/go/libraries/doltcore/sqle/index/index_reader.go
@@ -25,6 +25,7 @@ import (
 	"github.com/dolthub/dolt/go/libraries/doltcore/doltdb"
 	"github.com/dolthub/dolt/go/libraries/doltcore/doltdb/durable"
 	"github.com/dolthub/dolt/go/libraries/doltcore/row"
+	"github.com/dolthub/dolt/go/libraries/doltcore/schema"
 	"github.com/dolthub/dolt/go/libraries/doltcore/table/typed/noms"
 	"github.com/dolthub/dolt/go/store/prolly"
 	"github.com/dolthub/dolt/go/store/prolly/tree"
@@ -249,6 +250,9 @@ type IndexScanBuilder interface {
 
 	// Key returns the table root for caching purposes
 	Key() doltdb.DataCacheKey
+
+	// OutputSchema returns the output KV tuple schema
+	OutputSchema() schema.Schema
 }
 
 func NewIndexReaderBuilder(
@@ -374,6 +378,10 @@ func (ib *baseIndexImplBuilder) Key() doltdb.DataCacheKey {
 	return ib.key
 }
 
+func (ib *baseIndexImplBuilder) OutputSchema() schema.Schema {
+	return ib.idx.IndexSchema()
+}
+
 // NewPartitionRowIter implements IndexScanBuilder
 func (ib *baseIndexImplBuilder) NewPartitionRowIter(_ *sql.Context, _ sql.Partition) (sql.RowIter, error) {
 	panic("cannot call NewRowIter on baseIndexImplBuilder")
@@ -479,6 +487,10 @@ func (i *sequenceRangeIter) Next(ctx context.Context) (val.Tuple, val.Tuple, err
 	return k, v, nil
 }
 
+func (ib *coveringIndexImplBuilder) OutputSchema() schema.Schema {
+	return ib.idx.IndexSchema()
+}
+
 // NewRangeMapIter implements IndexScanBuilder
 func (ib *coveringIndexImplBuilder) NewRangeMapIter(ctx context.Context, r prolly.Range, reverse bool) (prolly.MapIter, error) {
 	if reverse {
@@ -562,6 +574,11 @@ func (i *nonCoveringMapIter) Next(ctx context.Context) (val.Tuple, val.Tuple, er
 	return pk, value, nil
 }
 
+func (ib *nonCoveringIndexImplBuilder) OutputSchema() schema.Schema {
+	// this refs the table schema
+	return ib.baseIndexImplBuilder.idx.Schema()
+}
+
 // NewRangeMapIter implements IndexScanBuilder
 func (ib *nonCoveringIndexImplBuilder) NewRangeMapIter(ctx context.Context, r prolly.Range, reverse bool) (prolly.MapIter, error) {
 	var secIter prolly.MapIter
@@ -619,6 +636,10 @@ type keylessIndexImplBuilder struct {
 	s *durableIndexState
 }
 
+func (ib *keylessIndexImplBuilder) OutputSchema() schema.Schema {
+	return ib.idx.Schema()
+}
+
 // NewRangeMapIter implements IndexScanBuilder
 func (ib *keylessIndexImplBuilder) NewRangeMapIter(ctx context.Context, r prolly.Range, reverse bool) (prolly.MapIter, error) {
 	rows := ib.s.Primary
@@ -709,6 +730,10 @@ type nomsIndexImplBuilder struct {
 	s *durableIndexState
 }
 
+func (ib *nomsIndexImplBuilder) OutputSchema() schema.Schema {
+	return ib.baseIndexImplBuilder.idx.Schema()
+}
+
 // NewPartitionRowIter implements IndexScanBuilder
 func (ib *nomsIndexImplBuilder) NewPartitionRowIter(ctx *sql.Context, part sql.Partition) (sql.RowIter, error) {
 	p := part.(rangePartition)

diff --git a/go/libraries/doltcore/sqle/kvexec/builder.go b/go/libraries/doltcore/sqle/kvexec/builder.go
@@ -43,7 +43,7 @@ func (b Builder) Build(ctx *sql.Context, n sql.Node, r sql.Row) (sql.RowIter, er
 			if ita, ok := getIta(n.Right()); ok && len(r) == 0 && simpleLookupExpressions(ita.Expressions()) {
 				if _, _, dstIter, _, dstTags, dstFilter, err := getSourceKv(ctx, n.Right(), false); err == nil && dstIter != nil {
 					if srcMap, srcIter, _, srcSchema, srcTags, srcFilter, err := getSourceKv(ctx, n.Left(), true); err == nil && srcSchema != nil {
-						if keyLookupMapper := newLookupKeyMapping(ctx, srcSchema, srcMap, dstIter.InputKeyDesc(), ita.Expressions()); keyLookupMapper.valid() {
+						if keyLookupMapper := newLookupKeyMapping(ctx, srcSchema, dstIter.InputKeyDesc(), ita.Expressions(), srcMap.NodeStore()); keyLookupMapper.valid() {
 							// conditions:
 							// (1) lookup or left lookup join
 							// (2) left-side is something we read KVs from (table or indexscan, ex: no subqueries)
@@ -315,6 +315,8 @@ func getSourceKv(ctx *sql.Context, n sql.Node, isSrc bool) (prolly.Map, prolly.M
 		}
 		indexMap = durable.ProllyMapFromIndex(rowData)
 
+		sch = lb.OutputSchema()
+
 		if isSrc {
 			l, err := n.GetLookup(ctx, nil)
 			if err != nil {

diff --git a/go/libraries/doltcore/sqle/kvexec/lookup_join.go b/go/libraries/doltcore/sqle/kvexec/lookup_join.go
@@ -212,7 +212,7 @@ type lookupMapping struct {
 	pool pool.BuffPool
 }
 
-func newLookupKeyMapping(ctx context.Context, sourceSch schema.Schema, src prolly.Map, tgtKeyDesc val.TupleDesc, keyExprs []sql.Expression) *lookupMapping {
+func newLookupKeyMapping(ctx context.Context, sourceSch schema.Schema, tgtKeyDesc val.TupleDesc, keyExprs []sql.Expression, ns tree.NodeStore) *lookupMapping {
 	keyless := schema.IsKeyless(sourceSch)
 	// |split| is an index into the schema separating the key and value fields
 	var split int
@@ -254,24 +254,24 @@ func newLookupKeyMapping(ctx context.Context, sourceSch schema.Schema, src proll
 	litDesc := val.NewTupleDescriptor(litTypes...)
 	litTb := val.NewTupleBuilder(litDesc)
 	for i, j := range litMappings {
-		tree.PutField(ctx, src.NodeStore(), litTb, i, keyExprs[j].(*expression.Literal).Value())
+		tree.PutField(ctx, ns, litTb, i, keyExprs[j].(*expression.Literal).Value())
 	}
 
 	var litTuple val.Tuple
 	if litDesc.Count() > 0 {
-		litTuple = litTb.Build(src.Pool())
+		litTuple = litTb.Build(ns.Pool())
 	}
 
 	return &lookupMapping{
 		split:      split,
 		srcMapping: srcMapping,
 		litTuple:   litTuple,
 		litKd:      litDesc,
-		srcKd:      src.KeyDesc(),
-		srcVd:      src.ValDesc(),
+		srcKd:      sourceSch.GetKeyDescriptor(),
+		srcVd:      sourceSch.GetValueDescriptor(),
 		targetKb:   val.NewTupleBuilder(tgtKeyDesc),
-		ns:         src.NodeStore(),
-		pool:       src.Pool(),
+		ns:         ns,
+		pool:       ns.Pool(),
 	}
 }