Merge pull request #1046 from dolthub/zachmu/index-pushdown

Introduced OrderedIndex interface to deal with indexes that don't return ordered results

Author: zachmu

enginetest/evaluation.go

@@ -45,7 +45,7 @@ func RunQuery(t *testing.T, e *sqle.Engine, harness Harness, query string) {
 func RunQueryWithContext(t *testing.T, e *sqle.Engine, harness Harness, ctx *sql.Context, query string) {
 	ctx = ctx.WithQuery(query)
 	sch, iter, err := e.Query(ctx, query)
-	require.NoError(t, err)
+	require.NoError(t, err, "error running query %s: %v", query, err)
 	_, err = sql.RowIterToRows(ctx, sch, iter)
 	require.NoError(t, err)
 	validateEngine(t, ctx, harness, e)

memory/index.go

@@ -38,6 +38,7 @@ type Index struct {
 
 var _ sql.Index = (*Index)(nil)
 var _ sql.FilteredIndex = (*Index)(nil)
+var _ sql.OrderedIndex = (*Index)(nil)
 
 func (idx *Index) Database() string                    { return idx.DB }
 func (idx *Index) Driver() string                      { return idx.DriverName }
@@ -232,6 +233,10 @@ func getType(val interface{}) (interface{}, sql.Type) {
 	}
 }
 
+func (idx *Index) Order() sql.IndexOrder {
+	return sql.IndexOrderAsc
+}
+
 func or(expressions ...sql.Expression) sql.Expression {
 	if len(expressions) == 1 {
 		return expressions[0]

sql/analyzer/process.go

@@ -50,9 +50,9 @@ func trackProcess(ctx *sql.Context, a *Analyzer, n sql.Node, scope *Scope, sel R
 	n, _, err := transform.Node(n, func(n sql.Node) (sql.Node, transform.TreeIdentity, error) {
 		switch n := n.(type) {
 		case *plan.IndexedTableAccess:
-			// Only add a process table if ResolvedTable implements ParallelizedIndexAddressableTable
+			// Only parallelize indexed table accesses if the underlying table supports it
 			parallelizedTable, ok := n.ResolvedTable.Table.(sql.ParallelizedIndexAddressableTable)
-			if !ok {
+			if !ok || !parallelizedTable.ShouldParallelizeAccess() {
 				return n, transform.SameTree, nil
 			}
 
@@ -83,11 +83,7 @@ func trackProcess(ctx *sql.Context, a *Analyzer, n sql.Node, scope *Scope, sel R
 				}
 			}
 
-			// Wrap with ProcessTable
-			t := plan.NewProcessTable(parallelizedTable, onPartitionDone, onPartitionStart, onRowNext)
-
-			// Replace child
-			n, err := n.WithTable(t)
+			n, err := n.WithTable(plan.NewProcessTable(parallelizedTable, onPartitionDone, onPartitionStart, onRowNext))
 			return n, transform.NewTree, err
 		case *plan.ResolvedTable:
 			switch n.Table.(type) {

sql/analyzer/pushdown.go

@@ -680,21 +680,10 @@ func pushdownIndexesToTable(a *Analyzer, tableNode NameableNode, indexes map[str
 			if table == nil {
 				return n, transform.SameTree, nil
 			}
-			if indexAddressableTable, ok := table.(sql.IndexAddressableTable); ok {
+			if _, ok := table.(sql.IndexAddressableTable); ok {
 				indexLookup, ok := indexes[tableNode.Name()]
 				if ok {
 					a.Log("table %q transformed with pushdown of index", tableNode.Name())
-
-					// Only pass lookup into resolved table if it's Parallelizable
-					if _, ok := indexAddressableTable.(sql.ParallelizedIndexAddressableTable); ok {
-						indexedTable := indexAddressableTable.WithIndexLookup(indexLookup.lookup)
-						newResolvedTable, err := n.WithTable(indexedTable)
-						if err != nil {
-							return nil, transform.SameTree, err
-						}
-						return plan.NewStaticIndexedTableAccess(newResolvedTable, indexLookup.lookup, indexLookup.indexes[0], indexLookup.fields), transform.NewTree, nil
-
-					}
 					return plan.NewStaticIndexedTableAccess(n, indexLookup.lookup, indexLookup.indexes[0], indexLookup.fields), transform.NewTree, nil
 				}
 			}
@@ -951,6 +940,10 @@ func replacePkSort(ctx *sql.Context, a *Analyzer, n sql.Node, scope *Scope, sel
 		// Extract primary index
 		var pkIndex sql.Index
 		for _, idx := range idxs {
+			oi, ok := idx.(sql.OrderedIndex)
+			if !ok || oi.Order() != sql.IndexOrderAsc {
+				continue
+			}
 			if idx.ID() == "PRIMARY" {
 				pkIndex = idx
 				break

sql/index.go

@@ -51,13 +51,31 @@ type Index interface {
 	ColumnExpressionTypes(ctx *Context) []ColumnExpressionType
 }
 
+// FilteredIndex is an extension of |Index| that allows an index to declare certain filter predicates handled,
+// allowing them to be removed from the overall plan for greater execution efficiency
 type FilteredIndex interface {
 	Index
 	// HandledFilters returns a subset of |filters| that are satisfied
 	// by index lookups to this index.
 	HandledFilters(filters []Expression) (handled []Expression)
 }
 
+type IndexOrder byte
+
+const (
+	IndexOrderNone IndexOrder = iota
+	IndexOrderAsc
+	IndexOrderDesc
+)
+
+// OrderedIndex is an extension of |Index| that allows indexes to declare their return order. The query engine can
+// optimize certain queries if the order of an index is guaranteed, e.g. removing a sort operation.
+type OrderedIndex interface {
+	Index
+	// Order returns the order of results for reads from this index
+	Order() IndexOrder
+}
+
 // IndexLookup is the implementation-specific definition of an index lookup. The IndexLookup must contain all necessary
 // information to retrieve exactly the rows in the table as specified by the ranges given to their parent index.
 // Implementors are responsible for all semantics of correctly returning rows that match an index lookup.

sql/plan/exchange.go

@@ -173,7 +173,7 @@ func (e *Exchange) RowIter2(ctx *sql.Context, f *sql.RowFrame) (sql.RowIter2, er
 
 func (e *Exchange) String() string {
 	p := sql.NewTreePrinter()
-	_ = p.WriteNode("Exchange(parallelism=%d)", e.Parallelism)
+	_ = p.WriteNode("Exchange")
 	_ = p.WriteChildren(e.Child.String())
 	return p.String()
 }

sql/plan/indexed_table_access.go

@@ -305,16 +305,27 @@ func (i *IndexedTableAccess) Partitions(ctx *sql.Context) (sql.PartitionIter, er
 		return i.ResolvedTable.Partitions(ctx)
 	}
 
+	table := i.baseTable()
+	return table.Partitions(ctx)
+}
+
+// baseTable returns the underlying sql.Table with any static index lookup applied
+func (i *IndexedTableAccess) baseTable() sql.Table {
 	table := i.ResolvedTable.Table
-	if indexAddressableTable, ok := i.ResolvedTable.Table.(sql.IndexAddressable); ok {
+	// This won't work if we add another layer of wrapping on top
+	if tw, ok := table.(sql.TableWrapper); ok {
+		table = tw.Underlying()
+	}
+
+	if indexAddressableTable, ok := table.(sql.IndexAddressable); ok {
 		table = indexAddressableTable.WithIndexLookup(i.lookup)
 	}
-	return table.Partitions(ctx)
+	return table
 }
 
 // PartitionRows implements sql.Table
 func (i *IndexedTableAccess) PartitionRows(ctx *sql.Context, partition sql.Partition) (sql.RowIter, error) {
-	return i.ResolvedTable.PartitionRows(ctx, partition)
+	return i.baseTable().PartitionRows(ctx, partition)
 }
 
 // GetIndexLookup returns the sql.IndexLookup from an IndexedTableAccess.