[WIP] - force push over me.

Author: efritz

internal/execution/queries/plan/access.go

@@ -5,6 +5,7 @@ import (
 	"github.com/efritz/gostgres/internal/execution/expressions"
 	"github.com/efritz/gostgres/internal/execution/queries/nodes"
 	"github.com/efritz/gostgres/internal/execution/queries/nodes/access"
+	"github.com/efritz/gostgres/internal/execution/queries/plan/cost"
 	"github.com/efritz/gostgres/internal/shared/fields"
 	"github.com/efritz/gostgres/internal/shared/impls"
 )
@@ -59,8 +60,8 @@ func (n *logicalAccessNode) Optimize(ctx impls.OptimizationContext) {
 	n.order = nil
 }
 
-func (n *logicalAccessNode) EstimateCost() Cost {
-	return Cost{} // TODO
+func (n *logicalAccessNode) EstimateCost() impls.NodeCost {
+	return cost.ApplyFilterToCost(n.strategy.EstimateCost(), cost.EstimateFilterSelectivity(n.filter))
 }
 
 func (n *logicalAccessNode) Filter() impls.Expression {
@@ -162,6 +163,17 @@ func (s *logicalTableAccessStrategy) Ordering() impls.OrderExpression {
 	return nil
 }
 
+var tableAccessCostPerRow = impls.ResourceCost{CPU: 0.01, IO: 0.1}
+
+func (s *logicalTableAccessStrategy) EstimateCost() impls.NodeCost {
+	stats, _ := s.table.Statistics()
+
+	return impls.NodeCost{
+		EstimatedRows: stats.RowCount,
+		VariableCost:  tableAccessCostPerRow.ScaleUniform(float64(stats.RowCount)),
+	}
+}
+
 func (s *logicalTableAccessStrategy) Build() nodes.AccessStrategy {
 	return access.NewTableAccessStrategy(s.table)
 }
@@ -186,6 +198,20 @@ func (s *logicalIndexAccessStrategy[O]) Filter() impls.Expression {
 	return expressions.UnionFilters(append(expressions.Conjunctions(filterExpression), expressions.Conjunctions(condition)...)...)
 }
 
+var indexAccessCostPerRow = impls.ResourceCost{CPU: 0.01, IO: 0.1}
+
+func (s *logicalIndexAccessStrategy[O]) EstimateCost() impls.NodeCost {
+	stats, _ := s.table.Statistics()
+
+	// TODO - support parital indexes
+	// TODO - determine selectivity based on index cond
+
+	return impls.NodeCost{
+		EstimatedRows: stats.RowCount,
+		VariableCost:  indexAccessCostPerRow.ScaleUniform(float64(stats.RowCount)),
+	}
+}
+
 func (s *logicalIndexAccessStrategy[O]) Build() nodes.AccessStrategy {
 	return access.NewIndexAccessStrategy(s.table, s.index, s.opts, s.Filter())
 }

internal/execution/queries/plan/analyze.go

@@ -21,7 +21,7 @@ func (n *logicalAnalyze) Fields() []fields.Field
 func (n *logicalAnalyze) AddFilter(ctx impls.OptimizationContext, filter impls.Expression)    {} // top-level
 func (n *logicalAnalyze) AddOrder(ctx impls.OptimizationContext, order impls.OrderExpression) {} // top-level
 func (n *logicalAnalyze) Optimize(ctx impls.OptimizationContext)                              {}
-func (n *logicalAnalyze) EstimateCost() Cost                                                  { return Cost{} }
+func (n *logicalAnalyze) EstimateCost() impls.NodeCost                                        { return impls.NodeCost{} }
 func (n *logicalAnalyze) Filter() impls.Expression                                            { return nil }
 func (n *logicalAnalyze) Ordering() impls.OrderExpression                                     { return nil }
 func (n *logicalAnalyze) SupportsMarkRestore() bool                                           { return false }

internal/execution/queries/plan/cost/aggregation.go

@@ -0,0 +1,28 @@
+package cost
+
+import "github.com/efritz/gostgres/internal/shared/impls"
+
+var (
+	buildAggregateTableCostPerRow    = impls.ResourceCost{CPU: 0.1}
+	buildAggregateTableCostPerBucket = impls.ResourceCost{Memory: 0.1}
+)
+
+func ApplyAggregationToCost(innerCost impls.NodeCost, estimatedNumBuckets int) impls.NodeCost {
+	// Aggregation reads the entire inner relation on startup
+	cost := MaterializeCost(innerCost)
+
+	// Add the cost of creating aggregation buckets and hashing all rows from the inner relation
+	n := float64(innerCost.EstimatedRows)
+	cost.FixedCost = cost.FixedCost.Add(buildAggregateTableCostPerRow.ScaleUniform(n))
+	cost.FixedCost = cost.FixedCost.Add(buildAggregateTableCostPerBucket.ScaleUniform(float64(estimatedNumBuckets)))
+
+	// One row is projected for each bucket
+	cost.EstimatedRows = estimatedNumBuckets
+	cost = ApplyProjectionToCost(cost)
+
+	return cost
+}
+
+func EstimateDistinctCount(estimatedRows int) int {
+	return estimatedRows // TODO
+}

internal/execution/queries/plan/cost/filter.go

@@ -0,0 +1,26 @@
+package cost
+
+import "github.com/efritz/gostgres/internal/shared/impls"
+
+var filterEvaluationCostPerRow = impls.ResourceCost{CPU: 0.01}
+
+func ApplyFilterToCost(innerCost impls.NodeCost, filterSelectivity float64) impls.NodeCost {
+	cost := innerCost
+
+	// Evaluate a filter expression for every row
+	n := float64(innerCost.EstimatedRows)
+	cost.VariableCost = cost.VariableCost.Add(filterEvaluationCostPerRow.ScaleUniform(n))
+
+	// Only rows selected by the filter are emitted
+	cost.EstimatedRows = int(n * filterSelectivity)
+
+	return cost
+}
+
+func EstimateFilterSelectivity(filter impls.Expression) float64 {
+	if filter == nil {
+		return 1.0
+	}
+
+	return 1.0 // TODO - estimate
+}

internal/execution/queries/plan/cost/helpers.go

@@ -0,0 +1,9 @@
+package cost
+
+func coalesce(value *int, defaultValue int) int {
+	if value != nil {
+		return *value
+	}
+
+	return defaultValue
+}

internal/execution/queries/plan/cost/joins.go

@@ -0,0 +1,41 @@
+package cost
+
+import "github.com/efritz/gostgres/internal/shared/impls"
+
+var (
+	joinMergeCostPerRow  = impls.ResourceCost{CPU: 0.2}
+	joinFilterCostPerRow = impls.ResourceCost{CPU: 0.1}
+)
+
+func EstimateNestedLoopJoinCost(
+	leftCost impls.NodeCost,
+	rightCost impls.NodeCost,
+	joinSelectivity float64,
+	hasCondition bool,
+) impls.NodeCost {
+	estimatedLeftRows := float64(leftCost.EstimatedRows)
+	estimatedRightRows := float64(rightCost.EstimatedRows)
+	estimatedCandidateRows := estimatedLeftRows * estimatedRightRows
+	estimatedResultRows := estimatedCandidateRows * joinSelectivity
+
+	costPerCandidateRow := joinMergeCostPerRow
+	if hasCondition {
+		costPerCandidateRow = costPerCandidateRow.Add(joinFilterCostPerRow)
+	}
+
+	// On startup, we only initialize the left scanner
+	fixedCost := leftCost.FixedCost
+
+	variableCost := impls.SumCosts(
+		leftCost.VariableCost.ScaleUniform(estimatedLeftRows),       // Cost to scan each row from left relation
+		rightCost.FixedCost.ScaleUniform(estimatedLeftRows),         // Cost to re-initialized right scanner for every row from left relation
+		rightCost.VariableCost.ScaleUniform(estimatedCandidateRows), // Cost to scan each row from right relation
+		costPerCandidateRow.ScaleUniform(estimatedCandidateRows),    // Cost to merge row pairs and evaluate the join condition
+	)
+
+	return impls.NodeCost{
+		EstimatedRows: int(estimatedResultRows),
+		FixedCost:     fixedCost,
+		VariableCost:  variableCost,
+	}
+}

internal/execution/queries/plan/cost/limitoffset.go

@@ -0,0 +1,22 @@
+package cost
+
+import "github.com/efritz/gostgres/internal/shared/impls"
+
+func AlterCostByLimitOffset(innerCost impls.NodeCost, limit, offset *int) impls.NodeCost {
+	if limit == nil && offset == nil {
+		return innerCost
+	}
+
+	cost := innerCost
+
+	// Scale the variable cost of the inner relation by reading only limit + offset rows
+	o := coalesce(offset, 0)
+	l := coalesce(limit, innerCost.EstimatedRows-o)
+	cost.VariableCost = cost.VariableCost.ScaleUniform(float64(l+o) / float64(innerCost.EstimatedRows))
+
+	// Adjust number of output rows; this may be less than limit + offset, so we may end up
+	// "smearing" the variable cost of the inner relation over fewer rows of the outer relation.
+	cost.EstimatedRows = l
+
+	return cost
+}

internal/execution/queries/plan/cost/materialize.go

@@ -0,0 +1,15 @@
+package cost
+
+import "github.com/efritz/gostgres/internal/shared/impls"
+
+func MaterializeCost(innerCost impls.NodeCost) impls.NodeCost {
+	cost := innerCost
+
+	// We pay the variable cost of the inner relation at startup
+	cost.FixedCost = cost.FixedCost.Add(innerCost.VariableCost)
+
+	// Reset the varaible cost so we're not double counting it
+	cost.VariableCost = impls.ResourceCost{}
+
+	return cost
+}

internal/execution/queries/plan/cost/ordering.go

@@ -0,0 +1,24 @@
+package cost
+
+import (
+	"math"
+
+	"github.com/efritz/gostgres/internal/shared/impls"
+)
+
+var (
+	sortCostPerRow        = impls.ResourceCost{Memory: 0.1}
+	sortCostPerComparison = impls.ResourceCost{CPU: 0.1}
+)
+
+func ApplySortToCost(innerCost impls.NodeCost) impls.NodeCost {
+	// Sorting reads the entire inner relation on startup
+	cost := MaterializeCost(innerCost)
+
+	// Add the cost of storing each row in-memory and performing the comparison sort
+	n := float64(innerCost.EstimatedRows)
+	cost.FixedCost = cost.FixedCost.Add(sortCostPerRow.ScaleUniform(n))
+	cost.FixedCost = cost.FixedCost.Add(sortCostPerComparison.ScaleUniform(n * math.Log2(n)))
+
+	return cost
+}

internal/execution/queries/plan/cost/projection.go

@@ -0,0 +1,15 @@
+package cost
+
+import "github.com/efritz/gostgres/internal/shared/impls"
+
+var projectionEvaluationCostPerRow = impls.ResourceCost{CPU: 0.01}
+
+func ApplyProjectionToCost(innerCost impls.NodeCost) impls.NodeCost {
+	cost := innerCost
+
+	// Evaluate a projection expression for every row
+	n := float64(innerCost.EstimatedRows)
+	cost.VariableCost = cost.VariableCost.Add(projectionEvaluationCostPerRow.ScaleUniform(n))
+
+	return cost
+}