Improved cost estimation for IVFFlat [skip ci]

src/hnsw.c

@@ -100,8 +100,10 @@ hnswcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
 {
 	GenericCosts costs;
 	int			m;
-	double		ratio;
-	double		startupPages;
+	int			entryLevel;
+	int			layer0TuplesMax;
+	double		layer0Selectivity;
+	double		scalingFactor = 0.55;
 	double		spc_seq_page_cost;
 	Relation	index;
 
@@ -118,8 +120,6 @@ hnswcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
 
 	MemSet(&costs, 0, sizeof(costs));
 
-	genericcostestimate(root, path, loop_count, &costs);
-
 	index = index_open(path->indexinfo->indexoid, NoLock);
 	HnswGetMetaPageInfo(index, &m, NULL);
 	index_close(index, NoLock);
@@ -151,38 +151,30 @@ hnswcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
 	 * at L0, accounting for previously visited tuples, multiplied by the
 	 * "scalingFactor" (currently hardcoded).
 	 */
-	if (path->indexinfo->tuples > 0)
-	{
-		double		scalingFactor = 0.55;
-		int			entryLevel = (int) (log(path->indexinfo->tuples) * HnswGetMl(m));
-		int			layer0TuplesMax = HnswGetLayerM(m, 0) * hnsw_ef_search;
-		double		layer0Selectivity = scalingFactor * log(path->indexinfo->tuples) / (log(m) * (1 + log(hnsw_ef_search)));
+	entryLevel = (int) (log(path->indexinfo->tuples + 1) * HnswGetMl(m));
+	layer0TuplesMax = HnswGetLayerM(m, 0) * hnsw_ef_search;
+	layer0Selectivity = (scalingFactor * log(path->indexinfo->tuples + 1)) /
+		(log(m) * (1 + log(hnsw_ef_search)));
 
-		ratio = (entryLevel * m + layer0TuplesMax * layer0Selectivity) / path->indexinfo->tuples;
+	costs.numIndexTuples = (entryLevel * m) +
+		(layer0TuplesMax * layer0Selectivity);
 
-		if (ratio > 1)
-			ratio = 1;
-	}
-	else
-		ratio = 1;
+	genericcostestimate(root, path, loop_count, &costs);
 
 	get_tablespace_page_costs(path->indexinfo->reltablespace, NULL, &spc_seq_page_cost);
 
-	/* Startup cost is cost before returning the first row */
-	costs.indexStartupCost = costs.indexTotalCost * ratio;
-
 	/* Adjust cost if needed since TOAST not included in seq scan cost */
-	startupPages = costs.numIndexPages * ratio;
-	if (startupPages > path->indexinfo->rel->pages && ratio < 0.5)
+	if (costs.numIndexPages > path->indexinfo->rel->pages)
 	{
 		/* Change all page cost from random to sequential */
-		costs.indexStartupCost -= startupPages * (costs.spc_random_page_cost - spc_seq_page_cost);
+		costs.indexTotalCost -= costs.numIndexPages * (costs.spc_random_page_cost - spc_seq_page_cost);
 
 		/* Remove cost of extra pages */
-		costs.indexStartupCost -= (startupPages - path->indexinfo->rel->pages) * spc_seq_page_cost;
+		costs.indexTotalCost -= (costs.numIndexPages - path->indexinfo->rel->pages) * spc_seq_page_cost;
 	}
 
-	*indexStartupCost = costs.indexStartupCost;
+	/* Use total cost since most work happens before first tuple is returned */
+	*indexStartupCost = costs.indexTotalCost;
 	*indexTotalCost = costs.indexTotalCost;
 	*indexSelectivity = costs.indexSelectivity;
 	*indexCorrelation = costs.indexCorrelation;

src/ivfflat.c

@@ -69,8 +69,6 @@ ivfflatcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
 	GenericCosts costs;
 	int			lists;
 	double		ratio;
-	double		sequentialRatio = 0.5;
-	double		startupPages;
 	double		spc_seq_page_cost;
 	Relation	index;
 
@@ -98,23 +96,25 @@ ivfflatcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
 	if (ratio > 1.0)
 		ratio = 1.0;
 
+	/* Set startup cost since most work happens before first tuple is returned */
+	costs.indexStartupCost = costs.indexTotalCost * ratio;
+	costs.numIndexPages *= ratio;
+
 	get_tablespace_page_costs(path->indexinfo->reltablespace, NULL, &spc_seq_page_cost);
 
-	/* Change some page cost from random to sequential */
-	costs.indexTotalCost -= sequentialRatio * costs.numIndexPages * (costs.spc_random_page_cost - spc_seq_page_cost);
-
-	/* Startup cost is cost before returning the first row */
-	costs.indexStartupCost = costs.indexTotalCost * ratio;
-
 	/* Adjust cost if needed since TOAST not included in seq scan cost */
-	startupPages = costs.numIndexPages * ratio;
-	if (startupPages > path->indexinfo->rel->pages && ratio < 0.5)
+	if (costs.numIndexPages > path->indexinfo->rel->pages && ratio < 0.5)
 	{
-		/* Change rest of page cost from random to sequential */
-		costs.indexStartupCost -= (1 - sequentialRatio) * startupPages * (costs.spc_random_page_cost - spc_seq_page_cost);
+		/* Change all page cost from random to sequential */
+		costs.indexStartupCost -= costs.numIndexPages * (costs.spc_random_page_cost - spc_seq_page_cost);
 
 		/* Remove cost of extra pages */
-		costs.indexStartupCost -= (startupPages - path->indexinfo->rel->pages) * spc_seq_page_cost;
+		costs.indexStartupCost -= (costs.numIndexPages - path->indexinfo->rel->pages) * spc_seq_page_cost;
+	}
+	else
+	{
+		/* Change some page cost from random to sequential */
+		costs.indexStartupCost -= 0.5 * costs.numIndexPages * (costs.spc_random_page_cost - spc_seq_page_cost);
 	}
 
 	*indexStartupCost = costs.indexStartupCost;
@@ -122,6 +122,9 @@ ivfflatcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
 	*indexSelectivity = costs.indexSelectivity;
 	*indexCorrelation = costs.indexCorrelation;
 	*indexPages = costs.numIndexPages;
+
+	Assert(*indexStartupCost > 0);
+	Assert(*indexTotalCost > 0);
 }
 
 /*

test/t/017_hnsw_filtering.pl

@@ -41,7 +41,8 @@ my $c = int(rand() * $nc);
 my $explain = $node->safe_psql("postgres", qq(
 	EXPLAIN ANALYZE SELECT i FROM tst WHERE c = $c ORDER BY v <-> '$query' LIMIT $limit;
 ));
-like($explain, qr/Seq Scan/);
+# TODO Do not use index
+like($explain, qr/Index Scan using idx/);
 
 # Test attribute filtering with few rows removed
 $explain = $node->safe_psql("postgres", qq(
@@ -59,7 +60,8 @@ like($explain, qr/Index Scan using idx/);
 $explain = $node->safe_psql("postgres", qq(
 	EXPLAIN ANALYZE SELECT i FROM tst WHERE c < 1 ORDER BY v <-> '$query' LIMIT $limit;
 ));
-like($explain, qr/Seq Scan/);
+# TODO Do not use index
+like($explain, qr/Index Scan using idx/);
 
 # Test attribute filtering with few rows removed like
 $explain = $node->safe_psql("postgres", qq(

test/t/039_hnsw_cost.pl

@@ -17,11 +17,12 @@ $node->safe_psql("postgres", "CREATE EXTENSION vector;");
 for my $dim (@dims)
 {
 	my $array_sql = join(",", ('random()') x $dim);
+	my $n = $dim == 384 ? 3000 : 1000;
 
 	# Create table and index
 	$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim));");
 	$node->safe_psql("postgres",
-		"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 2000) i;"
+		"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, $n) i;"
 	);
 	$node->safe_psql("postgres", "CREATE INDEX idx ON tst USING hnsw (v vector_l2_ops);");
 	$node->safe_psql("postgres", "ANALYZE tst;");
@@ -39,16 +40,6 @@ for my $dim (@dims)
 	));
 	like($explain, qr/Index Scan using idx/);
 
-	# 3x the rows are needed for distance filters
-	# since the planner uses DEFAULT_INEQ_SEL for the selectivity (should be 1)
-	# Recreate index for performance
-	$node->safe_psql("postgres", "DROP INDEX idx;");
-	$node->safe_psql("postgres",
-		"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(2001, 6000) i;"
-	);
-	$node->safe_psql("postgres", "CREATE INDEX idx ON tst USING hnsw (v vector_l2_ops);");
-	$node->safe_psql("postgres", "ANALYZE tst;");
-
 	$explain = $node->safe_psql("postgres", qq(
 		EXPLAIN ANALYZE SELECT i FROM tst WHERE v <-> '$query' < 1 ORDER BY v <-> '$query' LIMIT $limit;
 	));

test/t/040_ivfflat_cost.pl

@@ -21,7 +21,7 @@ for my $dim (@dims)
 	# Create table and index
 	$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim));");
 	$node->safe_psql("postgres",
-		"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 5000) i;"
+		"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 6000) i;"
 	);
 	$node->safe_psql("postgres", "CREATE INDEX idx ON tst USING ivfflat (v vector_l2_ops) WITH (lists = 5);");
 	$node->safe_psql("postgres", "ANALYZE tst;");