Add tests

tie · Mar 7, 2021 · 9339181 · 9339181
1 parent e8791ec
commit 9339181
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 5 deletions.
diff --git a/planner/optimizer.go b/planner/optimizer.go
@@ -499,7 +499,7 @@ func getCandidateFromfilterNode(f *stream.FilterOperator, tableName string, info
 	// we'll start with checking if the path is the primary key of the table
 	if pk := info.GetPrimaryKey(); pk != nil && pk.Path.IsEqual(path) {
 		cd.isPk = true
-		cd.priority = 2
+		cd.priority = 3
 
 		ranges, err := getRangesFromOp(op, e)
 		if err != nil {
@@ -514,7 +514,11 @@ func getCandidateFromfilterNode(f *stream.FilterOperator, tableName string, info
 	// if not, check if an index exists for that path
 	if idx := indexes.GetIndexByPath(document.Path(path)); idx != nil {
 		cd.isIndex = true
-		cd.priority = 1
+		if idx.Info.Unique {
+			cd.priority = 2
+		} else {
+			cd.priority = 1
+		}
 
 		ranges, err := getRangesFromOp(op, e)
 		if err != nil {

diff --git a/planner/optimizer_test.go b/planner/optimizer_test.go
@@ -361,9 +361,38 @@ func TestUseIndexBasedOnSelectionNodeRule(t *testing.T) {
 			stream.New(stream.SeqScan("foo")).Pipe(stream.Filter(parser.MustParseExpr("1 IN a"))),
 		},
 		{
-			"FROM foo WHERE a > 10",
-			stream.New(stream.SeqScan("foo")).Pipe(stream.Filter(parser.MustParseExpr("a > 10"))),
-			stream.New(stream.IndexScan("idx_foo_a", st.Range{Min: parser.MustParseExpr("10"), Exclusive: true})),
+			"FROM foo WHERE a >= 10",
+			stream.New(stream.SeqScan("foo")).Pipe(stream.Filter(parser.MustParseExpr("a >= 10"))),
+			stream.New(stream.IndexScan("idx_foo_a", st.Range{Min: parser.MustParseExpr("10")})),
+		},
+		{
+			"FROM foo WHERE k = 1",
+			stream.New(stream.SeqScan("foo")).Pipe(stream.Filter(parser.MustParseExpr("k = 1"))),
+			stream.New(stream.PkScan("foo", st.Range{Min: parser.MustParseExpr("1"), Exact: true})),
+		},
+		{
+			"FROM foo WHERE k = 1 AND b = 2",
+			stream.New(stream.SeqScan("foo")).
+				Pipe(stream.Filter(parser.MustParseExpr("k = 1"))).
+				Pipe(stream.Filter(parser.MustParseExpr("b = 2"))),
+			stream.New(stream.PkScan("foo", st.Range{Min: parser.MustParseExpr("1"), Exact: true})).
+				Pipe(stream.Filter(parser.MustParseExpr("b = 2"))),
+		},
+		{
+			"FROM foo WHERE a = 1 AND k = 2",
+			stream.New(stream.SeqScan("foo")).
+				Pipe(stream.Filter(parser.MustParseExpr("a = 1"))).
+				Pipe(stream.Filter(parser.MustParseExpr("2 = k"))),
+			stream.New(stream.PkScan("foo", st.Range{Min: parser.MustParseExpr("2"), Exact: true})).
+				Pipe(stream.Filter(parser.MustParseExpr("a = 1"))),
+		},
+		{
+			"FROM foo WHERE a = 1 AND k < 2",
+			stream.New(stream.SeqScan("foo")).
+				Pipe(stream.Filter(parser.MustParseExpr("a = 1"))).
+				Pipe(stream.Filter(parser.MustParseExpr("k < 2"))),
+			stream.New(stream.IndexScan("idx_foo_a", st.Range{Min: parser.MustParseExpr("1"), Exact: true})).
+				Pipe(stream.Filter(parser.MustParseExpr("k < 2"))),
 		},
 	}