special case point range lookups

go/libraries/doltcore/sqle/enginetest/dolt_engine_test.go

@@ -39,19 +39,22 @@ func TestQueries(t *testing.T) {
 }
 
 func TestSingleQuery(t *testing.T) {
-	t.Skip()
-
 	var test enginetest.QueryTest
 	test = enginetest.QueryTest{
-		Query:    `SELECT * from mytable`,
-		Expected: []sql.Row{},
+		Query: `SELECT a.pk1, a.pk2, b.pk1, b.pk2
+				FROM two_pk a JOIN two_pk b
+				ON a.pk1+1=b.pk1 AND a.pk2+1=b.pk2
+				ORDER BY 1,2,3`,
+		Expected: []sql.Row{
+			{0, 0, 1, 1},
+		},
 	}
 
 	harness := newDoltHarness(t)
 	engine := enginetest.NewEngine(t, harness)
 	enginetest.CreateIndexes(t, harness, engine)
-	engine.Analyzer.Debug = true
-	engine.Analyzer.Verbose = true
+	//engine.Analyzer.Debug = true
+	//engine.Analyzer.Verbose = true
 
 	enginetest.TestQuery(t, harness, engine, test.Query, test.Expected, test.ExpectedColumns, test.Bindings)
 }

go/store/prolly/map.go

@@ -170,7 +170,28 @@ func (m Map) IterAll(ctx context.Context) (MapRangeIter, error) {
 
 // IterRange returns a MutableMapRangeIter that iterates over a Range.
 func (m Map) IterRange(ctx context.Context, rng Range) (MapRangeIter, error) {
-	return m.iterFromRange(ctx, rng)
+	if rng.isPointLookup(m.keyDesc) {
+		return m.pointLookupFromRange(ctx, rng)
+	} else {
+		return m.iterFromRange(ctx, rng)
+	}
+}
+
+func (m Map) pointLookupFromRange(ctx context.Context, rng Range) (*pointLookup, error) {
+	search := pointLookupSearchFn(rng)
+	cur, err := newCursorFromSearchFn(ctx, m.ns, m.root, search)
+	if err != nil {
+		return nil, err
+	}
+
+	key := val.Tuple(cur.currentKey())
+	value := val.Tuple(cur.currentValue())
+	if compareBound(rng.Start, key, m.keyDesc) != 0 {
+		// map does not contain this point lookup
+		key, value = nil, nil
+	}
+
+	return &pointLookup{k: key, v: value}, nil
 }
 
 func (m Map) iterFromRange(ctx context.Context, rng Range) (*prollyRangeIter, error) {

go/store/prolly/map_test.go

@@ -121,6 +121,8 @@ func prollyMapFromTuples(t *testing.T, kd, vd val.TupleDesc, tuples [][2]val.Tup
 
 func testGet(t *testing.T, om orderedMap, tuples [][2]val.Tuple) {
 	ctx := context.Background()
+
+	// test get
 	for _, kv := range tuples {
 		err := om.Get(ctx, kv[0], func(key, val val.Tuple) (err error) {
 			assert.NotNil(t, kv[0])
@@ -130,6 +132,27 @@ func testGet(t *testing.T, om orderedMap, tuples [][2]val.Tuple) {
 		})
 		require.NoError(t, err)
 	}
+
+	desc := keyDescFromMap(om)
+
+	// test point lookup
+	for _, kv := range tuples {
+		rng := pointRangeFromTuple(kv[0], desc)
+		require.True(t, rng.isPointLookup(desc))
+
+		iter, err := om.IterRange(ctx, rng)
+		require.NoError(t, err)
+
+		k, v, err := iter.Next(ctx)
+		require.NoError(t, err)
+		assert.Equal(t, kv[0], k)
+		assert.Equal(t, kv[1], v)
+
+		k, v, err = iter.Next(ctx)
+		assert.Error(t, err, io.EOF)
+		assert.Nil(t, k)
+		assert.Nil(t, v)
+	}
 }
 
 func testHas(t *testing.T, om Map, tuples [][2]val.Tuple) {
@@ -166,3 +189,19 @@ func testIterAll(t *testing.T, om orderedMap, tuples [][2]val.Tuple) {
 		assert.Equal(t, tuples[i][1], kv[1])
 	}
 }
+
+func pointRangeFromTuple(tup val.Tuple, desc val.TupleDesc) Range {
+	start := make([]RangeCut, len(desc.Types))
+	stop := make([]RangeCut, len(desc.Types))
+	for i := range start {
+		start[i].Value = tup.GetField(i)
+		start[i].Inclusive = true
+	}
+	copy(stop, start)
+
+	return Range{
+		Start: start,
+		Stop:  stop,
+		Desc:  desc,
+	}
+}

go/store/prolly/range.go

@@ -183,6 +183,26 @@ func (r Range) merge(other Range) Range {
 	}
 }
 
+func (r Range) isPointLookup(desc val.TupleDesc) bool {
+	if len(r.Start) < len(desc.Types) || len(r.Stop) < len(desc.Types) {
+		return false
+	}
+	for i := range r.Start {
+		if !r.Start[i].Inclusive || !r.Stop[i].Inclusive {
+			return false
+		}
+	}
+
+	compare := desc.Comparator()
+	for i, typ := range desc.Types {
+		lo, hi := r.Start[i].Value, r.Stop[i].Value
+		if compare.CompareValues(lo, hi, typ) != 0 {
+			return false
+		}
+	}
+	return true
+}
+
 func (r Range) format() string {
 	return formatRange(r)
 }
@@ -215,15 +235,17 @@ func (c RangeCut) lesserValue(other RangeCut, typ val.Type, tc val.TupleComparat
 	return cmp < 0
 }
 
+func compareBound(bound []RangeCut, tup val.Tuple, desc val.TupleDesc) int {
+	for i, cut := range bound {
+		cmp := desc.CompareField(cut.Value, i, tup)
+		if cmp != 0 {
+			return cmp
+		}
+	}
+	return 0
+}
+
 func rangeStartSearchFn(rng Range) searchFn {
-	return binarySearchRangeStart(rng)
-}
-
-func rangeStopSearchFn(rng Range) searchFn {
-	return binarySearchRangeStop(rng)
-}
-
-func binarySearchRangeStart(rng Range) searchFn {
 	return func(nd Node) int {
 		// todo(andy): inline sort.Search()
 		return sort.Search(int(nd.count), func(i int) (in bool) {
@@ -235,7 +257,7 @@ func binarySearchRangeStart(rng Range) searchFn {
 	}
 }
 
-func binarySearchRangeStop(rng Range) searchFn {
+func rangeStopSearchFn(rng Range) searchFn {
 	return func(nd Node) (idx int) {
 		// todo(andy): inline sort.Search()
 		return sort.Search(int(nd.count), func(i int) (out bool) {
@@ -247,6 +269,17 @@ func binarySearchRangeStop(rng Range) searchFn {
 	}
 }
 
+func pointLookupSearchFn(rng Range) searchFn {
+	return func(nd Node) (idx int) {
+		// todo(andy): inline sort.Search()
+		return sort.Search(int(nd.count), func(i int) (out bool) {
+			tup := val.Tuple(nd.getKey(i))
+			// |rng.Start| <= |tup|
+			return compareBound(rng.Start, tup, rng.Desc) <= 0
+		})
+	}
+}
+
 // GreaterRange defines a Range of Tuples greater than |start|.
 func GreaterRange(start val.Tuple, desc val.TupleDesc) Range {
 	return Range{

go/store/prolly/range_iter.go

@@ -26,8 +26,23 @@ type MapRangeIter interface {
 }
 
 var _ MapRangeIter = &prollyRangeIter{}
+var _ MapRangeIter = &pointLookup{}
 var _ MapRangeIter = &MutableMapRangeIter{}
 
+type pointLookup struct {
+	k, v val.Tuple
+}
+
+func (p *pointLookup) Next(context.Context) (key, value val.Tuple, err error) {
+	if p.k == nil || p.v == nil {
+		err = io.EOF
+	} else {
+		key, value = p.k, p.v
+		p.k, p.v = nil, nil
+	}
+	return
+}
+
 type rangeIter interface {
 	iterate(ctx context.Context) error
 	current() (key, value val.Tuple)